[Mlir-commits] [mlir] 89ecd8c - Teach the MLIR AsmPrinter to correctly escape asm names that use invalid characters.

[Chris Lattner]

Author: Chris Lattner
Date: 2020-03-12T22:36:41-07:00
New Revision: 89ecd8c149e482cca814d56113e141457fa90677

URL: https://github.com/llvm/llvm-project/commit/89ecd8c149e482cca814d56113e141457fa90677
DIFF: https://github.com/llvm/llvm-project/commit/89ecd8c149e482cca814d56113e141457fa90677.diff

LOG: Teach the MLIR AsmPrinter to correctly escape asm names that use invalid characters.

Reviewers: rriddle!

Subscribers: mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D75919

Added: 
    

Modified: 
    mlir/lib/IR/AsmPrinter.cpp
    mlir/test/IR/pretty-region-args.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 776792ec2f17..b51760044294 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -760,7 +760,43 @@ void SSANameState::setValueName(Value value, StringRef name) {
   valueNames[value] = uniqueValueName(name);
 }
 
+// Returns true if 'c' is an allowable punctuation character: [$._-]
+// Returns false otherwise.
+static bool isPunct(char c) {
+  return c == '$' || c == '.' || c == '_' || c == '-';
+}
+  
 StringRef SSANameState::uniqueValueName(StringRef name) {
+  assert(!name.empty() && "Shouldn't have an empty name here");
+  
+  // Check to see if this name is valid.  If it starts with a digit, then it
+  // could conflict with the autogenerated numeric ID's (we unique them in a
+  // 
diff erent map), so add an underscore prefix to avoid problems.
+  if (isdigit(name[0])) {
+    SmallString<16> tmpName("_");
+    tmpName += name;
+    return uniqueValueName(tmpName);
+  }
+  
+  // Check to see if the name consists of all-valid identifiers.  If not, we
+  // need to escape them.
+  for (auto ch : name) {
+    if (isalpha(ch) || isPunct(ch) || isdigit(ch))
+      continue;
+    
+    SmallString<16> tmpName;
+    for (auto ch : name) {
+      if (isalpha(ch) || isPunct(ch) || isdigit(ch))
+        tmpName += ch;
+      else if (ch == ' ')
+        tmpName += '_';
+      else {
+        tmpName += llvm::utohexstr((unsigned char)ch);
+      }
+    }
+    return uniqueValueName(tmpName);
+  }
+  
   // Check to see if this name is already unique.
   if (!usedNames.count(name)) {
     name = name.copy(usedNameAllocator);

diff  --git a/mlir/test/IR/pretty-region-args.mlir b/mlir/test/IR/pretty-region-args.mlir
index 59a9ebce092a..6e980fa2c3fb 100644
--- a/mlir/test/IR/pretty-region-args.mlir
+++ b/mlir/test/IR/pretty-region-args.mlir
@@ -10,3 +10,17 @@ func @custom_region_names() -> () {
   // CHECK-NEXT: ^bb{{.*}}(%i: index, %j: index, %k: index):
   return
 }
+
+// CHECK-LABEL: func @weird_names
+// Make sure the asmprinter handles weird names correctly.
+func @weird_names() -> () {
+  "test.polyfor"() ( {
+  ^bb0(%arg0: i32, %arg1: i32, %arg2: index):
+    "foo"() : () -> i32
+  }) { arg_names = ["a .^x", "0"] } : () -> ()
+  // CHECK: test.polyfor
+  // CHECK-NEXT: ^bb{{.*}}(%a_.5Ex: i32, %_0: i32, %arg0: index):
+  // CHECK-NEXT: %0 = "foo"()
+  return
+}
+