Don't sink ptrtoint/inttoptr sequences into non-noop addrspacecasts.

[maleadt]

Summary

In https://reviews.llvm.org/D30114, support for mismatching address spaces was introduced to CodeGenPrepare's optimizeMemoryInst, using addrspacecast as it was argued that only no-op addrspacecasts would be considered when constructing the address mode. However, by doing inttoptr/ptrtoint, it's possible to get CGP to emit an addrspace that's not actually no-op, introducing a miscompilation:

define void @kernel(i8* %julia_ptr) {
  %intptr = ptrtoint i8* %julia_ptr to i64
  %ptr = inttoptr i64 %intptr to i32 addrspace(3)*

  br label %end
end:

  store atomic i32 1, i32 addrspace(3)* %ptr unordered, align 4
  ret void
}

Gets compiled to:

define void @kernel(i8* %julia_ptr) {
end:
  %0 = addrspacecast i8* %julia_ptr to i32 addrspace(3)*
  store atomic i32 1, i32 addrspace(3)* %0 unordered, align 4
  ret void
}

In the case of NVPTX, this introduces a cvta.to.shared, whereas leaving out the %end block and branch doesn't trigger this optimization. This results in illegal memory accesses as seen in https://github.com/JuliaGPU/CUDA.jl/issues/558

In this change, I introduced a check before doing the pointer cast that verifies address spaces are the same. If not, it emits a ptrtoint/inttoptr combination to get a no-op cast between address spaces. I decided against disallowing ptrtoint/inttoptr with non-default AS in matchOperationAddr, because now its still possible to look through multiple sequences of them that ultimately do not result in a address space mismatch (i.e. the second lit test).

[jlebar]

This looks reasonable to me (and I appreciate all the debugging in the Julia bug!), but I have never touched this code, so I don't 100% feel comfortable approving the change.

[tra]

In D92210#2421637, @jlebar wrote:

This looks reasonable to me (and I appreciate all the debugging in the Julia bug!), but I have never touched this code, so I don't 100% feel comfortable approving the change.

Same here. The change makes sense, but there's more going on that I'd like the authors of the CGP to take a look.

[efriedma]

llvm/lib/CodeGen/CodeGenPrepare.cpp

5049

I'd state this comment differently:

There are two reasons the types might not match: a no-op addrspacecast, or a ptrtoint/inttoptr pair. Either way, we emit a ptrtoint/inttoptr pair, to ensure we match the original semantics.

---

Ideally, we don't want to convert an addrspacecast to a ptrtoint/inttoptr pair; it's semantically valid, but we lose information. At this point in the pipeline, it doesn't affect that much, but it does hurt alias analysis a bit. I'm okay with starting with a conservative fix, and revisiting later.

[arsenm]

llvm/lib/CodeGen/CodeGenPrepare.cpp

5049

Ideally we would allow bitcast between different address space pointers with the same size, which would avoid the need for this special case

[maleadt]

Update comment.

[maleadt]

Thanks for the review, I've updated the comment.
FWIW, we have already applied this patch to the LLVM tree used by Julia, and it (unsurprisingly) hasn't resulted in any issues.

llvm/lib/CodeGen/CodeGenPrepare.cpp

5049

You mean "two reasons the *address spaces of* the types might not match"?

I'll update a comment and add a TODO about bitcasts between different address space pointers.

[maleadt]

Bump, anything else needed here?

[jlebar]

In D92210#2559113, @maleadt wrote:

Bump, anything else needed here?

@arsenm I think this is to you.

[vchuravy]

bump

[arsenm]

LGTM

[maleadt]

Can somebody land this for me? I don't have commit access.

[arsenm]

a323dfc0152a11c6934695b1ae67920dae9fad9b

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
   // that we have to sink it into this block.  Check to see if we have already
   // done this for some other load/store instr in this block.  If so, reuse
   // the computation.  Before attempting reuse, check if the address is valid
   // as it may have been erased.
 
   WeakTrackingVH SunkAddrVH = SunkAddrs[Addr];
 
   Value * SunkAddr = SunkAddrVH.pointsToAliveValue() ? SunkAddrVH : nullptr;
