[Mlir-commits] [flang] [mlir] [mlir] add option to print SSA IDs using `NameLoc`s as prefixes (PR #119996)

Maksim Levental llvmlistbot at llvm.org
Sun Dec 15 15:42:05 PST 2024


https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/119996

>From 447bce768ab03da2a0adfb31389c3bda2b20017f Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Fri, 13 Dec 2024 23:39:19 -0500
Subject: [PATCH 1/5] [mlir] retain identifier names

---
 mlir/include/mlir/IR/OperationSupport.h |  7 ++++
 mlir/lib/IR/AsmPrinter.cpp              | 45 +++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 1b93f3d3d04fe8..c6a603ae518785 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1221,6 +1221,10 @@ class OpPrintingFlags {
   /// Return if printer should use unique SSA IDs.
   bool shouldPrintUniqueSSAIDs() const;
 
+  /// Returns if the printer should retain identifier names collected using
+  /// parsing.
+  bool shouldUseNameLocAsPrefix() const;
+
 private:
   /// Elide large elements attributes if the number of elements is larger than
   /// the upper limit.
@@ -1254,6 +1258,9 @@ class OpPrintingFlags {
 
   /// Print unique SSA IDs for values, block arguments and naming conflicts
   bool printUniqueSSAIDsFlag : 1;
+
+  /// Print the retained original names of identifiers
+  bool useNameLocAsPrefix : 1;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 61b90bc9b0a7bb..18db2dc555c766 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -195,6 +195,10 @@ struct AsmPrinterOptions {
       "mlir-print-unique-ssa-ids", llvm::cl::init(false),
       llvm::cl::desc("Print unique SSA ID numbers for values, block arguments "
                      "and naming conflicts across all regions")};
+
+  llvm::cl::opt<bool> useNameLocAsPrefix{
+      "mlir-use-nameloc-as-prefix", llvm::cl::init(false),
+      llvm::cl::desc("TODO")};
 };
 } // namespace
 
@@ -212,7 +216,8 @@ OpPrintingFlags::OpPrintingFlags()
     : printDebugInfoFlag(false), printDebugInfoPrettyFormFlag(false),
       printGenericOpFormFlag(false), skipRegionsFlag(false),
       assumeVerifiedFlag(false), printLocalScope(false),
-      printValueUsersFlag(false), printUniqueSSAIDsFlag(false) {
+      printValueUsersFlag(false), printUniqueSSAIDsFlag(false),
+      useNameLocAsPrefix(false) {
   // Initialize based upon command line options, if they are available.
   if (!clOptions.isConstructed())
     return;
@@ -231,6 +236,7 @@ OpPrintingFlags::OpPrintingFlags()
   skipRegionsFlag = clOptions->skipRegionsOpt;
   printValueUsersFlag = clOptions->printValueUsers;
   printUniqueSSAIDsFlag = clOptions->printUniqueSSAIDs;
+  useNameLocAsPrefix = clOptions->useNameLocAsPrefix;
 }
 
 /// Enable the elision of large elements attributes, by printing a '...'
@@ -362,6 +368,11 @@ bool OpPrintingFlags::shouldPrintUniqueSSAIDs() const {
   return printUniqueSSAIDsFlag || shouldPrintGenericOpForm();
 }
 
+/// TODO
+bool OpPrintingFlags::shouldUseNameLocAsPrefix() const {
+  return useNameLocAsPrefix;
+}
+
 //===----------------------------------------------------------------------===//
 // NewLineCounter
 //===----------------------------------------------------------------------===//
@@ -1514,10 +1525,22 @@ void SSANameState::numberValuesInRegion(Region &region) {
     setValueName(arg, name);
   };
 
+  bool alreadySetNames = false;
   if (!printerFlags.shouldPrintGenericOpForm()) {
     if (Operation *op = region.getParentOp()) {
-      if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op))
+      if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op)) {
         asmInterface.getAsmBlockArgumentNames(region, setBlockArgNameFn);
+        alreadySetNames = true;
+      }
+    }
+  }
+
+  if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames) {
+    for (BlockArgument arg : region.getArguments()) {
+      if (isa<NameLoc>(arg.getLoc())) {
+        auto nameLoc = cast<NameLoc>(arg.getLoc());
+        setBlockArgNameFn(arg, nameLoc.getName());
+      }
     }
   }
 
@@ -1553,7 +1576,12 @@ void SSANameState::numberValuesInBlock(Block &block) {
       specialNameBuffer.resize(strlen("arg"));
       specialName << nextArgumentID++;
     }
-    setValueName(arg, specialName.str());
+    if (printerFlags.shouldUseNameLocAsPrefix() && isa<NameLoc>(arg.getLoc())) {
+      auto nameLoc = cast<NameLoc>(arg.getLoc());
+      setValueName(arg, nameLoc.getName());
+    } else {
+      setValueName(arg, specialName.str());
+    }
   }
 
   // Number the operations in this block.
@@ -1589,10 +1617,21 @@ void SSANameState::numberValuesInOp(Operation &op) {
     blockNames[block] = {-1, name};
   };
 
+  bool alreadySetNames = false;
   if (!printerFlags.shouldPrintGenericOpForm()) {
     if (OpAsmOpInterface asmInterface = dyn_cast<OpAsmOpInterface>(&op)) {
       asmInterface.getAsmBlockNames(setBlockNameFn);
       asmInterface.getAsmResultNames(setResultNameFn);
+      alreadySetNames = true;
+    }
+  }
+
+  if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames) {
+    for (Value opResult : op.getResults()) {
+      if (isa<NameLoc>(opResult.getLoc())) {
+        auto nameLoc = cast<NameLoc>(opResult.getLoc());
+        setResultNameFn(opResult, nameLoc.getName());
+      }
     }
   }
 

>From bbca902568bfb047ef67a96dad631314669fd6f9 Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Sat, 14 Dec 2024 23:21:15 -0500
Subject: [PATCH 2/5] add mlir-use-nameloc-as-prefix and set to true and "fix"
 tests that break

---
 flang/test/Lower/OpenMP/map-component-ref.f90 | 41 ++++++++++++-------
 mlir/lib/IR/AsmPrinter.cpp                    |  6 +--
 mlir/test/IR/locations.mlir                   | 14 +++----
 mlir/test/IR/pretty-locations.mlir            |  2 +-
 mlir/test/IR/pretty_printed_region_op.mlir    | 12 +++---
 mlir/test/IR/wrapping_op.mlir                 |  6 +--
 mlir/test/mlir-tblgen/pattern.mlir            |  2 +-
 7 files changed, 48 insertions(+), 35 deletions(-)

