[Mlir-commits] [mlir] [OpenMP][MLIR] Add num_teams clause with dims modifier support (PR #169883)

[Michael Klemm]

================
@@ -1532,4 +1532,76 @@ class OpenMP_UseDevicePtrClauseSkip<
 
 def OpenMP_UseDevicePtrClause : OpenMP_UseDevicePtrClauseSkip<>;
 
+//===----------------------------------------------------------------------===//
+// V6.2: Multidimensional `num_teams` clause with dims modifier
+//===----------------------------------------------------------------------===//
+
+class OpenMP_NumTeamsMultiDimClauseSkip<
+    bit traits = false, bit arguments = false, bit assemblyFormat = false,
+    bit description = false, bit extraClassDeclaration = false
+  > : OpenMP_Clause<traits, arguments, assemblyFormat, description,
+            extraClassDeclaration> {
+  let arguments = (ins
+    ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_dims,
+    Variadic<AnyInteger>:$num_teams_values
+  );
+
+  let optAssemblyFormat = [{
+    `num_teams_multi_dim` `(` custom<NumTeamsMultiDimClause>($num_teams_dims,
+                                                    $num_teams_values,
+                                                    type($num_teams_values)) `)`
+  }];
+
+  let description = [{
+    The `num_teams_multi_dim` clause with dims modifier support specifies the limit on
+    the number of teams to be created in a multidimensional team space.
+
+    The dims modifier for the num_teams_multi_dim clause specifies the number of
+    dimensions for the league space (team space) that the clause arranges.
+    The dimensions argument in the dims modifier specifies the number of
+    dimensions and determines the length of the list argument. The list items
+    are specified in ascending order according to the ordinal number of the
+    dimensions (dimension 0, 1, 2, ..., N-1).
+
+    - If `dims` is not specified: The space is unidimensional (1D) with a single value
+    - If `dims(1)` is specified: The space is explicitly unidimensional (1D)
+    - If `dims(N)` where N > 1: The space is strictly multidimensional (N-D)
+
+    **Examples:**
+    - `num_teams_multi_dim(dims(3): %nt0, %nt1, %nt2 : i32, i32, i32)` creates a
+      3-dimensional team space with limits nt0, nt1, nt2 for dimensions 0, 1, 2.
+    - `num_teams_multi_dim(%nt : i32)` creates a unidimensional team space with limit nt.
+  }];
+
+  let extraClassDeclaration = [{
+    /// Returns true if the dims modifier is explicitly present
+    bool hasDimsModifier() {
+      return getNumTeamsDims().has_value();
+    }
+
+    /// Returns the number of dimensions specified by dims modifier
+    /// Returns 1 if dims modifier is not present (unidimensional by default)
+    unsigned getNumDimensions() {
+      if (!hasDimsModifier())
+        return 1;
+      return static_cast<unsigned>(*getNumTeamsDims());
+    }
+
+    /// Returns all dimension values as an operand range
+    ::mlir::OperandRange getDimensionValues() {
+      return getNumTeamsValues();
+    }
+
+    /// Returns the value for a specific dimension index
+    /// Index must be less than getNumDimensions()
+    ::mlir::Value getDimensionValue(unsigned index) {
+    assert(index < getDimensionValues().size() &&
+        "Dimension index out of bounds");
+      return getDimensionValues()[index];
+    }
+  }];
+}
+
+def OpenMP_NumTeamsMultiDimClause : OpenMP_NumTeamsMultiDimClauseSkip<>;
----------------
mjklemm wrote:

Should this be rather called modifier instead of clause?  The clause still is `num_threads`, but the modifier is `dims`.

https://github.com/llvm/llvm-project/pull/169883