reader/spirv: Handle the MatrixStride decoration

Add `transform::DecomposeStridedMatrix`, which replaces matrix members of storage or uniform buffer structures, that have a [[stride]] decoration, into an array
of N column vectors.

This is required to correctly handle `mat2x2` matrices in UBOs, as std140 rules will expect a default stride of 16 bytes, when in WGSL the default structure layout expects a stride of 8 bytes.

Bug: tint:1047
Change-Id: If5ca3c6ec087bbc1ac31a8d9a657b99bf34042a4
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/59840
Reviewed-by: David Neto <dneto@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>

test/BUILD.gn

@@ -289,6 +289,7 @@ tint_unittests_source_set("tint_unittests_core_src") {
     "../src/transform/calculate_array_length_test.cc",
     "../src/transform/canonicalize_entry_point_io_test.cc",
     "../src/transform/decompose_memory_access_test.cc",
+    "../src/transform/decompose_strided_matrix_test.cc",
     "../src/transform/external_texture_transform_test.cc",
     "../src/transform/first_index_offset_test.cc",
     "../src/transform/fold_constants_test.cc",

test/layout/storage/mat2x2/f32.wgsl

@@ -0,0 +1,11 @@
+[[block]]
+struct SSBO {
+  m : mat2x2<f32>;
+};
+[[group(0), binding(0)]] var<storage, read_write> ssbo : SSBO;
+
+[[stage(compute), workgroup_size(1)]]
+fn f() {
+    let v = ssbo.m;
+    ssbo.m = v;
+}

test/layout/storage/mat2x2/f32.wgsl.expected.hlsl

@@ -0,0 +1,17 @@
+RWByteAddressBuffer ssbo : register(u0, space0);
+
+float2x2 tint_symbol(RWByteAddressBuffer buffer, uint offset) {
+  return float2x2(asfloat(buffer.Load2((offset + 0u))), asfloat(buffer.Load2((offset + 8u))));
+}
+
+void tint_symbol_2(RWByteAddressBuffer buffer, uint offset, float2x2 value) {
+  buffer.Store2((offset + 0u), asuint(value[0u]));
+  buffer.Store2((offset + 8u), asuint(value[1u]));
+}
+
+[numthreads(1, 1, 1)]
+void f() {
+  const float2x2 v = tint_symbol(ssbo, 0u);
+  tint_symbol_2(ssbo, 0u, v);
+  return;
+}

test/layout/storage/mat2x2/f32.wgsl.expected.msl

@@ -0,0 +1,13 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct SSBO {
+  /* 0x0000 */ float2x2 m;
+};
+
+kernel void f(device SSBO& ssbo [[buffer(0)]]) {
+  float2x2 const v = ssbo.m;
+  ssbo.m = v;
+  return;
+}
+

test/layout/storage/mat2x2/f32.wgsl.expected.spvasm

@@ -0,0 +1,38 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 17
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %f "f"
+               OpExecutionMode %f LocalSize 1 1 1
+               OpName %SSBO "SSBO"
+               OpMemberName %SSBO 0 "m"
+               OpName %ssbo "ssbo"
+               OpName %f "f"
+               OpDecorate %SSBO Block
+               OpMemberDecorate %SSBO 0 Offset 0
+               OpMemberDecorate %SSBO 0 ColMajor
+               OpMemberDecorate %SSBO 0 MatrixStride 8
+               OpDecorate %ssbo DescriptorSet 0
+               OpDecorate %ssbo Binding 0
+      %float = OpTypeFloat 32
+    %v2float = OpTypeVector %float 2
+%mat2v2float = OpTypeMatrix %v2float 2
+       %SSBO = OpTypeStruct %mat2v2float
+%_ptr_StorageBuffer_SSBO = OpTypePointer StorageBuffer %SSBO
+       %ssbo = OpVariable %_ptr_StorageBuffer_SSBO StorageBuffer
+       %void = OpTypeVoid
+          %7 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_StorageBuffer_mat2v2float = OpTypePointer StorageBuffer %mat2v2float
+          %f = OpFunction %void None %7
+         %10 = OpLabel
+         %14 = OpAccessChain %_ptr_StorageBuffer_mat2v2float %ssbo %uint_0
+         %15 = OpLoad %mat2v2float %14
+         %16 = OpAccessChain %_ptr_StorageBuffer_mat2v2float %ssbo %uint_0
+               OpStore %16 %15
+               OpReturn
+               OpFunctionEnd

test/layout/storage/mat2x2/f32.wgsl.expected.wgsl

@@ -0,0 +1,12 @@
+[[block]]
+struct SSBO {
+  m : mat2x2<f32>;
+};
+
+[[group(0), binding(0)]] var<storage, read_write> ssbo : SSBO;
+
+[[stage(compute), workgroup_size(1)]]
+fn f() {
+  let v = ssbo.m;
+  ssbo.m = v;
+}

test/layout/storage/mat2x2/stride/16.spvasm

@@ -0,0 +1,33 @@
+                                  OpCapability Shader
+                                  OpMemoryModel Logical GLSL450
+                                  OpEntryPoint GLCompute %f "f"
+                                  OpExecutionMode %f LocalSize 1 1 1
+                                  OpName %SSBO "SSBO"
+                                  OpMemberName %SSBO 0 "m"
+                                  OpName %ssbo "ssbo"
+                                  OpName %f "f"
+                                  OpDecorate %SSBO Block
+                                  OpMemberDecorate %SSBO 0 Offset 0
+                                  OpMemberDecorate %SSBO 0 ColMajor
+                                  OpMemberDecorate %SSBO 0 MatrixStride 16
+                                  OpDecorate %ssbo DescriptorSet 0
+                                  OpDecorate %ssbo Binding 0
+                         %float = OpTypeFloat 32
+                       %v2float = OpTypeVector %float 2
+                   %mat2v2float = OpTypeMatrix %v2float 2
+                          %SSBO = OpTypeStruct %mat2v2float
+       %_ptr_StorageBuffer_SSBO = OpTypePointer StorageBuffer %SSBO
+                          %ssbo = OpVariable %_ptr_StorageBuffer_SSBO StorageBuffer
+                          %void = OpTypeVoid
+                             %7 = OpTypeFunction %void
+                          %uint = OpTypeInt 32 0
+                        %uint_0 = OpConstant %uint 0
+%_ptr_StorageBuffer_mat2v2float = OpTypePointer StorageBuffer %mat2v2float
+                             %f = OpFunction %void None %7
+                            %10 = OpLabel
+                            %14 = OpAccessChain %_ptr_StorageBuffer_mat2v2float %ssbo %uint_0
+                            %15 = OpLoad %mat2v2float %14
+                            %16 = OpAccessChain %_ptr_StorageBuffer_mat2v2float %ssbo %uint_0
+                                  OpStore %16 %15
+                                  OpReturn
+                                  OpFunctionEnd

test/layout/storage/mat2x2/stride/16.spvasm.expected.hlsl

@@ -0,0 +1,47 @@
+struct tint_padded_array_element {
+  float2 el;
+};
+
+RWByteAddressBuffer ssbo : register(u0, space0);
+
+float2x2 arr_to_mat2x2_stride_16(tint_padded_array_element arr[2]) {
+  return float2x2(arr[0u].el, arr[1u].el);
+}
+
+typedef tint_padded_array_element mat2x2_stride_16_to_arr_ret[2];
+mat2x2_stride_16_to_arr_ret mat2x2_stride_16_to_arr(float2x2 mat) {
+  const tint_padded_array_element tint_symbol_4[2] = {{mat[0u]}, {mat[1u]}};
+  return tint_symbol_4;
+}
+
+typedef tint_padded_array_element tint_symbol_ret[2];
+tint_symbol_ret tint_symbol(RWByteAddressBuffer buffer, uint offset) {
+  tint_padded_array_element arr_1[2] = (tint_padded_array_element[2])0;
+  {
+    for(uint i = 0u; (i < 2u); i = (i + 1u)) {
+      arr_1[i].el = asfloat(buffer.Load2((offset + (i * 16u))));
+    }
+  }
+  return arr_1;
+}
+
+void tint_symbol_2(RWByteAddressBuffer buffer, uint offset, tint_padded_array_element value[2]) {
+  tint_padded_array_element array[2] = value;
+  {
+    for(uint i_1 = 0u; (i_1 < 2u); i_1 = (i_1 + 1u)) {
+      buffer.Store2((offset + (i_1 * 16u)), asuint(array[i_1].el));
+    }
+  }
+}
+
+void f_1() {
+  const float2x2 x_15 = arr_to_mat2x2_stride_16(tint_symbol(ssbo, 0u));
+  tint_symbol_2(ssbo, 0u, mat2x2_stride_16_to_arr(x_15));
+  return;
+}
+
+[numthreads(1, 1, 1)]
+void f() {
+  f_1();
+  return;
+}

test/layout/storage/mat2x2/stride/16.spvasm.expected.msl

@@ -0,0 +1,34 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct tint_padded_array_element {
+  /* 0x0000 */ packed_float2 el;
+  /* 0x0008 */ int8_t tint_pad[8];
+};
+struct tint_array_wrapper {
+  /* 0x0000 */ tint_padded_array_element arr[2];
+};
+struct SSBO {
+  /* 0x0000 */ tint_array_wrapper m;
+};
+
+float2x2 arr_to_mat2x2_stride_16(tint_array_wrapper arr) {
+  return float2x2(arr.arr[0u].el, arr.arr[1u].el);
+}
+
+tint_array_wrapper mat2x2_stride_16_to_arr(float2x2 mat) {
+  tint_array_wrapper const tint_symbol = {.arr={{.el=mat[0u]}, {.el=mat[1u]}}};
+  return tint_symbol;
+}
+
+void f_1(device SSBO& ssbo) {
+  float2x2 const x_15 = arr_to_mat2x2_stride_16(ssbo.m);
+  ssbo.m = mat2x2_stride_16_to_arr(x_15);
+  return;
+}
+
+kernel void f(device SSBO& ssbo [[buffer(0)]]) {
+  f_1(ssbo);
+  return;
+}
+

test/layout/storage/mat2x2/stride/16.spvasm.expected.spvasm

@@ -0,0 +1,70 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 39
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %f "f"
+               OpExecutionMode %f LocalSize 1 1 1
+               OpName %SSBO "SSBO"
+               OpMemberName %SSBO 0 "m"
+               OpName %ssbo "ssbo"
+               OpName %arr_to_mat2x2_stride_16 "arr_to_mat2x2_stride_16"
+               OpName %arr "arr"
+               OpName %mat2x2_stride_16_to_arr "mat2x2_stride_16_to_arr"
+               OpName %mat "mat"
+               OpName %f_1 "f_1"
+               OpName %f "f"
+               OpDecorate %SSBO Block
+               OpMemberDecorate %SSBO 0 Offset 0
+               OpDecorate %_arr_v2float_uint_2 ArrayStride 16
+               OpDecorate %ssbo DescriptorSet 0
+               OpDecorate %ssbo Binding 0
+      %float = OpTypeFloat 32
+    %v2float = OpTypeVector %float 2
+       %uint = OpTypeInt 32 0
+     %uint_2 = OpConstant %uint 2
+%_arr_v2float_uint_2 = OpTypeArray %v2float %uint_2
+       %SSBO = OpTypeStruct %_arr_v2float_uint_2
+%_ptr_StorageBuffer_SSBO = OpTypePointer StorageBuffer %SSBO
+       %ssbo = OpVariable %_ptr_StorageBuffer_SSBO StorageBuffer
+%mat2v2float = OpTypeMatrix %v2float 2
+          %9 = OpTypeFunction %mat2v2float %_arr_v2float_uint_2
+     %uint_0 = OpConstant %uint 0
+     %uint_1 = OpConstant %uint 1
+         %19 = OpTypeFunction %_arr_v2float_uint_2 %mat2v2float
+       %void = OpTypeVoid
+         %26 = OpTypeFunction %void
+%_ptr_StorageBuffer__arr_v2float_uint_2 = OpTypePointer StorageBuffer %_arr_v2float_uint_2
+%arr_to_mat2x2_stride_16 = OpFunction %mat2v2float None %9
+        %arr = OpFunctionParameter %_arr_v2float_uint_2
+         %13 = OpLabel
+         %15 = OpCompositeExtract %v2float %arr 0
+         %17 = OpCompositeExtract %v2float %arr 1
+         %18 = OpCompositeConstruct %mat2v2float %15 %17
+               OpReturnValue %18
+               OpFunctionEnd
+%mat2x2_stride_16_to_arr = OpFunction %_arr_v2float_uint_2 None %19
+        %mat = OpFunctionParameter %mat2v2float
+         %22 = OpLabel
+         %23 = OpCompositeExtract %v2float %mat 0
+         %24 = OpCompositeExtract %v2float %mat 1
+         %25 = OpCompositeConstruct %_arr_v2float_uint_2 %23 %24
+               OpReturnValue %25
+               OpFunctionEnd
+        %f_1 = OpFunction %void None %26
+         %29 = OpLabel
+         %32 = OpAccessChain %_ptr_StorageBuffer__arr_v2float_uint_2 %ssbo %uint_0
+         %33 = OpLoad %_arr_v2float_uint_2 %32
+         %30 = OpFunctionCall %mat2v2float %arr_to_mat2x2_stride_16 %33
+         %34 = OpAccessChain %_ptr_StorageBuffer__arr_v2float_uint_2 %ssbo %uint_0
+         %35 = OpFunctionCall %_arr_v2float_uint_2 %mat2x2_stride_16_to_arr %30
+               OpStore %34 %35
+               OpReturn
+               OpFunctionEnd
+          %f = OpFunction %void None %26
+         %37 = OpLabel
+         %38 = OpFunctionCall %void %f_1
+               OpReturn
+               OpFunctionEnd

test/layout/storage/mat2x2/stride/16.spvasm.expected.wgsl

@@ -0,0 +1,25 @@
+[[block]]
+struct SSBO {
+  m : [[stride(16)]] array<vec2<f32>, 2>;
+};
+
+[[group(0), binding(0)]] var<storage, read_write> ssbo : SSBO;
+
+fn arr_to_mat2x2_stride_16(arr : [[stride(16)]] array<vec2<f32>, 2>) -> mat2x2<f32> {
+  return mat2x2<f32>(arr[0u], arr[1u]);
+}
+
+fn mat2x2_stride_16_to_arr(mat : mat2x2<f32>) -> [[stride(16)]] array<vec2<f32>, 2> {
+  return [[stride(16)]] array<vec2<f32>, 2>(mat[0u], mat[1u]);
+}
+
+fn f_1() {
+  let x_15 : mat2x2<f32> = arr_to_mat2x2_stride_16(ssbo.m);
+  ssbo.m = mat2x2_stride_16_to_arr(x_15);
+  return;
+}
+
+[[stage(compute), workgroup_size(1, 1, 1)]]
+fn f() {
+  f_1();
+}