[instcombine] Simplify separate_storage hints.

[davidtgoldblatt]

Summary

Before this change, we call getUnderlyingObject on each separate_storage
operand on every alias() call (potentially requiring lots of pointer
chasing). Instead, we rewrite the assumptions in instcombine to do this
pointer-chasing once.

We still leave the getUnderlyingObject calls in alias(), just expecting
them to be no-ops much of the time. This is relatively fast (just a
couple dyn_casts with no pointer chasing) and avoids making alias
analysis results depend on whether or not instcombine has been run.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2523

This should be done unconditionally, independently of the BasicAA flag.

2532

Use replaceUse() instead, drop code below.

[StephenFan]

llvm/test/Transforms/InstCombine/assume-separate_storage.ll

3

Is 2>&1 necessary?

[davidtgoldblatt]

Unconditionally do fold, use replaceUse, remove unnecessary pipe redirection.

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2529

2530

Drop braces.

2531

Sounds like OBO.Inputs should be a MutableArrayRef...

[davidtgoldblatt]

Nits.

[davidtgoldblatt]

Would you mind committing for me? I don't have commit access.

I'll take a look at doing the MutableArrayRef change as well.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

for (auto &AssumeVH : AC.assumptions()) {
       AssumeInst *Assume = cast<AssumeInst>(AssumeVH);
 
       for (unsigned Idx = 0; Idx < Assume->getNumOperandBundles(); Idx++) {
         OperandBundleUse OBU = Assume->getOperandBundleAt(Idx);
         if (OBU.getTagName() == "separate_storage") {
           assert(OBU.Inputs.size() == 2);
           const Value *Hint1 = OBU.Inputs[0].get();
           const Value *Hint2 = OBU.Inputs[1].get();
+          // This is often a no-op; instcombine rewrites this for us. No-op
+          // getUnderlyingObject calls are fast, though.
           const Value *HintO1 = getUnderlyingObject(Hint1);
           const Value *HintO2 = getUnderlyingObject(Hint2);
 
           if (((O1 == HintO1 && O2 == HintO2) ||
                (O1 == HintO2 && O2 == HintO1)) &&
               isValidAssumeForContext(Assume, CtxI, DT))
             return AliasResult::NoAlias;
         }

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (match(IIOperand, m_ICmp(Pred, m_Instruction(LHS), m_Zero())) &&
       MDNode *MD = MDNode::get(II->getContext(), std::nullopt);
       LHS->setMetadata(LLVMContext::MD_nonnull, MD);
       return RemoveConditionFromAssume(II);
 
       // TODO: apply nonnull return attributes to calls and invokes
       // TODO: apply range metadata for range check patterns?
     }
 
+    // Separate storage assumptions apply to the underlying allocations, not any
+    // particular pointer within them. When evaluating the hints for AA purposes
+    // we getUnderlyingObject them; by precomputing the answers here we can
+    // avoid having to do so repeatedly there.
+    for (unsigned Idx = 0; Idx < II->getNumOperandBundles(); Idx++) {
+      OperandBundleUse OBU = II->getOperandBundleAt(Idx);
+      if (OBU.getTagName() == "separate_storage") {
+        assert(OBU.Inputs.size() == 2);
+        auto MaybeSimplifyHint = [&](const Use &U) {
+          Value *Hint = U.get();
+          // Not having a limit is safe because InstCombine removes unreachable
+          // code.
+          Value *UnderlyingObject = getUnderlyingObject(Hint, /*MaxLookup*/ 0);
+          if (Hint != UnderlyingObject)
+            replaceUse(const_cast<Use &>(U), UnderlyingObject);
+        };
+        MaybeSimplifyHint(OBU.Inputs[0]);
+        MaybeSimplifyHint(OBU.Inputs[1]);
+      }
+    }
+
     // Convert nonnull assume like:
     // %A = icmp ne i32* %PTR, null
     // call void @llvm.assume(i1 %A)
     // into
     // call void @llvm.assume(i1 true) [ "nonnull"(i32* %PTR) ]
     if (EnableKnowledgeRetention &&
         match(IIOperand, m_Cmp(Pred, m_Value(A), m_Zero())) &&
         Pred == CmpInst::ICMP_NE && A->getType()->isPointerTy()) {

llvm/test/Transforms/InstCombine/assume-separate_storage.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=instcombine < %s | FileCheck %s
+
+declare void @llvm.assume(i1 noundef)
+
+; Just something to let us check that separate_storage bundles don't break
+; anything when given an operand that's not an instruction to fold.
+@some_global = global i32 777
+
+define void @simple_folding(ptr %a, ptr %b) {
+; CHECK-LABEL: @simple_folding(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "separate_storage"(ptr [[A:%.*]], ptr [[B:%.*]]) ]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %p1 = getelementptr i8, ptr %a, i64 123
+  %p2 = getelementptr i8, ptr %b, i64 777
+  call void @llvm.assume(i1 1) ["separate_storage"(ptr %p1, ptr %p2)]
+  ret void
+}
+
+define i64 @folds_removed_operands(ptr %a, ptr %b, i64 %n1, i64 %n2) {
+; CHECK-LABEL: @folds_removed_operands(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[REASS_ADD:%.*]] = shl i64 [[N2:%.*]], 1
+; CHECK-NEXT:    [[Y:%.*]] = add i64 [[REASS_ADD]], [[N1:%.*]]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "separate_storage"(ptr [[A:%.*]], ptr [[B:%.*]]) ]
+; CHECK-NEXT:    ret i64 [[Y]]
+;
+entry:
+  ; Ordinarily, n1 + n2 + n2 would get canonicalized into n1 + (n2 << 1) unless
+  ; there's another use of n1 + n2. Make sure that we remember to put removed
+  ; arguments to separate_storage bundles back on the worklist.
+  %x = add i64 %n1, %n2
+  %y = add i64 %x, %n2
+  %p1 = getelementptr i8, ptr %a, i64 %x
+  call void @llvm.assume(i1 1) ["separate_storage"(ptr %p1, ptr %b)]
+  ret i64 %y
+}
+
+define void @handles_globals(ptr %a) {
+; CHECK-LABEL: @handles_globals(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "separate_storage"(ptr [[A:%.*]], ptr @some_global) ]
+; CHECK-NEXT:    ret void
+;
+  %derived = getelementptr i8, ptr @some_global, i65 3
+  call void @llvm.assume(i1 1) ["separate_storage"(ptr %a, ptr %derived)]
+  ret void
+}