From 19941b0468c746a8d7ea543600fb491fe9808b04 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov@fb.com>
Date: Mon, 8 May 2023 18:48:40 -0700
Subject: [PATCH] [BOLT] Use MCInstPrinter in createRetpolineFunctionTag

Make retpoline functions invariant of X86 register numbers.
retpoline-synthetic.test is known to fail NFC testing due to shifting
register numbers. Use canonical register names instead of tablegen
numbers.

Before:
```
__retpoline_r51_
__retpoline_mem_r58+DATAat0x200fe8
__retpoline_mem_r51+0
__retpoline_mem_r132+0+8*53
```

After:
```
__retpoline_%rax_
__retpoline_mem_%rip+DATAat0x200fe8
__retpoline_mem_%rax+0
__retpoline_mem_%r12+0+8*%rbx
```

Test Plan:
- Revert 67bd3c58c0c7389e39c5a2f4d3b1a30459ccf5b7 that touches X86RegisterInfo.td.
- retpoline-synthetic.test passes in NFC mode with this diff, fails without it.

Reviewed By: #bolt, rafauler

Differential Revision: https://reviews.llvm.org/D150138
---
 bolt/lib/Passes/RetpolineInsertion.cpp        | 53 +++++++++++--------
 bolt/test/CMakeLists.txt                      |  1 +
 bolt/test/lit.cfg.py                          |  1 +
 .../test/runtime/X86/retpoline-synthetic.test |  7 +++
 4 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/bolt/lib/Passes/RetpolineInsertion.cpp b/bolt/lib/Passes/RetpolineInsertion.cpp
index 4d83fb3903a7..c8a42725c974 100644
--- a/bolt/lib/Passes/RetpolineInsertion.cpp
+++ b/bolt/lib/Passes/RetpolineInsertion.cpp
@@ -22,6 +22,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/RetpolineInsertion.h"
+#include "llvm/MC/MCInstPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "bolt-retpoline"
@@ -173,39 +174,45 @@ BinaryFunction *createNewRetpoline(BinaryContext &BC,
 std::string createRetpolineFunctionTag(BinaryContext &BC,
                                        const IndirectBranchInfo &BrInfo,
                                        bool R11Available) {
-  if (BrInfo.isReg())
-    return "__retpoline_r" + to_string(BrInfo.BranchReg) + "_";
+  std::string Tag;
+  llvm::raw_string_ostream TagOS(Tag);
+  TagOS << "__retpoline_";
+
+  if (BrInfo.isReg()) {
+    BC.InstPrinter->printRegName(TagOS, BrInfo.BranchReg);
+    TagOS << "_";
+    TagOS.flush();
+    return Tag;
+  }
 
   // Memory Branch
   if (R11Available)
     return "__retpoline_r11";
 
-  std::string Tag = "__retpoline_mem_";
-
   const IndirectBranchInfo::MemOpInfo &MemRef = BrInfo.Memory;
 
-  std::string DispExprStr;
-  if (MemRef.DispExpr) {
-    llvm::raw_string_ostream Ostream(DispExprStr);
-    MemRef.DispExpr->print(Ostream, BC.AsmInfo.get());
-    Ostream.flush();
+  TagOS << "mem_";
+
+  if (MemRef.BaseRegNum != BC.MIB->getNoRegister())
+    BC.InstPrinter->printRegName(TagOS, MemRef.BaseRegNum);
+
+  TagOS << "+";
+  if (MemRef.DispExpr)
+    MemRef.DispExpr->print(TagOS, BC.AsmInfo.get());
+  else
+    TagOS << MemRef.DispImm;
+
+  if (MemRef.IndexRegNum != BC.MIB->getNoRegister()) {
+    TagOS << "+" << MemRef.ScaleImm << "*";
+    BC.InstPrinter->printRegName(TagOS, MemRef.IndexRegNum);
   }
 
-  Tag += MemRef.BaseRegNum != BC.MIB->getNoRegister()
-             ? "r" + to_string(MemRef.BaseRegNum)
-             : "";
-
-  Tag += MemRef.DispExpr ? "+" + DispExprStr : "+" + to_string(MemRef.DispImm);
-
-  Tag += MemRef.IndexRegNum != BC.MIB->getNoRegister()
-             ? "+" + to_string(MemRef.ScaleImm) + "*" +
-                   to_string(MemRef.IndexRegNum)
-             : "";
-
-  Tag += MemRef.SegRegNum != BC.MIB->getNoRegister()
-             ? "_seg_" + to_string(MemRef.SegRegNum)
-             : "";
+  if (MemRef.SegRegNum != BC.MIB->getNoRegister()) {
+    TagOS << "_seg_";
+    BC.InstPrinter->printRegName(TagOS, MemRef.SegRegNum);
+  }
 
+  TagOS.flush();
   return Tag;
 }
 
diff --git a/bolt/test/CMakeLists.txt b/bolt/test/CMakeLists.txt
index 1c7814155bda..216a785b7d69 100644
--- a/bolt/test/CMakeLists.txt
+++ b/bolt/test/CMakeLists.txt
@@ -47,6 +47,7 @@ list(APPEND BOLT_TEST_DEPS
   llvm-objdump
   llvm-readelf
   llvm-readobj
+  llvm-strings
   llvm-strip
   llvm-objcopy
   merge-fdata
diff --git a/bolt/test/lit.cfg.py b/bolt/test/lit.cfg.py
index 5838699157db..f28afd591227 100644
--- a/bolt/test/lit.cfg.py
+++ b/bolt/test/lit.cfg.py
@@ -87,6 +87,7 @@ tools = [
     ToolSubst('llvm-nm', unresolved='fatal'),
     ToolSubst('llvm-objdump', unresolved='fatal'),
     ToolSubst('llvm-objcopy', unresolved='fatal'),
+    ToolSubst('llvm-strings', unresolved='fatal'),
     ToolSubst('llvm-strip', unresolved='fatal'),
     ToolSubst('llvm-readelf', unresolved='fatal'),
     ToolSubst('link_fdata', command=sys.executable, unresolved='fatal', extra_args=[link_fdata_cmd]),
diff --git a/bolt/test/runtime/X86/retpoline-synthetic.test b/bolt/test/runtime/X86/retpoline-synthetic.test
index 6cae6a85fbe7..46daf3fdacde 100644
--- a/bolt/test/runtime/X86/retpoline-synthetic.test
+++ b/bolt/test/runtime/X86/retpoline-synthetic.test
@@ -20,5 +20,12 @@ CHECK-CALL-NOT: callq  *
 RUN: llvm-objdump -d -j ".text" %t | FileCheck %s -check-prefix=CHECK-JUMP
 CHECK-JUMP-NOT: jmpq   *
 
+# Check generated retpoline stub names
+RUN: llvm-strings %t | FileCheck %s -check-prefix=CHECK-STRINGS
+CHECK-STRINGS-DAG: __retpoline_%rax_
+CHECK-STRINGS-DAG: __retpoline_mem_%rip+DATAat0x[[#]]
+CHECK-STRINGS-DAG: __retpoline_mem_%rax+0
+CHECK-STRINGS-DAG: __retpoline_mem_%r[[#]]+0+8*%rbx
+
 RUN: %t 1000 3 | FileCheck %s
 CHECK: 30000000