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

@@ -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()