[mlir][GPU] Add gpu.create_queue op and add queue as an optional argument to gpu ops

[Hardcode84]

Summary

Supersedes https://reviews.llvm.org/D146556

The above PR introduces a new op called gpu.create_queue and adds queue as an optional arg to some gpu dialect ops.
Some GPU runtimes (OpenCL/SYCL) are using queue as explicit argument to all ops.
Explicit queue allows more flexibility on scheduling kernels, e.g interleave execution on multiple devices, allow to create separate queues for kernel execution and data copying, etc.
Queues doesn't introduce any additional sychronization semantics to ops, aync dependencies are still controlled by async tokens.

The link to the discussion is:
https://discourse.llvm.org/t/proposal-to-add-stream-queue-as-an-optional-argument-to-few-gpu-dialect-ops/67920/13

[Hardcode84]

check for stream in gpu-to-llvm

[Hardcode84]

update doc

[Hardcode84]

update doc

[Hardcode84]

update doc

[aartbik]

mlir/docs/Dialects/GPU.md

69 ↗

(On Diff #538561)

I don't feel super strong about this, but should we adjust our syntax a bit.
(1) it feels like queue should be up front, just like tokens
(2) since we use [ ... ] for tokens we could use something similar for queue, e.g. <%queue1>,

or, alternatively, use something keyword based like 
      OP tokens %token1 queue %queue

Again, not a blocker, and I have not clean design in mind yet, but it feels like this is mixing and matching styles a bit

[Hardcode84]

rebase, use <> for queue arg

[Hardcode84]

mlir/docs/Dialects/GPU.md

69 ↗

(On Diff #538561)

Yes, I think it looks better.

[Hardcode84]

update doc

[Hardcode84]

ping

[mehdi_amini]

I don't think the discussion we had in the thread converged?

[Hardcode84]

I don't think the discussion we had in the thread converged?

Right, no activity in thread

[mehdi_amini]

mlir/docs/Dialects/GPU.md

44 ↗

(On Diff #540033)

This is back to the discussion in the thread, where I don't think I got clear answers, and I'd rather see it address there.

Right now it is claimed that the queues "allow more fine-grained control over kernel scheduling", but I don't quite get this argument: since they are entirely out-of-order it's not clear to me where is the "fine-grained control" here?

• Have different queues for kernel execution and data copying.

Why is it useful? Since queues are out-of-order, other than mapping to a given device, I don't see the semantic different between one queue or two here.

[Hardcode84]

mlir/docs/Dialects/GPU.md

44 ↗

(On Diff #540033)

Our main usecase is running on multiple devices, which you cannot express with current dialect at all. Parallelizing execution and copying is handled automatically in SYCL, yes, but for some more low level APIs (Intel Level Zero) you will need to create two separate queues (which SYCL does this under the hood).

[mehdi_amini]

mlir/docs/Dialects/GPU.md

44 ↗

(On Diff #540033)

Our main usecase is running on multiple devices, which you cannot express with current dialect at all.

OK, but we're mixing two things here here: 1) how to map launches to a specific device and 2) scheduling aspects.

You're still not addressing my questions clearly: neither here, nor in the discourse thread, it's not clear to me how we'll converge on this patch right now.

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

 //===----------------------------------------------------------------------===//
 // GPU Types.
 //===----------------------------------------------------------------------===//
 
 def GPU_AsyncToken : DialectType<
   GPU_Dialect, CPred<"::llvm::isa<::mlir::gpu::AsyncTokenType>($_self)">, "async token type">,
              BuildableType<"mlir::gpu::AsyncTokenType::get($_builder.getContext())">;
 
+def GPU_Queue : DialectType<
+  GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::QueueType>()">, "queue type">,
+             BuildableType<"mlir::gpu::QueueType::get($_builder.getContext())">;
+
 // Predicat to check if type is gpu::MMAMatrixType.
 def IsMMAMatrixTypePred : CPred<"::llvm::isa<::mlir::gpu::MMAMatrixType>($_self)">;
 
 def GPU_MMAMatrix : DialectType<
   GPU_Dialect, IsMMAMatrixTypePred, "MMAMatrix type">;
 
 // Memref type acceptable to gpu.subgroup_mma_{load|store}_matrix ops.
 def GPU_MMAMemRef : MemRefOf<[I8, I32, F16, F32, VectorOfRankAndType<[1], [I8, I32, F16, F32]>]>;

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

 
 class AsyncTokenType
     : public Type::TypeBase<AsyncTokenType, Type, TypeStorage> {
 public:
   // Used for generic hooks in TypeBase.
   using Base::Base;
 };
 
+class QueueType : public Type::TypeBase<QueueType, Type, TypeStorage> {
+public:
+  // Used for generic hooks in TypeBase.
+  using Base::Base;
+};
+
 /// MMAMatrixType storage and uniquing. Array is uniqued based on its shape
 /// and type.
 struct MMAMatrixStorageType : public TypeStorage {
   MMAMatrixStorageType(unsigned numDims, const int64_t *dimShapes,
                        Type elementType, StringRef operand)
       : dimShapes(dimShapes), numDims(numDims), elementType(elementType),
         operand(operand) {}
 

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

 
 def GPU_LaunchFuncOp : GPU_Op<"launch_func",
                               [GPU_AsyncOpInterface, AttrSizedOperandSegments]>,
     Arguments<(ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                SymbolRefAttr:$kernel,
                Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
                Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,
                Optional<I32>:$dynamicSharedMemorySize,
-               Variadic<AnyType>:$kernelOperands)>,
+               Variadic<AnyType>:$kernelOperands,
+               Optional<GPU_Queue>:$queue)>,
     Results<(outs Optional<GPU_AsyncToken>:$asyncToken)> {
   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
module attributes {gpu.container_module} {
                %arg1 : memref<?xf32, 1>)
     }
     ```
   }];
 
   let skipDefaultBuilders = 1;
 
   let builders = [
-    OpBuilder<(ins "GPUFuncOp":$kernelFunc, "KernelDim3":$gridSize,
-      "KernelDim3":$blockSize, "Value":$dynamicSharedMemorySize,
+    OpBuilder<(ins "GPUFuncOp":$kernelFunc,
+      "KernelDim3":$gridSize, "KernelDim3":$blockSize,
+      "Value":$dynamicSharedMemorySize,
       "ValueRange":$kernelOperands,
       CArg<"Type", "nullptr">:$asyncTokenType,
-      CArg<"ValueRange", "{}">:$asyncDependencies)>
-  ];
+      CArg<"ValueRange", "{}">:$asyncDependencies,
+      CArg<"Value", "Value{}">:$queue)>];
 
   let extraClassDeclaration = [{
     /// The name of the kernel's containing module.
     StringAttr getKernelModuleName();
 
     /// The name of the kernel.
     StringAttr getKernelName();
 
friend LogicalResult GPUDialect::verifyOperationAttribute(Operation *,
                                                               NamedAttribute);
   }];
 
   let assemblyFormat = [{
       custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
       $kernel
       `blocks` `in` ` ` `(`$gridSizeX`,` $gridSizeY`,` $gridSizeZ`)`
       `threads` `in` ` ` `(`$blockSizeX`,` $blockSizeY`,` $blockSizeZ`)`
+      (`queue` $queue^)?
       (`dynamic_shared_memory_size` $dynamicSharedMemorySize^)?
       custom<LaunchFuncOperands>($kernelOperands, type($kernelOperands)) attr-dict
   }];
   let hasVerifier = 1;
 }
 
 def GPU_LaunchOp : GPU_Op<"launch", [
       AutomaticAllocationScope, AttrSizedOperandSegments, GPU_AsyncOpInterface,
let description = [{
 
       This operation may not be supported in every environment, there is not yet a
           way to check at runtime whether this feature is supported.
   }];
 
   let assemblyFormat = "$value attr-dict `:` type($value)";
 }
 
-def GPU_WaitOp : GPU_Op<"wait", [GPU_AsyncOpInterface]> {
+def GPU_WaitOp : GPU_Op<"wait", [GPU_AsyncOpInterface, AttrSizedOperandSegments]> {
   let summary = "Wait for async gpu ops to complete.";
   let description = [{
     This op synchronizes the host or the device with a list of dependent ops.
 
     If the op contains the `async` keyword, it returns a new async token which
     is synchronized with the op arguments. This new token is merely a shortcut
     to the argument list, and one could replace the uses of the result with the
     arguments for the same effect. The async version of this op is primarily
let description = [{
     ```mlir
     %t0 = gpu.foo async : !gpu.async.token
     %t1 = gpu.bar async : !gpu.async.token
     // The gpu.wait op blocks until gpu.foo and gpu.bar have completed.
     gpu.wait [%t0, %t1]
     ```
   }];
 
-  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+                   Optional<GPU_Queue>:$queue);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
+  let skipDefaultBuilders = 1;
+
+  let builders = [
+    OpBuilder<(ins
+      CArg<"Type", "nullptr">:$asyncTokenType,
+      CArg<"ValueRange", "{}">:$asyncDependencies,
+      CArg<"Value", "Value{}">:$queue)>];
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) attr-dict
+    (`queue` $queue^)?
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def GPU_AllocOp : GPU_Op<"alloc", [
     GPU_AsyncOpInterface,
     AttrSizedOperandSegments
let description = [{
 
     ```mlir
     %memref, %token = gpu.alloc async [%dep] host_shared (%width) : memref<64x?xf32, 1>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Variadic<Index>:$dynamicSizes, Variadic<Index>:$symbolOperands,
-                   UnitAttr:$hostShared);
+                   UnitAttr:$hostShared, Optional<GPU_Queue>:$queue);
   let results = (outs Res<AnyMemRef, "", [MemAlloc]>:$memref,
                  Optional<GPU_AsyncToken>:$asyncToken);
 
   let extraClassDeclaration = [{
     MemRefType getType() { return ::llvm::cast<MemRefType>(getMemref().getType()); }
   }];
 
+  let skipDefaultBuilders = 1;
+
+  let builders = [
+    OpBuilder<(ins
+      "Type":$memref,
+      "Type":$asyncTokenType,
+      "ValueRange":$asyncDependencies,
+      "ValueRange":$dynamicSizes,
+      "ValueRange":$symbolOperands,
+      CArg<"bool", "false">:$hostShared,
+      CArg<"Value", "Value{}">:$queue)>];
+
   let assemblyFormat = [{
-    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) (` ` `host_shared` $hostShared^)? ` `
-    `(` $dynamicSizes `)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
+    (` ` `host_shared` $hostShared^)? ` `
+    `(` $dynamicSizes `)` (`` `[` $symbolOperands^ `]`)?
+    (` ` `queue` $queue^)? ` ` attr-dict `:` type($memref)
   }];
 
   let hasVerifier = 1;
   let hasCanonicalizer = 1;
 }
 
-def GPU_DeallocOp : GPU_Op<"dealloc", [GPU_AsyncOpInterface]> {
+def GPU_DeallocOp : GPU_Op<"dealloc", [GPU_AsyncOpInterface, AttrSizedOperandSegments]> {
 
   let summary = "GPU memory deallocation operation";
 
   let description = [{
     The `gpu.dealloc` operation frees the region of memory referenced by a
     memref which was originally created by the `gpu.alloc` operation. It is
     similar to the `memref.dealloc` op, but supports asynchronous GPU execution.
 
     The op does not execute before all async dependencies have finished
     executing.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token.
 
     Example:
 
     ```mlir
     %token = gpu.dealloc async [%dep] %memref : memref<8x64xf32, 1>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<AnyMemRef, "", [MemFree]>:$memref);
+                   Arg<AnyMemRef, "", [MemFree]>:$memref,
+                   Optional<GPU_Queue>:$queue);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
+  let skipDefaultBuilders = 1;
+
+  let builders = [
+    OpBuilder<(ins
+      "Type":$asyncTokenType,
+      "ValueRange":$asyncDependencies,
+      "Value":$memref,
+      CArg<"Value", "Value{}">:$queue)>];
+
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $memref attr-dict `:` type($memref)
+    $memref
+    (`queue` $queue^)? attr-dict `:` type($memref)
   }];
 }
 
