[MLIR][OpenMP] Add omp.private op

[ergawy]

This PR adds a new op to the OpenMP dialect: PrivateClauseOp. This op will be later used to model [first]private clauses for differnt OpenMP directives.

This is part of productizing the "delayed privatization" PoC which can be found in #79862.

[llvmbot]

@llvm/pr-subscribers-flang-openmp
@llvm/pr-subscribers-mlir-openmp

@llvm/pr-subscribers-mlir

Author: Kareem Ergawy (ergawy)

Changes

This PR adds a new op to the OpenMP dialect: PrivateClauseOp. This op will be later used to model [first]private clauses for differnt OpenMP directives.

This is part of productizing the "delayed privatization" PoC wich can be found in #79862.

---

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

4 Files Affected:

• (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+80-1)
• (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+55)
• (modified) mlir/test/Dialect/OpenMP/invalid.mlir (+46)
• (added) mlir/test/Dialect/OpenMP/roundtrip.mlir (+13)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index ca36350548577..711faa94bf11a 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -16,6 +16,7 @@
 
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/FunctionInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/IR/SymbolInterfaces.td"
@@ -133,6 +134,84 @@ def DeclareTargetAttr : OpenMP_Attr<"DeclareTarget", "declaretarget"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+//===----------------------------------------------------------------------===//
+// 2.19.4 Data-Sharing Attribute Clauses
+//===----------------------------------------------------------------------===//
+
+def PrivateClauseOp : OpenMP_Op<"private", [
+    IsolatedFromAbove, FunctionOpInterface
+  ]> {
+  let summary = "Outline [first]private logic in a separate op.";
+  let description = [{
+    Using this operation, the dialect can model the data-sharing attributes of
+    `private` and `firstprivate` variables on the IR level. This means that of
+    "eagerly" privatizing variables in the frontend, we can instead model which
+    variables should be privatized and only materialze the privatization when
+    necessary; e.g. directly before lowring to LLVM IR.
+
+    Examples:
+    ---------
+    * `private(x)` would be emitted as:
+    ```mlir
+    omp.private @x.privatizer : (!fir.ref<i32>) -> !fir.ref<i32> {
+    ^bb0(%arg0: !fir.ref<i32>):
+      %0 = fir.alloca i32
+      omp.yield(%0 : !fir.ref<i32>)
+    }
+    ```
+
+    * `firstprivate(x)` would be emitted as:
+    ```mlir
+    omp.private @x.privatizer : (!fir.ref<i32>) -> !fir.ref<i32> {
+    ^bb0(%arg0: !fir.ref<i32>):
+      %0 = fir.alloca i32
+      %1 = fir.load %arg0 : !fir.ref<i32>
+      fir.store %1 to %0 : !fir.ref<i32>
+      omp.yield(%0 : !fir.ref<i32>)
+    }
+    ```
+
+    However, the body of the `omp.private` op really depends on the code-gen
+    done by the emitting frontend. There are no restrictions on the body except
+    for having the yield a value of the same type as the operand.
+
+    Instances of this op would then be used by ops that model directives that
+    accept data-sharing attribute clauses.
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttrOf<FunctionType>:$function_type);
+
+  let regions = (region AnyRegion:$body);
+
+  let builders = [OpBuilder<(ins
+    "::mlir::Type":$privateVarType,
+    "::llvm::StringRef":$privatizerName
+  )>];
+
+  let assemblyFormat = [{
+    $sym_name `:` $function_type $body attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    ::mlir::Region *getCallableRegion() {
+      return &getBody();
+    }
+
+    /// Returns the argument types of this function.
+    ArrayRef<Type> getArgumentTypes() {
+      return getFunctionType().getInputs();
+    }
+
+    /// Returns the result types of this function.
+    ArrayRef<Type> getResultTypes() {
+      return getFunctionType().getResults();
+    }
+  }];
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.6 parallel Construct
 //===----------------------------------------------------------------------===//
@@ -612,7 +691,7 @@ def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
 def YieldOp : OpenMP_Op<"yield",
     [Pure, ReturnLike, Terminator,
      ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
-     "AtomicUpdateOp", "SimdLoopOp"]>]> {
+     "AtomicUpdateOp", "SimdLoopOp", "PrivateClauseOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "omp.yield" yields SSA values from the OpenMP dialect op region and
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 381f17d080419..ed3f78d73af1f 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1594,6 +1594,61 @@ LogicalResult DataBoundsOp::verify() {
   return success();
 }
 
+LogicalResult PrivateClauseOp::verify() {
+  Region &body = getBody();
+  auto argumentTypes = getArgumentTypes();
+  auto resultTypes = getResultTypes();
+
+  if (argumentTypes.empty()) {
+    return emitError() << "'" << getOperationName()
+                       << "' must accept at least one argument.";
+  }
+
+  if (resultTypes.empty()) {
+    return emitError() << "'" << getOperationName()
+                       << "' must return at least one result.";
+  }
+
+  for (Block &block : body) {
+    if (block.empty() || !block.mightHaveTerminator())
+      return mlir::emitError(block.empty() ? getLoc() : block.back().getLoc())
+             << "expected all blocks to have terminators.";
+
+    Operation *terminator = block.getTerminator();
+
+    if (!terminator->hasSuccessors() && !llvm::isa<YieldOp>(terminator))
+      return mlir::emitError(terminator->getLoc())
+             << "expected exit block terminator to be an `omp.yield` op.";
+
+    YieldOp yieldOp = llvm::cast<YieldOp>(terminator);
+    auto yieldedTypes = yieldOp.getResults().getTypes();
+
+    if (yieldedTypes.empty())
+      return mlir::emitError(yieldOp.getLoc())
+             << "'" << getOperationName() << "' must yield a value.";
+
+    bool yieldIsValid = [&]() {
+      if (yieldedTypes.size() != resultTypes.size()) {
+        return false;
+      }
+      for (size_t typeIdx = 0; typeIdx < yieldedTypes.size(); ++typeIdx) {
+        if (yieldedTypes[typeIdx] != resultTypes[typeIdx]) {
+          return false;
+        }
+      }
+
+      return true;
+    }();
+
+    if (!yieldIsValid)
+      return mlir::emitError(yieldOp.getLoc())
+             << "Invalid yielded value. Expected type: " << resultTypes
+             << ", got: " << yieldedTypes;
+  }
+
+  return success();
+}
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
 
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 812b79e35595f..563dfb9223eca 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1738,3 +1738,49 @@ func.func @omp_distribute(%data_var : memref<i32>) -> () {
       "omp.terminator"() : () -> ()
     }) : (memref<i32>) -> ()
 }
