[Backend] Refactor MMAv5 lowering and put mma_scaled in an if (triton-lang#6478)

This PR refactors the lowering of MMAv5 to share more of the code
between tc_gen5_mma and tc_gen5_mma_scaled, while also applying the same
optimization that tc_gen5_mma has that places the `tcgen05.mma`
instructions in an if block. This was previously checked to improve
performance.

Author: Mogball

test/Conversion/tritongpu_to_llvm_blackwell.mlir

@@ -197,9 +197,12 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shar
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // CHECK-LABEL: @tc_gen5_mma_block_scale
   // CHECK-SAME: (%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[USE_ACC:.+]]: i1, %{{.*}}: i1, %{{.*}})
-  // CHECK-DAG: %[[TMEM_BASE:.+]] = llvm.ptrtoint %{{.*}} : !llvm.ptr<3> to i32
-  // CHECK-DAG: %[[C0:.+]] = llvm.mlir.constant(0 : i32) : i32
-  // CHECK-DAG: %[[C32:.+]] = llvm.mlir.constant(32 : i32) : i32
+  // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %{{.*}} : !llvm.ptr<3> to i32
+  // CHECK: %[[WID:.+]] = nvgpu.warp_id
+  // CHECK: %[[C0:.+]] = llvm.mlir.constant(0 : i32) : i32
+  // CHECK: %[[P0:.+]] = llvm.icmp "eq" %[[WID]], %[[C0]] : i32
+  // CHECK: %[[P1:.+]] = llvm.and %{{.*}}, %[[P0]]  : i1
+  // CHECK: llvm.cond_br %[[P1]]
   // CHECK: %[[T0:.+]] = llvm.add %[[TMEM_BASE]], %[[C0]] : i32
   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144708608 : i32) : i32
   // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[T0]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %[[USE_ACC]]

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAHelpers.h

@@ -28,7 +28,7 @@ class DotOpMmaMemLoader {
 public:
   virtual ~DotOpMmaMemLoader() = default;
   virtual Value memLoad(int a, int b, ConversionPatternRewriter &rewriter,
-                        Location loc) = 0;
+                        Location loc) const = 0;
 };
 
 // Helper class to load shared memory slices following MMAv3 layout.
@@ -44,10 +44,10 @@ class DotOpMmaV3SmemLoader : public DotOpMmaMemLoader {
   // Return a descriptor pointing to the shared memory slice at coordinates (a,
   // b)
   Value smemLoad(int a, int b, ConversionPatternRewriter &rewriter,
-                 Location loc);
+                 Location loc) const;
 
   Value memLoad(int a, int b, ConversionPatternRewriter &rewriter,
-                Location loc) override {
+                Location loc) const override {
     return smemLoad(a, b, rewriter, loc);
   }
 
@@ -74,10 +74,10 @@ class DotOpMmaV5TmemLoader : public DotOpMmaMemLoader {
                        SmallVector<unsigned int> instrShape, bool interleaved,
                        bool trans);
   Value tmemLoad(int a, int b, ConversionPatternRewriter &rewriter,
-                 Location loc);
+                 Location loc) const;
 
   Value memLoad(int a, int b, ConversionPatternRewriter &rewriter,
-                Location loc) override {
+                Location loc) const override {
     return tmemLoad(a, b, rewriter, loc);
   }