Add Query API for llvm.assume holding attributes

[jdoerfert]

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

75

Documentation missing.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

191

Commonly we have messages in an assert, also below.

216

I mentioned in D72475 now that the name should not be a part of the key as it can change.

243

The above seems complicated. Can you use helpers or comments to make it easier to read?

270

Where is this used?

[Tyker]

addressed comments.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

243

i added and used a helper to simplify this and other pieces of code is it enough ?

270

this is used for testing ShouldPreserveAllAttributes option from unittests.

[jdoerfert]

I won't review this properly till next week. Generally looks good so you can work on the next patch if you want to.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

255

What is the use case for IsOn = nullptr? If there is none, we should make it a reference so it's obvious.

270

If it's only used in the unit test can we move it there?

[Tyker]

In D72885#1827778, @jdoerfert wrote:

I won't review this properly till next week. Generally looks good so you can work on the next patch if you want to.

i started working on ignoring or dropping uses in llvm.assume when a transformation is use sensitive.

In D72475#1821013, @jdoerfert wrote:

Also because of that, and we do not have settled on this scheme in the RFC thread.

i would like to have confirmation that the operand bundles representation is what we are going for ?
i don't recall it being a point of contention but it wasn't explicitly approved. and you mentioned in that the issue as not settled.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

255

currently IsOn = nullptr is for attributes that where on the call of a function or the declaration of a function.

this case only occurs when ShouldPreserveAllAttributes is set.

270

it needs to have access to the ShouldPreserveAllAttributes flag. which we probably shouldn't make global.

[jdoerfert]

In D72885#1832612, @Tyker wrote:

In D72885#1827778, @jdoerfert wrote:

I won't review this properly till next week. Generally looks good so you can work on the next patch if you want to.

i started working on ignoring or dropping uses in llvm.assume when a transformation is use sensitive.

Great! Can you share it on phab so I can reference it in my email (see below).

In D72475#1821013, @jdoerfert wrote:

Also because of that, and we do not have settled on this scheme in the RFC thread.

i would like to have confirmation that the operand bundles representation is what we are going for ?
i don't recall it being a point of contention but it wasn't explicitly approved. and you mentioned in that the issue as not settled.

Correct. I'll send another email today, describing our progress and hoping ppl agree this is a good solution :)

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

255

I would have expected the IsOn to be the called value or respectively function then. Is there a reason for one or the other?

270

Now I understand but I still think this is problematic.

What if we add a second call into KnowledgeRetention.h where the value of ShouldPreserveAllAttributes is passed in directly. The existing interface will call that internally and use the command line flag value. Instead of the command line flag value, existing uses will then use the new argument. That way we can determine the value from the call site or via the command line.

What do you think?

[Tyker]

In D72885#1834094, @jdoerfert wrote:

In D72885#1832612, @Tyker wrote:

In D72885#1827778, @jdoerfert wrote:

I won't review this properly till next week. Generally looks good so you can work on the next patch if you want to.

i started working on ignoring or dropping uses in llvm.assume when a transformation is use sensitive.

Great! Can you share it on phab so I can reference it in my email (see below).

i added it to phab (D73404) after cleaning it up. sorry i toke soo long i haven't had the time i wanted to work on this.

In D72475#1821013, @jdoerfert wrote:

Also because of that, and we do not have settled on this scheme in the RFC thread.

i would like to have confirmation that the operand bundles representation is what we are going for ?
i don't recall it being a point of contention but it wasn't explicitly approved. and you mentioned in that the issue as not settled.

Correct. I'll send another email today, describing our progress and hoping ppl agree this is a good solution :)

thanks for the mail.

[Tyker]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

255

i am staring to think the attribute on call and function is perhaps too small of a use case to be worth supporting. i haven't yet found any other use case than the cold attribute that said i haven't be looking for it.

but the idea was that Users of that API wouldn't need to care or know what function was called here. but just that this function has X attribute. i think it loose its last bit of usefulness if users need to know ahead of time that there was a calls to function F before querying it.

