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

[Ivan R. Ivanov via flang-commits]

================
@@ -1248,55 +1287,22 @@ void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
       dynamicSharedMemorySize ? 1 : 0;
   prop.operandSegmentSizes[segmentSizesLen - 2] =
       static_cast<int32_t>(kernelOperands.size());
-  prop.operandSegmentSizes[segmentSizesLen - 1] = 0;
+  prop.operandSegmentSizes[segmentSizesLen - 1] = asyncObject ? 1 : 0;
 }
 
 void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
                          GPUFuncOp kernelFunc, KernelDim3 gridSize,
                          KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
                          ValueRange kernelOperands, Type asyncTokenType,
-                         ValueRange asyncDependencies,
+                         ValueRange asyncDependencies, Value asyncObject,
                          std::optional<KernelDim3> clusterSize) {
   auto kernelModule = kernelFunc->getParentOfType<GPUModuleOp>();
   auto kernelSymbol =
       SymbolRefAttr::get(kernelModule.getNameAttr(),
                          {SymbolRefAttr::get(kernelFunc.getNameAttr())});
   build(builder, result, kernelSymbol, gridSize, getBlockSize,
         dynamicSharedMemorySize, kernelOperands, asyncTokenType,
-        asyncDependencies, clusterSize);
-}
-
-void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
-                         SymbolRefAttr kernel, KernelDim3 gridSize,
-                         KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
-                         ValueRange kernelOperands, Value asyncObject,
-                         std::optional<KernelDim3> clusterSize) {
-  // 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 (clusterSize.has_value())
-    result.addOperands({clusterSize->x, clusterSize->y, clusterSize->z});
-  if (dynamicSharedMemorySize)
-    result.addOperands(dynamicSharedMemorySize);
-  result.addOperands(kernelOperands);
-  if (asyncObject)
-    result.addOperands(asyncObject);
-  Properties &prop = result.getOrAddProperties<Properties>();
-  prop.kernel = kernel;
-  size_t segmentSizesLen = std::size(prop.operandSegmentSizes);
-  // Initialize the segment sizes to 1.
-  llvm::fill(prop.operandSegmentSizes, 1);
-  prop.operandSegmentSizes[0] = 0;
-  if (!clusterSize.has_value()) {
-    prop.operandSegmentSizes[segmentSizesLen - 4] = 0;
-    prop.operandSegmentSizes[segmentSizesLen - 5] = 0;
-    prop.operandSegmentSizes[segmentSizesLen - 6] = 0;
-  }
-  prop.operandSegmentSizes[segmentSizesLen - 3] =
-      dynamicSharedMemorySize ? 1 : 0;
-  prop.operandSegmentSizes[segmentSizesLen - 2] =
-      static_cast<int32_t>(kernelOperands.size());
-  prop.operandSegmentSizes[segmentSizesLen - 1] = asyncObject ? 1 : 0;
+        asyncDependencies, asyncObject, clusterSize);
 }
----------------
ivanradanov wrote:

Currently we had three builders:

```
`gpu::FuncOp`, dep version
`SymbolAttr`, dep version
`SymbolAttr`, stream version
```

This now reduces these to two 

```
`gpu::FuncOp`, common for both dep and stream version
`SymbolAttr`, common for both dep and stream version
```

Having common builder for dep + stream version, allows us to rebuild this op by passing the appropriate arguments one by one (except the argument one needs to change) to the common builder, so I wanted to have that available.

The builder I deleted duplicated a lot of code, and I feel like having a lot of builders starts becoming messy so I decided to simplify them and reduce the code duplication.


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