[llvm-branch-commits] [flang] [flang][fir] Add `fir.local` op for locality specifiers (PR #138505)
Kareem Ergawy via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue May 6 23:49:19 PDT 2025
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/138505
>From 09f3a12a25be35d119c02164ef6a4d21a024a940 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Mon, 5 May 2025 04:20:40 -0500
Subject: [PATCH 1/3] [flang][fir] Add `fir.local` op for locality specifiers
Adds a new `fir.local` op to model `local` and `local_init` locality
specifiers. This op is a clone of `omp.private`. In particular, this new
op also models the privatization/localization logic of an SSA value in
the `fir` dialect just like `omp.private` does for OpenMP.
---
.../flang/Optimizer/Dialect/FIRAttr.td | 19 +++
.../include/flang/Optimizer/Dialect/FIROps.td | 131 ++++++++++++++++++
flang/lib/Optimizer/Dialect/FIROps.cpp | 99 +++++++++++++
flang/test/Fir/do_concurrent.fir | 19 +++
4 files changed, 268 insertions(+)
diff --git a/flang/include/flang/Optimizer/Dialect/FIRAttr.td b/flang/include/flang/Optimizer/Dialect/FIRAttr.td
index 3ebc24951cfff..2845080030b92 100644
--- a/flang/include/flang/Optimizer/Dialect/FIRAttr.td
+++ b/flang/include/flang/Optimizer/Dialect/FIRAttr.td
@@ -200,4 +200,23 @@ def fir_OpenMPSafeTempArrayCopyAttr : fir_Attr<"OpenMPSafeTempArrayCopy"> {
}];
}
+def LocalitySpecTypeLocal : I32EnumAttrCase<"Local", 0, "local">;
+def LocalitySpecTypeLocalInit
+ : I32EnumAttrCase<"LocalInit", 1, "local_init">;
+
+def LocalitySpecifierType : I32EnumAttr<
+ "LocalitySpecifierType",
+ "Type of a locality specifier", [
+ LocalitySpecTypeLocal,
+ LocalitySpecTypeLocalInit
+ ]> {
+ let genSpecializedAttr = 0;
+ let cppNamespace = "::fir";
+}
+
+def LocalitySpecifierTypeAttr : EnumAttr<FIROpsDialect, LocalitySpecifierType,
+ "locality_specifier_type"> {
+ let assemblyFormat = "`{` `type` `=` $value `}`";
+}
+
#endif // FIR_DIALECT_FIR_ATTRS
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 0ba985641069b..f87ffce72192d 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -3485,6 +3485,137 @@ def fir_BoxTotalElementsOp
let hasCanonicalizer = 1;
}
+def YieldOp : fir_Op<"yield",
+ [Pure, ReturnLike, Terminator,
+ ParentOneOf<["LocalitySpecifierOp"]>]> {
+ let summary = "loop yield and termination operation";
+ let description = [{
+ "fir.yield" yields SSA values from the fir dialect op region and
+ terminates the region. The semantics of how the values are yielded is
+ defined by the parent operation.
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$results);
+
+ let builders = [
+ OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
+ ];
+
+ let assemblyFormat = "( `(` $results^ `:` type($results) `)` )? attr-dict";
+}
+
+def fir_LocalitySpecifierOp : fir_Op<"local", [IsolatedFromAbove]> {
+ let summary = "Provides declaration of [first]private logic.";
+ let description = [{
+ This operation provides a declaration of how to implement the
+ localization of a variable. The dialect users should provide
+ which type should be allocated for this variable. The allocated (usually by
+ alloca) variable is passed to the initialization region which does everything
+ else (e.g. initialization of Fortran runtime descriptors). Information about
+ how to initialize the copy from the original item should be given in the
+ copy region, and if needed, how to deallocate memory (allocated by the
+ initialization region) in the dealloc region.
+
+ Examples:
+
+ * `local(x)` would not need any regions because no initialization is
+ required by the standard for i32 variables and this is not firstprivate.
+ ```
+ fir.local {type = local} @x.localizer : i32
+ ```
+
+ * `local_init(x)` would be emitted as:
+ ```
+ fir.local {type = local_init} @x.localizer : i32 copy {
+ ^bb0(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i32>):
+ // %arg0 is the original host variable.
+ // %arg1 represents the memory allocated for this private variable.
+ ... copy from host to the localized clone ....
+ fir.yield(%arg1 : !fir.ref<i32>)
+ }
+ ```
+
+ * `local(x)` for "allocatables" would be emitted as:
+ ```
+ fir.local {type = local} @x.privatizer : !some.type init {
+ ^bb0(%arg0: !some.pointer<!some.type>, %arg1: !some.pointer<!some.type>):
+ // initialize %arg1, using %arg0 as a mold for allocations.
+ // For example if %arg0 is a heap allocated array with a runtime determined
+ // length and !some.type is a runtime type descriptor, the init region
+ // will read the array length from %arg0, and heap allocate an array of the
+ // right length and initialize %arg1 to contain the array allocation and
+ // length.
+ fir.yield(%arg1 : !some.pointer<!some.type>)
+ } dealloc {
+ ^bb0(%arg0: !some.pointer<!some.type>):
+ // ... deallocate memory allocated by the init region...
+ // In the example above, this will free the heap allocated array data.
+ fir.yield
+ }
+ ```
+
+ There are no restrictions on the body except for:
+ - The `dealloc` regions has a single argument.
+ - The `init` & `copy` regions have 2 arguments.
+ - All three regions are terminated by `fir.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 all 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 localized. This type
+ will be implicitly allocated in MLIR->LLVMIR conversion and passed as the
+ second argument to the init region. Therefore the type of arguments to
+ the regions should be a type which represents a pointer to `type`.
+
+ The `locality_specifier_type` attribute specifies whether the localized
+ corresponds to a `local` or a `local_init` specifier.
+ }];
+
+ let arguments = (ins SymbolNameAttr:$sym_name,
+ TypeAttrOf<AnyType>:$type,
+ LocalitySpecifierTypeAttr:$locality_specifier_type);
+
+ let regions = (region AnyRegion:$init_region,
+ AnyRegion:$copy_region,
+ AnyRegion:$dealloc_region);
+
+ let assemblyFormat = [{
+ $locality_specifier_type $sym_name `:` $type
+ (`init` $init_region^)?
+ (`copy` $copy_region^)?
+ (`dealloc` $dealloc_region^)?
+ attr-dict
+ }];
+
+ let builders = [
+ OpBuilder<(ins CArg<"mlir::TypeRange">:$result,
+ CArg<"mlir::StringAttr">:$sym_name,
+ CArg<"mlir::TypeAttr">:$type)>
+ ];
+
+ let extraClassDeclaration = [{
+ /// Get the type for arguments to nested regions. This should
+ /// generally be either the same as getType() or some pointer
+ /// type (pointing to the type allocated by this op).
+ /// This method will return Type{nullptr} if there are no nested
+ /// regions.
+ mlir::Type getArgType() {
+ for (mlir::Region *region : getRegions())
+ for (mlir::Type ty : region->getArgumentTypes())
+ return ty;
+ return nullptr;
+ }
+ }];
+
+ let hasRegionVerifier = 1;
+}
+
def fir_DoConcurrentOp : fir_Op<"do_concurrent",
[SingleBlock, AutomaticAllocationScope]> {
let summary = "do concurrent loop wrapper";
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 05ef69169bae5..65ec730e134c2 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -4909,6 +4909,105 @@ void fir::BoxTotalElementsOp::getCanonicalizationPatterns(
patterns.add<SimplifyBoxTotalElementsOp>(context);
}
+//===----------------------------------------------------------------------===//
+// LocalitySpecifierOp
+//===----------------------------------------------------------------------===//
+
+llvm::LogicalResult fir::LocalitySpecifierOp::verifyRegions() {
+ mlir::Type argType = getArgType();
+ auto verifyTerminator = [&](mlir::Operation *terminator,
+ bool yieldsValue) -> llvm::LogicalResult {
+ if (!terminator->getBlock()->getSuccessors().empty())
+ return llvm::success();
+
+ if (!llvm::isa<fir::YieldOp>(terminator))
+ return mlir::emitError(terminator->getLoc())
+ << "expected exit block terminator to be an `fir.yield` op.";
+
+ YieldOp yieldOp = llvm::cast<YieldOp>(terminator);
+ mlir::TypeRange yieldedTypes = yieldOp.getResults().getTypes();
+
+ if (!yieldsValue) {
+ if (yieldedTypes.empty())
+ return llvm::success();
+
+ return mlir::emitError(terminator->getLoc())
+ << "Did not expect any values to be yielded.";
+ }
+
+ if (yieldedTypes.size() == 1 && yieldedTypes.front() == argType)
+ return llvm::success();
+
+ auto error = mlir::emitError(yieldOp.getLoc())
+ << "Invalid yielded value. Expected type: " << argType
+ << ", got: ";
+
+ if (yieldedTypes.empty())
+ error << "None";
+ else
+ error << yieldedTypes;
+
+ return error;
+ };
+
+ auto verifyRegion = [&](mlir::Region ®ion, unsigned expectedNumArgs,
+ llvm::StringRef regionName,
+ bool yieldsValue) -> llvm::LogicalResult {
+ assert(!region.empty());
+
+ if (region.getNumArguments() != expectedNumArgs)
+ return mlir::emitError(region.getLoc())
+ << "`" << regionName << "`: "
+ << "expected " << expectedNumArgs
+ << " region arguments, got: " << region.getNumArguments();
+
+ for (mlir::Block &block : region) {
+ // MLIR will verify the absence of the terminator for us.
+ if (!block.mightHaveTerminator())
+ continue;
+
+ if (failed(verifyTerminator(block.getTerminator(), yieldsValue)))
+ return llvm::failure();
+ }
+
+ return llvm::success();
+ };
+
+ // Ensure all of the region arguments have the same type
+ for (mlir::Region *region : getRegions())
+ for (mlir::Type ty : region->getArgumentTypes())
+ if (ty != argType)
+ return emitError() << "Region argument type mismatch: got " << ty
+ << " expected " << argType << ".";
+
+ mlir::Region &initRegion = getInitRegion();
+ if (!initRegion.empty() &&
+ failed(verifyRegion(getInitRegion(), /*expectedNumArgs=*/2, "init",
+ /*yieldsValue=*/true)))
+ return llvm::failure();
+
+ LocalitySpecifierType dsType = getLocalitySpecifierType();
+
+ if (dsType == LocalitySpecifierType::Local && !getCopyRegion().empty())
+ return emitError("`local` specifiers do not require a `copy` region.");
+
+ if (dsType == LocalitySpecifierType::LocalInit && getCopyRegion().empty())
+ return emitError(
+ "`local_init` specifier require at least a `copy` region.");
+
+ if (dsType == LocalitySpecifierType::LocalInit &&
+ failed(verifyRegion(getCopyRegion(), /*expectedNumArgs=*/2, "copy",
+ /*yieldsValue=*/true)))
+ return llvm::failure();
+
+ if (!getDeallocRegion().empty() &&
+ failed(verifyRegion(getDeallocRegion(), /*expectedNumArgs=*/1, "dealloc",
+ /*yieldsValue=*/false)))
+ return llvm::failure();
+
+ return llvm::success();
+}
+
//===----------------------------------------------------------------------===//
// DoConcurrentOp
//===----------------------------------------------------------------------===//
diff --git a/flang/test/Fir/do_concurrent.fir b/flang/test/Fir/do_concurrent.fir
index 8e80ffb9c7b0b..4e55777402428 100644
--- a/flang/test/Fir/do_concurrent.fir
+++ b/flang/test/Fir/do_concurrent.fir
@@ -90,3 +90,22 @@ func.func @dc_2d_reduction(%i_lb: index, %i_ub: index, %i_st: index,
// CHECK: fir.store %[[J_IV_CVT]] to %[[J]] : !fir.ref<i32>
// CHECK: }
// CHECK: }
+
+
+fir.local {type = local} @local_privatizer : i32
+
+// CHECK: fir.local {type = local} @[[LOCAL_PRIV_SYM:local_privatizer]] : i32
+
+fir.local {type = local_init} @local_init_privatizer : i32 copy {
+^bb0(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i32>):
+ %0 = fir.load %arg0 : !fir.ref<i32>
+ fir.store %0 to %arg1 : !fir.ref<i32>
+ fir.yield(%arg1 : !fir.ref<i32>)
+}
+
+// CHECK: fir.local {type = local_init} @[[LOCAL_INIT_PRIV_SYM:local_init_privatizer]] : i32
+// CHECK: ^bb0(%[[ORIG_VAL:.*]]: !fir.ref<i32>, %[[LOCAL_VAL:.*]]: !fir.ref<i32>):
+// CHECK: %[[ORIG_VAL_LD:.*]] = fir.load %[[ORIG_VAL]]
+// CHECK: fir.store %[[ORIG_VAL_LD]] to %[[LOCAL_VAL]] : !fir.ref<i32>
+// CHECK: fir.yield(%[[LOCAL_VAL]] : !fir.ref<i32>)
+// CHECK: }
>From da39596e52e046d94d8a0bf387b75a6fac04252b Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Tue, 6 May 2025 14:09:22 -0500
Subject: [PATCH 2/3] handle review comments
---
flang/include/flang/Optimizer/Dialect/FIROps.td | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index f87ffce72192d..ee82325ac55b7 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -3505,7 +3505,7 @@ def YieldOp : fir_Op<"yield",
}
def fir_LocalitySpecifierOp : fir_Op<"local", [IsolatedFromAbove]> {
- let summary = "Provides declaration of [first]private logic.";
+ let summary = "Provides declaration of local and local_init logic.";
let description = [{
This operation provides a declaration of how to implement the
localization of a variable. The dialect users should provide
@@ -3519,7 +3519,7 @@ def fir_LocalitySpecifierOp : fir_Op<"local", [IsolatedFromAbove]> {
Examples:
* `local(x)` would not need any regions because no initialization is
- required by the standard for i32 variables and this is not firstprivate.
+ required by the standard for i32 variables and this is not local_init.
```
fir.local {type = local} @x.localizer : i32
```
@@ -3538,16 +3538,16 @@ def fir_LocalitySpecifierOp : fir_Op<"local", [IsolatedFromAbove]> {
* `local(x)` for "allocatables" would be emitted as:
```
fir.local {type = local} @x.privatizer : !some.type init {
- ^bb0(%arg0: !some.pointer<!some.type>, %arg1: !some.pointer<!some.type>):
+ ^bb0(%arg0: !fir.ref<!some.type>, %arg1: !fir.ref<!some.type>):
// initialize %arg1, using %arg0 as a mold for allocations.
// For example if %arg0 is a heap allocated array with a runtime determined
// length and !some.type is a runtime type descriptor, the init region
// will read the array length from %arg0, and heap allocate an array of the
// right length and initialize %arg1 to contain the array allocation and
// length.
- fir.yield(%arg1 : !some.pointer<!some.type>)
+ fir.yield(%arg1 : !fir.ref<!some.type>)
} dealloc {
- ^bb0(%arg0: !some.pointer<!some.type>):
+ ^bb0(%arg0: !fir.ref<!some.type>):
// ... deallocate memory allocated by the init region...
// In the example above, this will free the heap allocated array data.
fir.yield
>From fcb6b5d4893147bc67a9a21dc903fdc19c71a427 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Wed, 7 May 2025 01:49:02 -0500
Subject: [PATCH 3/3] add verifier tests
---
.../include/flang/Optimizer/Dialect/FIROps.td | 4 +-
flang/lib/Optimizer/Dialect/FIROps.cpp | 2 +-
flang/test/Fir/invalid.fir | 78 ++++++++++++++++++-
3 files changed, 79 insertions(+), 5 deletions(-)
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index ee82325ac55b7..acc0c6967c739 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -3490,7 +3490,7 @@ def YieldOp : fir_Op<"yield",
ParentOneOf<["LocalitySpecifierOp"]>]> {
let summary = "loop yield and termination operation";
let description = [{
- "fir.yield" yields SSA values from the fir dialect op region and
+ "fir.yield" yields SSA values from a fir dialect op region and
terminates the region. The semantics of how the values are yielded is
defined by the parent operation.
}];
@@ -3537,7 +3537,7 @@ def fir_LocalitySpecifierOp : fir_Op<"local", [IsolatedFromAbove]> {
* `local(x)` for "allocatables" would be emitted as:
```
- fir.local {type = local} @x.privatizer : !some.type init {
+ fir.local {type = local} @x.localizer : !some.type init {
^bb0(%arg0: !fir.ref<!some.type>, %arg1: !fir.ref<!some.type>):
// initialize %arg1, using %arg0 as a mold for allocations.
// For example if %arg0 is a heap allocated array with a runtime determined
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 65ec730e134c2..955acbe7018d3 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -4993,7 +4993,7 @@ llvm::LogicalResult fir::LocalitySpecifierOp::verifyRegions() {
if (dsType == LocalitySpecifierType::LocalInit && getCopyRegion().empty())
return emitError(
- "`local_init` specifier require at least a `copy` region.");
+ "`local_init` specifiers require at least a `copy` region.");
if (dsType == LocalitySpecifierType::LocalInit &&
failed(verifyRegion(getCopyRegion(), /*expectedNumArgs=*/2, "copy",
diff --git a/flang/test/Fir/invalid.fir b/flang/test/Fir/invalid.fir
index f9f5e267dd9bc..733227339bc39 100644
--- a/flang/test/Fir/invalid.fir
+++ b/flang/test/Fir/invalid.fir
@@ -1,5 +1,3 @@
-
-
// RUN: fir-opt -split-input-file -verify-diagnostics --strict-fir-volatile-verifier %s
// expected-error at +1{{custom op 'fir.string_lit' must have character type}}
@@ -1311,3 +1309,79 @@ func.func @bad_convert_volatile6(%arg0: !fir.ref<i32>) -> !fir.ref<i64> {
%0 = fir.volatile_cast %arg0 : (!fir.ref<i32>) -> !fir.ref<i64>
return %0 : !fir.ref<i64>
}
+
+// -----
+
+fir.local {type = local} @x.localizer : i32 init {
+^bb0(%arg0: i32, %arg1: i32):
+ %0 = arith.constant 0.0 : f32
+ // expected-error @below {{Invalid yielded value. Expected type: 'i32', got: 'f32'}}
+ fir.yield(%0 : f32)
+}
+
+// -----
+
+// expected-error @below {{Region argument type mismatch: got 'f32' expected 'i32'.}}
+fir.local {type = local} @x.localizer : i32 init {
+^bb0(%arg0: i32, %arg1: f32):
+ fir.yield
+}
+
+// -----
+
+fir.local {type = local} @x.localizer : f32 init {
+^bb0(%arg0: f32, %arg1: f32):
+ fir.yield(%arg0: f32)
+} dealloc {
+^bb0(%arg0: f32):
+ // expected-error @below {{Did not expect any values to be yielded.}}
+ fir.yield(%arg0 : f32)
+}
+
+// -----
+
+fir.local {type = local} @x.localizer : i32 init {
+^bb0(%arg0: i32, %arg1: i32):
+ // expected-error @below {{expected exit block terminator to be an `fir.yield` op.}}
+ fir.unreachable
+}
+
+// -----
+
+// expected-error @below {{`init`: expected 2 region arguments, got: 1}}
+fir.local {type = local} @x.localizer : f32 init {
+^bb0(%arg0: f32):
+ fir.yield(%arg0 : f32)
+}
+
+// -----
+
+// expected-error @below {{`copy`: expected 2 region arguments, got: 1}}
+fir.local {type = local_init} @x.privatizer : f32 copy {
+^bb0(%arg0: f32):
+ fir.yield(%arg0 : f32)
+}
+
+// -----
+
+// expected-error @below {{`dealloc`: expected 1 region arguments, got: 2}}
+fir.local {type = local} @x.localizer : f32 dealloc {
+^bb0(%arg0: f32, %arg1: f32):
+ fir.yield
+}
+
+// -----
+
+// expected-error @below {{`local` specifiers do not require a `copy` region.}}
+fir.local {type = local} @x.localizer : f32 copy {
+^bb0(%arg0: f32, %arg1 : f32):
+ fir.yield(%arg0 : f32)
+}
+
+// -----
+
+// expected-error @below {{`local_init` specifiers require at least a `copy` region.}}
+fir.local {type = local_init} @x.localizer : f32 init {
+^bb0(%arg0: f32, %arg1: f32):
+ fir.yield(%arg0 : f32)
+}
More information about the llvm-branch-commits
mailing list