-def GPU_MemcpyOp : GPU_Op<"memcpy", [GPU_AsyncOpInterface]> {
+def GPU_MemcpyOp : GPU_Op<"memcpy", [GPU_AsyncOpInterface, AttrSizedOperandSegments]> {
 
   let summary = "GPU memcpy operation";
 
   let description = [{
     The `gpu.memcpy` operation copies the content of one memref to another.
 
     The op does not execute before all async dependencies have finished
     executing.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token.
 
     Example:
 
     ```mlir
     %token = gpu.memcpy async [%dep] %dst, %src : memref<?xf32, 1>, memref<?xf32>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Arg<AnyMemRef, "", [MemWrite]>:$dst,
-                   Arg<AnyMemRef, "", [MemRead]>:$src);
+                   Arg<AnyMemRef, "", [MemRead]>:$src,
+                   Optional<GPU_Queue>:$queue);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
+  let skipDefaultBuilders = 1;
+
+  let builders = [
+    OpBuilder<(ins
+      "Type":$asyncTokenType,
+      "ValueRange":$asyncDependencies,
+      "Value":$dst,
+      "Value":$src,
+      CArg<"Value", "Value{}">:$queue)>];
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $dst`,` $src `:` type($dst)`,` type($src) attr-dict
+    $dst`,` $src  (`queue` $queue^)? `:` type($dst)`,` type($src) attr-dict
+
   }];
   let hasFolder = 1;
   let hasVerifier = 1;
   let hasCanonicalizer = 1;
 }
 
 def GPU_MemsetOp : GPU_Op<"memset",
-  [GPU_AsyncOpInterface, AllElementTypesMatch<["dst", "value"]>]> {
+  [GPU_AsyncOpInterface, AttrSizedOperandSegments,
+   AllElementTypesMatch<["dst", "value"]>]> {
 
   let summary = "GPU memset operation";
 
   let description = [{
     The `gpu.memset` operation sets the content of memref to a scalar value.
 
     The op does not execute before all async dependencies have finished
     executing.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token.
 
     Example:
 
     ```mlir
     %token = gpu.memset async [%dep] %dst, %value : memref<?xf32, 1>, f32
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Arg<AnyMemRef, "", [MemWrite]>:$dst,
-                   Arg<AnyType, "">:$value);
+                   Arg<AnyType, "">:$value,
+                   Optional<GPU_Queue>:$queue);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $dst`,` $value `:` type($dst)`,` type($value) attr-dict
+    $dst`,` $value (`queue` $queue^)? `:` type($dst)`,` type($value) attr-dict
   }];
   let hasFolder = 1;
 }
 
 def GPU_SetDefaultDeviceOp : GPU_Op<"set_default_device",
                                     [MemoryEffects<[MemWrite]>]>,
     Arguments<(ins I32:$devIndex)> {
   let summary = "Set default GPU for operations after this by index";
let description = [{
     %spmat, %token = gpu.create_2to4_spmat async [%dep] %rows, %cols, %mem : memref<?xf64>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        Index:$rows,
                        Index:$cols,
                        AnyMemRef:$memref);
-  let results = (outs Res<GPU_SparseSpMatHandle>:$spMat, 
+  let results = (outs Res<GPU_SparseSpMatHandle>:$spMat,
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $rows `,` $cols `,` $memref attr-dict `:` type($memref)
   }];
 }
 
def GPU_SDDMMOp : GPU_Op<"sddmm", [GPU_AsyncOpInterface]> {
   ];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $buffer attr-dict `:` type($buffer) `into` $computeType
   }];
 }
 
