[AssumeBundles] Add API to query a bundles from a use

[Tyker]

Summary

Finding what information is know about a value from a use is generally useful and can be done quickly.

[jdoerfert]

llvm/lib/IR/Instructions.cpp

387

Please explain this in a comment at the beginning of the function. Also comment the magic constants.
Add text to asserts please.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

306

Is this more for testing or do you expect another use case?

llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp

489

What is the runtime we are expecting here?

[Tyker]

fixed formating in test
added comments as requested.

[Tyker]

llvm/lib/IR/Instructions.cpp

387

added comments and messages to assets. hope its enough.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

306

there is another use case for getBundleOpInfoForOperand coming in D73404 in the implementation of Value::dropDroppableUses.

and i think that getKnowledgeFromOperandInAssume is a generally useful functionality especially with the getKnowledgeFromUseInAssume. you can gather all knowledge we have on a value by going through the use list.

llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp

489

On my computer the average for AssumeQueryAPI.getKnowledgeFromUseInAssume is around 116 ms
and without the large test the average is below 1ms.

do you think it is reasonable ? or the test should be reduced in size.

[jdoerfert]

LGTM. One nit below.

llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h

96

Can you add a short comment here as well please.

llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp

489

Should be fine.

llvm/include/llvm/IR/InstrTypes.h

public:
   /// Bundles.  Return the op_iterator pointing to the Use& one past the last
   /// last bundle operand use.
   ///
   /// Each \p OperandBundleDef instance is tracked by a OperandBundleInfo
   /// instance allocated in this User's descriptor.
   op_iterator populateBundleOperandInfos(ArrayRef<OperandBundleDef> Bundles,
                                          const unsigned BeginIndex);
 
+public:
   /// Return the BundleOpInfo for the operand at index OpIdx.
   ///
   /// It is an error to call this with an OpIdx that does not correspond to an
   /// bundle operand.
+  BundleOpInfo &getBundleOpInfoForOperand(unsigned OpIdx);
   const BundleOpInfo &getBundleOpInfoForOperand(unsigned OpIdx) const {
-    for (auto &BOI : bundle_op_infos())
-      if (BOI.Begin <= OpIdx && OpIdx < BOI.End)
-        return BOI;
-
-    llvm_unreachable("Did not find operand bundle for operand!");
+    return const_cast<CallBase *>(this)->getBundleOpInfoForOperand(OpIdx);
   }
 
 protected:
   /// Return the total number of values used in \p Bundles.
   static unsigned CountBundleInputs(ArrayRef<OperandBundleDef> Bundles) {
     unsigned Total = 0;
     for (auto &B : Bundles)
       Total += B.input_size();

llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h

 
 /// Insert into the map all the informations contained in the operand bundles of
 /// the llvm.assume. This should be used instead of hasAttributeInAssume when
 /// many queries are going to be made on the same llvm.assume.
 /// String attributes are not inserted in the map.
 /// If the IR changes the map will be outdated.
 void fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result);
 
+/// Represent one information held inside an operand bundle of an llvm.assume.
+/// AttrKind is the property that hold.
+/// WasOn if not null is that Value for which AttrKind holds.
+/// ArgValue is optionally an argument.
+struct RetainedKnowledge {
+  Attribute::AttrKind AttrKind = Attribute::None;
+  Value *WasOn = nullptr;
+  unsigned ArgValue = 0;
+};
+
+/// Retreive the information help by Assume on the operand at index Idx.
+/// Assume should be an llvm.assume and Idx should be in the operand bundle.
+RetainedKnowledge getKnowledgeFromOperandInAssume(CallInst &Assume,
+                                                  unsigned Idx);
+
+/// Retreive the information help by the Use U of an llvm.assume. the use should
+/// be in the operand bundle.
+inline RetainedKnowledge getKnowledgeFromUseInAssume(const Use *U) {
+  return getKnowledgeFromOperandInAssume(*cast<CallInst>(U->getUser()),
+                                         U->getOperandNo());
+}
+
 //===----------------------------------------------------------------------===//
 // Utilities for testing
 //===----------------------------------------------------------------------===//
 
 /// This pass will try to build an llvm.assume for every instruction in the
 /// function. Its main purpose is testing.
 struct AssumeBuilderPass : public PassInfoMixin<AssumeBuilderPass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 } // namespace llvm
 
 #endif

llvm/lib/IR/Instructions.cpp

for (auto &BOI : bundle_op_infos()) {
     BI++;
   }
 
   assert(BI == Bundles.end() && "Incorrect allocation?");
 
   return It;
 }
 
