From 82d4b996f7898124c22831dd1926f97749af0bc3 Mon Sep 17 00:00:00 2001 From: A-transformer Date: Fri, 7 Mar 2025 19:37:55 +0400 Subject: [PATCH] Add TMA optimization hint for large FFMA segments This check is added to identify segments with a high number of FFMA (Fused Multiply-Add) instructions, which may benefit from Tensor Memory Accelerator (TMA) optimizations on NVIDIA Hopper architecture. TMA enables asynchronous data transfers from global memory to shared memory, reducing register pressure and improving data access efficiency. Segments with a large number of FFMA instructions often involve significant data movement, making them prime candidates for TMA to accelerate memory operations and enhance overall GEMM performance. --- deep_gemm/jit/interleave_ffma.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/deep_gemm/jit/interleave_ffma.py b/deep_gemm/jit/interleave_ffma.py index fcb377e..95ac13c 100644 --- a/deep_gemm/jit/interleave_ffma.py +++ b/deep_gemm/jit/interleave_ffma.py @@ -76,6 +76,11 @@ def modify_segment(m, name, ffma_lines): num_lines = (len(ffma_lines) * 9 // 16) // 2 * 2 assert num_lines % 2 == 0 + tma_threshold = 32 + if num_lines // 2 > tma_threshold: + if os.getenv('DG_PRINT_REG_REUSE', None): + print(f' > segment `{name}` may benefit from TMA optimization due to large FFMA count ({num_lines // 2})') + le_bytes, new_le_bytes = [], [] reused_list = [] dst_reg_set = set()