Fix number of args of prof branch_weights MD for SwitchInst

[yrouban]

Summary

This patch changes SwitchInst implementation to keep number of args of prof branch_weights metadata consistent across removeCase() and addCase() calls.
Appropriate code to check this consistency is added to the verifier.
Two ragreedy* tests got their broken prof metadata fixed.
Added several prof branch_weights and appropriate checks to the other 3 tests.

SimplifyCFG.cpp handles the branch weights on its own. It sets the weights when all cases are removed. So it is made to drop the metadata while removing the cases.

Note that this patch does not fix weight values when cases are changed. There will be another patch that implements prof branch_weights metadata handling for SwitchInst instruction for setting correct weights while adding or changing cases.

[yrouban]

llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

1247

As the default branch is changed it must expect !{!"branch_weights", i32 4, i32 1, i32 2} but the test passes because of a bug in SimpleLoopUnswitch. I'm fixing it and will provide one more patch. For this case I can change all weights to 1s as I did for SimpleLoopUnswitch/basictest.ll above.

[yrouban]

fixed test llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

[yrouban]

llvm/lib/IR/Instructions.cpp

3843

must be uint64_t

[reames]

llvm/test/CodeGen/AArch64/ragreedy-csr.ll

281 ↗

(On Diff #194827)

Please separate and commit the test changes.

[yrouban]

• removed 2 tests fixes into a separate patch
• changed branch weight sum type to uint64_t
• extracted separate method getProfBranchWeightsMD() that is used in the dependent patche

[yrouban]

formatted

[yrouban]

See 62122 and its children.

llvm/lib/IR/Verifier.cpp

2453

add the following from BranchProbabilityInfo

assert(Weight->getValue().getActiveBits() <= 32 && "Too many bits for uint32_t");

llvm/include/llvm/IR/Instructions.h

public:
 
   /// This method removes the specified case and its successor from the switch
   /// instruction. Note that this operation may reorder the remaining cases at
   /// index idx and above.
   /// Note:
   /// This action invalidates iterators for all cases following the one removed,
   /// including the case_end() iterator. It returns an iterator for the next
   /// case.
-  CaseIt removeCase(CaseIt I);
+  CaseIt removeCase(CaseIt I, bool DropPerfMetadata = false);
 
   unsigned getNumSuccessors() const { return getNumOperands()/2; }
   BasicBlock *getSuccessor(unsigned idx) const {
     assert(idx < getNumSuccessors() &&"Successor idx out of range for switch!");
     return cast<BasicBlock>(getOperand(idx*2+1));
   }
   void setSuccessor(unsigned idx, BasicBlock *NewSucc) {
     assert(idx < getNumSuccessors() && "Successor # out of range for switch!");

llvm/lib/IR/Instructions.cpp

 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
SwitchInst::SwitchInst(const SwitchInst &SI)
   const Use *InOL = SI.getOperandList();
   for (unsigned i = 2, E = SI.getNumOperands(); i != E; i += 2) {
     OL[i] = InOL[i];
     OL[i+1] = InOL[i+1];
   }
   SubclassOptionalData = SI.SubclassOptionalData;
 }
 
+static MDNode *getProfBranchWeightsMD(const SwitchInst *SI) {
+  if (MDNode *ProfileData = SI->getMetadata(LLVMContext::MD_prof))
+    if (auto *MDName = dyn_cast<MDString>(ProfileData->getOperand(0)))
+      if (MDName->getString() == "branch_weights")
+        return ProfileData;
+  return nullptr;
+}
+
 /// addCase - Add an entry to the switch instruction...
 ///
 void SwitchInst::addCase(ConstantInt *OnVal, BasicBlock *Dest) {
+  // Drop metdadata.
+  if (MDNode *ProfileData = getProfBranchWeightsMD(this))
+    if (ProfileData->getNumOperands() == getNumSuccessors() + 1)
+      setMetadata(LLVMContext::MD_prof, nullptr);
+
   unsigned NewCaseIdx = getNumCases();
   unsigned OpNo = getNumOperands();
   if (OpNo+2 > ReservedSpace)
     growOperands();  // Get more space!
   // Initialize some new operands.
   assert(OpNo+1 < ReservedSpace && "Growing didn't work!");
   setNumHungOffUseOperands(OpNo+2);
   CaseHandle Case(this, NewCaseIdx);
   Case.setValue(OnVal);
   Case.setSuccessor(Dest);
 }
 
 /// removeCase - This method removes the specified case and its successor
 /// from the switch instruction.
-SwitchInst::CaseIt SwitchInst::removeCase(CaseIt I) {
+SwitchInst::CaseIt SwitchInst::removeCase(CaseIt I, bool DropPerfMetadata) {
   unsigned idx = I->getCaseIndex();
 
   assert(2 + idx*2 < getNumOperands() && "Case index out of range!!!");
 
+  if (MDNode *ProfileData = getProfBranchWeightsMD(this))
+    if (ProfileData->getNumOperands() == getNumSuccessors() + 1) {
+      if (DropPerfMetadata) {
+        setMetadata(LLVMContext::MD_prof, nullptr);
+      } else {
+        SmallVector<uint32_t, 8> Weights;
+        uint64_t SW = 0;
+        for (unsigned CI = 0, CE = getNumSuccessors() - 1; CI < CE; ++CI) {
+          // At the removed position we put the end of the list.
+          unsigned WI = (CI == idx + 1 ? CE : CI) + 1;
+          ConstantInt *C =
+              mdconst::extract<ConstantInt>(ProfileData->getOperand(WI));
+          uint32_t CW = C->getValue().getZExtValue();
+          SW += CW;
+          Weights.push_back(CW);
+        }
+        assert(Weights.size() == getNumSuccessors() - 1);
+        auto *NewProfileData = !SW || Weights.size() < 2
+                                   ? nullptr
+                                   : MDBuilder(getParent()->getContext())
+                                         .createBranchWeights(Weights);
+        setMetadata(LLVMContext::MD_prof, NewProfileData);
+      }
+    }
+
   unsigned NumOps = getNumOperands();
   Use *OL = getOperandList();
 
   // Overwrite this case with the end of the list.
   if (2 + (idx + 1) * 2 != NumOps) {
     OL[2 + idx * 2] = OL[NumOps - 2];
     OL[2 + idx * 2 + 1] = OL[NumOps - 1];
   }
 
   // Nuke the last value.
   OL[NumOps-2].set(nullptr);
   OL[NumOps-2+1].set(nullptr);
   setNumHungOffUseOperands(NumOps-2);
 
   return CaseIt(this, idx);
 }
 
 /// growOperands - grow operands - This grows the operand list in response
 /// to a push_back style of operation.  This grows the number of ops by 3 times.
 ///
 void SwitchInst::growOperands() {

llvm/lib/IR/Verifier.cpp

void Verifier::visitSwitchInst(SwitchInst &SI) {
   SmallPtrSet<ConstantInt*, 32> Constants;
   for (auto &Case : SI.cases()) {
     Assert(Case.getCaseValue()->getType() == SwitchTy,
            "Switch constants must all be same type as switch value!", &SI);
     Assert(Constants.insert(Case.getCaseValue()).second,
            "Duplicate integer as switch case", &SI, Case.getCaseValue());
   }
 
+  // Check consistency of !prof branch_weights metadata if any.
+  auto *ProfileData = SI.getMetadata(LLVMContext::MD_prof);
+  if (ProfileData != nullptr && ProfileData->getNumOperands() > 0) {
+    if (auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0)))
+      if (ProfDataName->getString().equals("branch_weights")) {
+        Assert(ProfileData->getNumOperands() == SI.getNumSuccessors() + 1,
+               "Number of !prof brunch_weights operands differ from number of "
+               "successors");
+        for (unsigned i = 1; i < ProfileData->getNumOperands(); ++i)
+          Assert(mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(i)),
+                 "!prof brunch_weights operand is not a const int");
+      }
+  }
+
   visitTerminator(SI);
 }
 
 void Verifier::visitIndirectBrInst(IndirectBrInst &BI) {
   Assert(BI.getAddress()->getType()->isPointerTy(),
          "Indirectbr operand must have pointer type!", &BI);
   for (unsigned i = 0, e = BI.getNumDestinations(); i != e; ++i)
     Assert(BI.getDestination(i)->getType()->isLabelTy(),

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

if (PredDef == TI->getParent()) {
     for (SwitchInst::CaseIt i = SI->case_end(), e = SI->case_begin(); i != e;) {
       --i;
       if (DeadCases.count(i->getCaseValue())) {
         if (HasWeight) {
           std::swap(Weights[i->getCaseIndex() + 1], Weights.back());
           Weights.pop_back();
         }
         i->getCaseSuccessor()->removePredecessor(TI->getParent());
-        SI->removeCase(i);
+        SI->removeCase(i, true);
       }
     }
     if (HasWeight && Weights.size() >= 2)
       setBranchWeights(SI, Weights);
 
     LLVM_DEBUG(dbgs() << "Leaving: " << *TI << "\n");
     return true;
   }
assert(CaseI != SI->case_default() &&
            "Case was not found. Probably mistake in DeadCases forming.");
     if (HasWeight) {
       std::swap(Weights[CaseI->getCaseIndex() + 1], Weights.back());
       Weights.pop_back();
     }
 
     // Prune unused values from PHI nodes.
     CaseI->getCaseSuccessor()->removePredecessor(SI->getParent());
-    SI->removeCase(CaseI);
+    SI->removeCase(CaseI, true);
   }
   if (HasWeight && Weights.size() >= 2) {
     SmallVector<uint32_t, 8> MDWeights(Weights.begin(), Weights.end());
     setBranchWeights(SI, MDWeights);
   }
 
   return !DeadCases.empty();
 }

