Add the "gpu.launch_func" op to the GPU dialect.

    The idea is to lower `gpu.launch` operations into `gpu.launch_func` operations by outlining the kernel body into a function, which is closer to the NVVM model.

--

PiperOrigin-RevId: 246806890

diff --git a/mlir/g3doc/Dialects/GPU.md b/mlir/g3doc/Dialects/GPU.md
index ddf3c3e..471ac2c 100644 (file)
--- a/mlir/g3doc/Dialects/GPU.md
+++ b/mlir/g3doc/Dialects/GPU.md
@@ -84,3 +84,41 @@ understanding that a value has additional semantics (e.g., we will need to know
 what value corresponds to threadIdx.x for coalescing). We can recover these
 properties by analyzing the operations producing values, but it is easier just
 to have that information by construction.
+
+### `gpu.launch_func`
+
+Launch a kernel given as a function on the specified grid of thread blocks.
+`gpu.launch` operations are lowered to `gpu.launch_func` operations by outlining
+the kernel body into a function, which is closer to the NVVM model. The
+`gpu.launch_func` operation has a function attribute namend `kernel` to specify
+the kernel function to launch. The kernel function itself has a `nvvm.kernel`
+attribute.
+
+The operation takes at least six operands, with the first three operands being
+grid sizes along x,y,z dimensions and the following three being block sizes
+along x,y,z dimensions. When a lower-dimensional kernel is required, unused
+sizes must be explicitly set to `1`. The remaining operands are passed as
+arguments to the kernel function.
+
+A custom syntax for this operation is currently not available.
+
+Example:
+
+```mlir {.mlir}
+func @kernel_1(%arg0 : f32, %arg1 : !llvm<"float*">)
+    attributes { nvvm.kernel: true } {
+
+  // Operations that produce block/thread IDs and sizes will be injected when
+  // outlining the `gpu.launch` body to a function called by `gpu.launch_func`.
+  // TODO(tjoerg): Update this example when outlining is implemented.
+
+  "some_op"(%bx, %tx) : (index, index) -> ()
+  %42 = load %arg1[%bx] : memref<?xf32, 1>
+}
+
+"gpu.launch_func"(%cst, %cst, %cst,  // Grid sizes.
+                  %cst, %cst, %cst,  // Block sizes.
+                  %arg0, %arg1)      // Arguments passed to the kernel function.
+      {kernel: @kernel_1 : (f32, !llvm<"float*">) -> ()}  // Kernel function.
+      : (index, index, index, index, index, index, f32, !llvm<"float*">) -> ()
+```

diff --git a/mlir/include/mlir/GPU/GPUDialect.h b/mlir/include/mlir/GPU/GPUDialect.h
index 4a9b0f6..e555e00 100644 (file)
--- a/mlir/include/mlir/GPU/GPUDialect.h
+++ b/mlir/include/mlir/GPU/GPUDialect.h
@@ -27,6 +27,7 @@
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
+namespace gpu {
 
 /// The dialect containing GPU kernel launching operations and related
 /// facilities.
@@ -97,6 +98,29 @@ private:
   static constexpr unsigned kNumConfigRegionAttributes = 12;
 };
 
+/// Operation to launch a kernel given as outlined function.
+class LaunchFuncOp : public Op<LaunchFuncOp, OpTrait::AtLeastNOperands<6>::Impl,
+                               OpTrait::ZeroResult> {
+public:
+  using Op::Op;
+
+  /// The kernel function specified by the operation's `kernel` attribute.
+  Function *kernel();
+  /// The number of operands passed to the kernel function.
+  unsigned getNumKernelOperands();
+  /// The i-th operand passed to the kernel function.
+  Value *getKernelOperand(unsigned i);
+
+  LogicalResult verify();
+
+  static StringRef getOperationName() { return "gpu.launch_func"; }
+
+  /// The number of launch configuration operands, placed at the leading
+  /// positions of the operand list.
+  static constexpr unsigned kNumConfigOperands = 6;
+};
+
+} // end namespace gpu
 } // end namespace mlir
 
 #endif // MLIR_GPUKERNEL_GPUDIALECT_H

diff --git a/mlir/lib/GPU/IR/DialectRegistration.cpp b/mlir/lib/GPU/IR/DialectRegistration.cpp
index e777133..8d00032 100644 (file)
--- a/mlir/lib/GPU/IR/DialectRegistration.cpp
+++ b/mlir/lib/GPU/IR/DialectRegistration.cpp
@@ -18,4 +18,4 @@
 #include "mlir/GPU/GPUDialect.h"
 
 // Static initialization for GPU dialect registration.
-static mlir::DialectRegistration<mlir::GPUDialect> kernelDialect;
+static mlir::DialectRegistration<mlir::gpu::GPUDialect> kernelDialect;

diff --git a/mlir/lib/GPU/IR/GPUDialect.cpp b/mlir/lib/GPU/IR/GPUDialect.cpp
index cc440b7..a55e925 100644 (file)
--- a/mlir/lib/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/GPU/IR/GPUDialect.cpp
@@ -25,12 +25,13 @@
 #include "mlir/IR/StandardTypes.h"
 
 using namespace mlir;
