[DirectX] Add a GEP to loads and stores on array allocas (#148059)

Fixes #147114 by inserting a GEP between any direct loads and stores on
an alloca.

Author: Icohedron

llvm/lib/Target/DirectX/DXILLegalizePass.cpp

@@ -98,9 +98,9 @@ static void fixI8UseChain(Instruction &I,
       ElementType = AI->getAllocatedType();
     if (auto *GEP = dyn_cast<GetElementPtrInst>(NewOperands[0])) {
       ElementType = GEP->getSourceElementType();
-      if (ElementType->isArrayTy())
-        ElementType = ElementType->getArrayElementType();
     }
+    if (ElementType->isArrayTy())
+      ElementType = ElementType->getArrayElementType();
     LoadInst *NewLoad = Builder.CreateLoad(ElementType, NewOperands[0]);
     ReplacedValues[Load] = NewLoad;
     ToRemove.push_back(Load);
@@ -562,6 +562,45 @@ legalizeGetHighLowi64Bytes(Instruction &I,
   }
 }
 
+static void
+legalizeLoadStoreOnArrayAllocas(Instruction &I,
+                                SmallVectorImpl<Instruction *> &ToRemove,
+                                DenseMap<Value *, Value *> &) {
+
+  Value *PtrOp;
+  unsigned PtrOpIndex;
+  [[maybe_unused]] Type *LoadStoreTy;
+  if (auto *LI = dyn_cast<LoadInst>(&I)) {
+    PtrOp = LI->getPointerOperand();
+    PtrOpIndex = LI->getPointerOperandIndex();
+    LoadStoreTy = LI->getType();
+  } else if (auto *SI = dyn_cast<StoreInst>(&I)) {
+    PtrOp = SI->getPointerOperand();
+    PtrOpIndex = SI->getPointerOperandIndex();
+    LoadStoreTy = SI->getValueOperand()->getType();
+  } else
+    return;
+
+  assert(LoadStoreTy->isSingleValueType() &&
+         "Expected load/store type to be a single-valued type");
+
+  auto *AllocaPtrOp = dyn_cast<AllocaInst>(PtrOp);
+  if (!AllocaPtrOp)
+    return;
+
+  Type *Ty = AllocaPtrOp->getAllocatedType();
+  if (!isa<ArrayType>(Ty))
+    return;
+  assert(!isa<ArrayType>(Ty->getArrayElementType()) &&
+         "Expected allocated type of AllocaInst to be a flat ArrayType");
+
+  IRBuilder<> Builder(&I);
+  Value *Zero = Builder.getInt32(0);
+  Value *GEP = Builder.CreateGEP(Ty, AllocaPtrOp, {Zero, Zero}, "",
+                                 GEPNoWrapFlags::all());
+  I.setOperand(PtrOpIndex, GEP);
+}
+
 namespace {
 class DXILLegalizationPipeline {
 
@@ -612,6 +651,7 @@ class DXILLegalizationPipeline {
     // downcastI64toI32InsertExtractElements needs to handle.
     LegalizationPipeline[Stage2].push_back(
         downcastI64toI32InsertExtractElements);
+    LegalizationPipeline[Stage2].push_back(legalizeLoadStoreOnArrayAllocas);
   }
 };
 

llvm/test/CodeGen/DirectX/legalize-i8.ll

@@ -127,10 +127,9 @@ define i32 @i8_geps_index0() {
   ; CHECK:         [[LOAD:%.*]] = load i32, ptr [[GEP]], align 4
   ; CHECK-NEXT:    ret i32 [[LOAD]]
   %1 = alloca [2 x i32], align 8
-  %2 = getelementptr inbounds nuw i8, ptr %1, i32 0
-  %3 = load i8, ptr %2
-  %4 = sext i8 %3 to i32
-  ret i32 %4
+  %2 = load i8, ptr %1
+  %3 = sext i8 %2 to i32
+  ret i32 %3
 }
 
 define i32 @i8_geps_index1() {
@@ -149,11 +148,14 @@ define i32 @i8_geps_index1() {
 define i32 @i8_gep_store() {
   ; CHECK-LABEL: define i32 @i8_gep_store(
   ; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [2 x i32], align 8
+  ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds nuw [2 x i32], ptr [[ALLOCA]], i32 0, i32 0
+  ; CHECK-NEXT:    store i32 0, ptr [[GEP]], align 4
   ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds nuw [2 x i32], ptr [[ALLOCA]], i32 0, i32 1
   ; CHECK-NEXT:    store i32 1, ptr [[GEP]], align 4
   ; CHECK:         [[LOAD:%.*]] = load i32, ptr [[GEP]], align 4
   ; CHECK-NEXT:    ret i32 [[LOAD]]
   %1 = alloca [2 x i32], align 8
+  store i8 0, ptr %1
   %2 = getelementptr inbounds nuw i8, ptr %1, i32 4
   store i8 1, ptr %2
   %3 = load i8, ptr %2

llvm/test/CodeGen/DirectX/legalize-load-store-array-alloca.ll

@@ -0,0 +1,23 @@
+; RUN: opt -S -passes='dxil-legalize' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s
+
+define float @load() {
+; CHECK-LABEL: define float @load
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [2 x float], align 4
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds nuw [2 x float], ptr [[ALLOCA]], i32 0, i32 0
+; CHECK-NEXT:    [[LOAD:%.*]] = load float, ptr [[GEP]], align 4
+; CHECK-NEXT:    ret float [[LOAD]]
+  %a = alloca [2 x float], align 4
+  %b = load float, ptr %a, align 4
+  ret float %b
+}
+
+define void @store() {
+; CHECK-LABEL: define void @store
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [3 x i32], align 4
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds nuw [3 x i32], ptr [[ALLOCA]], i32 0, i32 0
+; CHECK-NEXT:    store i32 0, ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+  %a = alloca [3 x i32], align 4
+  store i32 0, ptr %a, align 4
+  ret void
+}