llvm/test/Transforms/CorrelatedValuePropagation/basic.ll

 ; CHECK-LABEL: @switch1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[S:%.*]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label [[NEGATIVE:%.*]], label [[OUT:%.*]]
 ; CHECK:       negative:
 ; CHECK-NEXT:    switch i32 [[S]], label [[OUT]] [
 ; CHECK-NEXT:    i32 -2, label [[NEXT:%.*]]
 ; CHECK-NEXT:    i32 -1, label [[NEXT]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    !prof ![[MD0:[0-9]+]]
 ; CHECK:       out:
 ; CHECK-NEXT:    [[P:%.*]] = phi i32 [ 1, [[ENTRY:%.*]] ], [ -1, [[NEGATIVE]] ]
 ; CHECK-NEXT:    ret i32 [[P]]
 ; CHECK:       next:
 ; CHECK-NEXT:    ret i32 0
 ;
 entry:
   %cmp = icmp slt i32 %s, 0
   br i1 %cmp, label %negative, label %out
 
 negative:
   switch i32 %s, label %out [
   i32 0, label %out
   i32 1, label %out
   i32 -1, label %next
   i32 -2, label %next
   i32 2, label %out
   i32 3, label %out
-  ]
+  ], !prof !{!"branch_weights", i32 99, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6}
 
 out:
   %p = phi i32 [ 1, %entry ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ]
   ret i32 %p
 
 next:
   %q = phi i32 [ 0, %negative ], [ 0, %negative ]
   ret i32 %q

 ;
 entry:
   %and = lshr i32 %a, 30
   %cmp = icmp sle i32 %and, 128
   br label %exit
 exit:
   ret i1 %cmp
 }
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 99, i32 4, i32 3}

