handle ir.Operation._CAPICreate on the C side

Xilinx · Nov 3, 2023 · d91855d · d91855d
1 parent b5cee69
commit d91855d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git 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
@@ -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 {