[ValueTracking] Add tracking of the alignment assume bundle

[scui]

Summary

This patch is to add the support of the value tracking of the alignment assume bundle.

[etiotto]

Look OK to me but you should also add other reviewers.

[scui]

Anyone has any comments for this patch? Any comments or suggestions are appreciated.

[efriedma]

What does the third operand of an align bundle mean? It doesn't seem to be documented in LangRef.

[jdoerfert]

llvm/lib/Analysis/ValueTracking.cpp

799

You don't want to look for the bundle yourself. Use

SmallVector<Attribute::AttrKind, 2> AttrKinds{Attribute::Align};
if (RetainedKnowledge RK = getKnowledgeValidInContext(V, AttrKinds, CtxI, DT, AC)) {
  ...
}

instead to account for all the interesting cases, e.g., multiple align bundles. Also, use braces for non trivial conditionals. I don't think the return is what you want here, other deductions might still apply. Also, I think we have alignment bundle -> Align information in the AlignmentFromAssumptions.cpp pass. We should create a helper function so this is less magic.

[scui]

In D88669#2328334, @efriedma wrote:

What does the third operand of an align bundle mean? It doesn't seem to be documented in LangRef.

You are right. My understanding is that align bundle of (p, 32, 24) assumes (p+24) is 32-byte aligned. - typo corrected.

[jdoerfert]

In D88669#2328607, @scui wrote:

In D88669#2328334, @efriedma wrote:

What does the third operand of an align bundle mean? It doesn't seem to be documented in LangRef.

You are right. My understanding is that align bundle of (p, 32, 24) assumes (p-24) is 32-byte aligned.

That's correct. It matches the old llvm.assume alignment encoding via instructions.
I mentioned before that we should probably just create a gep with the offset instead of keeping it as a third operand but as long as we do we should have a helper to extract it from the bundle.

[scui]

Rewrite the code using getKnowledgeValidInContext. The code is much cleaner than the previous version, but the result computed is less accurate if the size of the alignment bundle is three. I'll leave the improvement for the future, together with the improvement of value tracking of GEP (an example is that, with the future improvement, we can fold away the branch "br i1 %3, ...") in f2 of the added case assume-align.ll).

[scui]

llvm/lib/Analysis/ValueTracking.cpp

799

You don't want to look for the bundle yourself. Use

SmallVector<Attribute::AttrKind, 2> AttrKinds{Attribute::Align};
if (RetainedKnowledge RK = getKnowledgeValidInContext(V, AttrKinds, CtxI, DT, AC)) {
  ...
}

instead to account for all the interesting cases, e.g., multiple align bundles. Also, use braces for non trivial conditionals. I don't think the return is what you want here, other deductions might still apply. Also, I think we have alignment bundle -> Align information in the AlignmentFromAssumptions.cpp pass. We should create a helper function so this is less magic.

Thanks for the comments and suggestions. I have rewrite the code using getKnowledgeValidInContext. Please take a look. Thanks

[jdoerfert]

Some nits you should address, otherwise LGTM.

In D88669#2330433, @scui wrote:

Rewrite the code using getKnowledgeValidInContext. The code is much cleaner than the previous version, but the result computed is less accurate if the size of the alignment bundle is three. I'll leave the improvement for the future, together with the improvement of value tracking of GEP (an example is that, with the future improvement, we can fold away the branch "br i1 %3, ...") in f2 of the added case assume-align.ll).

FWIW, we should either not encode the offset but use a GEP, or make getKnowledgeValidInContext aware. Either should be a follow up.

llvm/lib/Analysis/ValueTracking.cpp

753

I don't think the parenthesis part in the comment should be here.

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

49

Instead of the comment above, add a TODO/FIXME here describing the two potential solutions to this problem.

[jdoerfert]

Reverse ping. More comments (nits) and notice that we should merge this, D89830 reimplemented it in the meantime, see my comment there too.

