Revert "[mlir][NVVM] Disallow results on kernel functions (#96399)"

[Mogball]

NVVM IR itself doesn't place any restriction that a function annotated as nvvm.kernel actually has no results, so this is a mismatch at the NVVMDialect level and NVVMIR. The GPU dialect might enforce that kernel functions have no results, but it doesn't make sense to propagate this constraint downstream.

Reverts #96399

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Jeff Niu (Mogball)

Changes

NVVM IR itself doesn't place any restriction that a function annotated as nvvm.kernel actually has no results, so this is a mismatch at the NVVMDialect level and NVVMIR. The GPU dialect might enforce that kernel functions have no results, but it doesn't make sense to propagate this constraint downstream.

Reverts llvm/llvm-project#96399

---

Full diff: https://github.com/llvm/llvm-project/pull/97074.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+3-10)
• (modified) mlir/test/Target/LLVMIR/nvvmir.mlir (-7)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 3d6a911f36541..94197e473ce01 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -214,8 +214,7 @@ void MmaOp::print(OpAsmPrinter &p) {
   p.printOptionalAttrDict(this->getOperation()->getAttrs(), ignoreAttrNames);
 
   // Print the types of the operands and result.
-  p << " : "
-    << "(";
+  p << " : " << "(";
   llvm::interleaveComma(SmallVector<Type, 3>{frags[0].regs[0].getType(),
                                              frags[1].regs[0].getType(),
                                              frags[2].regs[0].getType()},
@@ -956,9 +955,7 @@ std::string NVVM::WgmmaMmaAsyncOp::getPtx() {
   ss << "},";
   // Need to map read/write registers correctly.
   regCnt = (regCnt * 2);
-  ss << " $" << (regCnt) << ","
-     << " $" << (regCnt + 1) << ","
-     << " p";
+  ss << " $" << (regCnt) << "," << " $" << (regCnt + 1) << "," << " p";
   if (getTypeD() != WGMMATypes::s32) {
     ss << ", $" << (regCnt + 3) << ",  $" << (regCnt + 4);
   }
@@ -1056,14 +1053,10 @@ LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
   StringAttr attrName = attr.getName();
   // Kernel function attribute should be attached to functions.
   if (attrName == NVVMDialect::getKernelFuncAttrName()) {
-    auto funcOp = dyn_cast<LLVM::LLVMFuncOp>(op);
-    if (!funcOp) {
+    if (!isa<LLVM::LLVMFuncOp>(op)) {
       return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
                              << "' attribute attached to unexpected op";
     }
-    if (!funcOp.getResultTypes().empty()) {
-      return op->emitError() << "kernel function cannot have results";
-    }
   }
   // If maxntid and reqntid exist, it must be an array with max 3 dim
   if (attrName == NVVMDialect::getMaxntidAttrName() ||
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 26ba80cba6ed5..a8ae4d97888c9 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -574,10 +574,3 @@ llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant})
 llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant}, %arg1: f32, %arg2: !llvm.ptr {llvm.byval = f32, nvvm.grid_constant}) attributes {nvvm.kernel} {
   llvm.return
 }
-
-// -----
-
-// expected-error @below{{kernel function cannot have results}}
-llvm.func @kernel_with_result(%i: i32) -> i32 attributes {nvvm.kernel}  {
-  llvm.return %i : i32
-}

[grypp]

The GPU dialect might enforce that kernel functions have no results, but it doesn't make sense to propagate this constraint downstream.

Can you please show me an example of a kernel can return a value? So we can agree on that revert :)

[Mogball]

I agree that kernel functions return void. However, the design of the LLVMDialect and its sister dialects, like NVVMDialect, are such that they are meant to be trivial mappings to LLVMIR. Thus, these dialects should not pose additional constraints that are not present in LLVMIR, unless they are MLIR specific.

[grypp]

Thanks for the design information. If MLIR doesn’t have a verifier there, somewhere in llvm or ptxas will surely fail. But isn’t that like letting a small leak turn into a flood?

I thought MLIR has opportunity to catch issues early and elegantly.

[Mogball]

LLVM and PTXAS happily accept this 🤷 . If they did enforce the invariant, then it would make sense for the LLVMDialect to as well.

[joker-eph]

these dialects should not pose additional constraints that are not present in LLVMIR, unless they are MLIR specific.

Does the nvvm.kernel attribute exist in LLVM? Otherwise this could qualify as "MLIR Specific" to me.

[Mogball]

Yeah it does. It's part of the NVVM spec.

[joker-eph]

This could be seen as LLVM being too lose as well: probably because there is no infrastructure in LLVM to add custom verifier through attributes?