+  Type *IntPtrTy = DL->getIntPtrType(Addr->getType());
   if (SunkAddr) {
     LLVM_DEBUG(dbgs() << "CGP: Reusing nonlocal addrmode: " << AddrMode
                       << " for " << *MemoryInst << "\n");
-    if (SunkAddr->getType() != Addr->getType())
+    if (SunkAddr->getType() != Addr->getType()) {
+      if (SunkAddr->getType()->getPointerAddressSpace() !=
+              Addr->getType()->getPointerAddressSpace() &&
+          !DL->isNonIntegralPointerType(Addr->getType())) {
+        // There are two reasons the address spaces might not match: a no-op
+        // addrspacecast, or a ptrtoint/inttoptr pair. Either way, we emit a
+        // ptrtoint/inttoptr pair to ensure we match the original semantics.
+        // TODO: allow bitcast between different address space pointers with the
+        // same size.
+        SunkAddr = Builder.CreatePtrToInt(SunkAddr, IntPtrTy, "sunkaddr");
+        SunkAddr =
+            Builder.CreateIntToPtr(SunkAddr, Addr->getType(), "sunkaddr");
+      } else
         SunkAddr = Builder.CreatePointerCast(SunkAddr, Addr->getType());
+    }
   } else if (AddrSinkUsingGEPs || (!AddrSinkUsingGEPs.getNumOccurrences() &&
                                    SubtargetInfo->addrSinkUsingGEPs())) {
     // By default, we use the GEP-based method when AA is used later. This
     // prevents new inttoptr/ptrtoint pairs from degrading AA capabilities.
     LLVM_DEBUG(dbgs() << "CGP: SINKING nonlocal addrmode: " << AddrMode
                       << " for " << *MemoryInst << "\n");
-    Type *IntPtrTy = DL->getIntPtrType(Addr->getType());
     Value *ResultPtr = nullptr, *ResultIndex = nullptr;
 
     // First, find the pointer.
     if (AddrMode.BaseReg && AddrMode.BaseReg->getType()->isPointerTy()) {
       ResultPtr = AddrMode.BaseReg;
       AddrMode.BaseReg = nullptr;
     }
 
if (!ResultPtr &&
           ResultPtr = Builder.CreatePointerCast(ResultPtr, I8PtrTy);
         SunkAddr =
             AddrMode.InBounds
                 ? Builder.CreateInBoundsGEP(I8Ty, ResultPtr, ResultIndex,
                                             "sunkaddr")
                 : Builder.CreateGEP(I8Ty, ResultPtr, ResultIndex, "sunkaddr");
       }
 
-      if (SunkAddr->getType() != Addr->getType())
+      if (SunkAddr->getType() != Addr->getType()) {
+        if (SunkAddr->getType()->getPointerAddressSpace() !=
+                Addr->getType()->getPointerAddressSpace() &&
+            !DL->isNonIntegralPointerType(Addr->getType())) {
+          // There are two reasons the address spaces might not match: a no-op
+          // addrspacecast, or a ptrtoint/inttoptr pair. Either way, we emit a
+          // ptrtoint/inttoptr pair to ensure we match the original semantics.
+          // TODO: allow bitcast between different address space pointers with
+          // the same size.
+          SunkAddr = Builder.CreatePtrToInt(SunkAddr, IntPtrTy, "sunkaddr");
+          SunkAddr =
+              Builder.CreateIntToPtr(SunkAddr, Addr->getType(), "sunkaddr");
+        } else
           SunkAddr = Builder.CreatePointerCast(SunkAddr, Addr->getType());
       }
+    }
   } else {
     // We'd require a ptrtoint/inttoptr down the line, which we can't do for
     // non-integral pointers, so in that case bail out now.
     Type *BaseTy = AddrMode.BaseReg ? AddrMode.BaseReg->getType() : nullptr;
     Type *ScaleTy = AddrMode.Scale ? AddrMode.ScaledReg->getType() : nullptr;
     PointerType *BasePtrTy = dyn_cast_or_null<PointerType>(BaseTy);
     PointerType *ScalePtrTy = dyn_cast_or_null<PointerType>(ScaleTy);
     if (DL->isNonIntegralPointerType(Addr->getType()) ||

llvm/test/Transforms/CodeGenPrepare/NVPTX/dont-introduce-addrspacecast.ll

+; RUN: opt -S -codegenprepare < %s | FileCheck %s
+
+target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-nvidia-cuda"
+
+
+; ptrtoint/inttoptr combinations can introduce semantically-meaningful address space casts
+; which we can't sink into an addrspacecast
+
+; CHECK-LABEL: @test
+define void @test(i8* %input_ptr) {
+  ; CHECK-LABEL: l1:
+  ; CHECK-NOT: addrspacecast
+  %intptr = ptrtoint i8* %input_ptr to i64
+  %ptr = inttoptr i64 %intptr to i32 addrspace(3)*
+
+  br label %l1
+l1:
+
+  store atomic i32 1, i32 addrspace(3)* %ptr unordered, align 4
+  ret void
+}
+
+
+; we still should be able to look through multiple sequences of inttoptr/ptrtoint
+
+; CHECK-LABEL: @test2
+define void @test2(i8* %input_ptr) {
+  ; CHECK-LABEL: l2:
+  ; CHECK: bitcast
+  ; CHECK-NEXT: store
+  %intptr = ptrtoint i8* %input_ptr to i64
+  %ptr = inttoptr i64 %intptr to i32 addrspace(3)*
+
+  %intptr2 = ptrtoint i32 addrspace(3)* %ptr to i64
+  %ptr2 = inttoptr i64 %intptr2 to i32*
+
+  br label %l2
+l2:
+
+  store atomic i32 1, i32* %ptr2 unordered, align 4
+  ret void
+}