[Mlir-commits] [mlir] 0ad1f83 - [mlir] Python: Extend print large elements limit to resources (#125738)

[llvmlistbot at llvm.org]

Author: Nikhil Kalra
Date: 2025-02-05T11:48:11-08:00
New Revision: 0ad1f8369c8637287367112dd777c525d3a8e383

URL: https://github.com/llvm/llvm-project/commit/0ad1f8369c8637287367112dd777c525d3a8e383
DIFF: https://github.com/llvm/llvm-project/commit/0ad1f8369c8637287367112dd777c525d3a8e383.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/test/python/ir/operation.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8e351cb22eb9489..47a85c2a486fd46 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1296,8 +1296,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 (enableDebugInfo)
     mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
                                        /*prettyForm=*/prettyDebugInfo);

diff  --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index d5d72c98b66ad0e..c2d3aed8808b477 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<resource1> : 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,