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*">) -> ()
+```
#include "mlir/IR/OpDefinition.h"
namespace mlir {
+namespace gpu {
/// The dialect containing GPU kernel launching operations and related
/// facilities.
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
#include "mlir/GPU/GPUDialect.h"
// Static initialization for GPU dialect registration.
-static mlir::DialectRegistration<mlir::GPUDialect> kernelDialect;
+static mlir::DialectRegistration<mlir::gpu::GPUDialect> kernelDialect;
#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>();
}
//===----------------------------------------------------------------------===//
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();
+}
}) : (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
+}
}
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
+}