diff --git a/mlir/lib/Bindings/Python/Rewrite.cpp b/mlir/lib/Bindings/Python/Rewrite.cpp index 9830c277ac147..faab66d5ce4e5 100644 --- a/mlir/lib/Bindings/Python/Rewrite.cpp +++ b/mlir/lib/Bindings/Python/Rewrite.cpp @@ -509,26 +509,42 @@ void populateRewriteSubmodule(nb::module_ &m) { &PyFrozenRewritePatternSet::createFromCapsule); m.def( "apply_patterns_and_fold_greedily", - [](PyModule &module, PyFrozenRewritePatternSet &set) { + [](PyModule &module, PyFrozenRewritePatternSet &set, nb::object config) { + if (config.is_none()) { + config = nb::cast(PyGreedyRewriteDriverConfig()); + } + auto status = mlirApplyPatternsAndFoldGreedily( - module.get(), set.get(), mlirGreedyRewriteDriverConfigCreate()); + module.get(), set.get(), + nb::cast(config).get()); if (mlirLogicalResultIsFailure(status)) throw std::runtime_error("pattern application failed to converge"); }, - "module"_a, "set"_a, + "module"_a, "set"_a, "config"_a = nb::none(), + // clang-format off + nb::sig("def apply_patterns_and_fold_greedily(module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Module") ", set: FrozenRewritePatternSet, config: GreedyRewriteDriverConfig | None = None) -> None"), + // clang-format on "Applys the given patterns to the given module greedily while folding " "results.") .def( "apply_patterns_and_fold_greedily", - [](PyOperationBase &op, PyFrozenRewritePatternSet &set) { + [](PyOperationBase &op, PyFrozenRewritePatternSet &set, + nb::object config) { + if (config.is_none()) { + config = nb::cast(PyGreedyRewriteDriverConfig()); + } + auto status = mlirApplyPatternsAndFoldGreedilyWithOp( op.getOperation(), set.get(), - mlirGreedyRewriteDriverConfigCreate()); + nb::cast(config).get()); if (mlirLogicalResultIsFailure(status)) throw std::runtime_error( "pattern application failed to converge"); }, - "op"_a, "set"_a, + "op"_a, "set"_a, "config"_a = nb::none(), + // clang-format off + nb::sig("def apply_patterns_and_fold_greedily(op: " MAKE_MLIR_PYTHON_QUALNAME("ir._OperationBase") ", set: FrozenRewritePatternSet, config: GreedyRewriteDriverConfig | None = None) -> None"), + // clang-format on "Applys the given patterns to the given op greedily while folding " "results.") .def( diff --git a/mlir/test/python/rewrite.py b/mlir/test/python/rewrite.py index a2fbbde38b8c0..43e9b761a0ea2 100644 --- a/mlir/test/python/rewrite.py +++ b/mlir/test/python/rewrite.py @@ -182,3 +182,47 @@ def testGreedySimplifyRegionLevelEnum(): config.region_simplification_level = GreedySimplifyRegionLevel.AGGRESSIVE level = config.region_simplification_level print(f"region_level AGGRESSIVE: {level}") + + +# CHECK-LABEL: TEST: testRewriteWithGreedyRewriteDriverConfig +@run +def testRewriteWithGreedyRewriteDriverConfig(): + def constant_1_to_2(op, rewriter): + c = op.value.value + if c != 1: + return True # failed to match + with rewriter.ip: + new_op = arith.constant(op.type, 2, loc=op.location) + rewriter.replace_op(op, [new_op]) + + with Context(): + patterns = RewritePatternSet() + patterns.add(arith.ConstantOp, constant_1_to_2) + frozen = patterns.freeze() + + module = ModuleOp.parse( + r""" + module { + func.func @const() -> (i64, i64) { + %0 = arith.constant 1 : i64 + %1 = arith.constant 1 : i64 + return %0, %1 : i64, i64 + } + } + """ + ) + + config = GreedyRewriteDriverConfig() + config.enable_constant_cse = False + apply_patterns_and_fold_greedily(module, frozen, config) + # CHECK: %c2_i64 = arith.constant 2 : i64 + # CHECK: %c2_i64_0 = arith.constant 2 : i64 + # CHECK: return %c2_i64, %c2_i64_0 : i64, i64 + print(module) + + config = GreedyRewriteDriverConfig() + config.enable_constant_cse = True + apply_patterns_and_fold_greedily(module, frozen, config) + # CHECK: %c2_i64 = arith.constant 2 : i64 + # CHECK: return %c2_i64, %c2_i64 : i64 + print(module)