[mlir][python] support taking ops instead of values in op constructors

[ftynse]

Summary

Introduce support for accepting ops instead of values when constructing ops. A
single-result op can be used instead of a value, including in lists of values,
and any op can be used instead of a list of values. This is similar to, but
more powerful, than the C++ API that allows for implicitly casting an OpType to
Value if it is statically known to have a single result - the cast in Python is
based on the op dynamically having a single result, and also handles the
multi-result case. This allows to build IR in a more concise way:

op = dialect.produce_multiple_results()
other = dialect.produce_single_result()
dialect.consume_multiple_results(other, op)

instead of having to access the results manually

op = dialect.produce.multiple_results()
other = dialect.produce_single_result()
dialect.consume_multiple_results(other.result, op.operation.results)

The dispatch is implemented directly in Python and is triggered automatically
for autogenerated OpView subclasses. Extension OpView classes should use the
functions provided in ods_common.py if they want to implement this behavior.
An alternative could be to implement the dispatch in the C++ bindings code, but
it would require to forward opaque types through all Python functions down to a
binding call, which makes it hard to inspect them in Python, e.g., to obtain
the types of values.

[gysit]

mlir/python/mlir/dialects/_ods_common.py

133

typo: reuslt -> result

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

48

Should we support inputs of type Union[Value, Operation, OpView] as well?

