Fix OOB in subscript indexing of col-major matrix in cbuffer (#7866)

A bug in HLOperationLower code that lowers the matrix subscript
operation when the matrix is in a cbuffer could cause out-of-bounds
alloca indexing.

The indexing creates an alloca to store a column-vector loaded from one
cbuffer vector. This should have one element per row of the matrix, so
it can then index along that dimension using the local array.

The bug was using the number of columns when sizing the alloca instead
of the number of rows. This caused it to write to out-of-bounds index
when the number of rows exceeded the number of columns (like for
float3x2).

Fixes #7865.

Author: tex3d

lib/HLSL/HLOperationLower.cpp

@@ -8405,7 +8405,7 @@ void TranslateCBAddressUserLegacy(Instruction *user, Value *handle,
           // row.z = c[2].[idx]
           // row.w = c[3].[idx]
           Value *Elts[4];
-          ArrayType *AT = ArrayType::get(EltTy, MatTy.getNumColumns());
+          ArrayType *AT = ArrayType::get(EltTy, MatTy.getNumRows());
 
           IRBuilder<> AllocaBuilder(user->getParent()
                                         ->getParent()

tools/clang/test/DXC/Passes/DxilGen/cb-col-matrix-subscript-dxilgen.ll

@@ -0,0 +1,130 @@
+; RUN: %dxopt %s -hlsl-passes-resume -dxilgen -S | FileCheck %s
+
+; Generated with this, and debug info manually stripped:
+; utils/hct/ExtractIRForPassTest.py -p dxilgen -o tools/clang/test/DXC/Passes/DxilGen/cb-col-matrix-subscript-dxilgen.ll <source>.hlsl -- -T vs_6_0
+
+; Source HLSL:
+; float3x2 f3x2;
+; float main(int row : ROW) : OUT {
+;     return f3x2[row].y;
+; }
+
+; Check cbuffer column-major matrix subscript codegen uses correct alloca size.
+
+; Column-major matrix subscript is lowered in HLOperationLower by:
+; Loading one column into an alloca, dynamically indexing that by row,
+; then doing the same for the next column.
+; Each column is 3 floats for 3 rows. Thus alloca size should be 3 floats.
+
+; Just check the chain for the returned component, skipping the first cbuffer load.
+; alloca size should be number of rows (3), not number of columns (2)
+; CHECK: %[[alloca:[^ ]+]] = alloca [3 x float]
+; Get row index from input:
+; CHECK: %[[row:[^ ]+]] = call i32 @dx.op.loadInput.i32(i32 4, i32 0, i32 0, i8 0, i32 undef)
+; Load of column 1, skipping column 0:
+; CHECK: %[[cbld:[^ ]+]] = call %dx.types.CBufRet.f32 @dx.op.cbufferLoadLegacy.f32(i32 59, %dx.types.Handle %{{[^,]+}}, i32 1)
+; Reconstruct vector from original HL op:
+; CHECK-NEXT: %[[cbld_x:[^ ]+]] = extractvalue %dx.types.CBufRet.f32 %[[cbld]], 0
+; CHECK-NEXT: %[[v3x:[^ ]+]] = insertelement <3 x float> undef, float %[[cbld_x]], i64 0
+; CHECK-NEXT: %[[cbld_y:[^ ]+]] = extractvalue %dx.types.CBufRet.f32 %[[cbld]], 1
+; CHECK-NEXT: %[[v3xy:[^ ]+]] = insertelement <3 x float> %[[v3x]], float %[[cbld_y]], i64 1
+; CHECK-NEXT: %[[cbld_z:[^ ]+]] = extractvalue %dx.types.CBufRet.f32 %[[cbld]], 2
+; CHECK-NEXT: %[[v3xyz:[^ ]+]] = insertelement <3 x float> %[[v3xy]], float %[[cbld_z]], i64 2
+; Store y column to alloca:
+; CHECK-NEXT: %[[ee:[^ ]+]] = extractelement <3 x float> %[[v3xyz]], i32 0
+; CHECK-NEXT: %[[gep_alloca_0:[^ ]+]] = getelementptr inbounds [3 x float], [3 x float]* %[[alloca]], i32 0, i32 0
+; CHECK-NEXT: store float %[[ee]], float* %[[gep_alloca_0]]
+; CHECK-NEXT: %[[ee:[^ ]+]] = extractelement <3 x float> %[[v3xyz]], i32 1
+; CHECK-NEXT: %[[gep_alloca_1:[^ ]+]] = getelementptr inbounds [3 x float], [3 x float]* %[[alloca]], i32 0, i32 1
+; CHECK-NEXT: store float %[[ee]], float* %[[gep_alloca_1]]
+; CHECK-NEXT: %[[ee:[^ ]+]] = extractelement <3 x float> %[[v3xyz]], i32 2
+; CHECK-NEXT: %[[gep_alloca_2:[^ ]+]] = getelementptr inbounds [3 x float], [3 x float]* %[[alloca]], i32 0, i32 2
+; CHECK-NEXT: store float %[[ee]], float* %[[gep_alloca_2]]
+; Load [row].y from alloca:
+; CHECK-NEXT: %[[gep_alloca_row:[^ ]+]] = getelementptr inbounds [3 x float], [3 x float]* %[[alloca]], i32 0, i32 %[[row]]
+; CHECK-NEXT: %[[load_row:[^ ]+]] = load float, float* %[[gep_alloca_row]]
+; Insert [row].x into vec2 (unused)
+; CHECK-NEXT: %[[v2x:[^ ]+]] = insertelement <2 x float> undef, float %{{[^,]+}}, i64 0
+; Insert [row].y into vec2
+; CHECK-NEXT: %[[v2xy:[^ ]+]] = insertelement <2 x float> %[[v2x]], float %[[load_row]], i64 1
+; Extract component for .y output:
+; CHECK-NEXT: %[[ee:[^ ]+]] = extractelement <2 x float> %[[v2xy]], i32 1
+; CHECK-NEXT: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float %[[ee]])
+
+target datalayout = "e-m:e-p:32:32-i1:32-i8:32-i16:32-i32:32-i64:64-f16:32-f32:32-f64:64-n8:16:32:64"
+target triple = "dxil-ms-dx"
+
+%"$Globals" = type { %class.matrix.float.3.2 }
+%class.matrix.float.3.2 = type { [3 x <2 x float>] }
+%dx.types.Handle = type { i8* }
+%dx.types.ResourceProperties = type { i32, i32 }
+
+@"$Globals" = external constant %"$Globals"
+
+; Function Attrs: nounwind readnone
+declare <2 x float>* @"dx.hl.subscript.colMajor[].rn.<2 x float>* (i32, %class.matrix.float.3.2*, i32, i32)"(i32, %class.matrix.float.3.2*, i32, i32) #0
+
+; Function Attrs: nounwind readnone
+declare %"$Globals"* @"dx.hl.subscript.cb.rn.%\22$Globals\22* (i32, %dx.types.Handle, i32)"(i32, %dx.types.Handle, i32) #0
+
+; Function Attrs: nounwind readnone
+declare %dx.types.Handle @"dx.hl.createhandle..%dx.types.Handle (i32, %\22$Globals\22*, i32)"(i32, %"$Globals"*, i32) #0
+
+; Function Attrs: nounwind readnone
+declare %dx.types.Handle @"dx.hl.annotatehandle..%dx.types.Handle (i32, %dx.types.Handle, %dx.types.ResourceProperties, %\22$Globals\22)"(i32, %dx.types.Handle, %dx.types.ResourceProperties, %"$Globals") #0
+
+; Function Attrs: nounwind
+define void @main(float* noalias, i32) #1 {
+entry:
+  %2 = call %dx.types.Handle @"dx.hl.createhandle..%dx.types.Handle (i32, %\22$Globals\22*, i32)"(i32 0, %"$Globals"* @"$Globals", i32 0)
+  %3 = call %dx.types.Handle @"dx.hl.annotatehandle..%dx.types.Handle (i32, %dx.types.Handle, %dx.types.ResourceProperties, %\22$Globals\22)"(i32 14, %dx.types.Handle %2, %dx.types.ResourceProperties { i32 13, i32 28 }, %"$Globals" undef)
+  %4 = call %"$Globals"* @"dx.hl.subscript.cb.rn.%\22$Globals\22* (i32, %dx.types.Handle, i32)"(i32 6, %dx.types.Handle %3, i32 0)
+  %5 = getelementptr inbounds %"$Globals", %"$Globals"* %4, i32 0, i32 0
+  %6 = add i32 3, %1
+  %7 = call <2 x float>* @"dx.hl.subscript.colMajor[].rn.<2 x float>* (i32, %class.matrix.float.3.2*, i32, i32)"(i32 1, %class.matrix.float.3.2* %5, i32 %1, i32 %6)
+  %8 = load <2 x float>, <2 x float>* %7
+  %9 = extractelement <2 x float> %8, i32 1
+  store float %9, float* %0
+  ret void
+}
+
+attributes #0 = { nounwind readnone }
+attributes #1 = { nounwind }
+
+!llvm.module.flags = !{!0}
+!pauseresume = !{!1}
+!llvm.ident = !{!2}
+!dx.version = !{!3}
+!dx.valver = !{!4}
+!dx.shaderModel = !{!5}
+!dx.typeAnnotations = !{!6, !10}
+!dx.entryPoints = !{!19}
+!dx.fnprops = !{!23}
+!dx.options = !{!24, !25}
+
+!0 = !{i32 2, !"Debug Info Version", i32 3}
+!1 = !{!"hlsl-hlemit", !"hlsl-hlensure"}
+!2 = !{!"dxc(private) 1.9.0.5175 (dxc-version-print, 93ba0122d-dirty)"}
+!3 = !{i32 1, i32 0}
+!4 = !{i32 1, i32 10}
+!5 = !{!"vs", i32 6, i32 0}
+!6 = !{i32 0, %"$Globals" undef, !7}
+!7 = !{i32 28, !8}
+!8 = !{i32 6, !"f3x2", i32 2, !9, i32 3, i32 0, i32 7, i32 9}
+!9 = !{i32 3, i32 2, i32 2}
+!10 = !{i32 1, void (float*, i32)* @main, !11}
+!11 = !{!12, !14, !17}
+!12 = !{i32 0, !13, !13}
+!13 = !{}
+!14 = !{i32 1, !15, !16}
+!15 = !{i32 4, !"OUT", i32 7, i32 9}
+!16 = !{i32 0}
+!17 = !{i32 0, !18, !16}
+!18 = !{i32 4, !"ROW", i32 7, i32 4}
+!19 = !{void (float*, i32)* @main, !"main", null, !20, null}
+!20 = !{null, null, !21, null}
+!21 = !{!22}
+!22 = !{i32 0, %"$Globals"* @"$Globals", !"$Globals", i32 0, i32 -1, i32 1, i32 28, null}
+!23 = !{void (float*, i32)* @main, i32 1}
+!24 = !{i32 64}
+!25 = !{i32 -1}