llvm/lib/Analysis/ValueTracking.cpp

748

No need for the vector, replace it below with {Attribute::Alignment}.

In D89830 @arichardson uses Log2_32 instead of countTrailingZeros. I think that makes sense since we have a power of 2.

[scui]

Updated to address review comments: I adopted the use of Log2_32 from D89830, removed the parenthesis comment, and added a TODO comment in the test case.

[scui]

In D88669#2345905, @jdoerfert wrote:

Reverse ping. More comments (nits) and notice that we should merge this, D89830 reimplemented it in the meantime, see my comment there too.

Thanks for the comments. I don’t have commit access. If we commit this first, I will need your help to commit for me (“Shimin Cui <[e-mail]>”). Otherwise @arichardson can merge this with D89830 to commit both.

llvm/lib/Analysis/ValueTracking.cpp

static void computeKnownBitsFromAssume(const Value *V, KnownBits &Known,
   if (!Q.AC || !Q.CxtI)
     return;
 
   unsigned BitWidth = Known.getBitWidth();
 
   // Note that the patterns below need to be kept in sync with the code
   // in AssumptionCache::updateAffectedValues.
 
   for (auto &AssumeVH : Q.AC->assumptionsFor(V)) {
     if (!AssumeVH)
       continue;
     CallInst *I = cast<CallInst>(AssumeVH);
     assert(I->getParent()->getParent() == Q.CxtI->getParent()->getParent() &&
            "Got assumption for the wrong function!");
     if (Q.isExcluded(I))
       continue;
 
     // Warning: This loop can end up being somewhat performance sensitive.
     // We're running this loop for once for each value queried resulting in a
     // runtime of ~O(#assumes * #values).
 
     assert(I->getCalledFunction()->getIntrinsicID() == Intrinsic::assume &&
            "must be an assume intrinsic");
 
     Value *Arg = I->getArgOperand(0);
 
     if (Arg == V && isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
       assert(BitWidth == 1 && "assume operand is not i1?");
       Known.setAllOnes();
       return;
     }
     if (match(Arg, m_Not(m_Specific(V))) &&
         isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
       assert(BitWidth == 1 && "assume operand is not i1?");
       Known.setAllZero();
       return;
     }
 
+    // Refine Known set if the pointer alignment is set by assume("align").
+    if (auto AlignOB = I->getOperandBundle("align"))
+      if (AlignOB->Inputs[0].get() == V &&
+          isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+        unsigned Size = AlignOB->Inputs.size();
+        assert((Size == 2 || Size == 3) && "Incorrect align assume");
+        Value *AlignValue = AlignOB->Inputs[1].get();
+        assert(isa<ConstantInt>(AlignValue) && "Expect constant align value");
+        uint64_t Alignment = cast<ConstantInt>(AlignValue)->getZExtValue();
+        assert(isPowerOf2_64(Alignment) && "Alignment should be power of 2");
+        Known.Zero.setLowBits(countTrailingZeros(Alignment));
+        if (Size == 3) {
+          Value *OffsetValue = AlignOB->Inputs[2].get();
+          assert(isa<ConstantInt>(OffsetValue) &&
+                 "Expect constant offset value");
+          uint64_t Offset = cast<ConstantInt>(OffsetValue)->getZExtValue();
+          Known.Zero &= ~(Alignment - Offset) | ~(Alignment - 1);
+          Known.One.setLowBits(countTrailingZeros(Alignment));
+          Known.One &= (Alignment - Offset) | ~(Alignment - 1);
+        }
+        return;
+      }
+
     // The remaining tests are all recursive, so bail out if we hit the limit.
     if (Depth == MaxAnalysisRecursionDepth)
       continue;
 
     ICmpInst *Cmp = dyn_cast<ICmpInst>(Arg);
     if (!Cmp)
       continue;
 

llvm/test/Analysis/BasicAA/featuretest.ll

 ; NO_ASSUME-NEXT:    [[C1:%.*]] = getelementptr i32, i32* [[A:%.*]], i64 1
 ; NO_ASSUME-NEXT:    [[C:%.*]] = bitcast i32* [[C1]] to i8*
 ; NO_ASSUME-NEXT:    store i8 42, i8* [[C]], align 1
 ; NO_ASSUME-NEXT:    ret i32 0
 ;
 ; USE_ASSUME-LABEL: @gep_distance_test3(
 ; USE_ASSUME-NEXT:    [[C1:%.*]] = getelementptr i32, i32* [[A:%.*]], i64 1
 ; USE_ASSUME-NEXT:    [[C:%.*]] = bitcast i32* [[C1]] to i8*
-; USE_ASSUME-NEXT:    store i8 42, i8* [[C]], align 1
+; USE_ASSUME-NEXT:    store i8 42, i8* [[C]], align 4
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[A]], i64 4), "nonnull"(i32* [[A]]), "align"(i32* [[A]], i64 4) ]
 ; USE_ASSUME-NEXT:    ret i32 0
 ;
   %X = load i32, i32* %A
   %B = bitcast i32* %A to i8*
   %C = getelementptr i8, i8* %B, i64 4
   store i8 42, i8* %C
   %Y = load i32, i32* %A

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

