[Mlir-commits] [mlir] [MLIR][OpenMP] Add `omp.private` op (PR #80955)
Kareem Ergawy
llvmlistbot at llvm.org
Thu Feb 8 03:10:00 PST 2024
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/80955
>From 5e2503785699acb3d9bcb9170dddcd106719564e 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] [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>)
+})
+
More information about the Mlir-commits
mailing list