+using namespace mlir::gpu;
 
 StringRef GPUDialect::getDialectName() { return "gpu"; }
 
 GPUDialect::GPUDialect(MLIRContext *context)
     : Dialect(getDialectName(), context) {
-  addOperations<LaunchOp>();
+  addOperations<LaunchOp, LaunchFuncOp>();
 }
 
 //===----------------------------------------------------------------------===//
@@ -257,3 +258,43 @@ bool LaunchOp::parse(OpAsmParser *parser, OperationState *result) {
   return parser->parseRegion(*body, regionArgs, dataTypes) ||
          parser->parseOptionalAttributeDict(result->attributes);
 }
+
+
+//===----------------------------------------------------------------------===//
+// LaunchFuncOp
+//===----------------------------------------------------------------------===//
+Function *LaunchFuncOp::kernel() {
+  return this->getAttr("kernel").dyn_cast<FunctionAttr>().getValue();
+}
+
+unsigned LaunchFuncOp::getNumKernelOperands() {
+  return getNumOperands() - kNumConfigOperands;
+}
+
+Value *LaunchFuncOp::getKernelOperand(unsigned i) {
+  return getOperation()->getOperand(i + kNumConfigOperands);
+}
+
+LogicalResult LaunchFuncOp::verify() {
+  auto kernelAttr = this->getAttr("kernel");
+  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();
+  unsigned numKernelFuncArgs = kernelFunc->getNumArguments();
+  if (getNumKernelOperands() != numKernelFuncArgs) {
+    return emitOpError("got " + Twine(getNumKernelOperands()) +
+                       " kernel operands but expected " +
+                       Twine(numKernelFuncArgs));
+  }
+  for (unsigned i = 0; i < numKernelFuncArgs; ++i) {
+    if (getKernelOperand(i)->getType() !=
+        kernelFunc->getArgument(i)->getType()) {
+      return emitOpError("type of function argument " + Twine(i) +
+                         " does not match");
+    }
+  }
+  return success();
+}

diff --git a/mlir/test/GPU/invalid.mlir b/mlir/test/GPU/invalid.mlir
index 217a656..42db60c 100644 (file)
--- a/mlir/test/GPU/invalid.mlir
+++ b/mlir/test/GPU/invalid.mlir
@@ -71,3 +71,59 @@ func @nested_isolation(%sz : index) {
   }) : (index, index, index, index, index, index) -> ()
   return
 }
+
+// -----
+
+func @launch_func_too_few_operands(%sz : index) {
+  // expected-error@+1 {{expected 6 or more operands}}
+  "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz)
+      : (index, index, index, index, index) -> ()
+  return
+}
+
+// -----
+
+func @launch_func_missing_callee_attribute(%sz : index) {
+  // expected-error@+1 {{attribute 'kernel' must be specified}}
+  "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz) {foo: "bar"}
+      : (index, index, index, index, index, index) -> ()
+  return
+}
+
+// -----
+
+func @launch_func_no_function_attribute(%sz : index) {
+  // expected-error@+1 {{attribute 'kernel' must be a function}}
+  "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz) {kernel: "bar"}
+      : (index, index, index, index, index, index) -> ()
+  return
+}
+
+// -----
+
+func @kernel_1(%arg1 : !llvm<"float*">) attributes { nvvm.kernel: true } {
+  return
+}
+
+func @launch_func_kernel_operand_size(%sz : index, %arg : !llvm<"float*">) {
+  // expected-error@+1 {{got 2 kernel operands but expected 1}}
+  "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg, %arg)
+      {kernel: @kernel_1 : (!llvm<"float*">) -> ()}
+      : (index, index, index, index, index, index, !llvm<"float*">,
+         !llvm<"float*">) -> ()
+  return
+}
+
+// -----
+
+func @kernel_1(%arg1 : !llvm<"float*">) attributes { nvvm.kernel: true } {
+  return
+}
+
+func @launch_func_kernel_operand_types(%sz : index, %arg : f32) {
+  // expected-error@+1 {{type of function argument 0 does not match}}
+  "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg)
+      {kernel: @kernel_1 : (!llvm<"float*">) -> ()}
+      : (index, index, index, index, index, index, f32) -> ()
+  return
+}

diff --git a/mlir/test/GPU/ops.mlir b/mlir/test/GPU/ops.mlir
index 74b3705..1e99661 100644 (file)
--- a/mlir/test/GPU/ops.mlir
+++ b/mlir/test/GPU/ops.mlir
@@ -52,3 +52,21 @@ func @nested_isolation(%sz : index) {
   }
   return
 }
+
+func @kernel_1(%arg0 : f32, %arg1 : !llvm<"float*">)
+    attributes { nvvm.kernel: true } {
+  return
+}
+
+func @foo() {
+  %0 = "op"() : () -> (f32)
+  %1 = "op"() : () -> (!llvm<"float*">)
+// CHECK: %c8 = constant 8
+  %cst = constant 8 : index
+
+// CHECK: "gpu.launch_func"(%c8, %c8, %c8, %c8, %c8, %c8, %0, %1) {kernel: @kernel_1 : (f32, !llvm<"float*">) -> ()} : (index, index, index, index, index, index, f32, !llvm<"float*">) -> ()
+  "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1)
+      {kernel: @kernel_1 : (f32, !llvm<"float*">) -> ()}
+      : (index, index, index, index, index, index, f32, !llvm<"float*">) -> ()
+  return
+}