[webgpu] Add Alias def for Flatten (microsoft#23038)

### Description Add `Alias` definition for Flatten in WebGPU EP. also add int32/uint32 in type constraint T.
intel · Dec 11, 2024 · 8dd8bc5 · 8dd8bc5
1 parent 6cd9a06
commit 8dd8bc5
Showing 1 changed file with 21 additions and 6 deletions.
diff --git a/onnxruntime/core/providers/webgpu/tensor/flatten.cc b/onnxruntime/core/providers/webgpu/tensor/flatten.cc
@@ -13,40 +13,55 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     kOnnxDomain,
     1, 8,
     kWebGpuExecutionProvider,
-    (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedFloatTypes()).InputMemoryType(OrtMemTypeCPU, 1),
+    (*KernelDefBuilder::Create())
+        .Alias(0, 0)
+        .TypeConstraint("T", WebGpuSupportedNumberTypes())
+        .InputMemoryType(OrtMemTypeCPU, 1),
     Flatten);
 
 ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     Flatten,
     kOnnxDomain,
     9, 10,
     kWebGpuExecutionProvider,
-    (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedFloatTypes()).InputMemoryType(OrtMemTypeCPU, 1),
+    (*KernelDefBuilder::Create())
+        .Alias(0, 0)
+        .TypeConstraint("T", WebGpuSupportedNumberTypes())
+        .InputMemoryType(OrtMemTypeCPU, 1),
     Flatten);
 
 ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     Flatten,
     kOnnxDomain,
     11, 12,
     kWebGpuExecutionProvider,
-    (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedFloatTypes()).InputMemoryType(OrtMemTypeCPU, 1),
+    (*KernelDefBuilder::Create())
+        .Alias(0, 0)
+        .TypeConstraint("T", WebGpuSupportedNumberTypes())
+        .InputMemoryType(OrtMemTypeCPU, 1),
     Flatten);
 
 ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     Flatten,
     kOnnxDomain,
     13, 20,
     kWebGpuExecutionProvider,
-    (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedFloatTypes()).InputMemoryType(OrtMemTypeCPU, 1),
+    (*KernelDefBuilder::Create())
+        .Alias(0, 0)
+        .TypeConstraint("T", WebGpuSupportedNumberTypes())
+        .InputMemoryType(OrtMemTypeCPU, 1),
     Flatten);
 
 ONNX_OPERATOR_KERNEL_EX(
     Flatten,
     kOnnxDomain,
     21,
     kWebGpuExecutionProvider,
-    (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedFloatTypes()).InputMemoryType(OrtMemTypeCPU, 1),
+    (*KernelDefBuilder::Create())
+        .Alias(0, 0)
+        .TypeConstraint("T", WebGpuSupportedNumberTypes())
+        .InputMemoryType(OrtMemTypeCPU, 1),
     Flatten);
 
 }  // namespace webgpu
-}  // namespace onnxruntime
+}  // namespace onnxruntime