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

[Steven Wu via Phabricator via llvm-commits]

steven_wu updated this revision to Diff 279312.
steven_wu added a comment.

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.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D83699/new/

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,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}
Index: llvm/lib/Bitcode/Reader/BitcodeReader.cpp
===================================================================
--- llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -5394,6 +5394,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.279312.patch
Type: text/x-patch
Size: 2512 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200721/9073b220/attachment-0001.bin>