shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[MLIR][OpenMP] Enable BlockArgOpenMPOpInterface accessing operands (#130769)

Browse Source 
This patch makes additions to the `BlockArgOpenMPOpInterface` to
simplify its use by letting it handle the matching between operands and
their associated entry block arguments. Most significantly, the
following is now possible:

```c++
SmallVector<std::pair<Value, BlockArgument>> pairs;
cast<BlockArgOpenMPOpInterface>(op).getBlockArgsPairs(pairs);
for (auto [var, arg] : pairs) {
  // var points to the operand (outside value) and arg points to the entry
  // block argument associated to that value.
}
```

This is achieved by making the interface define and use `getXyzVars()`
methods, which by default return empty `OperandRange`s and are overriden
by getters automatically produced for the `Variadic<...> $xyz_vars`
tablegen argument of the corresponding clause. These definitions can
then be simplified, since they no longer need to manually define
`numXyzBlockArgs` functions as a result.

A side-effect of this is that all ops implementing this interface will
now publicly define `getXyzVars()` functions for all entry block
argument-generating clauses, even if they don't actually accept all
clauses. However, these would just return empty ranges, so it shouldn't
cause issues.

This change uncovered some incorrect definitions of class declarations
related to the `ReductionClauseInterface`, and the `OpenMP_DetachClause`
incorrectly implementing the `BlockArgOpenMPOpInterface`, so these
issues are also addressed.
This commit is contained in:
Sergio Afonso 2025-03-12 11:50:12 +00:00 committed by GitHub
parent c851ee38ad
commit 032f83b743
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 122 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
mlir/docs/Dialects/OpenMPDialect/_index.md
View File

@ -352,12 +352,29 @@ let assemblyFormat = clausesAssemblyFormat # [{
```
The `BlockArgOpenMPOpInterface` has been introduced to simplify the addition and
handling of these kinds of clauses. It holds `num<ClauseName>BlockArgs()`
functions that by default return 0, to be overriden by each clause through the
`extraClassDeclaration` property. Based on these functions and the expected
alphabetical sorting between entry block argument-defining clauses, it
implements `get<ClauseName>BlockArgs()` functions that are the intended method
of accessing clause-defined block arguments.
handling of these kinds of clauses. Adding it to an operation directly, or
indirectly through a clause, results in the addition of overridable
`get<ClauseName>Vars()` and `num<ClauseName>BlockArgs()` public functions for
all entry block argument-generating clauses. By default, the reported number of
block arguments defined by a clause will correspond to the number of operands
taken by the operation for that clause. This list of operands will be empty by
default, and will automatically be overriden by getters of the corresponding
`Variadic<...> $<clause_name>_vars` argument of the same clause's definition.
In addition to these methods added to the actual operations, the
`BlockArgOpenMPOpInterface` itself defines a set of methods based on the
previous ones and on the convention that entry block arguments for multiple
clauses are sorted alphabetically by clause name. These are listed below, and
they represent the main way in which clause-defined block arguments should be
accessed:
- `get<ClauseName>BlockArgsStart()`: Returns the index within the list of
entry block arguments where the first element defined by the given clause
should be located.
- `get<ClauseName>BlockArgs()`: Returns the list of entry block arguments
defined by the given clause.
- `getBlockArgsPairs()`: Returns a list of pairs where the first element is
the outside value, or operand, and the second element is the corresponding
entry block argument.
## Loop-Associated Directives

44
mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
View File

@ -470,12 +470,6 @@ class OpenMP_HasDeviceAddrClauseSkip<
Variadic<OpenMP_PointerLikeType>:$has_device_addr_vars
);
let extraClassDeclaration = [{
unsigned numHasDeviceAddrBlockArgs() {
return getHasDeviceAddrVars().size();
}
}];
let description = [{
The optional `has_device_addr_vars` indicates that list items already have
device addresses, so they may be directly accessed from the target device.
@ -565,12 +559,6 @@ class OpenMP_HostEvalClauseSkip<
Variadic<AnyType>:$host_eval_vars
);
let extraClassDeclaration = [{
unsigned numHostEvalBlockArgs() {
return getHostEvalVars().size();
}
}];
let description = [{
The optional `host_eval_vars` holds values defined outside of the region of
the `IsolatedFromAbove` operation for which a corresponding entry block
@ -629,12 +617,10 @@ class OpenMP_InReductionClauseSkip<
let extraClassDeclaration = [{
/// Returns the reduction variables.
SmallVector<Value> getReductionVars() {
SmallVector<Value> getAllReductionVars() {
return SmallVector<Value>(getInReductionVars().begin(),
getInReductionVars().end());
}
unsigned numInReductionBlockArgs() { return getInReductionVars().size(); }
}];
// Description varies depending on the operation. Assembly format not defined
@ -749,6 +735,9 @@ class OpenMP_MapClauseSkip<
Variadic<OpenMP_PointerLikeType>:$map_vars
);
// This assembly format should only be used by operations where `map` does not
// define entry block arguments. Otherwise, it must be printed and parsed
// together with the corresponding region.
let optAssemblyFormat = [{
`map_entries` `(` $map_vars `:` type($map_vars) `)`
}];
@ -1060,8 +1049,6 @@ class OpenMP_DetachClauseSkip<
: OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
let traits = [BlockArgOpenMPOpInterface];
let arguments = (ins Optional<OpenMP_PointerLikeType>:$event_handle);
let optAssemblyFormat = [{
@ -1126,10 +1113,6 @@ class OpenMP_PrivateClauseSkip<
OptionalAttr<SymbolRefArrayAttr>:$private_syms
);
let extraClassDeclaration = [{
unsigned numPrivateBlockArgs() { return getPrivateVars().size(); }
}];
// TODO: Add description.
// Assembly format not defined because this clause must be processed together
// with the first region of the operation, as it defines entry block
@ -1186,7 +1169,6 @@ class OpenMP_ReductionClauseSkip<
let extraClassDeclaration = [{
/// Returns the number of reduction variables.
unsigned getNumReductionVars() { return getReductionVars().size(); }
unsigned numReductionBlockArgs() { return getReductionVars().size(); }
}];
// Description varies depending on the operation.
@ -1316,14 +1298,10 @@ class OpenMP_TaskReductionClauseSkip<
let extraClassDeclaration = [{
/// Returns the reduction variables.
SmallVector<Value> getReductionVars() {
SmallVector<Value> getAllReductionVars() {
return SmallVector<Value>(getTaskReductionVars().begin(),
getTaskReductionVars().end());
}
unsigned numTaskReductionBlockArgs() {
return getTaskReductionVars().size();
}
}];
let description = [{
@ -1413,12 +1391,6 @@ class OpenMP_UseDeviceAddrClauseSkip<
Variadic<OpenMP_PointerLikeType>:$use_device_addr_vars
);
let extraClassDeclaration = [{
unsigned numUseDeviceAddrBlockArgs() {
return getUseDeviceAddrVars().size();
}
}];
let description = [{
The optional `use_device_addr_vars` specifies the address of the objects in
the device data environment.
@ -1448,12 +1420,6 @@ class OpenMP_UseDevicePtrClauseSkip<
Variadic<OpenMP_PointerLikeType>:$use_device_ptr_vars
);
let extraClassDeclaration = [{
unsigned numUseDevicePtrBlockArgs() {
return getUseDevicePtrVars().size();
}
}];
let description = [{
The optional `use_device_ptr_vars` specifies the device pointers to the
corresponding list items in the device data environment.

51
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
View File

@ -284,8 +284,8 @@ def SectionOp : OpenMP_Op<"section", traits = [
// Override BlockArgOpenMPOpInterface methods based on the parent
// omp.sections operation. Only forward-declare here because SectionsOp is
// not completely defined at this point.
unsigned numPrivateBlockArgs();
unsigned numReductionBlockArgs();
OperandRange getPrivateVars();
OperandRange getReductionVars();
}] # clausesExtraClassDeclaration;
let assemblyFormat = "$region attr-dict";
}
@ -824,11 +824,6 @@ def TaskloopOp : OpenMP_Op<"taskloop", traits = [
/// Returns the reduction variables
SmallVector<Value> getAllReductionVars();
// Define BlockArgOpenMPOpInterface methods here because they are not
// inherited from the respective clauses.
unsigned numInReductionBlockArgs() { return getInReductionVars().size(); }
unsigned numReductionBlockArgs() { return getReductionVars().size(); }
void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
}] # clausesExtraClassDeclaration;
@ -1151,6 +1146,14 @@ def TargetDataOp: OpenMP_Op<"target_data", traits = [
OpBuilder<(ins CArg<"const TargetDataOperands &">:$clauses)>
];
let extraClassDeclaration = [{
// Override BlockArgOpenMPOpInterface method because `map` clauses have no
// associated entry block arguments in this operation.
unsigned numMapBlockArgs() {
return 0;
}
}] # clausesExtraClassDeclaration;
let assemblyFormat = clausesAssemblyFormat # [{
custom<UseDeviceAddrUseDevicePtrRegion>(
$region, $use_device_addr_vars, type($use_device_addr_vars),
@ -1185,6 +1188,14 @@ def TargetEnterDataOp: OpenMP_Op<"target_enter_data", traits = [
OpBuilder<(ins CArg<"const TargetEnterExitUpdateDataOperands &">:$clauses)>
];
let extraClassDeclaration = [{
// Override BlockArgOpenMPOpInterface method because `map` clauses have no
// associated entry block arguments in this operation.
unsigned numMapBlockArgs() {
return 0;
}
}] # clausesExtraClassDeclaration;
let hasVerifier = 1;
}
@ -1213,6 +1224,14 @@ def TargetExitDataOp: OpenMP_Op<"target_exit_data", traits = [
OpBuilder<(ins CArg<"const TargetEnterExitUpdateDataOperands &">:$clauses)>
];
let extraClassDeclaration = [{
// Override BlockArgOpenMPOpInterface method because `map` clauses have no
// associated entry block arguments in this operation.
unsigned numMapBlockArgs() {
return 0;
}
}] # clausesExtraClassDeclaration;
let hasVerifier = 1;
}
@ -1249,6 +1268,14 @@ def TargetUpdateOp: OpenMP_Op<"target_update", traits = [
OpBuilder<(ins CArg<"const TargetEnterExitUpdateDataOperands &">:$clauses)>
];
let extraClassDeclaration = [{
// Override BlockArgOpenMPOpInterface method because `map` clauses have no
// associated entry block arguments in this operation.
unsigned numMapBlockArgs() {
return 0;
}
}] # clausesExtraClassDeclaration;
let hasVerifier = 1;
}
@ -1292,8 +1319,6 @@ def TargetOp : OpenMP_Op<"target", traits = [
];
let extraClassDeclaration = [{
unsigned numMapBlockArgs() { return getMapVars().size(); }
mlir::Value getMappedValueForPrivateVar(unsigned privVarIdx) {
std::optional<DenseI64ArrayAttr> privateMapIdices = getPrivateMapsAttr();
@ -1818,6 +1843,14 @@ def DeclareMapperInfoOp : OpenMP_Op<"declare_mapper.info", [
OpBuilder<(ins CArg<"const DeclareMapperInfoOperands &">:$clauses)>
];
let extraClassDeclaration = [{
// Override BlockArgOpenMPOpInterface method because `map` clauses have no
// associated entry block arguments in this operation.
unsigned numMapBlockArgs() {
return 0;
}
}] # clausesExtraClassDeclaration;
let hasVerifier = 1;
}

57
mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
View File

@ -23,19 +23,41 @@ include "mlir/IR/OpBase.td"
// arguments corresponding to each of these clauses.
class BlockArgOpenMPClause<string clauseNameSnake, string clauseNameCamel,
BlockArgOpenMPClause previousClause> {
// Default-implemented method to be overriden by the corresponding clause.
// Default-implemented method, overriden by the corresponding clause. It
// returns the range of operands passed to the operation associated to the
// clause.
//
// For the override to work, the clause tablegen definition must contain a
// `Variadic<...> $clause_name_vars` argument.
//
// Usage example:
//
// ```c++
// auto iface = cast<BlockArgOpenMPOpInterface>(op);
// unsigned numInReductionArgs = iface.numInReductionBlockArgs();
// OperandRange reductionVars = op.getReductionVars();
// ```
InterfaceMethod varsMethod = InterfaceMethod<
"Get operation operands associated to `" # clauseNameSnake # "`.",
"::mlir::OperandRange", "get" # clauseNameCamel # "Vars", (ins), [{}], [{
return {0, 0};
}]
>;
// It returns the number of entry block arguments introduced by the given
// clause.
//
// By default, it will be the number of operands corresponding to that clause,
// but it can be overriden by operations where this might not be the case
// (e.g. `map` clause in `omp.target_update`).
//
// Usage example:
//
// ```c++
// unsigned numInReductionArgs = op.numInReductionBlockArgs();
// ```
InterfaceMethod numArgsMethod = InterfaceMethod<
"Get number of block arguments defined by `" # clauseNameSnake # "`.",
"unsigned", "num" # clauseNameCamel # "BlockArgs", (ins), [{}], [{
return 0;
}]
"unsigned", "num" # clauseNameCamel # "BlockArgs", (ins), [{}],
"return $_op." # varsMethod.name # "().size();"
>;
// Unified access method for the start index of clause-associated entry block
@ -52,7 +74,7 @@ class BlockArgOpenMPClause<string clauseNameSnake, string clauseNameCamel,
"unsigned", "get" # clauseNameCamel # "BlockArgsStart", (ins),
!if(!initialized(previousClause), [{
auto iface = ::llvm::cast<BlockArgOpenMPOpInterface>(*$_op);
}] # "return iface." # previousClause.startMethod.name # "() + $_op."
}] # "return iface." # previousClause.startMethod.name # "() + iface."
# previousClause.numArgsMethod.name # "();",
"return 0;"
)
@ -72,7 +94,7 @@ class BlockArgOpenMPClause<string clauseNameSnake, string clauseNameCamel,
"get" # clauseNameCamel # "BlockArgs", (ins), [{
auto iface = ::llvm::cast<BlockArgOpenMPOpInterface>(*$_op);
return $_op->getRegion(0).getArguments().slice(
}] # "iface." # startMethod.name # "(), $_op." # numArgsMethod.name # "());"
}] # "iface." # startMethod.name # "(), iface." # numArgsMethod.name # "());"
>;
}
@ -109,9 +131,26 @@ def BlockArgOpenMPOpInterface : OpInterface<"BlockArgOpenMPOpInterface"> {
BlockArgUseDeviceAddrClause, BlockArgUseDevicePtrClause ];
let methods = !listconcat(
!foreach(clause, clauses, clause.varsMethod),
!foreach(clause, clauses, clause.numArgsMethod),
!foreach(clause, clauses, clause.startMethod),
!foreach(clause, clauses, clause.blockArgsMethod)
!foreach(clause, clauses, clause.blockArgsMethod),
[
InterfaceMethod<
"Populate a vector of pairs representing the matching between operands "
"and entry block arguments.", "void", "getBlockArgsPairs",
(ins "::llvm::SmallVectorImpl<std::pair<::mlir::Value, ::mlir::BlockArgument>> &" : $pairs),
[{
auto iface = ::llvm::cast<BlockArgOpenMPOpInterface>(*$_op);
}] # !interleave(!foreach(clause, clauses, [{
}] # "if (iface." # clause.numArgsMethod.name # "() > 0) {" # [{
}] # " for (auto [var, arg] : ::llvm::zip_equal(" #
"iface." # clause.varsMethod.name # "()," #
"iface." # clause.blockArgsMethod.name # "()))" # [{
pairs.emplace_back(var, arg);
} }]), "\n")
>
]
);
let verify = [{

8
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
View File

@ -2248,12 +2248,12 @@ LogicalResult TeamsOp::verify() {
// SectionOp
//===----------------------------------------------------------------------===//
unsigned SectionOp::numPrivateBlockArgs() {
return getParentOp().numPrivateBlockArgs();
OperandRange SectionOp::getPrivateVars() {
return getParentOp().getPrivateVars();
}
unsigned SectionOp::numReductionBlockArgs() {
return getParentOp().numReductionBlockArgs();
OperandRange SectionOp::getReductionVars() {
return getParentOp().getReductionVars();
}
//===----------------------------------------------------------------------===//