[Mlir-commits] [mlir] [mlir][python] Add `walk` method to PyOperationBase (PR #87962)

Hideto Ueno llvmlistbot at llvm.org
Thu Apr 11 06:38:19 PDT 2024


================
@@ -1249,6 +1250,21 @@ void PyOperationBase::writeBytecode(const py::object &fileObject,
                               .str());
 }
 
+void PyOperationBase::walk(
+    std::function<MlirWalkResult(MlirOperation)> callback,
+    MlirWalkOrder walkOrder) {
+  PyOperation &operation = getOperation();
+  operation.checkValid();
+  MlirOperationWalkCallback walkCallback = [](MlirOperation op,
+                                              void *userData) {
+    auto *fn =
+        static_cast<std::function<MlirWalkResult(MlirOperation)> *>(userData);
+    return (*fn)(op);
----------------
uenoku wrote:

I'm not sure it works since the function is type casted to  `std::function<MlirWalkResult(MlirOperation)> callback`. 
To allow void functions I think we could change the callback type to `std::variant<std::function<MlirWalkResult(MlirOperation)>, std::function<void(MlirOperation)>` or maybe `py::function` (and dynamically type casts the returned type). 

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


More information about the Mlir-commits mailing list