Address comments

Author: Cubevoid

mlir/python/mlir/dialects/scf.py

@@ -71,15 +71,16 @@ def inner_iter_args(self):
         return self.body.arguments[1:]
 
 
-def dispatch_index_op_fold_results(
-    ofrs: Sequence[Union[int, Value]],
+def _dispatch_index_op_fold_results(
+    ofrs: Sequence[Union[Operation, OpView, Value, int]],
 ) -> Tuple[List[Value], List[int]]:
     """`mlir::dispatchIndexOpFoldResults`"""
     dynamic_vals = []
     static_vals = []
     for ofr in ofrs:
-        if isinstance(ofr, Value):
-            dynamic_vals.append(ofr)
+        if isinstance(ofr, (Operation, OpView, Value)):
+            val = _get_op_result_or_value(ofr)
+            dynamic_vals.append(val)
             static_vals.append(ShapedType.get_dynamic_size())
         else:
             static_vals.append(ofr)
@@ -92,10 +93,10 @@ class ForallOp(ForallOp):
 
     def __init__(
         self,
-        lower_bounds: Sequence[Union[Value, int]],
-        upper_bounds: Sequence[Union[Value, int]],
+        lower_bounds: Sequence[Union[Operation, OpView, Value, int]],
+        upper_bounds: Sequence[Union[Operation, OpView, Value, int]],
         steps: Sequence[Union[Value, int]],
-        iter_args: Optional[Union[Operation, OpView, Sequence[Value]]] = None,
+        shared_outs: Optional[Union[Operation, OpView, Sequence[Value]]] = None,
         *,
         mapping=None,
         loc=None,
@@ -106,18 +107,21 @@ def __init__(
         - `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
+        - `shared_outs` 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]
+        assert (
+            len(lower_bounds) == len(upper_bounds) == len(steps)
+        ), "Mismatch in length of lower bounds, upper bounds, and steps"
+        if shared_outs is None:
+            shared_outs = []
+        shared_outs = _get_op_results_or_values(shared_outs)
+
+        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 shared_outs]
         super().__init__(
             results,
             dynamic_lbs,
@@ -126,7 +130,7 @@ def __init__(
             static_lbs,
             static_ubs,
             static_steps,
-            iter_args,
+            shared_outs,
             mapping=mapping,
             loc=loc,
             ip=ip,
@@ -151,18 +155,17 @@ def induction_variables(self) -> BlockArgumentList:
         return self.body.arguments[: self.rank]
 
     @property
-    def inner_iter_args(self):
+    def inner_iter_args(self) -> BlockArgumentList:
         """Returns the loop-carried arguments usable within the loop.
 
         To obtain the loop-carried operands, use `iter_args`.
         """
         return self.body.arguments[self.rank :]
 
-    @property
     def terminator(self) -> InParallelOp:
         """
         Returns the loop terminator if it exists.
-        Otherwise, create a new one.
+        Otherwise, creates a new one.
         """
         ops = self.body.operations
         with InsertionPoint(self.body):

mlir/test/python/dialects/scf.py

@@ -33,7 +33,7 @@ def forall_loop(tensor):
         with InsertionPoint(loop.body):
             i, j = loop.induction_variables
             arith.addi(i, j)
-            loop.terminator
+            loop.terminator()
         # The verifier will check that the regions have been created properly.
         assert loop.verify()