[InstCombine] Canonicalize `getelementptr` patterns to `@llvm.ptrmask`

[goldstein.w.n]

Summary

[InstCombine] Fold getelementptr patterns to @llvm.ptrmask

@llvm.ptrmask is documented as being equivilent to:
(getelementptr i8 Ptr, (sub (and (ptrtoint Ptr), Mask), (ptrtoint Ptr)))

See: https://llvm.org/docs/LangRef.html#llvm-ptrmask-intrinsic

As well clang's __builtin_align_up and __builtin_align_down emit
patterns that can be folded to pointer mask.

Handle all three.

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2247

This is not legal, because it does not preserve pointer provenance. This one even fails to preserve it in a particularly catastrophic way, because the resulting pointer has nullary provenance (i.e. it is UB to perform any accesses through it).

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2247

I see, is the codes to create @llvm.ptrmask okay? Or should that too be abandoned?

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2247

If you start at the documented GEP pattern (rather than the one based on null), then that should be okay.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2247

Assuming I change the null version to be the sub + and version.

[goldstein.w.n]

Update impl to only handle clear GEP -> ptrmask patterns

[goldstein.w.n]

@nikic, okay only handles the align_up, align_down, and documented pattern.

[arichardson]

Would it also make sense to update the align_up codegen to use this intrinsic directly? I didn't do it at the time since there were concerns about using ptrmask, but it seems like optimization passes handle it now?

[goldstein.w.n]

In D154010#4457648, @arichardson wrote:

Would it also make sense to update the align_up codegen to use this intrinsic directly? I didn't do it at the time since there were concerns about using ptrmask, but it seems like optimization passes handle it now?

I think probably. What prompted this was the recent support added for it in knownbits + the fact that we do a better just assigning access attributes (and ingeneral tracking pointers) when there is no ptrtoint.

[nikic]

Can you please add proofs for cases 2 and 3?

This generally looks fine to me, but I think we need slightly better ptrmask support before we land this. My minimum expectation would be ability to fold away ptrmask on already aligned pointers.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2249

Missing , after Ptr.

[nikic]

FYI alive2 now has ptrmask support.

[goldstein.w.n]

In D154010#4538589, @nikic wrote:

FYI alive2 now has ptrmask support.

Oh nice :)
I'll add proofs when I get to properly supporting ptrmask + rebasing this.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

if (!CanBeNull && !CanBeFreed && DerefBytes != 0) {
         }
       }
     }
   }
 
   if (Instruction *R = foldSelectGEP(GEP, Builder))
     return R;
 
