[Mlir-commits] [mlir] [MLIR][IR] add -mlir-print-unique-ssa-ids to AsmPrinter (PR #91241)

[Scott Manley]

https://github.com/rscottmanley updated https://github.com/llvm/llvm-project/pull/91241

>From 11fe9873be6678a1dba58bfd78f1577446b64c77 Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Mon, 6 May 2024 09:41:46 -0700
Subject: [PATCH] [MLIR][IR] add -mlir-print-unique-ssa-ids to AsmPrinter

Add an option to print unique SSA IDs of values, block arguments and naming
conflicts when requested and/or printing generic op form. This is helpful
when debugging. For example, if you have:

    scf.for
      %0 =
      %1 = opA %0

    scf.for
      %0 =
      %1 = opB %0

And you get a verifier error which says opB's "operand #0 does not dominate
this use", it looks like %0 does dominate the use. This is not intuitive. If
these were numbered uniquely, it would look like:

    scf.for
      %0 =
      %1 = opA %0

    scf.for
      %2 =
      %3 = opB %0

And thus, much clearer as to why you are getting the error since %0 is out
of scope. Since generic op form should aim to give you the most possible
information, it seems like a good idea to use unique numbers in this
situation. Adding an option also gives those an option to use it outside
of generic op form.
---
 mlir/include/mlir/IR/OperationSupport.h |  6 +++++
 mlir/lib/IR/AsmPrinter.cpp              | 23 ++++++++++++++++---
 mlir/test/IR/print-unique-ssa-ids.mlir  | 30 +++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 mlir/test/IR/print-unique-ssa-ids.mlir

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index e661bb87a27ed0..f8ab5338107f2a 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1219,6 +1219,9 @@ class OpPrintingFlags {
   /// Return if the printer should print users of values.
   bool shouldPrintValueUsers() const;
 
+  /// Return if printer should use unique SSA IDs.
+  bool shouldPrintUniqueSSAIDs() const;
+
 private:
   /// Elide large elements attributes if the number of elements is larger than
   /// the upper limit.
@@ -1249,6 +1252,9 @@ class OpPrintingFlags {
 
   /// Print users of values.
   bool printValueUsersFlag : 1;
+
+  /// Print unique SSA IDs for values, block arguments and naming conflicts
+  bool printUniqueSSAIDsFlag : 1;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index e915b97d9ff17b..9a5c51ba738f9e 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -189,6 +189,11 @@ struct AsmPrinterOptions {
       "mlir-print-value-users", llvm::cl::init(false),
       llvm::cl::desc(
           "Print users of operation results and block arguments as a comment")};
+
+  llvm::cl::opt<bool> printUniqueSSAIDs{
+      "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")};
 };
 } // namespace
 
@@ -206,7 +211,7 @@ OpPrintingFlags::OpPrintingFlags()
     : printDebugInfoFlag(false), printDebugInfoPrettyFormFlag(false),
       printGenericOpFormFlag(false), skipRegionsFlag(false),
       assumeVerifiedFlag(false), printLocalScope(false),
-      printValueUsersFlag(false) {
+      printValueUsersFlag(false), printUniqueSSAIDsFlag(false) {
   // Initialize based upon command line options, if they are available.
   if (!clOptions.isConstructed())
     return;
@@ -224,6 +229,7 @@ OpPrintingFlags::OpPrintingFlags()
   printLocalScope = clOptions->printLocalScopeOpt;
   skipRegionsFlag = clOptions->skipRegionsOpt;
   printValueUsersFlag = clOptions->printValueUsers;
+  printUniqueSSAIDsFlag = clOptions->printUniqueSSAIDs;
 }
 
 /// Enable the elision of large elements attributes, by printing a '...'
@@ -350,6 +356,11 @@ bool OpPrintingFlags::shouldPrintValueUsers() const {
   return printValueUsersFlag;
 }
 
+/// Return if the printer should use unique IDs.
+bool OpPrintingFlags::shouldPrintUniqueSSAIDs() const {
+  return printUniqueSSAIDsFlag || shouldPrintGenericOpForm();
+}
+
 //===----------------------------------------------------------------------===//
 // NewLineCounter
 //===----------------------------------------------------------------------===//
@@ -1369,8 +1380,14 @@ SSANameState::SSANameState(Operation *op, const OpPrintingFlags &printerFlags)
   while (!nameContext.empty()) {
     Region *region;
     UsedNamesScopeTy *parentScope;
-    std::tie(region, nextValueID, nextArgumentID, nextConflictID, parentScope) =
-        nameContext.pop_back_val();
+
+    if (printerFlags.shouldPrintUniqueSSAIDs())
+      // To print unique SSA IDs, ignore saved ID counts from parent regions
+      std::tie(region, std::ignore, std::ignore, std::ignore, parentScope) =
+          nameContext.pop_back_val();
+    else
+      std::tie(region, nextValueID, nextArgumentID, nextConflictID,
+               parentScope) = nameContext.pop_back_val();
 
     // When we switch from one subtree to another, pop the scopes(needless)
     // until the parent scope.
diff --git a/mlir/test/IR/print-unique-ssa-ids.mlir b/mlir/test/IR/print-unique-ssa-ids.mlir
new file mode 100644
index 00000000000000..a2d2d9bb79076c
--- /dev/null
+++ b/mlir/test/IR/print-unique-ssa-ids.mlir
@@ -0,0 +1,30 @@
+// RUN: mlir-opt -mlir-print-unique-ssa-ids %s | FileCheck %s
+// RUN: mlir-opt -mlir-print-op-generic %s | FileCheck %s
+// RUN: mlir-opt %s | FileCheck %s --check-prefix=LOCAL_SCOPE
+
+// CHECK: %arg3
+// CHECK: %7
+// LOCAL_SCOPE-NOT: %arg3
+// LOCAL_SCOPE-NOT: %7
+module {
+  func.func @uniqueSSAIDs(%arg0 : memref<i32>, %arg1 : memref<i32>) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c8 = arith.constant 8 : index
+    scf.for %arg2 = %c0 to %c8 step %c1 {
+      %a = memref.load %arg0[] : memref<i32>
+      %b = memref.load %arg1[] : memref<i32>
+      %0 = arith.addi %a, %b : i32
+      %1 = arith.subi %a, %b : i32
+      scf.yield
+    }
+    scf.for %arg2 = %c0 to %c8 step %c1 {
+      %a = memref.load %arg0[] : memref<i32>
+      %b = memref.load %arg1[] : memref<i32>
+      %0 = arith.addi %a, %b : i32
+      %1 = arith.subi %a, %b : i32
+      scf.yield
+    }
+    return
+  }
+}