[js/webgpu] Float16Array polyfill for uniform

js/common/lib/tensor-impl-type-mapping.ts

@@ -14,7 +14,6 @@ export const NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP = new Map<string, SupportedTy
   ['uint8', Uint8Array],
   ['int8', Int8Array],
   ['uint16', Uint16Array],
-  ['float16', Uint16Array],
   ['int16', Int16Array],
   ['int32', Int32Array],
   ['bool', Uint8Array],
@@ -34,16 +33,22 @@ export const NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP = new Map<SupportedTypedArray
   [Uint32Array, 'uint32'],
 ]);
 
-// the following code allows delaying execution of BigInt checking. This allows lazy initialization for
-// NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP and NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP, which allows BigInt polyfill
-// if available.
-let isBigIntChecked = false;
-export const checkBigInt = () => {
-  if (!isBigIntChecked) {
-    isBigIntChecked = true;
-    const isBigInt64ArrayAvailable = typeof BigInt64Array !== 'undefined' && typeof BigInt64Array.from === 'function';
-    const isBigUint64ArrayAvailable =
-        typeof BigUint64Array !== 'undefined' && typeof BigUint64Array.from === 'function';
+// a dummy type declaration for Float16Array in case any polyfill is available.
+declare global {
+  // eslint-disable-next-line @typescript-eslint/naming-convention, @typescript-eslint/no-explicit-any
+  const Float16Array: any;
+}
+
+// the following code allows delaying execution of BigInt/Float16Array checking. This allows lazy initialization for
+// NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP and NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP, which allows BigInt/Float16Array
+// polyfill if available.
+let isTypedArrayChecked = false;
+export const checkTypedArray = () => {
+  if (!isTypedArrayChecked) {
+    isTypedArrayChecked = true;
+    const isBigInt64ArrayAvailable = typeof BigInt64Array !== 'undefined' && BigInt64Array.from;
+    const isBigUint64ArrayAvailable = typeof BigUint64Array !== 'undefined' && BigUint64Array.from;
+    const isFloat16ArrayAvailable = typeof Float16Array !== 'undefined' && Float16Array.from;
 
     if (isBigInt64ArrayAvailable) {
       NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP.set('int64', BigInt64Array);
@@ -53,5 +58,12 @@ export const checkBigInt = () => {
       NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP.set('uint64', BigUint64Array);
       NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP.set(BigUint64Array, 'uint64');
     }
+    if (isFloat16ArrayAvailable) {
+      NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP.set('float16', Float16Array);
+      NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP.set(Float16Array, 'float16');
+    } else {
+      // if Float16Array is not available, use 'Uint16Array' to store the data.
+      NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP.set('float16', Uint16Array);
+    }
   }
 };

js/common/lib/tensor-impl.ts

@@ -5,7 +5,7 @@ import {tensorToDataURL, tensorToImageData} from './tensor-conversion-impl.js';
 import {TensorToDataUrlOptions, TensorToImageDataOptions} from './tensor-conversion.js';
 import {tensorFromGpuBuffer, tensorFromImage, tensorFromPinnedBuffer, tensorFromTexture} from './tensor-factory-impl.js';
 import {CpuPinnedConstructorParameters, GpuBufferConstructorParameters, TensorFromGpuBufferOptions, TensorFromImageBitmapOptions, TensorFromImageDataOptions, TensorFromImageElementOptions, TensorFromTextureOptions, TensorFromUrlOptions, TextureConstructorParameters} from './tensor-factory.js';
-import {checkBigInt, NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP, NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP, SupportedTypedArray, SupportedTypedArrayConstructors} from './tensor-impl-type-mapping.js';
+import {checkTypedArray, NUMERIC_TENSOR_TYPE_TO_TYPEDARRAY_MAP, NUMERIC_TENSOR_TYPEDARRAY_TO_TYPE_MAP, SupportedTypedArray, SupportedTypedArrayConstructors} from './tensor-impl-type-mapping.js';
 import {calculateSize, tensorReshape} from './tensor-utils-impl.js';
 import {Tensor as TensorInterface} from './tensor.js';
 
@@ -67,8 +67,8 @@ export class Tensor implements TensorInterface {
       arg0: TensorType|TensorDataType|readonly string[]|readonly boolean[]|CpuPinnedConstructorParameters|
       TextureConstructorParameters|GpuBufferConstructorParameters,
       arg1?: TensorDataType|readonly number[]|readonly string[]|readonly boolean[], arg2?: readonly number[]) {
-    // perform one-time check for BigInt support
-    checkBigInt();
+    // perform one-time check for BigInt/Float16Array support
+    checkTypedArray();
 
     let type: TensorType;
     let dims: readonly number[];
@@ -142,7 +142,8 @@ export class Tensor implements TensorInterface {
             throw new TypeError(`Unsupported tensor type: ${arg0}.`);
           }
           if (Array.isArray(arg1)) {
-            if (arg0 === 'float16') {
+            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.

js/common/package.json

@@ -18,7 +18,8 @@
     "test": "mocha ./test/**/*.js --timeout 30000"
   },
   "devDependencies": {
-    "typedoc": "^0.23.22"
+    "typedoc": "^0.23.22",
+    "@petamoriken/float16": "^3.8.4"
   },
   "main": "dist/cjs/index.js",
   "exports": {

js/common/test/unit-tests/common.ts

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+import {Float16Array} from '@petamoriken/float16';
 import assert from 'assert/strict';
 import {Tensor} from 'onnxruntime-common';
 
@@ -36,12 +37,18 @@ export const BIGINT_TYPES = [
  */
 export const FLOAT16_TYPE = ['float16', Uint16Array, false] as const;
 
+/**
+ * float16 type, data represented by Float16Array polyfill
+ */
+export const FLOAT16_TYPE_POLYFILL = ['float16', Float16Array, false] as const;
+
 /**
  * A list of all numerical types.
  *
  * not including string and bool.
  */
-export const ALL_NUMERICAL_TYPES = [...NUMBER_COMPATIBLE_NUMERICAL_TYPES, ...BIGINT_TYPES, FLOAT16_TYPE];
+export const ALL_NUMERICAL_TYPES =
+    [...NUMBER_COMPATIBLE_NUMERICAL_TYPES, ...BIGINT_TYPES, FLOAT16_TYPE, FLOAT16_TYPE_POLYFILL];
 
 /**
  * a helper function to assert that a value is an array of a certain type

js/web/lib/onnxjs/backends/webgl/ops/resize-packed.ts

@@ -241,8 +241,8 @@ const prepareInputs = (inputs: Tensor[], attributes: UpsampleAttributes): [reado
 
 const parseScalesData = (scale: Tensor, mode: string, isResize: boolean): number[] => {
   const scales = Array.from(scale.floatData);
-  scalesValidation(scales, mode, isResize);
-  return scales;
+  scalesValidation(scales as number[], mode, isResize);
+  return scales as number[];
 };
 
 const parseScalesDataFromOutputSize =

js/web/lib/onnxjs/tensor.ts

@@ -4,6 +4,8 @@
 import {Guid} from 'guid-typescript';
 import Long from 'long';
 
+import {Float16ArrayType} from '../../lib/wasm/wasm-common';
+
 import {onnxruntime} from './ort-schema/flatbuffers/ort-generated';
 import {onnx} from './ort-schema/protobuf/onnx';
 import {decodeUtf8String, ProtoUtil, ShapeUtil} from './util';
@@ -13,6 +15,7 @@ import ortFbs = onnxruntime.experimental.fbs;
 export declare namespace Tensor {
   export interface DataTypeMap {
     bool: Uint8Array;
+    float16: Float16ArrayType;
     float32: Float32Array;
     float64: Float64Array;
     string: string[];
@@ -31,7 +34,7 @@ export declare namespace Tensor {
   export type BooleanType = Tensor.DataTypeMap['bool'];
   export type IntegerType = Tensor.DataTypeMap['int8']|Tensor.DataTypeMap['uint8']|Tensor.DataTypeMap['int16']|
                             Tensor.DataTypeMap['uint16']|Tensor.DataTypeMap['int32']|Tensor.DataTypeMap['uint32'];
-  export type FloatType = Tensor.DataTypeMap['float32']|Tensor.DataTypeMap['float64'];
+  export type FloatType = Tensor.DataTypeMap['float16']|Tensor.DataTypeMap['float32']|Tensor.DataTypeMap['float64'];
   export type NumberType = BooleanType|IntegerType|FloatType;
 
   export type Id = Guid;
@@ -93,6 +96,7 @@ export class Tensor {
    */
   get floatData() {
     switch (this.type) {
+      case 'float16':
       case 'float32':
       case 'float64':
         return this.data as Tensor.FloatType;
@@ -188,7 +192,7 @@ export class Tensor {
     } else {
       if (cache !== undefined) {
         const constructor = dataviewConstructor(type);
-        if (!(cache instanceof constructor)) {
+        if (!(cache instanceof constructor) && !(cache instanceof Float16Array)) {
           throw new TypeError(`cache should be type ${constructor.name}`);
         }
       }
@@ -357,6 +361,7 @@ function sizeof(type: Tensor.DataType): number {
       return 1;
     case 'int16':
     case 'uint16':
+    case 'float16':
       return 2;
     case 'int32':
     case 'uint32':
@@ -412,6 +417,8 @@ function dataviewConstructor(type: Tensor.DataType) {
       return Uint32Array;
     case 'int64':
       return BigInt64Array;
+    case 'float16':
+      return typeof Float16Array !== 'undefined' && Float16Array.from ? Float16Array : Uint16Array;
     case 'float32':
       return Float32Array;
     case 'float64':

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

@@ -488,8 +488,12 @@ 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.
+            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,21 @@ 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 hasFloat16Array = typeof Float16Array !== 'undefined';
+    if (hasFloat16Array) {
+      return elementCount === 0 ? new Float16Array() :
+                                  new Float16Array(this.module.HEAP8.buffer, this.data, elementCount);
+    } else {
+      return elementCount === 0 ? new Uint16Array() :
+                                  new Uint16Array(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

@@ -3,6 +3,13 @@
 
 import {Tensor} from 'onnxruntime-common';
 
+// a dummy type declaration for Float16Array in case any polyfill is available.
+declare global {
+  // eslint-disable-next-line no-var,@typescript-eslint/naming-convention,@typescript-eslint/no-explicit-any
+  var Float16Array: any;
+}
+
+export type Float16ArrayType = InstanceType<typeof Float16Array>;
 // This file includes common definitions. They do NOT have dependency on the WebAssembly instance.
 
 /**
@@ -117,7 +124,8 @@ export const tensorTypeToTypedArrayConstructor = (type: Tensor.Type): Float32Arr
     Uint8ArrayConstructor|Float64ArrayConstructor|Uint32ArrayConstructor|BigUint64ArrayConstructor => {
       switch (type) {
         case 'float16':
-          return Uint16Array;
+          // allow Float16Array polyfill.
+          return typeof Float16Array !== 'undefined' && Float16Array.from ? Float16Array : Uint16Array;
         case 'float32':
           return Float32Array;
         case 'uint8':

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": 1.2, "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.0],
+            "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.0, 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.0],
+            "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.0, 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/test-main.ts

@@ -13,6 +13,10 @@ import {Logger} from '../lib/onnxjs/instrument';
 
 import {Test} from './test-types';
 
+import {Float16Array} from '@petamoriken/float16';
+
+globalThis.Float16Array = Float16Array;
+
 if (ORT_WEB_TEST_CONFIG.model.some(testGroup => testGroup.tests.some(test => test.backend === 'cpu'))) {
   // require onnxruntime-node
   require('../../node');