[js/webgpu] Float16Array polyfill for uniform

js/common/lib/tensor-impl.ts

@@ -142,15 +142,7 @@ export class Tensor implements TensorInterface {
             throw new TypeError(`Unsupported tensor type: ${arg0}.`);
           }
           if (Array.isArray(arg1)) {
-            if (arg0 === 'float16' && typedArrayConstructor === Uint16Array) {
-              // When no Float16Array polyfill is used, we cannot create 'float16' tensor from number array.
-              //
-              // Throw error here because when user try to use number array as data,
-              // e.g. new Tensor('float16', [1, 2, 3, 4], dims)), it will actually call
-              // Uint16Array.from(arg1) which generates wrong data.
-              throw new TypeError(
-                  'Creating a float16 tensor from number array is not supported. Please use Uint16Array as data.');
-            } else if (arg0 === 'uint64' || arg0 === 'int64') {
+            if (arg0 === 'uint64' || arg0 === 'int64') {
               // use 'as any' here because:
               // 1. TypeScript's check on type of 'Array.isArray()' does not work with readonly arrays.
               // see https://github.com/microsoft/TypeScript/issues/17002

js/web/lib/wasm/jsep/backend-webgpu.ts

@@ -512,8 +512,14 @@ export class WebGpuBackend {
         } else if (v.type === DataType.uint32) {
           new Uint32Array(arrayBuffer, offset, data.length).set(data);
         } else if (v.type === DataType.float16) {
-          // TODO: use Float16Array.
-          new Uint16Array(arrayBuffer, offset, data.length).set(data);
+          if (typeof Float16Array !== 'undefined') {
+            new Float16Array(arrayBuffer, offset, data.length).set(data);
+          } else {
+            // Fallback to Uint16Array when Float16Array polyfill is not available, unit test only.
+            // eslint-disable-next-line no-console
+            console.warn('Unit test only, please make sure the float16 data has been encoded as float 16 bits.');
+            new Uint16Array(arrayBuffer, offset, data.length).set(data);
+          }
         } else if (v.type === DataType.float) {
           new Float32Array(arrayBuffer, offset, data.length).set(data);
         } else {

js/web/lib/wasm/jsep/init.ts

@@ -4,7 +4,7 @@
 import {Env} from 'onnxruntime-common';
 
 import {OrtWasmModule} from '../binding/ort-wasm';
-import {DataType, getTensorElementSize} from '../wasm-common';
+import {DataType, Float16ArrayType, getTensorElementSize} from '../wasm-common';
 
 import {WebGpuBackend} from './backend-webgpu';
 import {LOG_DEBUG} from './log';
@@ -19,6 +19,16 @@ class TensorViewImpl implements TensorView {
       private module: OrtWasmModule, public readonly dataType: number, public readonly data: number,
       public readonly dims: readonly number[]) {}
 
+  getFloat16Array(): Float16ArrayType {
+    if (this.dataType !== DataType.float16) {
+      throw new Error('Invalid data type');
+    }
+    const elementCount = ShapeUtil.size(this.dims);
+    const float16ViewConstructor = typeof Float16Array !== 'undefined' ? Float16Array : Uint16Array;
+    return elementCount === 0 ? new float16ViewConstructor() :
+                                new float16ViewConstructor(this.module.HEAP8.buffer, this.data, elementCount);
+  }
+
   getFloat32Array(): Float32Array {
     if (this.dataType !== DataType.float) {
       throw new Error('Invalid data type');

js/web/lib/wasm/jsep/tensor-view.ts

@@ -3,7 +3,7 @@
 
 import {Tensor} from 'onnxruntime-common';
 
-import {tensorTypeToTypedArrayConstructor} from '../wasm-common';
+import {Float16ArrayType, tensorTypeToTypedArrayConstructor} from '../wasm-common';
 
 export const createView = (dataBuffer: ArrayBuffer, type: Tensor.Type): Int32Array|Uint32Array|BigInt64Array|
     BigUint64Array|Uint8Array|Float32Array|Float64Array|Int8Array|Int16Array|Uint16Array =>
@@ -17,6 +17,11 @@ export interface TensorView {
   readonly dataType: number;
   readonly dims: readonly number[];
 
+  /**
+   * get a Float16Array data view of the tensor data. tensor data must be on CPU.
+   */
+  getFloat16Array(): Float16ArrayType;
+
   /**
    * get a Float32Array data view of the tensor data. tensor data must be on CPU.
    */

js/web/lib/wasm/jsep/webgpu/ops/pad.ts

@@ -200,7 +200,9 @@ const createPadProgramInfo = (inputs: readonly TensorView[], attributes: PadAttr
 const createPadAttributesFromInputs = (inputs: readonly TensorView[], attributes: PadAttributes): PadAttributes => {
   if (inputs.length > 1) {
     const bigInt64Pads = inputs[1].getBigInt64Array();
-    const value = (inputs.length >= 3 && inputs[2].data) ? inputs[2].getFloat32Array()[0] : 0.0;
+    const value = (inputs.length >= 3 && inputs[2].data) ?
+        (inputs[2].dataType === DataType.float16 ? inputs[2].getFloat16Array()[0] : inputs[2].getFloat32Array()[0]) :
+        0.0;
 
     const inputRank = inputs[0].dims.length;
     const updatePads = new Int32Array(2 * inputRank).fill(0);

js/web/lib/wasm/wasm-common.ts

@@ -8,6 +8,7 @@ declare global {
   // eslint-disable-next-line @typescript-eslint/naming-convention, @typescript-eslint/no-explicit-any
   const Float16Array: any;
 }
+export type Float16ArrayType = InstanceType<typeof Float16Array>;
 
 // This file includes common definitions. They do NOT have dependency on the WebAssembly instance.
 

js/web/package.json

@@ -63,7 +63,8 @@
     "minimatch": "^7.4.2",
     "minimist": "^1.2.8",
     "numpy-parser": "^1.2.3",
-    "strip-json-comments": "^5.0.0"
+    "strip-json-comments": "^5.0.0",
+    "@petamoriken/float16": "^3.8.4"
   },
   "main": "dist/ort-web.node.js",
   "exports": {

js/web/test/data/ops/pad-f16.jsonc

@@ -0,0 +1,74 @@
+[
+  {
+    "name": "constant 2D float16",
+    "operator": "Pad",
+    "opset": { "domain": "", "version": 10 },
+    "attributes": [
+      { "name": "mode", "data": "constant", "type": "string" },
+      { "name": "value", "data": 15565, "type": "float" },
+      { "name": "pads", "data": [3, 2, 2, 3], "type": "ints" }
+    ],
+    "cases": [
+      {
+        "name": "[2,2]->[7,7]",
+        "inputs": [
+          {
+            "data": [1.0, 2.0, 3.0, 4.5],
+            "dims": [2, 2],
+            "type": "float16"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [
+              1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2,
+              1.2, 1.2, 1.0, 2.0, 1.2, 1.2, 1.2, 1.2, 1.2, 3.0, 4.5, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2,
+              1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2
+            ],
+            "dims": [7, 7],
+            "type": "float16"
+          }
+        ]
+      }
+    ]
+  },
+  {
+    "name": "constant 2D float16",
+    "operator": "Pad",
+    "opset": { "domain": "", "version": 19 },
+    "attributes": [{ "name": "mode", "data": "constant", "type": "string" }],
+    "cases": [
+      {
+        "name": "[2,2]->[7,7]",
+        "inputs": [
+          {
+            "data": [1.0, 2.0, 3.0, 4.5],
+            "dims": [2, 2],
+            "type": "float16"
+          },
+          {
+            "data": [3, 2, 2, 3],
+            "dims": [4],
+            "type": "int64"
+          },
+          {
+            "data": [1.2],
+            "dims": [1],
+            "type": "float16"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [
+              1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2,
+              1.2, 1.2, 1.0, 2.0, 1.2, 1.2, 1.2, 1.2, 1.2, 3.0, 4.5, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2,
+              1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2
+            ],
+            "dims": [7, 7],
+            "type": "float16"
+          }
+        ]
+      }
+    ]
+  }
+]

js/web/test/suite-test-list.jsonc

@@ -1381,6 +1381,7 @@
       "gelu.jsonc",
       "pad.jsonc",
       "pad-big.jsonc",
+      "pad-f16.jsonc",
       "pow.jsonc",
       "pow_int32.jsonc",
       "pow-big-number.jsonc",

js/web/test/test-runner.ts

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+import {Float16Array} from '@petamoriken/float16';
 import {expect} from 'chai';
 import * as ort from 'onnxruntime-common';
 import {extname} from 'path';
@@ -390,6 +391,7 @@ export class TensorResultValidator {
       case 'string':
         return this.strictEqual(actual.data, expected.data);
 
+      case 'float16':
       case 'float32':
       case 'float64':
         return this.floatEqual(
@@ -890,8 +892,11 @@ async function runProtoOpTestcase(
   const fetches: Record<string, Pick<ort.Tensor, 'dims'|'type'>> = {};
   testCase.inputs.forEach((input, i) => {
     if (input.data) {
-      let data: number[]|BigUint64Array|BigInt64Array = input.data;
-      if (input.type === 'uint64') {
+      let data: number[]|BigUint64Array|BigInt64Array|Uint16Array = input.data;
+      if (input.type === 'float16') {
+        const floata16Array = Float16Array.from(input.data);
+        data = new Uint16Array(floata16Array.buffer, 0, floata16Array.length);
+      } else if (input.type === 'uint64') {
         data = BigUint64Array.from(input.data.map(BigInt));
       } else if (input.type === 'int64') {
         data = BigInt64Array.from(input.data.map(BigInt));
@@ -904,8 +909,11 @@ async function runProtoOpTestcase(
   const expectedOutputNames: string[] = [];
   testCase.outputs.forEach((output, i) => {
     if (output.data) {
-      let data: number[]|BigUint64Array|BigInt64Array = output.data;
-      if (output.type === 'uint64') {
+      let data: number[]|BigUint64Array|BigInt64Array|Uint16Array = output.data;
+      if (output.type === 'float16') {
+        const floata16Array = Float16Array.from(output.data);
+        data = new Uint16Array(floata16Array.buffer, 0, floata16Array.length);
+      } else if (output.type === 'uint64') {
         data = BigUint64Array.from(output.data.map(BigInt));
       } else if (output.type === 'int64') {
         data = BigInt64Array.from(output.data.map(BigInt));