[llvm] [BPF] Handle certain mem intrinsic functions with addr-space arguments (PR #160025)

[via llvm-commits]

================
@@ -493,21 +559,69 @@ bool BPFCheckAndAdjustIR::insertASpaceCasts(Module &M) {
   for (Function &F : M) {
     DenseMap<Value *, Value *> CastsCache;
     for (BasicBlock &BB : F) {
-      for (Instruction &I : BB) {
+      for (Instruction &I : llvm::make_early_inc_range(BB)) {
         unsigned PtrOpNum;
 
-        if (auto *LD = dyn_cast<LoadInst>(&I))
+        if (auto *LD = dyn_cast<LoadInst>(&I)) {
           PtrOpNum = LD->getPointerOperandIndex();
-        else if (auto *ST = dyn_cast<StoreInst>(&I))
+          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
+          continue;
+        }
+        if (auto *ST = dyn_cast<StoreInst>(&I)) {
           PtrOpNum = ST->getPointerOperandIndex();
-        else if (auto *CmpXchg = dyn_cast<AtomicCmpXchgInst>(&I))
+          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
+          continue;
+        }
+        if (auto *CmpXchg = dyn_cast<AtomicCmpXchgInst>(&I)) {
           PtrOpNum = CmpXchg->getPointerOperandIndex();
-        else if (auto *RMW = dyn_cast<AtomicRMWInst>(&I))
+          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
+          continue;
+        }
+        if (auto *RMW = dyn_cast<AtomicRMWInst>(&I)) {
           PtrOpNum = RMW->getPointerOperandIndex();
-        else
+          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
           continue;
+        }
+
+        auto *CI = dyn_cast<CallInst>(&I);
+        if (!CI)
+          continue;
+
+        Function *Callee = CI->getCalledFunction();
+        if (!Callee || !Callee->isIntrinsic())
+          continue;
+
+        // Check memset/memcpy/memmove
+        Intrinsic::ID ID = Callee->getIntrinsicID();
+        bool IsSet = ID == Intrinsic::memset;
+        bool IsCpy = ID == Intrinsic::memcpy;
+        bool IsMove = ID == Intrinsic::memmove;
+        if (!IsSet && !IsCpy && !IsMove)
+          continue;
+
+        auto isAS = [&](unsigned ArgIdx) {
+          Value *V = CI->getArgOperand(ArgIdx);
+          if (auto *PTy = dyn_cast<PointerType>(V->getType()))
+            return PTy->getAddressSpace() != 0;
+          return false;
+        };
+
+        // For memset: only dest is a pointer; for memcpy/memmove: dest & src.
+        bool HasAS = IsSet ? isAS(0) : (isAS(0) || isAS(1));
----------------
eddyz87 wrote:

Nit: I'd check operands and bail out of transformation inside `aspaceMem{Set,Cpy,Move}`.
E.g. by making `wrapPtrIfASNotZero()` return NULL and checking it's return value.

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