[llvm] [InferAS] Infer the address space of inttoptr (PR #173244)

[Shilei Tian via llvm-commits]

================
@@ -375,15 +396,79 @@ getPointerOperands(const Value &V, const DataLayout &DL,
     return {II.getArgOperand(0)};
   }
   case Instruction::IntToPtr: {
-    assert(isNoopPtrIntCastPair(&Op, DL, TTI));
-    auto *P2I = cast<Operator>(Op.getOperand(0));
-    return {P2I->getOperand(0)};
+    if (isNoopPtrIntCastPair(&Op, DL, TTI)) {
+      auto *P2I = cast<Operator>(Op.getOperand(0));
+      return {P2I->getOperand(0)};
+    }
+    assert(isSafeToCastIntToPtrAddrSpace(&Op, DL));
+    return {PtrIntCastPairs[&Op]};
   }
   default:
     llvm_unreachable("Unexpected instruction type.");
   }
 }
 
+// Return mask. The 1 in mask indicate the bit is changed.
+// This helper function is to compute the max know changed bits for ptr1 and
+// ptr2 after the operation `ptr2 = ptr1 Op Mask`.
+static APInt computeMaxChangedPtrBits(const Operator *Op, const Value *Mask,
+                                      const DataLayout &DL, AssumptionCache *AC,
+                                      const DominatorTree *DT) {
+  KnownBits Known = computeKnownBits(Mask, DL, AC, nullptr, DT);
+  switch (Op->getOpcode()) {
+  case Instruction::Xor:
+  case Instruction::Or:
+    return ~Known.Zero;
+  case Instruction::And:
+    return ~Known.One;
+  default:
+    return APInt::getAllOnes(Known.getBitWidth());
+  }
+}
+
+bool InferAddressSpacesImpl::isSafeToCastIntToPtrAddrSpace(
+    const Operator *I2P, const DataLayout &DL) const {
+  assert(I2P->getOpcode() == Instruction::IntToPtr);
+  // If I2P has been checked before and has the corresponding old pointer value,
+  // just return true.
+  if (PtrIntCastPairs.count(I2P))
+    return true;
+
+  if (I2P->getType()->isVectorTy())
+    return false;
+
+  Value *LogicalOp = I2P->getOperand(0);
+  Value *OldPtr, *Mask;
+  if (!match(LogicalOp,
+             m_c_BitwiseLogic(m_PtrToInt(m_Value(OldPtr)), m_Value(Mask))))
+    return false;
+
+  Operator *AsCast = dyn_cast<AddrSpaceCastOperator>(OldPtr);
+  if (!AsCast)
+    return false;
+
+  unsigned SrcAS = I2P->getType()->getPointerAddressSpace();
+  unsigned DstAS = AsCast->getOperand(0)->getType()->getPointerAddressSpace();
+  APInt PreservedPtrMask = TTI->getAddrSpaceCastPreservedPtrMask(SrcAS, DstAS);
+  if (PreservedPtrMask.isZero())
+    return false;
+  APInt ChangedPtrBits =
+      computeMaxChangedPtrBits(cast<Operator>(LogicalOp), Mask, DL, &AC, DT);
+  // Check if the address bits change is within the preserved mask. If the bits
+  // change is not preserved, it is not safe to perform address space cast.
+  // The following pattern is not safe to cast address space.
+  //   %1 = ptrtoint ptr addrspace(3) %sp to i32
+  //   %2 = zext i32 %1 to i64
+  //   %gp = inttoptr i64 %2 to ptr
+  assert(ChangedPtrBits.getBitWidth() == PreservedPtrMask.getBitWidth());
+  if (ChangedPtrBits.isSubsetOf(PreservedPtrMask)) {
----------------
shiltian wrote:

I don't think this is a well designed interface. On one hand, it is supposed to be a query function, but on the other hand, it also modifies a data member. I think this function should only do the query, and the data update should be in the caller.

https://github.com/llvm/llvm-project/pull/173244