[Mlir-commits] [mlir] [MLIR][OpenMP] Add `omp.private` op (PR #80955)
Kareem Ergawy
llvmlistbot at llvm.org
Sun Feb 11 03:03:40 PST 2024
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/80955
>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/2] [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 ®ion, 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 5426d5f6ad6a40eeff21a201791a451164672f5a 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/2] Handle review comments.
---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 40 ++++++++++---------
mlir/test/Dialect/OpenMP/invalid.mlir | 32 +++++++--------
mlir/test/Dialect/OpenMP/roundtrip.mlir | 26 ++++++------
3 files changed, 50 insertions(+), 48 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 7c5cfa98dc7118..f58d60485558c7 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,21 +184,27 @@ 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 $sym_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,
@@ -212,10 +215,9 @@ def PrivateClauseOp : OpenMP_Op<"private", [
AnyRegion:$copy_region);
let assemblyFormat = [{
- $data_sharing_type $sym_name `:` $sym_type `(`
+ $data_sharing_type $sym_name `:` $sym_type
`alloc` $alloc_region
(`copy` $copy_region^)?
- `)`
attr-dict
}];
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)
+}
More information about the Mlir-commits
mailing list