[Mlir-commits] [mlir] [MLIR][OpenMP] Add `omp.private` op (PR #80955)
Kareem Ergawy
llvmlistbot at llvm.org
Wed Feb 7 02:23:11 PST 2024
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/80955
>From 728806bec462c7cda61b5d8e6408dadb241ab8cf 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 | 81 ++++++++++++++++++-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 55 +++++++++++++
mlir/test/Dialect/OpenMP/invalid.mlir | 46 +++++++++++
mlir/test/Dialect/OpenMP/roundtrip.mlir | 13 +++
4 files changed, 194 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 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>)
+}
+
More information about the Mlir-commits
mailing list