[WebNN EP] WebNN CPU backend only support up to 4 Split outputs (micr…

…osoft#20350)
AmusementClub · Apr 19, 2024 · 7c80c39 · 7c80c39
1 parent 4d1963c
commit 7c80c39
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/onnxruntime/core/providers/webnn/builders/impl/split_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/split_op_builder.cc
@@ -27,7 +27,7 @@ class SplitOpBuilder : public BaseOpBuilder {
   // Operator support related.
  private:
   bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node,
-                         const WebnnDeviceType /* device_type */, const logging::Logger& logger) const override;
+                         const WebnnDeviceType device_type, const logging::Logger& logger) const override;
 };
 
 // Add operator related.
@@ -94,7 +94,7 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
 
 bool SplitOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers,
                                        const Node& node,
-                                       const WebnnDeviceType /* device_type */,
+                                       const WebnnDeviceType device_type,
                                        const logging::Logger& logger) const {
   const auto& input_defs = node.InputDefs();
   std::vector<int64_t> input_shape;
@@ -126,6 +126,10 @@ bool SplitOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers,
       LOGS(logger, VERBOSE) << "Cannot get split.";
       return false;
     }
+    if (split.size() > 4 && device_type == WebnnDeviceType::CPU) {
+      LOGS(logger, VERBOSE) << "WebNN CPU backend only supports up to 4 outputs.";
+      return false;
+    }
   } else {
     if (helper.HasAttr("num_outputs")) {
       // Split has 'num_outputs' attribute when opset is 18.
@@ -134,6 +138,10 @@ bool SplitOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers,
         LOGS(logger, VERBOSE) << "The 'num_outputs' must be a positive integer.";
         return false;
       }
+      if (num_outputs > 4 && device_type == WebnnDeviceType::CPU) {
+        LOGS(logger, VERBOSE) << "WebNN CPU backend only supports up to 4 outputs.";
+        return false;
+      }
     } else {
       const auto opset = node.SinceVersion();
       if (opset >= 18) {