+; RUN: opt -S -passes=instcombine,simplify-cfg < %s 2>&1 | FileCheck %s
+
+declare void @llvm.assume(i1 noundef)
+
+define void @f1(i8* %a) {
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[PTR:%.*]] = getelementptr inbounds i8, i8* [[A:%.*]], i64 4
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[PTR]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = and i64 [[TMP0]], 3
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i64 [[TMP1]], 0
+; CHECK-NEXT:    br i1 [[TMP2]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i8* [[PTR]], i64 4) ]
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i8* [[PTR]] to i32*
+; CHECK-NEXT:    store i32 4, i32* [[TMP3]], align 4
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       if.end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %ptr = getelementptr inbounds i8, i8* %a, i64 4
+  %0 = ptrtoint i8* %ptr to i64
+  %1 = and i64 %0, 3
+  %2 = icmp eq i64 %1, 0
+  br i1 %2, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @llvm.assume(i1 true) [ "align"(i8* %ptr, i64 4) ]
+  %3 = ptrtoint i8* %ptr to i64
+  %4 = and i64 %3, 3
+  %5 = icmp eq i64 %4, 0
+  br i1 %5, label %if.then1, label %if.else1
+
+if.then1:                                         ; preds = %if.then
+  %6 = bitcast i8* %ptr to i32*
+  store i32 4, i32* %6, align 4
+  br label %if.end
+
+if.else1:                                         ; preds = %if.then
+  store i8 1, i8* %ptr, align 1
+  br label %if.end
+
+if.end:                                           ; preds = %if.then1, %if.else1, %entry
+  ret void
+}
+
+define void @f2(i8* %a) {
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i8* [[A:%.*]], i64 32, i32 24) ]
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i8, i8* [[A]], i64 8
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i8* [[TMP0]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = and i64 [[TMP1]], 8
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i64 [[TMP2]], 0
+; CHECK-NEXT:    br i1 [[TMP3]], label [[IF_THEN1:%.*]], label [[IF_ELSE1:%.*]]
+;
+entry:
+  call void @llvm.assume(i1 true) [ "align"(i8* %a, i64 32, i32 24) ]
+  %0 = getelementptr inbounds i8, i8* %a, i64 8
+  %1 = ptrtoint i8* %0 to i64
+  %2 = and i64 %1, 15
+  %3 = icmp eq i64 %2, 0
+  br i1 %3, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  %4 = bitcast i8* %0 to i64*
+  store i64 16, i64* %4, align 4
+  br label %if.end
+
+if.else:                                          ; preds = %entry
+  store i8 1, i8* %0, align 1
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then
+  ret void
+}
+