[PATCH] D83699: [Bitcode] Drop invalid branch_weight in BitcodeReader

[Steven Wu via Phabricator via llvm-commits]

steven_wu created this revision.
steven_wu added reviewers: yrouban, t.p.northover, dexonsmith.
Herald added subscribers: ributzka, jkorous, hiraditya.
Herald added a project: LLVM.

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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D83699

Files:
  llvm/lib/Bitcode/Reader/BitcodeReader.cpp
  llvm/test/Bitcode/Inputs/branch-weight.bc
  llvm/test/Bitcode/branch-weight.ll


Index: llvm/test/Bitcode/branch-weight.ll
===================================================================
--- /dev/null
+++ llvm/test/Bitcode/branch-weight.ll
@@ -0,0 +1,14 @@
+;; 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:
+;;   unreachable
+;; 3:
+;;   ret void
+;; }
+;; !0 = !{!"branch_weights", i32 1, i32 2, i32 3}
+
+; RUN: llvm-dis %S/Inputs/branch-weight.bc -o - | FileCheck %s
+; CHECK-NOT: !prof
+; CHECK-NOT: !"branch_weights"
Index: llvm/lib/Bitcode/Reader/BitcodeReader.cpp
===================================================================
--- llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -5390,6 +5390,37 @@
     }
   }
 
+  // Check the branch weight on the instructions are correct. If not, drop them.
+  for (auto &I : instructions(F)) {
+    if (auto *MD = I.getMetadata(LLVMContext::MD_prof)) {
+      if (MD->getNumOperands() >= 2 || 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")) {
+          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 doens'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);
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D83699.277489.patch
Type: text/x-patch
Size: 2317 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200713/5164bfb7/attachment-0001.bin>