add f16 to webgpu/shader/types.ts utilities (#3407)

* add f16 to webgpu/shader/types.ts utilities Udpate affected tests; - webgpu:shader,execution,zero_init: skip f16 cases - webgpu:shader,execution,robust_access: generalize to cover most f16 cases Bug: #3405 * Move TODO to test description * Rewrite the array generation Remove unused 'atomic' container type. Create an AlignmentAndSize type Document the generator type, and add a bunch of comments. * Update src/webgpu/shader/types.ts rewrite line comments as block comments Co-authored-by: Kai Ninomiya <[email protected]> * Update src/webgpu/shader/types.ts Co-authored-by: Kai Ninomiya <[email protected]> * Update src/webgpu/shader/types.ts Co-authored-by: Kai Ninomiya <[email protected]> * Sized arrays only need one yield statement --------- Co-authored-by: Kai Ninomiya <[email protected]>
gpuweb · Mar 22, 2024 · c2e75b4 · c2e75b4
1 parent 2abbfa3
commit c2e75b4
Show file tree

Hide file tree

Showing 4 changed files with 190 additions and 55 deletions.
diff --git a/src/webgpu/shader/execution/robust_access.spec.ts b/src/webgpu/shader/execution/robust_access.spec.ts
@@ -7,6 +7,7 @@ TODO: add tests to check that textureLoad operations stay in-bounds.
 
 import { makeTestGroup } from '../../../common/framework/test_group.js';
 import { assert } from '../../../common/util/util.js';
+import { Float16Array } from '../../../external/petamoriken/float16/float16.js';
 import { GPUTest } from '../../gpu_test.js';
 import { align } from '../../util/math.js';
 import { generateTypes, supportedScalarTypes, supportsAtomics } from '../types.js';
@@ -25,6 +26,7 @@ const kMinI32 = -0x8000_0000;
  */
 async function runShaderTest(
   t: GPUTest,
+  enables: string,
   stage: GPUShaderStageFlags,
   testSource: string,
   layout: GPUPipelineLayout,
@@ -41,7 +43,7 @@ async function runShaderTest(
     usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE,
   });
 
-  const source = `
+  const source = `${enables}
 struct Constants {
   zero: u32
 };
@@ -96,10 +98,12 @@ fn main() {
 /** Fill an ArrayBuffer with sentinel values, except clear a region to zero. */
 function testFillArrayBuffer(
   array: ArrayBuffer,
-  type: 'u32' | 'i32' | 'f32',
+  type: 'u32' | 'i32' | 'f16' | 'f32',
   { zeroByteStart, zeroByteCount }: { zeroByteStart: number; zeroByteCount: number }
 ) {
-  const constructor = { u32: Uint32Array, i32: Int32Array, f32: Float32Array }[type];
+  const constructor = { u32: Uint32Array, i32: Int32Array, f16: Float16Array, f32: Float32Array }[
+    type
+  ];
   assert(zeroByteCount % constructor.BYTES_PER_ELEMENT === 0);
   new constructor(array).fill(42);
   new constructor(array, zeroByteStart, zeroByteCount / constructor.BYTES_PER_ELEMENT).fill(0);
@@ -122,6 +126,7 @@ g.test('linear_memory')
     TODO: Test types like vec2<atomic<i32>>, if that's allowed.
     TODO: Test exprIndexAddon as constexpr.
     TODO: Test exprIndexAddon as pipeline-overridable constant expression.
+    TODO: Adjust test logic to support array of f16 in the uniform address space
   `
   )
   .params(u =>
@@ -168,10 +173,15 @@ g.test('linear_memory')
         { shadowingMode: 'function-scope' },
       ])
       .expand('isAtomic', p => (supportsAtomics(p) ? [false, true] : [false]))
-      .beginSubcases()
       .expand('baseType', supportedScalarTypes)
+      .beginSubcases()
       .expandWithParams(generateTypes)
   )
+  .beforeAllSubcases(t => {
+    if (t.params.baseType === 'f16') {
+      t.selectDeviceOrSkipTestCase('shader-f16');
+    }
+  })
   .fn(async t => {
     const {
       addressSpace,
@@ -189,6 +199,13 @@ g.test('linear_memory')
     assert(_kTypeInfo !== undefined, 'not an indexable type');
     assert('arrayLength' in _kTypeInfo);
 
+    if (baseType === 'f16' && addressSpace === 'uniform' && containerType === 'array') {
+      // Array elements must be aligned to 16 bytes, but the logic in generateTypes
+      // creates an array of vec4 of the baseType. But for f16 that's only 8 bytes.
+      // We would need to write more complex logic for that.
+      t.skip('Test logic does not handle array of f16 in the uniform address space');
+    }
+
     let usesCanary = false;
     let globalSource = '';
     let testFunctionSource = '';
@@ -429,6 +446,8 @@ fn runTest() -> u32 {
       ],
     });
 
+    const enables = t.params.baseType === 'f16' ? 'enable f16;' : '';
+
     // Run it.
     if (bufferBindingSize !== undefined && baseType !== 'bool') {
       const expectedData = new ArrayBuffer(testBufferSize);
@@ -450,6 +469,7 @@ fn runTest() -> u32 {
       // Run the shader, accessing the buffer.
       await runShaderTest(
         t,
+        enables,
         GPUShaderStage.COMPUTE,
         testSource,
         layout,
@@ -475,6 +495,6 @@ fn runTest() -> u32 {
         bufferBindingEnd
       );
     } else {
-      await runShaderTest(t, GPUShaderStage.COMPUTE, testSource, layout, []);
+      await runShaderTest(t, enables, GPUShaderStage.COMPUTE, testSource, layout, []);
     }
   });
diff --git a/src/webgpu/shader/execution/zero_init.spec.ts b/src/webgpu/shader/execution/zero_init.spec.ts
@@ -107,6 +107,10 @@ g.test('compute,zero_init')
               ? [true, false]
               : [false]) {
               for (const scalarType of supportedScalarTypes({ isAtomic, ...p })) {
+                // Fewer subcases: supportedScalarTypes was expanded to include f16
+                // but that may take too much time. It would require more complex code.
+                if (scalarType === 'f16') continue;
+
                 // Fewer subcases: For nested types, skip atomic u32 and non-atomic i32.
                 if (p._containerDepth > 0) {
                   if (scalarType === 'u32' && isAtomic) continue;