[MLIR][Python] Support eliding large resource strings in PassManager#149187
Merged
makslevental merged 1 commit intollvm:mainfrom Jul 17, 2025
Merged
[MLIR][Python] Support eliding large resource strings in PassManager#149187makslevental merged 1 commit intollvm:mainfrom
makslevental merged 1 commit intollvm:mainfrom
Conversation
|
✅ With the latest revision this PR passed the Python code formatter. |
akshaykhadse
force-pushed
the
fix-resource-string-printing
branch
2 times, most recently
from
a904185 to
658a9ef
Compare
akshaykhadse
requested review from
ftynse,
makslevental and
stellaraccident
as code owners
Member
|
@llvm/pr-subscribers-mlir Author: Akshay Khadse (akshaykhadse) Changes
Full diff: https://github.com/llvm/llvm-project/pull/149187.diff 7 Files Affected:
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
+@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
+@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():
|
makslevental
approved these changes
Jul 17, 2025
Contributor
There was a problem hiding this comment.
Looks good to me but this is a "breaking" change - can you take a look and see if anyone is using the previous API to handle resource strings (you can try https://sourcegraph.com/search). If you don't find anything it's probably safe to change but you'll need to make a post on the discourse about it (like this one https://discourse.llvm.org/t/psa-changes-in-mlir-python-type-casters/87383).
akshaykhadse
force-pushed
the
fix-resource-string-printing
branch
from
658a9ef to
60479f3
Compare
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
- Introduces a `large_resource_limit` parameter across Python bindings, enabling the eliding of resource strings exceeding a specified character limit during IR printing. - To maintain backward compatibilty, when using `operation.print()` API, if `large_resource_limit` is None and the `large_elements_limit` is set, the later will be used to elide the resource string as well. This change was introduced by llvm#125738. - For printing using pass manager, the `large_resource_limit` and `large_elements_limit` are completely independent of each other.
akshaykhadse
force-pushed
the
fix-resource-string-printing
branch
from
60479f3 to
6323eeb
Compare
akshaykhadse
changed the title
akshaykhadse
changed the title
Contributor
|
Cool since this is backwards compatible we can merge this - let me know when you're ready and I'll merge it. |
This was referenced Jul 23, 2025
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
large_resource_limitparameter across Python bindings, enabling the eliding of resource strings exceeding a specified character limit during IR printing.operation.print()API, iflarge_resource_limitis None and thelarge_elements_limitis set, the later will be used to elide the resource string as well. This change was introduced by [mlir] Python: Extend print large elements limit to resources #125738.large_resource_limitandlarge_elements_limitare completely independent of each other.