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,8 +93,8 @@ 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,
         *,
@@ -109,13 +110,16 @@ def __init__(
         - `iter_args` is a list of additional loop-carried arguments or an operation
           producing them as results.
         """
+        assert (
+            len(lower_bounds) == len(upper_bounds) == len(steps)
+        ), "Mismatch in length of lower bounds, upper bounds, and steps"
         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)
+        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__(
@@ -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):