From 3d75e3ac6771f25a6e4f7f20d629bf71be0693d9 Mon Sep 17 00:00:00 2001 From: Nikhil Kalra Date: Tue, 4 Feb 2025 10:15:59 -0800 Subject: [PATCH 1/2] [mlir] Python: Extend print large elemnts limit to resources 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. --- mlir/lib/Bindings/Python/IRCore.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 8e351cb22eb94..47a85c2a486fd 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -1296,8 +1296,10 @@ void PyOperationBase::print(std::optional largeElementsLimit, fileObject = nb::module_::import_("sys").attr("stdout"); MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate(); - if (largeElementsLimit) + if (largeElementsLimit) { mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit); + mlirOpPrintingFlagsElideLargeResourceString(flags, *largeElementsLimit); + } if (enableDebugInfo) mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true, /*prettyForm=*/prettyDebugInfo); From 15e66232bb528bb6fa8da95f90f44448ceea926d Mon Sep 17 00:00:00 2001 From: Nikhil Kalra Date: Tue, 4 Feb 2025 16:59:12 -0800 Subject: [PATCH 2/2] add test --- mlir/test/python/ir/operation.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py index d5d72c98b66ad..c2d3aed8808b4 100644 --- a/mlir/test/python/ir/operation.py +++ b/mlir/test/python/ir/operation.py @@ -583,14 +583,24 @@ def testOperationPrint(): r""" func.func @f1(%arg0: i32) -> i32 { %0 = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> loc("nom") + %1 = arith.constant dense_resource : tensor<3xi64> return %arg0 : i32 } + + {-# + dialect_resources: { + builtin: { + resource1: "0x08000000010000000000000002000000000000000300000000000000" + } + } + #-} """, ctx, ) # Test print to stdout. # CHECK: return %arg0 : i32 + # CHECK: resource1: "0x08 module.operation.print() # Test print to text file. @@ -607,7 +617,8 @@ def testOperationPrint(): module.operation.write_bytecode(bytecode_stream, desired_version=1) bytecode = bytecode_stream.getvalue() assert bytecode.startswith(b"ML\xefR"), "Expected bytecode to start with MLïR" - module_roundtrip = Module.parse(bytecode, ctx) + ctx2 = Context() + module_roundtrip = Module.parse(bytecode, ctx2) f = io.StringIO() module_roundtrip.operation.print(file=f) roundtrip_value = f.getvalue() @@ -633,7 +644,8 @@ def testOperationPrint(): # Test print with options. # CHECK: value = dense_resource<__elided__> : tensor<4xi32> - # CHECK: "func.return"(%arg0) : (i32) -> () -:4:7 + # CHECK: "func.return"(%arg0) : (i32) -> () -:5:7 + # CHECK-NOT: resource1: "0x08 module.operation.print( large_elements_limit=2, enable_debug_info=True,