[Mlir-commits] [mlir] [MLIR][GPU] Support grid constant, byval, byref on gpu.func (PR #172037)

[Asher Mancinelli]

https://github.com/ashermancinelli updated https://github.com/llvm/llvm-project/pull/172037

>From 9e89396e79da4200c9bf7bd7dc221b301edd1395 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Fri, 12 Dec 2025 08:09:05 -0800
Subject: [PATCH 1/2] Support grid constant, byval, byref on gpu.func

---
 mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp |  4 ++++
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp       |  3 ++-
 mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir  | 16 ++++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
index eb662a1b056de..d0a03098044c7 100644
--- a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
+++ b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
@@ -11,6 +11,7 @@
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/LLVMCommon/VectorPattern.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -369,7 +370,10 @@ GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
     }
 
     if (lowersToPointer) {
+      copyPointerAttribute(mlir::NVVM::NVVMDialect::getGridConstantAttrName());
       copyPointerAttribute(LLVM::LLVMDialect::getNoAliasAttrName());
+      copyPointerAttribute(LLVM::LLVMDialect::getByValAttrName());
+      copyPointerAttribute(LLVM::LLVMDialect::getByRefAttrName());
       copyPointerAttribute(LLVM::LLVMDialect::getNoCaptureAttrName());
       copyPointerAttribute(LLVM::LLVMDialect::getNoFreeAttrName());
       copyPointerAttribute(LLVM::LLVMDialect::getAlignAttrName());
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index a3ff904146b92..603dc6d91b166 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -5443,7 +5443,8 @@ LogicalResult NVVMDialect::verifyRegionArgAttribute(Operation *op,
   if (!funcOp)
     return success();
 
-  bool isKernel = op->hasAttr(NVVMDialect::getKernelFuncAttrName());
+  const bool isKernel = op->hasAttr(NVVMDialect::getKernelFuncAttrName()) ||
+                        op->hasAttr(gpu::GPUDialect::getKernelFuncAttrName());
   StringAttr attrName = argAttr.getName();
   if (attrName == NVVM::NVVMDialect::getGridConstantAttrName()) {
     if (!isKernel) {
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index f1cc1eb983267..28613ee33cac6 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -1149,3 +1149,19 @@ gpu.module @test_module_56 {
     func.return %sin16, %cos16, %sin32, %cos32, %sin64, %cos64 : f16, f16, f32, f32, f64, f64
   }
 }
+
+// Check that nvvm.grid_constant is a valid argument attribute on gpu.kernel.
+gpu.module @test_module_57 {
+  // CHECK:       gpu.module @test_module_57
+  // CHECK-LABEL:   llvm.func @test_kernel(
+  // CHECK-SAME:      %[[VAL_0:.*]]: !llvm.ptr {llvm.byval = i64, nvvm.grid_constant}
+  // CHECK-SAME:      %[[VAL_1:.*]]: !llvm.ptr {llvm.byref = i64}
+  // CHECK:           llvm.return
+  // CHECK:         }
+  // CHECK:       }
+  gpu.func @test_kernel(
+      %arg0: !llvm.ptr {nvvm.grid_constant, llvm.byval = i64},
+      %arg1: !llvm.ptr {llvm.byref = i64}) kernel {
+    gpu.return
+  }
+}

>From ff5228ecee3ff4928f31815db9c98258136ed416 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Fri, 12 Dec 2025 08:24:38 -0800
Subject: [PATCH 2/2] Prefer early-exit

---
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp    | 33 ++++++++++---------
 .../Conversion/GPUToNVVM/gpu-to-nvvm.mlir     |  2 +-
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 603dc6d91b166..c4eb77ecc693b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -5443,24 +5443,25 @@ LogicalResult NVVMDialect::verifyRegionArgAttribute(Operation *op,
   if (!funcOp)
     return success();
 
+  StringAttr attrName = argAttr.getName();
+  if (attrName != NVVM::NVVMDialect::getGridConstantAttrName())
+    return success();
+
   const bool isKernel = op->hasAttr(NVVMDialect::getKernelFuncAttrName()) ||
                         op->hasAttr(gpu::GPUDialect::getKernelFuncAttrName());
-  StringAttr attrName = argAttr.getName();
-  if (attrName == NVVM::NVVMDialect::getGridConstantAttrName()) {
-    if (!isKernel) {
-      return op->emitError()
-             << "'" << attrName
-             << "' attribute must be present only on kernel arguments";
-    }
-    if (!isa<UnitAttr>(argAttr.getValue()))
-      return op->emitError() << "'" << attrName << "' must be a unit attribute";
-    if (!funcOp.getArgAttr(argIndex, LLVM::LLVMDialect::getByValAttrName())) {
-      return op->emitError()
-             << "'" << attrName
-             << "' attribute requires the argument to also have attribute '"
-             << LLVM::LLVMDialect::getByValAttrName() << "'";
-    }
-  }
+  if (!isKernel)
+    return op->emitError()
+           << "'" << attrName
+           << "' attribute must be present only on kernel arguments";
+
+  if (!isa<UnitAttr>(argAttr.getValue()))
+    return op->emitError() << "'" << attrName << "' must be a unit attribute";
+
+  if (!funcOp.getArgAttr(argIndex, LLVM::LLVMDialect::getByValAttrName()))
+    return op->emitError()
+           << "'" << attrName
+           << "' attribute requires the argument to also have attribute '"
+           << LLVM::LLVMDialect::getByValAttrName() << "'";
 
   return success();
 }
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index 28613ee33cac6..68e6025651a8a 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -1150,7 +1150,7 @@ gpu.module @test_module_56 {
   }
 }
 
-// Check that nvvm.grid_constant is a valid argument attribute on gpu.kernel.
+// Check that nvvm.grid_constant, llvm.byref, and llvm.byval are valid argument attributes on gpu.kernel.
 gpu.module @test_module_57 {
   // CHECK:       gpu.module @test_module_57
   // CHECK-LABEL:   llvm.func @test_kernel(