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

[Majid Dadashi]

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

Do not erase location info when moving an op within the same block.

This change avoids resetting the location of an operation when it is moved to the front of a block it is already in. The location is only reset when the operation is moved from a different block. END_PUBLIC

>From fc53a40ca28214269476cc227107fb6c074a61f5 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] BEGIN_PUBLIC Do not erase location info when moving an op
 within the same block.

This change avoids resetting the location of an operation when it is moved to the front of a block it is already in. The location is only reset when the operation is moved from a different block.
END_PUBLIC
---
 mlir/lib/Transforms/Utils/FoldUtils.cpp          | 6 ++++--
 mlir/test/Transforms/canonicalize-debuginfo.mlir | 8 ++++----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index e9adda0cd01db..d22d697ac78df 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -155,10 +155,12 @@ bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   // 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.
   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..fc2b522e86f99 100644
--- a/mlir/test/Transforms/canonicalize-debuginfo.mlir
+++ b/mlir/test/Transforms/canonicalize-debuginfo.mlir
@@ -2,7 +2,7 @@
 
 // CHECK-LABEL: func @merge_constants
 func.func @merge_constants() -> (index, index, index, index) {
-  // CHECK-NEXT: arith.constant 42 : index loc(#[[UnknownLoc:.*]])
+  // CHECK: arith.constant 42 : index loc(#[[UnknownLoc:.*]])
   %0 = arith.constant 42 : index loc("merge_constants":0:0)
   %1 = arith.constant 42 : index loc("merge_constants":1:0)
   %2 = arith.constant 42 : index loc("merge_constants":2:0)
@@ -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,15 +26,15 @@ 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)
 
 // -----
 
 // CHECK-LABEL: func @hoist_and_merge
 func.func @hoist_and_merge(%arg0: memref<8xi32>) {
-  // CHECK-NEXT: arith.constant 42 : i32 loc(#[[UnknownLoc:.*]])
+  // CHECK: arith.constant 42 : i32 loc(#[[UnknownLoc:.*]])
   affine.for %arg1 = 0 to 8 {
     %0 = arith.constant 42 : i32 loc("hoist_and_merge":0:0)
     %1 = arith.constant 42 : i32 loc("hoist_and_merge":1:0)