[jdoerfert]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

255

I can see your argument. Maybe the "map API" will allow your use case even if we retain the functions that carried the attributes.
I am hesitant to throw away information here if there is another way. When we query *all* information the assume stores in a map the keys could either be attributes or values (probably both make sense). In the keys = attributes version the user can look for "cold" and act on it. Arguably, we don't even need the map API for this. We can have a define that if IsOn is null any occurrence of the attribute is returned even though we store the original "IsOn" value internally.

[jdoerfert]

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

81

This is not needed anymore, or is it? (I hope we can replace it somehow if it is.)

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

69

Typo.

[Tyker]

addressed comments.

[jdoerfert]

This looks for with minor comments and one remaining question: Do we use the value name in the sorting? If so, are we sure arbitrary renaming of values will not be a problem? If so, does it help to use the value name in the key?

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

60

Nit: Can you move the enum class (with a short doc) above the documentation of the function. It is a bit confusing this way and I'm not sure Doxygen will pick up on it they way we want it to.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

88

This still uses the value name in the key, right? Does this work if the values are renamed arbitrarily?

194

Nit: Why do we allow None here?

205

Nit: Move this down closer to the use sites. Maybe also replace Elem with Value to make it clear it returns an llvm value.

[Tyker]

addressed comments

the Name is used in the key to prevent non-deterministic order of properties in llvm.assume.

the DenseSet used for uniquing, to prevent reliance on the actual hashing implementation is randomized on each run of the program. so if we don't fully order element as we do BuildAssumeFromInst can build llvm.assume where the order of constaint in the bundle may change from run to run. this is not usually visible. but this is going to be an issue for the tests as they are usually simple string matches on the output.

[Tyker]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

88

The Name is still used in the key. but this is only to get a full ordering without depending on pointer values or other sources of non-determinism. the lookup algorithm doesn't depend on Names.

if there is an other deterministic source that can be used for ordering that you would feel more comfortable with, it can be used instead.

for example the current unit test passes when we use the reverse of the name. also unit test have a test with replaceAllUsesWith.

[jdoerfert]

In D72885#1879532, @Tyker wrote:

addressed comments

the Name is used in the key to prevent non-deterministic order of properties in llvm.assume.

the DenseSet used for uniquing, to prevent reliance on the actual hashing implementation is randomized on each run of the program. so if we don't fully order element as we do BuildAssumeFromInst can build llvm.assume where the order of constaint in the bundle may change from run to run. this is not usually visible. but this is going to be an issue for the tests as they are usually simple string matches on the output.

[Apologies for the wall of text, I wrote down my thoughts as I think it might help (in the future).]

At first this made sense to me and put my worries to rest. I agree that if we only use the names to make it deterministic but not rely on them in any other way we should be OK. That said, it is not deterministic in the absence of names, correct? I mean, if we remove the names of *all* variables, every getName() invocation will return the empty string. While determinism is a describable property for testing, it is crucial for compilation in general. However, I think the compilation part is not a problem right now as we can only query the extreme values (largest/lowest) for a pair attribute + value. In which order all the operand bundles that store a value for such a pair are internally doesn't matter, the largest/lowest is always the same. Now for testing I am also not too worried. We can test with the CHECK-DAG feature if we have ambiguous cases or force values to have names to make it deterministic.

Long story short, LGTM.

[jdoerfert]

Now we need to query the operand bundle information somewhere and actually emit it ;)

[nikic]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

232

This causes a warning:

warning: comparison of unsigned expression < 0 is always false [-Wtype-limits]

Can you please check whether this condition can be simply dropped, or some other change is needed?

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

 // transformation may destroy or worsen information that can be derived from the
 // IR.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_ASSUMEBUILDER_H
 #define LLVM_TRANSFORMS_UTILS_ASSUMEBUILDER_H
 
