[Mlir-commits] [mlir] Avoid unnecessary erasing of constant Locs (PR #151573)

[Majid Dadashi]

https://github.com/majiddadashi updated https://github.com/llvm/llvm-project/pull/151573

>From fb4b5637e74b43ac93eeada1aed1cdc5b6d7001d Mon Sep 17 00:00:00 2001
From: Majid Dadashi <majiddadashi at google.com>
Date: Thu, 31 Jul 2025 18:16:19 +0000
Subject: [PATCH] Do not erase location info when moving an op within the same
 block.

The MLIR constant folding utility erases location information when moving an operation. This was being done even when an operation was moved to the front of a block it was already in.

In TFLite, this location information is used to provide meaningful names for tensors, which aids in debugging and mapping compiled tensors back to their original layers. The aggressive erasure of location info caused many tensors in TFLite models to receive generic names (e.g., tfl.pseudo_qconst), making the models harder to inspect.

This change modifies the logic to preserve the location of an operation when it is moved within the same block. The location is now only erased when the operation is moved from a different block entirely. This ensures that most tensor names are preserved, improving the debugging experience for TFLite models.
---
 mlir/lib/Transforms/Utils/FoldUtils.cpp          | 7 +++++--
 mlir/test/Transforms/canonicalize-debuginfo.mlir | 4 ++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index e9adda0cd01db..5e07509871ea2 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -154,11 +154,14 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   // already at the front of the block, or the previous operation is already a
   // constant we unique'd (i.e. one we inserted), then we don't need to do
   // anything. Otherwise, we move the constant to the insertion block.
+  // The location info is erased if the constant is moved to a different block.
   Block *insertBlock = &insertRegion->front();
-  if (opBlock != insertBlock || (&insertBlock->front() != op &&
-                                 !isFolderOwnedConstant(op->getPrevNode()))) {
+  if (opBlock != insertBlock) {
     op->moveBefore(&insertBlock->front());
     op->setLoc(erasedFoldedLocation);
+  } else if (&insertBlock->front() != op &&
+             !isFolderOwnedConstant(op->getPrevNode())) {
+    op->moveBefore(&insertBlock->front());
   }
 
   folderConstOp = op;
diff --git a/mlir/test/Transforms/canonicalize-debuginfo.mlir b/mlir/test/Transforms/canonicalize-debuginfo.mlir
index 30c8022daa76b..5af0023fed801 100644
--- a/mlir/test/Transforms/canonicalize-debuginfo.mlir
+++ b/mlir/test/Transforms/canonicalize-debuginfo.mlir
@@ -15,7 +15,7 @@ func.func @merge_constants() -> (index, index, index, index) {
 
 // CHECK-LABEL: func @simple_hoist
 func.func @simple_hoist(%arg0: memref<8xi32>) -> i32 {
-  // CHECK: arith.constant 88 : i32 loc(#[[UnknownLoc:.*]])
+  // CHECK: arith.constant 88 : i32 loc(#[[ConstLoc2:.*]])
   // CHECK: arith.constant 42 : i32 loc(#[[ConstLoc0:.*]])
   // CHECK: arith.constant 0 : index loc(#[[ConstLoc1:.*]])
   %0 = arith.constant 42 : i32 loc("simple_hoist":0:0)
@@ -26,9 +26,9 @@ func.func @simple_hoist(%arg0: memref<8xi32>) -> i32 {
 
   return %2 : i32
 }
-// CHECK-DAG: #[[UnknownLoc]] = loc(unknown)
 // CHECK-DAG: #[[ConstLoc0]] = loc("simple_hoist":0:0)
 // CHECK-DAG: #[[ConstLoc1]] = loc("simple_hoist":1:0)
+// CHECK-DAG: #[[ConstLoc2]] = loc("simple_hoist":2:0)
 
 // -----