[Attributor] Propagate alignment through ptrmask

llvm/include/llvm/Transforms/IPO/Attributor.h

@@ -5339,6 +5339,17 @@ struct AAPotentialConstantValues
     return nullptr;
   }
 
+  /// Return the minimum trailing zeros of potential constants
+  unsigned getAssumedMinTrailingZeros() const {
+    unsigned TrailingZeros = getAssumedSet().begin()->getBitWidth() + 1;
+    for (const APInt &It : getAssumedSet()) {
+      if (It.countTrailingZeros() < TrailingZeros)
+        TrailingZeros = It.countTrailingZeros();
+    }
+    if (TrailingZeros > getAssumedSet().begin()->getBitWidth())
+      return 0;
+    return TrailingZeros;
+  }
   /// See AbstractAttribute::getName()
   StringRef getName() const override { return "AAPotentialConstantValues"; }
 

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -5185,6 +5185,33 @@ struct AADereferenceableCallSiteReturned final
 // ------------------------ Align Argument Attribute ------------------------
 
 namespace {
+
+static Align getAssumedAlignForIntrinsic(Attributor &A, AAAlign &QueryingAA,
+                                         const IntrinsicInst &II) {
+  Align Alignment;
+  switch (II.getIntrinsicID()) {
+  case Intrinsic::ptrmask: {
+    const auto *ConstVals = A.getAAFor<AAPotentialConstantValues>(
+        QueryingAA, IRPosition::value(*II.getOperand(1)), DepClassTy::REQUIRED);
+    const auto *AlignAA =
+        A.getAAFor<AAAlign>(QueryingAA, IRPosition::value(*(II.getOperand(0))),
+                            DepClassTy::REQUIRED);
+    if (ConstVals && ConstVals->isValidState()) {
+      unsigned ShiftValue = std::min(ConstVals->getAssumedMinTrailingZeros(),
+                                     Value::MaxAlignmentExponent);
+      Alignment = Align(UINT64_C(1) << ShiftValue);
+    }
+    if (AlignAA && AlignAA->isValidState())
+      Alignment = std::max(AlignAA->getAssumedAlign(), Alignment);
+
+    return std::min(QueryingAA.getAssumedAlign(), Alignment);
+  }
+  default:
+    break;
+  }
+  return Alignment;
+}
+
 static unsigned getKnownAlignForUse(Attributor &A, AAAlign &QueryingAA,
                                     Value &AssociatedValue, const Use *U,
                                     const Instruction *I, bool &TrackUse) {
@@ -5200,6 +5227,28 @@ static unsigned getKnownAlignForUse(Attributor &A, AAAlign &QueryingAA,
       TrackUse = true;
     return 0;
   }
+  if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I))
+    switch (II->getIntrinsicID()) {
+    case Intrinsic::ptrmask: {
+      // Is it appropriate to pull attribute in initialization?
+      const auto *ConstVals = A.getAAFor<AAPotentialConstantValues>(
+          QueryingAA, IRPosition::value(*II->getOperand(1)), DepClassTy::NONE);
+      const auto *AlignAA = A.getAAFor<AAAlign>(
+          QueryingAA, IRPosition::value(*II), DepClassTy::NONE);
+      if (ConstVals && ConstVals->isValidState() && ConstVals->isAtFixpoint()) {
+        unsigned ShiftValue = std::min(ConstVals->getAssumedMinTrailingZeros(),
+                                       Value::MaxAlignmentExponent);
+        Align ConstAlign(UINT64_C(1) << ShiftValue);
+        if (ConstAlign >= AlignAA->getKnownAlign())
+          return Align(1).value();
+      }
+      if (AlignAA)
+        return AlignAA->getKnownAlign().value();
+      break;
+    }
+    default:
+      break;
+    }
 
   MaybeAlign MA;
   if (const auto *CB = dyn_cast<CallBase>(I)) {
@@ -5499,6 +5548,17 @@ struct AAAlignCallSiteReturned final
   AAAlignCallSiteReturned(const IRPosition &IRP, Attributor &A)
       : Base(IRP, A) {}
 
+  ChangeStatus updateImpl(Attributor &A) override {
+    Instruction *I = getIRPosition().getCtxI();
+    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
+      Align Align = getAssumedAlignForIntrinsic(A, *this, *II);
+      if (Align.value() > 1) {
+        return clampStateAndIndicateChange<StateType>(this->getState(),
+                                                      Align.value());
+      }
+    }
+    return Base::updateImpl(A);
+  };
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override { STATS_DECLTRACK_CS_ATTR(align); }
 };