[flang-commits] [flang] [mlir] [mlir][gpu] Add stream-based async mode to gpu.launch (PR #213031)

[Fabian Mora via flang-commits]

================
@@ -479,31 +479,31 @@ struct CUFLaunchOpConversion
       args.push_back(arg);
     }
     mlir::Value dynamicShmemSize = op.getBytes() ? op.getBytes() : zero;
+    mlir::Type tokenType = nullptr;
+    SmallVector<Value, 1> tokens;
+    if (op.getStream()) {
+      tokens.push_back(
+          cuf::StreamCastOp::create(rewriter, loc, op.getStream()));
+      tokenType = tokens.front().getType();
+    }
     auto gpuLaunchOp = mlir::gpu::LaunchFuncOp::create(
         rewriter, loc, kernelName,
         mlir::gpu::KernelDim3{gridSizeX, gridSizeY, gridSizeZ},
         mlir::gpu::KernelDim3{blockSizeX, blockSizeY, blockSizeZ},
-        dynamicShmemSize, args);
+        dynamicShmemSize, args, tokenType, tokens);
     if (clusterDimX && clusterDimY && clusterDimZ) {
       gpuLaunchOp.getClusterSizeXMutable().assign(clusterDimX);
       gpuLaunchOp.getClusterSizeYMutable().assign(clusterDimY);
       gpuLaunchOp.getClusterSizeZMutable().assign(clusterDimZ);
     }
-    if (op.getStream()) {
-      mlir::OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.setInsertionPoint(gpuLaunchOp);
-      mlir::Value stream =
-          cuf::StreamCastOp::create(rewriter, loc, op.getStream());
-      gpuLaunchOp.getAsyncDependenciesMutable().append(stream);
-    }
     if (procAttr)
       gpuLaunchOp->setAttr(cuf::getProcAttrName(), procAttr);
     else
       // Set default global attribute of the original was not found.
       gpuLaunchOp->setAttr(cuf::getProcAttrName(),
                            cuf::ProcAttributeAttr::get(
                                op.getContext(), cuf::ProcAttribute::Global));
-    rewriter.replaceOp(op, gpuLaunchOp);
+    rewriter.eraseOp(op);
----------------
fabianmcg wrote:

Why did it changed from replace to erase?

https://github.com/llvm/llvm-project/pull/213031