Additionally I wonder if it would be easier to do the StructuredOpOuts to ValueList conversion right in dsl.py (https://github.com/llvm/llvm-project/blob/94e2b0258a176c7451dd8291cdf060ea048fee44/mlir/python/mlir/dialects/linalg/opdsl/lang/dsl.py#L47) where the arguments are handed over to the operation?

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

654

nit: singleElementAppendTemplate -> multiOperandAppendPackTemplate

689

nit: singleElementAppendTemplate -> singleResultAppendTemplate

[ftynse]

Address review.

[jpienaar]

For ops producing multiple results, would this effectively flatten all results into args?

[gysit]

Thanks for these usability improvements!

[ftynse]

In D111306#3048422, @jpienaar wrote:

For ops producing multiple results, would this effectively flatten all results into args?

It would flatten them to one list. If the consumer of the values takes a single list of operands, it will be that. There are many cases where we have several variadic operand lists that are passed as independent arguments of the constructor. Each of those arguments may accept an op instead, but there is no intermixing between different lists. For example,

lbs = std.CallOp([IndexType.get()] * 3, "compute_lower_bounds", [])
ubs = std.CallOp([IndexType.get()] * 3, "compute_lower_bounds", [])
one = std.ConstantOp(IndexType.get(), 1)
scf.ParallelOp(lbs, ubs, [one] * 3)

each of lbs and ubs is an op that produces three values that are passed into individual groups of operands (or ODS-level arguments).

mlir/python/mlir/dialects/_linalg_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from typing import Optional, Sequence, Union
   from ..ir import *
   from ._ods_common import get_default_loc_context
   from .._mlir_libs._mlir.dialects.linalg import fill_builtin_region
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
+from ._ods_common import get_op_result_or_value as _get_op_result_or_value
 
 def isa(cls: Type, ty: Type):
   try:
     cls(ty)
     return True
   except ValueError:
     return False
 
 
 class FillOp:
   """Extends the linalg.fill op."""
 
   def __init__(self, output: Value, value: Value, *, loc=None, ip=None):
     results = []
     if isa(RankedTensorType, output.type):
       results = [output.type]
-    op = self.build_generic(results=results,
-                            operands=[value, output],
-                            attributes=None,
-                            loc=loc,
-                            ip=ip)
+    op = self.build_generic(
+        results=results,
+        operands=[_get_op_result_or_value(o) for o in [value, output]],
     OpView.__init__(self, op)
     linalgDialect = Context.current.get_dialect_descriptor("linalg")
     fill_builtin_region(linalgDialect, self.operation)
     # TODO: self.result is None. When len(results) == 1 we expect it to be
     # results[0] as per _linalg_ops_gen.py. This seems like an orthogonal bug
     # in the generator of _linalg_ops_gen.py where we have:
     # ```
     # def result(self):

mlir/python/mlir/dialects/_ods_common.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Provide a convenient name for sub-packages to resolve the main C-extension
 # with a relative import.
 from .._mlir_libs import _mlir as _cext
+from typing import Sequence as _Sequence, Union as _Union
 
 __all__ = [
     "equally_sized_accessor",
     "extend_opview_class",
     "get_default_loc_context",
+    "get_op_result_or_value",
+    "get_op_results_or_values",
     "segmented_accessor",
 ]
 
 
 def extend_opview_class(ext_module):
   """Decorator to extend an OpView class from an extension module.
 
   Extension modules can expose various entry-points:
def get_default_loc_context(location=None):
   Returns a context in which the defaulted location is created. If the location
   is None, takes the current location from the stack, raises ValueError if there
   is no location on the stack.
   """
   if location is None:
     # Location.current raises ValueError if there is no current location.
     return _cext.ir.Location.current.context
   return location.context
+
+
+def get_op_result_or_value(
+    arg: _Union[_cext.ir.OpView, _cext.ir.Operation, _cext.ir.Value]
+) -> _cext.ir.Value:
+  """Returns the given value or the single result of the given op.
+
+  This is useful to implement op constructors so that they can take other ops as
+  arguments instead of requiring the caller to extract results for every op.
+  Raises ValueError if provided with an op that doesn't have a single reuslt.
+  """
+  if isinstance(arg, _cext.ir.OpView):
+    return arg.operation.result
+  elif isinstance(arg, _cext.ir.Operation):
+    return arg.result
+  else:
+    assert isinstance(arg, _cext.ir.Value)
+    return arg
+
+
+def get_op_results_or_values(
+    arg: _Union[_cext.ir.OpView, _cext.ir.Operation, _Sequence[_cext.ir.Value]]
+) -> _Union[_Sequence[_cext.ir.Value], _cext.ir.OpResultList]:
+  """Returns the given sequence of values or the results of the given op.
+
+  This is useful to implement op constructors so that they can take other ops as
+  lists of arguments instead of requiring the caller to extract results for
+  every op.
+  """
+  if isinstance(arg, _cext.ir.OpView):
+    return arg.operation.results
+  elif isinstance(arg, _cext.ir.Operation):
+    return arg.results
+  else:
+    return arg

mlir/python/mlir/dialects/_scf_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from ..ir import *
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
-from typing import Any, Sequence
-
+from typing import Any, Optional, Sequence, Union
+from ._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
 
 class ForOp:
   """Specialization for the SCF for op class."""
 
   def __init__(self,
                lower_bound,
                upper_bound,
                step,
-               iter_args: Sequence[Any] = [],
+               iter_args: Optional[Union[Operation, OpView,
+                                         Sequence[Value]]] = None,
                *,
                loc=None,
                ip=None):
     """Creates an SCF `for` operation.
 
     - `lower_bound` is the value to use as lower bound of the loop.
     - `upper_bound` is the value to use as upper bound of the loop.
     - `step` is the value to use as loop step.
-    - `iter_args` is a list of additional loop-carried arguments.
+    - `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)
+
     results = [arg.type for arg in iter_args]
     super().__init__(
         self.build_generic(
             regions=1,
             results=results,
-            operands=[lower_bound, upper_bound, step] + list(iter_args),
+            operands=[
+                _get_op_result_or_value(o)
+                for o in [lower_bound, upper_bound, step]
+            ] + list(iter_args),
             loc=loc,
             ip=ip))
     self.regions[0].blocks.append(IndexType.get(), *results)
 
   @property
   def body(self):
     """Returns the body (block) of the loop."""
     return self.regions[0].blocks[0]

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Dict, Sequence
+from typing import Dict, List, Sequence, Tuple, Union
 
 from .....ir import *
 from ....._mlir_libs._mlir.dialects.linalg import fill_builtin_region
 
 from .... import linalg
 from .... import std
 from .... import math
+from ...._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
 
 from .scalar_expr import *
 from .config import *
 import numpy as np
 
 __all__ = [
     "emit_generic_structured_op",
     "emit_named_structured_op",
+    "StructuredOpOuts",
+    "ValueList",
 ]
 
+StructuredOpOuts = Union[Operation, OpView, OpResultList,
+                         Sequence[Union[Value, Operation, OpView]]]
+ValueList = Union[Sequence[Value], OpResultList]
 
 def isa(cls: Type, ty: Type):
   try:
     cls(ty)
     return True
   except ValueError:
     return False
 
 
+def prepare_structure_op_outs(outs: StructuredOpOuts) -> ValueList:
+  if isinstance(outs, (Operation, OpView)):
+    return _get_op_results_or_values(outs)
+  elif isinstance(outs, OpResultList):
+    return outs
+
+  return [_get_op_result_or_value(o) for o in outs]
+
+
 def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
-                                 *ins: Value, outs: Sequence[Value],
+                                 *ins: Value, outs: StructuredOpOuts,
                                  **attrs: Sequence[int]):
   all_arg_defs = op_config.ordered_operands
   in_arg_defs = [arg for arg in all_arg_defs if arg.usage == "InputOperand"]
   out_arg_defs = [arg for arg in all_arg_defs if arg.usage == "OutputOperand"]
   attr_arg_defs = [arg for arg in all_arg_defs if arg.usage == "IndexAttribute"]
+  unpacked_outs = prepare_structure_op_outs(outs)
 
-  # Verify outs is a sequence.
-  if not isinstance(outs, Sequence):
-    raise ValueError(f"Expected named argument outs to have type Sequence "
-                     f"but got {type(outs)}")
+  # Verify outs is a sequence or a list of results.
+  if not isinstance(unpacked_outs, (Sequence, OpResultList)):
+    raise ValueError(
+        f"Expected named argument outs to have type Sequence or OpResultLis but got {type(unpacked_outs)}"
+    )
 
   # Arity validation.
   if len(ins) != len(in_arg_defs):
     raise ValueError(f"Expected {len(in_arg_defs)} inputs but got "
                      f"{len(ins)} for {op_config}")
-  if outs and len(outs) != len(out_arg_defs):
+  if unpacked_outs and len(unpacked_outs) != len(out_arg_defs):
     raise ValueError(f"Expected {len(out_arg_defs)} outputs but got "
-                     f"{len(outs)} for {op_config}")
+                     f"{len(unpacked_outs)} for {op_config}")
 
   # Compute a replacement list for all attribute symbols.
   expressions = []  # type: Sequence[AffineExpr]
   replacements = []  # type: Sequence[AffineExpr]
   for attr in attr_arg_defs:
     if attr.name not in attrs:
       raise ValueError(f"Expected named argument for the attribute {attr.name}")
     attribute_values = attrs.get(attr.name)
def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
   # Compress them for now and verify none are left.
   indexing_maps = AffineMap.compress_unused_symbols(indexing_maps,
                                                     Context.current)
   if any(indexing_map.n_symbols != 0 for indexing_map in indexing_maps):
     raise ValueError(f"Expected indexing_maps to use no symbols after "
                      f"replacement and compression but got {indexing_maps}")
 
   outs, out_types = _infer_structured_outs(op_config, in_arg_defs, ins,
-                                           out_arg_defs, outs)
+                                           out_arg_defs, unpacked_outs)
 
   result_types = [t for t in out_types if isa(RankedTensorType, t)]
 
   # Initialize the type dictionary with the predefined types.
   type_mapping = dict()  # type: Dict[str, Type]
   type_mapping["F32"] = F32Type.get()
   type_mapping["F64"] = F64Type.get()
   type_mapping["I32"] = IntegerType.get_signless(32)
for attr in attr_arg_defs:
     index_attributes[attr.name] = DenseElementsAttr.get(array)
 
   return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
           type_mapping, indexing_maps_attr, iterator_types_attr,
           index_attributes, block_arg_types)
 
 
 def emit_generic_structured_op(op_config: LinalgStructuredOpConfig, *ins: Value,
-                               outs: Sequence[Value], **attrs: Sequence[int]):
+                               outs: StructuredOpOuts, **attrs: Sequence[int]):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
   indexing_maps_attr, iterator_types_attr, index_attributes, block_arg_types = \
      prepare_common_structured_op(op_config, *ins, outs = outs, **attrs)
 
   generic_op = linalg.GenericOp(
       result_tensors=result_types,
       inputs=ins,
       outputs=outs,
def emit_generic_structured_op(op_config: LinalgStructuredOpConfig, *ins: Value,
   if len(result_types) == 1:
     return generic_op.result
   else:
     return generic_op.results
 
 
 def emit_named_structured_op(op_config: LinalgStructuredOpConfig, op_name: str,
                              op_class_name: str, *ins: Value,
-                             outs: Sequence[Value], **attrs: Sequence[int]):
+                             outs: StructuredOpOuts, **attrs: Sequence[int]):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
   indexing_maps_attr, iterator_types_attr, index_attributes, block_arg_types = \
      prepare_common_structured_op(op_config, *ins, outs = outs, **attrs)
 
   # If we get here, there must exist a builtin class `op_class_name`.
   ctx = Context.current
   fully_qualified_name = "linalg." + op_name
   if (not ctx.is_registered_operation(fully_qualified_name) or
class _BodyBuilder:
   def _eval_min(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return std.MinFOp(lhs.type, lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return std.MinSIOp(lhs.type, lhs, rhs).result
     raise NotImplementedError("Unsupported 'min' operand: {lhs}")
 
 
-def _infer_structured_outs(op_config: LinalgStructuredOpConfig,
-                           in_arg_defs: Sequence[OperandDefConfig],
-                           ins: Sequence[Value],
-                           out_arg_defs: Sequence[OperandDefConfig],
-                           outs: Sequence[Value]):
+def _infer_structured_outs(
+    op_config: LinalgStructuredOpConfig,
+    in_arg_defs: Sequence[OperandDefConfig], ins: Sequence[Value],
+    outs: Union[Sequence[Value], OpResultList]) -> Tuple[ValueList, List[Type]]:
   """Infers implicit outs and output types.
 
   Respects existing contents of outs if not empty.
 
   Returns:
     normalized outs, output types
   """
   # If outs were explicitly provided, we accept them verbatim.

mlir/test/mlir-tblgen/op-python-bindings.td

 // CHECK-LABEL: OPERATION_NAME = "test.attr_sized_operands"
 // CHECK: _ODS_OPERAND_SEGMENTS = [-1,1,-1,]
 def AttrSizedOperandsOp : TestOp<"attr_sized_operands",
                                  [AttrSizedOperandSegments]> {
   // CHECK: def __init__(self, variadic1, non_variadic, variadic2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
-  // CHECK:   operands.append(variadic1)
-  // CHECK:   operands.append(non_variadic)
-  // CHECK:   if variadic2 is not None: operands.append(variadic2)
+  // CHECK:   operands.append(_get_op_results_or_values(variadic1))
+  // CHECK:   operands.append(_get_op_result_or_value(non_variadic))
+  // CHECK:   if variadic2 is not None: operands.append(_get_op_result_or_value(variadic2))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):
   // CHECK:   operand_range = _ods_segmented_accessor(

 // CHECK-LABEL: OPERATION_NAME = "test.attributed_op_with_operands"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def AttributedOpWithOperands : TestOp<"attributed_op_with_operands"> {
   // CHECK: def __init__(self, _gen_arg_0, in_, _gen_arg_2, is_, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
-  // CHECK:   operands.append(_gen_arg_0)
-  // CHECK:   operands.append(_gen_arg_2)
+  // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_0))
+  // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_2))
   // CHECK:   if bool(in_): attributes["in"] = _ods_ir.UnitAttr.get(
   // CHECK:     _ods_get_default_loc_context(loc))
   // CHECK:   if is_ is not None: attributes["is"] = is_
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 

 def MissingNamesOp : TestOp<"missing_names"> {
   // CHECK: def __init__(self, i32, _gen_res_1, i64, _gen_arg_0, f32, _gen_arg_2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   results.append(i32)
   // CHECK:   results.append(_gen_res_1)
   // CHECK:   results.append(i64)
-  // CHECK:   operands.append(_gen_arg_0)
-  // CHECK:   operands.append(f32)
-  // CHECK:   operands.append(_gen_arg_2)
+  // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_0))
+  // CHECK:   operands.append(_get_op_result_or_value(f32))
+  // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_2))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def f32(self):
   // CHECK:   return self.operation.operands[1]

 // CHECK-LABEL: OPERATION_NAME = "test.one_variadic_operand"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def OneVariadicOperandOp : TestOp<"one_variadic_operand"> {
   // CHECK: def __init__(self, non_variadic, variadic, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
-  // CHECK:   operands.append(non_variadic)
-  // CHECK:   operands.extend(variadic)
+  // CHECK:   operands.append(_get_op_result_or_value(non_variadic))
+  // CHECK:   operands.extend(_get_op_results_or_values(variadic))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def non_variadic(self):
   // CHECK:   return self.operation.operands[0]

 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class PythonKeywordOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.python_keyword"
 def PythonKeywordOp : TestOp<"python_keyword"> {
   // CHECK: def __init__(self, in_, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
-  // CHECK:   operands.append(in_)
+  // CHECK:   operands.append(_get_op_result_or_value(in_))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def in_(self):
   // CHECK:   return self.operation.operands[0]

 // CHECK-LABEL: OPERATION_NAME = "test.simple"
 def SimpleOp : TestOp<"simple"> {
   // CHECK: def __init__(self, i64, f64, i32, f32, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   results.append(i64)
   // CHECK:   results.append(f64)
-  // CHECK:   operands.append(i32)
-  // CHECK:   operands.append(f32)
+  // CHECK:   operands.append(_get_op_result_or_value(i32))
+  // CHECK:   operands.append(_get_op_result_or_value(f32))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def i32(self):
   // CHECK:   return self.operation.operands[0]

mlir/test/python/dialects/linalg/ops.py

with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                    RankedTensorType.get((16, 8), f32))
       def generic_form(lhs, rhs):
         init_result = linalg.InitTensorOp([4, 8], f32)
         # CHECK: linalg.generic
         return linalg.matmul(lhs, rhs, outs=[init_result.result], emit_generic=True)
 
   print(module)
+
+
+# CHECK-LABEL: TEST: testOpResultFromOtherOp
+@run
+def testOpResultFromOtherOp():
+  with Context(), Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    with InsertionPoint(module.body):
+
+      @builtin.FuncOp.from_py_func(
+          RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
+                                                                   f32))
+      def pass_an_op_directly(lhs, rhs):
+        # CHECK: %[[INIT:.*]] = linalg.init_tensor
+        init = linalg.InitTensorOp([4, 8], f32)
+        # CHECK: linalg.matmul
+        # CHECK: outs(%[[INIT]]
+        return linalg.matmul(lhs, rhs, outs=init)
+
+  print(module)

mlir/test/python/dialects/scf.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import scf
+from mlir.dialects import std
 from mlir.dialects import builtin
 
 
-def run(f):
+def constructAndPrintInModule(f):
   print("\nTEST:", f.__name__)
+  with Context(), Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
       f()
+    print(module)
   return f
 
 
 # CHECK-LABEL: TEST: testSimpleLoop
-@run
+@constructAndPrintInModule
 def testSimpleLoop():
-  with Context(), Location.unknown():
-    module = Module.create()
-    index_type = IndexType.get()
-    with InsertionPoint(module.body):
 
-      @builtin.FuncOp.from_py_func(index_type, index_type, index_type)
-      def simple_loop(lb, ub, step):
-        loop = scf.ForOp(lb, ub, step, [lb, lb])
-        with InsertionPoint(loop.body):
-          scf.YieldOp(loop.inner_iter_args)
-        return
 
-  # CHECK: func @simple_loop(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
-  # CHECK: scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
-  # CHECK: iter_args(%[[I1:.*]] = %[[ARG0]], %[[I2:.*]] = %[[ARG0]])
-  # CHECK: scf.yield %[[I1]], %[[I2]]
-  print(module)
+
 
 
 # CHECK-LABEL: TEST: testInductionVar
-@run
+@constructAndPrintInModule
 def testInductionVar():
-  with Context(), Location.unknown():
-    module = Module.create()
-    index_type = IndexType.get()
-    with InsertionPoint(module.body):
 
-      @builtin.FuncOp.from_py_func(index_type, index_type, index_type)
-      def induction_var(lb, ub, step):
-        loop = scf.ForOp(lb, ub, step, [lb])
-        with InsertionPoint(loop.body):
-          scf.YieldOp([loop.induction_variable])
-        return
 
-  # CHECK: func @induction_var(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
-  # CHECK: scf.for %[[IV:.*]] = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
-  # CHECK: scf.yield %[[IV]]
-  print(module)
+
+
+
+@constructAndPrintInModule
+def testOpsAsArguments():
+  index_type = IndexType.get()
+  callee = builtin.FuncOp(
+      "callee", ([], [index_type, index_type]), visibility="private")
+  func = builtin.FuncOp("ops_as_arguments", ([], []))
+  with InsertionPoint(func.add_entry_block()):
+    lb = std.ConstantOp.create_index(0)
+    ub = std.ConstantOp.create_index(42)
+    step = std.ConstantOp.create_index(2)
+    iter_args = std.CallOp(callee, [])
+    loop = scf.ForOp(lb, ub, step, iter_args)
+    with InsertionPoint(loop.body):
+      scf.YieldOp(loop.inner_iter_args)
+    std.ReturnOp([])
+
+
+# CHECK-LABEL: TEST: testOpsAsArguments
+# CHECK: func private @callee() -> (index, index)
+# CHECK: func @ops_as_arguments() {
+# CHECK:   %[[LB:.*]] = constant 0
+# CHECK:   %[[UB:.*]] = constant 42
+# CHECK:   %[[STEP:.*]] = constant 2
+# CHECK:   %[[ARGS:.*]]:2 = call @callee()
+# CHECK:   scf.for %arg0 = %c0 to %c42 step %c2
+# CHECK:   iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)
+# CHECK:     scf.yield %{{.*}}, %{{.*}}
+# CHECK:   return

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

 using namespace mlir::tblgen;
 
 /// File header and includes.
 ///   {0} is the dialect namespace.
 constexpr const char *fileHeader = R"Py(
 # Autogenerated by mlir-tblgen; don't manually edit.
 
 from ._ods_common import _cext as _ods_cext
-from ._ods_common import extend_opview_class as _ods_extend_opview_class, segmented_accessor as _ods_segmented_accessor, equally_sized_accessor as _ods_equally_sized_accessor, get_default_loc_context as _ods_get_default_loc_context
+from ._ods_common import extend_opview_class as _ods_extend_opview_class, segmented_accessor as _ods_segmented_accessor, equally_sized_accessor as _ods_equally_sized_accessor, get_default_loc_context as _ods_get_default_loc_context, get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
 _ods_ir = _ods_cext.ir
 
 try:
   from . import _{0}_ops_ext as _ods_ext_module
 except ImportError:
   _ods_ext_module = None
 
 import builtins
def __init__(self, {0}):
     attributes = {{}
     {1}
     super().__init__(self.build_generic(
       attributes=attributes, results=results, operands=operands,
       successors=_ods_successors, loc=loc, ip=ip))
 )Py";
 
 /// Template for appending a single element to the operand/result list.
-///   {0} is either 'operand' or 'result';
-///   {1} is the field name.
-constexpr const char *singleElementAppendTemplate = "{0}s.append({1})";
+///   {0} is the field name.
+constexpr const char *singleOperandAppendTemplate =
+    "operands.append(_get_op_result_or_value({0}))";
+constexpr const char *singleResultAppendTemplate = "results.append({0})";
 
 /// Template for appending an optional element to the operand/result list.
-///   {0} is either 'operand' or 'result';
-///   {1} is the field name.
-constexpr const char *optionalAppendTemplate =
-    "if {1} is not None: {0}s.append({1})";
-
-/// Template for appending a a list of elements to the operand/result list.
-///   {0} is either 'operand' or 'result';
-///   {1} is the field name.
-constexpr const char *multiElementAppendTemplate = "{0}s.extend({1})";
+///   {0} is the field name.
+constexpr const char *optionalAppendOperandTemplate =
+    "if {0} is not None: operands.append(_get_op_result_or_value({0}))";
+constexpr const char *optionalAppendResultTemplate =
+    "if {0} is not None: results.append({0})";
+
+/// Template for appending a list of elements to the operand/result list.
+///   {0} is the field name.
+constexpr const char *multiOperandAppendTemplate =
+    "operands.extend(_get_op_results_or_values({0}))";
+constexpr const char *multiOperandAppendPackTemplate =
+    "operands.append(_get_op_results_or_values({0}))";
+constexpr const char *multiResultAppendTemplate = "results.extend({0})";
 
 /// Template for setting an attribute in the operation builder.
 ///   {0} is the attribute name;
 ///   {1} is the builder argument name.
 constexpr const char *initAttributeTemplate = R"Py(attributes["{0}"] = {1})Py";
 
 /// Template for setting an optional attribute in the operation builder.
 ///   {0} is the attribute name;
for (int i = 0, e = successorArgNames.size(); i < e; ++i) {
     const NamedSuccessor &successor = op.getSuccessor(i);
     builderLines.push_back(
         llvm::formatv(addSuccessorTemplate,
                       successor.isVariadic() ? "extend" : "append", argName));
   }
 }
 
 /// Populates `builderLines` with additional lines that are required in the
-/// builder. `kind` must be either "operand" or "result". `names` contains the
-/// names of init arguments that correspond to the elements.
-static void populateBuilderLines(
-    const Operator &op, const char *kind, llvm::ArrayRef<std::string> names,
-    llvm::SmallVectorImpl<std::string> &builderLines,
-    llvm::function_ref<int(const Operator &)> getNumElements,
-    llvm::function_ref<const NamedTypeConstraint &(const Operator &, int)>
-        getElement) {
-  bool sizedSegments = op.getTrait(attrSizedTraitForKind(kind)) != nullptr;
+/// builder to set up op operands.
+static void
+populateBuilderLinesOperand(const Operator &op,
+                            llvm::ArrayRef<std::string> names,
+                            llvm::SmallVectorImpl<std::string> &builderLines) {
+  bool sizedSegments = op.getTrait(attrSizedTraitForKind("operand")) != nullptr;
 
   // For each element, find or generate a name.
-  for (int i = 0, e = getNumElements(op); i < e; ++i) {
-    const NamedTypeConstraint &element = getElement(op, i);
+  for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
+    const NamedTypeConstraint &element = op.getOperand(i);
     std::string name = names[i];
 
     // Choose the formatting string based on the element kind.
     llvm::StringRef formatString;
     if (!element.isVariableLength()) {
-      formatString = singleElementAppendTemplate;
+      formatString = singleOperandAppendTemplate;
     } else if (element.isOptional()) {
-      formatString = optionalAppendTemplate;
+      formatString = optionalAppendOperandTemplate;
     } else {
       assert(element.isVariadic() && "unhandled element group type");
       // If emitting with sizedSegments, then we add the actual list typed
       // element using the singleElementAppendTemplate. Otherwise, we extend
       // the actual operands.
       if (sizedSegments) {
-        // Append the list as is.
-        formatString = singleElementAppendTemplate;
+        formatString = multiOperandAppendPackTemplate;
       } else {
-        // Append the list elements.
-        formatString = multiElementAppendTemplate;
+        formatString = multiOperandAppendTemplate;
       }
     }
 
-    // Add the lines.
-    builderLines.push_back(llvm::formatv(formatString.data(), kind, name));
+    builderLines.push_back(llvm::formatv(formatString.data(), name));
+  }
+}
+
+/// Populates `builderLines` with additional lines that are required in the
+/// builder to set up op results.
+static void
+populateBuilderLinesResult(const Operator &op,
+                           llvm::ArrayRef<std::string> names,
+                           llvm::SmallVectorImpl<std::string> &builderLines) {
+  bool sizedSegments = op.getTrait(attrSizedTraitForKind("result")) != nullptr;
+
+  // For each element, find or generate a name.
+  for (int i = 0, e = op.getNumResults(); i < e; ++i) {
+    const NamedTypeConstraint &element = op.getResult(i);
+    std::string name = names[i];
+
+    // Choose the formatting string based on the element kind.
+    llvm::StringRef formatString;
+    if (!element.isVariableLength()) {
+      formatString = singleResultAppendTemplate;
+    } else if (element.isOptional()) {
+      formatString = optionalAppendResultTemplate;
+    } else {
+      assert(element.isVariadic() && "unhandled element group type");
+      // If emitting with sizedSegments, then we add the actual list typed
+      // element using the singleElementAppendTemplate. Otherwise, we extend
+      // the actual operands.
+      if (sizedSegments) {
+        formatString = singleResultAppendTemplate;
+      } else {
+        formatString = multiResultAppendTemplate;
+      }
+    }
+
+    builderLines.push_back(llvm::formatv(formatString.data(), name));
   }
 }
 
 /// Emits a default builder constructing an operation from the list of its
 /// result types, followed by a list of its operands.
 static void emitDefaultOpBuilder(const Operator &op, raw_ostream &os) {
   // If we are asked to skip default builders, comply.
   if (op.skipDefaultBuilders())
     return;
 
   llvm::SmallVector<std::string> builderArgs;
   llvm::SmallVector<std::string> builderLines;
   llvm::SmallVector<std::string> operandArgNames;
   llvm::SmallVector<std::string> successorArgNames;
   builderArgs.reserve(op.getNumOperands() + op.getNumResults() +
                       op.getNumNativeAttributes() + op.getNumSuccessors());
   populateBuilderArgs(op, builderArgs, operandArgNames, successorArgNames);
 
-  populateBuilderLines(
-      op, "result",
-      llvm::makeArrayRef(builderArgs).take_front(op.getNumResults()),
-      builderLines, getNumResults, getResult);
-  populateBuilderLines(op, "operand", operandArgNames, builderLines,
-                       getNumOperands, getOperand);
+  populateBuilderLinesResult(
+      op, llvm::makeArrayRef(builderArgs).take_front(op.getNumResults()),
+      builderLines);
+  populateBuilderLinesOperand(op, operandArgNames, builderLines);
   populateBuilderLinesAttr(
       op, llvm::makeArrayRef(builderArgs).drop_front(op.getNumResults()),
       builderLines);
   populateBuilderLinesSuccessors(op, successorArgNames, builderLines);
 
   builderArgs.push_back("*");
   builderArgs.push_back("loc=None");
   builderArgs.push_back("ip=None");