Added system to handle when we use default names

Author: Wheest

mlir/lib/IR/AsmPrinter.cpp

@@ -981,7 +981,8 @@ class DummyAliasDialectAsmPrinter : public DialectAsmPrinter {
 /// store the new copy,
 static StringRef sanitizeIdentifier(StringRef name, SmallString<16> &buffer,
                                     StringRef allowedPunctChars = "$._-",
-                                    bool allowTrailingDigit = true) {
+                                    bool allowTrailingDigit = true,
+                                    bool allowNumeric = false) {
   assert(!name.empty() && "Shouldn't have an empty name here");
 
   auto copyNameToBuffer = [&] {
@@ -997,16 +998,17 @@ static StringRef sanitizeIdentifier(StringRef name, SmallString<16> &buffer,
 
   // Check to see if this name is valid. If it starts with a digit, then it
   // could conflict with the autogenerated numeric ID's, so add an underscore
-  // prefix to avoid problems.
-  if (isdigit(name[0])) {
+  // prefix to avoid problems. This can be overridden by setting allowNumeric.
+  if (isdigit(name[0]) && !allowNumeric) {
     buffer.push_back('_');
     copyNameToBuffer();
     return buffer;
   }
 
   // If the name ends with a trailing digit, add a '_' to avoid potential
-  // conflicts with autogenerated ID's.
-  if (!allowTrailingDigit && isdigit(name.back())) {
+  // conflicts with autogenerated ID's. This can be overridden by setting
+  // allowNumeric.
+  if (!allowTrailingDigit && isdigit(name.back()) && !allowNumeric) {
     copyNameToBuffer();
     buffer.push_back('_');
     return buffer;
@@ -1292,11 +1294,11 @@ class SSANameState {
                             std::optional<int> &lookupResultNo) const;
 
   /// Set a special value name for the given value.
-  void setValueName(Value value, StringRef name);
+  void setValueName(Value value, StringRef name, bool allowNumeric = false);
 
   /// Uniques the given value name within the printer. If the given name
   /// conflicts, it is automatically renamed.
-  StringRef uniqueValueName(StringRef name);
+  StringRef uniqueValueName(StringRef name, bool allowNumeric = false);
 
   /// Set the original identifier names if available. Used in debugging with
   /// `--retain-identifier-names`/`shouldRetainIdentifierNames` in ParserConfig
@@ -1541,7 +1543,10 @@ void SSANameState::numberValuesInOp(Operation &op) {
   // Function used to set the special result names for the operation.
   SmallVector<int, 2> resultGroups(/*Size=*/1, /*Value=*/0);
   auto setResultNameFn = [&](Value result, StringRef name) {
-    assert(!valueIDs.count(result) && "result numbered multiple times");
+    // Case where the result has already been named
+    if (valueIDs.count(result))
+      return;
+    // assert(!valueIDs.count(result) && "result numbered multiple times");
     assert(result.getDefiningOp() == &op && "result not defined by 'op'");
     setValueName(result, name);
 
@@ -1565,17 +1570,17 @@ void SSANameState::numberValuesInOp(Operation &op) {
     blockNames[block] = {-1, name};
   };
 
+  // Set the original identifier names if available. Used in debugging with
+  // `--retain-identifier-names`/`shouldRetainIdentifierNames` in ParserConfig
+  setRetainedIdentifierNames(op, resultGroups);
+
   if (!printerFlags.shouldPrintGenericOpForm()) {
     if (OpAsmOpInterface asmInterface = dyn_cast<OpAsmOpInterface>(&op)) {
       asmInterface.getAsmBlockNames(setBlockNameFn);
       asmInterface.getAsmResultNames(setResultNameFn);
     }
   }
 
-  // Set the original identifier names if available. Used in debugging with
-  // `--retain-identifier-names`/`shouldRetainIdentifierNames` in ParserConfig
-  setRetainedIdentifierNames(op, resultGroups);
-
   unsigned numResults = op.getNumResults();
   if (numResults == 0) {
     // If value users should be printed, operations with no result need an id.
@@ -1610,7 +1615,7 @@ void SSANameState::setRetainedIdentifierNames(
       auto resultName = resultNames[i].cast<StringAttr>().strref();
       if (!resultName.empty()) {
         if (!usedNames.count(resultName))
-          setValueName(results[i], resultName);
+          setValueName(results[i], resultName, /*allowNumeric=*/true);
         // If a result has a name, it is the start of a result group.
         if (i > 0)
           resultGroups.push_back(i);
@@ -1628,7 +1633,7 @@ void SSANameState::setRetainedIdentifierNames(
     for (size_t i = 0; i < opArgs.size() && i < opArgNames.size(); ++i) {
       auto opArgName = opArgNames[i].cast<StringAttr>().strref();
       if (!usedNames.count(opArgName))
-        setValueName(opArgs[i], opArgName);
+        setValueName(opArgs[i], opArgName, /*allowNumeric=*/true);
     }
     op.removeDiscardableAttr("mlir.opArgNames");
   }
@@ -1649,7 +1654,7 @@ void SSANameState::setRetainedIdentifierNames(
     for (size_t i = 0; i < blockArgs.size() && i < blockArgNames.size(); ++i) {
       auto blockArgName = blockArgNames[i].cast<StringAttr>().strref();
       if (!usedNames.count(blockArgName))
-        setValueName(blockArgs[i], blockArgName);
+        setValueName(blockArgs[i], blockArgName, /*allowNumeric=*/true);
     }
     op.removeDiscardableAttr("mlir.blockArgNames");
   }
@@ -1695,20 +1700,22 @@ void SSANameState::getResultIDAndNumber(
   lookupValue = owner->getResult(groupResultNo);
 }
 
-void SSANameState::setValueName(Value value, StringRef name) {
+void SSANameState::setValueName(Value value, StringRef name,
+                                bool allowNumeric) {
   // If the name is empty, the value uses the default numbering.
   if (name.empty()) {
     valueIDs[value] = nextValueID++;
     return;
   }
 
   valueIDs[value] = NameSentinel;
-  valueNames[value] = uniqueValueName(name);
+  valueNames[value] = uniqueValueName(name, allowNumeric);
 }
 
-StringRef SSANameState::uniqueValueName(StringRef name) {
+StringRef SSANameState::uniqueValueName(StringRef name, bool allowNumeric) {
   SmallString<16> tmpBuffer;
-  name = sanitizeIdentifier(name, tmpBuffer);
+  name = sanitizeIdentifier(name, tmpBuffer, /*allowedPunctChars=*/"$._-",
+                            /*allowTrailingDigit=*/true, allowNumeric);
 
   // Check to see if this name is already unique.
   if (!usedNames.count(name)) {