[DirectX] Allow vector Allocas to be transformed into arrays

[farzonl]

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.

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: Farzon Lotfi (farzonl)

Changes

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.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXILDataScalarization.cpp (+9-5)
• (modified) llvm/test/CodeGen/DirectX/scalarize-alloca.ll (+19-1)

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

[farzonl]

tools/dxil-dis/shuffle.ll test seems to indicate we don't want to scalarize vector allocas. This pr might need to be modified to drop the vector transformation to arrays and just keep the array recursion.

[farzonl]

tools/dxil-dis/shuffle.ll test seems to indicate we don't want to scalarize vector allocas. This pr might need to be modified to drop the vector transformation to arrays and just keep the array recursion.

Test needed to be update to avoid the Data scalarizer pass.