InstCombine: Gracefully handle more allocas in the wrong address space

[arsenm]

Summary

Officially this is currently required to always use the datalayout's
alloca address space. This may change in the future, and it's cleaner
to propagate the existing alloca's addrspace anyway.

This is a triple fix. Initially the change in simplifyAllocaArraySize
would drop the address space, but produce output. Fixing this hit an
assertion in the cast combine.

This patch also makes the changes to handle this situation from
a33e12801279a947c74fdee2655b24480941fb39 dead, so eliminate
it. InstCombine should not take it upon itself to introduce
addrspacecasts, and preserve the original address space instead.

[jdoerfert]

IC keeping the AS as it is, I like it.

[arsenm]

286237962ad35bc4756947c32310838ad4c76451

llvm/include/llvm/IR/Instructions.h

public:
   const Value *getArraySize() const { return getOperand(0); }
   Value *getArraySize() { return getOperand(0); }
 
   /// Overload to return most specific pointer type.
   PointerType *getType() const {
     return cast<PointerType>(Instruction::getType());
   }
 
+  /// Return the address space for the allocation.
+  unsigned getAddressSpace() const {
+    return getType()->getAddressSpace();
+  }
+
   /// Get allocation size in bits. Returns None if size can't be determined,
   /// e.g. in case of a VLA.
   Optional<TypeSize> getAllocationSizeInBits(const DataLayout &DL) const;
 
   /// Return the type that is being allocated by the instruction.
   Type *getAllocatedType() const { return AllocatedType; }
   /// for use only in special circumstances that need to generically
   /// transform a whole instruction (eg: IR linking and vectorization).

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

Instruction *InstCombinerImpl::PromoteCastOfAllocation(BitCastInst &CI,
   }
 
   if (uint64_t Offset = (AllocElTySize*ArrayOffset)/CastElTySize) {
     Value *Off = ConstantInt::get(AI.getArraySize()->getType(),
                                   Offset, true);
     Amt = Builder.CreateAdd(Amt, Off);
   }
 
