From b9f9290765efa4b08b6f4e7e4b6304d83304a8b9 Mon Sep 17 00:00:00 2001
From: Phillis Tang <[email protected]>
Date: Wed, 13 May 2026 15:07:18 -0700
Subject: [PATCH] webnn: limit concat and split inputs/outputs

Validate that split inputs and concat outputs can't exceed 8192.

Bug: 505137978, 505136542
Change-Id: I0caad753c6ef2685e3ccbfe3dc70821c78cb7d8c
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7842056
Commit-Queue: Phillis Tang <[email protected]>
Reviewed-by: Reilly Grant <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1630257}
---

diff --git a/services/webnn/public/cpp/graph_validation_utils.cc b/services/webnn/public/cpp/graph_validation_utils.cc
index 2b312c4c..6c023be0 100644
--- a/services/webnn/public/cpp/graph_validation_utils.cc
+++ b/services/webnn/public/cpp/graph_validation_utils.cc
@@ -32,6 +32,9 @@
 
 namespace {
 
+// Used by concat/split to validate operand count limit.
+constexpr uint32_t kMaxValidTensorCount = 8192;
+
 struct Conv2dInputOutputInfo {
   uint32_t batches;
   uint32_t channels;
@@ -583,6 +586,13 @@
         ErrorWithLabel(label, "The inputs should not be empty."));
   }
 
+  if (inputs.size() > kMaxValidTensorCount) {
+    return base::unexpected(ErrorWithLabel(
+        label, base::StringPrintf(
+                   "The number of inputs must be less than or equal to %u.",
+                   kMaxValidTensorCount)));
+  }
+
   for (const auto& input : inputs) {
     if (!context_properties.data_type_limits.concat_inputs.Supports(input)) {
       return base::unexpected(ErrorWithLabel(
@@ -2864,6 +2874,13 @@
           ErrorWithLabel(label, "The splits must be greater than zero."));
     }
 
+    if (splits > kMaxValidTensorCount) {
+      return base::unexpected(ErrorWithLabel(
+          label,
+          base::StringPrintf("The splits must be less than or equal to %u.",
+                             kMaxValidTensorCount)));
+    }
+
     if (input.shape()[attributes.axis] % splits != 0) {
       return base::unexpected(
           ErrorWithLabel(label,
@@ -2888,6 +2905,18 @@
                  attributes.splits)) {
     const auto& splits =
         std::get<base::span<const uint32_t>>(attributes.splits);
+    if (splits.empty()) {
+      return base::unexpected(
+          ErrorWithLabel(label, "The splits should not be empty."));
+    }
+
+    if (splits.size() > kMaxValidTensorCount) {
+      return base::unexpected(ErrorWithLabel(
+          label, base::StringPrintf(
+                     "The number of splits must be less than or equal to %u.",
+                     kMaxValidTensorCount)));
+    }
+
     if (std::ranges::any_of(splits,
                             [](uint32_t split) { return split == 0; })) {
       return base::unexpected(
diff --git a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js
index a9286d4e..099dd8c9 100644
--- a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js
+++ b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js
@@ -147,3 +147,13 @@
 
   assert_throws_js(TypeError, () => builder.concat(input1, input2));
 }, '[concat] throw if the output number of elements is too large');
+
+promise_test(async t => {
+  const builder = new MLGraphBuilder(context);
+  const operandDescriptor = {dataType: 'float32', shape: [1]};
+  const inputs = [];
+  for (let i = 0; i < 8193; ++i) {
+    inputs.push(builder.input(`input${i}`, operandDescriptor));
+  }
+  assert_throws_js(TypeError, () => builder.concat(inputs));
+}, '[concat] throw if the number of inputs exceeds limit');
diff --git a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js
index 21b790bc..609b3b3 100644
--- a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js
+++ b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js
@@ -101,6 +101,23 @@
       label: label,
     }
   },
+  {
+    name: '[split] Throw if splits (scalar) exceeds limit.',
+    input: {dataType: 'float32', shape: [2, 6]},
+    splits: 8193,
+    options: {
+      label: label,
+    },
+  },
+  {
+    name: '[split] Throw if number of splits exceeds limit.',
+    input: {dataType: 'float32', shape: [2, 8193]},
+    splits: Array(8193).fill(1),
+    options: {
+      axis: 1,
+      label: label,
+    },
+  },
 ];
 
 tests.forEach(

diff --git a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js
index a9286d4e..099dd8c9 100644
--- a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js
+++ b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/concat.https.any.js
@@ -147,3 +147,13 @@
 
   assert_throws_js(TypeError, () => builder.concat(input1, input2));
 }, '[concat] throw if the output number of elements is too large');
+
+promise_test(async t => {
+  const builder = new MLGraphBuilder(context);
+  const operandDescriptor = {dataType: 'float32', shape: [1]};
+  const inputs = [];
+  for (let i = 0; i < 8193; ++i) {
+    inputs.push(builder.input(`input${i}`, operandDescriptor));
+  }
+  assert_throws_js(TypeError, () => builder.concat(inputs));
+}, '[concat] throw if the number of inputs exceeds limit');
diff --git a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js
index 21b790bc..609b3b3 100644
--- a/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js
+++ b/third_party/blink/web_tests/external/wpt/webnn/validation_tests/split.https.any.js
@@ -101,6 +101,23 @@
       label: label,
     }
   },
+  {
+    name: '[split] Throw if splits (scalar) exceeds limit.',
+    input: {dataType: 'float32', shape: [2, 6]},
+    splits: 8193,
+    options: {
+      label: label,
+    },
+  },
+  {
+    name: '[split] Throw if number of splits exceeds limit.',
+    input: {dataType: 'float32', shape: [2, 8193]},
+    splits: Array(8193).fill(1),
+    options: {
+      axis: 1,
+      label: label,
+    },
+  },
 ];
 
 tests.forEach(