[MLIR][OpenMP] Support clause-based representation of operations (#92519)

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.

Author: skatrak

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