+CallBase::BundleOpInfo &CallBase::getBundleOpInfoForOperand(unsigned OpIdx) {
+  /// When there isn't many bundles, we do a simple linear search.
+  /// Else fallback to a binary-search that use the fact that bundles usually
+  /// have similar number of argument to get faster convergence.
+  if (bundle_op_info_end() - bundle_op_info_begin() < 8) {
+    for (auto &BOI : bundle_op_infos())
+      if (BOI.Begin <= OpIdx && OpIdx < BOI.End)
+        return BOI;
+
+    llvm_unreachable("Did not find operand bundle for operand!");
+  }
+
+  assert(OpIdx >= arg_size() && "the Idx is not in the operand bundles");
+  assert(bundle_op_info_end() - bundle_op_info_begin() > 0 &&
+         OpIdx < std::prev(bundle_op_info_end())->End &&
+         "The Idx isn't in the operand bundle");
+
+  /// We need a decimal number below and to prevent using floating point numbers
+  /// we use an intergal value multiplied by this constant.
+  constexpr unsigned NumberScaling = 1024;
+
+  bundle_op_iterator Begin = bundle_op_info_begin();
+  bundle_op_iterator End = bundle_op_info_end();
+  bundle_op_iterator Current;
+
+  while (Begin != End) {
+    unsigned ScaledOperandPerBundle =
+        NumberScaling * (std::prev(End)->End - Begin->Begin) / (End - Begin);
+    Current = Begin + (((OpIdx - Begin->Begin) * NumberScaling) /
+                       ScaledOperandPerBundle);
+    if (Current >= End)
+      Current = std::prev(End);
+    assert(Current < End && Current >= Begin &&
+           "the operand bundle doesn't cover every value in the range");
+    if (OpIdx >= Current->Begin && OpIdx < Current->End)
+      break;
+    if (OpIdx >= Current->End)
+      Begin = Current + 1;
+    else
+      End = Current;
+  }
+
+  assert(OpIdx >= Current->Begin && OpIdx < Current->End &&
+         "the operand bundle doesn't cover every value in the range");
+  return *Current;
+}
+
 //===----------------------------------------------------------------------===//
 //                        CallInst Implementation
 //===----------------------------------------------------------------------===//
 
 void CallInst::init(FunctionType *FTy, Value *Func, ArrayRef<Value *> Args,
                     ArrayRef<OperandBundleDef> Bundles, const Twine &NameStr) {
   this->FTy = FTy;
   assert(getNumOperands() == Args.size() + CountBundleInputs(Bundles) + 1 &&

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

if (Lookup == Result.end()) {
       Result[Key] = {Val, Val};
       continue;
     }
     Lookup->second.Min = std::min(Val, Lookup->second.Min);
     Lookup->second.Max = std::max(Val, Lookup->second.Max);
   }
 }
 
+RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(CallInst &AssumeCI,
+                                                        unsigned Idx) {
+  IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
+  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
+         "this function is intended to be used on llvm.assume");
+  CallBase::BundleOpInfo BOI = Assume.getBundleOpInfoForOperand(Idx);
+  RetainedKnowledge Result;
+  Result.AttrKind = Attribute::getAttrKindFromName(BOI.Tag->getKey());
+  Result.WasOn = getValueFromBundleOpInfo(Assume, BOI, BOIE_WasOn);
+  if (BOI.End - BOI.Begin > BOIE_Argument)
+    Result.ArgValue =
+        cast<ConstantInt>(getValueFromBundleOpInfo(Assume, BOI, BOIE_Argument))
+            ->getZExtValue();
+
+  return Result;
+}
+
 PreservedAnalyses AssumeBuilderPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   for (Instruction &I : instructions(F))
     if (Instruction *Assume = BuildAssumeFromInst(&I))
       Assume->insertBefore(&I);
   return PreservedAnalyses::all();
 }

llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp

 //===- KnowledgeRetention.h - utilities to preserve informations *- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/KnowledgeRetention.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/CommandLine.h"
 #include "gtest/gtest.h"
+#include <random>
 
 using namespace llvm;
 
 extern cl::opt<bool> ShouldPreserveAllAttributes;
 
 static void RunTest(
     StringRef Head, StringRef Tail,
     std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
Tests.push_back(std::make_pair(
         Map.clear();
         fillMapFromAssume(*Assume, Map);
         AssertFindExactlyAttributes(Map, New,
                                        "(nonnull|align|dereferenceable)");
         AssertFindExactlyAttributes(Map, Old, "");
       }));
   RunTest(Head, Tail, Tests);
 }
