diff --git a/python/PythonPass.cpp b/python/PythonPass.cpp index 029dc393dd..9e6daa7b71 100644 --- a/python/PythonPass.cpp +++ b/python/PythonPass.cpp @@ -25,12 +25,18 @@ struct PythonPassDemo PythonPassDemo(PyObject *func) : func(func) {} StringRef getArgument() const final { return "python-pass-demo"; } + void runOnOperation() override { this->getOperation()->walk([this](Operation *op) { - PyObject *arglist = - Py_BuildValue("(O)", mlirPythonOperationToCapsule(wrap(op))); - auto result = PyObject_CallObject(func, arglist); - Py_DECREF(arglist); + PyObject *mlirModule = + PyImport_ImportModule(MAKE_MLIR_PYTHON_QUALNAME("ir")); + PyObject *cAPIFactory = PyObject_GetAttrString( + PyObject_GetAttrString(mlirModule, "Operation"), + MLIR_PYTHON_CAPI_FACTORY_ATTR); + PyObject *opApiObject = PyObject_CallFunction( + cAPIFactory, "(O)", mlirPythonOperationToCapsule(wrap(op))); + auto result = PyObject_CallFunction(func, "(O)", opApiObject); + Py_DECREF(opApiObject); }); } diff --git a/test/python/aie_ops.py b/test/python/aie_ops.py index f44b0b3534..817a42e355 100644 --- a/test/python/aie_ops.py +++ b/test/python/aie_ops.py @@ -126,8 +126,8 @@ def testPythonPassDemo(): # CHECK-LABEL: testPythonPassDemo print("\nTEST: testPythonPassDemo") - def print_ops(op_caps): - print(Operation._CAPICreate(op_caps).name) + def print_ops(op): + print(op.name) module = """ module {