handle fp16 for where op (microsoft#19969)

this prevents falling back from webgpu to cpu, aka helps performance
TedThemistokleous · May 7, 2024 · 9f589d8 · 9f589d8
1 parent e247fbf
commit 9f589d8
Showing 1 changed file with 14 additions and 12 deletions.
diff --git a/onnxruntime/core/providers/js/operators/where.cc b/onnxruntime/core/providers/js/operators/where.cc
@@ -6,18 +6,19 @@
 namespace onnxruntime {
 namespace js {
 
-#define REG_ELEMENTWISE_KERNEL(OP_TYPE, VERSION, KERNEL_CLASS)      \
-  ONNX_OPERATOR_KERNEL_EX(                                          \
-      OP_TYPE,                                                      \
-      kOnnxDomain,                                                  \
-      VERSION,                                                      \
-      kJsExecutionProvider,                                         \
-      KernelDefBuilder()                                            \
-          .TypeConstraint("T",                                      \
-                          {DataTypeImpl::GetTensorType<float>(),    \
-                           DataTypeImpl::GetTensorType<int32_t>(),  \
-                           DataTypeImpl::GetTensorType<uint32_t>(), \
-                           DataTypeImpl::GetTensorType<bool>()}),   \
+#define REG_ELEMENTWISE_KERNEL(OP_TYPE, VERSION, KERNEL_CLASS)       \
+  ONNX_OPERATOR_KERNEL_EX(                                           \
+      OP_TYPE,                                                       \
+      kOnnxDomain,                                                   \
+      VERSION,                                                       \
+      kJsExecutionProvider,                                          \
+      KernelDefBuilder()                                             \
+          .TypeConstraint("T",                                       \
+                          {DataTypeImpl::GetTensorType<float>(),     \
+                           DataTypeImpl::GetTensorType<MLFloat16>(), \
+                           DataTypeImpl::GetTensorType<int32_t>(),   \
+                           DataTypeImpl::GetTensorType<uint32_t>(),  \
+                           DataTypeImpl::GetTensorType<bool>()}),    \
       KERNEL_CLASS);
 
 #define REG_ELEMENTWISE_VERSIONED_KERNEL(OP_TYPE, VERSION_FROM, VERSION_TO, KERNEL_CLASS) \
@@ -29,6 +30,7 @@ namespace js {
       KernelDefBuilder()                                                                  \
           .TypeConstraint("T",                                                            \
                           {DataTypeImpl::GetTensorType<float>(),                          \
+                           DataTypeImpl::GetTensorType<MLFloat16>(),                      \
                            DataTypeImpl::GetTensorType<int32_t>(),                        \
                            DataTypeImpl::GetTensorType<uint32_t>(),                       \
                            DataTypeImpl::GetTensorType<bool>()}),                         \