+  // The definition of llvm.ptrmask(Ptr, Mask) is:
+  // (getelementptr i8, Ptr, (sub (and (ptrtoint Ptr), Mask), (ptrtoint Ptr)))
+  //
+  // llvm.ptrmask is preferable as it doesn't require transitioning from ptr ->
+  // int (and is generally easier to analyze).
+  // Don't do this for for vector of pointers, as llvm.ptrmask has weird
+  // definition for vectors.
+  if (!GEPType->isVectorTy() && GEPEltType->isIntegerTy(8) &&
+      GEP.getNumIndices() == 1) {
+    Value *Index = GEP.idx_begin()->get();
+    if (Index->hasOneUse()) {
+      Value *Ptr = GEP.getPointerOperand();
+      Value *Mask;
+      const APInt *CMask, *NotCMask;
+      Value *R = nullptr;
+      // (getelementptr i8, Ptr, (sub (and (ptrtoint Ptr), Mask), (ptrtoint Ptr)))
+      //      -> llvm.ptrmask(Ptr, Mask)
+      if (match(Index, m_Sub(m_OneUse(m_c_And(m_PtrToInt(m_Specific(Ptr)),
+                                              m_Value(Mask))),
+                             m_PtrToInt(m_Specific(Ptr)))))
+        R = Builder.CreateIntrinsic(Ptr->getType(), Intrinsic::ptrmask,
+                                    {Ptr, Mask});
+      // (getelementptr i8, Ptr (sub 0, (and (ptrtoint Ptr), CMask)))
+      //      -> llvm.ptrmask(Ptr, ~CMask)
+      //
+      // This is what __builtin_align_down generates.
+      else if (match(Index, m_Neg(m_OneUse(m_And(m_PtrToInt(m_Specific(Ptr)),
+                                                 m_APInt(CMask))))) &&
+               CMask->isMask())
+        R = Builder.CreateIntrinsic(
+            Ptr->getType(), Intrinsic::ptrmask,
+            {Ptr, ConstantInt::get(Index->getType(), ~(*CMask))});
+
+      // (getelementptr i8, Ptr (sub (and (add (ptrtoint Ptr), ~CMask), CMask), (ptrtoint Ptr)))
+      //      -> llvm.ptrmask((getelementptr i8, Ptr, ~CMask), CMask)
+      //
+      // This is what __builtin_align_up generates.
+      else if (match(Index,
+                     m_Sub(m_OneUse(
+                               m_And(m_OneUse(m_Add(m_PtrToInt(m_Specific(Ptr)),
+                                                    m_APInt(NotCMask))),
+                                     m_APInt(CMask))),
+                           m_PtrToInt(m_Specific(Ptr)))) &&
+               (*CMask) == ~(*NotCMask) && NotCMask->isMask())
+        R = Builder.CreateIntrinsic(
+            Ptr->getType(), Intrinsic::ptrmask,
+            {Builder.CreateGEP(GEP.getSourceElementType(), Ptr,
+                               ConstantInt::get(Index->getType(), *NotCMask)),
+             ConstantInt::get(Index->getType(), *CMask)});
+      if (R)
+        return replaceInstUsesWith(GEP, R);
+    }
+  }
   return nullptr;
 }
 
 static bool isNeverEqualToUnescapedAlloc(Value *V, const TargetLibraryInfo &TLI,
                                          Instruction *AI) {
   if (isa<ConstantPointerNull>(V))
     return true;
   if (auto *LI = dyn_cast<LoadInst>(V))

llvm/test/Transforms/InstCombine/to-ptrmask.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 declare ptr @llvm.ptrmask(ptr, i64)
 declare void @use.i64(i64)
 
 define ptr @GEP_default_pattern_to_ptrmask(ptr %p, i64 %m) {
 ; CHECK-LABEL: define ptr @GEP_default_pattern_to_ptrmask
 ; CHECK-SAME: (ptr [[P:%.*]], i64 [[M:%.*]]) {
-; CHECK-NEXT:    [[PI:%.*]] = ptrtoint ptr [[P]] to i64
-; CHECK-NEXT:    [[PI_M:%.*]] = and i64 [[PI]], [[M]]
-; CHECK-NEXT:    [[PI_OFF:%.*]] = sub i64 [[PI_M]], [[PI]]
-; CHECK-NEXT:    [[PM:%.*]] = getelementptr i8, ptr [[P]], i64 [[PI_OFF]]
+; CHECK-NEXT:    [[PM:%.*]] = call ptr @llvm.ptrmask.p0.i64(ptr [[P]], i64 [[M]])
 ; CHECK-NEXT:    ret ptr [[PM]]
 ;
   %pi = ptrtoint ptr %p to i64
   %pi_m = and i64 %pi, %m
   %pi_off = sub i64 %pi_m, %pi
   %pm = getelementptr i8, ptr %p, i64 %pi_off
   ret ptr %pm
 }
;
   %pm = getelementptr i16, ptr %p, i64 %pi_off
   ret ptr %pm
 }
 
 define ptr @builtin_align_up_to_ptrmask(ptr %x) {
 ; CHECK-LABEL: define ptr @builtin_align_up_to_ptrmask
 ; CHECK-SAME: (ptr [[X:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[INTPTR:%.*]] = ptrtoint ptr [[X]] to i64
-; CHECK-NEXT:    [[OVER_BOUNDARY:%.*]] = add i64 [[INTPTR]], 7
-; CHECK-NEXT:    [[ALIGNED_INTPTR:%.*]] = and i64 [[OVER_BOUNDARY]], -8
-; CHECK-NEXT:    [[DIFF:%.*]] = sub i64 [[ALIGNED_INTPTR]], [[INTPTR]]
-; CHECK-NEXT:    [[ALIGNED_RESULT:%.*]] = getelementptr inbounds i8, ptr [[X]], i64 [[DIFF]]
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[X]], i64 7
+; CHECK-NEXT:    [[ALIGNED_RESULT:%.*]] = call ptr @llvm.ptrmask.p0.i64(ptr [[TMP0]], i64 -8)
 ; CHECK-NEXT:    ret ptr [[ALIGNED_RESULT]]
 ;
 entry:
   %intptr = ptrtoint ptr %x to i64
   %over_boundary = add i64 %intptr, 7
   %aligned_intptr = and i64 %over_boundary, -8
   %diff = sub i64 %aligned_intptr, %intptr
   %aligned_result = getelementptr inbounds i8, ptr %x, i64 %diff
entry:
   %aligned_result = getelementptr inbounds i8, ptr %x, i64 %diff
   ret ptr %aligned_result
 }
 
 define ptr @builtin_align_down_to_ptrmask(ptr %x) {
 ; CHECK-LABEL: define ptr @builtin_align_down_to_ptrmask
 ; CHECK-SAME: (ptr [[X:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[INTPTR:%.*]] = ptrtoint ptr [[X]] to i64
-; CHECK-NEXT:    [[AND:%.*]] = and i64 [[INTPTR]], 7
-; CHECK-NEXT:    [[DIFF:%.*]] = sub nsw i64 0, [[AND]]
-; CHECK-NEXT:    [[ALIGNED_RESULT:%.*]] = getelementptr inbounds i8, ptr [[X]], i64 [[DIFF]]
+; CHECK-NEXT:    [[ALIGNED_RESULT:%.*]] = call ptr @llvm.ptrmask.p0.i64(ptr [[X]], i64 -8)
 ; CHECK-NEXT:    ret ptr [[ALIGNED_RESULT]]
 ;
 entry:
   %intptr = ptrtoint ptr %x to i64
   %and = and i64 %intptr, 7
   %diff = sub nsw i64 0, %and
   %aligned_result = getelementptr inbounds i8, ptr %x, i64 %diff
   ret ptr %aligned_result