webnn: Limit the tensor size with the opSupportLimits

The tensor byte length is UINT32_MAX for DirectML backend on Windows,
INT32_MAX for CoreML and TFLite backend, `OperandDescriptor::Create`
function will be used to validate the limit.

This limitation also prevents the mask tensor of triangular to allocate
large amounts of memory.
The `OperandDescriptor::CreateForDeserialization` only be used by mojom
traits that need to be validated tensor size limit later.
Add some unittests and WPT tests to validate the byte length limit.

Bug: 359729258
Change-Id: Ic174fe47da3c110a6d263d22318738f5656a07cc
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6079317
Commit-Queue: Junwei Fu <[email protected]>
Reviewed-by: Reilly Grant <[email protected]>
Reviewed-by: ningxin hu <[email protected]>
Reviewed-by: Alex Gough <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1407058}

Author: fujunwei

webnn/conformance_tests/tensor.https.any.js

@@ -128,6 +128,17 @@ const testCreateTensorFails = (testName, tensorDescriptor) => {
   }, `${testName} / ${tensorDescriptor.dataType}`);
 };
 
+
+promise_test(async t => {
+  const tensorDescriptor = {
+    dataType: 'int32',
+    shape: [(context.opSupportLimits().maxTensorByteLength + 1) / 4],
+    writable: true,
+  };
+  await promise_rejects_js(
+    t, TypeError, context.createTensor(tensorDescriptor));
+}, `create too large tensor byte length that exceeds limit`);
+
 /**
  * Asserts the tensor data in MLTensor matches expected.
  * @param {MLContext} mlContext - The context used to create the tensor.

webnn/resources/utils_validation.js

@@ -248,6 +248,24 @@ function validateTwoInputsBroadcastable(operationName, label) {
   }, `[${operationName}] TypeError is expected if two inputs aren't broadcastable`);
 }
 
+function validateTwoBroadcastableInputsTensorLimit(operationName, label) {
+  if (navigator.ml === undefined) {
+    return;
+  }
+  promise_test(async t => {
+    const builder = new MLGraphBuilder(context);
+
+    const a = builder.input('a', {dataType: 'float32',
+        shape: [context.opSupportLimits().maxTensorByteLength / 4, 1]});
+    const b = builder.input('b', {dataType: 'float32', shape: [1, 5] });
+
+    const options = {label};
+    const regrexp = new RegExp('\\[' + label + '\\]');
+    assert_throws_with_label(
+        () => builder[operationName](a, b, options), regrexp);
+  }, `[${operationName}] throw if the output tensor byte length exceeds limit`);
+}
+
 function validateTwoInputsOfSameDataType(operationName, label) {
   if (navigator.ml === undefined) {
     return;

webnn/validation_tests/cast.https.any.js

@@ -14,3 +14,12 @@ multi_builder_test(async (t, builder, otherBuilder) => {
   assert_throws_js(
       TypeError, () => builder.cast(inputFromOtherBuilder, 'int64'));
 }, '[cast] throw if input is from another builder');
+
+promise_test(async t => {
+    const builder = new MLGraphBuilder(context);
+    const input = builder.input('input', {
+        dataType: 'int8',
+        shape: [context.opSupportLimits().maxTensorByteLength / 2]});
+    assert_throws_js(
+        TypeError, () => builder.cast(input, 'int64'));
+  }, '[cast] throw if the output tensor byte length exceeds limit');

webnn/validation_tests/concat.https.any.js

@@ -83,7 +83,6 @@ const tests = [
     ],
     axis: 1,
   },
-
 ];
 
 tests.forEach(
@@ -120,3 +119,18 @@ multi_builder_test(async (t, builder, otherBuilder) => {
       TypeError,
       () => builder.concat([input1, input2, inputFromOtherBuilder, input3]));
 }, '[concat] throw if any input is from another builder');
+
+promise_test(async t => {
+  const builder = new MLGraphBuilder(context);
+
+  const operandDescriptor = {
+    dataType: 'float32',
+    shape: [context.opSupportLimits().maxTensorByteLength / 4]
+  };
+  const input1 = builder.input('input1', operandDescriptor);
+  const input2 = builder.input('input2', operandDescriptor);
+  const input3 = builder.input('input3', operandDescriptor);
+
+  assert_throws_js(
+      TypeError, () => builder.concat(input1, input2, input3));
+}, '[concat] throw if the output tensor byte length exceeds limit');

webnn/validation_tests/dequantizeLinear.https.any.js

@@ -7,6 +7,8 @@
 
 'use strict';
 
+const label = 'dequantize_linear_123';
+const regrexp = new RegExp('\\[' + label + '\\]');
 const tests = [
   {
     name:
@@ -94,9 +96,7 @@ tests.forEach(
         assert_equals(output.dataType, test.output.dataType);
         assert_array_equals(output.shape, test.output.shape);
       } else {
-        const label = 'dequantize_linear_123';
         const options = {label};
-        const regrexp = new RegExp('\\[' + label + '\\]');
         assert_throws_with_label(
             () => builder.dequantizeLinear(input, scale, zeroPoint, options),
             regrexp);
@@ -143,3 +143,17 @@ multi_builder_test(async (t, builder, otherBuilder) => {
       TypeError,
       () => builder.dequantizeLinear(input, scale, zeroPointFromOtherBuilder));
 }, '[dequantizeLinear] throw if zeroPoint is from another builder');
+
+promise_test(async t => {
+  const builder = new MLGraphBuilder(context);
+
+  const input = builder.input('input', {
+      dataType: 'int8',
+      shape: [context.opSupportLimits().maxTensorByteLength / 5, 5]});
+  const scale = builder.input('scale', {dataType: 'float32', shape: [5]});
+  const zeroPoint = builder.input('zeroPoint', {dataType: 'int8', shape: [5]});
+
+  const options = {label};
+  assert_throws_with_label(
+      () => builder.dequantizeLinear(input, scale, zeroPoint, options), regrexp);
+}, '[dequantizeLinear] throw if the output tensor byte length exceeds limit');

webnn/validation_tests/elementwise-binary.https.any.js

@@ -88,5 +88,6 @@ kElementwiseBinaryOperators.forEach((operatorName) => {
   validateTwoInputsOfSameDataType(operatorName, label);
   validateTwoInputsBroadcastable(operatorName, label);
   validateTwoInputsFromMultipleBuilders(operatorName);
+  validateTwoBroadcastableInputsTensorLimit(operatorName, label);
   runElementWiseBinaryTests(operatorName, tests);
 });

webnn/validation_tests/expand.https.any.js

@@ -17,7 +17,7 @@ multi_builder_test(async (t, builder, otherBuilder) => {
 }, '[expand] throw if input is from another builder');
 
 const label = 'xxx_expand';
-
+const regrexp = new RegExp('\\[' + label + '\\]');
 const tests = [
   {
     name: '[expand] Test with 0-D scalar to 3-D tensor.',
@@ -76,7 +76,6 @@ tests.forEach(
       } else {
         const options = {...test.options};
         if (options.label) {
-          const regrexp = new RegExp('\\[' + label + '\\]');
           assert_throws_with_label(
               () => builder.expand(input, test.newShape, options), regrexp);
         } else {
@@ -104,3 +103,15 @@ promise_test(async t => {
     }
   }
 }, `[expand] Test expand with all of the data types.`);
+
+promise_test(async t => {
+  const builder = new MLGraphBuilder(context);
+
+  const input = builder.input('input', {
+      dataType: 'float32', shape: [1, 2, 1, 1]});
+  const newShape = [1, 2, context.opSupportLimits().maxTensorByteLength, 1];
+
+  const options = {label};
+  assert_throws_with_label(
+      () => builder.expand(input, newShape, options), regrexp);
+}, '[expand] throw if the output tensor byte length exceeds limit');

webnn/validation_tests/gather.https.any.js

@@ -7,6 +7,8 @@
 
 'use strict';
 
+const label = 'gather_'
+const regrexp = new RegExp('\\[' + label + '\\]');
 const tests = [
   {
     name: '[gather] Test gather with default options and 0-D indices',
@@ -58,7 +60,7 @@ const tests = [
         '[gather] TypeError is expected if the data type of indices is uint64 which is invalid',
     input: {dataType: 'float16', shape: [1, 2, 3, 4]},
     indices: {dataType: 'uint64', shape: [5, 6]},
-  }
+  },
 ];
 
 tests.forEach(
@@ -77,9 +79,7 @@ tests.forEach(
         assert_equals(output.dataType, test.output.dataType);
         assert_array_equals(output.shape, test.output.shape);
       } else {
-        const label = 'gather_'
         options.label = label;
-        const regrexp = new RegExp('\\[' + label + '\\]');
         assert_throws_with_label(
             () => builder.gather(input, indices, options), regrexp);
       }
@@ -102,3 +102,19 @@ multi_builder_test(async (t, builder, otherBuilder) => {
   assert_throws_js(
       TypeError, () => builder.gather(input, indicesFromOtherBuilder));
 }, '[gather] throw if indices is from another builder');
+
+promise_test(async t => {
+  const builder = new MLGraphBuilder(context);
+
+  const input = builder.input('input', {
+      dataType: 'float32', shape: [1, 3, 3, 4]});
+  const indices = builder.input('indices', {
+      dataType: 'int32',
+      shape: [context.opSupportLimits().maxTensorByteLength / 4] });
+
+  const options = {};
+  options.label = label;
+  options.axis = 2;
+  assert_throws_with_label(
+      () => builder.gather(input, indices, options), regrexp);
+}, '[gather] throw if the output tensor byte length exceeds limit');

webnn/validation_tests/gatherND.https.any.js

@@ -7,6 +7,8 @@
 
 'use strict';
 
+const label = 'gatherND_';
+const regexp = new RegExp('\\[' + label + '\\]');
 const tests = [
   {
     name: '[gatherND] Test gatherND with 5D input 3D indices',
@@ -49,9 +51,7 @@ tests.forEach(test => promise_test(async t => {
                   assert_equals(output.dataType, test.output.dataType);
                   assert_array_equals(output.shape, test.output.shape);
                 } else {
-                  const label = 'gatherND_';
                   const options = {label: label};
-                  const regexp = new RegExp('\\[' + label + '\\]');
                   assert_throws_with_label(
                       () => builder.gatherND(input, indices, options), regexp);
                 }
@@ -74,3 +74,17 @@ multi_builder_test(async (t, builder, otherBuilder) => {
   assert_throws_js(
       TypeError, () => builder.gatherND(input, indicesFromOtherBuilder));
 }, '[gatherND] Throw if indices is from another builder');
+
+promise_test(async t => {
+  const builder = new MLGraphBuilder(context);
+
+  const input = builder.input('input', {
+      dataType: 'float32', shape: [2, 2, 3, 3, 4]});
+  const indices = builder.input('indices', {
+    dataType: 'int32',
+    shape: [context.opSupportLimits().maxTensorByteLength / 4, 1, 1]});
+
+  const options = {label};
+  assert_throws_with_label(
+      () => builder.gatherND(input, indices, options), regexp);
+}, '[gatherND] throw if the output tensor byte length exceeds limit');

webnn/validation_tests/gemm.https.any.js

@@ -7,12 +7,14 @@
 
 'use strict';
 
+const label = 'gemm_xxx';
 const kExampleInputDescriptor = {
   dataType: 'float32',
   shape: [2, 2]
 };
 
 validateTwoInputsFromMultipleBuilders('gemm');
+validateTwoBroadcastableInputsTensorLimit('gemm', label);
 
 multi_builder_test(async (t, builder, otherBuilder) => {
   const cFromOtherBuilder = otherBuilder.input('c', kExampleInputDescriptor);
@@ -23,8 +25,6 @@ multi_builder_test(async (t, builder, otherBuilder) => {
   assert_throws_js(TypeError, () => builder.gemm(a, b, options));
 }, '[gemm] throw if c option is from another builder');
 
-const label = 'gemm_xxx';
-
 const tests = [
   {
     name: '[gemm] Test building gemm with default option.',