Make Python MLIR Operation not iterable

[mehdi_amini]

Summary

The current behavior is conveniently allowing to iterate on the regions of an operation
implicitly by exposing an operation as Iterable. However this is also error prone and
code that may intend to iterate on the results or the operands could end up "working"
apparently instead of throwing a runtime error.
The lack of static type checking in Python contributes to the ambiguity here, it seems
safer to not do this and require and explicit qualification to iterate (op.results, op.regions, ...).

[ftynse]

I was tricked by this, too. Please update the docs/Python.md to remove the mention of operations being iterable.

[stellaraccident]

mlir/python/mlir/dialects/_builtin_ops_ext.py

201

Oof, this is a nasty bug you uncovered. In general, I get worried about these ad-hoc conversions between op<->results and am undecided on whether we should try to silently convert more (and whether there is a more principled utility function we should use for that, which we can test more invariants of) or raise an error (i.e. "Passed an Operation to an API which expected value(s): Please use either the .result or .results property if intending to access a result.").

Thoughts?

mlir/test/python/dialects/math.py

19

I would argue that this test is wrong and should be fixed. Regardless of whether there is some automatic conversion, this should be:

return mlir_math.SqrtOp(F32Type.get(), arg).result

[ftynse]

mlir/python/mlir/dialects/_builtin_ops_ext.py

201

Utility function like this one https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/dialects/_ods_common.py#L126? We are already using those in the code generated from ODS.

mlir/test/python/dialects/math.py

19

It isn't surprising to me that there is an automatic conversion to value here. Single-result operations are implicitly convertible to Value in C++, and in Python when used as arguments. (We also convert multi-result operations to lists of arguments in Python). Given that we can write std.ReturnOp(mlir_math.SqrtOp(F32Type.get(), arg)) and it does what I expect it to, having it for the DSL-style embedding makes even more sense.

[mehdi_amini]

mlir/test/python/dialects/math.py

19

It seems that we're being pulled in between two use-cases here:

• IR building API
• DSL-style.

