[Mlir-commits] [mlir] [OpenACC] add acc.combine operation (PR #181853)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-openacc

Author: Scott Manley (rscottmanley)

<details>
<summary>Changes</summary>

To facilitate codegen decisions, we need to create an operation that can abstract the final update of the original and partial sum from a reduction. This is represented within the combiner recipes. Having an operator allows future lowering to clearly identify how to handle the final accumulation. This is currently an NFC.

The format of this operation is:

```
  acc.combine %srcMemref into %destMemref <reductionOperator> : type
```

---
Full diff: https://github.com/llvm/llvm-project/pull/181853.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td (+41) 
- (modified) mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp (+17) 
- (modified) mlir/test/Dialect/OpenACC/ops.mlir (+10) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
index 8f7c8c2a89be9..99f51db6544c9 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
@@ -20,6 +20,47 @@
 // the necessary includes and definitions. The operations defined here use
 // types and definitions from that file.
 
+//===----------------------------------------------------------------------===//
+// acc.combine
+//===----------------------------------------------------------------------===//
+
+def OpenACC_CombineOp: OpenACC_Op<"combine", [SameTypeOperands, DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Combine a reduction partial sum with its original value";
+  let description = [{
+    This operation is a composite to do a typical update of a reduction
+    variable. The intention of this operator is to facilitate codegen
+    decisions (such as generate an atomic update). E.g.
+
+    ```
+      lro.combine %src into %dest <addi> : memref<i32>
+    ```
+
+    Might lower to something similar to
+
+    ```
+      %loadSrc = memref.load %src[] : memref<i32>
+      %loadDest = memref.load %dest[] : memref<i32>
+      %combine = arith.addi %loadSrc, %loadDest : i32
+      memref.store %combine, %dest[] : memref<i32>
+    ```
+    
+    The $destMemref operand is a "pointer" to the original reduction 
+    variable (typically shared). The $srcMemref operand is a "pointer"
+    to the partial sum of the reduction (typically private).
+
+    The $kind is the OpenACC reduction operator that determines how to 
+    accumulate the two values.
+  }];
+
+  let arguments = (ins AnyType:$destMemref,
+                       AnyType:$srcMemref,
+                       OpenACC_ReductionOperatorAttr:$reductionOperator);
+  
+  let assemblyFormat = [{
+    $srcMemref `into` $destMemref $reductionOperator `:` type($destMemref) attr-dict
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // acc.kernel_environment
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 460314f8f678f..33c285d4bfb27 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -5129,6 +5129,23 @@ LogicalResult acc::WaitOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// CombineOp
+//===----------------------------------------------------------------------===//
+void acc::CombineOp::getEffects(
+    llvm::SmallVectorImpl<
+        mlir::SideEffects::EffectInstance<mlir::MemoryEffects::Effect>>
+        &effects) {
+  effects.emplace_back(mlir::MemoryEffects::Read::get(), &getSrcMemrefMutable(),
+                       mlir::SideEffects::DefaultResource::get());
+  effects.emplace_back(mlir::MemoryEffects::Read::get(),
+                       &getDestMemrefMutable(),
+                       mlir::SideEffects::DefaultResource::get());
+  effects.emplace_back(mlir::MemoryEffects::Write::get(),
+                       &getDestMemrefMutable(),
+                       mlir::SideEffects::DefaultResource::get());
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenACC/OpenACCOps.cpp.inc"
 
diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index d31397c15769b..cc9cae98aa277 100644
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/mlir/test/Dialect/OpenACC/ops.mlir
@@ -2494,3 +2494,13 @@ func.func @test_kernel_environment_with_async(%arg0: memref<1024xf32>) {
 // CHECK:           gpu.launch
 // CHECK:             memref.store %{{.*}}, %[[CREATE]]
 // CHECK:         acc.copyout accPtr(%[[CREATE]] : memref<1024xf32>) async(%[[ASYNC]] : i32) to varPtr(%{{.*}} : memref<1024xf32>)
+
+// -----
+
+func.func @test_acc_combine(%arg0 : memref<i32>, %arg1 : memref<i32>) {
+  acc.combine %arg0 into %arg1 <add> : memref<i32>
+  return
+}
+
+// CHECK-LABEL: func @test_acc_combine
+// CHECK:       acc.combine %arg0 into %arg1 <add> : memref<i32>
\ No newline at end of file

``````````

</details>


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