+#include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// Build a call to llvm.assume to preserve informations that can be derived
 /// from the given instruction.
 /// If no information derived from \p I, this call returns null.
 /// The returned instruction is not inserted anywhere.
 CallInst *BuildAssumeFromInst(const Instruction *I, Module *M);
 inline CallInst *BuildAssumeFromInst(Instruction *I) {
   return BuildAssumeFromInst(I, I->getModule());
 }
 
+/// It is possible to have multiple Value for the argument of an attribute in
+/// the same llvm.assume on the same llvm::Value. This is rare but need to be
+/// dealt with.
+enum class AssumeQuery {
+  Highest, ///< Take the highest value available.
+  Lowest,  ///< Take the lowest value available.
+};
+
+/// Query the operand bundle of an llvm.assume to find a single attribute of
+/// the specified kind applied on a specified Value.
+///
+/// This has a non-constant complexity. It should only be used when a single
+/// attribute is going to be queried.
+///
+/// Return true iff the queried attribute was found.
+/// If ArgVal is set. the argument will be stored to ArgVal.
+bool hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn, StringRef AttrName,
+                          uint64_t *ArgVal = nullptr,
+                          AssumeQuery AQR = AssumeQuery::Highest);
+inline bool hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
+                                 Attribute::AttrKind Kind,
+                                 uint64_t *ArgVal = nullptr,
+                                 AssumeQuery AQR = AssumeQuery::Highest) {
+  return hasAttributeInAssume(
+      AssumeCI, IsOn, Attribute::getNameFromAttrKind(Kind), ArgVal, AQR);
+}
+
+/// TODO: Add an function to create/fill a map from the bundle when users intend
+/// to make many different queries on the same bundles. to be used for example
+/// in the Attributor.
+
+//===----------------------------------------------------------------------===//
+// 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/Transforms/Utils/KnowledgeRetention.cpp

 #include "llvm/ADT/DenseSet.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
 
-namespace {
-
 cl::opt<bool> ShouldPreserveAllAttributes(
     "assume-preserve-all", cl::init(false), cl::Hidden,
     cl::desc("enable preservation of all attrbitues. even those that are "
              "unlikely to be usefull"));
 
+namespace {
+
 struct AssumedKnowledge {
   const char *Name;
   Value *Argument;
   enum {
     None,
     Empty,
     Tombstone,
   };
return LHS.WasOn == RHS.WasOn && LHS.Name == RHS.Name &&
            LHS.Argument == RHS.Argument;
   }
 };
 
 } // namespace llvm
 
 namespace {
 
+/// Index of elements in the operand bundle.
+/// If the element exist it is guaranteed to be what is specified in this enum
+/// but it may not exist.
+enum BundleOpInfoElem {
+  BOIE_WasOn = 0,
+  BOIE_Argument = 1,
+};
+
 /// Deterministically compare OperandBundleDef.
 /// The ordering is:
-/// - by the name of the attribute, (doesn't change)
-/// - then by the Value of the argument, (doesn't change)
+/// - by the attribute's name aka operand bundle tag, (doesn't change)
+/// - then by the numeric Value of the argument, (doesn't change)
 /// - lastly by the Name of the current Value it WasOn. (may change)
 /// This order is deterministic and allows looking for the right kind of
 /// attribute with binary search. However finding the right WasOn needs to be
-/// done via linear search because values can get remplaced.
+/// done via linear search because values can get replaced.
 bool isLowerOpBundle(const OperandBundleDef &LHS, const OperandBundleDef &RHS) {
   auto getTuple = [](const OperandBundleDef &Op) {
     return std::make_tuple(
         Op.getTag(),
-        Op.input_size() < 2
+        Op.input_size() <= BOIE_Argument
             ? 0
-            : cast<ConstantInt>(*std::next(Op.input_begin()))->getZExtValue(),
-        Op.input_size() < 1 ? StringRef("") : (*Op.input_begin())->getName());
+            : cast<ConstantInt>(*(Op.input_begin() + BOIE_Argument))
+                  ->getZExtValue(),
+         Op.input_size() <= BOIE_WasOn
+            ? StringRef("")
+            : (*(Op.input_begin() + BOIE_WasOn))->getName());
   };
   return getTuple(LHS) < getTuple(RHS);
 }
 
 /// This class contain all knowledge that have been gather while building an
 /// llvm.assume and the function to manipulate it.
 struct AssumeBuilderState {
   Module *M;

 } // namespace
 
 CallInst *llvm::BuildAssumeFromInst(const Instruction *I, Module *M) {
   AssumeBuilderState Builder(M);
   Builder.addInstruction(I);
   return Builder.build();
 }
 