+def GPU_CreateQueueOp : GPU_Op<"create_queue"> {
+  let description = [{
+   The `gpu.create_queue` takes an optional attribute `desc` as input and
+   returns a new queue. `desc` meaning is implementation specific, e.g. it can
+   be a device name.
+   The queue is then used for launching/queuing kernels
+   on the GPU.
+
+   Example:
+
+   ```mlir
+   %queue = gpu.create_queue "test"
+
+   OR
+
+   %queue = gpu.create_queue
+   ```
+
+  }];
+
+  let skipDefaultBuilders = 1;
+
+  let arguments = (ins OptionalAttr<AnyAttr> : $desc);
+  let results = (outs GPU_Queue : $queue);
+  let builders = [OpBuilder<(ins CArg<"Attribute", "{}">:$desc)>];
+
+  let assemblyFormat = "($desc^)? attr-dict";
+}
+
+def GPU_DestroyQueueOp : GPU_Op<"destroy_queue"> {
+  let description = [{
+   The `gpu.destroy_queue` op destoys the GPU queue.
+
+   Example:
+
+   ```mlir
+   gpu.destroy_queue %queue
+   ```
+  }];
+
+  let arguments = (ins GPU_Queue : $queue);
+
+  let assemblyFormat = "attr-dict $queue";
+}
+
 #endif // GPU_OPS

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

 };
 } // namespace
 
 void GPUDialect::initialize() {
   addTypes<AsyncTokenType>();
   addTypes<MMAMatrixType>();
   addTypes<SparseDnTensorHandleType>();
   addTypes<SparseSpMatHandleType>();
+  addTypes<QueueType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/GPU/IR/GPUOps.cpp.inc"
       >();
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc"
       >();
Type GPUDialect::parseType(DialectAsmParser &parser) const {
   if (parser.parseKeyword(&keyword))
     return Type();
   MLIRContext *context = getContext();
 
   // Handle 'async token' types.
   if (keyword == "async.token")
     return AsyncTokenType::get(context);
 
+  if (keyword == "queue")
+    return QueueType::get(context);
+
   if (keyword == "mma_matrix") {
     SMLoc beginLoc = parser.getNameLoc();
 
     // Parse '<'.
     if (parser.parseLess())
       return nullptr;
 
     // Parse the size and elementType.

 void GPUDialect::printType(Type type, DialectAsmPrinter &os) const {
   TypeSwitch<Type>(type)
       .Case<AsyncTokenType>([&](Type) { os << "async.token"; })
       .Case<SparseDnTensorHandleType>([&](Type) {
         os << getSparseHandleKeyword(SparseHandleKind::DnTensor);
       })
       .Case<SparseSpMatHandleType>(
           [&](Type) { os << getSparseHandleKeyword(SparseHandleKind::SpMat); })
+      .Case<QueueType>([&](Type) { os << "queue"; })
       .Case<MMAMatrixType>([&](MMAMatrixType fragTy) {
         os << "mma_matrix<";
         auto shape = fragTy.getShape();
         for (auto dim = shape.begin(), e = shape.end() - 1; dim != e; ++dim)
           os << *dim << 'x';
         os << shape.back() << 'x' << fragTy.getElementType();
         os << ", \"" << fragTy.getOperand() << "\"" << '>';
       })

 /// private memory.
 BlockArgument LaunchOp::addPrivateAttribution(Type type, Location loc) {
   // Buffers on the private memory always come after buffers on the workgroup
   // memory.
   return getBody().addArgument(type, loc);
 }
 
 //===----------------------------------------------------------------------===//
+// CreateQueueOp
+//===----------------------------------------------------------------------===//
+
+void CreateQueueOp::build(OpBuilder &odsBuilder, OperationState &odsState,
+                          /*optional*/ Attribute desc) {
+  if (desc)
+    odsState.getOrAddProperties<Properties>().desc = desc;
+
+  odsState.addTypes(QueueType::get(odsBuilder.getContext()));
+}
+
+//===----------------------------------------------------------------------===//
 // LaunchFuncOp
 //===----------------------------------------------------------------------===//
 
 void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
                          GPUFuncOp kernelFunc, KernelDim3 gridSize,
                          KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
                          ValueRange kernelOperands, Type asyncTokenType,
-                         ValueRange asyncDependencies) {
+                         ValueRange asyncDependencies, Value queue) {
   result.addOperands(asyncDependencies);
   if (asyncTokenType)
     result.types.push_back(builder.getType<AsyncTokenType>());
 
   // Add grid and block sizes as op operands, followed by the data operands.
   result.addOperands({gridSize.x, gridSize.y, gridSize.z, getBlockSize.x,
                       getBlockSize.y, getBlockSize.z});
   if (dynamicSharedMemorySize)
     result.addOperands(dynamicSharedMemorySize);
   result.addOperands(kernelOperands);
   auto kernelModule = kernelFunc->getParentOfType<GPUModuleOp>();
   auto kernelSymbol =
       SymbolRefAttr::get(kernelModule.getNameAttr(),
                          {SymbolRefAttr::get(kernelFunc.getNameAttr())});
   result.addAttribute(getKernelAttrName(result.name), kernelSymbol);
-  SmallVector<int32_t, 9> segmentSizes(9, 1);
+
+  if (queue)
+    result.addOperands(queue);
+
+  SmallVector<int32_t, 10> segmentSizes(10, 1);
   segmentSizes.front() = asyncDependencies.size();
-  segmentSizes[segmentSizes.size() - 2] = dynamicSharedMemorySize ? 1 : 0;
-  segmentSizes.back() = static_cast<int32_t>(kernelOperands.size());
+  segmentSizes[segmentSizes.size() - 3] = dynamicSharedMemorySize ? 1 : 0;
+  segmentSizes[segmentSizes.size() - 2] =
+      static_cast<int32_t>(kernelOperands.size());
+  segmentSizes.back() = queue ? 1 : 0;
   result.addAttribute(getOperandSegmentSizeAttr(),
                       builder.getDenseI32ArrayAttr(segmentSizes));
 }
 
 StringAttr LaunchFuncOp::getKernelModuleName() {
   return getKernel().getRootReference();
 }
 

 
 } // end anonymous namespace
 
 void MemcpyOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
   results.add<EraseTrivialCopyOp>(context);
 }
 
+void MemcpyOp::build(OpBuilder &odsBuilder, OperationState &odsState,
+                     /*optional*/ Type asyncToken, ValueRange asyncDependencies,
+                     Value dst, Value src, /*optional*/ Value queue) {
+  odsState.addOperands(asyncDependencies);
+  odsState.addOperands(dst);
+  odsState.addOperands(src);
+  if (queue)
+    odsState.addOperands(queue);
+  (odsState.getOrAddProperties<Properties>().operand_segment_sizes =
+       odsBuilder.getDenseI32ArrayAttr(
+           {static_cast<int32_t>(asyncDependencies.size()), 1, 1,
+            (queue ? 1 : 0)}));
+  if (asyncToken)
+    odsState.addTypes(asyncToken);
+}
+
 //===----------------------------------------------------------------------===//
 // GPU_SubgroupMmaLoadMatrixOp
 //===----------------------------------------------------------------------===//
 
 /// Return true if the last dimension of the MemRefType has unit stride. Also
 /// return true for memrefs with no strides.
 static bool isLastMemrefDimUnitStride(MemRefType type) {
   int64_t offset;
LogicalResult MemsetOp::fold(FoldAdaptor adaptor,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return memref::foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // GPU_WaitOp
 //===----------------------------------------------------------------------===//
 
+void WaitOp::build(OpBuilder &builder, OperationState &result,
+                   Type asyncTokenType, ValueRange asyncDependencies,
+                   Value queue) {
+  result.addOperands(asyncDependencies);
+  if (asyncTokenType)
+    result.types.push_back(builder.getType<AsyncTokenType>());
+
+  if (queue)
+    result.addOperands(queue);
+
+  SmallVector<int32_t, 2> segmentSizes(2, 1);
+  segmentSizes.front() = asyncDependencies.size();
+  segmentSizes.back() = queue ? 1 : 0;
+  result.addAttribute(getOperandSegmentSizeAttr(),
+                      builder.getDenseI32ArrayAttr(segmentSizes));
+}
+
 namespace {
 
 /// Remove gpu.wait op use of gpu.wait op def without async dependencies.
 /// %t = gpu.wait async []       // No async dependencies.
 /// ...  gpu.wait ... [%t, ...]  // %t can be removed.
 struct EraseRedundantGpuWaitOpPairs : public OpRewritePattern<WaitOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(WaitOp op,
                                 PatternRewriter &rewriter) const final {
     auto predicate = [](Value value) {
       auto waitOp = value.getDefiningOp<WaitOp>();
-      return waitOp && waitOp->getNumOperands() == 0;
+      return waitOp && waitOp.getAsyncDependencies().size() == 0;
     };
+
     if (llvm::none_of(op.getAsyncDependencies(), predicate))
       return failure();
+
     SmallVector<Value> validOperands;
-    for (Value operand : op->getOperands()) {
+    for (Value operand : op.getAsyncDependencies()) {
       if (predicate(operand))
         continue;
       validOperands.push_back(operand);
     }
-    rewriter.updateRootInPlace(op, [&]() { op->setOperands(validOperands); });
+
+    rewriter.updateRootInPlace(
+        op, [&]() { op.getAsyncDependenciesMutable().assign(validOperands); });
     return success();
   }
 };
 
 /// Simplify trivial gpu.wait ops for the following patterns.
 /// 1. %t = gpu.wait async ... ops, where %t has no uses (regardless of async
 /// dependencies).
 /// 2. %t1 = gpu.wait async [%t0], in this case, we can replace uses of %t1 with

 
 } // namespace
 
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<SimplifyDimOfAllocOp>(context);
 }
 
+void AllocOp::build(OpBuilder &odsBuilder, OperationState &odsState,
+                    Type memref,
+                    /*optional*/ Type asyncToken, ValueRange asyncDependencies,
+                    ValueRange dynamicSizes, ValueRange symbolOperands,
+                    /*optional*/ bool hostShared,
+                    /*optional*/ Value queue) {
+  odsState.addOperands(asyncDependencies);
+  odsState.addOperands(dynamicSizes);
+  odsState.addOperands(symbolOperands);
+  if (queue)
+    odsState.addOperands(queue);
+  (odsState.getOrAddProperties<Properties>().operand_segment_sizes =
+       odsBuilder.getDenseI32ArrayAttr(
+           {static_cast<int32_t>(asyncDependencies.size()),
+            static_cast<int32_t>(dynamicSizes.size()),
+            static_cast<int32_t>(symbolOperands.size()), (queue ? 1 : 0)}));
+  if (hostShared) {
+    odsState.getOrAddProperties<Properties>().hostShared =
+        ((hostShared) ? odsBuilder.getUnitAttr() : nullptr);
+  }
+  odsState.addTypes(memref);
+  if (asyncToken)
+    odsState.addTypes(asyncToken);
+}
+
+//===----------------------------------------------------------------------===//
+// GPU_DeallocOp
+//===----------------------------------------------------------------------===//
+
+void DeallocOp::build(OpBuilder &odsBuilder, OperationState &odsState,
+                      /*optional*/ Type asyncToken,
+                      ValueRange asyncDependencies, Value memref,
+                      /*optional*/ Value queue) {
+  odsState.addOperands(asyncDependencies);
+  odsState.addOperands(memref);
+  if (queue)
+    odsState.addOperands(queue);
+  (odsState.getOrAddProperties<Properties>().operand_segment_sizes =
+       odsBuilder.getDenseI32ArrayAttr(
+           {static_cast<int32_t>(asyncDependencies.size()), 1,
+            (queue ? 1 : 0)}));
+  if (asyncToken)
+    odsState.addTypes(asyncToken);
+}
+
 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.cpp.inc"
 #include "mlir/Dialect/GPU/IR/GPUOpsEnums.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOps.cpp.inc"

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

static gpu::AllocOp genAllocMemRef(OpBuilder &builder, Location loc, Value mem,
   auto memTp = MemRefType::get(shape, elemTp);
   SmallVector<Value> dynamicSizes;
   for (unsigned r = 0, rank = tp.getRank(); r < rank; r++) {
     if (shape[r] == ShapedType::kDynamic) {
       Value dimOp = linalg::createOrFoldDimOp(builder, loc, mem, r);
       dynamicSizes.push_back(dimOp);
     }
   }
-  return builder.create<gpu::AllocOp>(loc, TypeRange({memTp, token.getType()}),
-                                      token, dynamicSizes, ValueRange());
+  return builder.create<gpu::AllocOp>(loc, memTp, token.getType(), token,
+                                      dynamicSizes,
+                                      /*symbolOperands*/ ValueRange());
 }
 
 // Allocates a void buffer on the device with given size.
 static gpu::AllocOp genAllocBuffer(OpBuilder &builder, Location loc, Value size,
                                    Value token) {
   const auto memTp =
       MemRefType::get({ShapedType::kDynamic}, builder.getI8Type());
-  return builder.create<gpu::AllocOp>(loc, TypeRange({memTp, token.getType()}),
-                                      token, size, ValueRange());
+  return builder.create<gpu::AllocOp>(loc, memTp, token.getType(), token, size,
+                                      /*symbolOperands*/ ValueRange());
 }
 
 /// Deallocates memory from the device.
 static Value genDeallocMemRef(OpBuilder &builder, Location loc, Value mem,
                               Value token) {
   return builder.create<gpu::DeallocOp>(loc, token.getType(), token, mem)
       .getAsyncToken();
 }

mlir/test/Dialect/GPU/ops.mlir

module attributes {gpu.container_module} {
   func.func @foo() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<?xf32, 1>)
     // CHECK: %{{.*}} = arith.constant 8
     %cst = arith.constant 8 : index
     %c0 = arith.constant 0 : i32
     %t0 = gpu.wait async
 
+    // CHECK: %[[queue:.*]] = gpu.create_queue
+    %queue = gpu.create_queue
+
     // CHECK: gpu.launch_func @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
     gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%0 : f32, %1 : memref<?xf32, 1>)
 
+    gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) queue %queue args(%0 : f32, %1 : memref<?xf32, 1>)
+
+    gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) queue %queue dynamic_shared_memory_size %c0 args(%0 : f32, %1 : memref<?xf32, 1>)
+
     gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) dynamic_shared_memory_size %c0 args(%0 : f32, %1 : memref<?xf32, 1>)
 
     // CHECK: gpu.launch_func @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})
     gpu.launch_func @kernels::@kernel_2 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
 
     // CHECK: %{{.*}} = gpu.launch_func async [%{{.*}}] @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})
     %t1 = gpu.launch_func async [%t0] @kernels::@kernel_2  blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
 
     // CHECK: %[[VALUES:.*]]:2 = call
     %values:2 = func.call @two_value_generator() : () -> (f32, memref<?xf32, 1>)
     // CHECK: gpu.launch_func @kernels::@kernel_1 {{.*}} args(%[[VALUES]]#0 : f32, %[[VALUES]]#1 : memref<?xf32, 1>)
     gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%values#0 : f32, %values#1 : memref<?xf32, 1>)
 
