[Mlir-commits] [mlir] [mlir][python] Make the Context/Operation capsule creation methods work as documented. (PR #76010)

[Stella Laurenzo]

stellaraccident wrote:

> Yes, without `forOpeation`, users _are_ broken. This basically makes it impossible for "normal" users to return a non-owning reference to an operation. I don't think that's what we want. Here's the minimal example that breaks upstream. Note that anybody downstream cannot return something else than `MlirOperation` AFAIK. Upstream doesn't have a problem because it can, and does, return a `PyOpreationRef`.
> 
> ```
> diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
> index 39757dfad5be..2ce640674245 100644
> --- a/mlir/lib/Bindings/Python/IRCore.cpp
> +++ b/mlir/lib/Bindings/Python/IRCore.cpp
> @@ -3071,6 +3071,11 @@ void mlir::python::populateIRCore(py::module &m) {
>                    py::arg("successors") = py::none(), py::arg("regions") = 0,
>                    py::arg("loc") = py::none(), py::arg("ip") = py::none(),
>                    py::arg("infer_type") = false, kOperationCreateDocstring)
> +      .def("_get_first_in_block", [](PyOperation &self) -> MlirOperation {
> +        MlirBlock block = mlirOperationGetBlock(self.get());
> +        MlirOperation first = mlirBlockGetFirstOperation(block);
> +        return first;
> +      })
>        .def_static(
>            "parse",
>            [](const std::string &sourceStr, const std::string &sourceName,
> diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
> index f59b1a26ba48..6b12b8da5c24 100644
> --- a/mlir/test/python/ir/operation.py
> +++ b/mlir/test/python/ir/operation.py
> @@ -24,6 +24,25 @@ def expect_index_error(callback):
>      except IndexError:
>          pass
>  
> + at run
> +def testCustomBind():
> +    ctx = Context()
> +    ctx.allow_unregistered_dialects = True
> +    module = Module.parse(
> +        r"""
> +    func.func @f1(%arg0: i32) -> i32 {
> +      %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
> +      return %1 : i32
> +    }
> +  """,
> +        ctx,
> +    )
> +    add = module.body.operations[0].regions[0].blocks[0].operations[0]
> +    op = add.operation
> +    # This will get a reference to itself.
> +    f1 = op._get_first_in_block()
> +
> +
>  
>  # Verify iterator based traversal of the op/region/block hierarchy.
>  # CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators
> ```

Darn - I was trying to not break that case... Which was the reason for the disjoint implementations between context and op import. Thanks for the test case. That will help in rolling this forward. Probably next week at this point.

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