+#ifndef NDEBUG
+
+static bool isExistingAttribute(StringRef Name) {
+  return StringSwitch<bool>(Name)
+#define GET_ATTR_NAMES
+#define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) .Case(#DISPLAY_NAME, true)
+#include "llvm/IR/Attributes.inc"
+      .Default(false);
+}
+
+#endif
+
+bool llvm::hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
+                                StringRef AttrName, uint64_t *ArgVal,
+                                AssumeQuery AQR) {
+  IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
+  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
+         "this function is intended to be used on llvm.assume");
+  assert(isExistingAttribute(AttrName) && "this attribute doesn't exist");
+  assert((ArgVal == nullptr || Attribute::doesAttrKindHaveArgument(
+                                   Attribute::getAttrKindFromName(AttrName))) &&
+         "requested value for an attribute that has no argument");
+  if (Assume.bundle_op_infos().empty())
+    return false;
+
+  CallInst::bundle_op_iterator Lookup;
+
+  /// The right attribute can be found by binary search. After this finding the
+  /// right WasOn needs to be done via linear search.
+  /// Element have been ordered by argument value so the first we find is the
+  /// one we need.
+  if (AQR == AssumeQuery::Lowest)
+    Lookup =
+        llvm::lower_bound(Assume.bundle_op_infos(), AttrName,
+                          [](const CallBase::BundleOpInfo &BOI, StringRef RHS) {
+                            assert(isExistingAttribute(BOI.Tag->getKey()) &&
+                                   "this attribute doesn't exist");
+                            return BOI.Tag->getKey() < RHS;
+                          });
+  else
+    Lookup = std::prev(
+        llvm::upper_bound(Assume.bundle_op_infos(), AttrName,
+                          [](StringRef LHS, const CallBase::BundleOpInfo &BOI) {
+                            assert(isExistingAttribute(BOI.Tag->getKey()) &&
+                                   "this attribute doesn't exist");
+                            return LHS < BOI.Tag->getKey();
+                          }));
+
+  auto getValueFromBundleOpInfo = [&Assume](const CallBase::BundleOpInfo &BOI,
+                                            unsigned Idx) {
+    assert(BOI.End - BOI.Begin > Idx && "index out of range");
+    return (Assume.op_begin() + BOI.Begin + Idx)->get();
+  };
+
+  if (Lookup == Assume.bundle_op_info_end() ||
+      Lookup->Tag->getKey() != AttrName)
+    return false;
+  if (IsOn) {
+    if (Lookup->End - Lookup->Begin < BOIE_WasOn)
+      return false;
+    while (true) {
+      if (Lookup == Assume.bundle_op_info_end() ||
+          Lookup->Tag->getKey() != AttrName)
+        return false;
+      if (getValueFromBundleOpInfo(*Lookup, BOIE_WasOn) == IsOn)
+        break;
+      if (AQR == AssumeQuery::Highest &&
+          Lookup == Assume.bundle_op_info_begin())
+        return false;
+      Lookup = Lookup + (AQR == AssumeQuery::Lowest ? 1 : -1);
+    }
+  }
+
+  if (Lookup->End - Lookup->Begin < BOIE_Argument)
+    return true;
+  if (ArgVal)
+    *ArgVal =
+        cast<ConstantInt>(getValueFromBundleOpInfo(*Lookup, BOIE_Argument))
+            ->getZExtValue();
+  return true;
+}
+
 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/CMakeLists.txt