+    // CHECK: gpu.destroy_queue %[[queue]]
+    gpu.destroy_queue %queue
     return
   }
 
   gpu.module @gpu_funcs {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
     // CHECK:       workgroup
     // CHECK:       private
     // CHECK:       attributes
gpu.module @explicit_attributions {
     ^bb0(%arg0: f32, %arg1: memref<?xf32>, %arg2: memref<5xf32, 3>, %arg3: memref<5xf32, 5>):
       "gpu.return"() : () -> ()
     } ) {function_type = (f32, memref<?xf32>) -> (), gpu.kernel, sym_name = "kernel_1", workgroup_attributions = 1: i64} : () -> ()
   }
 
   func.func @alloc() {
     // CHECK-LABEL: func @alloc()
 
+
     // CHECK: %[[m0:.*]] = gpu.alloc () : memref<13xf32, 1>
     %m0 = gpu.alloc () : memref<13xf32, 1>
     // CHECK: gpu.dealloc %[[m0]] : memref<13xf32, 1>
     gpu.dealloc %m0 : memref<13xf32, 1>
 
     %t0 = gpu.wait async
     // CHECK: %[[m1:.*]], %[[t1:.*]] = gpu.alloc async [{{.*}}] () : memref<13xf32, 1>
     %m1, %t1 = gpu.alloc async [%t0] () : memref<13xf32, 1>
     // CHECK: gpu.dealloc async [%[[t1]]] %[[m1]] : memref<13xf32, 1>
     %t2 = gpu.dealloc async [%t1] %m1 : memref<13xf32, 1>
 
     // CHECK: %[[m2:.*]] = gpu.alloc host_shared () : memref<13xf32, 1>
     %m2 = gpu.alloc host_shared () : memref<13xf32, 1>
     // CHECK: gpu.dealloc %[[m2]] : memref<13xf32, 1>
     gpu.dealloc %m2 : memref<13xf32, 1>
 
+    // CHECK: %[[queue:.*]] = gpu.create_queue
+    %queue = gpu.create_queue
+    // CHECK: %[[m3:.*]] = gpu.alloc () queue %[[queue]] : memref<13xf32, 1>
+    %m3 = gpu.alloc () queue %queue  : memref<13xf32, 1>
+    // CHECK: gpu.dealloc %[[m3]] queue %[[queue]] : memref<13xf32, 1>
+    gpu.dealloc %m3 queue %queue : memref<13xf32, 1>
+
     return
   }
 
   func.func @async_token(%arg0 : !gpu.async.token) -> !gpu.async.token {
     // CHECK-LABEL: func @async_token({{.*}}: !gpu.async.token)
     // CHECK: return {{.*}} : !gpu.async.token
     return %arg0 : !gpu.async.token
   }
