[MLIR][OpenMP] Support clause-based representation of operations

[skatrak]

Currently, OpenMP operations are defined independently of each other. However, one property of the OpenMP specification is that many clauses can be applied to multiple constructs.

Keeping the MLIR representation of clauses consistent across all operations that can accept them is important, but since this information is scattered into multiple operation definitions, it is currently prone to divergence as new features and changes are added to the dialect. Furthermore, centralizing this information allows for a single source of truth and avoids redundancy in the dialect.

The proposal in this patch is to make OpenMP clauses independent top level definitions which can then be passed in a template argument list to OpenMP operation definitions, just as it's done for traits. Clauses can define these properties, which are joined together in order to make a default initialization for the fields of the same name of the OpenMP operation:

• traits: Optional. It gets added to the list of traits of the operation.
• arguments: Mandatory. It defines how the clause is represented.
• assemblyFormat: Optional (though it should almost always be defined). This is the declarative definition of the printer/parser for the arguments. How these are combined depends on whether this is an optional or required clause.
• description: Optional. It's used to populate a clausesDescription field, so each operation definition must still define a description itself. That field is intended to be appended to the end of the OpenMP_Op's description.
• extraClassDeclaration: Optional. It can define some C++ code to be added to every OpenMP operation that includes that clause.

In order to give operation definitions fine-grained control over features of a certain clause might need to be inhibited, the OpenMP_Clause class takes "skipTraits", "skipArguments", "skipAssemblyFormat", "skipDescription" and "skipExtraClassDeclaration" bit template arguments. These are intended to be used very sparingly for cases where some of the clauses might collide in some way otherwise.

[llvmbot]

@llvm/pr-subscribers-mlir-openmp
@llvm/pr-subscribers-flang-openmp

@llvm/pr-subscribers-mlir

Author: Sergio Afonso (skatrak)

Changes

Currently, OpenMP operations are defined independently of each other. However, one property of the OpenMP specification is that many clauses can be applied to multiple constructs.

Keeping the MLIR representation of clauses consistent across all operations that can accept them is important, but since this information is scattered into multiple operation definitions, it is currently prone to divergence as new features and changes are added to the dialect. Furthermore, centralizing this information allows for a single source of truth and avoids redundancy in the dialect.

The proposal in this patch is to make OpenMP clauses independent top level definitions which can then be passed in a template argument list to OpenMP operation definitions, just as it's done for traits. Clauses can define these properties, which are joined together in order to make a default initialization for the fields of the same name of the OpenMP operation:

• traits: Optional. It gets added to the list of traits of the operation.
• arguments: Mandatory. It defines how the clause is represented.
• assemblyFormat: Optional (though it should almost always be defined). This is the declarative definition of the printer/parser for the arguments. How these are combined depends on whether this is an optional or required clause.
• description: Optional. It's used to populate a clausesDescription field, so each operation definition must still define a description itself. That field is intended to be appended to the end of the OpenMP_Op's description.
• extraClassDeclaration: Optional. It can define some C++ code to be added to every OpenMP operation that includes that clause.

In order to give operation definitions fine-grained control over features of a certain clause might need to be inhibited, the OpenMP_Clause class takes "skipTraits", "skipArguments", "skipAssemblyFormat", "skipDescription" and "skipExtraClassDeclaration" bit template arguments. These are intended to be used very sparingly for cases where some of the clauses might collide in some way otherwise.

---

Full diff: https://github.com/llvm/llvm-project/pull/92519.diff

1 Files Affected:

