Add `gpu.create_token` op.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/GPUBase.td

60

Is this interface more generic than the GPU dialect? Can you include GPU in the name to avoid confusion with the async dialect.

62

Please extend the documentation.

83

if constexpr is a c++17 feature, LLVM is using C++14.

92

Nit: this is a non-trivial block of C++, can this go inside an implementation file?

mlir/include/mlir/Dialect/GPU/GPUOps.td

759

Is ForkOp intended here?

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

327

This hunk seems unrelated to this revision, please split and commit ahead.

[mehdi_amini]

Actually seems like you generated a mixed diff that includes changes from https://reviews.llvm.org/D88972 ?

[csigg]

In D88972#2317357, @mehdi_amini wrote:

Actually seems like you generated a mixed diff that includes changes from https://reviews.llvm.org/D88972 ?

This is a follow-up change. Is there a better way than creating a new revision and adding the parent revision in phabricator?

[mehdi_amini]

Phab works with "diff", somehow you create this revision as a diff against master? In general I do an interactive rebase, edit individually each commit and arc diff HEAD~

[csigg]

Remove base revision, address comments.

[csigg]

mlir/include/mlir/Dialect/GPU/GPUBase.td

60

Sorry, this was meant to be in the ::mlir::gpu namespace. I left the name unchanged.

mlir/include/mlir/Dialect/GPU/GPUOps.td

759

Oops. Thanks for catching that.

[csigg]

In D88972#2317773, @mehdi_amini wrote:

Phab works with "diff", somehow you create this revision as a diff against master? In general I do an interactive rebase, edit individually each commit and arc diff HEAD~

Got it, thanks.

mlir/include/mlir/Dialect/GPU/CMakeLists.txt

 add_mlir_dialect(GPUOps gpu)
 add_mlir_doc(GPUOps -gen-op-doc GPUOps Dialects/)
 
+set(LLVM_TARGET_DEFINITIONS GPUBase.td)
+mlir_tablegen(GPUOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(GPUOpInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRGPUOpInterfacesIncGen)
+
 set(LLVM_TARGET_DEFINITIONS ParallelLoopMapperAttr.td)
 mlir_tablegen(ParallelLoopMapperAttr.h.inc -gen-struct-attr-decls)
 mlir_tablegen(ParallelLoopMapperAttr.cpp.inc -gen-struct-attr-defs)
 add_public_tablegen_target(MLIRParallelLoopMapperAttrGen)
 
 set(LLVM_TARGET_DEFINITIONS ParallelLoopMapperAttr.td)
 mlir_tablegen(ParallelLoopMapperEnums.h.inc -gen-enum-decls)
 mlir_tablegen(ParallelLoopMapperEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRParallelLoopMapperEnumsGen)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name GPU)
 add_public_tablegen_target(MLIRGPUPassIncGen)
 
 add_mlir_doc(Passes -gen-pass-doc GPUPasses ./)

mlir/include/mlir/Dialect/GPU/GPUBase.td

let extraClassDeclaration = [{
     static unsigned getWorkgroupAddressSpace() { return 3; }
 
     /// Returns the numeric value used to identify the private memory address
     /// space.
     static unsigned getPrivateAddressSpace() { return 5; }
   }];
 }
 
+def GPU_AsyncToken : DialectType<
+  GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::AsyncTokenType>()">, "async token">,
+             BuildableType<"mlir::gpu::AsyncTokenType::get($_builder.getContext())">;
+
+def GPU_AsyncOpInterface : OpInterface<"AsyncOpInterface"> {
+  let description = [{
+    Interface for GPU operations that execute asynchronously on the device.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Query the operands that represent async dependency tokens.
+      }],
+      "OperandRange", "getAsyncDependencies", (ins), [{}], [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.asyncDependencies();
+      }]
+    >,
+    InterfaceMethod<[{
+        Set the operands that represent async dependency tokens.
+      }],
+      "void", "setAsyncDependencies", (ins "ValueRange":$asyncDependencies), [{}], [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto asyncDepOdsIndex = 0; // Async dependencies always come first.
+        auto indexAndLength = op.getODSOperandIndexAndLength(asyncDepOdsIndex);
+        op.getOperation()->setOperands(indexAndLength.first, indexAndLength.second, asyncDependencies);
+        if constexpr (op.template hasTrait<OpTrait::AttrSizedOperandSegments>()) {
+          auto attrName = ConcreteOp::getOperandSegmentSizeAttr();
+          auto sizeAttr = op.template getAttrOfType<DenseIntElementsAttr>(attrName);
+          assert(sizeAttr && "operand segment size attribute not found");
+          SmallVector<int32_t, 8> sizes;
+          for (auto size : sizeAttr.getIntValues())
+            sizes.push_back(size.getSExtValue());
+          sizes[asyncDepOdsIndex] = asyncDependencies.size();
+          op.setAttr(attrName, Builder(op.getContext()).getI32VectorAttr(sizes));
+        }
+      }]
+    >,
+    InterfaceMethod<[{
+        Query the result that represents the async token to depend on.
+      }],
+      "OpResult", "getAsyncToken", (ins), [{}], [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.asyncToken().template cast<OpResult>();
+      }]
+    >,
+  ];
+}
+
 #endif // GPU_BASE

