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

[nbpatel]

Summary

The above PR introduces a new op called gpu.create_stream and adds stream as an optional arg to some gpu dialect
ops. This will give the user the flexibility to launch kernels on their custom stream. We require this to enable
the MLIR pipeline on Intel GPU.

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

[rengolin]

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

973

If streams are optional, you should have builders without them, but this will force all builder calls to pass a nullptr or something for the stream, and could break existing downstream code.

Is there a reason to drop the existing default builders?

1425

What creates a device? How can I get hold of it?

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

744

No validation for what a device is or should be?

mlir/test/Dialect/GPU/ops.mlir

134

You add functionality to pass a device but offers no test to how a device should be created, used or destroyed.

[bondhugula]

But the gpu.wait async op already creates a stream. Can you expand the commit summary on why this is needed? If you already have a custom stream that you want to pass to your IR, a !gpu.async_token type could be passed.

[Hardcode84]

In D146556#4229385, @bondhugula wrote:

But the gpu.wait async op already creates a stream. Can you expand the commit summary on why this is needed? If you already have a custom stream that you want to pass to your IR, a !gpu.async_token type could be passed.

Stream is separate concept from synchronization, please see my comment in https://discourse.llvm.org/t/proposal-to-add-stream-queue-as-an-optional-argument-to-few-gpu-dialect-ops/67920/17?u=hardcode84

[nbpatel]

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

1425

Device is optional and just a placeholder for now. The user can pass an opaque pointer as device which will be used by create_stream op. We are just exposing the API.

mlir/test/Dialect/GPU/ops.mlir

134

Will add a test.

[nbpatel]

Added a test for creating stream with device

[nbpatel]

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

744

what sort of validation you expect here?

[nbpatel]

Ping for review

[bondhugula]

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

97–101

Both of these need proper doc comments.

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

1418–1420

Reflow - use whole width.

1425

Please document what device's type is/can be and how it's created.

[bondhugula]

What's the current behavior if a gpu dialect op with the stream argument is passed to the gpu-to-llvm pass? Please ensure it doesn't assert/crash.

mlir/test/Dialect/GPU/ops.mlir

228

Drop blank line.

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

 //===----------------------------------------------------------------------===//
 // GPU Types.
 //===----------------------------------------------------------------------===//
 
 def GPU_AsyncToken : DialectType<
   GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::AsyncTokenType>()">, "async token type">,
              BuildableType<"mlir::gpu::AsyncTokenType::get($_builder.getContext())">;
 
+def GPU_Stream : DialectType<
+  GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::StreamType>()">, "stream type">,
+             BuildableType<"mlir::gpu::StreamType::get($_builder.getContext())">;
+
+def GPU_Device : DialectType<
+  GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::DeviceType>()">, "device type">,
+             BuildableType<"mlir::gpu::DeviceType::get($_builder.getContext())">;
+
 // Predicat to check if type is gpu::MMAMatrixType.
 def IsMMAMatrixTypePred : CPred<"$_self.isa<::mlir::gpu::MMAMatrixType>()">;
 
 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 StreamType : public Type::TypeBase<StreamType, Type, TypeStorage> {
+public:
+  // Used for generic hooks in TypeBase.
+  using Base::Base;
+};
+
+class DeviceType : public Type::TypeBase<DeviceType, 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_Stream>:$stream)>,
     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{}">:$stream)>];
 
   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`)`
