[Bitcode] Drop invalid branch_weight in BitcodeReader

[steven_wu]

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

[aprantl]

So this is a bitcode "upgrade" of sorts. An alternative would be to implement this in Verifier and strip invalid metadata similar to how invalid debug info is stripped, but treating this as an upgrade in the bitcode reader is probably fine too.

llvm/test/Bitcode/branch-weight.ll

15

Can you add one positive CHECK, too, to make sure this doesn't just print nothing at all?

[steven_wu]

Address review feedback and also add a testcase to make sure good branch
weights are not stripped.

I don't think the debug info approach is good here since it is a lot of
information to pass through to only strip the bad branch weights. Also
not sure if it is really a good idea to run verifier in upgrader.

[aprantl]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

5397

Perhaps even more explicit:
// "Upgrade" older incorrect branch weights by dropping them.

5400

Are the || operators correct here?

I feel like this should be

MD->getOperand(0) != nullptr && isa<MDString>(MD->getOperand(0) ?

5405

if (!ProfName.equals("branch_weights"))

continue;

5420

doesn't

[steven_wu]

Address review feedback

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

for (auto &I : instructions(F)) {
       MDNode *TBAA = I.getMetadata(LLVMContext::MD_tbaa);
       if (!TBAA || TBAAVerifyHelper.visitTBAAMetadata(I, TBAA))
         continue;
       MDLoader->setStripTBAA(true);
       stripTBAA(F->getParent());
     }
   }
 
+  // "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);
 
   // Bring in any functions that this function forward-referenced via
   // blockaddresses.
   return materializeForwardReferencedFunctions();
 }
 

llvm/test/Bitcode/branch-weight.ll

+;; 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}