[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

diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index 86b8f5ca5b84db2f4752dfa2a7a8c145a3b20e75..2c7882b6ad515d1be252489f7e45bf6fda3cf295 100644 (file)
--- 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 (file)
index 0000000..fbb77a3
Binary files /dev/null and b/llvm/test/Bitcode/Inputs/branch-weight.bc differ

diff --git a/llvm/test/Bitcode/branch-weight.ll b/llvm/test/Bitcode/branch-weight.ll
new file mode 100644 (file)
index 0000000..d93d79e
--- /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}