[DirectX] Fix handling of i8 GEPs in DXILDataScalarization

[Icohedron]

The DXIL data scalarization pass incorrectly replaced i8 GEPs.
This PR makes the data scalarization pass only change the type of the GEP if it is not an i8.
i8 GEPs are handled by the later DXIL legalization pass.

Partially fixes #145370, specifically this case: https://godbolt.org/z/9fesTfKre

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: Deric C. (Icohedron)

Changes

The DXIL data scalarization pass incorrectly replaced i8 GEPs.
This PR makes the data scalarization pass change only the type of the GEP if it is not an i8.
i8 GEPs are handled by the later DXIL legalization pass.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXILDataScalarization.cpp (+7-1)
• (added) llvm/test/CodeGen/DirectX/issue-145370-scalarize-i8-gep.ll (+16)

diff --git a/llvm/lib/Target/DirectX/DXILDataScalarization.cpp b/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
index 0a9b2bb99f7eb..d0abda1fc4002 100644
--- a/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
+++ b/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
@@ -298,7 +298,7 @@ bool DataScalarizerVisitor::visitExtractElementInst(ExtractElementInst &EEI) {
 
 bool DataScalarizerVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
   Value *PtrOperand = GEPI.getPointerOperand();
-  Type *OrigGEPType = GEPI.getPointerOperandType();
+  Type *OrigGEPType = GEPI.getSourceElementType();
   Type *NewGEPType = OrigGEPType;
   bool NeedsTransform = false;
 
@@ -316,6 +316,12 @@ bool DataScalarizerVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
     }
   }
 
+  // The new GEP type will be the flattened type unless the original GEP is
+  // an i8 GEP which will be handled by the DXIL legalization pass instead
+  if (OrigGEPType->isIntegerTy(8)) {
+    NewGEPType = OrigGEPType;
+  }
+
   // Note: We bail if this isn't a gep touched via alloca or global
   // transformations
   if (!NeedsTransform)
diff --git a/llvm/test/CodeGen/DirectX/issue-145370-scalarize-i8-gep.ll b/llvm/test/CodeGen/DirectX/issue-145370-scalarize-i8-gep.ll
new file mode 100644
index 0000000000000..2a0fa96eb7e9a
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/issue-145370-scalarize-i8-gep.ll
@@ -0,0 +1,16 @@
+; RUN: opt -S -passes='dxil-data-scalarization' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s --check-prefixes=SCHECK,CHECK
+; RUN: opt -S -passes='dxil-data-scalarization,dxil-flatten-arrays,function(dxil-legalize)' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s --check-prefixes=LCHECK,CHECK
+
+; SCHECK: @g.scalarized = local_unnamed_addr addrspace(3) global [2 x [2 x i32]] zeroinitializer, align 8
+; LCHECK: @g.scalarized.1dim = local_unnamed_addr addrspace(3) global [4 x i32] zeroinitializer, align 8
+@g = local_unnamed_addr addrspace(3) global [2 x <2 x i32>] zeroinitializer, align 8
+
+; CHECK-LABEL: test_store
+define void @test_store(<4 x float> noundef %a) #0 {
+  ; SCHECK: store i32 0, ptr addrspace(3) getelementptr inbounds nuw (i8, ptr addrspace(3) @g.scalarized, i32 12), align 4
+  ; LCHECK: store i32 0, ptr addrspace(3) getelementptr inbounds nuw ([4 x i32], ptr addrspace(3) @g.scalarized.1dim, i32 0, i32 3), align 4
+  store i32 0, ptr addrspace(3) getelementptr inbounds nuw (i8, ptr addrspace(3) @g, i32 12), align 4
+  ret void
+} 
+
+attributes #0 = { convergent norecurse nounwind "hlsl.export"}