[WebNN] Fallback int64 indices to int32 (microsoft#26308)

ONNX's ScatterND and ScatterElements limit their indices input to int64,
but some WebNN backends only support int32 indices. As a workaround for
such backends, we can insert a Cast operation to convert the data type.

Author: Honry

onnxruntime/core/providers/webnn/builders/helper.cc

@@ -439,5 +439,16 @@ uint16_t PackFloat32ToUint16AsFloat16(float value) {
   return sign_float16 | (exponent_float16 << 10) | mantissa_float16;
 }
 
+// Check if it can fallback to int32 if the input of WebNN op doesn't support int64.
+bool CanFallbackInt64ToInt32(const emscripten::val& wnn_limits,
+                             const std::string& webnn_op_type,
+                             const std::string& input_name) {
+  emscripten::val supported_data_types = wnn_limits[webnn_op_type][input_name]["dataTypes"];
+
+  return !supported_data_types.isUndefined() &&
+         !supported_data_types.call<emscripten::val>("includes", emscripten::val("int64")).as<bool>() &&
+         supported_data_types.call<emscripten::val>("includes", emscripten::val("int32")).as<bool>();
+}
+
 }  // namespace webnn
 }  // namespace onnxruntime

onnxruntime/core/providers/webnn/builders/helper.h

@@ -297,5 +297,9 @@ bool IsMLTensorSupported();
 uint8_t PackInt8ToUint8DoubledNibbles(int8_t value, const int32_t& data_type);
 uint16_t PackFloat32ToUint16AsFloat16(float value);
 
+bool CanFallbackInt64ToInt32(const emscripten::val& wnn_limits,
+                             const std::string& webnn_op_type,
+                             const std::string& input_name);
+
 }  // namespace webnn
 }  // namespace onnxruntime

onnxruntime/core/providers/webnn/builders/impl/scatterElements_op_builder.cc

@@ -24,6 +24,7 @@ class ScatterElementsOpBuilder : public BaseOpBuilder {
                          const WebnnDeviceType /* device_type */, const logging::Logger& logger) const override;
   bool HasSupportedInputsImpl(const GraphViewer& graph_viewer, const Node& node,
                               const emscripten::val& wnn_limits, const logging::Logger& logger) const override;
+  mutable bool can_fallback_int64_to_int32_ = false;
 };
 
 // Add operator related.
@@ -35,14 +36,21 @@ Status ScatterElementsOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_build
   emscripten::val indices = model_builder.GetOperand(input_defs[1]->Name());
   emscripten::val updates = model_builder.GetOperand(input_defs[2]->Name());
   emscripten::val options = emscripten::val::object();
-  options.set("label", node.Name());
+
+  // ONNX specifies that indices must use int64, but some WebNN backends only support int32.
+  // As a workaround for such backends, we can insert a Cast operation to convert the data type.
+  if (can_fallback_int64_to_int32_) {
+    options.set("label", node.Name() + "_cast_indices_to_int32");
+    indices = model_builder.GetBuilder().call<emscripten::val>("cast", indices, emscripten::val("int32"), options);
+  }
 
   std::vector<int64_t> input_shape;
   ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_shape, logger), "Cannot get shape");
   const size_t rank = input_shape.size();
   NodeAttrHelper helper(node);
   const uint32_t axis = static_cast<uint32_t>(HandleNegativeAxis(helper.Get("axis", 0), rank));
   options.set("axis", axis);
+  options.set("label", node.Name());
 
   emscripten::val output =
       model_builder.GetBuilder().call<emscripten::val>("scatterElements", data, indices, updates, options);
@@ -86,9 +94,16 @@ bool ScatterElementsOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const
 
   const std::string_view op_type = node.OpType();
 
-  return IsDataTypeSupportedByOp(op_type, data_type, wnn_limits, "input", "data", logger) &&
-         IsDataTypeSupportedByOp(op_type, indices_type, wnn_limits, "indices", "indices", logger) &&
-         IsInputRankSupportedByOp(node, wnn_limits, logger);
+  if (!IsDataTypeSupportedByOp(op_type, data_type, wnn_limits, "input", "data", logger) ||
+      !IsInputRankSupportedByOp(node, wnn_limits, logger)) {
+    return false;
+  }
+
+  // ONNX specifies that indices must use int64, but some WebNN backends only support int32.
+  // Allows to use int32 as a workaround for such backends.
+  can_fallback_int64_to_int32_ = CanFallbackInt64ToInt32(wnn_limits, "scatterElements", "indices");
+  return can_fallback_int64_to_int32_ ||
+         IsDataTypeSupportedByOp(op_type, indices_type, wnn_limits, "indices", "indices", logger);
 }
 
 void CreateScatterElementsOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) {

onnxruntime/core/providers/webnn/builders/impl/scatterND_op_builder.cc

@@ -24,6 +24,7 @@ class ScatterNDOpBuilder : public BaseOpBuilder {
                          const WebnnDeviceType /* device_type */, const logging::Logger& logger) const override;
   bool HasSupportedInputsImpl(const GraphViewer& graph_viewer, const Node& node,
                               const emscripten::val& wnn_limits, const logging::Logger& logger) const override;
+  mutable bool can_fallback_int64_to_int32_ = false;
 };
 
 // Add operator related.
@@ -35,6 +36,14 @@ Status ScatterNDOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, co
   emscripten::val indices = model_builder.GetOperand(input_defs[1]->Name());
   emscripten::val updates = model_builder.GetOperand(input_defs[2]->Name());
   emscripten::val options = emscripten::val::object();
+
+  // ONNX specifies that indices must use int64, but some WebNN backends only support int32.
+  // As a workaround for such backends, we can insert a Cast operation to convert the data type.
+  if (can_fallback_int64_to_int32_) {
+    options.set("label", node.Name() + "_cast_indices_to_int32");
+    indices = model_builder.GetBuilder().call<emscripten::val>("cast", indices, emscripten::val("int32"), options);
+  }
+
   options.set("label", node.Name());
   emscripten::val output =
       model_builder.GetBuilder().call<emscripten::val>("scatterND", data, indices, updates, options);
@@ -76,9 +85,17 @@ bool ScatterNDOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node&
     return false;
   }
   const std::string_view op_type = node.OpType();
-  return IsDataTypeSupportedByOp(op_type, data_type, wnn_limits, "input", "data", logger) &&
-         IsDataTypeSupportedByOp(op_type, indices_type, wnn_limits, "indices", "indices", logger) &&
-         IsInputRankSupportedByOp(node, wnn_limits, logger);
+
+  if (!IsDataTypeSupportedByOp(op_type, data_type, wnn_limits, "input", "data", logger) ||
+      !IsInputRankSupportedByOp(node, wnn_limits, logger)) {
+    return false;
+  }
+
+  // ONNX specifies that indices must use int64, but some WebNN backends only support int32.
+  // Allows to use int32 as a workaround for such backends.
+  can_fallback_int64_to_int32_ = CanFallbackInt64ToInt32(wnn_limits, "scatterND", "indices");
+  return can_fallback_int64_to_int32_ ||
+         IsDataTypeSupportedByOp(op_type, indices_type, wnn_limits, "indices", "indices", logger);
 }
 
 void CreateScatterNDOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) {