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

[Farzon Lotfi via llvm-commits]

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

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.

>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] [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