mlir/include/mlir/Dialect/GPU/GPUDialect.h

 //===- GPUDialect.h - MLIR Dialect for GPU Kernels --------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the GPU kernel-related operations and puts them in the
 // corresponding dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_GPU_GPUDIALECT_H
 #define MLIR_DIALECT_GPU_GPUDIALECT_H
 
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 class FuncOp;
 
 namespace gpu {
 
 /// Utility class for the GPU dialect to represent triples of `Value`s
 /// accessible through `.x`, `.y`, and `.z` similarly to CUDA notation.
 struct KernelDim3 {
   Value x;
   Value y;
   Value z;
 };
 
+class AsyncTokenType
+    : public Type::TypeBase<AsyncTokenType, Type, TypeStorage> {
+public:
+  // Used for generic hooks in TypeBase.
+  using Base::Base;
+};
+
 } // end namespace gpu
 } // end namespace mlir
 
 #include "mlir/Dialect/GPU/GPUOpsDialect.h.inc"
 
+#include "mlir/Dialect/GPU/GPUOpInterfaces.h.inc"
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.h.inc"
 
-
 #endif // MLIR_DIALECT_GPU_GPUDIALECT_H

mlir/include/mlir/Dialect/GPU/GPUOps.td

let extraClassDeclaration = [{
     /// workgroup memory.
     BlockArgument addWorkgroupAttribution(Type type);
 
     /// Returns the number of buffers located in the private memory.
     unsigned getNumPrivateAttributions() {
       return getBody().getNumArguments() - getType().getNumInputs() -
           getNumWorkgroupAttributions();
     }
- 
     /// Returns a list of block arguments that correspond to buffers located in
     /// the private memory.
     ArrayRef<BlockArgument> getPrivateAttributions() {
       // Buffers on the private memory always come after buffers on the workgroup
       // memory.
       auto begin =
           std::next(getBody().args_begin(),
                     getType().getNumInputs() + getNumWorkgroupAttributions());

 }
 
 def GPU_LaunchFuncOp : GPU_Op<"launch_func">,
     Arguments<(ins IntOrIndex:$gridSizeX, IntOrIndex:$gridSizeY,
                IntOrIndex:$gridSizeZ, IntOrIndex:$blockSizeX,
                IntOrIndex:$blockSizeY, IntOrIndex:$blockSizeZ,
                Variadic<AnyType>:$operands)>,
     Results<(outs)> {
-  let summary = "Launches a function as a GPU kerneel";
+  let summary = "Launches a function as a GPU kernel";
 
   let description = [{
     Launch a kernel 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 in a dedicated module, which
     reflects the separate compilation process. The kernel function is required
     to have the `gpu.kernel` attribute. The module containing the kernel
     function is required to be a gpu.module. And finally, the module containing
let description = [{
     launched afterwards. Writes from the device are guaranteed to be visible on
     the host after synchronizing with the device kernel completion.
   }];
 
   let assemblyFormat = "$value attr-dict `:` type($value)";
   let verifier = [{ return success(); }];
 }
 
+def GPU_ForkOp : GPU_Op<"create_token", [GPU_AsyncOpInterface]> {
+  let summary = "Creates a GPU async token.";
+  let description = [{
+    This op creates a new async token from a list of async dependencies.
+
+    The op is inserted during lowering so that each async token is used only
+    once. This makes forks in async execution explicit.
+
+    The op does not imply any host synchronization.
+  }];
+
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
+  let results = (outs GPU_AsyncToken:$asyncToken);
+
+  let assemblyFormat = "(`[` $asyncDependencies^ `]`)? attr-dict";
+}
+
 #endif // GPU_OPS

mlir/lib/Dialect/GPU/CMakeLists.txt

 add_mlir_dialect_library(MLIRGPU
   IR/GPUDialect.cpp
   Transforms/AllReduceLowering.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
   Transforms/ParallelLoopMapper.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU
 
   DEPENDS
   MLIRGPUOpsIncGen
+  MLIRGPUOpInterfacesIncGen
   MLIRGPUPassIncGen
   MLIRParallelLoopMapperAttrGen
   MLIRParallelLoopMapperEnumsGen
 
   LINK_LIBS PUBLIC
   MLIREDSC
   MLIRIR
   MLIRLLVMIR
   MLIRSCF
   MLIRPass
   MLIRSideEffectInterfaces
   MLIRStandard
   MLIRSupport
   MLIRTransformUtils
   )

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/GPUDialect.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/Function.h"
+#include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::gpu;
 
 //===----------------------------------------------------------------------===//
 // GPUDialect
 //===----------------------------------------------------------------------===//
 
 bool GPUDialect::isKernel(Operation *op) {
   UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName());
   return static_cast<bool>(isKernelAttr);
 }
 
 void GPUDialect::initialize() {
+  addTypes<AsyncTokenType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
       >();
 }
 
+Type GPUDialect::parseType(DialectAsmParser &parser) const {
+  // Parse the main keyword for the type.
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword))
+    return Type();
+  MLIRContext *context = getContext();
+
+  // Handle 'async token' types.
+  if (keyword == "async token")
+    return AsyncTokenType::get(context);
+
+  parser.emitError(parser.getNameLoc(), "unknown gpu type: " + keyword);
+  return Type();
+}
+
+void GPUDialect::printType(Type type, DialectAsmPrinter &os) const {
+  TypeSwitch<Type>(type)
+      .Case<AsyncTokenType>([&](Type) { os << "async token"; })
+      .Default([](Type) { llvm_unreachable("unexpected 'gpu' type kind"); });
+}
+
 LogicalResult GPUDialect::verifyOperationAttribute(Operation *op,
                                                    NamedAttribute attr) {
   if (!attr.second.isa<UnitAttr>() ||
       attr.first != getContainerModuleAttrName())
     return success();
 
   auto module = dyn_cast<ModuleOp>(op);
   if (!module)
static LogicalResult verify(LaunchOp op) {
   return success();
 }
 
 // Pretty-print the kernel grid/block size assignment as
 //   (%iter-x, %iter-y, %iter-z) in
 //   (%size-x = %ssa-use, %size-y = %ssa-use, %size-z = %ssa-use)
 // where %size-* and %iter-* will correspond to the body region arguments.
 static void printSizeAssignment(OpAsmPrinter &p, KernelDim3 size,
-                                ValueRange operands, KernelDim3 ids) {
+                                KernelDim3 operands, KernelDim3 ids) {
   p << '(' << ids.x << ", " << ids.y << ", " << ids.z << ") in (";
-  p << size.x << " = " << operands[0] << ", ";
-  p << size.y << " = " << operands[1] << ", ";
-  p << size.z << " = " << operands[2] << ')';
+  p << size.x << " = " << operands.x << ", ";
+  p << size.y << " = " << operands.y << ", ";
+  p << size.z << " = " << operands.z << ')';
 }
 
 static void printLaunchOp(OpAsmPrinter &p, LaunchOp op) {
-  ValueRange operands = op.getOperands();
-
   // Print the launch configuration.
   p << LaunchOp::getOperationName() << ' ' << op.getBlocksKeyword();
-  printSizeAssignment(p, op.getGridSize(), operands.take_front(3),
+  printSizeAssignment(p, op.getGridSize(), op.getGridSizeOperandValues(),
                       op.getBlockIds());
   p << ' ' << op.getThreadsKeyword();
-  printSizeAssignment(p, op.getBlockSize(), operands.slice(3, 3),
+  printSizeAssignment(p, op.getBlockSize(), op.getBlockSizeOperandValues(),
                       op.getThreadIds());
 
   p.printRegion(op.body(), /*printEntryBlockArgs=*/false);
   p.printOptionalAttrDict(op.getAttrs());
 }
 
 // Parse the size assignment blocks for blocks and threads.  These have the form
 //   (%region_arg, %region_arg, %region_arg) in
 //   (%region_arg = %operand, %region_arg = %operand, %region_arg = %operand)
 // where %region_arg are percent-identifiers for the region arguments to be
 // introduced further (SSA defs), and %operand are percent-identifiers for the
 // SSA value uses.
 static ParseResult
static void print(OpAsmPrinter &p, GPUModuleOp op) {
   p << op.getOperationName() << ' ';
   p.printSymbolName(op.getName());
   p.printOptionalAttrDictWithKeyword(op.getAttrs(),
                                      {SymbolTable::getSymbolAttrName()});
   p.printRegion(op.getOperation()->getRegion(0), /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/false);
 }
 
+#include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"

mlir/test/Dialect/GPU/ops.mlir

module attributes {gpu.container_module} {
 
   gpu.module @explicit_attributions {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32, {{.*}}: memref<?xf32>) workgroup({{.*}}: memref<5xf32, 3>) private({{.*}}: memref<5xf32, 5>)
     "gpu.func"() ( {
     ^bb0(%arg0: f32, %arg1: memref<?xf32>, %arg2: memref<5xf32, 3>, %arg3: memref<5xf32, 5>):
       "gpu.return"() : () -> ()
     } ) {gpu.kernel, sym_name = "kernel_1", type = (f32, memref<?xf32>) -> (), workgroup_attributions = 1: i64} : () -> ()
   }
+
+  func @async() {
+    // CHECK-LABEL: func @async
+    // CHECK: %[[t0:.*]] = gpu.create_token
+    %0 = gpu.create_token
+    // CHECK: %[[t1:.*]] = gpu.create_token[%[[t0]]]
+    %1 = gpu.create_token[%0]
+    // CHECK: gpu.create_token[%[[t0]], %[[t1]]]
+    %2 = gpu.create_token[%0, %1]
+    return
+  }
 }