[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 = []
----------------
ftynse wrote:

```suggestion
```

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