add_llvm_unittest(UtilsTests
   ASanStackFrameLayoutTest.cpp
   BasicBlockUtilsTest.cpp
   CallPromotionUtilsTest.cpp
   CloningTest.cpp
   CodeExtractorTest.cpp
   CodeMoverUtilsTest.cpp
   FunctionComparatorTest.cpp
   IntegerDivisionTest.cpp
+  KnowledgeRetentionTest.cpp
   LocalTest.cpp
   LoopRotationUtilsTest.cpp
   LoopUtilsTest.cpp
   SizeOptsTest.cpp
   SSAUpdaterBulkTest.cpp
   UnrollLoopTest.cpp
   ValueMapperTest.cpp
   VFABIUtils.cpp
   )

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/Support/Regex.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/CommandLine.h"
+#include "gtest/gtest.h"
+
+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) {
+  std::string IR;
+  IR.append(Head.begin(), Head.end());
+  for (auto &Elem : Tests)
+    IR.append(Elem.first.begin(), Elem.first.end());
+  IR.append(Tail.begin(), Tail.end());
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+  if (!Mod)
+    Err.print("AssumeQueryAPI", errs());
+  unsigned Idx = 0;
+  for (Instruction &I : (*Mod->getFunction("test")->begin())) {
+    if (Idx < Tests.size())
+      Tests[Idx].second(&I);
+    Idx++;
+  }
+}
+
+void AssertMatchesExactlyAttributes(CallInst *Assume, Value *WasOn,
+                                    StringRef AttrToMatch) {
+  Regex Reg(AttrToMatch);
+  SmallVector<StringRef, 1> Matches;
+  for (StringRef Attr : {
+#define GET_ATTR_NAMES
+#define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) StringRef(#DISPLAY_NAME),
+#include "llvm/IR/Attributes.inc"
+       }) {
+    bool ShouldHaveAttr = Reg.match(Attr, &Matches) && Matches[0] == Attr;
+    if (ShouldHaveAttr != hasAttributeInAssume(*Assume, WasOn, Attr))
+      ASSERT_TRUE(false);
+  }
+}
+
+void AssertHasTheRightValue(CallInst *Assume, Value *WasOn,
+                            Attribute::AttrKind Kind, unsigned Value, bool Both,
+                            AssumeQuery AQ = AssumeQuery::Highest) {
+  if (!Both) {
+    uint64_t ArgVal = 0;
+    ASSERT_TRUE(hasAttributeInAssume(*Assume, WasOn, Kind, &ArgVal, AQ));
+    ASSERT_EQ(ArgVal, Value);
+    return;
+  }
+  uint64_t ArgValLow = 0;
+  uint64_t ArgValHigh = 0;
+  bool ResultLow = hasAttributeInAssume(*Assume, WasOn, Kind, &ArgValLow,
+                                        AssumeQuery::Lowest);
+  bool ResultHigh = hasAttributeInAssume(*Assume, WasOn, Kind, &ArgValHigh,
+                                         AssumeQuery::Highest);
+  if (ResultLow != ResultHigh)
+    ASSERT_TRUE(false);
+  if (ArgValLow != Value || ArgValLow != ArgValHigh)
+    ASSERT_EQ(ArgValLow, Value);
+}
+
+TEST(AssumeQueryAPI, Basic) {
+  StringRef Head =
+      "declare void @llvm.assume(i1)\n"
+      "declare void @func(i32*, i32*)\n"
+      "declare void @func1(i32*, i32*, i32*, i32*)\n"
+      "declare void @func_many(i32*) \"no-jump-tables\" nounwind "
+      "\"less-precise-fpmad\" willreturn norecurse\n"
+      "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3) {\n";
+  StringRef Tail = "ret void\n"
+                   "}";
+  std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
+      Tests;
+  Tests.push_back(std::make_pair(
+      "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
+      "8 noalias %P1)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(0),
+                                       "(nonnull|align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1),
+                                       "(noalias|align)");
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Dereferenceable, 16, true);
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Alignment, 4, true);
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Alignment, 4, true);
+      }));
+  Tests.push_back(std::make_pair(
+      "call void @func1(i32* nonnull align 32 dereferenceable(48) %P, i32* "
+      "nonnull "
+      "align 8 dereferenceable(28) %P, i32* nonnull align 64 "
+      "dereferenceable(4) "
+      "%P, i32* nonnull align 16 dereferenceable(12) %P)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(0),
+                                       "(nonnull|align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1),
+                                       "(nonnull|align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(2),
+                                       "(nonnull|align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(3),
+                                       "(nonnull|align|dereferenceable)");
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Dereferenceable, 48, false,
+                               AssumeQuery::Highest);
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Alignment, 64, false,
+                               AssumeQuery::Highest);
+        AssertHasTheRightValue(Assume, I->getOperand(1),
+                               Attribute::AttrKind::Alignment, 64, false,
+                               AssumeQuery::Highest);
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Dereferenceable, 4, false,
+                               AssumeQuery::Lowest);
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Alignment, 8, false,
+                               AssumeQuery::Lowest);
+        AssertHasTheRightValue(Assume, I->getOperand(1),
+                               Attribute::AttrKind::Alignment, 8, false,
+                               AssumeQuery::Lowest);
+      }));
+  Tests.push_back(std::make_pair(
+      "call void @func_many(i32* align 8 %P1) cold\n", [](Instruction *I) {
+        ShouldPreserveAllAttributes.setValue(true);
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+        AssertMatchesExactlyAttributes(
+            Assume, nullptr,
+            "(align|no-jump-tables|less-precise-fpmad|"
+            "nounwind|norecurse|willreturn|cold)");
+        ShouldPreserveAllAttributes.setValue(false);
+      }));
+  Tests.push_back(
+      std::make_pair("call void @llvm.assume(i1 true)\n", [](Instruction *I) {
+        CallInst *Assume = cast<CallInst>(I);
+        AssertMatchesExactlyAttributes(Assume, nullptr, "");
+      }));
+  Tests.push_back(std::make_pair(
+      "call void @func1(i32* readnone align 32 "
+      "dereferenceable(48) noalias %P, i32* "
+      "align 8 dereferenceable(28) %P1, i32* align 64 "
+      "dereferenceable(4) "
+      "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+        AssertMatchesExactlyAttributes(
+            Assume, I->getOperand(0),
+            "(readnone|align|dereferenceable|noalias)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1),
+                                       "(align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(2),
+                                       "(align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(3),
+                                       "(nonnull|align|dereferenceable)");
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Alignment, 32, true);
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Dereferenceable, 48, true);
+        AssertHasTheRightValue(Assume, I->getOperand(1),
+                               Attribute::AttrKind::Dereferenceable, 28, true);
+        AssertHasTheRightValue(Assume, I->getOperand(1),
+                               Attribute::AttrKind::Alignment, 8, true);
+        AssertHasTheRightValue(Assume, I->getOperand(2),
+                               Attribute::AttrKind::Alignment, 64, true);
+        AssertHasTheRightValue(Assume, I->getOperand(2),
+                               Attribute::AttrKind::Dereferenceable, 4, true);
+        AssertHasTheRightValue(Assume, I->getOperand(3),
+                               Attribute::AttrKind::Alignment, 16, true);
+        AssertHasTheRightValue(Assume, I->getOperand(3),
+                               Attribute::AttrKind::Dereferenceable, 12, true);
+      }));
+
+  /// Keep this test last as it modifies the function.
+  Tests.push_back(std::make_pair(
+      "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
+      "8 noalias %P1)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+        Value *New = I->getFunction()->getArg(3);
+        Value *Old = I->getOperand(0);
+        AssertMatchesExactlyAttributes(Assume, New, "");
+        AssertMatchesExactlyAttributes(Assume, Old,
+                                       "(nonnull|align|dereferenceable)");
+        Old->replaceAllUsesWith(New);
+        AssertMatchesExactlyAttributes(Assume, New,
+                                       "(nonnull|align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, Old, "");
+      }));
+  RunTest(Head, Tail, Tests);
+}