[AMD] Optimize address increments for buffer loads in loops (#8464)

This PR transfers address computation from offsets in buffer loads to
base pointers, which reuses amount of required computations and lowers
register pressure.

---------

Co-authored-by: Alexander Efimov <efimov.alexander@gmail.com>

Author: alefimov-amd

bin/RegisterTritonDialects.h

@@ -122,6 +122,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::registerTritonAMDGPUCanonicalizePointers();
   mlir::registerTritonAMDGPUConvertToBufferOps();
   mlir::registerTritonAMDGPUConvertToTensorOps();
+  mlir::registerTritonAMDGPUOptimizeBufferOpPtr();
   mlir::registerTritonAMDGPUInThreadTranspose();
   mlir::registerTritonAMDGPUCoalesceAsyncCopy();
   mlir::registerTritonAMDGPUUpdateAsyncWaitCount();

python/src/passes.h

@@ -1,6 +1,11 @@
 #define ADD_PASS_WRAPPER_0(name, builder)                                      \
   m.def(name, [](mlir::PassManager &pm) { pm.addPass(builder()); })
 
+#define ADD_FUNC_PASS_WRAPPER_0(name, builder)                                 \
+  m.def(name, [](mlir::PassManager &pm) {                                      \
+    pm.addNestedPass<mlir::triton::FuncOp>(builder());                         \
+  });
+
 #define ADD_PASS_WRAPPER_1(name, builder, ty0)                                 \
   m.def(name,                                                                  \
         [](mlir::PassManager &pm, ty0 val0) { pm.addPass(builder(val0)); })

test/TritonGPU/amd/amd-optimize-buffer-ops-base-ptr-increment.mlir

@@ -281,6 +281,7 @@ def make_ttgir(mod, metadata, options):
                 knobs.amd.use_buffer_atomics,
                 knobs.amd.buffer_ops_analyze_small_tensor_range,
             )
+            amd.passes.ttgpuir.add_optimize_buffer_op_ptr(pm)
 
         amd.passes.ttgpuir.add_fold_true_cmpi(pm)
         amd.passes.ttgpuir.add_prepare_if_combining(pm)

third_party/amd/backend/compiler.py

@@ -186,6 +186,14 @@ def TritonAMDGPUConvertToBufferOps : Pass<"tritonamdgpu-convert-buffer-ops", "ml
   ];
 }
 
+def TritonAMDGPUOptimizeBufferOpPtr : Pass<"tritonamdgpu-optimize-buffer-op-ptr", "mlir::triton::FuncOp"> {
+  let summary = "Optimize address operands of buffer operations";
+
+  let description = "This pass optimizes address computation for buffer operations";
+
+  let dependentDialects = ["mlir::triton::amdgpu::TritonAMDGPUDialect"];
+}
+
 def TritonAMDGPUBlockPingpong: Pass<"tritonamdgpu-block-pingpong", "mlir::ModuleOp"> {
   let summary = "Interleaving instructions from two warps on the same SIMD to better utilize matrix core";
 

third_party/amd/include/TritonAMDGPUTransforms/Passes.td

@@ -5,6 +5,7 @@ add_triton_library(TritonAMDGPUTransforms
   CoalesceAsyncCopy.cpp
   ConvertToBufferOps.cpp
   ConvertToTensorOps.cpp
+  OptimizeBufferOpPtr.cpp
   OptimizeEpilogue.cpp
   OptimizeDotOperands.cpp
   HoistLayoutConversions.cpp