-  AllocaInst *New = Builder.CreateAlloca(CastElTy, Amt);
+  AllocaInst *New = Builder.CreateAlloca(CastElTy, AI.getAddressSpace(), Amt);
   New->setAlignment(AI.getAlign());
   New->takeName(&AI);
   New->setUsedWithInAlloca(AI.isUsedWithInAlloca());
 
   // If the allocation has multiple real uses, insert a cast and change all
   // things that used it to use the new cast.  This will also hack on CI, but it
   // will die soon.
   if (!AI.hasOneUse()) {

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

if (!AI.isArrayAllocation()) {
     // Canonicalize it.
     return IC.replaceOperand(AI, 0, IC.Builder.getInt32(1));
   }
 
   // Convert: alloca Ty, C - where C is a constant != 1 into: alloca [C x Ty], 1
   if (const ConstantInt *C = dyn_cast<ConstantInt>(AI.getArraySize())) {
     if (C->getValue().getActiveBits() <= 64) {
       Type *NewTy = ArrayType::get(AI.getAllocatedType(), C->getZExtValue());
-      AllocaInst *New = IC.Builder.CreateAlloca(NewTy, nullptr, AI.getName());
+      AllocaInst *New = IC.Builder.CreateAlloca(NewTy, AI.getAddressSpace(),
+                                                nullptr, AI.getName());
       New->setAlignment(AI.getAlign());
 
       // Scan to the end of the allocation instructions, to skip over a block of
       // allocas if possible...also skip interleaved debug info
       //
       BasicBlock::iterator It(New);
       while (isa<AllocaInst>(*It) || isa<DbgInfoIntrinsic>(*It))
         ++It;
 
       // Now that I is pointing to the first non-allocation-inst in the block,
       // insert our getelementptr instruction...
       //
       Type *IdxTy = IC.getDataLayout().getIntPtrType(AI.getType());
       Value *NullIdx = Constant::getNullValue(IdxTy);
       Value *Idx[2] = {NullIdx, NullIdx};
-      Instruction *NewI = GetElementPtrInst::CreateInBounds(
+      Instruction *GEP = GetElementPtrInst::CreateInBounds(
           NewTy, New, Idx, New->getName() + ".sub");
-      IC.InsertNewInstBefore(NewI, *It);
-
-      // Gracefully handle allocas in other address spaces.
-      if (AI.getType()->getPointerAddressSpace() !=
-          NewI->getType()->getPointerAddressSpace()) {
-        NewI =
-            CastInst::CreatePointerBitCastOrAddrSpaceCast(NewI, AI.getType());
-        IC.InsertNewInstBefore(NewI, *It);
-      }
+      IC.InsertNewInstBefore(GEP, *It);
 
       // Now make everything use the getelementptr instead of the original
       // allocation.
-      return IC.replaceInstUsesWith(AI, NewI);
+      return IC.replaceInstUsesWith(AI, GEP);
     }
   }
 
   if (isa<UndefValue>(AI.getArraySize()))
     return IC.replaceInstUsesWith(AI, Constant::getNullValue(AI.getType()));
 
   // Ensure that the alloca array size argument has type intptr_t, so that
   // any casting is exposed early.

llvm/test/Transforms/InstCombine/alloca-in-non-alloca-as.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; Gracefully handle the alloca that is not in the alloca AS (=5)
 
 target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7"
 target triple = "amdgcn-amd-amdhsa"
 
 declare void @use(i8*, i32**)
 
 define weak amdgpu_kernel void @__omp_offloading_802_ea0109_main_l8(i32* %a) {
 ; CHECK-LABEL: @__omp_offloading_802_ea0109_main_l8(
 ; CHECK-NEXT:  .master:
-; CHECK-NEXT:    [[TMP0:%.*]] = alloca i32*, align 1, addrspace(5)
-; CHECK-NEXT:    [[DOTSUB:%.*]] = bitcast i32* addrspace(5)* [[TMP0]] to i8 addrspace(5)*
-; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i8 addrspace(5)* [[DOTSUB]] to i8*
-; CHECK-NEXT:    [[A_ON_STACK:%.*]] = addrspacecast i32* addrspace(5)* [[TMP0]] to i32**
-; CHECK-NEXT:    call void @use(i8* [[TMP1]], i32** [[A_ON_STACK]])
+; CHECK-NEXT:    [[TMP0:%.*]] = alloca i32*, align 1
+; CHECK-NEXT:    [[DOTSUB:%.*]] = bitcast i32** [[TMP0]] to i8*
+; CHECK-NEXT:    call void @use(i8* [[DOTSUB]], i32** [[TMP0]])
 ; CHECK-NEXT:    ret void
 ;
 .master:
   %0 = alloca i8, i64 8, align 1
   %a_on_stack = bitcast i8* %0 to i32**
   store i32* undef, i32** %a_on_stack, align 8
   call void @use(i8* %0, i32** %a_on_stack)
   ret void
 }
 
 %struct.widget = type { [8 x i8] }
 
 define void @spam(i64* %arg1) {
 ; CHECK-LABEL: @spam(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[ALLOCA1:%.*]] = alloca [0 x [30 x %struct.widget]], align 16, addrspace(5)
-; CHECK-NEXT:    [[ALLOCA1_SUB:%.*]] = getelementptr inbounds [0 x [30 x %struct.widget]], [0 x [30 x %struct.widget]] addrspace(5)* [[ALLOCA1]], i32 0, i32 0
-; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast [30 x %struct.widget] addrspace(5)* [[ALLOCA1_SUB]] to [30 x %struct.widget]*
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [30 x %struct.widget], [30 x %struct.widget]* [[TMP0]], i64 0, i64 0
+; CHECK-NEXT:    [[ALLOCA1:%.*]] = alloca [0 x [30 x %struct.widget]], align 16
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [0 x [30 x %struct.widget]], [0 x [30 x %struct.widget]]* [[ALLOCA1]], i64 0, i64 0, i64 0
 ; CHECK-NEXT:    call void @zot(%struct.widget* [[GEP]])
 ; CHECK-NEXT:    ret void
 ;
 bb:
   %alloca = alloca [30 x %struct.widget], i32 0, align 16
   %gep = getelementptr inbounds [30 x %struct.widget], [30 x %struct.widget]* %alloca, i64 0, i64 0
   call void @zot(%struct.widget* %gep)
   ret void
 }
 
 declare hidden void @zot(%struct.widget*)