[llvm] [DirectX] Allow vector Allocas to be transformed into arrays (PR #145972)

[Farzon Lotfi via llvm-commits]

https://github.com/farzonl updated https://github.com/llvm/llvm-project/pull/145972

>From 2cfc27454cecb68eb5c2360fd93de0c5e1c651be Mon Sep 17 00:00:00 2001
From: Farzon Lotfi <farzonlotfi at microsoft.com>
Date: Thu, 26 Jun 2025 16:54:59 -0400
Subject: [PATCH 1/2] [DirectX] Allow vector Allocas to be transformed into
 arrays fixes #145782

This change modifies isArrayOfVectors into isVectorOrArrayOfVectors.
The previous implementation did not support vector to array
transformations. Further it was too simplistic and didn't assume allocas
would create multidimensional arrays.
---
 .../Target/DirectX/DXILDataScalarization.cpp  | 14 ++++++++-----
 llvm/test/CodeGen/DirectX/scalarize-alloca.ll | 20 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/DirectX/DXILDataScalarization.cpp b/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
index 71eb1349314ea..c97c604fdbf77 100644
--- a/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
+++ b/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
@@ -117,19 +117,23 @@ DataScalarizerVisitor::lookupReplacementGlobal(Value *CurrOperand) {
   return nullptr; // Not found
 }
 
-static bool isArrayOfVectors(Type *T) {
+// Helper function to check if a type is a vector or an array of vectors
+static bool isVectorOrArrayOfVectors(Type *T) {
+  if (isa<VectorType>(T))
+    return true;
   if (ArrayType *ArrType = dyn_cast<ArrayType>(T))
-    return isa<VectorType>(ArrType->getElementType());
+    return isa<VectorType>(ArrType->getElementType()) ||
+           isVectorOrArrayOfVectors(ArrType->getElementType());
   return false;
 }
 
 bool DataScalarizerVisitor::visitAllocaInst(AllocaInst &AI) {
-  if (!isArrayOfVectors(AI.getAllocatedType()))
+  Type *AllocatedType = AI.getAllocatedType();
+  if (!isVectorOrArrayOfVectors(AllocatedType))
     return false;
 
-  ArrayType *ArrType = cast<ArrayType>(AI.getAllocatedType());
   IRBuilder<> Builder(&AI);
-  Type *NewType = equivalentArrayTypeFromVector(ArrType);
+  Type *NewType = equivalentArrayTypeFromVector(AllocatedType);
   AllocaInst *ArrAlloca =
       Builder.CreateAlloca(NewType, nullptr, AI.getName() + ".scalarize");
   ArrAlloca->setAlignment(AI.getAlign());
diff --git a/llvm/test/CodeGen/DirectX/scalarize-alloca.ll b/llvm/test/CodeGen/DirectX/scalarize-alloca.ll
index b589136d6965c..51495e19001f4 100644
--- a/llvm/test/CodeGen/DirectX/scalarize-alloca.ll
+++ b/llvm/test/CodeGen/DirectX/scalarize-alloca.ll
@@ -2,7 +2,7 @@
 ; RUN: opt -S -passes='dxil-data-scalarization,dxil-flatten-arrays' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s --check-prefixes=FCHECK,CHECK
 
 ; CHECK-LABEL: alloca_2d__vec_test
-define void @alloca_2d__vec_test() local_unnamed_addr #2 {
+define void @alloca_2d__vec_test() {
   ; SCHECK:  alloca [2 x [4 x i32]], align 16
   ; FCHECK:  alloca [8 x i32], align 16
   ; CHECK: ret void
@@ -10,6 +10,24 @@ define void @alloca_2d__vec_test() local_unnamed_addr #2 {
   ret void
 }
 
+; CHECK-LABEL: alloca_4d__vec_test
+define void @alloca_4d__vec_test() {
+  ; SCHECK:  alloca [2 x [2 x [2 x [2 x i32]]]], align 16
+  ; FCHECK:  alloca [16 x i32], align 16
+  ; CHECK: ret void
+  %1 = alloca [2 x [2 x [2 x <2 x i32>]]], align 16
+  ret void
+}
+
+; CHECK-LABEL: alloca_vec_test
+define void @alloca_vec_test() {
+  ; SCHECK:  alloca [4 x i32], align 16
+  ; FCHECK:  alloca [4 x i32], align 16
+  ; CHECK: ret void
+  %1 = alloca <4 x i32>, align 16
+  ret void
+}
+
 ; CHECK-LABEL: alloca_2d_gep_test
 define void @alloca_2d_gep_test() {
   ; SCHECK:  [[alloca_val:%.*]] = alloca [2 x [2 x i32]], align 16

>From b74010a2adb37dc0cc58e48012b3c4ee101023e9 Mon Sep 17 00:00:00 2001
From: Farzon Lotfi <farzonlotfi at microsoft.com>
Date: Fri, 27 Jun 2025 11:15:49 -0400
Subject: [PATCH 2/2] The dxil-dis/shuffle.ll test is not valid sm6.7 DXIL, but
 there seems to be an effort to make sure shufflevector works in the bitcode
 writer so lets start llc after the scalarization passes

---
 llvm/test/tools/dxil-dis/shuffle.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/tools/dxil-dis/shuffle.ll b/llvm/test/tools/dxil-dis/shuffle.ll
index 0739afc87ff2b..5f7ecba45edd0 100644
--- a/llvm/test/tools/dxil-dis/shuffle.ll
+++ b/llvm/test/tools/dxil-dis/shuffle.ll
@@ -1,4 +1,4 @@
-; RUN: llc --filetype=obj %s -o - 2>&1 | dxil-dis -o - | FileCheck %s
+; RUN: llc -start-after=dxil-flatten-arrays --filetype=obj %s -o - 2>&1 | dxil-dis -o - | FileCheck %s
 target datalayout = "e-m:e-p:32:32-i1:32-i8:8-i16:16-i32:32-i64:64-f16:16-f32:32-f64:64-n8:16:32:64"
 target triple = "dxil-unknown-shadermodel6.7-library"