Skip to content

Commit d97e8cd

Browse files
committed
[mlir][python] Include anchor op in PassManager constructor
This adds an extra argument for specifying the pass manager's anchor op, with a default of `any`. Previously the anchor was always defaulted to `builtin.module`. Reviewed By: mehdi_amini Differential Revision: https://reviews.llvm.org/D136406
1 parent 66645a0 commit d97e8cd

File tree

2 files changed

+17
-3
lines changed

2 files changed

+17
-3
lines changed

mlir/lib/Bindings/Python/Pass.cpp

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,11 +56,14 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
5656
// Mapping of the top-level PassManager
5757
//----------------------------------------------------------------------------
5858
py::class_<PyPassManager>(m, "PassManager", py::module_local())
59-
.def(py::init<>([](DefaultingPyMlirContext context) {
60-
MlirPassManager passManager =
61-
mlirPassManagerCreate(context->get());
59+
.def(py::init<>([](const std::string &anchorOp,
60+
DefaultingPyMlirContext context) {
61+
MlirPassManager passManager = mlirPassManagerCreateOnOperation(
62+
context->get(),
63+
mlirStringRefCreate(anchorOp.data(), anchorOp.size()));
6264
return new PyPassManager(passManager);
6365
}),
66+
py::arg("anchor_op") = py::str("any"),
6467
py::arg("context") = py::none(),
6568
"Create a new PassManager for the current (or provided) Context.")
6669
.def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,

mlir/test/python/pass_manager.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,17 @@ def testCapsule():
2828
assert pm1 is not None # And does not crash.
2929
run(testCapsule)
3030

31+
# CHECK-LABEL: TEST: testConstruct
32+
@run
33+
def testConstruct():
34+
with Context():
35+
# CHECK: pm1: 'any()'
36+
# CHECK: pm2: 'builtin.module()'
37+
pm1 = PassManager()
38+
pm2 = PassManager("builtin.module")
39+
log(f"pm1: '{pm1}'")
40+
log(f"pm2: '{pm2}'")
41+
3142

3243
# Verify successful round-trip.
3344
# CHECK-LABEL: TEST: testParseSuccess

0 commit comments

Comments
 (0)