[NVIDIA] Replace some NVGPU ops with equivalent NVVM ops (#7420)

This change updates the lowering of WGMMAFenceOp, WGMMACommitGroupOp and
ClusterWaitOp to generate NVVM dialect operations instead of inline
assembly strings. The NVVM ops will be lowered to LLVM intrinsics in
subsequent passes, providing better optimization opportunities.
Additionally, unused constant kFenceMbarrierInitOp is cleaned up.

Author: Pecco-314

test/Conversion/nvgpu_to_llvm.mlir

@@ -2,25 +2,11 @@
 
 // CHECK-LABEL: @nvvm_syncs
 llvm.func @nvvm_syncs() {
-  // CHECK: wgmma.fence.sync.aligned;
-  nvgpu.wgmma_fence
-
-  // CHECK: wgmma.commit_group.sync.aligned;
-  nvgpu.wgmma_commit_group
-
-  // CHECK: barrier.cluster.wait.aligned;
-  nvgpu.cluster_wait
-
   // CHECK: fence.proxy.async.shared::cta;
   nvgpu.fence_async_shared {bCluster = false}
   // CHECK: fence.proxy.async.shared::cluster;
   nvgpu.fence_async_shared {bCluster = true}
 
-  // CHECK: barrier.cluster.arrive.aligned;
-  nvgpu.cluster_arrive {relaxed = false}
-  // CHECK: barrier.cluster.arrive.relaxed.aligned;
-  nvgpu.cluster_arrive {relaxed = true}
-
   llvm.return
 }
 

third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td

@@ -63,14 +63,6 @@ def NVGPU_MemSyncScopeAttr : I32EnumAttr<
 class NVGPU_Op<string mnemonic, list<Trait> traits = []> :
     LLVM_OpBase<NVGPU_Dialect, mnemonic, traits>;
 
-def NVGPU_WGMMAFenceOp : NVGPU_Op<"wgmma_fence", []> {
-  let assemblyFormat = "attr-dict";
-}
-
-def NVGPU_WGMMACommitGroupOp : NVGPU_Op<"wgmma_commit_group", []> {
-  let assemblyFormat = "attr-dict";
-}
-
 def NVGPU_WGMMAWaitGroupOp : NVGPU_Op<"wgmma_wait_group", [DeclareOpInterfaceMethods<InferTypeOpInterface>,
                                                            AllTypesMatch<["input", "output"]>]> {
   let arguments = (ins LLVM_AnyStruct:$input, I32Attr:$pendings);
@@ -118,16 +110,6 @@ def NVGPU_FenceAsyncSharedOp : NVGPU_Op<"fence_async_shared", []> {
   let assemblyFormat = "attr-dict";
 }
 
-def NVGPU_ClusterArriveOp : NVGPU_Op<"cluster_arrive", []> {
-  let arguments = (ins I1Attr:$relaxed);
-
-  let assemblyFormat = "attr-dict";
-}
-
-def NVGPU_ClusterWaitOp : NVGPU_Op<"cluster_wait", []> {
-  let assemblyFormat = "attr-dict";
-}
-
 def NVGPU_StoreMatrixOp : NVGPU_Op<"stmatrix", [MemoryEffects<[MemWrite]>]> {
   let arguments = (
     ins LLVM_PointerShared:$addr,

third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp

@@ -23,10 +23,6 @@ namespace triton {
 
 namespace {
 
-const std::string kWgmmaFenceOp = "wgmma.fence.sync.aligned;";
-const std::string kWgmmaCommitGroupOp = "wgmma.commit_group.sync.aligned;";
-const std::string kClusterWaitOp = "barrier.cluster.wait.aligned;";
-const std::string kFenceMbarrierInitOp = "fence.mbarrier_init.release.cluster;";
 const std::string kClusterCtaIdOp = "{\n"
                                     ".reg .u32 a<5>;              \n"
                                     "mov.u32 a0, %cluster_ctaid.x;\n"  // x
@@ -255,19 +251,6 @@ class WarpIdOpPattern : public OpRewritePattern<ttn::WarpIdOp> {
   }
 };
 
-class ClusterArriveOpPattern : public OpRewritePattern<ttn::ClusterArriveOp> {
-public:
-  using OpRewritePattern<ttn::ClusterArriveOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(ttn::ClusterArriveOp op,
-                                PatternRewriter &rewriter) const override {
-    std::string ptxAsm = op.getRelaxed()
-                             ? "barrier.cluster.arrive.relaxed.aligned;"
-                             : "barrier.cluster.arrive.aligned;";
-    return rewriteAsPtxAsm(op, rewriter, std::move(ptxAsm));
-  }
-};
-
 // Base class for Matrix Operation Patterns
 template <typename MatrixOpType, typename ConcreteMatrixOpPattern>
 class MatrixOpPattern : public OpRewritePattern<MatrixOpType> {
@@ -788,21 +771,12 @@ class ConvertNVGPUToLLVM
     ModuleOp mod = getOperation();
     RewritePatternSet patterns(context);
 
-#define POPULATE_NVGPU_OP(SRC_OP, ASM)                                         \
-  patterns.add<NVGPUOpGenericPattern<SRC_OP>>(context, ASM, Constraints(),     \
-                                              Constraints());
-    POPULATE_NVGPU_OP(ttn::WGMMAFenceOp, kWgmmaFenceOp)
-    POPULATE_NVGPU_OP(ttn::WGMMACommitGroupOp, kWgmmaCommitGroupOp)
-    POPULATE_NVGPU_OP(ttn::ClusterWaitOp, kClusterWaitOp)
-#undef POPULATE_NVGPU_OP
     patterns.add<NVGPUOpGenericPattern<ttn::ClusterCTAIdOp>>(
         context, kClusterCtaIdOp, Constraints({"=r"}), Constraints());
 
-    patterns
-        .add<FenceAsyncSharedOpPattern, LoadMatrixOpPattern,
-             StoreMatrixOpPattern, ClusterArriveOpPattern, WGMMAOpPattern,
-             LoadAcquireOpPattern, WGMMAWaitGroupOpPattern, WarpIdOpPattern>(
-            context);
+    patterns.add<FenceAsyncSharedOpPattern, LoadMatrixOpPattern,
+                 StoreMatrixOpPattern, WGMMAOpPattern, LoadAcquireOpPattern,
+                 WGMMAWaitGroupOpPattern, WarpIdOpPattern>(context);
 
     if (applyPatternsGreedily(mod, std::move(patterns)).failed())
       signalPassFailure();

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ClusterOpsToLLVM.cpp

@@ -24,6 +24,7 @@
 #include "Dialect/NVGPU/IR/Dialect.h"
 #include "PatternTritonGPUOpToLLVM.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
 using namespace mlir;
@@ -38,8 +39,13 @@ struct ClusterArriveOpConversion
   LogicalResult
   matchAndRewrite(triton::nvidia_gpu::ClusterArriveOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<triton::nvgpu::ClusterArriveOp>(
-        op, op.getRelaxed());
+    auto ctx = rewriter.getContext();
+    auto unitAttr = UnitAttr::get(ctx);
+    if (op.getRelaxed()) {
+      rewriter.replaceOpWithNewOp<NVVM::ClusterArriveRelaxedOp>(op, unitAttr);
+    } else {
+      rewriter.replaceOpWithNewOp<NVVM::ClusterArriveOp>(op, unitAttr);
+    }
     return success();
   }
 };
@@ -52,7 +58,8 @@ struct ClusterWaitOpConversion
   LogicalResult
   matchAndRewrite(triton::nvidia_gpu::ClusterWaitOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<triton::nvgpu::ClusterWaitOp>(op);
+    auto ctx = rewriter.getContext();
+    rewriter.replaceOpWithNewOp<NVVM::ClusterWaitOp>(op, UnitAttr::get(ctx));
     return success();
   }
 };

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp

@@ -23,6 +23,7 @@
 
 #include "MMAHelpers.h"
 #include "Utility.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
@@ -408,7 +409,7 @@ LogicalResult convertDot(const LLVMTypeConverter *typeConverter,
                                               : triton::nvgpu::WGMMALayout::col;
 
   auto func = op->getParentOfType<LLVM::LLVMFuncOp>();
-  Operation *startSequence = rewriter.create<triton::nvgpu::WGMMAFenceOp>(loc);
+  Operation *startSequence = rewriter.create<NVVM::WgmmaFenceAlignedOp>(loc);
   SmallVector<Value> mmaResults;
   for (int m = 0; m < numRepM; ++m) {
     for (int n = 0; n < numRepN; ++n) {
@@ -479,7 +480,7 @@ LogicalResult convertDot(const LLVMTypeConverter *typeConverter,
       }
     }
   }
-  rewriter.create<triton::nvgpu::WGMMACommitGroupOp>(loc);
+  rewriter.create<NVVM::WgmmaGroupSyncAlignedOp>(loc);
 
   if (sync)
     mmaResults = emitWait(rewriter, loc, mmaResults, 0);