Add chained and compound swizzle assignment tests (#4596)

Author: sudonatalie

src/webgpu/listing_meta.json

@@ -2065,7 +2065,6 @@
   "webgpu:shader,execution,statement,increment_decrement:vec4_element_decrement:*": { "subcaseMS": 5.300 },
   "webgpu:shader,execution,statement,increment_decrement:vec4_element_increment:*": { "subcaseMS": 6.300 },
   "webgpu:shader,execution,statement,phony:executes:*": { "subcaseMS": 129.949 },
-  "webgpu:shader,execution,statement,swizzle_assignment:swizzle_assignment_chained:*": { "subcaseMS": 0.112 },
   "webgpu:shader,execution,statement,swizzle_assignment:swizzle_assignment_vars:*": { "subcaseMS": 1200.970 },
   "webgpu:shader,execution,statement,swizzle_assignment:swizzle_compound_assignment:*": { "subcaseMS": 0.091 },
   "webgpu:shader,execution,value_init:array,nested:*": { "subcaseMS": 3004.523 },

src/webgpu/shader/execution/statement/swizzle_assignment.spec.ts

@@ -216,6 +216,36 @@ const kSwizzleAssignmentCases: Record<string, SwizzleAssignmentCase> = {
     rhs: 'vec2<bool>(false, false)',
     expected: [0, 1, 0],
   },
+  // v = vec4u(1, 2, 3, 4)
+  // v.xy.x = 5;
+  vec4u_xy_x_literal: {
+    elemType: 'u32',
+    vecSize: 4,
+    initial: [1, 2, 3, 4],
+    swizzle: 'xy.x',
+    rhs: '5',
+    expected: [5, 2, 3, 4],
+  },
+  // v = vec3f(1.0, 2.0, 3.0)
+  // v.zyx.yz = vec2f(5.0, 6.0);
+  vec3f_zyx_yz_vec2f: {
+    elemType: 'f32',
+    vecSize: 3,
+    initial: [1.0, 2.0, 3.0],
+    swizzle: 'zyx.yz',
+    rhs: 'vec2f(5.0, 6.0)',
+    expected: [6.0, 5.0, 3.0],
+  },
+  // v = vec3i(-1, 0, -1)
+  // v.xz.yx = vec2i(2);
+  vec2i_xz_yx_vec2i: {
+    elemType: 'i32',
+    vecSize: 3,
+    initial: [-1, 0, -1],
+    swizzle: 'xz.yx',
+    rhs: 'vec2i(2,3)',
+    expected: [3, 0, 2],
+  },
 };
 
 g.test('swizzle_assignment_vars')
@@ -285,10 +315,108 @@ fn main() {
     runSwizzleAssignmentTest(t, elemType, expected, wgsl);
   });
 
-g.test('swizzle_assignment_chained')
-  .desc('Tests the value of a vector after swizzle assignment on a chained swizzle.')
-  .unimplemented();
+interface SwizzleCompoundAssignmentCase extends SwizzleAssignmentCase {
+  op: string;
+}
+
+const kSwizzleCompoundAssignmentCases: Record<string, SwizzleCompoundAssignmentCase> = {
+  // v = vec4u(1, 2, 3, 4)
+  // v.w += 5;
+  vec4u_w_add_5: {
+    elemType: 'u32',
+    vecSize: 4,
+    initial: [1, 2, 3, 4],
+    swizzle: 'w',
+    op: '+=',
+    rhs: '5',
+    expected: [1, 2, 3, 9],
+  },
+  // v = vec4u(1, 2, 3, 4)
+  // v.xy *= vec2u(6, 7);
+  vec4u_xy_mul_vec2u: {
+    elemType: 'u32',
+    vecSize: 4,
+    initial: [1, 2, 3, 4],
+    swizzle: 'xy',
+    op: '*=',
+    rhs: 'vec2u(6, 7)',
+    expected: [6, 14, 3, 4],
+  },
+  // v = vec3i(10, 20, 30)
+  // v.zx += vec2i(100);
+  vec3i_zx_add_vec2i: {
+    elemType: 'i32',
+    vecSize: 3,
+    initial: [10, 20, 30],
+    swizzle: 'zx',
+    op: '+=',
+    rhs: 'vec2i(100)',
+    expected: [110, 20, 130],
+  },
+  // v = vec3f(1.0, 2.0, 3.0)
+  // v.xy *= vec2f(0.5, 2.0);
+  vec3f_xy_mul_vec2f: {
+    elemType: 'f32',
+    vecSize: 3,
+    initial: [1.0, 2.0, 3.0],
+    swizzle: 'xy',
+    op: '*=',
+    rhs: 'vec2f(0.5, 2.0)',
+    expected: [0.5, 4.0, 3.0],
+  },
+};
 
 g.test('swizzle_compound_assignment')
   .desc('Tests the value of a vector after compound swizzle assignment.')
-  .unimplemented();
+  .params(u =>
+    u
+      .combine('case', keysOf(kSwizzleCompoundAssignmentCases))
+      .beginSubcases()
+      .combine('address_space', ['function', 'private', 'workgroup', 'storage'])
+      .combine('memory_view', ['ref', 'ptr'])
+  )
+  .fn(t => {
+    const { elemType, vecSize, initial, swizzle, op, rhs, expected } =
+      kSwizzleCompoundAssignmentCases[t.params.case];
+
+    const vecType = `vec${vecSize}<${elemType}>`;
+    const initialValues = initial.join(', ');
+
+    const var_ref = t.params.address_space === 'storage' ? 'outputs.v' : 'v';
+    const lhs =
+      t.params.memory_view === 'ptr'
+        ? `let ptr = &${var_ref}; ptr.${swizzle}`
+        : `${var_ref}.${swizzle}`;
+    const new_rhs = rhs.replaceAll(/\bv\b/g, `${var_ref}`);
+
+    const wgsl = `
+requires swizzle_assignment;
+${elemType === 'f16' ? 'enable f16;' : ''}
+
+struct Outputs {
+  ${t.params.address_space === 'storage' ? `v : ${vecType},` : ''}
+  data : array<${elemType}>,
+};
+
+@group(0) @binding(1) var<storage, read_write> outputs : Outputs;
+
+${
+  t.params.address_space === 'private' || t.params.address_space === 'workgroup'
+    ? `var<${t.params.address_space}> v : ${vecType};`
+    : ''
+}
+
+@compute @workgroup_size(1)
+fn main() {
+
+  ${t.params.address_space === 'function' ? `var v : ${vecType};` : ''}
+  ${var_ref} = ${vecType}(${initialValues});
+  ${lhs} ${op} ${new_rhs};
+
+  // Store result to Output
+  for (var i = 0; i < ${vecSize}; i++) {
+    outputs.data[i] = ${var_ref}[i];
+  }
+}`;
+    runSwizzleAssignmentTest(t, elemType, expected, wgsl);
+  });