Skip to content
This repository was archived by the owner on Oct 11, 2025. It is now read-only.

Commit 22355e9

Browse files
akshaykhadseakshaykhadse
authored
[MLIR][Python] Support eliding large resource strings in PassManager (#149187)
- 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/llvm-project#125738. - For printing using pass manager, the `large_resource_limit` and `large_elements_limit` are completely independent of each other.
1 parent d788edb commit 22355e9

File tree

5 files changed

+41
-15
lines changed

5 files changed

+41
-15
lines changed

‎mlir/lib/Bindings/Python/IRCore.cpp‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,10 @@ static const char kOperationPrintDocstring[] =
9797
binary: Whether to write bytes (True) or str (False). Defaults to False.
9898
large_elements_limit: Whether to elide elements attributes above this
9999
number of elements. Defaults to None (no limit).
100+
large_resource_limit: Whether to elide resource attributes above this
101+
number of characters. Defaults to None (no limit). If large_elements_limit
102+
is set and this is None, the behavior will be to use large_elements_limit
103+
as large_resource_limit.
100104
enable_debug_info: Whether to print debug/location information. Defaults
101105
to False.
102106
pretty_debug_info: Whether to format debug information for easier reading
@@ -1303,6 +1307,7 @@ void PyOperation::checkValid() const {
13031307
}
13041308

13051309
void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
1310+
std::optional<int64_t> largeResourceLimit,
13061311
bool enableDebugInfo, bool prettyDebugInfo,
13071312
bool printGenericOpForm, bool useLocalScope,
13081313
bool useNameLocAsPrefix, bool assumeVerified,
@@ -1314,10 +1319,10 @@ void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
13141319
fileObject = nb::module_::import_("sys").attr("stdout");
13151320

13161321
MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
1317-
if (largeElementsLimit) {
1322+
if (largeElementsLimit)
13181323
mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);
1319-
mlirOpPrintingFlagsElideLargeResourceString(flags, *largeElementsLimit);
1320-
}
1324+
if (largeResourceLimit)
1325+
mlirOpPrintingFlagsElideLargeResourceString(flags, *largeResourceLimit);
13211326
if (enableDebugInfo)
13221327
mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
13231328
/*prettyForm=*/prettyDebugInfo);
@@ -1405,6 +1410,7 @@ void PyOperationBase::walk(
14051410

14061411
nb::object PyOperationBase::getAsm(bool binary,
14071412
std::optional<int64_t> largeElementsLimit,
1413+
std::optional<int64_t> largeResourceLimit,
14081414
bool enableDebugInfo, bool prettyDebugInfo,
14091415
bool printGenericOpForm, bool useLocalScope,
14101416
bool useNameLocAsPrefix, bool assumeVerified,
@@ -1416,6 +1422,7 @@ nb::object PyOperationBase::getAsm(bool binary,
14161422
fileObject = nb::module_::import_("io").attr("StringIO")();
14171423
}
14181424
print(/*largeElementsLimit=*/largeElementsLimit,
1425+
/*largeResourceLimit=*/largeResourceLimit,
14191426
/*enableDebugInfo=*/enableDebugInfo,
14201427
/*prettyDebugInfo=*/prettyDebugInfo,
14211428
/*printGenericOpForm=*/printGenericOpForm,
@@ -3348,6 +3355,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
33483355
[](PyOperationBase &self) {
33493356
return self.getAsm(/*binary=*/false,
33503357
/*largeElementsLimit=*/std::nullopt,
3358+
/*largeResourceLimit=*/std::nullopt,
33513359
/*enableDebugInfo=*/false,
33523360
/*prettyDebugInfo=*/false,
33533361
/*printGenericOpForm=*/false,
@@ -3363,11 +3371,12 @@ void mlir::python::populateIRCore(nb::module_ &m) {
33633371
nb::arg("state"), nb::arg("file").none() = nb::none(),
33643372
nb::arg("binary") = false, kOperationPrintStateDocstring)
33653373
.def("print",
3366-
nb::overload_cast<std::optional<int64_t>, bool, bool, bool, bool,
3367-
bool, bool, nb::object, bool, bool>(
3368-
&PyOperationBase::print),
3374+
nb::overload_cast<std::optional<int64_t>, std::optional<int64_t>,
3375+
bool, bool, bool, bool, bool, bool, nb::object,
3376+
bool, bool>(&PyOperationBase::print),
33693377
// Careful: Lots of arguments must match up with print method.
33703378
nb::arg("large_elements_limit").none() = nb::none(),
3379+
nb::arg("large_resource_limit").none() = nb::none(),
33713380
nb::arg("enable_debug_info") = false,
33723381
nb::arg("pretty_debug_info") = false,
33733382
nb::arg("print_generic_op_form") = false,
@@ -3383,6 +3392,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
33833392
// Careful: Lots of arguments must match up with get_asm method.
33843393
nb::arg("binary") = false,
33853394
nb::arg("large_elements_limit").none() = nb::none(),
3395+
nb::arg("large_resource_limit").none() = nb::none(),
33863396
nb::arg("enable_debug_info") = false,
33873397
nb::arg("pretty_debug_info") = false,
33883398
nb::arg("print_generic_op_form") = false,

‎mlir/lib/Bindings/Python/IRModule.h‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -599,18 +599,18 @@ class PyOperationBase {
599599
public:
600600
virtual ~PyOperationBase() = default;
601601
/// Implements the bound 'print' method and helps with others.
602-
void print(std::optional<int64_t> largeElementsLimit, bool enableDebugInfo,
602+
void print(std::optional<int64_t> largeElementsLimit,
603+
std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
603604
bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope,
604605
bool useNameLocAsPrefix, bool assumeVerified,
605606
nanobind::object fileObject, bool binary, bool skipRegions);
606607
void print(PyAsmState &state, nanobind::object fileObject, bool binary);
607608

608-
nanobind::object getAsm(bool binary,
609-
std::optional<int64_t> largeElementsLimit,
610-
bool enableDebugInfo, bool prettyDebugInfo,
611-
bool printGenericOpForm, bool useLocalScope,
612-
bool useNameLocAsPrefix, bool assumeVerified,
613-
bool skipRegions);
609+
nanobind::object
610+
getAsm(bool binary, std::optional<int64_t> largeElementsLimit,
611+
std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
612+
bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope,
613+
bool useNameLocAsPrefix, bool assumeVerified, bool skipRegions);
614614

615615
// Implement the bound 'writeBytecode' method.
616616
void writeBytecode(const nanobind::object &fileObject,

‎mlir/lib/Bindings/Python/Pass.cpp‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,12 +78,19 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
7878
[](PyPassManager &passManager, bool printBeforeAll,
7979
bool printAfterAll, bool printModuleScope, bool printAfterChange,
8080
bool printAfterFailure, std::optional<int64_t> largeElementsLimit,
81-
bool enableDebugInfo, bool printGenericOpForm,
81+
std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
82+
bool printGenericOpForm,
8283
std::optional<std::string> optionalTreePrintingPath) {
8384
MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
84-
if (largeElementsLimit)
85+
if (largeElementsLimit) {
8586
mlirOpPrintingFlagsElideLargeElementsAttrs(flags,
8687
*largeElementsLimit);
88+
mlirOpPrintingFlagsElideLargeResourceString(flags,
89+
*largeElementsLimit);
90+
}
91+
if (largeResourceLimit)
92+
mlirOpPrintingFlagsElideLargeResourceString(flags,
93+
*largeResourceLimit);
8794
if (enableDebugInfo)
8895
mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
8996
/*prettyForm=*/false);
@@ -103,6 +110,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
103110
"print_module_scope"_a = false, "print_after_change"_a = false,
104111
"print_after_failure"_a = false,
105112
"large_elements_limit"_a.none() = nb::none(),
113+
"large_resource_limit"_a.none() = nb::none(),
106114
"enable_debug_info"_a = false, "print_generic_op_form"_a = false,
107115
"tree_printing_dir_path"_a.none() = nb::none(),
108116
"Enable IR printing, default as mlir-print-ir-after-all.")

‎mlir/python/mlir/_mlir_libs/_mlir/ir.pyi‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ class _OperationBase:
200200
def get_asm(
201201
binary: Literal[True],
202202
large_elements_limit: int | None = None,
203+
large_resource_limit: int | None = None,
203204
enable_debug_info: bool = False,
204205
pretty_debug_info: bool = False,
205206
print_generic_op_form: bool = False,
@@ -212,6 +213,7 @@ class _OperationBase:
212213
self,
213214
binary: bool = False,
214215
large_elements_limit: int | None = None,
216+
large_resource_limit: int | None = None,
215217
enable_debug_info: bool = False,
216218
pretty_debug_info: bool = False,
217219
print_generic_op_form: bool = False,
@@ -253,6 +255,7 @@ class _OperationBase:
253255
def print(
254256
self,
255257
large_elements_limit: int | None = None,
258+
large_resource_limit: int | None = None,
256259
enable_debug_info: bool = False,
257260
pretty_debug_info: bool = False,
258261
print_generic_op_form: bool = False,
@@ -270,6 +273,10 @@ class _OperationBase:
270273
binary: Whether to write bytes (True) or str (False). Defaults to False.
271274
large_elements_limit: Whether to elide elements attributes above this
272275
number of elements. Defaults to None (no limit).
276+
large_resource_limit: Whether to elide resource strings above this
277+
number of characters. Defaults to None (no limit). If large_elements_limit
278+
is set and this is None, the behavior will be to use large_elements_limit
279+
as large_resource_limit.
273280
enable_debug_info: Whether to print debug/location information. Defaults
274281
to False.
275282
pretty_debug_info: Whether to format debug information for easier reading

‎mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class PassManager:
2323
print_after_change: bool = False,
2424
print_after_failure: bool = False,
2525
large_elements_limit: int | None = None,
26+
large_resource_limit: int | None = None,
2627
enable_debug_info: bool = False,
2728
print_generic_op_form: bool = False,
2829
tree_printing_dir_path: str | None = None,

0 commit comments

Comments
 (0)