[llvm] 7870934 - [Bitcode] Drop invalid branch_weight in BitcodeReader

[Steven Wu via llvm-commits]

Author: Steven Wu
Date: 2020-07-23T09:07:15-07:00
New Revision: 78709345fb343123020a8615325f3943bd1cb420

URL: https://github.com/llvm/llvm-project/commit/78709345fb343123020a8615325f3943bd1cb420
DIFF: https://github.com/llvm/llvm-project/commit/78709345fb343123020a8615325f3943bd1cb420.diff

LOG: [Bitcode] Drop invalid branch_weight in BitcodeReader

Summary:
If bitcode reader gets an invalid branch weight, drop that from the
inputs. This allows us to read the broken modules we generated before
the verifier was able to catch this.

rdar://64870641

Reviewers: yrouban, t.p.northover, dexonsmith, arphaman, aprantl

Reviewed By: aprantl

Subscribers: aprantl, hiraditya, jkorous, ributzka, llvm-commits

Tags: #llvm

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

Added: 
    llvm/test/Bitcode/Inputs/branch-weight.bc
    llvm/test/Bitcode/branch-weight.ll

Modified: 
    llvm/lib/Bitcode/Reader/BitcodeReader.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index 86b8f5ca5b84..2c7882b6ad51 100644
--- a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -5394,6 +5394,36 @@ Error BitcodeReader::materialize(GlobalValue *GV) {
     }
   }
 
+  // "Upgrade" older incorrect branch weights by dropping them.
+  for (auto &I : instructions(F)) {
+    if (auto *MD = I.getMetadata(LLVMContext::MD_prof)) {
+      if (MD->getOperand(0) != nullptr && isa<MDString>(MD->getOperand(0))) {
+        MDString *MDS = cast<MDString>(MD->getOperand(0));
+        StringRef ProfName = MDS->getString();
+        // Check consistency of !prof branch_weights metadata.
+        if (!ProfName.equals("branch_weights"))
+          continue;
+        unsigned ExpectedNumOperands = 0;
+        if (BranchInst *BI = dyn_cast<BranchInst>(&I))
+          ExpectedNumOperands = BI->getNumSuccessors();
+        else if (SwitchInst *SI = dyn_cast<SwitchInst>(&I))
+          ExpectedNumOperands = SI->getNumSuccessors();
+        else if (isa<CallInst>(&I))
+          ExpectedNumOperands = 1;
+        else if (IndirectBrInst *IBI = dyn_cast<IndirectBrInst>(&I))
+          ExpectedNumOperands = IBI->getNumDestinations();
+        else if (isa<SelectInst>(&I))
+          ExpectedNumOperands = 2;
+        else
+          continue; // ignore and continue.
+
+        // If branch weight doesn't match, just strip branch weight.
+        if (MD->getNumOperands() != 1 + ExpectedNumOperands)
+          I.setMetadata(LLVMContext::MD_prof, nullptr);
+      }
+    }
+  }
+
   // Look for functions that rely on old function attribute behavior.
   UpgradeFunctionAttributes(*F);
 

diff  --git a/llvm/test/Bitcode/Inputs/branch-weight.bc b/llvm/test/Bitcode/Inputs/branch-weight.bc
new file mode 100644
index 000000000000..fbb77a39b1de
Binary files /dev/null and b/llvm/test/Bitcode/Inputs/branch-weight.bc 
diff er

diff  --git a/llvm/test/Bitcode/branch-weight.ll b/llvm/test/Bitcode/branch-weight.ll
new file mode 100644
index 000000000000..d93d79e93922
--- /dev/null
+++ b/llvm/test/Bitcode/branch-weight.ll
@@ -0,0 +1,19 @@
+;; Test strip branch_weight if operand number doesn't match.
+;; Input bitcode is generated from:
+;; define void @test(i1 %0) {
+;;   br i1 %0, label %2, label %3, !prof !0
+;; 2:
+;;   br i1 %0, label %4, label %3, !prof !1
+;; 3:
+;;   unreachable
+;; 4:
+;;   ret void
+;; }
+;;!0 = !{!"branch_weights", i32 1, i32 2}
+;;!1 = !{!"branch_weights", i32 1, i32 2, i32 3}
+
+; RUN: llvm-dis %S/Inputs/branch-weight.bc -o - | FileCheck %s
+; CHECK: !prof !0
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK-NOT: !prof !1
+; CHECK-NOT: !1 = !{!"branch_weights", i32 1, i32 2, i32 3}