Make SwitchInstProfUpdateWrapper safer

[yrouban]

Summary

While prof branch_weights inconsistencies are being fixed patch by patch (pass by pass) we need SwitchInstProfUpdateWrapper to be safe with respect to inconsistent metadata that can come from passes that have not been fixed yet. See the bug found by @nikic in D62126.
This patch introduces one more state (called Invalid) to the wrapper class that allows users to work with the underlying SwitchInst ignoring the prof metadata changes.
Created a unit test for the SwitchInstProfUpdateWrapper class.

[davidxl]

llvm/lib/IR/Instructions.cpp

3907 ↗

(On Diff #202166)

should we fix the underlying problem that leads to this inconsistent state?

[yrouban]

llvm/lib/IR/Instructions.cpp

3907 ↗

(On Diff #202166)

That is what I'm trying to do.
The problem could have been solved by one of the proposed patches. See the dependency graph.
It would be great to identify where the problem stems from.

[davidxl]

What I meant is that we should never see inconsistent state and instead of setting invalid state, we assert there. Ideally, the profile setting/update APIs should help user avoid creating the inconsistent state in the first place (i.e., catch the problem on the spot).

llvm/unittests/IR/InstructionsTest.cpp

797 ↗

(On Diff #202166)

looks like exposing the raw interface to setMetadata for SI is not safe. Using wrapper interface can help catch error that leads to inconsistent state.

813 ↗

(On Diff #202166)

why it is not 39?

[xbolva00]

llvm/include/llvm/IR/Instructions.h

3444 ↗

(On Diff #202166)

Typo

[yrouban]

In D62656#1523764, @davidxl wrote:

What I meant is that we should never see inconsistent state and instead of setting invalid state, we assert there.

Yes. But in LLVM there are many places to fix. So we have to fix all of them at once and set the assertion here. The other option is to fix them one by one and tolerate such inconsistencies.
For example, SimplifyCFG is safe: it checks that the prof data is ok and only then makes its prof changes.

[yrouban]

llvm/unittests/IR/InstructionsTest.cpp

797 ↗

(On Diff #202166)

it is just a test and SI is just a SwitchInst.

813 ↗

(On Diff #202166)

That is because 39 was set when SIW was inconsistent. So SwitchInstProfUpdateWrapper::addCase() skipped setting the weight 39 and delegated addCase() to the underlying SwitchInst. The prof data attached to the underlying SwicthInst was kept intact and had 4 operands: 99, 11, 22, 33. So the next time we created a new wrapper it started with valid state.

[yrouban]

Formatted and fixed one typo.

[davidxl]

I suggest the following path:

1. add an option to enforce internal error on inconsistency; by default the inconsistency is ignored.
2. fix the exposed root cause one by one
3. turn on the option by default

llvm/unittests/IR/InstructionsTest.cpp

797 ↗

(On Diff #202166)

What I meant is that this kind of interface can be source of inconsistency.

813 ↗

(On Diff #202166)

Such interface behavior is not desirable. It is certainly surprising to the user which calls addCase as there is no way to know that the weight is ignored.

[nikic]

@davidxl I believe the changes here are only temporary until D61179 lands, which will enforce this at the verifier level. As such I don't think there's much value in introducing a verification option here, because this is not really the right place to do that.

What I would suggest though is to maybe drop the branch weight metadata entirely if it's not valid. I think it's pretty weird (as was commented above) that you might end up with "valid" metadata after performing additional transformations (removing/adding cases), but that metadata will just be an accidental assignment, more or less. I think it would be better to discard it wholesale than risk an incorrect reinterpretation.

[yrouban]

In D62656#1524574, @davidxl wrote:

I suggest the following path:

1. add an option to enforce internal error on inconsistency; by default the inconsistency is ignored.

Ok. This option will probably simplify catching such bugs.

llvm/unittests/IR/InstructionsTest.cpp

813 ↗

(On Diff #202166)

In future there will be no inconsistent states allowed and I suggest that the users write their code as if the stae is always consistent. During the transition period we just propagate the Undefined Behavior related to the prof data.

[yrouban]

added option switch-inst-prof-update-wrapper-strict to assert if prof data is inconsistent. It is off by default.
Probably it makes sense to enable prof branch_weights validation (patch D61179) in the same way.

[davidxl]

once the cleanup is done, the state tracking should be removed as well.

llvm/lib/IR/Instructions.cpp

48 ↗

(On Diff #202393)

typo: Strinct

[asbirlea]

llvm/lib/IR/Instructions.cpp

48 ↗

(On Diff #202393)

also typo: Swicth, var name and in the cl::desc below.

[yrouban]

fixed typos

[davidxl]

lgtm assuming the state tracking code gets removed (or put under NDEBUG in some way) after the cleanup.

llvm/trunk/include/llvm/IR/Instructions.h

static bool classof(const Value *V) {
     return isa<Instruction>(V) && classof(cast<Instruction>(V));
   }
 };
 
 /// A wrapper class to simplify modification of SwitchInst cases along with
 /// their prof branch_weights metadata.
 class SwitchInstProfUpdateWrapper {
   SwitchInst &SI;
-  Optional<SmallVector<uint32_t, 8> > Weights;
-  bool Changed = false;
+  Optional<SmallVector<uint32_t, 8> > Weights = None;
+
+  // Sticky invalid state is needed to safely ignore operations with prof data
+  // in cases where SwitchInstProfUpdateWrapper is created from SwitchInst
+  // with inconsistent prof data. TODO: once we fix all prof data
+  // inconsistencies we can turn invalid state to assertions.
+  enum {
+    Invalid,
+    Initialized,
+    Changed
+  } State = Invalid;
 
 protected:
   static MDNode *getProfBranchWeightsMD(const SwitchInst &SI);
 
   MDNode *buildProfBranchWeightsMD();
 
-  Optional<SmallVector<uint32_t, 8> > getProfBranchWeights();
+  void init();
 
 public:
   using CaseWeightOpt = Optional<uint32_t>;
   SwitchInst *operator->() { return &SI; }
   SwitchInst &operator*() { return SI; }
   operator SwitchInst *() { return &SI; }
 
-  SwitchInstProfUpdateWrapper(SwitchInst &SI)
-      : SI(SI), Weights(getProfBranchWeights()) {}
+  SwitchInstProfUpdateWrapper(SwitchInst &SI) : SI(SI) { init(); }
 
   ~SwitchInstProfUpdateWrapper() {
-    if (Changed)
+    if (State == Changed)
       SI.setMetadata(LLVMContext::MD_prof, buildProfBranchWeightsMD());
   }
 
   /// Delegate the call to the underlying SwitchInst::removeCase() and remove
   /// correspondent branch weight.
   SwitchInst::CaseIt removeCase(SwitchInst::CaseIt I);
 
   /// Delegate the call to the underlying SwitchInst::addCase() and set the

llvm/trunk/lib/IR/Instructions.cpp

 #include "llvm/Support/MathExtras.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <vector>
 
 using namespace llvm;
 
+static cl::opt<bool> SwitchInstProfUpdateWrapperStrict(
+    "switch-inst-prof-update-wrapper-strict", cl::Hidden,
+    cl::desc("Assert that prof branch_weights metadata is valid when creating "
+             "an instance of SwitchInstProfUpdateWrapper"),
+    cl::init(false));
+
 //===----------------------------------------------------------------------===//
 //                            AllocaInst Class
 //===----------------------------------------------------------------------===//
 
 Optional<uint64_t>
 AllocaInst::getAllocationSizeInBits(const DataLayout &DL) const {
   uint64_t Size = DL.getTypeAllocSizeInBits(getAllocatedType());
   if (isArrayAllocation()) {
SwitchInstProfUpdateWrapper::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;
 }
 
 MDNode *SwitchInstProfUpdateWrapper::buildProfBranchWeightsMD() {
-  assert(Changed && "called only if metadata has changed");
+  assert(State == Changed && "called only if metadata has changed");
 
   if (!Weights)
     return nullptr;
 
   assert(SI.getNumSuccessors() == Weights->size() &&
          "num of prof branch_weights must accord with num of successors");
 
   bool AllZeroes =
       all_of(Weights.getValue(), [](uint32_t W) { return W == 0; });
 
   if (AllZeroes || Weights.getValue().size() < 2)
     return nullptr;
 
   return MDBuilder(SI.getParent()->getContext()).createBranchWeights(*Weights);
 }
 
-Optional<SmallVector<uint32_t, 8> >
-SwitchInstProfUpdateWrapper::getProfBranchWeights() {
+void SwitchInstProfUpdateWrapper::init() {
   MDNode *ProfileData = getProfBranchWeightsMD(SI);
-  if (!ProfileData)
-    return None;
+  if (!ProfileData) {
+    State = Initialized;
+    return;
+  }
+
+  if (ProfileData->getNumOperands() != SI.getNumSuccessors() + 1) {
+    State = Invalid;
+    if (SwitchInstProfUpdateWrapperStrict)
+      assert(!"number of prof branch_weights metadata operands corresponds to"
+              " number of succesors");
+    return;
+  }
 
   SmallVector<uint32_t, 8> Weights;
   for (unsigned CI = 1, CE = SI.getNumSuccessors(); CI <= CE; ++CI) {
     ConstantInt *C = mdconst::extract<ConstantInt>(ProfileData->getOperand(CI));
     uint32_t CW = C->getValue().getZExtValue();
     Weights.push_back(CW);
   }
-  return Weights;
+  State = Initialized;
+  this->Weights = std::move(Weights);
 }
 
 SwitchInst::CaseIt
 SwitchInstProfUpdateWrapper::removeCase(SwitchInst::CaseIt I) {
   if (Weights) {
     assert(SI.getNumSuccessors() == Weights->size() &&
            "num of prof branch_weights must accord with num of successors");
-    Changed = true;
+    State = Changed;
     // Copy the last case to the place of the removed one and shrink.
     // This is tightly coupled with the way SwitchInst::removeCase() removes
     // the cases in SwitchInst::removeCase(CaseIt).
     Weights.getValue()[I->getCaseIndex() + 1] = Weights.getValue().back();
     Weights.getValue().pop_back();
   }
   return SI.removeCase(I);
 }
 
 void SwitchInstProfUpdateWrapper::addCase(
     ConstantInt *OnVal, BasicBlock *Dest,
     SwitchInstProfUpdateWrapper::CaseWeightOpt W) {
   SI.addCase(OnVal, Dest);
 
+  if (State == Invalid)
+    return;
+
   if (!Weights && W && *W) {
-    Changed = true;
+    State = Changed;
     Weights = SmallVector<uint32_t, 8>(SI.getNumSuccessors(), 0);
     Weights.getValue()[SI.getNumSuccessors() - 1] = *W;
   } else if (Weights) {
-    Changed = true;
+    State = Changed;
     Weights.getValue().push_back(W ? *W : 0);
   }
   if (Weights)
     assert(SI.getNumSuccessors() == Weights->size() &&
            "num of prof branch_weights must accord with num of successors");
 }
 
 SymbolTableList<Instruction>::iterator
 SwitchInstProfUpdateWrapper::eraseFromParent() {
   // Instruction is erased. Mark as unchanged to not touch it in the destructor.
-  Changed = false;
-
+  if (State != Invalid) {
+    State = Initialized;
     if (Weights)
       Weights->resize(0);
+  }
   return SI.eraseFromParent();
 }
 
 SwitchInstProfUpdateWrapper::CaseWeightOpt
 SwitchInstProfUpdateWrapper::getSuccessorWeight(unsigned idx) {
   if (!Weights)
     return None;
   return Weights.getValue()[idx];
 }
 
 void SwitchInstProfUpdateWrapper::setSuccessorWeight(
     unsigned idx, SwitchInstProfUpdateWrapper::CaseWeightOpt W) {
-  if (!W)
+  if (!W || State == Invalid)
     return;
 
   if (!Weights && *W)
     Weights = SmallVector<uint32_t, 8>(SI.getNumSuccessors(), 0);
 
   if (Weights) {
     auto &OldW = Weights.getValue()[idx];
     if (*W != OldW) {
-      Changed = true;
+      State = Changed;
       OldW = *W;
     }
   }
 }
 
 SwitchInstProfUpdateWrapper::CaseWeightOpt
 SwitchInstProfUpdateWrapper::getSuccessorWeight(const SwitchInst &SI,
                                                 unsigned idx) {
   if (MDNode *ProfileData = getProfBranchWeightsMD(SI))
+    if (ProfileData->getNumOperands() == SI.getNumSuccessors() + 1)
       return mdconst::extract<ConstantInt>(ProfileData->getOperand(idx + 1))
           ->getValue()
           .getZExtValue();
 
   return None;
 }
 
 //===----------------------------------------------------------------------===//
 //                        IndirectBrInst Implementation
 //===----------------------------------------------------------------------===//
 

llvm/trunk/unittests/IR/InstructionsTest.cpp

TEST(InstructionsTest, SwitchInst) {
   EXPECT_EQ(3, std::distance(CCI, CCE));
 
   // Make sure that the const iterator is compatible with a const auto ref.
   const auto &Handle = *CCI;
   EXPECT_EQ(1, Handle.getCaseValue()->getSExtValue());
   EXPECT_EQ(BB1.get(), Handle.getCaseSuccessor());
 }
 
+TEST(InstructionsTest, SwitchInstProfUpdateWrapper) {
+  LLVMContext C;
+
+  std::unique_ptr<BasicBlock> BB1, BB2, BB3;
+  BB1.reset(BasicBlock::Create(C));
+  BB2.reset(BasicBlock::Create(C));
+  BB3.reset(BasicBlock::Create(C));
+
+  // We create block 0 after the others so that it gets destroyed first and
+  // clears the uses of the other basic blocks.
+  std::unique_ptr<BasicBlock> BB0(BasicBlock::Create(C));
+
+  auto *Int32Ty = Type::getInt32Ty(C);
+
+  SwitchInst *SI =
+      SwitchInst::Create(UndefValue::get(Int32Ty), BB0.get(), 4, BB0.get());
+  SI->addCase(ConstantInt::get(Int32Ty, 1), BB1.get());
+  SI->addCase(ConstantInt::get(Int32Ty, 2), BB2.get());
+  SI->setMetadata(LLVMContext::MD_prof,
+                  MDBuilder(C).createBranchWeights({ 9, 1, 22 }));
+
+  {
+    SwitchInstProfUpdateWrapper SIW(*SI);
+    EXPECT_EQ(*SIW.getSuccessorWeight(0), 9u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(1), 1u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(2), 22u);
+    SIW.setSuccessorWeight(0, 99u);
+    SIW.setSuccessorWeight(1, 11u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(0), 99u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(1), 11u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(2), 22u);
+  }
+
+  { // Create another wrapper and check that the data persist.
+    SwitchInstProfUpdateWrapper SIW(*SI);
+    EXPECT_EQ(*SIW.getSuccessorWeight(0), 99u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(1), 11u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(2), 22u);
+  }
+
+  // Make prof data invalid by adding one extra weight.
+  SI->setMetadata(LLVMContext::MD_prof, MDBuilder(C).createBranchWeights(
+                                            { 99, 11, 22, 33 })); // extra
+  { // Invalid prof data makes wrapper act as if there were no prof data.
+    SwitchInstProfUpdateWrapper SIW(*SI);
+    ASSERT_FALSE(SIW.getSuccessorWeight(0).hasValue());
+    ASSERT_FALSE(SIW.getSuccessorWeight(1).hasValue());
+    ASSERT_FALSE(SIW.getSuccessorWeight(2).hasValue());
+    SIW.addCase(ConstantInt::get(Int32Ty, 3), BB3.get(), 39);
+    ASSERT_FALSE(SIW.getSuccessorWeight(3).hasValue()); // did not add weight 39
+  }
+
+  { // With added 3rd case the prof data become consistent with num of cases.
+    SwitchInstProfUpdateWrapper SIW(*SI);
+    EXPECT_EQ(*SIW.getSuccessorWeight(0), 99u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(1), 11u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(2), 22u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(3), 33u);
+  }
+
+  // Make prof data invalid by removing one extra weight.
+  SI->setMetadata(LLVMContext::MD_prof,
+                  MDBuilder(C).createBranchWeights({ 99, 11, 22 })); // shorter
+  { // Invalid prof data makes wrapper act as if there were no prof data.
+    SwitchInstProfUpdateWrapper SIW(*SI);
+    ASSERT_FALSE(SIW.getSuccessorWeight(0).hasValue());
+    ASSERT_FALSE(SIW.getSuccessorWeight(1).hasValue());
+    ASSERT_FALSE(SIW.getSuccessorWeight(2).hasValue());
+    SIW.removeCase(SwitchInst::CaseIt(SI, 2));
+  }
+
+  { // With removed 3rd case the prof data become consistent with num of cases.
+    SwitchInstProfUpdateWrapper SIW(*SI);
+    EXPECT_EQ(*SIW.getSuccessorWeight(0), 99u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(1), 11u);
+    EXPECT_EQ(*SIW.getSuccessorWeight(2), 22u);
+  }
+}
+
 TEST(InstructionsTest, CommuteShuffleMask) {
   SmallVector<int, 16> Indices({-1, 0, 7});
   ShuffleVectorInst::commuteShuffleMask(Indices, 4);
   EXPECT_THAT(Indices, testing::ContainerEq(ArrayRef<int>({-1, 4, 3})));
 }
 
 TEST(InstructionsTest, ShuffleMaskQueries) {
   // Create the elements for various constant vectors.