From 3cc4fc267b35b300ea86331a48d1be5ff417a37d Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Fri, 22 Nov 2019 14:53:20 -0800
Subject: [PATCH] [BOLT] Proper support for -trap-avx512 option

Summary:
If -trap-avx512 option is not set, verify that we correctly encode
AVX-512 instructions and treat them as ordinary instructions.

(cherry picked from FBD18666427)
---
 bolt/src/BinaryContext.cpp       | 18 +++++++++++++++
 bolt/src/BinaryContext.h         |  5 ++++
 bolt/src/BinaryFunction.cpp      | 39 ++++++++++++++++----------------
 bolt/src/Passes/BinaryPasses.cpp | 10 ++++----
 4 files changed, 48 insertions(+), 24 deletions(-)

diff --git a/bolt/src/BinaryContext.cpp b/bolt/src/BinaryContext.cpp
index 7b5e9524e5df..18277d7d2518 100644
--- a/bolt/src/BinaryContext.cpp
+++ b/bolt/src/BinaryContext.cpp
@@ -1863,6 +1863,24 @@ BinaryContext::calculateEmittedSize(BinaryFunction &BF, bool FixBranches) {
   return std::make_pair(HotSize, ColdSize);
 }
 
+bool BinaryContext::validateEncoding(const MCInst &Inst,
+                                     ArrayRef<uint8_t> InputEncoding) const {
+  SmallString<256> Code;
+  SmallVector<MCFixup, 4> Fixups;
+  raw_svector_ostream VecOS(Code);
+
+  MCE->encodeInstruction(Inst, VecOS, Fixups, *STI);
+  auto EncodedData = ArrayRef<uint8_t>((uint8_t *)Code.data(), Code.size());
+  if (InputEncoding != EncodedData) {
+    errs() << "BOLT-ERROR: mismatched encoding detected\n"
+           << "      input: " << InputEncoding << '\n'
+           << "     output: " << EncodedData << '\n';
+    return false;
+  }
+
+  return true;
+}
+
 BinaryFunction *
 BinaryContext::getBinaryFunctionContainingAddress(uint64_t Address,
                                                   bool CheckPastEnd,
diff --git a/bolt/src/BinaryContext.h b/bolt/src/BinaryContext.h
index 95e40ea1e1b9..990c89405bde 100644
--- a/bolt/src/BinaryContext.h
+++ b/bolt/src/BinaryContext.h
@@ -989,6 +989,11 @@ public:
     return Size;
   }
 
+  /// Verify that assembling instruction \p Inst results in the same sequence of
+  /// bytes as \p Encoding.
+  bool validateEncoding(const MCInst &Instruction,
+                        ArrayRef<uint8_t> Encoding) const;
+
   /// Return a function execution count threshold for determining whether
   /// the function is 'hot'. Consider it hot if count is above the average exec
   /// count of profiled functions.
diff --git a/bolt/src/BinaryFunction.cpp b/bolt/src/BinaryFunction.cpp
index f10d66216745..ff8fa016ae68 100644
--- a/bolt/src/BinaryFunction.cpp
+++ b/bolt/src/BinaryFunction.cpp
@@ -1027,36 +1027,35 @@ void BinaryFunction::disassemble(ArrayRef<uint8_t> FunctionData) {
     // Check integrity of LLVM assembler/disassembler.
     if (opts::CheckEncoding && !BC.MIB->isBranch(Instruction) &&
         !BC.MIB->isCall(Instruction) && !BC.MIB->isNoop(Instruction)) {
-      SmallString<256> Code;
-      SmallVector<MCFixup, 4> Fixups;
-      raw_svector_ostream VecOS(Code);
-      BC.MCE->encodeInstruction(Instruction, VecOS, Fixups, *BC.STI);
-      auto EncodedData = ArrayRef<uint8_t>((uint8_t *)Code.data(), Code.size());
-      if (FunctionData.slice(Offset, Size) != EncodedData) {
+      if (!BC.validateEncoding(Instruction, FunctionData.slice(Offset, Size))) {
         errs() << "BOLT-WARNING: mismatching LLVM encoding detected in "
-               << "function " << *this << ":\n";
+               << "function " << *this << " for instruction :\n";
         BC.printInstruction(errs(), Instruction, AbsoluteInstrAddr);
-        errs() << "       input: " << FunctionData.slice(Offset, Size)
-               << "\n      output: " << EncodedData << "\n\n";
+        errs() << '\n';
       }
     }
 
-    // Cannot process functions with AVX-512 instructions.
+    // Special handling for AVX-512 instructions.
     if (MIB->hasEVEXEncoding(Instruction)) {
-      if (opts::Verbosity >= 1) {
-        errs() << "BOLT-WARNING: function " << *this << " uses instruction"
-               " encoded with EVEX (AVX-512) at offset 0x"
-               << Twine::utohexstr(Offset) << ". Disassembly could be wrong."
-               " Skipping further processing.\n";
-      }
-
       if (BC.HasRelocations && opts::TrapOnAVX512) {
         setTrapOnEntry();
         BC.TrappedFunctions.push_back(this);
-      } else {
-        IsSimple = false;
+        break;
+      }
+
+      // Check if our disassembly is correct and matches the assembler output.
+      if (!BC.validateEncoding(Instruction, FunctionData.slice(Offset, Size))) {
+        if (BC.HasRelocations) {
+          errs() << "BOLT-ERROR: internal assembler/disassembler error "
+                    "detected for AVX512 instruction:\n";
+          BC.printInstruction(errs(), Instruction, AbsoluteInstrAddr);
+          errs() << " in function " << *this << '\n';
+          exit(1);
+        } else {
+          setSimple(false);
+          break;
+        }
       }
-      break;
     }
 
     // Check if there's a relocation associated with this instruction.
diff --git a/bolt/src/Passes/BinaryPasses.cpp b/bolt/src/Passes/BinaryPasses.cpp
index 5f085a098ce3..60143dcb4d2e 100644
--- a/bolt/src/Passes/BinaryPasses.cpp
+++ b/bolt/src/Passes/BinaryPasses.cpp
@@ -1434,13 +1434,15 @@ PrintProgramStats::runOnFunctions(BinaryContext &BC) {
 
   if (!BC.TrappedFunctions.empty()) {
     errs() << "BOLT-WARNING: " << BC.TrappedFunctions.size()
-           << " functions will trap on entry";
-    if (opts::Verbosity >= 1) {
-      errs() << ".\n";
+           << " function" << (BC.TrappedFunctions.size() > 1 ? "s" : "")
+           << " will trap on entry. Use -trap-avx512=0 to disable"
+              " traps.";
+    if (opts::Verbosity >= 1 || BC.TrappedFunctions.size() <= 5) {
+      errs() << '\n';
       for (const auto *Function : BC.TrappedFunctions)
         errs() << "  " << *Function << '\n';
     } else {
-      errs() << " (use -v=1 to see the list).\n";
+      errs() << " Use -v=1 to see the list.\n";
     }
   }