[Mlir-commits] [mlir] [MLIR][SCF] Add dedicated Python bindings for ForallOp (PR #149416)

[Oleksandr Alex Zinenko]

================
@@ -71,6 +71,120 @@ def inner_iter_args(self):
         return self.body.arguments[1:]
 
 
+def dispatch_index_op_fold_results(
+    ofrs: Sequence[Union[int, Value]],
+) -> Tuple[List[Value], List[int]]:
+    """`mlir::dispatchIndexOpFoldResults`"""
+    dynamic_vals = []
+    static_vals = []
+    for ofr in ofrs:
+        if isinstance(ofr, Value):
+            dynamic_vals.append(ofr)
+            static_vals.append(ShapedType.get_dynamic_size())
+        else:
+            static_vals.append(ofr)
+    return dynamic_vals, static_vals
+
+
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class ForallOp(ForallOp):
+    """Specialization for the SCF forall op class."""
+
+    def __init__(
+        self,
+        lower_bounds: Sequence[Union[Value, int]],
+        upper_bounds: Sequence[Union[Value, int]],
+        steps: Sequence[Union[Value, int]],
+        iter_args: Optional[Union[Operation, OpView, Sequence[Value]]] = None,
+        *,
+        mapping=None,
+        loc=None,
+        ip=None,
+    ):
+        """Creates an SCF `forall` operation.
+
+        - `lower_bounds` are the values to use as lower bounds of the loop.
+        - `upper_bounds` are the values to use as upper bounds of the loop.
+        - `steps` are the values to use as loop steps.
+        - `iter_args` is a list of additional loop-carried arguments or an operation
+          producing them as results.
+        """
+        if iter_args is None:
+            iter_args = []
+        iter_args = _get_op_results_or_values(iter_args)
+
+        dynamic_lbs, static_lbs = dispatch_index_op_fold_results(lower_bounds)
+        dynamic_ubs, static_ubs = dispatch_index_op_fold_results(upper_bounds)
+        dynamic_steps, static_steps = dispatch_index_op_fold_results(steps)
+
+        results = [arg.type for arg in iter_args]
+        super().__init__(
+            results,
+            dynamic_lbs,
+            dynamic_ubs,
+            dynamic_steps,
+            static_lbs,
+            static_ubs,
+            static_steps,
+            iter_args,
+            mapping=mapping,
+            loc=loc,
+            ip=ip,
+        )
+        rank = len(static_lbs)
+        iv_types = [IndexType.get()] * rank
+        self.regions[0].blocks.append(*iv_types, *results)
+
+    @property
+    def body(self) -> Block:
+        """Returns the body (block) of the loop."""
+        return self.regions[0].blocks[0]
+
+    @property
+    def rank(self) -> int:
+        """Returns the number of induction variables the loop has."""
+        return len(self.staticLowerBound)
+
+    @property
+    def induction_variables(self) -> BlockArgumentList:
+        """Returns the induction variables usable within the loop."""
+        return self.body.arguments[: self.rank]
+
+    @property
+    def inner_iter_args(self):
----------------
ftynse wrote:

```suggestion
    def inner_iter_args(self) -> BlockArgumentList:
```

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