The former would prefer to be explicit (we don't pretend to be a language) while the later would want the most "natural language" style.

This revision goes in the first direction, but if we agree that this is the kind of programmability that the bindings aim to offer, we should avoid the DSL-style in every aspect of the API, otherwise this will lead to user confusion (reminds me of TF eager vs TF function).

[stellaraccident]

mlir/python/mlir/dialects/_builtin_ops_ext.py

201

Yes, like that.

mlir/test/python/dialects/math.py

19

I didn't intend to argue the purity of DSL style vs other (I prefer some ergonomics and the op vs value brittleness had cost me a lot of time as a user -- and I built a lot of this stuff).

I would like to use a helper like what Alex points out: I'd prefer to not open code such checks/conversions.

19

Fair point. I guess my complaint is that this is the only thing that broke: it is an unrelated test. The implicit conversion should be tested in the builtin dialect tests along with the rest of the tests for this decorator.

[ftynse]

mlir/test/python/dialects/math.py

19

Re testing: when https://reviews.llvm.org/D111656 lands, it will enable the python_test dialect as a side effect. We can then use that dialect to test infrastructural features such as automatic conversions using that, the builtin dialect doesn't have enough ops to test all the things we need.

[mehdi_amini]

Rebase and error out on multi-result operations

[mehdi_amini]

mlir/python/mlir/dialects/_builtin_ops_ext.py

201

I don't think I can directly use this function from here, can I?

Anyway can clarify the desired behavior (and then I'll add tests), right now what I added is asserting if we return a multi-result operation here, but we still accept the single result case:

1. Is this what we want? This is consistent with C++ "somehow": in C++ we would do it on the OpView but not the Operation, I could also align on this but I'm not sure this all makes sense in Python.
2. We could also accept to unpack a multi-result operation assuming we return a single operation.
3. I added an assert because this is done like that a few lines before in this function, but shouldn't we raise an exception instead? (which kind?).

[stellaraccident]

mlir/python/mlir/dialects/_builtin_ops_ext.py

201

Yes, you can from ._ods_common import foobar from here.

I've been convinced and am generally in favor of some coercion so long as it isn't ad-hoc. Starting convservative and broadening as we see more seems safe, but I don't have a super strong opinion.

[mehdi_amini]

Rebase and update to implement coercion for single ops

mlir/lib/Bindings/Python/IRCore.cpp

py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
           "location",
           [](PyOperationBase &self) {
             PyOperation &operation = self.getOperation();
             return PyLocation(operation.getContext(),
                               mlirOperationGetLocation(operation.get()));
           },
           "Returns the source location the operation was defined or derived "
           "from.")
-      .def("__iter__",
-           [](PyOperationBase &self) {
-             return PyRegionIterator(self.getOperation().getRef());
-           })
       .def(
           "__str__",
           [](PyOperationBase &self) {
             return self.getAsm(/*binary=*/false,
                                /*largeElementsLimit=*/llvm::None,
                                /*enableDebugInfo=*/false,
                                /*prettyDebugInfo=*/false,
                                /*printGenericOpForm=*/false,

mlir/python/mlir/dialects/_builtin_ops_ext.py

def decorator(f):
               "Capturing a python function with explicit `results=` "
               "requires that the wrapped function returns None.")
         else:
           # Coerce return values, add ReturnOp and rewrite func type.
           if return_values is None:
             return_values = []
           elif isinstance(return_values, Value):
             return_values = [return_values]
+          elif isinstance(return_values, OpView):
+            assert len(return_values.operation.results) == 1, (
+              "Returning a multi-results operation should explicitly "
+              "unpack with `.results`.")
+            return_values = return_values.operation.result
+          elif isinstance(return_values, Operation):
+            assert len(return_values.results) == 1, (
+              "Returning a multi-results operation should explicitly "
+              "unpack with `.results`.")
+            return_values = return_values.result
           else:
             return_values = list(return_values)
           std.ReturnOp(return_values)
           # Recompute the function type.
           return_types = [v.type for v in return_values]
           function_type = FunctionType.get(inputs=inputs, results=return_types)
           func_op.attributes["type"] = TypeAttr.get(function_type)
 

mlir/test/python/dialects/math.py

 # CHECK-LABEL: TEST: testMathOps
 @run
 def testMathOps():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(F32Type.get())
       def emit_sqrt(arg):
         return mlir_math.SqrtOp(F32Type.get(), arg)
 
     # CHECK-LABEL: func @emit_sqrt(
-    # CHECK-SAME:                  %[[ARG:.*]]: f32) {
+    # CHECK-SAME:                  %[[ARG:.*]]: f32) -> f32 {
     # CHECK:         math.sqrt %[[ARG]] : f32
     # CHECK:         return
     # CHECK:       }
     print(module)

mlir/test/python/ir/operation.py

def testTraverseOpRegionBlockIterators():
   # CHECK: MODULE REGIONS=1 BLOCKS=1
   print(f"MODULE REGIONS={len(regions)} BLOCKS={len(blocks)}")
 
   # Should verify.
   # CHECK: .verify = True
   print(f".verify = {module.operation.verify()}")
 
   # Get the regions and blocks from the default collections.
-  default_regions = list(op)
+  default_regions = list(op.regions)
   default_blocks = list(default_regions[0])
   # They should compare equal regardless of how obtained.
   assert default_regions == regions
   assert default_blocks == blocks
 
   # Should be able to get the operations from either the named collection
   # or the block.
   operations = list(blocks[0].operations)
   default_operations = list(blocks[0])
   assert default_operations == operations
 
   def walk_operations(indent, op):
-    for i, region in enumerate(op):
+    for i, region in enumerate(op.regions):
       print(f"{indent}REGION {i}:")
       for j, block in enumerate(region):
         print(f"{indent}  BLOCK {j}:")
         for k, child_op in enumerate(block):
           print(f"{indent}    OP {k}: {child_op}")
           walk_operations(indent + "      ", child_op)
 
   # CHECK: REGION 0: