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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir-gpu

Author: Asher Mancinelli (ashermancinelli)

<details>
<summary>Changes</summary>

Adds support for the argument attributes required to mark pointer arguments on gpu.func ops as grid-constant.

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


3 Files Affected:

- (modified) mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp (+4) 
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+2-1) 
- (modified) mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir (+16) 


``````````diff
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
+  }
+}

``````````

</details>


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