[llvm] [DirectX] Fix handling of i8 GEPs in DXILDataScalarization (PR #145780)

[Deric C. via llvm-commits]

https://github.com/Icohedron created https://github.com/llvm/llvm-project/pull/145780

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.

>From 2d442d3e3295c37d149b6b76b69ce3705e4332b9 Mon Sep 17 00:00:00 2001
From: Icohedron <cheung.deric at gmail.com>
Date: Mon, 23 Jun 2025 21:24:13 +0000
Subject: [PATCH] Fix handling of i8 GEPs in data scalarizer

The DXIL data scalarization pass incorrectly replaced instances of i8
GEPs. This commit makes the data scalarization pass change only the type
of the GEP if it was not an i8. i8 GEPs will be handled by the later
DXIL legalization pass.
---
 .../lib/Target/DirectX/DXILDataScalarization.cpp |  8 +++++++-
 .../DirectX/issue-145370-scalarize-i8-gep.ll     | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/DirectX/issue-145370-scalarize-i8-gep.ll

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
+ at 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"}