[Mlir-commits] [mlir] [MLIR][OpenMP] Support `llvm` conversion for `omp.private` regions (PR #81414)

Kareem Ergawy llvmlistbot at llvm.org
Sun Feb 11 06:03:34 PST 2024


https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/81414

Introduces conversion of `omp.private`'s regions to the LLVM dialect.
This reuses the already existing conversion pattern for
`ReducetionDeclareOp` and repurposes it to be used for multi-region ops
as well.

#### This is a follow-up PR for #80955, only the top commit is relevant to this PR.

>From f4ad640189f4b1dad770b0fdaa54675057e448bd Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Tue, 6 Feb 2024 05:41:25 -0600
Subject: [PATCH 1/3] [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.
---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 91 ++++++++++++++++++-
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  | 67 ++++++++++++++
 mlir/test/Dialect/OpenMP/invalid.mlir         | 70 ++++++++++++++
 mlir/test/Dialect/OpenMP/roundtrip.mlir       | 21 +++++
 4 files changed, 248 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/OpenMP/roundtrip.mlir

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index ca363505485773..7c5cfa98dc7118 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -133,6 +133,95 @@ def DeclareTargetAttr : OpenMP_Attr<"DeclareTarget", "declaretarget"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+//===----------------------------------------------------------------------===//
+// 2.19.4 Data-Sharing Attribute Clauses
+//===----------------------------------------------------------------------===//
+
+def DataSharingTypePrivate      : I32EnumAttrCase<"Private", 0, "private">;
+def DataSharingTypeFirstPrivate : I32EnumAttrCase<"FirstPrivate", 1, "firstprivate">;
+
+def DataSharingClauseType : I32EnumAttr<
+    "DataSharingClauseType",
+    "Type of a data-sharing clause",
+    [DataSharingTypePrivate, DataSharingTypeFirstPrivate]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::omp";
+}
+
+def DataSharingClauseTypeAttr : EnumAttr<
+    OpenMP_Dialect, DataSharingClauseType, "data_sharing_type"> {
+  let assemblyFormat = "`{` `type` `=` $value `}`";
+}
+
+def PrivateClauseOp : OpenMP_Op<"private", [
+    IsolatedFromAbove
+  ]> {
+  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 lowering to LLVM IR.
+
+    Examples:
+    ---------
+    * `private(x)` would be emitted as:
+    ```mlir
+    omp.private {type = private} @x.privatizer : !fir.ref<i32> (alloc {
+    ^bb0(%arg0: !fir.ref<i32>):
+    %0 = ... allocate proper memory for the private clone ...
+    omp.yield(%0 : !fir.ref<i32>)
+    })
+    ```
+
+    * `firstprivate(x)` would be emitted as:
+    ```mlir
+    omp.private {type = firstprivate} @y.privatizer : !fir.ref<i32> (alloc {
+    ^bb0(%arg0: !fir.ref<i32>):
+    %0 = ... allocate proper memory for the private clone ...
+    omp.yield(%0 : !fir.ref<i32>)
+    } copy {
+    ^bb0(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i32>):
+    // %arg0 is the original host variable. Same as for `alloc`.
+    // %arg1 represents the memory allocated in `alloc`.
+    ... copy from host to the privatized clone ....
+    omp.yield(%arg1 : !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:
+    - The `alloc` region has a single argument.
+    - The `copy` region has 2 arguments.
+    - Both regions are existed by `omp.yield` ops.
+    The above restrictions and other obvious restrictions (e.g. verifying the
+    type of yielded values) are verified by the custom op verifier. The actual
+    contents of the blocks inside both regions are not verified.
+
+    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<AnyType>:$sym_type,
+                       DataSharingClauseTypeAttr:$data_sharing_type);
+
+  let regions = (region MinSizedRegion<1>:$alloc_region,
+                        AnyRegion:$copy_region);
+
+  let assemblyFormat = [{
+    $data_sharing_type $sym_name `:` $sym_type `(`
+      `alloc` $alloc_region
+      (`copy` $copy_region^)?
+      `)`
+      attr-dict
+  }];
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.6 parallel Construct
 //===----------------------------------------------------------------------===//
@@ -612,7 +701,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 381f17d0804191..96266f8304718c 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1594,6 +1594,73 @@ LogicalResult DataBoundsOp::verify() {
   return success();
 }
 
+LogicalResult PrivateClauseOp::verify() {
+  Type symType = getSymType();
+
+  auto verifyTerminator = [&](Operation *terminator) -> LogicalResult {
+    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);
+    TypeRange yieldedTypes = yieldOp.getResults().getTypes();
+
+    if (yieldedTypes.size() == 1 && yieldedTypes.front() == symType)
+      return success();
+
+    auto error = mlir::emitError(yieldOp.getLoc())
+                 << "Invalid yielded value. Expected type: " << symType
+                 << ", got: ";
+
+    if (yieldedTypes.empty())
+      error << "None";
+    else
+      error << yieldedTypes;
+
+    return error;
+  };
+
+  auto verifyRegion = [&](Region &region, unsigned expectedNumArgs,
+                          StringRef regionName) -> LogicalResult {
+    assert(!region.empty());
+
+    if (region.getNumArguments() != expectedNumArgs)
+      return mlir::emitError(region.getLoc())
+             << "`" << regionName << "`: "
+             << "expected " << expectedNumArgs
+             << " region arguments, got: " << region.getNumArguments();
+
+    for (Block &block : region) {
+      if (block.empty() || !block.mightHaveTerminator())
+        return mlir::emitError(block.empty() ? getLoc() : block.back().getLoc())
+               << "expected all blocks to have terminators.";
+
+      if (failed(verifyTerminator(block.getTerminator())))
+        return failure();
+    }
+
+    return success();
+  };
+
+  if (failed(verifyRegion(getAllocRegion(), /*expectedNumArgs=*/1, "alloc")))
+    return failure();
+
+  DataSharingClauseType dsType = getDataSharingType();
+
+  if (dsType == DataSharingClauseType::Private && !getCopyRegion().empty())
+    return emitError("`private` clauses require only an `alloc` region.");
+
+  if (dsType == DataSharingClauseType::FirstPrivate && getCopyRegion().empty())
+    return emitError(
+        "`firstprivate` clauses require both `alloc` and `copy` regions.");
+
+  if (dsType == DataSharingClauseType::FirstPrivate &&
+      failed(verifyRegion(getCopyRegion(), /*expectedNumArgs=*/2, "copy")))
+    return failure();
+
+  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 812b79e35595f0..9dc095ba80212a 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1738,3 +1738,73 @@ func.func @omp_distribute(%data_var : memref<i32>) -> () {
       "omp.terminator"() : () -> ()
     }) : (memref<i32>) -> ()
 }
+
+// -----
+
+omp.private {type = private} @x.privatizer : i32 (alloc {
+^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 {type = private} @x.privatizer : i32 (alloc {
+^bb0(%arg0: i32):
+  // expected-error @below {{Invalid yielded value. Expected type: 'i32', got: None}}
+  omp.yield
+})
+
+// -----
+
+// expected-error @below {{expected all blocks to have terminators.}}
+omp.private {type = private} @x.privatizer : i32 (alloc {
+^bb0(%arg0: i32):
+})
+
+// -----
+
+omp.private {type = private} @x.privatizer : i32 (alloc {
+^bb0(%arg0: i32):
+  // expected-error @below {{expected exit block terminator to be an `omp.yield` op.}}
+  omp.terminator
+})
+
+// -----
+
+// expected-error @below {{`alloc`: expected 1 region arguments, got: 2}}
+omp.private {type = private} @x.privatizer : f32 (alloc {
+^bb0(%arg0: f32, %arg1: f32):
+  omp.yield(%arg0 : f32)
+})
+
+// -----
+
+// expected-error @below {{`copy`: expected 2 region arguments, got: 1}}
+omp.private {type = firstprivate} @x.privatizer : f32 (alloc {
+^bb0(%arg0: f32):
+  omp.yield(%arg0 : f32)
+} copy {
+^bb0(%arg0: f32):
+  omp.yield(%arg0 : f32)
+})
+
+// -----
+
+// expected-error @below {{`private` clauses require only an `alloc` region.}}
+omp.private {type = private} @x.privatizer : f32 (alloc {
+^bb0(%arg0: f32):
+  omp.yield(%arg0 : f32)
+} copy {
+^bb0(%arg0: f32, %arg1 : f32):
+  omp.yield(%arg0 : f32)
+})
+
+// -----
+
+// expected-error @below {{`firstprivate` clauses require both `alloc` and `copy` regions.}}
+omp.private {type = firstprivate} @x.privatizer : f32 (alloc {
+^bb0(%arg0: f32):
+  omp.yield(%arg0 : f32)
+})
diff --git a/mlir/test/Dialect/OpenMP/roundtrip.mlir b/mlir/test/Dialect/OpenMP/roundtrip.mlir
new file mode 100644
index 00000000000000..b707c210eddec9
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/roundtrip.mlir
@@ -0,0 +1,21 @@
+// RUN: fir-opt -verify-diagnostics %s | fir-opt | FileCheck %s
+
+// CHECK: omp.private {type = private} @x.privatizer : !fir.ref<i32>(alloc {
+omp.private {type = private} @x.privatizer : !fir.ref<i32> (alloc {
+// CHECK: ^bb0(%arg0: {{.*}}):
+^bb0(%arg0: !fir.ref<i32>):
+  omp.yield(%arg0 : !fir.ref<i32>)
+})
+
+// CHECK: omp.private {type = firstprivate} @y.privatizer : !fir.ref<i32>(alloc {
+omp.private {type = firstprivate} @y.privatizer : !fir.ref<i32> (alloc {
+// CHECK: ^bb0(%arg0: {{.*}}):
+^bb0(%arg0: !fir.ref<i32>):
+  omp.yield(%arg0 : !fir.ref<i32>)
+// CHECK: } copy {
+} copy {
+// CHECK: ^bb0(%arg0: {{.*}}, %arg1: {{.*}}):
+^bb0(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i32>):
+  omp.yield(%arg0 : !fir.ref<i32>)
+})
+

>From 23e7c7d21e117688b5e3005da155250d0e2bf52e Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Sun, 11 Feb 2024 05:03:21 -0600
Subject: [PATCH 2/3] Handle review comments.

---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 42 ++++++++++---------
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  |  2 +-
 mlir/test/Dialect/OpenMP/invalid.mlir         | 32 +++++++-------
 mlir/test/Dialect/OpenMP/roundtrip.mlir       | 26 ++++++------
 4 files changed, 52 insertions(+), 50 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 7c5cfa98dc7118..d2696f6cca0af3 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -153,31 +153,28 @@ def DataSharingClauseTypeAttr : EnumAttr<
   let assemblyFormat = "`{` `type` `=` $value `}`";
 }
 
-def PrivateClauseOp : OpenMP_Op<"private", [
-    IsolatedFromAbove
-  ]> {
-  let summary = "Outline [first]private logic in a separate op.";
+def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove]> {
+  let summary = "Provides declaration of [first]private logic.";
   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 lowering to LLVM IR.
+    This operation provides a declaration of how to implement the
+    [first]privatization of a variable. The dialect users should provide
+    information about how to create an instance of the type in the alloc region
+    and how to initialize the copy from the original item in the copy region.
 
     Examples:
     ---------
     * `private(x)` would be emitted as:
     ```mlir
-    omp.private {type = private} @x.privatizer : !fir.ref<i32> (alloc {
+    omp.private {type = private} @x.privatizer : !fir.ref<i32> alloc {
     ^bb0(%arg0: !fir.ref<i32>):
     %0 = ... allocate proper memory for the private clone ...
     omp.yield(%0 : !fir.ref<i32>)
-    })
+    }
     ```
 
     * `firstprivate(x)` would be emitted as:
     ```mlir
-    omp.private {type = firstprivate} @y.privatizer : !fir.ref<i32> (alloc {
+    omp.private {type = firstprivate} @y.privatizer : !fir.ref<i32> alloc {
     ^bb0(%arg0: !fir.ref<i32>):
     %0 = ... allocate proper memory for the private clone ...
     omp.yield(%0 : !fir.ref<i32>)
@@ -187,35 +184,40 @@ def PrivateClauseOp : OpenMP_Op<"private", [
     // %arg1 represents the memory allocated in `alloc`.
     ... copy from host to the privatized clone ....
     omp.yield(%arg1 : !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:
+    There are no restrictions on the body except for:
     - The `alloc` region has a single argument.
     - The `copy` region has 2 arguments.
-    - Both regions are existed by `omp.yield` ops.
+    - Both regions are terminated by `omp.yield` ops.
     The above restrictions and other obvious restrictions (e.g. verifying the
     type of yielded values) are verified by the custom op verifier. The actual
     contents of the blocks inside both regions are not verified.
 
     Instances of this op would then be used by ops that model directives that
     accept data-sharing attribute clauses.
+
+    The $sym_name attribute provides a symbol by which the privatizer op can be
+    referenced by other dialect ops.
+
+    The $type attribute is the type of the value being privatized.
+
+    The $data_sharing_type attribute specifies whether privatizer corresponds
+    to a `private` or a `firstprivate` clause.
   }];
 
   let arguments = (ins SymbolNameAttr:$sym_name,
-                       TypeAttrOf<AnyType>:$sym_type,
+                       TypeAttrOf<AnyType>:$type,
                        DataSharingClauseTypeAttr:$data_sharing_type);
 
   let regions = (region MinSizedRegion<1>:$alloc_region,
                         AnyRegion:$copy_region);
 
   let assemblyFormat = [{
-    $data_sharing_type $sym_name `:` $sym_type `(`
+    $data_sharing_type $sym_name `:` $type
       `alloc` $alloc_region
       (`copy` $copy_region^)?
-      `)`
       attr-dict
   }];
 
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 96266f8304718c..b967a544581029 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1595,7 +1595,7 @@ LogicalResult DataBoundsOp::verify() {
 }
 
 LogicalResult PrivateClauseOp::verify() {
-  Type symType = getSymType();
+  Type symType = getType();
 
   auto verifyTerminator = [&](Operation *terminator) -> LogicalResult {
     if (!terminator->hasSuccessors() && !llvm::isa<YieldOp>(terminator))
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 9dc095ba80212a..9b4f9ef82f7a09 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1741,70 +1741,70 @@ func.func @omp_distribute(%data_var : memref<i32>) -> () {
 
 // -----
 
-omp.private {type = private} @x.privatizer : i32 (alloc {
+omp.private {type = private} @x.privatizer : i32 alloc {
 ^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 {type = private} @x.privatizer : i32 (alloc {
+omp.private {type = private} @x.privatizer : i32 alloc {
 ^bb0(%arg0: i32):
   // expected-error @below {{Invalid yielded value. Expected type: 'i32', got: None}}
   omp.yield
-})
+}
 
 // -----
 
 // expected-error @below {{expected all blocks to have terminators.}}
-omp.private {type = private} @x.privatizer : i32 (alloc {
+omp.private {type = private} @x.privatizer : i32 alloc {
 ^bb0(%arg0: i32):
-})
+}
 
 // -----
 
-omp.private {type = private} @x.privatizer : i32 (alloc {
+omp.private {type = private} @x.privatizer : i32 alloc {
 ^bb0(%arg0: i32):
   // expected-error @below {{expected exit block terminator to be an `omp.yield` op.}}
   omp.terminator
-})
+}
 
 // -----
 
 // expected-error @below {{`alloc`: expected 1 region arguments, got: 2}}
-omp.private {type = private} @x.privatizer : f32 (alloc {
+omp.private {type = private} @x.privatizer : f32 alloc {
 ^bb0(%arg0: f32, %arg1: f32):
   omp.yield(%arg0 : f32)
-})
+}
 
 // -----
 
 // expected-error @below {{`copy`: expected 2 region arguments, got: 1}}
-omp.private {type = firstprivate} @x.privatizer : f32 (alloc {
+omp.private {type = firstprivate} @x.privatizer : f32 alloc {
 ^bb0(%arg0: f32):
   omp.yield(%arg0 : f32)
 } copy {
 ^bb0(%arg0: f32):
   omp.yield(%arg0 : f32)
-})
+}
 
 // -----
 
 // expected-error @below {{`private` clauses require only an `alloc` region.}}
-omp.private {type = private} @x.privatizer : f32 (alloc {
+omp.private {type = private} @x.privatizer : f32 alloc {
 ^bb0(%arg0: f32):
   omp.yield(%arg0 : f32)
 } copy {
 ^bb0(%arg0: f32, %arg1 : f32):
   omp.yield(%arg0 : f32)
-})
+}
 
 // -----
 
 // expected-error @below {{`firstprivate` clauses require both `alloc` and `copy` regions.}}
-omp.private {type = firstprivate} @x.privatizer : f32 (alloc {
+omp.private {type = firstprivate} @x.privatizer : f32 alloc {
 ^bb0(%arg0: f32):
   omp.yield(%arg0 : f32)
-})
+}
diff --git a/mlir/test/Dialect/OpenMP/roundtrip.mlir b/mlir/test/Dialect/OpenMP/roundtrip.mlir
index b707c210eddec9..2553442638ee84 100644
--- a/mlir/test/Dialect/OpenMP/roundtrip.mlir
+++ b/mlir/test/Dialect/OpenMP/roundtrip.mlir
@@ -1,21 +1,21 @@
-// RUN: fir-opt -verify-diagnostics %s | fir-opt | FileCheck %s
+// RUN: mlir-opt -verify-diagnostics %s | mlir-opt | FileCheck %s
 
-// CHECK: omp.private {type = private} @x.privatizer : !fir.ref<i32>(alloc {
-omp.private {type = private} @x.privatizer : !fir.ref<i32> (alloc {
+// CHECK: omp.private {type = private} @x.privatizer : !llvm.ptr alloc {
+omp.private {type = private} @x.privatizer : !llvm.ptr alloc {
 // CHECK: ^bb0(%arg0: {{.*}}):
-^bb0(%arg0: !fir.ref<i32>):
-  omp.yield(%arg0 : !fir.ref<i32>)
-})
+^bb0(%arg0: !llvm.ptr):
+  omp.yield(%arg0 : !llvm.ptr)
+}
 
-// CHECK: omp.private {type = firstprivate} @y.privatizer : !fir.ref<i32>(alloc {
-omp.private {type = firstprivate} @y.privatizer : !fir.ref<i32> (alloc {
+// CHECK: omp.private {type = firstprivate} @y.privatizer : !llvm.ptr alloc {
+omp.private {type = firstprivate} @y.privatizer : !llvm.ptr alloc {
 // CHECK: ^bb0(%arg0: {{.*}}):
-^bb0(%arg0: !fir.ref<i32>):
-  omp.yield(%arg0 : !fir.ref<i32>)
+^bb0(%arg0: !llvm.ptr):
+  omp.yield(%arg0 : !llvm.ptr)
 // CHECK: } copy {
 } copy {
 // CHECK: ^bb0(%arg0: {{.*}}, %arg1: {{.*}}):
-^bb0(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i32>):
-  omp.yield(%arg0 : !fir.ref<i32>)
-})
+^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+  omp.yield(%arg0 : !llvm.ptr)
+}
 

>From 460f94bc90683e46204fbd8c9dd5d4d1b542b651 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Sun, 11 Feb 2024 07:57:53 -0600
Subject: [PATCH 3/3] [MLIR][OpenMP] Support `llvm` conversion for
 `omp.private` regions

Introduces conversion of `omp.private`'s regions to the LLVM dialect.
This reuses the already existing conversion pattern for
`ReducetionDeclareOp` and repurposes it to be used for multi-region ops
as well.
---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td |  6 ++++
 .../Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp  | 32 +++++++++++--------
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  |  9 ++++++
 .../OpenMPToLLVM/convert-to-llvmir.mlir       | 26 +++++++++++++++
 4 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index d2696f6cca0af3..5ddcfdbfb2ec58 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -221,6 +221,12 @@ def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove]> {
       attr-dict
   }];
 
+  let builders = [
+    OpBuilder<(ins CArg<"TypeRange">:$result,
+                   CArg<"StringAttr">:$sym_name,
+                   CArg<"TypeAttr">:$type)>
+  ];
+
   let hasVerifier = 1;
 }
 
diff --git a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
index 730858ffc67a71..2eba4fba105c7b 100644
--- a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
+++ b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
@@ -200,16 +200,20 @@ struct ReductionOpConversion : public ConvertOpToLLVMPattern<omp::ReductionOp> {
   }
 };
 
-struct ReductionDeclareOpConversion
-    : public ConvertOpToLLVMPattern<omp::ReductionDeclareOp> {
-  using ConvertOpToLLVMPattern<omp::ReductionDeclareOp>::ConvertOpToLLVMPattern;
+template <typename OpType>
+struct MultiRegionOpConversion : public ConvertOpToLLVMPattern<OpType> {
+  using ConvertOpToLLVMPattern<OpType>::ConvertOpToLLVMPattern;
   LogicalResult
-  matchAndRewrite(omp::ReductionDeclareOp curOp, OpAdaptor adaptor,
+  matchAndRewrite(OpType curOp, typename OpType::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto newOp = rewriter.create<omp::ReductionDeclareOp>(
+    auto newOp = rewriter.create<OpType>(
         curOp.getLoc(), TypeRange(), curOp.getSymNameAttr(),
         TypeAttr::get(this->getTypeConverter()->convertType(
             curOp.getTypeAttr().getValue())));
+
+    if constexpr (std::is_same_v<OpType, mlir::omp::PrivateClauseOp>)
+      newOp.setDataSharingType(curOp.getDataSharingType());
+
     for (unsigned idx = 0; idx < curOp.getNumRegions(); idx++) {
       rewriter.inlineRegionBefore(curOp.getRegion(idx), newOp.getRegion(idx),
                                   newOp.getRegion(idx).end());
@@ -231,11 +235,12 @@ void mlir::configureOpenMPToLLVMConversionLegality(
       mlir::omp::DataOp, mlir::omp::OrderedRegionOp, mlir::omp::ParallelOp,
       mlir::omp::WsLoopOp, mlir::omp::SimdLoopOp, mlir::omp::MasterOp,
       mlir::omp::SectionOp, mlir::omp::SectionsOp, mlir::omp::SingleOp,
-      mlir::omp::TaskGroupOp, mlir::omp::TaskOp>([&](Operation *op) {
-    return typeConverter.isLegal(&op->getRegion(0)) &&
-           typeConverter.isLegal(op->getOperandTypes()) &&
-           typeConverter.isLegal(op->getResultTypes());
-  });
+      mlir::omp::TaskGroupOp, mlir::omp::TaskOp, mlir::omp::PrivateClauseOp>(
+      [&](Operation *op) {
+        return typeConverter.isLegal(&op->getRegion(0)) &&
+               typeConverter.isLegal(op->getOperandTypes()) &&
+               typeConverter.isLegal(op->getResultTypes());
+      });
   target.addDynamicallyLegalOp<
       mlir::omp::AtomicReadOp, mlir::omp::AtomicWriteOp, mlir::omp::FlushOp,
       mlir::omp::ThreadprivateOp, mlir::omp::YieldOp, mlir::omp::EnterDataOp,
@@ -267,9 +272,10 @@ void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
 
   patterns.add<
       AtomicReadOpConversion, MapInfoOpConversion, ReductionOpConversion,
-      ReductionDeclareOpConversion, RegionOpConversion<omp::CriticalOp>,
-      RegionOpConversion<omp::MasterOp>, ReductionOpConversion,
-      RegionOpConversion<omp::OrderedRegionOp>,
+      MultiRegionOpConversion<omp::ReductionDeclareOp>,
+      MultiRegionOpConversion<omp::PrivateClauseOp>,
+      RegionOpConversion<omp::CriticalOp>, RegionOpConversion<omp::MasterOp>,
+      ReductionOpConversion, RegionOpConversion<omp::OrderedRegionOp>,
       RegionOpConversion<omp::ParallelOp>, RegionOpConversion<omp::WsLoopOp>,
       RegionOpConversion<omp::SectionsOp>, RegionOpConversion<omp::SectionOp>,
       RegionOpConversion<omp::SimdLoopOp>, RegionOpConversion<omp::SingleOp>,
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index b967a544581029..4ac154da2cb3ac 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1594,6 +1594,15 @@ LogicalResult DataBoundsOp::verify() {
   return success();
 }
 
+void PrivateClauseOp::build(OpBuilder &odsBuilder, OperationState &odsState,
+                            TypeRange /*result_types*/, StringAttr symName,
+                            TypeAttr type) {
+  PrivateClauseOp::build(
+      odsBuilder, odsState, symName, type,
+      DataSharingClauseTypeAttr::get(odsBuilder.getContext(),
+                                     DataSharingClauseType::Private));
+}
+
 LogicalResult PrivateClauseOp::verify() {
   Type symType = getType();
 
diff --git a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
index 3fbeaebb592a4d..eacdde9255032f 100644
--- a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
+++ b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
@@ -479,3 +479,29 @@ llvm.func @_QPtarget_map_with_bounds(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2:
   }
   llvm.return
 }
+
+// -----
+
+// CHECK: omp.private {type = private} @x.privatizer : !llvm.struct<{{.*}}> alloc {
+omp.private {type = private} @x.privatizer : memref<?xf32> alloc {
+// CHECK: ^bb0(%arg0: !llvm.struct<{{.*}}>):
+^bb0(%arg0: memref<?xf32>):
+  // CHECK: omp.yield(%arg0 : !llvm.struct<{{.*}}>)
+  omp.yield(%arg0 : memref<?xf32>)
+}
+
+// -----
+
+// CHECK: omp.private {type = firstprivate} @y.privatizer : i64 alloc {
+omp.private {type = firstprivate} @y.privatizer : index alloc {
+// CHECK: ^bb0(%arg0: i64):
+^bb0(%arg0: index):
+  // CHECK: omp.yield(%arg0 : i64)
+  omp.yield(%arg0 : index)
+// CHECK: } copy {
+} copy {
+// CHECK: ^bb0(%arg0: i64, %arg1: i64):
+^bb0(%arg0: index, %arg1: index):
+  // CHECK: omp.yield(%arg0 : i64)
+  omp.yield(%arg0 : index)
+}



More information about the Mlir-commits mailing list