module attributes {gpu.container_module} {
   func.func @memcpy(%dst : memref<3x7xf32>, %src : memref<3x7xf32, 1>) {
     // CHECK-LABEL: func @memcpy
     // CHECK: gpu.memcpy {{.*}}, {{.*}} : memref<3x7xf32>, memref<3x7xf32, 1>
     gpu.memcpy %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>
     // CHECK: %[[t0:.*]] = gpu.wait async
     %0 = gpu.wait async
     // CHECK: {{.*}} = gpu.memcpy async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, memref<3x7xf32, 1>
     %1 = gpu.memcpy async [%0] %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>
+
+    // CHECK: %[[queue:.*]] = gpu.create_queue
+    %queue = gpu.create_queue
+    // CHECK: gpu.memcpy {{.*}}, {{.*}} queue %[[queue]] : memref<3x7xf32>, memref<3x7xf32, 1>
+    gpu.memcpy %dst, %src queue %queue : memref<3x7xf32>, memref<3x7xf32, 1>
+    // CHECK: %[[t1:.*]] = gpu.wait async queue %[[queue]]
+    %2 = gpu.wait async queue %queue
+    // CHECK: {{.*}} = gpu.memcpy async [%[[t1]]] {{.*}}, {{.*}} queue %[[queue]] : memref<3x7xf32>, memref<3x7xf32, 1>
+    %3 = gpu.memcpy async [%2] %dst, %src queue %queue : memref<3x7xf32>, memref<3x7xf32, 1>
+
     return
   }
 