+      (`stream` $stream^)?
       (`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 = [{
     Writes from the host are guaranteed to be visible to device kernels that are
     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)";
 }
 
-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_Stream>:$stream);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
+  let skipDefaultBuilders = 1;
+
+  let builders = [
+    OpBuilder<(ins 
+      CArg<"Type", "nullptr">:$asyncTokenType,
+      CArg<"ValueRange", "{}">:$asyncDependencies,
+      CArg<"Value", "Value{}">:$stream)>];
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) attr-dict
+    (`stream` $stream^)?
   }];
 
   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_Stream>:$stream);
   let results = (outs Res<AnyMemRef, "", [MemAlloc]>:$memref,
                  Optional<GPU_AsyncToken>:$asyncToken);
 
   let extraClassDeclaration = [{
     MemRefType getType() { return getMemref().getType().cast<MemRefType>(); }
   }];
 
   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^ `]`)?
+    (` ` `stream` $stream^)? ` ` 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_Stream>:$stream);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $memref attr-dict `:` type($memref)
+    $memref
+    (`stream` $stream^)? 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_Stream>:$stream);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $dst`,` $src `:` type($dst)`,` type($src) attr-dict
+    $dst`,` $src  (`stream` $stream^)? `:` 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_Stream>:$stream);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $dst`,` $value `:` type($dst)`,` type($value) attr-dict
+    $dst`,` $value (`stream` $stream^)? `:` 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 extraClassDeclaration = [{
     }
   }];
 
   let assemblyFormat = [{
     $opType $args attr-dict `:` functional-type($args, $res)
   }];
 }
 
+
+def GPU_CreateStreamOp : GPU_Op<"create_stream", [SameVariadicOperandSize]> {
+
+  let description = [{ 
+   The `gpu.create_stream` takes an optional argument `device` as input and
+   returns a stream, based on the device. If no device is provided, a default
+   device will be created by the underlying runtime.
+   The stream is then used for launching/queuing kernels
+   on the GPU.
+
+   Example:
+
+   ```mlir
+   %stream = gpu.create_stream %device : !gpu.stream
+
+   OR 
+
+   %stream = gpu.create_stream : !gpu.stream
+   ```
+
+  }];
+  
+  let skipDefaultBuilders = 1;
+
+  let arguments = (ins Optional<GPU_Device> : $device);
+  let results = (outs GPU_Stream : $stream);
+  let builders = [OpBuilder<(ins CArg<"Value", "Value{}">:$device)>];
+
+  let assemblyFormat = [{
+      ($device^)? attr-dict `:` type($stream)
+  }];
+
+}
+
 #endif // GPU_OPS

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

bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 };
 } // namespace
 
 void GPUDialect::initialize() {
   addTypes<AsyncTokenType>();
   addTypes<MMAMatrixType>();
+  addTypes<StreamType>();
+  addTypes<DeviceType>();
   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"
       >();
   addInterfaces<GPUInlinerInterface>();
 }
 
 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);
 
+  if (keyword == "stream")
+    return StreamType::get(context);
+
+  if (keyword == "device")
+    return DeviceType::get(context);
+
   if (keyword == "mma_matrix") {
     SMLoc beginLoc = parser.getNameLoc();
 
     // Parse '<'.
     if (parser.parseLess())
       return nullptr;
 
     // Parse the size and elementType.
Type GPUDialect::parseType(DialectAsmParser &parser) const {
 
   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"; })
+      .Case<StreamType>([&](Type) { os << "stream"; })
+      .Case<DeviceType>([&](Type) { os << "device"; })
       .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() << "\"" << '>';
       })

 };
 
 void LaunchOp::getCanonicalizationPatterns(RewritePatternSet &rewrites,
                                            MLIRContext *context) {
   rewrites.add<FoldLaunchArguments>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// CreateStreamOp
+//===----------------------------------------------------------------------===//
+
+void CreateStreamOp::build(OpBuilder &odsBuilder, OperationState &result,
+                           Value device) {
+  if (device)
+    result.addOperands(device);
+
+  SmallVector<int32_t, 1> segmentSizes(1, 1);
+  segmentSizes.front() = device ? 1 : 0;
+}
+
+//===----------------------------------------------------------------------===//
 // 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 stream) {
   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 (stream)
+    result.addOperands(stream);
+
+  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() = stream ? 1 : 0;
   result.addAttribute(getOperandSegmentSizeAttr(),
                       builder.getDenseI32ArrayAttr(segmentSizes));
 }
 
 StringAttr LaunchFuncOp::getKernelModuleName() {
   return getKernel().getRootReference();
 }
 
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 stream) {
+  result.addOperands(asyncDependencies);
+  if (asyncTokenType)
+    result.types.push_back(builder.getType<AsyncTokenType>());
+
+  if (stream)
+    result.addOperands(stream);
+
+  SmallVector<int32_t, 2> segmentSizes(2, 1);
+  segmentSizes.front() = asyncDependencies.size();
+  segmentSizes.back() = stream ? 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

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: %[[stream:.*]] = gpu.create_stream : !gpu.stream
+    %stream = gpu.create_stream : !gpu.stream
+
     // 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) stream %stream args(%0 : f32, %1 : memref<?xf32, 1>)
+
+    gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) stream %stream 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)
 
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: %[[stream:.*]] = gpu.create_stream : !gpu.stream
+    %stream = gpu.create_stream : !gpu.stream
+    // CHECK: %[[m3:.*]] = gpu.alloc () stream %[[stream]] : memref<13xf32, 1>
+    %m3 = gpu.alloc () stream %stream  : memref<13xf32, 1>
+    // CHECK: gpu.dealloc %[[m3]] stream %[[stream]] : memref<13xf32, 1>
+    gpu.dealloc %m3 stream %stream : 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: %[[stream:.*]] = gpu.create_stream : !gpu.stream
+    %stream = gpu.create_stream : !gpu.stream
+    // CHECK: gpu.memcpy {{.*}}, {{.*}} stream %[[stream]] : memref<3x7xf32>, memref<3x7xf32, 1>
+    gpu.memcpy %dst, %src stream %stream : memref<3x7xf32>, memref<3x7xf32, 1>
+    // CHECK: %[[t1:.*]] = gpu.wait async stream %[[stream]]
+    %2 = gpu.wait async stream %stream
+    // CHECK: {{.*}} = gpu.memcpy async [%[[t1]]] {{.*}}, {{.*}} stream %[[stream]] : memref<3x7xf32>, memref<3x7xf32, 1>
+    %3 = gpu.memcpy async [%2] %dst, %src stream %stream : 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: %[[stream:.*]] = gpu.create_stream : !gpu.stream
+    %stream = gpu.create_stream : !gpu.stream
+    // CHECK: gpu.memset {{.*}}, {{.*}} stream %[[stream]] : memref<3x7xf32>, f32
+    gpu.memset %dst, %value stream %stream : memref<3x7xf32>, f32
+    // CHECK: %[[t1:.*]] = gpu.wait async stream %[[stream]]
+    %2 = gpu.wait async stream %stream
+    // CHECK: {{.*}} = gpu.memset async [%[[t1]]] {{.*}}, {{.*}} stream %[[stream]] : memref<3x7xf32>, f32
+    %3 = gpu.memset async [%2] %dst, %value stream %stream : 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