[llvm] [DirectX] Fix GEP flattening with 0-indexed GEPs on global variables (PR #149211)

[Deric C. via llvm-commits]

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

Fixes #149179 

The issue is that `Builder.CreateGEP` does not return a GEP Instruction or GEP ContantExpr when the pointer operand is a global variable and all indices are constant zeroes. 

This PR ensures that a GEP instruction is created if `Builder.CreateGEP` did not return a GEP.

>From 9726ea76b39bc39b6941164b18073239e7172e31 Mon Sep 17 00:00:00 2001
From: Deric Cheung <dericcheung at microsoft.com>
Date: Wed, 16 Jul 2025 22:45:29 +0000
Subject: [PATCH] Fix GEP flattening with 0-indexed GEPs on global variables

---
 llvm/lib/Target/DirectX/DXILFlattenArrays.cpp | 10 +++++++++
 llvm/test/CodeGen/DirectX/flatten-array.ll    | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/llvm/lib/Target/DirectX/DXILFlattenArrays.cpp b/llvm/lib/Target/DirectX/DXILFlattenArrays.cpp
index ce43645d005b0..f0e2e786dfaf4 100644
--- a/llvm/lib/Target/DirectX/DXILFlattenArrays.cpp
+++ b/llvm/lib/Target/DirectX/DXILFlattenArrays.cpp
@@ -343,6 +343,16 @@ bool DXILFlattenArraysVisitor::visitGetElementPtrInst(GetElementPtrInst &GEP) {
         Info.RootFlattenedArrayType, Info.RootPointerOperand,
         {ZeroIndex, FlattenedIndex}, GEP.getName(), GEP.getNoWrapFlags());
 
+    // If the pointer operand is a global variable and all indices are 0,
+    // IRBuilder::CreateGEP will return the global variable instead of creating
+    // a GEP instruction or GEP ConstantExpr. In this case we have to create and
+    // insert our own GEP instruction.
+    if (!isa<GEPOperator>(NewGEP))
+      NewGEP = GetElementPtrInst::Create(
+          Info.RootFlattenedArrayType, Info.RootPointerOperand,
+          {ZeroIndex, FlattenedIndex}, GEP.getNoWrapFlags(), GEP.getName(),
+          Builder.GetInsertPoint());
+
     // Replace the current GEP with the new GEP. Store GEPInfo into the map
     // for later use in case this GEP was not the end of the chain
     GEPChainInfoMap.insert({cast<GEPOperator>(NewGEP), std::move(Info)});
diff --git a/llvm/test/CodeGen/DirectX/flatten-array.ll b/llvm/test/CodeGen/DirectX/flatten-array.ll
index 1376a1db25975..a2e105537ab88 100644
--- a/llvm/test/CodeGen/DirectX/flatten-array.ll
+++ b/llvm/test/CodeGen/DirectX/flatten-array.ll
@@ -218,6 +218,28 @@ define void @two_index_gep_const() {
   ret void
 }
 
+define void @zero_index_global() {
+  ; CHECK-LABEL: define void @zero_index_global(
+  ; CHECK-NEXT: [[GEP:%.*]] = getelementptr inbounds nuw [4 x float], ptr addrspace(3) @g.1dim, i32 0, i32 0
+  ; CHECK-NEXT: load float, ptr addrspace(3) [[GEP]], align 4
+  ; CHECK-NEXT: ret void
+  %1 = getelementptr inbounds nuw [2 x [2 x float]], ptr addrspace(3) @g, i32 0, i32 0, i32 0
+  %2 = load float, ptr addrspace(3) %1, align 4
+  ret void
+}
+
+; Note: A ConstantExpr GEP with all 0 indices is equivalent to the pointer
+; operand of the GEP. Therefore the visitLoadInst will not see the pointer operand
+; as a ConstantExpr GEP and will not create a GEP instruction to be visited.
+; The later dxil-legalize pass will insert a GEP in this instance.
+define void @zero_index_global_const() {
+  ; CHECK-LABEL: define void @zero_index_global_const(
+  ; CHECK-NEXT: load float, ptr addrspace(3) @g.1dim, align 4
+  ; CHECK-NEXT: ret void
+  %1 = load float, ptr addrspace(3) getelementptr inbounds nuw ([2 x [2 x float]], ptr addrspace(3) @g, i32 0, i32 0, i32 0), align 4
+  ret void
+}
+
 define void @gep_4d_index_test()  {
     ; CHECK-LABEL: gep_4d_index_test
     ; CHECK: [[a:%.*]] = alloca [16 x i32], align 4