[MLIR][OpenMP] Add omp.private op

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.

Author: ergawy

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

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"
 

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
+}

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>)
+}
+