+
+static void RunRandTest(uint64_t Seed, int Size, int MinCount, int MaxCount,
+                        unsigned MaxValue) {
+  LLVMContext C;
+  SMDiagnostic Err;
+
+  std::random_device dev;
+  std::mt19937 Rng(Seed);
+  std::uniform_int_distribution<int> DistCount(MinCount, MaxCount);
+  std::uniform_int_distribution<unsigned> DistValue(0, MaxValue);
+  std::uniform_int_distribution<unsigned> DistAttr(0,
+                                                   Attribute::EndAttrKinds - 1);
+
+  std::unique_ptr<Module> Mod = std::make_unique<Module>("AssumeQueryAPI", C);
+  if (!Mod)
+    Err.print("AssumeQueryAPI", errs());
+
+  std::vector<Type *> TypeArgs;
+  for (int i = 0; i < (Size * 2); i++)
+    TypeArgs.push_back(Type::getInt32PtrTy(C));
+  FunctionType *FuncType =
+      FunctionType::get(Type::getVoidTy(C), TypeArgs, false);
+
+  Function *F =
+      Function::Create(FuncType, GlobalValue::ExternalLinkage, "test", &*Mod);
+  BasicBlock *BB = BasicBlock::Create(C);
+  BB->insertInto(F);
+  Instruction *Ret = ReturnInst::Create(C);
+  BB->getInstList().insert(BB->begin(), Ret);
+  Function *FnAssume = Intrinsic::getDeclaration(Mod.get(), Intrinsic::assume);
+
+  std::vector<Argument *> ShuffledArgs;
+  std::vector<bool> HasArg;
+  for (auto &Arg : F->args()) {
+    ShuffledArgs.push_back(&Arg);
+    HasArg.push_back(false);
+  }
+
+  std::shuffle(ShuffledArgs.begin(), ShuffledArgs.end(), Rng);
+
+  std::vector<OperandBundleDef> OpBundle;
+  OpBundle.reserve(Size);
+  std::vector<Value *> Args;
+  Args.reserve(2);
+  for (int i = 0; i < Size; i++) {
+    int count = DistCount(Rng);
+    int value = DistValue(Rng);
+    int attr = DistAttr(Rng);
+    std::string str;
+    raw_string_ostream ss(str);
+    ss << Attribute::getNameFromAttrKind(
+        static_cast<Attribute::AttrKind>(attr));
+    Args.clear();
+
+    if (count > 0) {
+      Args.push_back(ShuffledArgs[i]);
+      HasArg[i] = true;
+    }
+    if (count > 1)
+      Args.push_back(ConstantInt::get(Type::getInt32Ty(C), value));
+
+    OpBundle.push_back(OperandBundleDef{ss.str().c_str(), std::move(Args)});
+  }
+
+  Instruction *Assume =
+      CallInst::Create(FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}),
+                       std::move(OpBundle));
+  Assume->insertBefore(&F->begin()->front());
+  RetainedKnowledgeMap Map;
+  fillMapFromAssume(*cast<CallInst>(Assume), Map);
+  for (int i = 0; i < (Size * 2); i++) {
+    if (!HasArg[i])
+      continue;
+    RetainedKnowledge K =
+        getKnowledgeFromUseInAssume(&*ShuffledArgs[i]->use_begin());
+    auto LookupIt = Map.find(RetainedKnowledgeKey{K.WasOn, K.AttrKind});
+    ASSERT_TRUE(LookupIt != Map.end());
+    MinMax MM = LookupIt->second;
+    ASSERT_TRUE(MM.Min == MM.Max);
+    ASSERT_TRUE(MM.Min == K.ArgValue);
+  }
+}
+
+TEST(AssumeQueryAPI, getKnowledgeFromUseInAssume) {
+  // // For Fuzzing
+  // std::random_device dev;
+  // std::mt19937 Rng(dev());
+  // while (true) {
+  //   unsigned Seed = Rng();
+  //   dbgs() << Seed << "\n";
+  //   RunRandTest(Seed, 100000, 0, 2, 100);
+  // }
+  RunRandTest(23456, 4, 0, 2, 100);
+  RunRandTest(560987, 25, -3, 2, 100);
+
+  // Large bundles can lead to special cases. this is why this test is soo
+  // large.
+  RunRandTest(9876789, 100000, -0, 7, 100);
+}