[WebGPU] Webgpu im2col fused activation

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

@@ -164,12 +164,13 @@ Status Conv<is_channels_last, is_fused>::ComputeInternal(ComputeContext& context
 
   if (CanApplyIm2ColMatMulProgram(context,
                                   is_channels_last,
-                                  activation_.activation_kind_ != ActivationKind::None,
+                                  activation_,
                                   kernel_shape,
                                   onnxruntime::narrow<uint32_t>(conv_attrs_.group),
                                   kernel->DataType())) {
     return ApplyIm2ColMatMulProgram(context,
                                     is_channels_last,
+                                    activation_,
                                     dilations,
                                     pads,
                                     strides,
@@ -340,7 +341,7 @@ Status Conv<is_channels_last, is_fused>::PrePackInternal(ComputeContextBase& con
   // Im2ColMatMul path uses a different transpose (OIHW -> OHWI) and reads
   // kernel directly from context.Input(1), ignoring prepacked weights.
   // Skip prepacking when this path will be used at runtime.
-  if (CanApplyIm2ColMatMulProgram(context, is_channels_last, activation_.activation_kind_ != ActivationKind::None,
+  if (CanApplyIm2ColMatMulProgram(context, is_channels_last, activation_,
                                   kernel_shape, onnxruntime::narrow<uint32_t>(conv_attrs_.group),
                                   tensor.DataType())) {
     return Status::OK();

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

@@ -14,6 +14,7 @@ class OpKernelInfo;
 
 namespace webgpu {
 
+// Values are mirrored by im2col_matmul.wgsl.template; append without reordering.
 enum class ActivationKind {
   None,
   Relu,

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

@@ -54,8 +54,33 @@ bool IsDeviceSupported(const ComputeContextBase& context) {
   return false;
 }
 
+// Keep this list synchronized with the activation_kind branches in the WGSL template.
+bool IsActivationSupported(const Activation& activation) {
+  switch (activation.activation_kind_) {
+    case ActivationKind::None:
+    case ActivationKind::Relu:
+    case ActivationKind::Sigmoid:
+    case ActivationKind::Clip:
+    case ActivationKind::HardSigmoid:
+    case ActivationKind::LeakyRelu:
+    case ActivationKind::Tanh:
+      return true;
+    default:
+      return false;
+  }
+}
+
 }  // namespace
 
+// The template dispatches on the numeric enum values.
+static_assert(static_cast<int>(ActivationKind::None) == 0, "im2col_matmul.wgsl.template mirrors ActivationKind");
+static_assert(static_cast<int>(ActivationKind::Relu) == 1, "im2col_matmul.wgsl.template mirrors ActivationKind");
+static_assert(static_cast<int>(ActivationKind::Sigmoid) == 2, "im2col_matmul.wgsl.template mirrors ActivationKind");
+static_assert(static_cast<int>(ActivationKind::Clip) == 3, "im2col_matmul.wgsl.template mirrors ActivationKind");
+static_assert(static_cast<int>(ActivationKind::HardSigmoid) == 4, "im2col_matmul.wgsl.template mirrors ActivationKind");
+static_assert(static_cast<int>(ActivationKind::LeakyRelu) == 5, "im2col_matmul.wgsl.template mirrors ActivationKind");
+static_assert(static_cast<int>(ActivationKind::Tanh) == 6, "im2col_matmul.wgsl.template mirrors ActivationKind");
+
 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);
@@ -69,6 +94,7 @@ Status Im2ColMatMulProgram::GenerateShaderCode(ShaderHelper& shader) const {
   ORT_ENFORCE(vec_size_ == 1 || vec_size_ == 2 || vec_size_ == 4, "vec_size must be 1, 2 or 4.");
 
   return WGSL_TEMPLATE_APPLY(shader, "nn/im2col_matmul.wgsl.template",
+                             WGSL_TEMPLATE_PARAMETER(activation_kind, static_cast<uint32_t>(activation_kind_)),
                              WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_),
                              WGSL_TEMPLATE_PARAMETER(tile_m, tile_m_),
                              WGSL_TEMPLATE_PARAMETER(tile_n, tile_n_),
@@ -81,6 +107,7 @@ Status Im2ColMatMulProgram::GenerateShaderCode(ShaderHelper& shader) const {
 
 Status ApplyIm2ColMatMulProgram(ComputeContext& context,
                                 bool is_channels_last,
+                                const Activation& activation,
                                 const std::vector<uint32_t>& dilations,
                                 const std::vector<uint32_t>& pads,
                                 const std::vector<uint32_t>& strides,
@@ -125,7 +152,8 @@ Status ApplyIm2ColMatMulProgram(ComputeContext& context,
   // If the status of this condition is uncertain, the feature must be disabled.
   const bool use_subgroup = false;
   const uint32_t vec_size = channel_input % 4 == 0 ? 4 : (channel_input % 2 == 0 ? 2 : 1);
-  Im2ColMatMulProgram im2col_mm_program{has_bias, tile_m, tile_n, vec_size, use_subgroup};
+  Im2ColMatMulProgram im2col_mm_program{has_bias, tile_m, tile_n, vec_size, use_subgroup,
+                                        activation.activation_kind_};
   im2col_mm_program.SetWorkgroupSize(workgroup_size);
 
   const uint32_t M_tiles = CeilDiv(im2col_m, tile_m);
@@ -161,14 +189,15 @@ Status ApplyIm2ColMatMulProgram(ComputeContext& context,
                                          {dilations},
                                          {pads},
                                          {strides}});
-  im2col_mm_program.CacheHint(has_bias, tile_m, tile_n, vec_size, use_subgroup);
+  AppendActivationUniformsData(activation, im2col_mm_program);
+  im2col_mm_program.CacheHint(has_bias, tile_m, tile_n, vec_size, use_subgroup, activation.CacheKey());
 
   return context.RunProgram(im2col_mm_program);
 }
 
 bool CanApplyIm2ColMatMulProgram(ComputeContextBase& context,
                                  const bool is_channels_last,
-                                 const bool is_fused,
+                                 const Activation& activation,
                                  const TensorShape weight_shape,
                                  const uint32_t group,
                                  const MLDataType data_type) {
@@ -183,9 +212,12 @@ bool CanApplyIm2ColMatMulProgram(ComputeContextBase& context,
   }
 
   // TODO: Support !is_channels_last
-  // TODO: Support fuse
   // TODO: Support group conv
-  if (!is_channels_last || is_fused || group != 1) {
+  if (!is_channels_last || group != 1) {
+    return false;
+  }
+
+  if (!IsActivationSupported(activation)) {
     return false;
   }
 

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

@@ -24,12 +24,14 @@ class Im2ColMatMulProgram final : public Program<Im2ColMatMulProgram> {
                       uint32_t tile_m,
                       uint32_t tile_n,
                       uint32_t vec_size,
-                      bool use_subgroup) : Program("Im2ColMatMul"),
-                                           has_bias_(has_bias),
-                                           tile_m_(tile_m),
-                                           tile_n_(tile_n),
-                                           vec_size_(vec_size),
-                                           use_subgroup_(use_subgroup) {}
+                      bool use_subgroup,
+                      ActivationKind activation_kind) : Program("Im2ColMatMul"),
+                                                        has_bias_(has_bias),
+                                                        tile_m_(tile_m),
+                                                        tile_n_(tile_n),
+                                                        vec_size_(vec_size),
+                                                        use_subgroup_(use_subgroup),
+                                                        activation_kind_(activation_kind) {}
 
   Status GenerateShaderCode(ShaderHelper& shader) const override;
 
@@ -50,7 +52,8 @@ class Im2ColMatMulProgram final : public Program<Im2ColMatMulProgram> {
       {"K_tiles", ProgramUniformVariableDataType::Uint32},
       {"dilations", ProgramUniformVariableDataType::Uint32},
       {"pads", ProgramUniformVariableDataType::Uint32},
-      {"strides", ProgramUniformVariableDataType::Uint32});
+      {"strides", ProgramUniformVariableDataType::Uint32},
+      WEBGPU_PROGRAM_ACTIVATION_UNIFORM_VARIABLES);
 
  private:
   bool has_bias_;
@@ -59,17 +62,19 @@ class Im2ColMatMulProgram final : public Program<Im2ColMatMulProgram> {
   uint32_t tile_n_;
   uint32_t vec_size_;
   bool use_subgroup_;
+  ActivationKind activation_kind_;
 };
 
 bool CanApplyIm2ColMatMulProgram(ComputeContextBase& context,
                                  const bool is_channels_last,
-                                 const bool is_fused,
+                                 const Activation& activation,
                                  const TensorShape kernel_shape,
                                  const uint32_t group,
                                  const MLDataType data_type);
 
 Status ApplyIm2ColMatMulProgram(ComputeContext& context,
                                 const bool is_channels_last,
+                                const Activation& activation,
                                 const std::vector<uint32_t>& dilations,
                                 const std::vector<uint32_t>& pads,
                                 const std::vector<uint32_t>& strides,

onnxruntime/core/providers/webgpu/nn/im2col_matmul.wgsl.template

@@ -6,6 +6,10 @@
 #param tile_n
 #param use_subgroup
 #param vec_size
+// Mirrors ActivationKind; static_asserts in im2col_matmul.cc enforce these values.
+// 0=None, 1=Relu, 2=Sigmoid, 3=Clip, 4=HardSigmoid, 5=LeakyRelu, 6=Tanh.
+// Keep branches synchronized with IsActivationSupported().
+#param activation_kind
 
 #use .getByOffset .setByOffset
 
@@ -142,6 +146,19 @@ $MAIN {
   let bias = load_bias(n_base);
   for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {
     var output_data = results[m_idx] + bias;
+#if activation_kind == 1
+    output_data = max(output_data, output_element_t(0));
+#elif activation_kind == 2
+    output_data = output_element_t(1) / (output_element_t(1) + exp(-output_data));
+#elif activation_kind == 3
+    output_data = clamp(output_data, output_element_t(uniforms.activation_param_0), output_element_t(uniforms.activation_param_1));
+#elif activation_kind == 4
+    output_data = clamp(output_element_t(uniforms.activation_param_0) * output_data + output_element_t(uniforms.activation_param_1), output_element_t(0), output_element_t(1));
+#elif activation_kind == 5
+    output_data = select(output_element_t(uniforms.activation_param_0) * output_data, output_data, output_data >= output_element_t(0));
+#elif activation_kind == 6
+    output_data = tanh(output_data);
+#endif
     write_output(batch, m_base + m_idx, n_base, output_data);
   }
 }  // MAIN

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -6,6 +6,7 @@
 #endif
 
 #include <functional>
+#include <fstream>
 #include <random>
 #include <sstream>
 
@@ -3958,6 +3959,108 @@ ConvActivationBuilder SimpleActivation(const std::string& op_type,
   };
 }
 
+// Use fp16, group 1, NHWC, and a non-1x1 kernel to select Im2ColMatMulProgram.
+void RunWebGpuIm2ColActivationParity(const ConvActivationBuilder& add_activation,
+                                     const std::string& expected_activation,
+                                     int opset_version) {
+  if (!DefaultWebGpuExecutionProvider()) {
+    GTEST_SKIP() << "WebGPU EP unavailable in this build.";
+  }
+
+  const std::vector<int64_t> input_shape{1, 4, 14, 14};
+  const std::vector<int64_t> weight_shape{8, 4, 3, 3};
+
+  auto to_fp16 = [](const std::vector<float>& values) {
+    std::vector<MLFloat16> converted;
+    converted.reserve(values.size());
+    for (float v : values) {
+      converted.push_back(MLFloat16(v));
+    }
+    return converted;
+  };
+
+  // Use deterministic signed inputs for stable fp16 comparisons.
+  auto ramp = [](size_t count, float lo, float hi) {
+    std::vector<float> values(count);
+    for (size_t i = 0; i < count; ++i) {
+      values[i] = lo + (hi - lo) * (static_cast<float>(i % 32) / 31.0f);
+    }
+    return values;
+  };
+
+  auto build_test_case = [&](ModelTestBuilder& builder) {
+    auto* input = builder.MakeInput<MLFloat16>(input_shape, to_fp16(ramp(1 * 4 * 14 * 14, -3.0f, 3.0f)));
+    auto* weight = builder.MakeInitializer<MLFloat16>(weight_shape, to_fp16(ramp(8 * 4 * 3 * 3, -0.5f, 0.5f)));
+    auto* bias = builder.MakeInitializer<MLFloat16>({weight_shape[0]}, to_fp16(ramp(8, -0.5f, 0.5f)));
+    auto* conv_out = builder.MakeIntermediate();
+    auto* output = builder.MakeOutput();
+
+    builder.AddNode("Conv", {input, weight, bias}, {conv_out});
+    add_activation(builder, conv_out, output);
+  };
+
+  bool im2col_selected = false;
+  std::string observed_conv_programs;
+  auto check_transformed_graph = [&](InferenceSessionWrapper& session) {
+    bool fused = false;
+    std::ostringstream graph_description;
+    for (const auto& node : session.GetGraph().Nodes()) {
+      graph_description << " " << node.Domain() << "." << node.OpType()
+                        << "[" << node.GetExecutionProviderType() << "]";
+      if (node.OpType() != "Conv") {
+        continue;
+      }
+      const auto* activation_attr = graph_utils::GetNodeAttribute(node, "activation");
+      if (activation_attr != nullptr && activation_attr->s() == expected_activation) {
+        fused = true;
+      }
+    }
+    ASSERT_TRUE(fused) << "Conv did not absorb " << expected_activation
+                       << ", so no fused kernel ran at all. Graph was:" << graph_description.str();
+
+    // Confirm profiling observed Im2ColMatMulProgram rather than the fallback Conv path.
+    const std::string profile_path = session.EndProfiling();
+    ASSERT_FALSE(profile_path.empty()) << "profiling produced no file, so program selection is unverifiable";
+    std::ifstream profile_stream(profile_path);
+    ASSERT_TRUE(profile_stream.good()) << "cannot read profile " << profile_path;
+    std::stringstream buffer;
+    buffer << profile_stream.rdbuf();
+    const std::string profile_contents = buffer.str();
+
+    im2col_selected = profile_contents.find("Im2ColMatMul") != std::string::npos;
+    for (size_t pos = profile_contents.find("&Conv&"); pos != std::string::npos;
+         pos = profile_contents.find("&Conv&", pos + 1)) {
+      const size_t end = profile_contents.find('"', pos);
+      if (end != std::string::npos) {
+        observed_conv_programs += " " + profile_contents.substr(pos + 6, end - pos - 6);
+      }
+    }
+    // GPU timestamp events carry a "cache_key" argument. A device without timestamp query support
+    // emits none at all, which is not the same as dispatching an unexpected program.
+    if (profile_contents.find("cache_key") == std::string::npos) {
+      GTEST_SKIP() << "device reported no GPU timestamps, so program selection is unobservable here";
+    }
+    ASSERT_FALSE(observed_conv_programs.empty())
+        << "GPU kernels were profiled but none was a Conv dispatch, so this test can no longer tell "
+           "whether the im2col program was selected";
+  };
+
+  RunWebGpuFusionTransformerTest(build_test_case, check_transformed_graph, TransformerLevel::Level1, TransformerLevel::Level2, opset_version,
+                                 /*per_sample_tolerance=*/2e-2,
+                                 /*relative_per_sample_tolerance=*/2e-2,
+                                 /*transformer=*/nullptr, []() { return DefaultWebGpuExecutionProvider(); }, [](SessionOptions& session_options) {
+                                   session_options.enable_profiling = true;
+                                   session_options.profile_file_prefix = ORT_TSTR("webgpu_im2col_activation"); });
+
+  // Im2ColMatMulProgram is restricted to Intel Xe-2/Xe-3.
+  if (!im2col_selected) {
+    GTEST_SKIP() << "Im2ColMatMul did not run on this adapter, so the im2col activation epilogue was "
+                    "not exercised. Conv dispatched to:"
+                 << observed_conv_programs
+                 << ". Requires an Intel Xe-2/Xe-3 GPU per IsDeviceSupported().";
+  }
+}
+
 }  // namespace
 
 TEST_F(GraphTransformationTests, WebGpuConvReluFusionMatchesUnfusedResults) {
@@ -4047,6 +4150,38 @@ TEST_F(GraphTransformationTests, WebGpuConvLeakyReluParityAcrossAlphaValues) {
       "LeakyRelu", 17, input_shape, weight_shape);
 }
 
+TEST_F(GraphTransformationTests, WebGpuIm2ColConvReluFusionMatchesUnfusedResults) {
+  RunWebGpuIm2ColActivationParity(SimpleActivation("Relu"), "Relu", 17);
+}
+
+TEST_F(GraphTransformationTests, WebGpuIm2ColConvLeakyReluFusionMatchesUnfusedResults) {
+  RunWebGpuIm2ColActivationParity(
+      SimpleActivation("LeakyRelu", kOnnxDomain, [](Node& node) { node.AddAttribute("alpha", 0.25f); }),
+      "LeakyRelu", 17);
+}
+
+TEST_F(GraphTransformationTests, WebGpuIm2ColConvHardSigmoidFusionMatchesUnfusedResults) {
+  RunWebGpuIm2ColActivationParity(
+      SimpleActivation("HardSigmoid", kOnnxDomain,
+                       [](Node& node) {
+                         node.AddAttribute("alpha", 0.3f);
+                         node.AddAttribute("beta", 0.7f);
+                       }),
+      "HardSigmoid", 17);
+}
+
+// Clip is the only two-slot activation whose uniform slots mean {min, max} rather than
+// {alpha, beta}, so it is the one case where swapping or mis-indexing activation_param_0/1
+// would survive the HardSigmoid test above.
+TEST_F(GraphTransformationTests, WebGpuIm2ColConvClipFusionMatchesUnfusedResults) {
+  auto add_clip = [](ModelTestBuilder& builder, NodeArg* conv_out, NodeArg* output) {
+    // min/max must match the fp16 tensor type this path requires.
+    auto* min_value = builder.MakeScalarInitializer<MLFloat16>(MLFloat16(-0.25f));
+    auto* max_value = builder.MakeScalarInitializer<MLFloat16>(MLFloat16(0.75f));
+    builder.AddNode("Clip", {conv_out, min_value, max_value}, {output});
+  };
+  RunWebGpuIm2ColActivationParity(add_clip, "Clip", 17);
+}
 #endif  // defined(USE_WEBGPU)
 #endif  // !defined(DISABLE_CONTRIB_OPS)