+
+// -----
+
+omp.private @x.privatizer : (i32) -> i32 {
+^bb0(%arg0: i32):
+  %0 = arith.constant 0.0 : f32
+  // expected-error @below {{Invalid yielded value. Expected type: 'i32', got: 'f32'}}
+  omp.yield(%0 : f32)
+}
+
+// -----
+
+omp.private @x.privatizer : (i32) -> i32 {
+^bb0(%arg0: i32):
+  // expected-error @below {{'omp.private' must yield a value.}}
+  omp.yield
+}
+
+// -----
+
+// expected-error @below {{'omp.private' must accept at least one argument.}}
+omp.private @x.privatizer : () -> i32 {
+^bb0:
+  omp.yield
+}
+
+// -----
+
+// expected-error @below {{'omp.private' must return at least one result.}}
+omp.private @x.privatizer : (i32) -> () {
+}
+
+// -----
+
+// expected-error @below {{expected all blocks to have terminators.}}
+omp.private @x.privatizer : (i32) -> i32 {
+^bb0(%arg0: i32):
+}
+
+// -----
+
+omp.private @x.privatizer : (i32) -> i32 {
+^bb0(%arg0: i32):
+  // expected-error @below {{expected exit block terminator to be an `omp.yield` op.}}
+  omp.terminator
+}
diff --git a/mlir/test/Dialect/OpenMP/roundtrip.mlir b/mlir/test/Dialect/OpenMP/roundtrip.mlir
new file mode 100644
index 0000000000000..5c6bd63c8d8b5
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/roundtrip.mlir
@@ -0,0 +1,13 @@
+// RUN: fir-opt -verify-diagnostics %s | fir-opt | FileCheck %s
+
+// CHECK: omp.private @x.privatizer : (!fir.ref<i32>) -> !fir.ref<i32> {
+omp.private @x.privatizer : (!fir.ref<i32>) -> !fir.ref<i32> {
+// CHECK: ^bb0(%arg0: {{.*}}):
+^bb0(%arg0: !fir.ref<i32>):
+
+  // CHECK: %0 = fir.alloca i32
+  %0 = fir.alloca i32
+  // CHECK: omp.yield(%0 : !fir.ref<i32>)
+  omp.yield(%0 : !fir.ref<i32>)
+}
+

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[ergawy]

Ping! I think all comments were handled. Can you 🙏 have another look?

[clementval]

This looks mostly ok. I just have a comment about the syntax of the operation.

[kiranchandramohan]

A few comments.

[kiranchandramohan]

Why are the types different for alloc_region and copy_region?

[ergawy]

Because the alloc_region is mandatory for both private and firstprivate while the copy_region is mandatory only for firstprivate. Therefore, the alloc_region has to have at least one block while the copy_region might be empty.

[kiranchandramohan]

Is the mlir prefix here and in several places below required?

[ergawy]

Yes, because I don't want this to get resolved to this->emitError(..) in order to provide more accurate location information when needed.

[ergawy]

Thanks for the review! Handled the comments.

[ergawy]

ergawy deleted a comment from kiranchandramohan

I am not sure where this comes from :). @kiranchandramohan all your comments are either resolved or I replied to. Apologies if I accidentally deleted any of your comments.

[kiranchandramohan]

LG.

[NimishMishra]

Thank you for the work on this merge.

I was playing around with the patch, and have one small question. It seems for firstprivate, we are offloading the responsibility of ensuring correct copying elsewhere, which I am not able to pinpoint. Is it the responsibility of the lowering to MLIR to ensure correct copying?

For instance, currently, the following firstprivate op is getting accepted:

omp.private {type=firstprivate} @x.privatizer : i32 alloc{
      ^bb0(%arg0 : i32):
                %0 = arith.constant 0 : i32
                omp.yield(%0 : i32)
} copy {
     ^bb0(%arg0: i32, %arg1: i32):
                    %0 = llvm.add %arg1, %arg0 : i32
                     omp.yield(%0 : i32)
}

So which phase of compilation bears the responsibility of correctly creating the copy region?

[ergawy]

So which phase of compilation bears the responsibility of correctly creating the copy region?

Good question. The way it is now is to leave that responsibility on the front-end. So in this case flang would be responsible to fill in the alloc and copy regions with logic that makes sense to the particular symbol type being privatized.

I think this is the proper way since allocating or copying logic can be quite dependent on the semantics of the language being lowered.