Update

Author: daijh

onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc

@@ -5,6 +5,7 @@
 #include <vector>
 
 #include "core/providers/webgpu/webgpu_utils.h"
+#include "core/providers/webgpu/nn/im2col_matmul.h"
 #include "core/providers/webgpu/nn/conv.h"
 #include "core/providers/webgpu/nn/activation_util.h"
 
@@ -52,15 +53,6 @@ bool IsDeviceSupported(const ComputeContextBase& context) {
 
 }  // namespace
 
-Status OIHW2OHWIProgram::GenerateShaderCode(ShaderHelper& shader) const {
-  const auto& src = shader.AddInput("src", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
-  const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
-
-  return WGSL_TEMPLATE_APPLY(shader, "nn/oihw_to_ohwi.wgsl.template",
-                             WGSL_TEMPLATE_VARIABLE(output, output),
-                             WGSL_TEMPLATE_VARIABLE(src, src));
-}
-
 Status Im2ColMatMulProgram::GenerateShaderCode(ShaderHelper& shader) const {
   const auto& src = shader.AddInput("src", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
   const auto& weight = shader.AddInput("weight", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);

onnxruntime/core/providers/webgpu/nn/im2col_matmul.h

@@ -18,22 +18,6 @@
 namespace onnxruntime {
 namespace webgpu {
 
-// Transpose OIHW Weight to OHWI
-class OIHW2OHWIProgram final : public Program<OIHW2OHWIProgram> {
- public:
-  OIHW2OHWIProgram() : Program("OIHW2OHWI") {}
-
-  Status GenerateShaderCode(ShaderHelper& shader) const override;
-
-  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
-      {"O", ProgramUniformVariableDataType::Uint32},
-      {"I", ProgramUniformVariableDataType::Uint32},
-      {"H", ProgramUniformVariableDataType::Uint32},
-      {"W", ProgramUniformVariableDataType::Uint32},
-      {"Ci_tiles", ProgramUniformVariableDataType::Uint32},
-      {"H_W_tiles", ProgramUniformVariableDataType::Uint32});
-};
-
 class Im2ColMatMulProgram final : public Program<Im2ColMatMulProgram> {
  public:
   Im2ColMatMulProgram(bool has_bias,

…ers/webgpu/nn/oihw_to_ohwi.wgsl.template …webgpu/tensor/oihw_to_ohwi.wgsl.templateonnxruntime/core/providers/webgpu/nn/oihw_to_ohwi.wgsl.template renamed to onnxruntime/core/providers/webgpu/tensor/oihw_to_ohwi.wgsl.template onnxruntime/core/providers/webgpu/nn/oihw_to_ohwi.wgsl.template renamed to onnxruntime/core/providers/webgpu/tensor/oihw_to_ohwi.wgsl.template

@@ -48,6 +48,15 @@ ONNX_OPERATOR_KERNEL_EX(
         .TypeConstraint("T", WebGpuSupportedNumberTypes()),
     Transpose);
 
+Status OIHW2OHWIProgram::GenerateShaderCode(ShaderHelper& shader) const {
+  const auto& src = shader.AddInput("src", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+  const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+
+  return WGSL_TEMPLATE_APPLY(shader, "tensor/oihw_to_ohwi.wgsl.template",
+                             WGSL_TEMPLATE_VARIABLE(output, output),
+                             WGSL_TEMPLATE_VARIABLE(src, src));
+}
+
 auto SqueezeShape(const gsl::span<const int64_t>& shape,
                   const gsl::span<const size_t>& adjusted_perm,
                   TensorShapeVector& new_shape,

onnxruntime/core/providers/webgpu/tensor/transpose.cc

@@ -11,6 +11,22 @@
 namespace onnxruntime {
 namespace webgpu {
 
+// Transpose OIHW Weight to OHWI
+class OIHW2OHWIProgram final : public Program<OIHW2OHWIProgram> {
+ public:
+  OIHW2OHWIProgram() : Program("OIHW2OHWI") {}
+
+  Status GenerateShaderCode(ShaderHelper& shader) const override;
+
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
+      {"O", ProgramUniformVariableDataType::Uint32},
+      {"I", ProgramUniformVariableDataType::Uint32},
+      {"H", ProgramUniformVariableDataType::Uint32},
+      {"W", ProgramUniformVariableDataType::Uint32},
+      {"Ci_tiles", ProgramUniformVariableDataType::Uint32},
+      {"H_W_tiles", ProgramUniformVariableDataType::Uint32});
+};
+
 class Transpose final : public WebGpuKernel, public TransposeBase {
  public:
   Transpose(const OpKernelInfo& info) : WebGpuKernel{info}, TransposeBase{info} {