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..32e2c3ca2c302 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,23 @@ 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() {
+  ; CHECK:  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
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"