[js/webgpu] FP16 Cast, Resize (#18035)

### Description  Cast/Resize with f16 are missing in vae-decoder-f16. With this change, vae-decoder-f16 becomes 315 ms from over than 1 seconds.
microsoft · Oct 24, 2023 · eb47008 · eb47008
1 parent 688524a
commit eb47008
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 25 deletions.
diff --git a/onnxruntime/core/providers/js/operators/cast.cc b/onnxruntime/core/providers/js/operators/cast.cc
@@ -14,8 +14,7 @@ const std::vector<MLDataType>& CastOpTypeConstraints() {
   // https://gpuweb.github.io/gpuweb/wgsl/#plain-types-section
   //
   static std::vector<MLDataType> types{
-      // TODO(fs-eire): support f16 when it's ready
-      // DataTypeImpl::GetTensorType<MLFloat16>(),
+      DataTypeImpl::GetTensorType<MLFloat16>(),
       DataTypeImpl::GetTensorType<float>(),
       DataTypeImpl::GetTensorType<int32_t>(),
       DataTypeImpl::GetTensorType<uint32_t>(),

diff --git a/onnxruntime/core/providers/js/operators/resize.cc b/onnxruntime/core/providers/js/operators/resize.cc
@@ -5,15 +5,15 @@
 
 namespace onnxruntime {
 namespace js {
-#define REGISTER_RESIZE_VERSIONED_10_10_KERNEL(domain)                \
-  ONNX_OPERATOR_VERSIONED_KERNEL_EX(                                  \
-      Resize,                                                         \
-      domain,                                                         \
-      10, 10,                                                         \
-      kJsExecutionProvider,                                           \
-      (*KernelDefBuilder::Create())                                   \
-          .InputMemoryType(OrtMemTypeCPUInput, 1)                     \
-          .TypeConstraint("T", DataTypeImpl::GetTensorType<float>()), \
+#define REGISTER_RESIZE_VERSIONED_10_10_KERNEL(domain)     \
+  ONNX_OPERATOR_VERSIONED_KERNEL_EX(                       \
+      Resize,                                              \
+      domain,                                              \
+      10, 10,                                              \
+      kJsExecutionProvider,                                \
+      (*KernelDefBuilder::Create())                        \
+          .InputMemoryType(OrtMemTypeCPUInput, 1)          \
+          .TypeConstraint("T", JsepSupportedFloatTypes()), \
       Resize);
 
 #define REGISTER_RESIZE_VERSIONED_KERNEL(domain, sinceVersion, endVerion) \
@@ -26,22 +26,22 @@ namespace js {
           .InputMemoryType(OrtMemTypeCPUInput, 1)                         \
           .InputMemoryType(OrtMemTypeCPUInput, 2)                         \
           .InputMemoryType(OrtMemTypeCPUInput, 3)                         \
-          .TypeConstraint("T1", DataTypeImpl::GetTensorType<float>())     \
-          .TypeConstraint("T2", DataTypeImpl::GetTensorType<float>()),    \
+          .TypeConstraint("T1", JsepSupportedFloatTypes())                \
+          .TypeConstraint("T2", JsepSupportedFloatTypes()),               \
       Resize);
 
-#define REGISTER_RESIZE_KERNEL(domain, sinceVersion)                   \
-  ONNX_OPERATOR_KERNEL_EX(                                             \
-      Resize,                                                          \
-      domain,                                                          \
-      sinceVersion,                                                    \
-      kJsExecutionProvider,                                            \
-      (*KernelDefBuilder::Create())                                    \
-          .InputMemoryType(OrtMemTypeCPUInput, 1)                      \
-          .InputMemoryType(OrtMemTypeCPUInput, 2)                      \
-          .InputMemoryType(OrtMemTypeCPUInput, 3)                      \
-          .TypeConstraint("T1", DataTypeImpl::GetTensorType<float>())  \
-          .TypeConstraint("T2", DataTypeImpl::GetTensorType<float>()), \
+#define REGISTER_RESIZE_KERNEL(domain, sinceVersion)        \
+  ONNX_OPERATOR_KERNEL_EX(                                  \
+      Resize,                                               \
+      domain,                                               \
+      sinceVersion,                                         \
+      kJsExecutionProvider,                                 \
+      (*KernelDefBuilder::Create())                         \
+          .InputMemoryType(OrtMemTypeCPUInput, 1)           \
+          .InputMemoryType(OrtMemTypeCPUInput, 2)           \
+          .InputMemoryType(OrtMemTypeCPUInput, 3)           \
+          .TypeConstraint("T1", JsepSupportedFloatTypes())  \
+          .TypeConstraint("T2", JsepSupportedFloatTypes()), \
       Resize);
 
 #define REGISTER_RESIZE_KERNEL_DOMAIN(domain)       \