Added the GPU related operation

Author: Alwaysproblem

mlir/cuda-tile/Toy/include/toy/Ops.td

@@ -311,7 +311,7 @@ def ReshapeOp : Toy_Op<"reshape", [Pure]> {
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
-def ReturnOp : Toy_Op<"return", [Pure, HasParent<"FuncOp">,
+def ReturnOp : Toy_Op<"return", [Pure, HasParent<"FuncOp, GPUFuncOp">,
                                  Terminator]> {
   let summary = "return operation";
   let description = [{
@@ -405,4 +405,94 @@ def MatMulOp : Toy_Op<"matmul",
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// lauch GPU Op
+//===----------------------------------------------------------------------===//
+
+def LaunchGpuOp : Toy_Op<"launch_gpu",
+    [DeclareOpInterfaceMethods<CallOpInterface>]> {
+  let summary = "launch gpu kernel operation";
+  let description = [{
+    The "launch_gpu" operation launches a GPU kernel with given grid
+    dimensions.
+
+    ```mlir
+     %4 = toy.launch_gpu @my_func(%1, %3) {grid = [16, 16, 1]}
+           : (tensor<2x3xf32>, tensor<2x3xf32>)
+    ```
+
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$callee,
+    Variadic<F32Tensor>:$inputs,
+    OptionalAttr<DictArrayAttr>:$arg_attrs,
+    OptionalAttr<DictArrayAttr>:$res_attrs
+  );
+
+  let results = (outs Variadic<F32Tensor>:$results);
+
+  let assemblyFormat = [{
+    $callee `(` $inputs `)` attr-dict `:` functional-type($inputs, results)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "StringRef":$callee, "ArrayRef<Value>":$arguments)>
+  ];
+}
+
+//===----------------------------------------------------------------------===//
+// GPUFuncOp
+//===----------------------------------------------------------------------===//
+
+def GPUFuncOp : Toy_Op<"gpu_func", [
+    FunctionOpInterface, IsolatedFromAbove
+  ]> {
+  let summary = "GPU kernel function operation";
+  let description = [{
+    The "toy.gpu_func" operation represents a GPU kernel function. These are
+    callable SSA-region operations that contain toy computations to be run on
+    the GPU.
+
+    Example:
+
+    ```mlir
+    toy.gpu_func @my_kernel(tensor<*xf32> %arg0, tensor<*xf32> %arg1) {
+      ...
+      toy.return
+    }
+    ```
+  }];
+
+  let arguments = (ins
+    SymbolNameAttr:$sym_name,
+    TypeAttrOf<FunctionType>:$function_type,
+    OptionalAttr<DictArrayAttr>:$arg_attrs,
+    OptionalAttr<DictArrayAttr>:$res_attrs
+  );
+
+  let regions = (region AnyRegion:$body);
+
+  let builders = [OpBuilder<(ins
+    "StringRef":$name, "FunctionType":$type,
+    CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)
+  >];
+  let extraClassDeclaration = [{
+    //===------------------------------------------------------------------===//
+    // FunctionOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    /// Returns the argument types of this function.
+    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
+
+    /// Returns the result types of this function.
+    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }
+
+    /// Returns the region on the function operation that is callable.
+    Region *getCallableRegion() { return &getBody(); }
+  }];
+  let hasCustomAssemblyFormat = 1;
+  let skipDefaultBuilders = 1;
+}
+
 #endif // TOY_OPS

mlir/cuda-tile/Toy/mlir/Dialect.cpp

@@ -373,16 +373,20 @@ void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
 //===----------------------------------------------------------------------===//
 
 llvm::LogicalResult ReturnOp::verify() {
-  // We know that the parent operation is a function, because of the 'HasParent'
-  // trait attached to the operation definition.
-  auto function = cast<FuncOp>((*this)->getParentOp());
+  // Parent can be FuncOp or GPUFuncOp; both implement FunctionOpInterface.
+  auto *parent = (*this)->getParentOp();
+  auto function = dyn_cast<FunctionOpInterface>(parent);
+  if (!function)
+    return emitOpError() << "must be enclosed in a function-like op";
+
 
   /// ReturnOps can only have a single optional operand.
   if (getNumOperands() > 1)
     return emitOpError() << "expects at most 1 return operand";
 
   // The operand number and types must match the function signature.
-  const auto &results = function.getFunctionType().getResults();
+  auto funcType = llvm::cast<FunctionType>(function.getFunctionType());
+  const auto &results = funcType.getResults();
   if (getNumOperands() != results.size())
     return emitOpError() << "does not return the same number of values ("
                          << getNumOperands() << ") as the enclosing function ("
@@ -489,6 +493,79 @@ llvm::LogicalResult MatMulOp::verify() {
   return mlir::success();
 }
 
+//===----------------------------------------------------------------------===//
+// LaunchGpuOp
+//===----------------------------------------------------------------------===//
+
+void LaunchGpuOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
+                          StringRef callee, ArrayRef<mlir::Value> arguments) {
+  // Generic call always returns an unranked Tensor initially.
+  state.addTypes(UnrankedTensorType::get(builder.getF32Type()));
+  state.addOperands(arguments);
+  state.addAttribute("callee",
+                     mlir::SymbolRefAttr::get(builder.getContext(), callee));
+  state.addAttribute("grid", builder.getI64ArrayAttr({1, 1, 1}));
+}
+
+/// Return the callee of the generic call operation, this is required by the
+/// call interface.
+CallInterfaceCallable LaunchGpuOp::getCallableForCallee() {
+  return (*this)->getAttrOfType<SymbolRefAttr>("callee");
+}
+
+/// Set the callee for the generic call operation, this is required by the call
+/// interface.
+void LaunchGpuOp::setCalleeFromCallable(CallInterfaceCallable callee) {
+  (*this)->setAttr("callee", cast<SymbolRefAttr>(callee));
+}
+
+/// Get the argument operands to the called function, this is required by the
+/// call interface.
+Operation::operand_range LaunchGpuOp::getArgOperands() { return getInputs(); }
+
+/// Get the argument operands to the called function as a mutable range, this is
+/// required by the call interface.
+MutableOperandRange LaunchGpuOp::getArgOperandsMutable() {
+  return getInputsMutable();
+}
+
+
+//===----------------------------------------------------------------------===//
+// GPUFuncOp
+//===----------------------------------------------------------------------===//
+
+void GPUFuncOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
+                   llvm::StringRef name, mlir::FunctionType type,
+                   llvm::ArrayRef<mlir::NamedAttribute> attrs) {
+  // FunctionOpInterface provides a convenient `build` method that will populate
+  // the state of our GPUFuncOp, and create an entry block.
+  buildWithEntryBlock(builder, state, name, type, attrs, type.getInputs());
+}
+
+mlir::ParseResult GPUFuncOp::parse(mlir::OpAsmParser &parser,
+                                mlir::OperationState &result) {
+  // Dispatch to the FunctionOpInterface provided utility method that parses the
+  // function operation.
+  auto buildFuncType =
+      [](mlir::Builder &builder, llvm::ArrayRef<mlir::Type> argTypes,
+         llvm::ArrayRef<mlir::Type> results,
+         mlir::function_interface_impl::VariadicFlag,
+         std::string &) { return builder.getFunctionType(argTypes, results); };
+
+  return mlir::function_interface_impl::parseFunctionOp(
+      parser, result, /*allowVariadic=*/false,
+      getFunctionTypeAttrName(result.name), buildFuncType,
+      getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));
+}
+
+void GPUFuncOp::print(mlir::OpAsmPrinter &p) {
+  // Dispatch to the FunctionOpInterface provided utility method that prints the
+  // function operation.
+  mlir::function_interface_impl::printFunctionOp(
+      p, *this, /*isVariadic=*/false, getFunctionTypeAttrName(),
+      getArgAttrsAttrName(), getResAttrsAttrName());
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//

mlir/cuda-tile/sample/gpu.mlir

@@ -0,0 +1,13 @@
+toy.gpu_func @my_kernel(%arg0: tensor<2x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<2x2xf32> {
+  %2 = toy.matmul(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3x2xf32>) to tensor<2x2xf32>
+  toy.return %2 : tensor<2x2xf32>
+}
+
+toy.func @main() {
+  %1 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf32>
+  %3 = toy.constant dense<[[1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00], [5.000000e+00, 6.000000e+00]]> : tensor<3x2xf32>
+  %4 = toy.launch_gpu @my_kernel(%1, %3) {grid = [16, 16, 1]}
+        : (tensor<2x3xf32>, tensor<3x2xf32>) -> tensor<2x2xf32>
+  toy.print %4 : tensor<2x2xf32>
+  toy.return
+}