[Hopper] Verify WarpGroupDotWaitOp has at least 1 dep (#7732)

Also do a bit of code cleanup

Author: Mogball

include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td

@@ -123,6 +123,7 @@ def TTNG_WarpGroupDotWaitOp : TTNG_Op<"warp_group_dot_wait", [DeclareOpInterface
   }];
 
   let assemblyFormat = "$inputs attr-dict `:` type($inputs)";
+  let hasVerifier = 1;
 }
 
 def TTNG_InitBarrierOp : TTNG_Op<"init_barrier"> {

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -135,13 +135,18 @@ bool WarpGroupDotOp::verifyDims() {
 
 // -- WarpGroupDotWaitOp --
 LogicalResult WarpGroupDotWaitOp::inferReturnTypes(
-    ::mlir::MLIRContext *context, ::std::optional<::mlir::Location> location,
-    ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
-    ::mlir::OpaqueProperties properties, ::mlir::RegionRange regions,
-    ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+    MLIRContext *context, std::optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
   for (Value operand : operands)
     inferredReturnTypes.push_back(operand.getType());
-  return mlir::success();
+  return success();
+}
+
+LogicalResult WarpGroupDotWaitOp::verify() {
+  if (getOperands().empty())
+    return emitOpError("expected to be waiting on at least one dependency");
+  return success();
 }
 
 // -- InitBarrierOp --

python/src/gluon_ir.cc

@@ -485,7 +485,10 @@ void init_gluon_ir(py::module &&m) {
            })
       .def("create_warpgroup_mma_wait",
            [](GluonOpBuilder &self, std::vector<Value> &deps, int pendings) {
-             self.create<ttng::WarpGroupDotWaitOp>(deps, pendings);
+             std::vector<Value> results;
+             auto wait = self.create<ttng::WarpGroupDotWaitOp>(deps, pendings);
+             llvm::append_range(results, wait.getResults());
+             return results;
            })
       .def("create_tmem_alloc",
            [](GluonOpBuilder &self, Type resultTy, Value value) -> Value {

python/test/gluon/test_core.py

@@ -124,7 +124,7 @@ def warpgroup_mma_kernel(a, b, out, M: ttgl.constexpr, N: ttgl.constexpr, K: ttg
     acc = hopper.warpgroup_mma(a_shmem, b_shmem, acc, is_async=ASYNC)
 
     if ASYNC:
-        hopper.warpgroup_mma_wait(num_outstanding=1, deps=[acc])
+        acc = hopper.warpgroup_mma_wait(num_outstanding=1, deps=[acc])
 
     ttgl.store(out + out_offs_m * N + out_offs_n, acc)
 

python/test/gluon/test_frontend.py

@@ -618,7 +618,7 @@ def test_warpgroup_mma():
 def warpgroup_mma_wait_kernel():
     layout: ttgl.constexpr = ttgl.NVMMADistributedLayout(version=[3, 0], warps_per_cta=[4, 1], instr_shape=[16, 32, 16])
     acc = ttgl.full([128, 128], 0, dtype=ttgl.float16, layout=layout)
-    hopper.warpgroup_mma_wait(num_outstanding=1, deps=[acc])
+    acc = hopper.warpgroup_mma_wait(num_outstanding=1, deps=[acc])
 
 
 def test_warpgroup_mma_wait():

python/triton/experimental/gluon/language/nvidia/hopper/__init__.py

@@ -1,3 +1,4 @@
+from triton.compiler.code_generator import unflatten_ir_values
 from ..ampere import async_copy
 from . import mbarrier, tma
 from ... import _core
@@ -70,6 +71,10 @@ def warpgroup_mma_wait(num_outstanding=0, deps=None, _semantic=None):
         num_outstanding (int): Number of outstanding warpgroup MMA operations to wait for. Defaults to 0.
         deps (Sequence[tensor]): List of dependencies that need to be kept alive while the mma is unfinished.
     """
-    deps = [x.handle for x in deps] if deps is not None else []
+    deps_handles = [x.handle for x in deps] if deps is not None else []
     num_outstanding = _core._unwrap_if_constexpr(num_outstanding)
-    _semantic.builder.create_warpgroup_mma_wait(deps, num_outstanding)
+    results = _semantic.builder.create_warpgroup_mma_wait(deps_handles, num_outstanding)
+    results = tuple(unflatten_ir_values(results, [dep.type for dep in deps]))
+    if len(results) == 1:
+        return results[0]
+    return tuple(results)

test/Conversion/tritongpu_to_llvm_hopper.mlir

@@ -538,3 +538,24 @@ module attributes {ttg.global_scratch_memory_alignment = 1 : i32, ttg.global_scr
     tt.return
   }
 }
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
+
+module attributes {"ttg.target" = "cuda:90", "ttg.num-warps" = 4 : i32} {
+
+// CHECK-LABEL: @warpgroup_dot_wait_1_input
+tt.func @warpgroup_dot_wait_1_input(%arg0: tensor<128xf32, #blocked>) {
+  // CHECK: nvgpu.wgmma_wait_group
+  ttng.warp_group_dot_wait %arg0 {pendings = 0 : i32} : tensor<128xf32, #blocked>
+  tt.return
+}
+
+tt.func @warpgroup_dot_wait_2_inputs(%arg0: tensor<128xf32, #blocked>, %arg1: tensor<128xf32, #blocked>) {
+  // CHECK: nvgpu.wgmma_wait_group
+  ttng.warp_group_dot_wait %arg0, %arg1 {pendings = 0 : i32} : tensor<128xf32, #blocked>, tensor<128xf32, #blocked>
+  tt.return
+}
+
+}

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp

@@ -108,31 +108,28 @@ struct WarpGroupDotWaitOpConversion
                   ConversionPatternRewriter &rewriter) const override {
     auto pendings = op.getPendings();
     Location loc = op.getLoc();
-    if (adaptor.getInputs().size() <= 1) {
-      Value input =
-          adaptor.getInputs().size() == 1 ? adaptor.getInputs()[0] : Value();
-      rewriter.replaceOpWithNewOp<triton::nvgpu::WGMMAWaitGroupOp>(op, input,
-                                                                   pendings);
+    ValueRange inputs = adaptor.getInputs();
+    if (inputs.size() == 1) {
+      rewriter.replaceOpWithNewOp<triton::nvgpu::WGMMAWaitGroupOp>(
+          op, inputs.front(), pendings);
       return success();
     }
-    std::vector<Type> types;
+    SmallVector<Type> types;
     // Pack the inputs into a single struct.
-    for (Value input : adaptor.getInputs()) {
-      auto structType = dyn_cast<LLVM::LLVMStructType>(input.getType());
+    for (Type type : inputs.getTypes()) {
+      auto structType = dyn_cast<LLVM::LLVMStructType>(type);
       if (!structType)
         return failure();
-      for (Type type : structType.getBody())
-        types.push_back(type);
+      llvm::append_range(types, structType.getBody());
     }
     auto packedType =
         LLVM::LLVMStructType::getLiteral(rewriter.getContext(), types);
     Value packed = rewriter.create<LLVM::UndefOp>(loc, packedType);
     unsigned outputStructIndex = 0;
-    for (Value input : adaptor.getInputs()) {
-      auto structType = dyn_cast<LLVM::LLVMStructType>(input.getType());
-      for (unsigned i = 0; i < structType.getBody().size(); ++i) {
-        Value value = rewriter.create<LLVM::ExtractValueOp>(
-            loc, structType.getBody()[i], input, i);
+    for (Value input : inputs) {
+      for (auto [i, type] : llvm::enumerate(
+               cast<LLVM::LLVMStructType>(input.getType()).getBody())) {
+        Value value = rewriter.create<LLVM::ExtractValueOp>(loc, input, i);
         packed = rewriter.create<LLVM::InsertValueOp>(
             loc, packedType, packed, value, outputStructIndex++);
       }
@@ -142,14 +139,12 @@ struct WarpGroupDotWaitOpConversion
     // Unpack the output into the original struct types.
     SmallVector<Value> outputs;
     outputStructIndex = 0;
-    for (Value input : adaptor.getInputs()) {
-      auto structType = cast<LLVM::LLVMStructType>(input.getType());
+    for (Type type : inputs.getTypes()) {
+      auto structType = cast<LLVM::LLVMStructType>(type);
       Value unpacked = rewriter.create<LLVM::UndefOp>(loc, structType);
-      for (unsigned i = 0; i < structType.getBody().size(); ++i) {
+      for (auto [i, type] : llvm::enumerate(structType.getBody())) {
         Value value = rewriter.create<LLVM::ExtractValueOp>(
-            loc, packedType.getBody()[outputStructIndex], packedOutput,
-            outputStructIndex);
-        outputStructIndex++;
+            loc, packedOutput, outputStructIndex++);
         unpacked = rewriter.create<LLVM::InsertValueOp>(loc, structType,
                                                         unpacked, value, i);
       }