Skip to content

Commit 0ad1f83

Browse files
authored
[mlir] Python: Extend print large elements limit to resources (#125738)
If the large element limit is specified, large elements are hidden from the asm but large resources are not. This change extends the large elements limit to apply to printed resources as well.
1 parent d6cf04c commit 0ad1f83

File tree

2 files changed

+17
-3
lines changed

2 files changed

+17
-3
lines changed

mlir/lib/Bindings/Python/IRCore.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1296,8 +1296,10 @@ void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
12961296
fileObject = nb::module_::import_("sys").attr("stdout");
12971297

12981298
MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
1299-
if (largeElementsLimit)
1299+
if (largeElementsLimit) {
13001300
mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);
1301+
mlirOpPrintingFlagsElideLargeResourceString(flags, *largeElementsLimit);
1302+
}
13011303
if (enableDebugInfo)
13021304
mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
13031305
/*prettyForm=*/prettyDebugInfo);

mlir/test/python/ir/operation.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -583,14 +583,24 @@ def testOperationPrint():
583583
r"""
584584
func.func @f1(%arg0: i32) -> i32 {
585585
%0 = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> loc("nom")
586+
%1 = arith.constant dense_resource<resource1> : tensor<3xi64>
586587
return %arg0 : i32
587588
}
589+
590+
{-#
591+
dialect_resources: {
592+
builtin: {
593+
resource1: "0x08000000010000000000000002000000000000000300000000000000"
594+
}
595+
}
596+
#-}
588597
""",
589598
ctx,
590599
)
591600

592601
# Test print to stdout.
593602
# CHECK: return %arg0 : i32
603+
# CHECK: resource1: "0x08
594604
module.operation.print()
595605

596606
# Test print to text file.
@@ -607,7 +617,8 @@ def testOperationPrint():
607617
module.operation.write_bytecode(bytecode_stream, desired_version=1)
608618
bytecode = bytecode_stream.getvalue()
609619
assert bytecode.startswith(b"ML\xefR"), "Expected bytecode to start with MLïR"
610-
module_roundtrip = Module.parse(bytecode, ctx)
620+
ctx2 = Context()
621+
module_roundtrip = Module.parse(bytecode, ctx2)
611622
f = io.StringIO()
612623
module_roundtrip.operation.print(file=f)
613624
roundtrip_value = f.getvalue()
@@ -633,7 +644,8 @@ def testOperationPrint():
633644

634645
# Test print with options.
635646
# CHECK: value = dense_resource<__elided__> : tensor<4xi32>
636-
# CHECK: "func.return"(%arg0) : (i32) -> () -:4:7
647+
# CHECK: "func.return"(%arg0) : (i32) -> () -:5:7
648+
# CHECK-NOT: resource1: "0x08
637649
module.operation.print(
638650
large_elements_limit=2,
639651
enable_debug_info=True,

0 commit comments

Comments
 (0)