   func.func @memset(%dst : memref<3x7xf32>, %value : f32) {
     // CHECK-LABEL: func @memset
     // CHECK: gpu.memset {{.*}}, {{.*}} : memref<3x7xf32>, f32
     gpu.memset %dst, %value : memref<3x7xf32>, f32
     // CHECK: %[[t0:.*]] = gpu.wait async
     %0 = gpu.wait async
     // CHECK: {{.*}} = gpu.memset async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, f32
     %1 = gpu.memset async [%0] %dst, %value : memref<3x7xf32>, f32
+
+    // CHECK: %[[queue:.*]] = gpu.create_queue
+    %queue = gpu.create_queue
+    // CHECK: gpu.memset {{.*}}, {{.*}} queue %[[queue]] : memref<3x7xf32>, f32
+    gpu.memset %dst, %value queue %queue : memref<3x7xf32>, f32
+    // CHECK: %[[t1:.*]] = gpu.wait async queue %[[queue]]
+    %2 = gpu.wait async queue %queue
+    // CHECK: {{.*}} = gpu.memset async [%[[t1]]] {{.*}}, {{.*}} queue %[[queue]] : memref<3x7xf32>, f32
+    %3 = gpu.memset async [%2] %dst, %value queue %queue : memref<3x7xf32>, f32
+
     return
   }
 
   func.func @mmamatrix_valid_scalar_element_type(%src : memref<32x32xf16, affine_map<(d0, d1) -> (d0 * 64 + d1)>>){
     // CHECK-LABEL: func @mmamatrix_valid_scalar_element_type
     %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
     // CHECK: %[[wg:.*]] = memref.alloca()
     %i = arith.constant 16 : index
func.func @sparse_ops(%arg0: index) {
     // CHECK: gpu.destroy_sp_mat async
     %token15 = gpu.destroy_sp_mat async [%token14] %spmat
     // CHECK: gpu.destroy_dn_tensor async
     %token16 = gpu.destroy_dn_tensor async [%token15] %dnvec
     // CHECK: gpu.wait
     gpu.wait [%token16]
     return
   }
+
+  // CHECK-LABEL: func @create_queue_with_desc
+  func.func @create_queue_with_desc() {
+    // CHECK: gpu.create_queue "test"
+    %queue = gpu.create_queue "test"
+    return
+  }
 }
 
 // Just check that this doesn't crash.
 gpu.module @module {
   "gpu.func"() ({
     gpu.return
   }) {function_type = () -> (), sym_name = "func"} : () -> ()
 }