Verify that kernel functions referenced by gpu.launch_func have a gpu.kernel...

    Also extract gpu.launch_func's function attribute name into a constant.

--

PiperOrigin-RevId: 247595352

diff --git a/mlir/include/mlir/GPU/GPUDialect.h b/mlir/include/mlir/GPU/GPUDialect.h
index 86a1acf..7c8b2df 100644 (file)
--- a/mlir/include/mlir/GPU/GPUDialect.h
+++ b/mlir/include/mlir/GPU/GPUDialect.h
@@ -138,6 +138,10 @@ public:
   /// The number of launch configuration operands, placed at the leading
   /// positions of the operand list.
   static constexpr unsigned kNumConfigOperands = 6;
+
+private:
+  /// The name of the function attribute specifying the kernel to launch.
+  static StringRef getKernelAttrName() { return "kernel"; }
 };
 
 #define GET_OP_CLASSES

diff --git a/mlir/lib/GPU/IR/GPUDialect.cpp b/mlir/lib/GPU/IR/GPUDialect.cpp
index 2340dfa..4762add 100644 (file)
--- a/mlir/lib/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/GPU/IR/GPUDialect.cpp
@@ -304,7 +304,8 @@ void LaunchFuncOp::build(Builder *builder, OperationState *result,
   result->addOperands(
       {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
   result->addOperands(kernelOperands);
-  result->addAttribute("kernel", builder->getFunctionAttr(kernelFunc));
+  result->addAttribute(getKernelAttrName(),
+                       builder->getFunctionAttr(kernelFunc));
 }
 
 void LaunchFuncOp::build(Builder *builder, OperationState *result,
@@ -316,7 +317,7 @@ void LaunchFuncOp::build(Builder *builder, OperationState *result,
 }
 
 Function *LaunchFuncOp::kernel() {
-  return this->getAttr("kernel").dyn_cast<FunctionAttr>().getValue();
+  return this->getAttr(getKernelAttrName()).dyn_cast<FunctionAttr>().getValue();
 }
 
 unsigned LaunchFuncOp::getNumKernelOperands() {
@@ -328,13 +329,18 @@ Value *LaunchFuncOp::getKernelOperand(unsigned i) {
 }
 
 LogicalResult LaunchFuncOp::verify() {
-  auto kernelAttr = this->getAttr("kernel");
+  auto kernelAttr = this->getAttr(getKernelAttrName());
   if (!kernelAttr) {
     return emitOpError("attribute 'kernel' must be specified");
   } else if (!kernelAttr.isa<FunctionAttr>()) {
     return emitOpError("attribute 'kernel' must be a function");
   }
   Function *kernelFunc = this->kernel();
+  if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
+          GPUDialect::getKernelFuncAttrName())) {
+    return emitError("kernel function is missing the '" +
+                     GPUDialect::getKernelFuncAttrName() + "' attribute");
+  }
   unsigned numKernelFuncArgs = kernelFunc->getNumArguments();
   if (getNumKernelOperands() != numKernelFuncArgs) {
     return emitOpError("got " + Twine(getNumKernelOperands()) +

diff --git a/mlir/test/GPU/invalid.mlir b/mlir/test/GPU/invalid.mlir
index 42db60c..cbc7bfe 100644 (file)
--- a/mlir/test/GPU/invalid.mlir
+++ b/mlir/test/GPU/invalid.mlir
@@ -101,7 +101,21 @@ func @launch_func_no_function_attribute(%sz : index) {
 
 // -----
 
-func @kernel_1(%arg1 : !llvm<"float*">) attributes { nvvm.kernel: true } {
+func @kernel_1(%arg1 : !llvm<"float*">) {
+  return
+}
+
+func @launch_func_missing_kernel_attr(%sz : index, %arg : !llvm<"float*">) {
+  // expected-error@+1 {{kernel function is missing the 'gpu.kernel' attribute}}
+  "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg)
+      {kernel: @kernel_1 : (!llvm<"float*">) -> ()}
+      : (index, index, index, index, index, index, !llvm<"float*">) -> ()
+  return
+}
+
+// -----
+
+func @kernel_1(%arg1 : !llvm<"float*">) attributes { gpu.kernel } {
   return
 }
 
@@ -116,7 +130,7 @@ func @launch_func_kernel_operand_size(%sz : index, %arg : !llvm<"float*">) {
 
 // -----
 
-func @kernel_1(%arg1 : !llvm<"float*">) attributes { nvvm.kernel: true } {
+func @kernel_1(%arg1 : !llvm<"float*">) attributes { gpu.kernel } {
   return
 }
 

diff --git a/mlir/test/GPU/ops.mlir b/mlir/test/GPU/ops.mlir
index ac31d23..6d0105d 100644 (file)
--- a/mlir/test/GPU/ops.mlir
+++ b/mlir/test/GPU/ops.mlir
@@ -54,7 +54,7 @@ func @nested_isolation(%sz : index) {
 }
 
 func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>)
-    attributes { nvvm.kernel: true } {
+    attributes { gpu.kernel } {
   %tIdX = "gpu.thread_id"() {dimension: "x"} : () -> (index)
   %tIdY = "gpu.thread_id"() {dimension: "y"} : () -> (index)
   %tIdZ = "gpu.thread_id"() {dimension: "z"} : () -> (index)