From 78709345fb343123020a8615325f3943bd1cb420 Mon Sep 17 00:00:00 2001
From: Steven Wu <stevenwu@apple.com>
Date: Thu, 23 Jul 2020 09:06:47 -0700
Subject: [PATCH] [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
---
 llvm/lib/Bitcode/Reader/BitcodeReader.cpp |  30 ++++++++++++++++++++++
 llvm/test/Bitcode/Inputs/branch-weight.bc | Bin 0 -> 1488 bytes
 llvm/test/Bitcode/branch-weight.ll        |  19 ++++++++++++++
 3 files changed, 49 insertions(+)
 create mode 100644 llvm/test/Bitcode/Inputs/branch-weight.bc
 create mode 100644 llvm/test/Bitcode/branch-weight.ll

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 0000000000000000000000000000000000000000..fbb77a39b1dea5c520987255cbd5160655222ba6
GIT binary patch
literal 1488
zcmY*ZZ){Rm6u<OQ9$LJ-Dp9Zx_ONwn%+%M1woqUxeg5gTF2lGF-R2@-LAtd1N=ZxW
zvVh<oW_&>H2R``0CVtRwsLkRBv(dgznIWb#W^r57loYEj)sX2HlbLPjqFb_)_q+F=
z^Uk~HoZq?cE*3U5AP@s!x*nmwLAu4l;%+d2iUK;m<^7Mo2@U-@{XtW)WsgdzM^+s|
zTU7`()LBs-_-q8wiRLq&qWT>TBRI~gQo616T9ZZnRUoe0Ti1QIp}U_E40f&Yy|sdn
z?vduK&3o3%8OC$AuD__gV&KGmY3_TKAf$}jTDN7}KF=RI!zo4B_cVtPghRrXqL%AF
zr9Nl9n1W$BHI+52@BK2<km~Pc*k&}M_B1kgy9$N^t1^QBhD%ab9?eso;iMF&V?H?&
zPDW$$c!BDO!;@U+_u`>>JW<6%Sv=ZBPx$EQyyh`ON8<EEZ5&#{Biw8IwVGwKRkx<T
zabA1tcw>6F_phtE%cvVd{U;ncy{>NWL7?Pnmj<kYC8skr_GU9{-rRpvEk@QDJSO6N
z1rj_5vrQXk=gMrzpmL{0drHZ4$LaAHJ(3&Ntswmq20xEuo37h3J7c~x#fJF$np?7|
zM7fKiJ43lku}oB=PiX1!9$Av;35gnw$r2}<N{x6mi_5(J7_VQh;;I+hm^t&JmA!8@
z&-yJP{<QN((m5|WrYYBrq$@=^OI5?Dk3JEmJHkK}-4U08Kgmarw$h_7>S5clxO`YE
z4;$r3P9D!C-g;WaPnYoT1-!~rf#MWfnPN+|aly~twz6{x^KHN8$E@fmijG`uysp?M
zMSGfZPAWSg=*yGNdz3S;=xz<jQlA`=lF=|}RMaR(FUY1l2lY>p{x>hKWbtyBzRclO
zsDf)XFU(kGJ(jyxb|GO|SYj(>^RFsP$auh>7dy`<oh5I}>q30;Sd5P44)i|H>Inmu
z&n&n7>?5mramM_G%0d#hnyY@zho_SElH$yZZXo34tg9kA)1WGu@76-4@)+c!r6*eH
zNG@?`nZc_C{J@K!@y#=LJ(fy&)BUn}5fbPzIrB;Pbh0U>*fXM|q`0onIEo6?=qmM0
z4{K@2Y8VLYlTE+e#H&92lO;`+(?8(#4><iFUVMzf%RC14*^=h*60Ugh3tlY$Y&EY;
zu~k(qbN{}3KPh%-R$ke8U+gT2ooTTXs-73$Mwj7Ck;B!nHqr{cq9^*KoahG1^P*#3
zIl02#oLcwDLyppl)%=Qsjk}h8_@Ni5)mI3?`iEZpJReWc<E^#M$7INw)ZHhYSM1X&
z=f$LZl7bjNFbI+JO_y93lg{(dP|AHF>4N5F6}T3VpF?O9%p25uT;$&+*s{NvJ)lXc
zF{}d-|FI+C!(+jLqv7ZX@oUa;##!}@v0azo2e~>whpy=oOi9pn7dA5#iTR19i~Cxa
z@WQet!yH1yE?{O$3T4SPRFfbCdjnSMsmUrakQZuz0O$pf{Q3Yy$V_yJC?M?sau2=>
zU;yMhBLBW8YvNDjV*t4yi4Z?>Z=#l(iav!LnqUbzd`m3R;7GJh*un4Mj}An+jvyZp
z9JWx%(QgYL+0}0ocw5NkaI^*50>NOQJrL*!2)5mVKQh$5Yr9{P27}uNj*r;M_y1#V
P$Kc>$vh}N=ug8A?F$l>l

literal 0
HcmV?d00001

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}