[Mlir-commits] [mlir] [MLIR][Python] Add type filter to walk() binding and add get_ops_of_type() utility (PR #186131)

[Jakub Kuderski]

================
@@ -3858,16 +3858,41 @@ void populateIRCore(nb::module_ &m) {
 
             Note:
               After erasing, any Python references to the operation become invalid.)")
-      .def("walk", &PyOperationBase::walk, "callback"_a,
-           "walk_order"_a = PyWalkOrder::PostOrder,
-           // clang-format off
-          nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder) -> None"),
-           // clang-format on
-           R"(
+      .def(
+          "walk",
+          [](PyOperationBase &self,
+             std::function<PyWalkResult(MlirOperation)> callback,
+             std::optional<nb::object> opClass, PyWalkOrder walkOrder) {
+            if (opClass) {
+              self.walk(
+                  [&](MlirOperation mlirOp) -> PyWalkResult {
+                    nb::object opview =
+                        PyOperation::forOperation(
+                            self.getOperation().getContext(), mlirOp)
+                            ->createOpView();
+                    if (nb::isinstance(opview, *opClass))
+                      return callback(mlirOp);
+                    return PyWalkResult::Advance;
+                  },
+                  walkOrder);
+            } else {
+              self.walk(callback, walkOrder);
+            }
+          },
+          "callback"_a, "op_class"_a = nb::none(),
+          "walk_order"_a = PyWalkOrder::PostOrder,
+          // clang-format off
+           nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], op_class: type[OpView] | None = None, walk_order: WalkOrder = WalkOrder.POST_ORDER) -> None"),
+          // clang-format on
+          R"(
              Walks the operation tree with a callback function.
 
+             If op_class is provided, the callback is only invoked on operations
+             of that type; all other operations are skipped silently.
+
              Args:
                callback: A callable that takes an Operation and returns a WalkResult.
+               op_class: If provided, only operations of this type are passed to the callback.
----------------
kuhar wrote:

This seems like a major breaking change all the existing code that uses python bindings... While we don't guarantee much API stability, I think it would be preferred not to break existing code. Can we work around this somehow to either take op_class as an optional argugment or revert back to having a new function?

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