• (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td (+163-2)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
index b98d87aa74a6f..d93abd63977ef 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
@@ -42,7 +42,168 @@ def OpenMP_MapBoundsType : OpenMP_Type<"MapBounds", "map_bounds_ty"> {
 // Base classes for OpenMP dialect operations.
 //===----------------------------------------------------------------------===//
 
-class OpenMP_Op<string mnemonic, list<Trait> traits = []> :
-      Op<OpenMP_Dialect, mnemonic, traits>;
+// Base class for representing OpenMP clauses.
+//
+// Clauses are meant to be used in a mixin-style pattern to help define OpenMP
+// operations in a scalable way, since often the same clause can be applied to
+// multiple different operations.
+//
+// To keep the representation of clauses consistent across different operations,
+// each clause must define a set of arguments (values and attributes) which will
+// become input arguments of each OpenMP operation that accepts that clause.
+//
+// It is also recommended that an assembly format and description are defined
+// for each clause wherever posible, to make sure they are always printed,
+// parsed and described in the same way.
+//
+// Optionally, operation traits and extra class declarations might be attached
+// to clauses, which will be forwarded to all operations that include them.
+//
+// Each clause must specify whether it's required or optional. This impacts how
+// the `assemblyFormat` for operations including it get generated.
+//
+// An `OpenMP_Op` can inhibit the inheritance of `traits`, `arguments`,
+// `assemblyFormat`, `description` and `extraClassDeclaration` fields from any
+// given `OpenMP_Clause` by setting to 1 the corresponding "skip" template
+// argument bit.
+class OpenMP_Clause<bit isRequired, bit skipTraits, bit skipArguments,
+                    bit skipAssemblyFormat, bit skipDescription,
+                    bit skipExtraClassDeclaration> {
+  bit required = isRequired;
+
+  bit ignoreTraits = skipTraits;
+  list<Trait> traits = [];
+
+  bit ignoreArgs = skipArguments;
+  dag arguments;
+
+  bit ignoreAsmFormat = skipAssemblyFormat;
+  string assemblyFormat = "";
+
+  bit ignoreDesc = skipDescription;
+  string description = "";
+
+  bit ignoreExtraDecl = skipExtraClassDeclaration;
+  string extraClassDeclaration = "";
+}
+
+// Base class for representing OpenMP operations.
+//
+// This is a subclass of the builtin `Op` for the OpenMP dialect. By default,
+// some of its fields are initialized according to the list of OpenMP clauses
+// passed as template argument:
+//   - `traits`: It is a union of the traits list passed as template argument
+//     and those inherited from the `traits` field of all clauses.
+//   - `arguments`: They are a concatenation of clause-inherited arguments. They
+//     are saved to a `clausesArgs` field to allow overriding the arguments
+//     field in the definition of the operation and still being able to include
+//     those inherited from clauses.
+//   - `assemblyFormat`: It is a concatenation of the `assemblyFormat` of
+//     all required clauses followed by an `oilist()` containing the
+//     `assemblyFormat` of all optional clauses. The format string is completed
+//     with $region (if `singleRegion = true`) followed by attr-dict. This field
+//     remains uninitialized if no non-empty `assemblyFormat` strings are
+//     inherited from clauses. The `clausesAssemblyFormat` field holds
+//     all the format except for "$region attr-dict", so that an operation
+//     overriding `assemblyFormat` can still benefit from the auto-generated
+//     format for its clauses.
+//   - `description`: This is still required to be defined by the operation.
+//     However, a `clausesDescription` field is provided containing a
+//     concatenation of descriptions of all clauses, to be appended to the
+//     operation's `description` field.
+//   - `extraClassDeclaration`: It contains a concatenation of the
+//     `extraClassDeclaration` of all clauses. This string is also stored in
+//     `clausesExtraClassDeclaration`, so that an operation overriding this
+//     field can append the clause-inherited ones as well.
+//
+// The `regions` field will contain a single `AnyRegion:$region` element if the
+// `singleRegion` bit template argument is set to 1. Otherwise, it will be
+// empty.
+class OpenMP_Op<string mnemonic, list<Trait> traits = [],
+                list<OpenMP_Clause> clauses = [], bit singleRegion = false> :
+    Op<OpenMP_Dialect, mnemonic,
+    // The resulting operation's traits list will be the concatenation of
+    // explicit operation traits and all traits attached to the clauses of the
+    // operation. Repetitions are skipped.
+    !listconcat(traits,
+      !listremove(
+        !foldl([]<Trait>,
+               !foreach(clause,
+                        !filter(fClause, clauses, !not(fClause.ignoreTraits)),
+                        clause.traits),
+               acc, traitList, !listconcat(acc, !listremove(traitList, acc))),
+        traits
+      )
+    )> {
+  // Aggregate `arguments` fields of all clauses into a single dag, to be used
+  // by operations to populate their `arguments` field.
+  defvar argsFilteredClauses =
+    !filter(clause, clauses, !not(clause.ignoreArgs));
+
+  dag clausesArgs =
+    !foldl((ins), !foreach(clause, argsFilteredClauses, clause.arguments),
+           acc, argList, !con(acc, argList));
+
+  // Create assembly format string by concatenating format strings separately
+  // for required and optional clauses. Then, required clauses format strings
+  // are joined with spaces in between. Optional clauses format strings are
+  // wrapped into an unsorted list of optional values and separated by "|"
+  // characters.
+
+  // Required clauses.
+  defvar reqClauses = !filter(clause, clauses, clause.required);
+  defvar asmFormatFilteredReqClauses =
+    !filter(clause, reqClauses, !not(!or(clause.ignoreAsmFormat,
+                                     !empty(clause.assemblyFormat))));
+
+  defvar asmFormatReqClauseStrings =
+    !foreach(clause, asmFormatFilteredReqClauses, clause.assemblyFormat);
+
+  defvar asmFormatReqClauseBody = !interleave(asmFormatReqClauseStrings, " ");
+
+  // Optional clauses.
+  defvar optClauses = !filter(clause, clauses, !not(clause.required));
+  defvar asmFormatFilteredOptClauses =
+    !filter(clause, optClauses, !not(!or(clause.ignoreAsmFormat,
+                                     !empty(clause.assemblyFormat))));
+
+  defvar asmFormatOptClauseStrings =
+    !foreach(clause, asmFormatFilteredOptClauses, clause.assemblyFormat);
+
+  defvar asmFormatOptClauseBody = !interleave(asmFormatOptClauseStrings, "|");
+
+  string clausesAssemblyFormat =
+    !if(!empty(asmFormatReqClauseStrings), "", asmFormatReqClauseBody # " ") #
+    !if(!empty(asmFormatOptClauseStrings), "",
+        "oilist(" # asmFormatOptClauseBody # ")");
+
+  // Put together descriptions of all clauses into a single string.
+  defvar descFilteredClauses =
+    !filter(clause, clauses, !not(clause.ignoreDesc));
+
+  string clausesDescription =
+    !interleave(!foreach(clause, descFilteredClauses, clause.description), "");
+  
+  // Aggregate `extraClassDeclaration` of all clauses that define it.
+  defvar extraDeclFilteredClauses =
+    !filter(clause, clauses, !not(clause.ignoreExtraDecl));
+
+  string clausesExtraClassDeclaration =
+    !interleave(!foreach(clause, extraDeclFilteredClauses,
+                         clause.extraClassDeclaration), "\n");
+
+  // The default arguments, assembly format and extra class declarations for
+  // OpenMP operations are those defined by their args and clauses.
+  let arguments = clausesArgs;
+  let assemblyFormat =
+    !if(!empty(clausesAssemblyFormat), ?,
+        clausesAssemblyFormat # !if(singleRegion, " $region", "") #
+        " attr-dict");
+  let extraClassDeclaration = clausesExtraClassDeclaration;
+
+  // By default, the op will have zero regions. Setting `singleRegion = true`
+  // will result in a single region named `$region`.
+  let regions = !if(singleRegion, (region AnyRegion:$region), (region));
+}
 
 #endif  // OPENMP_OP_BASE

[kiranchandramohan]

@skatrak Could you copy the summary of the patch and create an RFC?

[Meinersbur]

I like the idea, but also have some questions:

1. How does an operation definition look like with all its clauses defined? Does it have to be repeated for each operation supporting the same clause(s)?
2. It seems this requires all properties of a clause specified in its constructor. Wouldn't it be better if you could subclass OpenMP_Clause and in there overwrite all the properties that are non-default?
3. Have you considered to define/derive from all this info in OpenMP.td of LLVMFrontend? Clause information should be language-independent.

[skatrak]

I like the idea, but also have some questions:

1. How does an operation definition look like with all its clauses defined? Does it have to be repeated for each operation supporting the same clause(s)?

You can see how things would look like in #92523, but I can copy one example here:

def TeamsOp : OpenMP_Op<"teams", traits = [
    AttrSizedOperandSegments, RecursiveMemoryEffects
  ], clauses = [
    OpenMP_NumTeamsClause, OpenMP_IfClause, OpenMP_ThreadLimitClause,
    OpenMP_AllocateClause, OpenMP_ReductionClause
  ], singleRegion = true> {
  let summary = "teams construct";
  let description = [{
    The teams construct defines a region of code that triggers the creation of a
    league of teams. Once created, the number of teams remains constant for the
    duration of its code region.
  }] # clausesDescription;

  let builders = [
    OpBuilder<(ins CArg<"const TeamsClauseOps &">:$clauses)>
  ];

  let hasVerifier = 1;
}

2. It seems this requires all properties of a clause specified in its constructor. Wouldn't it be better if you could subclass `OpenMP_Clause` and in there overwrite all the properties that are non-default?

Actually this only requires the list of clauses themselves to be passed in to the definition of the OpenMP_Op. They are the ones that define all these properties, and they indeed all subclass OpenMP_Clause (see in #92521).

3. Have you considered to define/derive from all this info in OpenMP.td of LLVMFrontend? Clause information should be language-independent.

I get that, the issue is that the information being used here is very much MLIR dialect specific, and OMP.td in LLVMFrontend doesn't seem like the right place to be defining these things, since it's outside of the MLIR project. Maybe descriptions could potentially be moved there, but then we'd have some sort of matching system between clauses in LLVMFrontend and the ones in MLIR. Not sure how feasible that would be.

[skatrak]

Here's a link to the RFC about this proposal, with links to all related PRs: https://discourse.llvm.org/t/rfc-clause-based-representation-of-openmp-dialect-operations/79053

[kiranchandramohan]

Looks OK in spirit. The RFC discourse post was well received and I don't see major objections. Please wait for further reviews.

[tblah]

+1 from me too

[skatrak]

Thank you, I just added some tablegen unit tests as recommended by @ftynse in the RFC.

[Meinersbur]

LGTM