[Mlir-commits] [mlir] Revert "[mlir][NVVM] Disallow results on kernel functions (#96399)" (PR #97074)

[Jeff Niu]

https://github.com/Mogball created https://github.com/llvm/llvm-project/pull/97074

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

>From 684b928f1075bd094ed2fc8cc001aa5cff889fb4 Mon Sep 17 00:00:00 2001
From: Jeff Niu <jeff at modular.com>
Date: Fri, 28 Jun 2024 08:32:18 -0700
Subject: [PATCH] Revert "[mlir][NVVM] Disallow results on kernel functions
 (#96399)"

This reverts commit 346c4a88afedcef3da40f68c83f0a5b3e0ac61ea.
---
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 13 +++----------
 mlir/test/Target/LLVMIR/nvvmir.mlir        |  7 -------
 2 files changed, 3 insertions(+), 17 deletions(-)

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
-}