[Mlir-commits] [mlir] Support eliding large resource strings in Python PassManager (PR #149187)
Akshay Khadse
llvmlistbot at llvm.org
Wed Jul 16 14:47:40 PDT 2025
https://github.com/akshaykhadse updated https://github.com/llvm/llvm-project/pull/149187
>From 658a9efc16c370942b893c064d07b9dc4868290c Mon Sep 17 00:00:00 2001
From: Akshay Khadse <akshayskhadse at gmail.com>
Date: Wed, 16 Jul 2025 14:23:42 -0400
Subject: [PATCH] Support eliding large resource strings in Python PassManager
- Introduces a `large_resource_limit` parameter across Python bindings, enabling the elision of resource strings exceeding a specified character limit during IR printing.
- Reverts the behavior where `large_elements_limit` would also control the resource string printing from operation's `print` method when using the Python binding. This is because the eliding of resource string takes the number of characters as input whereas the elision of elements takes number of elements as input. Both are separate and should be respected in Python bindings as well.
---
mlir/lib/Bindings/Python/IRCore.cpp | 20 +++++--
mlir/lib/Bindings/Python/IRModule.h | 14 ++---
mlir/lib/Bindings/Python/Pass.cpp | 7 ++-
mlir/python/mlir/_mlir_libs/_mlir/ir.pyi | 5 ++
.../mlir/_mlir_libs/_mlir/passmanager.pyi | 1 +
mlir/test/python/ir/operation.py | 9 +++
mlir/test/python/pass_manager.py | 57 +++++++++++++++++++
7 files changed, 99 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index d961482885300..c90b866e67ea0 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -97,6 +97,8 @@ static const char kOperationPrintDocstring[] =
binary: Whether to write bytes (True) or str (False). Defaults to False.
large_elements_limit: Whether to elide elements attributes above this
number of elements. Defaults to None (no limit).
+ large_resource_limit: Whether to elide resource attributes above this
+ number of characters. Defaults to None (no limit).
enable_debug_info: Whether to print debug/location information. Defaults
to False.
pretty_debug_info: Whether to format debug information for easier reading
@@ -1303,6 +1305,7 @@ void PyOperation::checkValid() const {
}
void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
+ std::optional<int64_t> largeResourceLimit,
bool enableDebugInfo, bool prettyDebugInfo,
bool printGenericOpForm, bool useLocalScope,
bool useNameLocAsPrefix, bool assumeVerified,
@@ -1314,10 +1317,10 @@ void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
fileObject = nb::module_::import_("sys").attr("stdout");
MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
- if (largeElementsLimit) {
+ if (largeElementsLimit)
mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);
- mlirOpPrintingFlagsElideLargeResourceString(flags, *largeElementsLimit);
- }
+ if (largeResourceLimit)
+ mlirOpPrintingFlagsElideLargeResourceString(flags, *largeResourceLimit);
if (enableDebugInfo)
mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
/*prettyForm=*/prettyDebugInfo);
@@ -1405,6 +1408,7 @@ void PyOperationBase::walk(
nb::object PyOperationBase::getAsm(bool binary,
std::optional<int64_t> largeElementsLimit,
+ std::optional<int64_t> largeResourceLimit,
bool enableDebugInfo, bool prettyDebugInfo,
bool printGenericOpForm, bool useLocalScope,
bool useNameLocAsPrefix, bool assumeVerified,
@@ -1416,6 +1420,7 @@ nb::object PyOperationBase::getAsm(bool binary,
fileObject = nb::module_::import_("io").attr("StringIO")();
}
print(/*largeElementsLimit=*/largeElementsLimit,
+ /*largeResourceLimit=*/largeResourceLimit,
/*enableDebugInfo=*/enableDebugInfo,
/*prettyDebugInfo=*/prettyDebugInfo,
/*printGenericOpForm=*/printGenericOpForm,
@@ -3348,6 +3353,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
[](PyOperationBase &self) {
return self.getAsm(/*binary=*/false,
/*largeElementsLimit=*/std::nullopt,
+ /*largeResourceLimit=*/std::nullopt,
/*enableDebugInfo=*/false,
/*prettyDebugInfo=*/false,
/*printGenericOpForm=*/false,
@@ -3363,11 +3369,12 @@ void mlir::python::populateIRCore(nb::module_ &m) {
nb::arg("state"), nb::arg("file").none() = nb::none(),
nb::arg("binary") = false, kOperationPrintStateDocstring)
.def("print",
- nb::overload_cast<std::optional<int64_t>, bool, bool, bool, bool,
- bool, bool, nb::object, bool, bool>(
- &PyOperationBase::print),
+ nb::overload_cast<std::optional<int64_t>, std::optional<int64_t>,
+ bool, bool, bool, bool, bool, bool, nb::object,
+ bool, bool>(&PyOperationBase::print),
// Careful: Lots of arguments must match up with print method.
nb::arg("large_elements_limit").none() = nb::none(),
+ nb::arg("large_resource_limit").none() = nb::none(),
nb::arg("enable_debug_info") = false,
nb::arg("pretty_debug_info") = false,
nb::arg("print_generic_op_form") = false,
@@ -3383,6 +3390,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
// Careful: Lots of arguments must match up with get_asm method.
nb::arg("binary") = false,
nb::arg("large_elements_limit").none() = nb::none(),
+ nb::arg("large_resource_limit").none() = nb::none(),
nb::arg("enable_debug_info") = false,
nb::arg("pretty_debug_info") = false,
nb::arg("print_generic_op_form") = false,
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 9befcce725bb7..0fdd2d1a7eff6 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -599,18 +599,18 @@ class PyOperationBase {
public:
virtual ~PyOperationBase() = default;
/// Implements the bound 'print' method and helps with others.
- void print(std::optional<int64_t> largeElementsLimit, bool enableDebugInfo,
+ void print(std::optional<int64_t> largeElementsLimit,
+ std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope,
bool useNameLocAsPrefix, bool assumeVerified,
nanobind::object fileObject, bool binary, bool skipRegions);
void print(PyAsmState &state, nanobind::object fileObject, bool binary);
- nanobind::object getAsm(bool binary,
- std::optional<int64_t> largeElementsLimit,
- bool enableDebugInfo, bool prettyDebugInfo,
- bool printGenericOpForm, bool useLocalScope,
- bool useNameLocAsPrefix, bool assumeVerified,
- bool skipRegions);
+ nanobind::object
+ getAsm(bool binary, std::optional<int64_t> largeElementsLimit,
+ std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
+ bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope,
+ bool useNameLocAsPrefix, bool assumeVerified, bool skipRegions);
// Implement the bound 'writeBytecode' method.
void writeBytecode(const nanobind::object &fileObject,
diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index 8d84864b9db4d..830e106499fc1 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -78,12 +78,16 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
[](PyPassManager &passManager, bool printBeforeAll,
bool printAfterAll, bool printModuleScope, bool printAfterChange,
bool printAfterFailure, std::optional<int64_t> largeElementsLimit,
- bool enableDebugInfo, bool printGenericOpForm,
+ std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
+ bool printGenericOpForm,
std::optional<std::string> optionalTreePrintingPath) {
MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
if (largeElementsLimit)
mlirOpPrintingFlagsElideLargeElementsAttrs(flags,
*largeElementsLimit);
+ if (largeResourceLimit)
+ mlirOpPrintingFlagsElideLargeResourceString(flags,
+ *largeResourceLimit);
if (enableDebugInfo)
mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
/*prettyForm=*/false);
@@ -103,6 +107,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
"print_module_scope"_a = false, "print_after_change"_a = false,
"print_after_failure"_a = false,
"large_elements_limit"_a.none() = nb::none(),
+ "large_resource_limit"_a.none() = nb::none(),
"enable_debug_info"_a = false, "print_generic_op_form"_a = false,
"tree_printing_dir_path"_a.none() = nb::none(),
"Enable IR printing, default as mlir-print-ir-after-all.")
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index ed476da28d6be..d0212c7a1e126 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -200,6 +200,7 @@ class _OperationBase:
def get_asm(
binary: Literal[True],
large_elements_limit: int | None = None,
+ large_resource_limit: int | None = None,
enable_debug_info: bool = False,
pretty_debug_info: bool = False,
print_generic_op_form: bool = False,
@@ -212,6 +213,7 @@ class _OperationBase:
self,
binary: bool = False,
large_elements_limit: int | None = None,
+ large_resource_limit: int | None = None,
enable_debug_info: bool = False,
pretty_debug_info: bool = False,
print_generic_op_form: bool = False,
@@ -253,6 +255,7 @@ class _OperationBase:
def print(
self,
large_elements_limit: int | None = None,
+ large_resource_limit: int | None = None,
enable_debug_info: bool = False,
pretty_debug_info: bool = False,
print_generic_op_form: bool = False,
@@ -270,6 +273,8 @@ class _OperationBase:
binary: Whether to write bytes (True) or str (False). Defaults to False.
large_elements_limit: Whether to elide elements attributes above this
number of elements. Defaults to None (no limit).
+ large_resource_limit: Whether to elide resource strings above this
+ number of characters. Defaults to None (no limit).
enable_debug_info: Whether to print debug/location information. Defaults
to False.
pretty_debug_info: Whether to format debug information for easier reading
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
index 0d2eaffe16d3e..1010daddae2aa 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -23,6 +23,7 @@ class PassManager:
print_after_change: bool = False,
print_after_failure: bool = False,
large_elements_limit: int | None = None,
+ large_resource_limit: int | None = None,
enable_debug_info: bool = False,
print_generic_op_form: bool = False,
tree_printing_dir_path: str | None = None,
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index b08fe98397fbc..ede1571f940f6 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -686,6 +686,15 @@ def testOperationPrint():
skip_regions=True,
)
+ # Test print with large_resource_limit.
+ # CHECK: func.func @f1(%arg0: i32) -> i32
+ # CHECK-NOT: resource1: "0x08
+ module.operation.print(large_resource_limit=2)
+
+ # Test large_elements_limit has no effect on resource string
+ # CHECK: func.func @f1(%arg0: i32) -> i32
+ # CHECK: resource1: "0x08
+ module.operation.print(large_elements_limit=2)
# CHECK-LABEL: TEST: testKnownOpView
@run
diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index 85d2eb304882e..4c992a3ab2e45 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -363,6 +363,63 @@ def testPrintIrLargeLimitElements():
pm.run(module)
+# CHECK-LABEL: TEST: testPrintIrLargeResourceLimit
+ at run
+def testPrintIrLargeResourceLimit():
+ with Context() as ctx:
+ module = ModuleOp.parse(
+ """
+ module {
+ func.func @main() -> tensor<3xi64> {
+ %0 = arith.constant dense_resource<blob1> : tensor<3xi64>
+ return %0 : tensor<3xi64>
+ }
+ }
+ {-#
+ dialect_resources: {
+ builtin: {
+ blob1: "0x010000000000000002000000000000000300000000000000"
+ }
+ }
+ #-}
+ """
+ )
+ pm = PassManager.parse("builtin.module(canonicalize)")
+ ctx.enable_multithreading(False)
+ pm.enable_ir_printing(large_resource_limit=4)
+ # CHECK-NOT: blob1: "0x01
+ pm.run(module)
+
+
+# CHECK-LABEL: TEST: testPrintIrLargeResourceLimitVsElementsLimit
+ at run
+def testPrintIrLargeResourceLimitVsElementsLimit():
+ """Test that large_elements_limit does not affect the printing of resources."""
+ with Context() as ctx:
+ module = ModuleOp.parse(
+ """
+ module {
+ func.func @main() -> tensor<3xi64> {
+ %0 = arith.constant dense_resource<blob1> : tensor<3xi64>
+ return %0 : tensor<3xi64>
+ }
+ }
+ {-#
+ dialect_resources: {
+ builtin: {
+ blob1: "0x010000000000000002000000000000000300000000000000"
+ }
+ }
+ #-}
+ """
+ )
+ pm = PassManager.parse("builtin.module(canonicalize)")
+ ctx.enable_multithreading(False)
+ pm.enable_ir_printing(large_elements_limit=1)
+ # CHECK: blob1: "0x01
+ pm.run(module)
+
+
# CHECK-LABEL: TEST: testPrintIrTree
@run
def testPrintIrTree():
More information about the Mlir-commits
mailing list