[DirectX] Add a GEP to loads and stores on array allocas

[Icohedron]

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

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: Deric C. (Icohedron)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/148059.diff

2 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXILLegalizePass.cpp (+43)
• (added) llvm/test/CodeGen/DirectX/legalize-load-store-array-alloca.ll (+23)

diff --git a/llvm/lib/Target/DirectX/DXILLegalizePass.cpp b/llvm/lib/Target/DirectX/DXILLegalizePass.cpp
index 76a46c7a2b760..724ee57c05a39 100644
--- a/llvm/lib/Target/DirectX/DXILLegalizePass.cpp
+++ b/llvm/lib/Target/DirectX/DXILLegalizePass.cpp
@@ -562,6 +562,48 @@ legalizeGetHighLowi64Bytes(Instruction &I,
   }
 }
 
+static void legalizeLoadStoreOnArrayAllocas(
+    Instruction &I, SmallVectorImpl<Instruction *> &ToRemove,
+    DenseMap<Value *, Value *> &) {
+
+  Value *PtrOp;
+  [[maybe_unused]] Type *LoadStoreTy;
+  if (auto *LI = dyn_cast<LoadInst>(&I)) {
+    PtrOp = LI->getPointerOperand();
+    LoadStoreTy = LI->getType();
+  } else if (auto *SI = dyn_cast<StoreInst>(&I)) {
+    PtrOp = SI->getPointerOperand();
+    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.CreateInBoundsGEP(Ty, AllocaPtrOp, {Zero, Zero});
+
+  Value *NewLoadStore = nullptr;
+  if (auto *LI = dyn_cast<LoadInst>(&I))
+    NewLoadStore = Builder.CreateLoad(LI->getType(), GEP, LI->getName());
+  else if (auto *SI = dyn_cast<StoreInst>(&I))
+    NewLoadStore =
+        Builder.CreateStore(SI->getValueOperand(), GEP, SI->isVolatile());
+
+  ToRemove.push_back(&I);
+  I.replaceAllUsesWith(NewLoadStore);
+}
+
 namespace {
 class DXILLegalizationPipeline {
 
@@ -605,6 +647,7 @@ class DXILLegalizationPipeline {
     LegalizationPipeline[Stage1].push_back(legalizeMemCpy);
     LegalizationPipeline[Stage1].push_back(removeMemSet);
     LegalizationPipeline[Stage1].push_back(updateFnegToFsub);
+    LegalizationPipeline[Stage1].push_back(legalizeLoadStoreOnArrayAllocas);
     // Note: legalizeGetHighLowi64Bytes and
     // downcastI64toI32InsertExtractElements both modify extractelement, so they
     // must run staggered stages. legalizeGetHighLowi64Bytes runs first b\c it
diff --git a/llvm/test/CodeGen/DirectX/legalize-load-store-array-alloca.ll b/llvm/test/CodeGen/DirectX/legalize-load-store-array-alloca.ll
new file mode 100644
index 0000000000000..703f569f4dfab
--- /dev/null
+++ b/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 [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 [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
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.