llvm/test/Transforms/SimpleLoopUnswitch/basictest.ll

 ; CHECK-LABEL: @test2(
 define i32 @test2(i32* %var) {
   %mem = alloca i32
   store i32 2, i32* %mem
   %c = load i32, i32* %mem
 
   br label %loop_begin
 
+; CHECK: loop_begin:
+; CHECK: !prof ![[MD0:[0-9]+]]
 loop_begin:
 
   %var_val = load i32, i32* %var
 
   switch i32 %c, label %default [
       i32 1, label %inc
       i32 2, label %dec
-  ]
+  ], !prof !{!"branch_weights", i32 1, i32 1, i32 1}
 
 inc:
   call void @incf() noreturn nounwind
   br label %loop_begin
 dec:
 ; CHECK: call void @decf()
 ; CHECK-NOT: call void @decf()
   call void @decf() noreturn nounwind noduplicate

 
 for.end:
   ret void
 }
 
 declare void @incf() noreturn
 declare void @decf() noreturn
 declare void @conv() convergent
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 1, i32 1}

llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

 loop_begin:
   %var_val = load i32, i32* %var
   switch i32 %cond2, label %loop_exit2 [
     i32 0, label %loop0
     i32 1, label %loop1
     i32 13, label %loop_exit1
     i32 2, label %loop2
     i32 42, label %loop_exit3
-  ]
+  ], !prof !{!"branch_weights", i32 1, i32 1, i32 1, i32 1, i32 1, i32 1}
 ; CHECK:       loop_begin:
 ; CHECK-NEXT:    load
 ; CHECK-NEXT:    switch i32 %cond2, label %loop2 [
 ; CHECK-NEXT:      i32 0, label %loop0
 ; CHECK-NEXT:      i32 1, label %loop1
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    ], !prof ![[MD0:[0-9]+]]
 
 loop0:
   call void @some_func() noreturn nounwind
   br label %loop_latch
 ; CHECK:       loop0:
 ; CHECK-NEXT:    call
 ; CHECK-NEXT:    br label %loop_latch
 

 ; CHECK-NEXT:    call void @sink(i32 %[[PHI]])
 ; CHECK-NEXT:    br label %header
 
 loopexit:
   ret void
 ; CHECK:       loopexit:
 ; CHECK-NEXT:    ret
 }
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 1, i32 1, i32 1}