diff --git a/flang/test/Lower/OpenMP/map-component-ref.f90 b/flang/test/Lower/OpenMP/map-component-ref.f90
index b7a7ee06b02f29..8abe7bd8658507 100644
--- a/flang/test/Lower/OpenMP/map-component-ref.f90
+++ b/flang/test/Lower/OpenMP/map-component-ref.f90
@@ -1,19 +1,32 @@
-! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
-! RUN: bbc -fopenmp -emit-hlfir %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s --check-prefix=CHECK-FLANG
+! RUN: bbc -fopenmp -emit-hlfir %s -o - | FileCheck %s --check-prefix=CHECK-BBC
 
 ! CHECK-LABEL: func.func @_QPfoo1
-! CHECK: %[[V0:[0-9]+]] = fir.alloca !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}> {bindc_name = "a", uniq_name = "_QFfoo1Ea"}
-! CHECK: %[[V1:[0-9]+]]:2 = hlfir.declare %[[V0]] {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
-! CHECK: %[[V2:[0-9]+]] = hlfir.designate %[[V1]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
-! CHECK: %[[V3:[0-9]+]] = omp.map.info var_ptr(%[[V2]] : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) -> !fir.ref<i32> {name = "a%a1"}
-! CHECK: %[[V4:[0-9]+]] = omp.map.info var_ptr(%[[V1]]#1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>) map_clauses(tofrom) capture(ByRef) members(%[[V3]] : [1] : !fir.ref<i32>) -> !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>> {name = "a", partial_map = true}
-! CHECK: omp.target map_entries(%[[V4]] -> %arg0, %[[V3]] -> %arg1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<i32>) {
-! CHECK:   %[[V5:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
-! CHECK:   %c0_i32 = arith.constant 0 : i32
-! CHECK:   %[[V6:[0-9]+]] = hlfir.designate %[[V5]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
-! CHECK:   hlfir.assign %c0_i32 to %[[V6]] : i32, !fir.ref<i32>
-! CHECK:   omp.terminator
-! CHECK: }
+! CHECK-FLANG: %[[V0:[0-9]+]] = fir.alloca !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}> {bindc_name = "a", uniq_name = "_QFfoo1Ea"}
+! CHECK-FLANG: %[[V1:[0-9]+]]:2 = hlfir.declare %[[V0]] {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
+! CHECK-FLANG: %[[V2:[0-9]+]] = hlfir.designate %[[V1]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
+! CHECK-FLANG: %[[V3:[0-9]+]] = omp.map.info var_ptr(%[[V2]] : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) -> !fir.ref<i32> {name = "a%a1"}
+! CHECK-FLANG: %[[V4:[0-9]+]] = omp.map.info var_ptr(%[[V1]]#1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>) map_clauses(tofrom) capture(ByRef) members(%[[V3]] : [1] : !fir.ref<i32>) -> !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>> {name = "a", partial_map = true}
+! CHECK-FLANG: omp.target map_entries(%[[V4]] -> %arg0, %[[V3]] -> %arg1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<i32>) {
+! CHECK-FLANG:   %[[V5:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
+! CHECK-FLANG:   %c0_i32 = arith.constant 0 : i32
+! CHECK-FLANG:   %[[V6:[0-9]+]] = hlfir.designate %[[V5]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
+! CHECK-FLANG:   hlfir.assign %c0_i32 to %[[V6]] : i32, !fir.ref<i32>
+! CHECK-FLANG:   omp.terminator
+! CHECK-FLANG: }
+
+! CHECK-BBC: %[[V0:[0-9]+]] = fir.alloca !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}> {bindc_name = "a", uniq_name = "_QFfoo1Ea"}
+! CHECK-BBC: %[[V1:[0-9]+]]:2 = hlfir.declare %[[V0]] {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
+! CHECK-BBC: %[[V2:[0-9]+]] = hlfir.designate %[[V1]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
+! CHECK-BBC: %a25a1 = omp.map.info var_ptr(%[[V2]] : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) -> !fir.ref<i32> {name = "a%a1"}
+! CHECK-BBC: %[[V4:[0-9]+]] = omp.map.info var_ptr(%[[V1]]#1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>) map_clauses(tofrom) capture(ByRef) members(%a25a1 : [1] : !fir.ref<i32>) -> !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>> {name = "a", partial_map = true}
+! CHECK-BBC: omp.target map_entries(%[[V4]] -> %arg0, %a25a1 -> %a25a1_3 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<i32>) {
+! CHECK-BBC:   %[[V5:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
+! CHECK-BBC:   %c0_i32 = arith.constant 0 : i32
+! CHECK-BBC:   %[[V6:[0-9]+]] = hlfir.designate %[[V5]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
+! CHECK-BBC:   hlfir.assign %c0_i32 to %[[V6]] : i32, !fir.ref<i32>
+! CHECK-BBC:   omp.terminator
+! CHECK-BBC: }
 
 subroutine foo1()
   implicit none
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 18db2dc555c766..2bab5afc2775f2 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -196,9 +196,9 @@ struct AsmPrinterOptions {
       llvm::cl::desc("Print unique SSA ID numbers for values, block arguments "
                      "and naming conflicts across all regions")};
 
-  llvm::cl::opt<bool> useNameLocAsPrefix{
-      "mlir-use-nameloc-as-prefix", llvm::cl::init(false),
-      llvm::cl::desc("TODO")};
+  llvm::cl::opt<bool> useNameLocAsPrefix{"mlir-use-nameloc-as-prefix",
+                                         llvm::cl::init(true),
+                                         llvm::cl::desc("TODO")};
 };
 } // namespace
 
diff --git a/mlir/test/IR/locations.mlir b/mlir/test/IR/locations.mlir
index b725307b420b79..dc39315a765a83 100644
--- a/mlir/test/IR/locations.mlir
+++ b/mlir/test/IR/locations.mlir
@@ -13,7 +13,7 @@ func.func @inline_notation() -> i32 {
   // CHECK: arith.constant 4 : index loc(callsite("foo" at "mysource.cc":10:8))
   %2 = arith.constant 4 : index loc(callsite("foo" at "mysource.cc":10:8))
 
-  // CHECK: affine.for %arg0 loc("IVlocation") = 0 to 8 {
+  // CHECK: affine.for %IVlocation loc("IVlocation") = 0 to 8 {
   // CHECK: } loc(fused["foo", "mysource.cc":10:8])
   affine.for %i0 loc("IVlocation") = 0 to 8 {
   } loc(fused["foo", "mysource.cc":10:8])
@@ -26,7 +26,7 @@ func.func @inline_notation() -> i32 {
   affine.if #set0(%2) {
   } loc(fused<"myPass">["foo"])
 
-  // CHECK: return %0 : i32 loc(unknown)
+  // CHECK: return %foo : i32 loc(unknown)
   return %1 : i32 loc(unknown)
 }
 
@@ -62,7 +62,7 @@ func.func @escape_strings() {
 // CHECK-LABEL: func @argLocs(
 // CHECK-SAME:  %arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+1]]:20),
 func.func @argLocs(%x: i32,
-// CHECK-SAME:  %arg1: i64 loc("hotdog")
+// CHECK-SAME:  %hotdog: i64 loc("hotdog")
               %y: i64 loc("hotdog")) {
   return
 }
@@ -73,11 +73,11 @@ func.func @argLocs(%x: i32,
 // CHECK-NEXT: ^bb0(%arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+2]]:7),
 // CHECK-ALIAS-NEXT: ^bb0(%arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+1]]:7),
  ^bb0(%x: i32,
-// CHECK-SAME: %arg1: i32 loc("cheetos"),
-// CHECK-ALIAS-SAME: %arg1: i32 loc("cheetos"),
+// CHECK-SAME: %cheetos: i32 loc("cheetos"),
+// CHECK-ALIAS-SAME: %cheetos: i32 loc("cheetos"),
       %y: i32 loc("cheetos"),
-// CHECK-SAME: %arg2: i32 loc("out_of_line_location2")):
-// CHECK-ALIAS-SAME: %arg2: i32 loc("out_of_line_location2")):
+// CHECK-SAME: %out_of_line_location2: i32 loc("out_of_line_location2")):
+// CHECK-ALIAS-SAME: %out_of_line_location2: i32 loc("out_of_line_location2")):
       %z: i32 loc("out_of_line_location2")):
     %1 = arith.addi %x, %y : i32
     "foo.yield"(%1) : (i32) -> ()
diff --git a/mlir/test/IR/pretty-locations.mlir b/mlir/test/IR/pretty-locations.mlir
index 598bebeb83aebd..30391252cabfbe 100644
--- a/mlir/test/IR/pretty-locations.mlir
+++ b/mlir/test/IR/pretty-locations.mlir
@@ -27,6 +27,6 @@ func.func @inline_notation() -> i32 {
   // CHECK: "foo.op"() : () -> () #test.custom_location<"foo.mlir" * 1234>
   "foo.op"() : () -> () loc(#test.custom_location<"foo.mlir" * 1234>)
 
-  // CHECK: return %0 : i32 [unknown]
+  // CHECK: return %foo : i32 [unknown]
   return %1 : i32 loc(unknown)
 }
diff --git a/mlir/test/IR/pretty_printed_region_op.mlir b/mlir/test/IR/pretty_printed_region_op.mlir
index 0a81f9d53a3097..8c6d71d1bf10f2 100644
--- a/mlir/test/IR/pretty_printed_region_op.mlir
+++ b/mlir/test/IR/pretty_printed_region_op.mlir
@@ -6,9 +6,9 @@
 
 func.func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
 // CHECK-CUSTOM:  test.pretty_printed_region %arg1, %arg0 start test.special.op end : (f32, f32) -> f32
-// CHECK-GENERIC: "test.pretty_printed_region"(%arg1, %arg0)
-// CHECK-GENERIC:   ^bb0(%arg[[x:[0-9]+]]: f32, %arg[[y:[0-9]+]]: f32
-// CHECK-GENERIC:     %[[RES:.*]] = "test.special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
+// CHECK-GENERIC: %some_NameLoc = "test.pretty_printed_region"(%arg1, %arg0)
+// CHECK-GENERIC:   ^bb0(%some_NameLoc_[[x:[0-9]+]]: f32, %some_NameLoc_[[y:[0-9]+]]: f32
+// CHECK-GENERIC:     %[[RES:.*]] = "test.special.op"(%some_NameLoc_[[x]], %some_NameLoc_[[y]]) : (f32, f32) -> f32
 // CHECK-GENERIC:     "test.return"(%[[RES]]) : (f32) -> ()
 // CHECK-GENERIC:  : (f32, f32) -> f32
 
@@ -37,9 +37,9 @@ func.func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
 // -----
 
 func.func @pretty_printed_region_op_deferred_loc(%arg0 : f32, %arg1 : f32) -> (f32) {
-// CHECK-LOCATION: "test.pretty_printed_region"(%arg1, %arg0)
-// CHECK-LOCATION:   ^bb0(%arg[[x:[0-9]+]]: f32 loc("foo"), %arg[[y:[0-9]+]]: f32 loc("foo")
-// CHECK-LOCATION:     %[[RES:.*]] = "test.special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
+// CHECK-LOCATION: %foo = "test.pretty_printed_region"(%arg1, %arg0)
+// CHECK-LOCATION:   ^bb0(%foo_[[x:[0-9]+]]: f32 loc("foo"), %foo_[[y:[0-9]+]]: f32 loc("foo")
+// CHECK-LOCATION:     %[[RES:.*]] = "test.special.op"(%foo_[[x]], %foo_[[y]]) : (f32, f32) -> f32
 // CHECK-LOCATION:     "test.return"(%[[RES]]) : (f32) -> ()
 // CHECK-LOCATION:  : (f32, f32) -> f32
 
diff --git a/mlir/test/IR/wrapping_op.mlir b/mlir/test/IR/wrapping_op.mlir
index 3f7a350ee654e7..7ddfca3e4bf64b 100644
--- a/mlir/test/IR/wrapping_op.mlir
+++ b/mlir/test/IR/wrapping_op.mlir
@@ -5,10 +5,10 @@
 // CHECK-GENERIC: "func.func"
 // CHECK-GENERIC-SAME: sym_name = "wrapping_op"
 func.func @wrapping_op(%arg0 : i32, %arg1 : f32) -> (i3, i2, i1) {
-// CHECK: %0:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
+// CHECK: %some_NameLoc, %some_NameLoc_0, %some_NameLoc_1 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
 // CHECK-GENERIC: "test.wrapping_region"() ({
-// CHECK-GENERIC:   %[[NESTED_RES:.*]]:3 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
-// CHECK-GENERIC:   "test.return"(%[[NESTED_RES]]#0, %[[NESTED_RES]]#1, %[[NESTED_RES]]#2) : (i1, i2, i3) -> () loc("some_NameLoc")
+// CHECK-GENERIC:   %some_NameLoc_2, %some_NameLoc_3, %some_NameLoc_4 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
+// CHECK-GENERIC:   "test.return"(%some_NameLoc_2, %some_NameLoc_3, %some_NameLoc_4) : (i1, i2, i3) -> () loc("some_NameLoc")
 // CHECK-GENERIC: }) : () -> (i1, i2, i3) loc("some_NameLoc")
   %res:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) { test.attr = "attr" } : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
   return %res#2, %res#1, %res#0 : i3, i2, i1
diff --git a/mlir/test/mlir-tblgen/pattern.mlir b/mlir/test/mlir-tblgen/pattern.mlir
index 60d46e676d2a33..5594e80420a722 100644
--- a/mlir/test/mlir-tblgen/pattern.mlir
+++ b/mlir/test/mlir-tblgen/pattern.mlir
@@ -678,7 +678,7 @@ func.func @returnTypeAndLocation(%arg0 : i32) -> i1 {
   %0 = "test.source_op"(%arg0) {tag = 66 : i32} : (i32) -> i1
   // CHECK: "test.op_x"(%arg0) : (i32) -> i32 loc("loc1")
   // CHECK: "test.op_x"(%arg0) : (i32) -> i32 loc("loc2")
-  // CHECK: "test.two_to_one"(%0, %1) : (i32, i32) -> i1
+  // CHECK: "test.two_to_one"(%loc1, %loc2) : (i32, i32) -> i1
   return %0 : i1
 }
 

>From ca084f95b2c1a424392049bc4a8555f4f23d83fe Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Sun, 15 Dec 2024 01:18:46 -0500
Subject: [PATCH 3/5] fix result groups

---
 mlir/lib/IR/AsmPrinter.cpp    | 14 ++++++++------
 mlir/test/IR/wrapping_op.mlir |  6 +++---
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 2bab5afc2775f2..e06ab003b0f6b9 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -73,7 +73,8 @@ OpAsmParser::~OpAsmParser() = default;
 MLIRContext *AsmParser::getContext() const { return getBuilder().getContext(); }
 
 /// Parse a type list.
-/// This is out-of-line to work-around https://github.com/llvm/llvm-project/issues/62918
+/// This is out-of-line to work-around
+/// https://github.com/llvm/llvm-project/issues/62918
 ParseResult AsmParser::parseTypeList(SmallVectorImpl<Type> &result) {
   return parseCommaSeparatedList(
       [&]() { return parseType(result.emplace_back()); });
@@ -1626,16 +1627,17 @@ void SSANameState::numberValuesInOp(Operation &op) {
     }
   }
 
+  unsigned numResults = op.getNumResults();
   if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames) {
-    for (Value opResult : op.getResults()) {
-      if (isa<NameLoc>(opResult.getLoc())) {
-        auto nameLoc = cast<NameLoc>(opResult.getLoc());
-        setResultNameFn(opResult, nameLoc.getName());
+    if (numResults > 0) {
+      Value resultBegin = op.getResult(0);
+      if (isa<NameLoc>(resultBegin.getLoc())) {
+        auto nameLoc = cast<NameLoc>(resultBegin.getLoc());
+        setResultNameFn(resultBegin, nameLoc.getName());
       }
     }
   }
 
-  unsigned numResults = op.getNumResults();
   if (numResults == 0) {
     // If value users should be printed, operations with no result need an id.
     if (printerFlags.shouldPrintValueUsers()) {
diff --git a/mlir/test/IR/wrapping_op.mlir b/mlir/test/IR/wrapping_op.mlir
index 7ddfca3e4bf64b..95623cda32ea0f 100644
--- a/mlir/test/IR/wrapping_op.mlir
+++ b/mlir/test/IR/wrapping_op.mlir
@@ -5,10 +5,10 @@
 // CHECK-GENERIC: "func.func"
 // CHECK-GENERIC-SAME: sym_name = "wrapping_op"
 func.func @wrapping_op(%arg0 : i32, %arg1 : f32) -> (i3, i2, i1) {
-// CHECK: %some_NameLoc, %some_NameLoc_0, %some_NameLoc_1 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
+// CHECK: %some_NameLoc:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
 // CHECK-GENERIC: "test.wrapping_region"() ({
-// CHECK-GENERIC:   %some_NameLoc_2, %some_NameLoc_3, %some_NameLoc_4 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
-// CHECK-GENERIC:   "test.return"(%some_NameLoc_2, %some_NameLoc_3, %some_NameLoc_4) : (i1, i2, i3) -> () loc("some_NameLoc")
+// CHECK-GENERIC:   %[[NESTED_RES:.*]]:3 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
+// CHECK-GENERIC:   "test.return"(%[[NESTED_RES]]#0, %[[NESTED_RES]]#1, %[[NESTED_RES]]#2) : (i1, i2, i3) -> () loc("some_NameLoc")
 // CHECK-GENERIC: }) : () -> (i1, i2, i3) loc("some_NameLoc")
   %res:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) { test.attr = "attr" } : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
   return %res#2, %res#1, %res#0 : i3, i2, i1

>From 26c0de4228770155c3ad2dc99fb7670c77561888 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Sun, 15 Dec 2024 01:23:58 -0500
Subject: [PATCH 4/5] turn off by default

---
 flang/test/Lower/OpenMP/map-component-ref.f90 | 41 +++++++------------
 mlir/lib/IR/AsmPrinter.cpp                    | 15 ++++---
 mlir/test/IR/locations.mlir                   | 14 +++----
 mlir/test/IR/pretty-locations.mlir            |  2 +-
 mlir/test/IR/pretty_printed_region_op.mlir    | 12 +++---
 mlir/test/IR/wrapping_op.mlir                 |  2 +-
 mlir/test/mlir-tblgen/pattern.mlir            |  2 +-
 7 files changed, 37 insertions(+), 51 deletions(-)

diff --git a/flang/test/Lower/OpenMP/map-component-ref.f90 b/flang/test/Lower/OpenMP/map-component-ref.f90
index 8abe7bd8658507..b7a7ee06b02f29 100644
--- a/flang/test/Lower/OpenMP/map-component-ref.f90
+++ b/flang/test/Lower/OpenMP/map-component-ref.f90
@@ -1,32 +1,19 @@
-! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s --check-prefix=CHECK-FLANG
-! RUN: bbc -fopenmp -emit-hlfir %s -o - | FileCheck %s --check-prefix=CHECK-BBC
+! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+! RUN: bbc -fopenmp -emit-hlfir %s -o - | FileCheck %s
 
 ! CHECK-LABEL: func.func @_QPfoo1
-! CHECK-FLANG: %[[V0:[0-9]+]] = fir.alloca !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}> {bindc_name = "a", uniq_name = "_QFfoo1Ea"}
-! CHECK-FLANG: %[[V1:[0-9]+]]:2 = hlfir.declare %[[V0]] {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
-! CHECK-FLANG: %[[V2:[0-9]+]] = hlfir.designate %[[V1]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
-! CHECK-FLANG: %[[V3:[0-9]+]] = omp.map.info var_ptr(%[[V2]] : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) -> !fir.ref<i32> {name = "a%a1"}
-! CHECK-FLANG: %[[V4:[0-9]+]] = omp.map.info var_ptr(%[[V1]]#1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>) map_clauses(tofrom) capture(ByRef) members(%[[V3]] : [1] : !fir.ref<i32>) -> !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>> {name = "a", partial_map = true}
-! CHECK-FLANG: omp.target map_entries(%[[V4]] -> %arg0, %[[V3]] -> %arg1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<i32>) {
-! CHECK-FLANG:   %[[V5:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
-! CHECK-FLANG:   %c0_i32 = arith.constant 0 : i32
-! CHECK-FLANG:   %[[V6:[0-9]+]] = hlfir.designate %[[V5]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
-! CHECK-FLANG:   hlfir.assign %c0_i32 to %[[V6]] : i32, !fir.ref<i32>
-! CHECK-FLANG:   omp.terminator
-! CHECK-FLANG: }
-
-! CHECK-BBC: %[[V0:[0-9]+]] = fir.alloca !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}> {bindc_name = "a", uniq_name = "_QFfoo1Ea"}
-! CHECK-BBC: %[[V1:[0-9]+]]:2 = hlfir.declare %[[V0]] {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
-! CHECK-BBC: %[[V2:[0-9]+]] = hlfir.designate %[[V1]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
-! CHECK-BBC: %a25a1 = omp.map.info var_ptr(%[[V2]] : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) -> !fir.ref<i32> {name = "a%a1"}
-! CHECK-BBC: %[[V4:[0-9]+]] = omp.map.info var_ptr(%[[V1]]#1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>) map_clauses(tofrom) capture(ByRef) members(%a25a1 : [1] : !fir.ref<i32>) -> !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>> {name = "a", partial_map = true}
-! CHECK-BBC: omp.target map_entries(%[[V4]] -> %arg0, %a25a1 -> %a25a1_3 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<i32>) {
-! CHECK-BBC:   %[[V5:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
-! CHECK-BBC:   %c0_i32 = arith.constant 0 : i32
-! CHECK-BBC:   %[[V6:[0-9]+]] = hlfir.designate %[[V5]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
-! CHECK-BBC:   hlfir.assign %c0_i32 to %[[V6]] : i32, !fir.ref<i32>
-! CHECK-BBC:   omp.terminator
-! CHECK-BBC: }
+! CHECK: %[[V0:[0-9]+]] = fir.alloca !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}> {bindc_name = "a", uniq_name = "_QFfoo1Ea"}
+! CHECK: %[[V1:[0-9]+]]:2 = hlfir.declare %[[V0]] {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
+! CHECK: %[[V2:[0-9]+]] = hlfir.designate %[[V1]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
+! CHECK: %[[V3:[0-9]+]] = omp.map.info var_ptr(%[[V2]] : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) -> !fir.ref<i32> {name = "a%a1"}
+! CHECK: %[[V4:[0-9]+]] = omp.map.info var_ptr(%[[V1]]#1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>) map_clauses(tofrom) capture(ByRef) members(%[[V3]] : [1] : !fir.ref<i32>) -> !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>> {name = "a", partial_map = true}
+! CHECK: omp.target map_entries(%[[V4]] -> %arg0, %[[V3]] -> %arg1 : !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<i32>) {
+! CHECK:   %[[V5:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFfoo1Ea"} : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>, !fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>)
+! CHECK:   %c0_i32 = arith.constant 0 : i32
+! CHECK:   %[[V6:[0-9]+]] = hlfir.designate %[[V5]]#0{"a1"}   : (!fir.ref<!fir.type<_QFfoo1Tt0{a0:i32,a1:i32}>>) -> !fir.ref<i32>
+! CHECK:   hlfir.assign %c0_i32 to %[[V6]] : i32, !fir.ref<i32>
+! CHECK:   omp.terminator
+! CHECK: }
 
 subroutine foo1()
   implicit none
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index e06ab003b0f6b9..8f9b78ccc72380 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -198,7 +198,7 @@ struct AsmPrinterOptions {
                      "and naming conflicts across all regions")};
 
   llvm::cl::opt<bool> useNameLocAsPrefix{"mlir-use-nameloc-as-prefix",
-                                         llvm::cl::init(true),
+                                         llvm::cl::init(false),
                                          llvm::cl::desc("TODO")};
 };
 } // namespace
@@ -1628,13 +1628,12 @@ void SSANameState::numberValuesInOp(Operation &op) {
   }
 
   unsigned numResults = op.getNumResults();
-  if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames) {
-    if (numResults > 0) {
-      Value resultBegin = op.getResult(0);
-      if (isa<NameLoc>(resultBegin.getLoc())) {
-        auto nameLoc = cast<NameLoc>(resultBegin.getLoc());
-        setResultNameFn(resultBegin, nameLoc.getName());
-      }
+  if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames &&
+      numResults > 0) {
+    Value resultBegin = op.getResult(0);
+    if (isa<NameLoc>(resultBegin.getLoc())) {
+      auto nameLoc = cast<NameLoc>(resultBegin.getLoc());
+      setResultNameFn(resultBegin, nameLoc.getName());
     }
   }
 
diff --git a/mlir/test/IR/locations.mlir b/mlir/test/IR/locations.mlir
index dc39315a765a83..b725307b420b79 100644
--- a/mlir/test/IR/locations.mlir
+++ b/mlir/test/IR/locations.mlir
@@ -13,7 +13,7 @@ func.func @inline_notation() -> i32 {
   // CHECK: arith.constant 4 : index loc(callsite("foo" at "mysource.cc":10:8))
   %2 = arith.constant 4 : index loc(callsite("foo" at "mysource.cc":10:8))
 
-  // CHECK: affine.for %IVlocation loc("IVlocation") = 0 to 8 {
+  // CHECK: affine.for %arg0 loc("IVlocation") = 0 to 8 {
   // CHECK: } loc(fused["foo", "mysource.cc":10:8])
   affine.for %i0 loc("IVlocation") = 0 to 8 {
   } loc(fused["foo", "mysource.cc":10:8])
@@ -26,7 +26,7 @@ func.func @inline_notation() -> i32 {
   affine.if #set0(%2) {
   } loc(fused<"myPass">["foo"])
 
-  // CHECK: return %foo : i32 loc(unknown)
+  // CHECK: return %0 : i32 loc(unknown)
   return %1 : i32 loc(unknown)
 }
 
@@ -62,7 +62,7 @@ func.func @escape_strings() {
 // CHECK-LABEL: func @argLocs(
 // CHECK-SAME:  %arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+1]]:20),
 func.func @argLocs(%x: i32,
-// CHECK-SAME:  %hotdog: i64 loc("hotdog")
+// CHECK-SAME:  %arg1: i64 loc("hotdog")
               %y: i64 loc("hotdog")) {
   return
 }
@@ -73,11 +73,11 @@ func.func @argLocs(%x: i32,
 // CHECK-NEXT: ^bb0(%arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+2]]:7),
 // CHECK-ALIAS-NEXT: ^bb0(%arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+1]]:7),
  ^bb0(%x: i32,
-// CHECK-SAME: %cheetos: i32 loc("cheetos"),
-// CHECK-ALIAS-SAME: %cheetos: i32 loc("cheetos"),
+// CHECK-SAME: %arg1: i32 loc("cheetos"),
+// CHECK-ALIAS-SAME: %arg1: i32 loc("cheetos"),
       %y: i32 loc("cheetos"),
-// CHECK-SAME: %out_of_line_location2: i32 loc("out_of_line_location2")):
-// CHECK-ALIAS-SAME: %out_of_line_location2: i32 loc("out_of_line_location2")):
+// CHECK-SAME: %arg2: i32 loc("out_of_line_location2")):
+// CHECK-ALIAS-SAME: %arg2: i32 loc("out_of_line_location2")):
       %z: i32 loc("out_of_line_location2")):
     %1 = arith.addi %x, %y : i32
     "foo.yield"(%1) : (i32) -> ()
diff --git a/mlir/test/IR/pretty-locations.mlir b/mlir/test/IR/pretty-locations.mlir
index 30391252cabfbe..598bebeb83aebd 100644
--- a/mlir/test/IR/pretty-locations.mlir
+++ b/mlir/test/IR/pretty-locations.mlir
@@ -27,6 +27,6 @@ func.func @inline_notation() -> i32 {
   // CHECK: "foo.op"() : () -> () #test.custom_location<"foo.mlir" * 1234>
   "foo.op"() : () -> () loc(#test.custom_location<"foo.mlir" * 1234>)
 
-  // CHECK: return %foo : i32 [unknown]
+  // CHECK: return %0 : i32 [unknown]
   return %1 : i32 loc(unknown)
 }
diff --git a/mlir/test/IR/pretty_printed_region_op.mlir b/mlir/test/IR/pretty_printed_region_op.mlir
index 8c6d71d1bf10f2..0a81f9d53a3097 100644
--- a/mlir/test/IR/pretty_printed_region_op.mlir
+++ b/mlir/test/IR/pretty_printed_region_op.mlir
@@ -6,9 +6,9 @@
 
 func.func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
 // CHECK-CUSTOM:  test.pretty_printed_region %arg1, %arg0 start test.special.op end : (f32, f32) -> f32
-// CHECK-GENERIC: %some_NameLoc = "test.pretty_printed_region"(%arg1, %arg0)
-// CHECK-GENERIC:   ^bb0(%some_NameLoc_[[x:[0-9]+]]: f32, %some_NameLoc_[[y:[0-9]+]]: f32
-// CHECK-GENERIC:     %[[RES:.*]] = "test.special.op"(%some_NameLoc_[[x]], %some_NameLoc_[[y]]) : (f32, f32) -> f32
+// CHECK-GENERIC: "test.pretty_printed_region"(%arg1, %arg0)
+// CHECK-GENERIC:   ^bb0(%arg[[x:[0-9]+]]: f32, %arg[[y:[0-9]+]]: f32
+// CHECK-GENERIC:     %[[RES:.*]] = "test.special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
 // CHECK-GENERIC:     "test.return"(%[[RES]]) : (f32) -> ()
 // CHECK-GENERIC:  : (f32, f32) -> f32
 
@@ -37,9 +37,9 @@ func.func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
 // -----
 
 func.func @pretty_printed_region_op_deferred_loc(%arg0 : f32, %arg1 : f32) -> (f32) {
-// CHECK-LOCATION: %foo = "test.pretty_printed_region"(%arg1, %arg0)
-// CHECK-LOCATION:   ^bb0(%foo_[[x:[0-9]+]]: f32 loc("foo"), %foo_[[y:[0-9]+]]: f32 loc("foo")
-// CHECK-LOCATION:     %[[RES:.*]] = "test.special.op"(%foo_[[x]], %foo_[[y]]) : (f32, f32) -> f32
+// CHECK-LOCATION: "test.pretty_printed_region"(%arg1, %arg0)
+// CHECK-LOCATION:   ^bb0(%arg[[x:[0-9]+]]: f32 loc("foo"), %arg[[y:[0-9]+]]: f32 loc("foo")
+// CHECK-LOCATION:     %[[RES:.*]] = "test.special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
 // CHECK-LOCATION:     "test.return"(%[[RES]]) : (f32) -> ()
 // CHECK-LOCATION:  : (f32, f32) -> f32
 
diff --git a/mlir/test/IR/wrapping_op.mlir b/mlir/test/IR/wrapping_op.mlir
index 95623cda32ea0f..3f7a350ee654e7 100644
--- a/mlir/test/IR/wrapping_op.mlir
+++ b/mlir/test/IR/wrapping_op.mlir
@@ -5,7 +5,7 @@
 // CHECK-GENERIC: "func.func"
 // CHECK-GENERIC-SAME: sym_name = "wrapping_op"
 func.func @wrapping_op(%arg0 : i32, %arg1 : f32) -> (i3, i2, i1) {
-// CHECK: %some_NameLoc:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
+// CHECK: %0:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
 // CHECK-GENERIC: "test.wrapping_region"() ({
 // CHECK-GENERIC:   %[[NESTED_RES:.*]]:3 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
 // CHECK-GENERIC:   "test.return"(%[[NESTED_RES]]#0, %[[NESTED_RES]]#1, %[[NESTED_RES]]#2) : (i1, i2, i3) -> () loc("some_NameLoc")
diff --git a/mlir/test/mlir-tblgen/pattern.mlir b/mlir/test/mlir-tblgen/pattern.mlir
index 5594e80420a722..60d46e676d2a33 100644
--- a/mlir/test/mlir-tblgen/pattern.mlir
+++ b/mlir/test/mlir-tblgen/pattern.mlir
@@ -678,7 +678,7 @@ func.func @returnTypeAndLocation(%arg0 : i32) -> i1 {
   %0 = "test.source_op"(%arg0) {tag = 66 : i32} : (i32) -> i1
   // CHECK: "test.op_x"(%arg0) : (i32) -> i32 loc("loc1")
   // CHECK: "test.op_x"(%arg0) : (i32) -> i32 loc("loc2")
-  // CHECK: "test.two_to_one"(%loc1, %loc2) : (i32, i32) -> i1
+  // CHECK: "test.two_to_one"(%0, %1) : (i32, i32) -> i1
   return %0 : i1
 }
 

>From c7ebcedb8d6a0d615c3f2afe987cd6b305fe8ac3 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Sun, 15 Dec 2024 15:31:16 -0500
Subject: [PATCH 5/5] add test

---
 mlir/include/mlir/IR/OperationSupport.h       |   6 +-
 mlir/lib/IR/AsmPrinter.cpp                    |  54 +++++----
 mlir/test/IR/print-use-nameloc-as-prefix.mlir | 105 ++++++++++++++++++
 3 files changed, 134 insertions(+), 31 deletions(-)
 create mode 100644 mlir/test/IR/print-use-nameloc-as-prefix.mlir

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index c6a603ae518785..9f2de582b03e56 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1221,8 +1221,8 @@ class OpPrintingFlags {
   /// Return if printer should use unique SSA IDs.
   bool shouldPrintUniqueSSAIDs() const;
 
-  /// Returns if the printer should retain identifier names collected using
-  /// parsing.
+  /// Return if the printer should use NameLocs as prefixes when printing SSA
+  /// IDs
   bool shouldUseNameLocAsPrefix() const;
 
 private:
@@ -1259,7 +1259,7 @@ class OpPrintingFlags {
   /// Print unique SSA IDs for values, block arguments and naming conflicts
   bool printUniqueSSAIDsFlag : 1;
 
-  /// Print the retained original names of identifiers
+  /// Print SSA IDs using NameLocs as prefixes
   bool useNameLocAsPrefix : 1;
 };
 
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 8f9b78ccc72380..d1b907e6cbbab1 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -197,9 +197,9 @@ struct AsmPrinterOptions {
       llvm::cl::desc("Print unique SSA ID numbers for values, block arguments "
                      "and naming conflicts across all regions")};
 
-  llvm::cl::opt<bool> useNameLocAsPrefix{"mlir-use-nameloc-as-prefix",
-                                         llvm::cl::init(false),
-                                         llvm::cl::desc("TODO")};
+  llvm::cl::opt<bool> useNameLocAsPrefix{
+      "mlir-use-nameloc-as-prefix", llvm::cl::init(false),
+      llvm::cl::desc("Print SSA IDs using NameLocs as prefixes")};
 };
 } // namespace
 
@@ -369,7 +369,7 @@ bool OpPrintingFlags::shouldPrintUniqueSSAIDs() const {
   return printUniqueSSAIDsFlag || shouldPrintGenericOpForm();
 }
 
-/// TODO
+/// Return if the printer should use NameLocs as prefixes when printing SSA IDs
 bool OpPrintingFlags::shouldUseNameLocAsPrefix() const {
   return useNameLocAsPrefix;
 }
@@ -1523,25 +1523,18 @@ void SSANameState::numberValuesInRegion(Region &region) {
     assert(!valueIDs.count(arg) && "arg numbered multiple times");
     assert(llvm::cast<BlockArgument>(arg).getOwner()->getParent() == &region &&
            "arg not defined in current region");
-    setValueName(arg, name);
+    if (printerFlags.shouldUseNameLocAsPrefix() && isa<NameLoc>(arg.getLoc())) {
+      auto nameLoc = cast<NameLoc>(arg.getLoc());
+      setValueName(arg, nameLoc.getName());
+    } else {
+      setValueName(arg, name);
+    }
   };
 
-  bool alreadySetNames = false;
   if (!printerFlags.shouldPrintGenericOpForm()) {
     if (Operation *op = region.getParentOp()) {
-      if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op)) {
+      if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op))
         asmInterface.getAsmBlockArgumentNames(region, setBlockArgNameFn);
-        alreadySetNames = true;
-      }
-    }
-  }
-
-  if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames) {
-    for (BlockArgument arg : region.getArguments()) {
-      if (isa<NameLoc>(arg.getLoc())) {
-        auto nameLoc = cast<NameLoc>(arg.getLoc());
-        setBlockArgNameFn(arg, nameLoc.getName());
-      }
     }
   }
 
@@ -1596,7 +1589,13 @@ void SSANameState::numberValuesInOp(Operation &op) {
   auto setResultNameFn = [&](Value result, StringRef name) {
     assert(!valueIDs.count(result) && "result numbered multiple times");
     assert(result.getDefiningOp() == &op && "result not defined by 'op'");
-    setValueName(result, name);
+    if (printerFlags.shouldUseNameLocAsPrefix() &&
+        isa<NameLoc>(result.getLoc())) {
+      auto nameLoc = cast<NameLoc>(result.getLoc());
+      setValueName(result, nameLoc.getName());
+    } else {
+      setValueName(result, name);
+    }
 
     // Record the result number for groups not anchored at 0.
     if (int resultNo = llvm::cast<OpResult>(result).getResultNumber())
@@ -1628,15 +1627,6 @@ void SSANameState::numberValuesInOp(Operation &op) {
   }
 
   unsigned numResults = op.getNumResults();
-  if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames &&
-      numResults > 0) {
-    Value resultBegin = op.getResult(0);
-    if (isa<NameLoc>(resultBegin.getLoc())) {
-      auto nameLoc = cast<NameLoc>(resultBegin.getLoc());
-      setResultNameFn(resultBegin, nameLoc.getName());
-    }
-  }
-
   if (numResults == 0) {
     // If value users should be printed, operations with no result need an id.
     if (printerFlags.shouldPrintValueUsers()) {
@@ -1645,8 +1635,16 @@ void SSANameState::numberValuesInOp(Operation &op) {
     }
     return;
   }
+
   Value resultBegin = op.getResult(0);
 
+  if (printerFlags.shouldUseNameLocAsPrefix() && !alreadySetNames) {
+    if (isa<NameLoc>(resultBegin.getLoc())) {
+      auto nameLoc = cast<NameLoc>(resultBegin.getLoc());
+      setResultNameFn(resultBegin, nameLoc.getName());
+    }
+  }
+
   // If the first result wasn't numbered, give it a default number.
   if (valueIDs.try_emplace(resultBegin, nextValueID).second)
     ++nextValueID;
diff --git a/mlir/test/IR/print-use-nameloc-as-prefix.mlir b/mlir/test/IR/print-use-nameloc-as-prefix.mlir
new file mode 100644
index 00000000000000..fb555d9708ee86
--- /dev/null
+++ b/mlir/test/IR/print-use-nameloc-as-prefix.mlir
@@ -0,0 +1,105 @@
+// RUN: mlir-opt %s -mlir-use-nameloc-as-prefix -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2' -mlir-use-nameloc-as-prefix -split-input-file | FileCheck %s --check-prefix=CHECK-PASS-PRESERVE
+
+// CHECK-LABEL: test_basic
+func.func @test_basic() {
+  %0 = memref.alloc() : memref<i32>
+  // CHECK: %alice
+  %1 = memref.load %0[] : memref<i32> loc("alice")
+  return
+}
+
+// -----
+
+// CHECK-LABEL: test_repeat_namelocs
+func.func @test_repeat_namelocs() {
+  %0 = memref.alloc() : memref<i32>
+  // CHECK: %alice
+  %1 = memref.load %0[] : memref<i32> loc("alice")
+  // CHECK: %alice_0
+  %2 = memref.load %0[] : memref<i32> loc("alice")
+  return
+}
+
+// -----
+
+// CHECK-LABEL: test_bb_args
+func.func @test_bb_args1(%arg0 : memref<i32> loc("foo")) {
+  // CHECK: %alice = memref.load %foo
+  %1 = memref.load %arg0[] : memref<i32> loc("alice")
+  return
+}
+
+// -----
+
+func.func private @make_two_results() -> (index, index)
+
+// CHECK-LABEL: test_multiple_results
+func.func @test_multiple_results(%cond: i1) {
+  // CHECK: %foo:2
+  %0:2 = call @make_two_results() : () -> (index, index) loc("foo")
+  // CHECK: %bar:2
+  %1, %2 = call @make_two_results() : () -> (index, index) loc("bar")
+
+  // CHECK: %kevin:2 = scf.while (%arg1 = %bar#0, %arg2 = %bar#0)
+  %5:2 = scf.while (%arg1 = %1, %arg2 = %1) : (index, index) -> (index, index) {
+    %6 = arith.cmpi slt, %arg1, %arg2 : index
+    scf.condition(%6) %arg1, %arg2 : index, index
+  } do {
+  // CHECK: ^bb0(%alice: index, %bob: index)
+  ^bb0(%arg3 : index loc("alice"), %arg4: index loc("bob")):
+    %c1, %c2 = func.call @make_two_results() : () -> (index, index) loc("harriet")
+    // CHECK: scf.yield %harriet#1, %harriet#1
+    scf.yield %c2, %c2 : index, index
+  } loc("kevin")
+  return
+}
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#trait = {
+  iterator_types = ["parallel"],
+  indexing_maps = [#map, #map, #map]
+}
+
+// CHECK-LABEL: test_op_asm_interface
+func.func @test_op_asm_interface(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {
+  // CHECK: %c0
+  %0 = arith.constant 0 : index
+  // CHECK: %foo
+  %1 = arith.constant 1 : index loc("foo")
+
+  linalg.generic #trait ins(%arg0: tensor<?xf32>) outs(%arg0, %arg1: tensor<?xf32>, tensor<?xf32>) {
+    // CHECK: ^bb0(%in: f32, %out: f32, %out_0: f32)
+    ^bb0(%a: f32, %b: f32, %c: f32):
+      linalg.yield %a, %a : f32, f32
+  } -> (tensor<?xf32>, tensor<?xf32>)
+
+  linalg.generic #trait ins(%arg0: tensor<?xf32>) outs(%arg0, %arg1: tensor<?xf32>, tensor<?xf32>) {
+    // CHECK: ^bb0(%bar: f32, %alice: f32, %steve: f32)
+    ^bb0(%a: f32 loc("bar"), %b: f32 loc("alice"), %c: f32 loc("steve")):
+      // CHECK: linalg.yield %alice, %steve
+      linalg.yield %b, %c : f32, f32
+  } -> (tensor<?xf32>, tensor<?xf32>)
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: test_pass
+func.func @test_pass(%arg0: memref<4xf32>, %arg1: memref<4xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  scf.for %arg2 = %c0 to %c4 step %c1 {
+    // CHECK-PASS-PRESERVE: %foo = memref.load
+    // CHECK-PASS-PRESERVE: memref.store %foo
+    // CHECK-PASS-PRESERVE: %foo_1 = memref.load
+    // CHECK-PASS-PRESERVE: memref.store %foo_1
+    %0 = memref.load %arg0[%arg2] : memref<4xf32> loc("foo")
+    memref.store %0, %arg1[%arg2] : memref<4xf32>
+  }
+  return
+}



More information about the Mlir-commits mailing list