Skip to content

Commit e4a3541

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

File tree

7 files changed

+107
-15
lines changed

7 files changed

+107
-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,

‎mlir/test/python/ir/operation.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -686,6 +686,15 @@ def testOperationPrint():
686686
skip_regions=True,
687687
)
688688

689+
# Test print with large_resource_limit.
690+
# CHECK: func.func @f1(%arg0: i32) -> i32
691+
# CHECK-NOT: resource1: "0x08
692+
module.operation.print(large_resource_limit=2)
693+
694+
# Test large_elements_limit has no effect on resource string
695+
# CHECK: func.func @f1(%arg0: i32) -> i32
696+
# CHECK: resource1: "0x08
697+
module.operation.print(large_elements_limit=2)
689698

690699
# CHECK-LABEL: TEST: testKnownOpView
691700
@run

‎mlir/test/python/pass_manager.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,63 @@ def testPrintIrLargeLimitElements():
363363
pm.run(module)
364364

365365

366+
# CHECK-LABEL: TEST: testPrintIrLargeResourceLimit
367+
@run
368+
def testPrintIrLargeResourceLimit():
369+
with Context() as ctx:
370+
module = ModuleOp.parse(
371+
"""
372+
module {
373+
func.func @main() -> tensor<3xi64> {
374+
%0 = arith.constant dense_resource<blob1> : tensor<3xi64>
375+
return %0 : tensor<3xi64>
376+
}
377+
}
378+
{-#
379+
dialect_resources: {
380+
builtin: {
381+
blob1: "0x010000000000000002000000000000000300000000000000"
382+
}
383+
}
384+
#-}
385+
"""
386+
)
387+
pm = PassManager.parse("builtin.module(canonicalize)")
388+
ctx.enable_multithreading(False)
389+
pm.enable_ir_printing(large_resource_limit=4)
390+
# CHECK-NOT: blob1: "0x01
391+
pm.run(module)
392+
393+
394+
# CHECK-LABEL: TEST: testPrintIrLargeResourceLimitVsElementsLimit
395+
@run
396+
def testPrintIrLargeResourceLimitVsElementsLimit():
397+
"""Test that large_elements_limit does not affect the printing of resources."""
398+
with Context() as ctx:
399+
module = ModuleOp.parse(
400+
"""
401+
module {
402+
func.func @main() -> tensor<3xi64> {
403+
%0 = arith.constant dense_resource<blob1> : tensor<3xi64>
404+
return %0 : tensor<3xi64>
405+
}
406+
}
407+
{-#
408+
dialect_resources: {
409+
builtin: {
410+
blob1: "0x010000000000000002000000000000000300000000000000"
411+
}
412+
}
413+
#-}
414+
"""
415+
)
416+
pm = PassManager.parse("builtin.module(canonicalize)")
417+
ctx.enable_multithreading(False)
418+
pm.enable_ir_printing(large_elements_limit=1)
419+
# CHECK-NOT: blob1: "0x01
420+
pm.run(module)
421+
422+
366423
# CHECK-LABEL: TEST: testPrintIrTree
367424
@run
368425
def testPrintIrTree():

0 commit comments

Comments
 (0)