diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm.cc b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm.cc index 34f73748a9603..198d6073dbed8 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm.cc +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm.cc @@ -8,25 +8,19 @@ #include #include #include -#include #include #include "core/common/narrow.h" #include "core/providers/webgpu/compute_context.h" +#include "core/providers/webgpu/subgroup_matrix_common.h" #include "core/providers/webgpu/math/subgroup_matrix_config.h" #include "core/providers/webgpu/shader_helper.h" -#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" namespace onnxruntime { namespace webgpu { namespace { -// Lanes per subgroup assumed by the subgroup-matrix kernel. The workgroup runs -// split_k subgroups, so its size is kSubgroupMatrixSubgroupSize * split_k. -// TODO: use subgroup-size-control to enforce the subgroup size is 32. -constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; - // Subgroup-matrix Gemm implementation. Loads A and B directly from global memory // (transposed operands via column-major loads) and runs the cooperative // subgroup-matrix kernel during Compute. Y = alpha * op(A) @ op(B) + beta * C. @@ -173,15 +167,6 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelpe WGSL_TEMPLATE_VARIABLE(output, output)); } -// Default tiling used on any vendor without a specialized policy: a fixed 32x32 -// output tile with no split-K. -SubgroupMatrixTilingSelector MakeDefaultTilingSelector() { - return [](const ComputeContext&, uint32_t /*M*/, uint32_t /*N*/, - uint32_t /*K*/, uint32_t /*batch*/) -> std::optional { - return SubgroupMatrixTiling{32, 32, 1}; - }; -} - } // namespace Status SubgroupMatrixGemmProgram::GenerateShaderCode(ShaderHelper& shader) const { @@ -203,19 +188,9 @@ Status SubgroupMatrixGemmProgram::GenerateShaderCode(ShaderHelper& shader) const std::unique_ptr CreateSubgroupMatrixGemmImpl( const Gemm& parent, const ComputeContextBase& context) { - // Only run on devices that report the fixed 8x16x16 F16 subgroup-matrix config - // this kernel is implemented for. int32_t config_index = 0; - if (!IsSubgroupMatrixConfigSupported(context, /*is_fp16=*/true, config_index) || - !supported_subgroup_matrix_configs[config_index].Is(8, 16, 16)) { - return nullptr; - } - // Intel GPUs use a tuned/heuristic tiling policy; every other vendor falls back - // to a fixed default tiling. - const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; - SubgroupMatrixTilingSelector tiling_selector = - is_intel ? intel::CreateSubgroupMatrixTilingSelector(context) : MakeDefaultTilingSelector(); - if (!tiling_selector) { + SubgroupMatrixTilingSelector tiling_selector; + if (!TrySelectSubgroupMatrixConfig(context, config_index, tiling_selector)) { return nullptr; } return std::make_unique(parent, config_index, std::move(tiling_selector)); diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc index ca632c8968245..a1df3f5ca4c63 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc @@ -10,43 +10,18 @@ #include #include #include -#include #include #include "core/common/narrow.h" #include "core/providers/webgpu/compute_context.h" #include "core/providers/webgpu/math/subgroup_matrix_config.h" #include "core/providers/webgpu/shader_helper.h" -#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" #include "core/providers/webgpu/webgpu_utils.h" namespace onnxruntime { namespace webgpu { namespace { -// Lanes per subgroup assumed by the subgroup-matrix kernel. The workgroup runs -// split_k subgroups, so its size is kSubgroupMatrixSubgroupSize * split_k. -// TODO: use subgroup-size-control to enforce the subgroup size is 32. -constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; - -// Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer -// (N_b >= N), zero-filling columns [N, N_b). Gives B an even row stride so the -// subgroup-matrix f16 load's 4-byte row-start alignment holds for odd N. -class SubgroupMatrixMatMulPadBProgram final : public Program { - public: - SubgroupMatrixMatMulPadBProgram() : Program{"SubgroupMatrixMatMulPadB"} {} - Status GenerateShaderCode(ShaderHelper& shader) const override { - const auto& input_b = shader.AddInput("input_b", ShaderUsage::UseValueTypeAlias); - const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); - return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul_pad_b.wgsl.template", - WGSL_TEMPLATE_VARIABLE(input_b, input_b), - WGSL_TEMPLATE_VARIABLE(output, output)); - } - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"N_b", ProgramUniformVariableDataType::Uint32}); -}; - // Subgroup-matrix MatMul implementation. Loads both A and B directly from global // memory and runs the subgroup-matrix kernel during Compute. The class is // intended to support all subgroup-matrix configs; for now only 8x16x16 is @@ -61,190 +36,20 @@ class SubgroupMatrixMatMulImpl final : public MatMul::MatMulOptImpl { tiling_selector_(std::move(tiling_selector)) {} Status Compute(ComputeContext& context, /*out*/ bool& handled) override { - handled = false; - const auto* a = context.Input(0); const auto* b = context.Input(1); - const auto& a_shape = a->Shape(); - const auto& b_shape = b->Shape(); - if (a_shape.NumDimensions() < 2 || !a->IsDataType() || !b->IsDataType()) { - return Status::OK(); - } - - const uint32_t K = narrow(a_shape[a_shape.NumDimensions() - 1]); - // The tiling selector is responsible for any subgroup-matrix alignment - // requirements (e.g. K % sg_mat_k == 0) and declines otherwise. - if (K == 0) { - return Status::OK(); - } - - // Two shapes are handled: - // * Shared 2D weight B [K, N]: all leading A dims fold into M and the whole - // problem runs as one z-slice (batch == 1). Optimal for weight matmuls - - // B is read once and M-parallelism is maximal. - // * Batched B [..., K, N] (true bmm): A is [..., M, K] with batch dims - // identical to B (no broadcasting). Each (A, B) pair is one z-dispatched - // slice; the kernel derives the per-slice A/B/output strides from M, N, K. - // Anything else (e.g. broadcasting A across B's batch) falls back. - uint32_t M = 0; - uint32_t N = 0; - uint32_t batch = 1; - if (b_shape.NumDimensions() == 2) { - ORT_ENFORCE(narrow(b_shape[0]) == K, - "MatMul contraction dim mismatch: A K=", K, " vs B rows=", b_shape[0]); - N = narrow(b_shape[1]); - M = narrow(a_shape.Size() / static_cast(K)); - } else { - // Batched B (true bmm): A is [..., M, K], B is [..., K, N]. The kernel pairs - // slice i of A with slice i of B and copies A's shape to the output, so it - // can only be correct when A and B have identical batch dims - it does not - // implement numpy broadcasting. Require equal rank and equal leading dims; - // anything broadcastable-but-not-identical (e.g. A=[2,1,M,K], B=[1,2,K,N]) - // falls back to the generic MatMul path. - const size_t rank = a_shape.NumDimensions(); - if (b_shape.NumDimensions() != rank) { - return Status::OK(); - } - ORT_ENFORCE(narrow(b_shape[rank - 2]) == K, - "MatMul contraction dim mismatch: A K=", K, " vs B rows=", b_shape[rank - 2]); - M = narrow(a_shape[rank - 2]); - N = narrow(b_shape[rank - 1]); - for (size_t i = 0; i + 2 < rank; ++i) { - if (a_shape[i] != b_shape[i]) { - return Status::OK(); - } - } - batch = narrow(a_shape.SizeToDimension(rank - 2)); - } - - // An empty M or N yields an empty output (ONNX permits zero-length dims). Let - // the generic MatMul path allocate the empty result rather than dispatch a - // degenerate (zero-workgroup) kernel. - if (M == 0 || N == 0) { - return Status::OK(); - } - - // The B right-operand is loaded with subgroupMatrixLoad using a row stride of - // N_b. Intel's f16 subgroup-matrix load reads columns in 32-bit (2xf16) pairs - // and requires each K-row to start 4-byte aligned, i.e. an even element stride. - // An odd N would offset every other K-row by 2 bytes and corrupt the odd output - // columns. For a constant weight (2D or batched) we can pad B to an even stride - // (N_b = N + 1); a non-constant odd-N B falls back here. This is only the cheap - // eligibility check - the actual padding is deferred until after tiling selection - // so we never allocate/dispatch a padded copy for a problem that will fall back - // anyway (e.g. K % 16 != 0, which the tiling selector declines). B has already - // been validated as 2D or a well-formed batched shape above. - if (N % 2 != 0 && !parent_.IsBConstant()) { - return Status::OK(); - } - - const std::optional tiling = tiling_selector_(context, M, N, K, batch); - if (!tiling) { - return Status::OK(); - } - - // The optimized path will run: now materialize the even-strided B for odd N. - const Tensor* b_used = b; - uint32_t N_b = N; - if (N % 2 != 0) { - ORT_RETURN_IF_ERROR(EnsurePaddedB(context, *b, N)); - b_used = padded_b_.get(); - N_b = padded_b_stride_; - } - - TensorShapeVector output_dims{a_shape.GetDims().begin(), a_shape.GetDims().end()}; - output_dims.back() = static_cast(N); - TensorShape output_shape{output_dims}; - auto* output = context.Output(0, output_shape); - - const bool has_bias = context.InputCount() > 2; - const Tensor* bias = has_bias ? context.Input(2) : nullptr; - - const auto& config = supported_subgroup_matrix_configs[config_index_]; - const uint32_t tile_m = tiling->tile_m; - const uint32_t tile_n = tiling->tile_n; - const uint32_t split_k = tiling->split_k; - const uint32_t sg_mat_count_m = tile_m / config.M; - const uint32_t sg_mat_count_n = tile_n / config.N; - ORT_ENFORCE(tile_m % config.M == 0 && tile_n % config.N == 0, - "Tiling must be a multiple of the subgroup-matrix shape: ", - tile_m, "x", tile_n, " vs ", config.M, "x", config.N); - const uint32_t dispatch_x = (N + tile_n - 1) / tile_n; - const uint32_t dispatch_y = (M + tile_m - 1) / tile_m; - - SubgroupMatrixMatMulProgram program{has_bias, config_index_, sg_mat_count_m, sg_mat_count_n, split_k}; - program.SetWorkgroupSize(kSubgroupMatrixSubgroupSize * split_k); - program.SetDispatchGroupSize(dispatch_x, dispatch_y, batch); - program.CacheHint(has_bias, config_index_, sg_mat_count_m, sg_mat_count_n, split_k) - .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b_used, ProgramTensorMetadataDependency::TypeAndRank, 1}}) - .AddOutput({output, ProgramTensorMetadataDependency::Rank, output->Shape(), 1}) - .AddUniformVariables({{M}, {N}, {K}, {dispatch_x}, {N_b}}); - if (has_bias) { - program.AddInput({bias, ProgramTensorMetadataDependency::None}); - } - ORT_RETURN_IF_ERROR(context.RunProgram(program)); - - handled = true; - return Status::OK(); + const Tensor* bias = context.InputCount() > 2 ? context.Input(2) : nullptr; + return DispatchSubgroupMatrixMatMul(context, config_index_, tiling_selector_, pad_cache_, + a, b, bias, /*output=*/nullptr, + a->Shape(), b->Shape(), parent_.IsBConstant(), handled); } private: - // Lazily builds an even-strided copy of a constant weight B [..., K, N] with odd N - // by widening its last dim to N_b = N + 1 (zero-filling the extra column) and - // caches it, so the per-run pad cost is paid once. Works for a 2D weight [K, N] - // and a batched weight [batch, K, N] alike: the pad pass treats B as a flat - // [rows, N] -> [rows, N_b] copy over rows = numel / N (= K, or batch*K), which is - // exactly the even-stride layout the kernel indexes via N_b. Runs the GPU pad pass - // on first use under call_once; the cached tensor is held for the kernel's - // lifetime. Only valid when B is a constant initializer (checked by the caller) - - // a runtime B changes per run and must not be cached. - Status EnsurePaddedB(ComputeContext& context, const Tensor& b, uint32_t N) const { - ORT_RETURN_IF_NOT(N < std::numeric_limits::max(), - "Cannot pad odd-N B because N+1 exceeds uint32_t range."); - const uint32_t n_b = N + 1; - TensorShapeVector padded_dims{b.Shape().GetDims().begin(), b.Shape().GetDims().end()}; - padded_dims.back() = static_cast(n_b); - const TensorShape padded_shape{padded_dims}; - const int64_t output_size_i64 = padded_shape.Size(); - ORT_RETURN_IF_NOT(output_size_i64 <= static_cast(std::numeric_limits::max()), - "Cannot pad odd-N B because the padded tensor has ", output_size_i64, - " elements, exceeding uint32_t shader indexing range."); - const uint32_t output_size = narrow(output_size_i64); - - std::call_once(pad_once_, [&]() { - auto padded = std::make_unique(context.CreateGPUTensor(b.DataType(), padded_shape)); - Status s = Status::OK(); - // A zero-element padded tensor (e.g. a zero-batch or empty constant B) needs no - // pad pass - dispatching 0 workgroups is pointless and some drivers reject it. - // Just cache the empty tensor; the main kernel dispatches nothing for it. - if (output_size != 0) { - SubgroupMatrixMatMulPadBProgram program; - program.SetWorkgroupSize(WORKGROUP_SIZE) - .SetDispatchGroupSize(CeilDiv(output_size, WORKGROUP_SIZE)) - .AddInput({&b, ProgramTensorMetadataDependency::TypeAndRank, 1}) - .AddOutput({padded.get(), ProgramTensorMetadataDependency::TypeAndRank, padded->Shape(), 1}) - .AddUniformVariables({{output_size}, {N}, {n_b}}); - s = context.RunProgram(program); - } - if (s.IsOK()) { - padded_b_ = std::move(padded); - padded_b_stride_ = n_b; - } - }); - // padded_b_ persists the outcome across calls: call_once runs the body only on - // the first call, so a failed pad stays failed (and null) on later calls. - return padded_b_ ? Status::OK() - : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to pad odd-N B for subgroup-matrix MatMul."); - } - const int32_t config_index_; SubgroupMatrixTilingSelector tiling_selector_; - // Cached even-strided B for odd N; built once by EnsurePaddedB. - mutable std::once_flag pad_once_; - mutable std::unique_ptr padded_b_; - mutable uint32_t padded_b_stride_ = 0; + // Odd-N even-strided B cache, shared with the Conv 1x1 path. + mutable SubgroupMatrixPadBCache pad_cache_; }; Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelper& output, @@ -261,16 +66,194 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelpe WGSL_TEMPLATE_VARIABLE(output, output)); } -// Default tiling used on any vendor without a specialized policy: a fixed 32x32 -// output tile with no split-K. -SubgroupMatrixTilingSelector MakeDefaultTilingSelector() { - return [](const ComputeContext&, uint32_t /*M*/, uint32_t /*N*/, - uint32_t /*K*/, uint32_t /*batch*/) -> std::optional { - return SubgroupMatrixTiling{32, 32, 1}; - }; +} // namespace + +Status DispatchSubgroupMatrixMatMul(ComputeContext& context, + int32_t config_index, + const SubgroupMatrixTilingSelector& tiling_selector, + const SubgroupMatrixPadBCache& pad_cache, + const Tensor* a, const Tensor* b, const Tensor* bias, + Tensor* output, + const TensorShape& a_shape, const TensorShape& b_shape, + bool b_is_constant, + /*out*/ bool& handled) { + handled = false; + + if (a_shape.NumDimensions() < 2 || !a->IsDataType() || !b->IsDataType()) { + return Status::OK(); + } + + const uint32_t K = narrow(a_shape[a_shape.NumDimensions() - 1]); + // The tiling selector is responsible for any subgroup-matrix alignment + // requirements (e.g. K % sg_mat_k == 0) and declines otherwise. + if (K == 0) { + return Status::OK(); + } + + // Two shapes are handled: + // * Shared 2D weight B [K, N]: all leading A dims fold into M and the whole + // problem runs as one z-slice (batch == 1). Optimal for weight matmuls - + // B is read once and M-parallelism is maximal. + // * Batched B [..., K, N] (true bmm): A is [..., M, K] with batch dims + // identical to B (no broadcasting). Each (A, B) pair is one z-dispatched + // slice; the kernel derives the per-slice A/B/output strides from M, N, K. + // Anything else (e.g. broadcasting A across B's batch) falls back. + uint32_t M = 0; + uint32_t N = 0; + uint32_t batch = 1; + if (b_shape.NumDimensions() == 2) { + ORT_ENFORCE(narrow(b_shape[0]) == K, + "MatMul contraction dim mismatch: A K=", K, " vs B rows=", b_shape[0]); + N = narrow(b_shape[1]); + M = narrow(a_shape.Size() / static_cast(K)); + } else { + // Batched B (true bmm): A is [..., M, K], B is [..., K, N]. The kernel pairs + // slice i of A with slice i of B and copies A's shape to the output, so it + // can only be correct when A and B have identical batch dims - it does not + // implement numpy broadcasting. Require equal rank and equal leading dims; + // anything broadcastable-but-not-identical (e.g. A=[2,1,M,K], B=[1,2,K,N]) + // falls back to the generic MatMul path. + const size_t rank = a_shape.NumDimensions(); + if (b_shape.NumDimensions() != rank) { + return Status::OK(); + } + ORT_ENFORCE(narrow(b_shape[rank - 2]) == K, + "MatMul contraction dim mismatch: A K=", K, " vs B rows=", b_shape[rank - 2]); + M = narrow(a_shape[rank - 2]); + N = narrow(b_shape[rank - 1]); + for (size_t i = 0; i + 2 < rank; ++i) { + if (a_shape[i] != b_shape[i]) { + return Status::OK(); + } + } + batch = narrow(a_shape.SizeToDimension(rank - 2)); + } + + // An empty M or N yields an empty output (ONNX permits zero-length dims). Let the + // caller allocate the empty result rather than dispatch a degenerate + // (zero-workgroup) kernel. + if (M == 0 || N == 0) { + return Status::OK(); + } + + // The B right-operand is loaded with subgroupMatrixLoad using a row stride of + // N_b. Intel's f16 subgroup-matrix load reads columns in 32-bit (2xf16) pairs + // and requires each K-row to start 4-byte aligned, i.e. an even element stride. + // An odd N would offset every other K-row by 2 bytes and corrupt the odd output + // columns. For a constant weight we can pad B to an even stride (N_b = N + 1); a + // non-constant odd-N B falls back here. The actual padding is deferred until after + // tiling selection so we never allocate a padded copy for a problem that will fall + // back anyway (e.g. K % 16 != 0, which the tiling selector declines). + if (N % 2 != 0 && !b_is_constant) { + return Status::OK(); + } + + const std::optional tiling = tiling_selector(context, M, N, K, batch); + if (!tiling) { + return Status::OK(); + } + + // The optimized path will run: now materialize the even-strided B for odd N. + const Tensor* b_used = b; + uint32_t N_b = N; + if (N % 2 != 0) { + ORT_RETURN_IF_ERROR(pad_cache.EnsurePaddedB(context, *b, N, b_used, N_b)); + } + + // Use the caller's pre-allocated output when provided (a reshaped flat buffer); + // otherwise allocate one shaped like A with its trailing dim replaced by N. + Tensor* out = output; + if (out == nullptr) { + TensorShapeVector output_dims{a_shape.GetDims().begin(), a_shape.GetDims().end()}; + output_dims.back() = static_cast(N); + out = context.Output(0, TensorShape{output_dims}); + } + + const bool has_bias = bias != nullptr; + const auto& config = supported_subgroup_matrix_configs[config_index]; + const uint32_t tile_m = tiling->tile_m; + const uint32_t tile_n = tiling->tile_n; + const uint32_t split_k = tiling->split_k; + const uint32_t sg_mat_count_m = tile_m / config.M; + const uint32_t sg_mat_count_n = tile_n / config.N; + ORT_ENFORCE(tile_m % config.M == 0 && tile_n % config.N == 0, + "Tiling must be a multiple of the subgroup-matrix shape: ", + tile_m, "x", tile_n, " vs ", config.M, "x", config.N); + const uint32_t dispatch_x = (N + tile_n - 1) / tile_n; + const uint32_t dispatch_y = (M + tile_m - 1) / tile_m; + + SubgroupMatrixMatMulProgram program{has_bias, config_index, sg_mat_count_m, sg_mat_count_n, split_k}; + program.SetWorkgroupSize(kSubgroupMatrixSubgroupSize * split_k); + program.SetDispatchGroupSize(dispatch_x, dispatch_y, batch); + program.CacheHint(has_bias, config_index, sg_mat_count_m, sg_mat_count_n, split_k) + .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b_used, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddOutput({out, ProgramTensorMetadataDependency::Rank, out->Shape(), 1}) + .AddUniformVariables({{M}, {N}, {K}, {dispatch_x}, {N_b}}); + if (has_bias) { + program.AddInput({bias, ProgramTensorMetadataDependency::None}); + } + ORT_RETURN_IF_ERROR(context.RunProgram(program)); + + handled = true; + return Status::OK(); } -} // namespace +Status SubgroupMatrixMatMulPadBProgram::GenerateShaderCode(ShaderHelper& shader) const { + const auto& input_b = shader.AddInput("input_b", ShaderUsage::UseValueTypeAlias); + const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); + return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul_pad_b.wgsl.template", + WGSL_TEMPLATE_VARIABLE(input_b, input_b), + WGSL_TEMPLATE_VARIABLE(output, output)); +} + +// Works for a 2D weight [K, N] and a batched weight [batch, K, N] alike: the pad +// pass treats B as a flat [rows, N] -> [rows, N_b] copy over rows = numel / N, +// which is exactly the even-stride layout the kernel indexes via N_b. Runs the GPU +// pad pass on first use under call_once; the cached tensor is held for this cache's +// lifetime. +Status SubgroupMatrixPadBCache::EnsurePaddedB(ComputeContext& context, const Tensor& b, uint32_t N, + const Tensor*& b_used, uint32_t& n_b_out) const { + ORT_RETURN_IF_NOT(N < std::numeric_limits::max(), + "Cannot pad odd-N B because N+1 exceeds uint32_t range."); + const uint32_t n_b = N + 1; + TensorShapeVector padded_dims{b.Shape().GetDims().begin(), b.Shape().GetDims().end()}; + padded_dims.back() = static_cast(n_b); + const TensorShape padded_shape{padded_dims}; + const int64_t output_size_i64 = padded_shape.Size(); + ORT_RETURN_IF_NOT(output_size_i64 <= static_cast(std::numeric_limits::max()), + "Cannot pad odd-N B because the padded tensor has ", output_size_i64, + " elements, exceeding uint32_t shader indexing range."); + const uint32_t output_size = narrow(output_size_i64); + + std::call_once(pad_once_, [&]() { + auto padded = std::make_unique(context.CreateGPUTensor(b.DataType(), padded_shape)); + Status s = Status::OK(); + // A zero-element padded tensor (e.g. a zero-batch or empty constant B) needs no + // pad pass - dispatching 0 workgroups is pointless and some drivers reject it. + if (output_size != 0) { + SubgroupMatrixMatMulPadBProgram program; + program.SetWorkgroupSize(WORKGROUP_SIZE) + .SetDispatchGroupSize(CeilDiv(output_size, WORKGROUP_SIZE)) + .AddInput({&b, ProgramTensorMetadataDependency::TypeAndRank, 1}) + .AddOutput({padded.get(), ProgramTensorMetadataDependency::TypeAndRank, padded->Shape(), 1}) + .AddUniformVariables({{output_size}, {N}, {n_b}}); + s = context.RunProgram(program); + } + if (s.IsOK()) { + padded_b_ = std::move(padded); + padded_b_stride_ = n_b; + } + }); + // call_once runs the body only on the first call, so a failed pad stays failed + // (and null) on later calls. + if (!padded_b_) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to pad odd-N B for the subgroup-matrix kernel."); + } + b_used = padded_b_.get(); + n_b_out = padded_b_stride_; + return Status::OK(); +} Status SubgroupMatrixMatMulProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("input_a", ShaderUsage::UseUniform); @@ -290,19 +273,9 @@ Status SubgroupMatrixMatMulProgram::GenerateShaderCode(ShaderHelper& shader) con std::unique_ptr CreateSubgroupMatrixMatMulImpl( const MatMul& parent, const ComputeContextBase& context) { - // Only run on devices that report the fixed 8x16x16 F16 subgroup-matrix config - // this kernel is implemented for. int32_t config_index = 0; - if (!IsSubgroupMatrixConfigSupported(context, /*is_fp16=*/true, config_index) || - !supported_subgroup_matrix_configs[config_index].Is(8, 16, 16)) { - return nullptr; - } - // Intel GPUs use a tuned/heuristic tiling policy; every other vendor falls back - // to a fixed default tiling. - const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; - SubgroupMatrixTilingSelector tiling_selector = - is_intel ? intel::CreateSubgroupMatrixTilingSelector(context) : MakeDefaultTilingSelector(); - if (!tiling_selector) { + SubgroupMatrixTilingSelector tiling_selector; + if (!TrySelectSubgroupMatrixConfig(context, config_index, tiling_selector)) { return nullptr; } return std::make_unique(parent, config_index, std::move(tiling_selector)); diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h index e6f00d9aaa4cd..bb00660730c60 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h @@ -7,33 +7,17 @@ #include #include +#include #include #include "core/providers/webgpu/math/matmul.h" +#include "core/providers/webgpu/subgroup_matrix_common.h" #include "core/providers/webgpu/program.h" #include "core/providers/webgpu/shader_helper.h" namespace onnxruntime { namespace webgpu { -// Per-workgroup output tiling for one MatMul problem: the tile shape and split-K -// factor chosen by a vendor-specific policy. The subgroup-matrix shape itself is -// separate from this selection. -struct SubgroupMatrixTiling { - uint32_t tile_m; // output rows per workgroup - uint32_t tile_n; // output cols per workgroup - uint32_t split_k; // subgroups cooperating along K (1 = no split) -}; - -// Vendor-supplied callback that selects the output tiling for a given problem. -// batch is the number of z-dispatched slices (1 for a shared 2D weight), used by -// the policy to scale occupancy. Returning nullopt declines the problem, so -// MatMul falls back to another compute path. An empty selector likewise yields -// no implementation. -using SubgroupMatrixTilingSelector = - std::function(const ComputeContext& context, - uint32_t M, uint32_t N, uint32_t K, uint32_t batch)>; - // Creates a MatMulOptImpl that runs the subgroup-matrix kernel on devices whose // vendor policy supports it. The per-problem output tiling comes from a // vendor-specific selector chosen internally from the device context. Returns @@ -74,6 +58,58 @@ class SubgroupMatrixMatMulProgram final : public Program= N), zero-filling columns [N, N_b). Gives B an even row stride so the +// subgroup-matrix f16 load's 4-byte row-start alignment holds for odd N. Shared by +// the subgroup-matrix MatMul and Conv 1x1 paths. +class SubgroupMatrixMatMulPadBProgram final : public Program { + public: + SubgroupMatrixMatMulPadBProgram() : Program{"SubgroupMatrixMatMulPadB"} {} + Status GenerateShaderCode(ShaderHelper& sh) const override; + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"N_b", ProgramUniformVariableDataType::Uint32}); +}; + +// Caches an even-strided (N_b = N + 1) copy of a constant weight B with odd N so +// the column-pad pass runs once and is reused across inference steps. Shared by the +// subgroup-matrix MatMul and Conv 1x1 paths, which need the same odd-N alignment +// fix. Only valid when B is a constant initializer - a runtime B changes per run +// and must not be cached. +class SubgroupMatrixPadBCache { + public: + // Builds the padded copy of `b` (whose last dim must equal N) on first use and + // returns it via `b_used` with its row stride via `n_b`. + Status EnsurePaddedB(ComputeContext& context, const Tensor& b, uint32_t N, + /*out*/ const Tensor*& b_used, /*out*/ uint32_t& n_b) const; + + private: + mutable std::once_flag pad_once_; + mutable std::unique_ptr padded_b_; + mutable uint32_t padded_b_stride_ = 0; +}; + +// Runs the subgroup-matrix kernel for Y = A @ B (+ optional bias) given resolved +// operands and their logical shapes. a_shape/b_shape may differ from the tensors' +// own shapes when the caller reshaped them (e.g. a 1x1 Conv folding N,H,W into M). +// Handles the shared 2D-weight and batched-B cases, odd-N even-stride padding (via +// pad_cache when b_is_constant), vendor tiling and dispatch. When `output` is +// non-null the caller's pre-allocated tensor is used as a flat buffer (its element +// count must equal batch*M*N); when null the result is allocated via +// context.Output(0, ...) shaped like a_shape with its trailing dim set to N. Sets +// handled=true on success; leaves handled=false (touching nothing) on an +// unsupported device/problem so the caller can fall back. Shared by the MatMul +// operator and the Conv 1x1 path. +Status DispatchSubgroupMatrixMatMul(ComputeContext& context, + int32_t config_index, + const SubgroupMatrixTilingSelector& tiling_selector, + const SubgroupMatrixPadBCache& pad_cache, + const Tensor* a, const Tensor* b, const Tensor* bias, + Tensor* output, + const TensorShape& a_shape, const TensorShape& b_shape, + bool b_is_constant, + /*out*/ bool& handled); + } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index d980f819a264f..710006ac3d451 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -10,6 +10,9 @@ #include "core/providers/webgpu/nn/grouped_conv.h" #include "core/providers/webgpu/webgpu_utils.h" #include "core/providers/webgpu/math/matmul.h" +#if !defined(__wasm__) +#include "core/providers/webgpu/nn/subgroup_matrix_conv.h" +#endif namespace onnxruntime { namespace webgpu { @@ -29,6 +32,23 @@ Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const Tens template Status Conv::ComputeInternal(ComputeContext& context) const { +#if !defined(__wasm__) + // Subgroup-matrix 1x1 Conv has top priority: create the impl once (device query), + // then let it decide inside Compute whether the kernel/attrs qualify. It declines + // (handled=false, allocating nothing) otherwise and we fall through to the normal + // path order below. + std::call_once(impl_init_flag_, [&]() { + impl_ = CreateSubgroupMatrixConvImpl(*this, context); + }); + if (impl_ != nullptr) { + bool handled = false; + ORT_RETURN_IF_ERROR(impl_->Compute(context, handled)); + if (handled) { + return Status::OK(); + } + } +#endif + bool has_bias = context.InputCount() > 2; const auto* input = context.Input(0); const Tensor* kernel = nullptr; @@ -204,8 +224,11 @@ Status Conv::ComputeInternal(ComputeContext& context return context.RunProgram(program); } - const auto same_size = is_channels_last && input_height == kernel_height && input_width == kernel_width && pads[0] == 0 && pads[1] == 0; - if (same_size || (kernel_height == 1 && kernel_width == 1 && pads[0] == 0 && pads[1] == 0 && strides[0] == 1 && strides[1] == 1)) { + // Both reshapes need all pads (not just the leading ones) and the inferred output + // geometry to qualify; see the predicates in conv.h. + const auto same_size = is_channels_last && IsConvSameSizeMatMul(input_height, input_width, kernel_height, + kernel_width, output_height, output_width, pads); + if (same_size || IsConv1x1MatMul(kernel_height, kernel_width, pads, strides)) { Tensor transposed_kernel; TensorShape input_reshape; TensorShape kernel_reshape; @@ -367,9 +390,10 @@ Status Conv::PrePackInternal(ComputeContextBase& con const int64_t kernel_height = dims[2]; const int64_t kernel_width = dims[3]; - const bool is_1x1_conv = - (kernel_height == 1 && kernel_width == 1 && pads[0] == 0 && pads[1] == 0 && - strides[0] == 1 && strides[1] == 1); + // Must stay in sync with the gate in ComputeInternal: if this predicate accepted a + // Conv that ComputeInternal sends down the general path instead, that path would read + // a kernel we deliberately left untransposed here. + const bool is_1x1_conv = IsConv1x1MatMul(kernel_height, kernel_width, pads, strides); if constexpr (!is_channels_last) { if (is_1x1_conv) { diff --git a/onnxruntime/core/providers/webgpu/nn/conv.h b/onnxruntime/core/providers/webgpu/nn/conv.h index 5bf94a459a44a..031bafaedc27a 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -3,6 +3,11 @@ #pragma once +#include +#include +#include +#include + #include "core/common/optional.h" #include "core/providers/webgpu/webgpu_kernel.h" #include "core/providers/cpu/nn/conv_attributes.h" @@ -13,13 +18,77 @@ namespace onnxruntime { namespace webgpu { +// Eligibility for the MatMul Conv fast paths, which reshape a Conv into a plain 2D +// matmul (see Conv::ComputeInternal and the subgroup-matrix Conv impl). A reshape is +// only equivalent to the Conv when nothing is padded: a leading pad shifts the window, +// and a trailing pad adds output positions the matmul never computes. Comparing the +// inferred output geometry on top of that also covers dilations, which change the +// effective kernel size without appearing anywhere in the reshape. +// The containers hold the spatial attributes in ONNX order, i.e. pads is +// [begin..., end...] and strides is one entry per spatial dim; both must already be +// defaulted/normalized to the 2D layout the callers use. + +// True when every pad is zero, leading and trailing, across all spatial dims. +template +bool HasNoConvPadding(const PadContainer& pads) { + return std::all_of(pads.begin(), pads.end(), [](auto p) { return p == 0; }); +} + +// 1x1 reshape: each N,H,W position becomes an independent matmul row, so the output +// grid has to match the input grid one-to-one in row-major order. With a 1x1 kernel, no +// padding and unit strides that holds for any dilation (the effective kernel stays +// 1x1), so this predicate needs only the attributes -- PrePackInternal, which cannot +// see the input shape, uses the same one to stay in sync with ComputeInternal. +template +bool IsConv1x1MatMul(int64_t kernel_height, int64_t kernel_width, + const PadContainer& pads, const StrideContainer& strides) { + return kernel_height == 1 && kernel_width == 1 && HasNoConvPadding(pads) && + strides[0] == 1 && strides[1] == 1; +} + +// same-size reshape: the whole window folds into a single matmul row per batch element, +// so the kernel must cover the entire input and leave exactly one output position. +// Only valid for channels-last (the caller checks that). +template +bool IsConvSameSizeMatMul(int64_t input_height, int64_t input_width, + int64_t kernel_height, int64_t kernel_width, + int64_t output_height, int64_t output_width, + const PadContainer& pads) { + return input_height == kernel_height && input_width == kernel_width && + output_height == 1 && output_width == 1 && HasNoConvPadding(pads); +} + template class Conv : public WebGpuKernel { public: +#if !defined(__wasm__) + // Abstract base class for alternative optimized Conv implementations (currently the + // subgroup-matrix 1x1 / same-size path). Implementations hold the parent Conv and read + // conv_attrs/activation/... from it, mirroring MatMul::MatMulOptImpl. + class ConvOptImpl { + public: + explicit ConvOptImpl(const Conv& parent) : parent_(parent) {} + virtual ~ConvOptImpl() = default; + + // Attempts the optimized path, reading the Conv operands from `context` and the Conv + // attributes from the parent. Sets handled=true when it ran; leaves handled=false + // (allocating nothing) so the caller falls back to the normal Conv path. + virtual Status Compute(ComputeContext& context, /*out*/ bool& handled) = 0; + + protected: + const Conv& parent_; + }; +#endif + Conv(const OpKernelInfo& info) : WebGpuKernel(info), conv_attrs_(info) { if (is_fused) { ORT_ENFORCE(GetFusedActivationAttr(info, activation_).IsOK()); } + // Whether the weight input (index 1) is a constant initializer. The 1x1 Conv + // subgroup-matrix path uses this to decide it can pad an odd-N weight once and + // cache the result; a non-constant weight must not be cached. + const Tensor* w = nullptr; + w_is_constant_ = info.TryGetConstantInput(1, &w); } Status ComputeInternal(ComputeContext& context) const override; @@ -29,10 +98,27 @@ class Conv : public WebGpuKernel { AllocatorPtr alloc, /*out*/ bool& is_packed) override; + // State an optimized implementation (ConvOptImpl) reads from its parent Conv. + const ConvAttributes& ConvAttrs() const { return conv_attrs_; } + const Activation& ConvActivation() const { return activation_; } + // True when input 1 (the weight) is a constant initializer. See w_is_constant_. + bool IsWeightConstant() const { return w_is_constant_; } + // The prepacked (OIHW->HWIO transposed) weight, or null when it is not prepacked. + const Tensor* PrepackedKernel() const { return transposed_kernel_.get(); } + protected: ConvAttributes conv_attrs_; Activation activation_; std::unique_ptr transposed_kernel_; // should only have value when `is_initializer` AND `is_4D` AND `is_NHWC` + bool w_is_constant_ = false; // whether the weight input (index 1) is a constant initializer + + // Optional subgroup-matrix implementation for the 1x1 Conv path, lazily created + // on the first Compute call (once device capabilities can be queried). Null after + // initialization means the device has no subgroup-matrix support. +#if !defined(__wasm__) + mutable std::unique_ptr impl_; + mutable std::once_flag impl_init_flag_; +#endif }; Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const TensorShape& kernel_shape, Tensor* transposed_kernel, const InlinedVector& perm); diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc new file mode 100644 index 0000000000000..f783cf874ebe7 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -0,0 +1,207 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(__wasm__) + +#include "core/providers/webgpu/nn/subgroup_matrix_conv.h" + +#include +#include +#include +#include +#include + +#include "core/providers/webgpu/compute_context.h" +#include "core/providers/webgpu/subgroup_matrix_common.h" +#include "core/providers/webgpu/math/subgroup_matrix_matmul.h" + +namespace onnxruntime { +namespace webgpu { + +namespace { + +// Subgroup-matrix implementation of the 1x1 / same-size Conv matmul. Runs its own +// Conv shape inference to decide eligibility (mirroring Conv::ComputeInternal), then +// builds the matmul operands and dispatches the shared subgroup-matrix kernel. +template +class SubgroupMatrixConvImpl final : public Conv::ConvOptImpl { + using ConvT = Conv; + using Base = typename ConvT::ConvOptImpl; + + public: + SubgroupMatrixConvImpl(const ConvT& parent, int32_t config_index, + SubgroupMatrixTilingSelector tiling_selector) + : Base(parent), config_index_(config_index), tiling_selector_(std::move(tiling_selector)) {} + + Status Compute(ComputeContext& context, /*out*/ bool& handled) override { + handled = false; + + const ConvT& parent = this->parent_; + const ConvAttributes& conv_attrs = parent.ConvAttrs(); + const Activation& activation = parent.ConvActivation(); + const bool w_is_constant = parent.IsWeightConstant(); + const Tensor* prepacked_kernel = parent.PrepackedKernel(); + + // Only a constant weight is served here (so a runtime transpose can be cached once; + // a prepacked weight is already transposed). It must also be a plain, non-grouped + // matmul with no fused activation; anything else falls back to the normal Conv path. + // (channels-last is guaranteed by CreateSubgroupMatrixConvImpl.) + if (!w_is_constant || + activation.activation_kind_ != ActivationKind::None || conv_attrs.group != 1) { + return Status::OK(); + } + + const auto* input = context.Input(0); + // A prepacked weight is the OIHW->HWIO transposed kernel; otherwise read the + // original OIHW weight from input 1. + const bool kernel_is_prepacked = prepacked_kernel != nullptr; + const Tensor* kernel = kernel_is_prepacked ? prepacked_kernel : context.Input(1); + const Tensor* bias = context.InputCount() > 2 ? context.Input(2) : nullptr; + + // Shape inference mirrors Conv::ComputeInternal (channels-last). Only Conv2D (a + // rank-4 NHWC input) is supported. + TensorShape input_shape = input->Shape(); + if (input_shape.NumDimensions() != 4) { + return Status::OK(); + } + TensorShape kernel_shape = kernel_is_prepacked + ? TensorShape(TensorShapeVector{kernel->Shape()[3], kernel->Shape()[2], kernel->Shape()[0], kernel->Shape()[1]}) + : kernel->Shape(); + ConvAttributes::ConvPadVector local_pads(conv_attrs.pads.begin(), conv_attrs.pads.end()); + TensorShapeVector local_dilations(conv_attrs.dilations.begin(), conv_attrs.dilations.end()); + TensorShapeVector local_strides(conv_attrs.strides.begin(), conv_attrs.strides.end()); + TensorShapeVector kernel_spacial_shape_vector; + ORT_RETURN_IF_ERROR(conv_attrs.ComputeKernelShape(kernel_shape, kernel_spacial_shape_vector, false)); + if (local_pads.empty()) { + local_pads.resize(kernel_spacial_shape_vector.size() * 2, 0); + } + if (local_dilations.empty()) { + local_dilations.resize(kernel_spacial_shape_vector.size(), 1); + } + if (local_strides.empty()) { + local_strides.resize(kernel_spacial_shape_vector.size(), 1); + } + TensorShapeVector input_shape_vector = input_shape.AsShapeVector(); + const auto batch = input_shape[0]; + TensorShapeVector output_shape_vector = {batch}; + TensorShape input_spacial_shape = TensorShape(TensorShapeVector(std::next(input_shape_vector.begin()), std::prev(input_shape_vector.end()))); + ORT_RETURN_IF_ERROR(conv_attrs.InferPadsAndOutputShape(input_spacial_shape, kernel_spacial_shape_vector, local_strides, local_dilations, local_pads, output_shape_vector)); + const auto output_channels = kernel_shape[0]; + output_shape_vector.push_back(output_channels); + const TensorShape output_shape = TensorShape(output_shape_vector); + + const auto input_height = input_shape[1]; + const auto input_width = input_shape[2]; + const auto input_channels = input_shape[3]; + const auto kernel_height = kernel_shape[2]; + const auto kernel_width = kernel_shape[3]; + const auto output_height = output_shape_vector[1]; + const auto output_width = output_shape_vector[2]; + + // Both reshapes need all pads (not just the leading ones) and the inferred output + // geometry to qualify; see the predicates in conv.h, which the normal Conv MatMul + // path shares. + const bool same_size = IsConvSameSizeMatMul(input_height, input_width, kernel_height, kernel_width, + output_height, output_width, local_pads); + const bool is_conv_matmul = + same_size || IsConv1x1MatMul(kernel_height, kernel_width, local_pads, local_strides); + if (!is_conv_matmul) { + return Status::OK(); + } + + auto* output = context.Output(0, output_shape); + + // B is the OIHW->HWIO transposed weight. A prepacked weight is already transposed; + // otherwise transpose the constant weight once and reuse the cached copy. + const Tensor* weight = prepacked_kernel; + if (weight == nullptr) { + ORT_RETURN_IF_ERROR(EnsureTransposedKernel(context, kernel, kernel_shape)); + weight = cached_transposed_kernel_.get(); + } + + // The operands are a plain 2D weight matmul (M x K) @ (K x N): a 1x1 conv makes each + // N,H,W position an independent row; same-size folds the whole window into one row + // per batch element. + TensorShape a_shape; + TensorShape b_shape; + if (same_size) { + const auto shared_dim = input_height * input_width * input_channels; + a_shape = TensorShape({batch, shared_dim}); + b_shape = TensorShape({shared_dim, output_channels}); + } else { + a_shape = TensorShape({batch * input_height * input_width, input_channels}); + b_shape = TensorShape({input_channels, output_channels}); + } + + return DispatchSubgroupMatrixMatMul(context, config_index_, tiling_selector_, pad_cache_, + input, weight, bias, output, a_shape, b_shape, + /*b_is_constant=*/true, handled); + } + + private: + // Builds the OIHW->HWIO transposed weight on first use and caches it; only valid + // because the weight is a constant initializer (checked by the caller). + Status EnsureTransposedKernel(ComputeContext& context, const Tensor* kernel, + const TensorShape& kernel_shape) const { + std::call_once(transpose_once_, [&]() { + const InlinedVector perm = {2, 3, 1, 0}; + auto transposed = std::make_unique(); + if (TransposeKernel(context, kernel, kernel_shape, transposed.get(), perm).IsOK()) { + cached_transposed_kernel_ = std::move(transposed); + } + }); + return cached_transposed_kernel_ ? Status::OK() + : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to transpose Conv 1x1 weight."); + } + + const int32_t config_index_; + const SubgroupMatrixTilingSelector tiling_selector_; + + // OIHW->HWIO transposed constant weight, built once by EnsureTransposedKernel. + mutable std::once_flag transpose_once_; + mutable std::unique_ptr cached_transposed_kernel_; + + // Odd-N even-strided B cache, shared with the subgroup-matrix MatMul path. + mutable SubgroupMatrixPadBCache pad_cache_; +}; + +} // namespace + +template +std::unique_ptr::ConvOptImpl> CreateSubgroupMatrixConvImpl( + const Conv& parent, + const ComputeContextBase& context) { + // Only channels-last Conv is served by the subgroup-matrix path; the caller falls back + // to the normal Conv path for channels-first. + if constexpr (!is_channels_last) { + ORT_UNUSED_PARAMETER(parent); + ORT_UNUSED_PARAMETER(context); + return nullptr; + } else { + int32_t config_index = 0; + SubgroupMatrixTilingSelector tiling_selector; + if (!TrySelectSubgroupMatrixConfig(context, config_index, tiling_selector)) { + return nullptr; + } + return std::make_unique>(parent, config_index, + std::move(tiling_selector)); + } +} + +// Explicit instantiation for every Conv instantiation (see conv.cc). +#define WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(CHANNELS_LAST, FUSED) \ + template std::unique_ptr::ConvOptImpl> \ + CreateSubgroupMatrixConvImpl(const Conv& parent, \ + const ComputeContextBase& context); + +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(false, false) +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(false, true) +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(true, false) +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(true, true) + +#undef WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL + +} // namespace webgpu +} // namespace onnxruntime + +#endif // !defined(__wasm__) diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h new file mode 100644 index 0000000000000..a2eaa9dd1e2ba --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#if !defined(__wasm__) + +#include + +#include "core/providers/webgpu/compute_context.h" +#include "core/providers/webgpu/nn/conv.h" + +namespace onnxruntime { +namespace webgpu { + +// Creates a subgroup-matrix Conv 1x1 implementation on devices whose vendor policy supports +// the subgroup-matrix kernel; returns nullptr otherwise, so the caller falls back to the +// normal Conv path. The impl reads the Conv attributes from `parent`. +template +std::unique_ptr::ConvOptImpl> CreateSubgroupMatrixConvImpl( + const Conv& parent, + const ComputeContextBase& context); + +} // namespace webgpu +} // namespace onnxruntime + +#endif // !defined(__wasm__) diff --git a/onnxruntime/core/providers/webgpu/subgroup_matrix_common.cc b/onnxruntime/core/providers/webgpu/subgroup_matrix_common.cc new file mode 100644 index 0000000000000..09338a84557dc --- /dev/null +++ b/onnxruntime/core/providers/webgpu/subgroup_matrix_common.cc @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(__wasm__) + +#include "core/providers/webgpu/subgroup_matrix_common.h" + +#include +#include + +#include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" + +namespace onnxruntime { +namespace webgpu { + +SubgroupMatrixTilingSelector MakeDefaultSubgroupMatrixTilingSelector() { + return [](const ComputeContext&, uint32_t /*M*/, uint32_t /*N*/, + uint32_t K, uint32_t /*batch*/) -> std::optional { + // Only K needs to align to the subgroup-matrix shape; M and N partial tiles are + // handled by bounds-checked stores in the kernel. Decline a misaligned K, which + // the kernel would otherwise silently truncate to whole blocks. + if (K % kSubgroupMatrixK != 0) { + return std::nullopt; + } + // 32x32 is a multiple of the subgroup-matrix M/N and fits the scratch budget. + return SubgroupMatrixTiling{32, 32, 1}; + }; +} + +bool TrySelectSubgroupMatrixConfig(const ComputeContextBase& context, + int32_t& config_index, + SubgroupMatrixTilingSelector& tiling_selector) { + // Only devices reporting the fixed 8x16x16 F16 subgroup-matrix config these kernels + // are implemented for are supported. + config_index = 0; + if (!IsSubgroupMatrixConfigSupported(context, /*is_fp16=*/true, config_index) || + !supported_subgroup_matrix_configs[config_index].Is(kSubgroupMatrixM, kSubgroupMatrixN, kSubgroupMatrixK)) { + return false; + } + // Intel GPUs use a tuned/heuristic tiling policy; every other vendor falls back to a + // fixed default tiling. + const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; + SubgroupMatrixTilingSelector selector = + is_intel ? intel::CreateSubgroupMatrixTilingSelector(context) : MakeDefaultSubgroupMatrixTilingSelector(); + if (!selector) { + return false; + } + tiling_selector = std::move(selector); + return true; +} + +} // namespace webgpu +} // namespace onnxruntime + +#endif // !defined(__wasm__) diff --git a/onnxruntime/core/providers/webgpu/subgroup_matrix_common.h b/onnxruntime/core/providers/webgpu/subgroup_matrix_common.h new file mode 100644 index 0000000000000..05724917ad3ad --- /dev/null +++ b/onnxruntime/core/providers/webgpu/subgroup_matrix_common.h @@ -0,0 +1,66 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#if !defined(__wasm__) + +#include +#include +#include + +#include "core/providers/webgpu/compute_context.h" + +namespace onnxruntime { +namespace webgpu { + +// Lanes per subgroup assumed by the subgroup-matrix kernels. The workgroup runs +// split_k subgroups, so its size is kSubgroupMatrixSubgroupSize * split_k. +// TODO: use subgroup-size-control to enforce the subgroup size is 32. +constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; + +// The subgroup-matrix shape these kernels are implemented for (see +// TrySelectSubgroupMatrixConfig, which rejects any device not reporting it). +// A tiling must be a multiple of M/N, and K must be a multiple of +// kSubgroupMatrixK: the kernel walks the reduction in whole blocks +// (k_blocks = K / kSubgroupMatrixK) and has no K tail handling. +constexpr uint32_t kSubgroupMatrixM = 8; +constexpr uint32_t kSubgroupMatrixN = 16; +constexpr uint32_t kSubgroupMatrixK = 16; + +// Per-workgroup output tiling for one subgroup-matrix problem: the tile shape and +// split-K factor chosen by a vendor-specific policy. The subgroup-matrix shape +// itself is separate from this selection. +struct SubgroupMatrixTiling { + uint32_t tile_m; // output rows per workgroup + uint32_t tile_n; // output cols per workgroup + uint32_t split_k; // subgroups cooperating along K (1 = no split) +}; + +// Vendor-supplied callback that selects the output tiling for a given problem. +// batch is the number of z-dispatched slices (1 for a shared 2D weight), used by +// the policy to scale occupancy. Returning nullopt declines the problem, so the +// caller falls back to another compute path. An empty selector likewise yields no +// implementation. +using SubgroupMatrixTilingSelector = + std::function(const ComputeContext& context, + uint32_t M, uint32_t N, uint32_t K, uint32_t batch)>; + +// Default tiling used on any vendor without a specialized policy: a fixed 32x32 +// output tile with no split-K. The fallback selector when no vendor policy applies. +// Declines (nullopt) problems whose K is not a multiple of kSubgroupMatrixK, which +// the kernel cannot compute. +SubgroupMatrixTilingSelector MakeDefaultSubgroupMatrixTilingSelector(); + +// Selects the device subgroup-matrix config and the vendor tiling selector shared by +// the MatMul, Gemm and Conv 1x1 factories. On a supported device sets config_index / +// tiling_selector and returns true; returns false when the device or vendor is +// unsupported, so the caller yields no implementation. +bool TrySelectSubgroupMatrixConfig(const ComputeContextBase& context, + /*out*/ int32_t& config_index, + /*out*/ SubgroupMatrixTilingSelector& tiling_selector); + +} // namespace webgpu +} // namespace onnxruntime + +#endif // !defined(__wasm__) diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.cc index 2e32f2a2d8830..ba16e1a928ea2 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.cc +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.cc @@ -14,7 +14,7 @@ #include "core/common/narrow.h" #include "core/providers/webgpu/compute_context.h" -#include "core/providers/webgpu/math/subgroup_matrix_matmul.h" +#include "core/providers/webgpu/subgroup_matrix_common.h" #include "core/providers/webgpu/vendor/intel/intel_device_info.h" // Pretuned tile + split-K table baked into the build; consulted before the @@ -27,10 +27,8 @@ namespace intel { namespace { -// Subgroup matrix configuration used by this implementation (Intel Xe2/Xe3, F16). -constexpr uint32_t kSubgroupMatrixM = 8; -constexpr uint32_t kSubgroupMatrixN = 16; -constexpr uint32_t kSubgroupMatrixK = 16; +// The subgroup-matrix shape (8x16x16, F16) comes from subgroup_matrix_common.h: +// kSubgroupMatrixM / kSubgroupMatrixN / kSubgroupMatrixK. constexpr uint32_t kTileMCandidates[] = {8, 16, 32, 64}; constexpr uint32_t kTileNCandidates[] = {16, 32, 64}; diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h b/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h index a4fd68ec761a3..c2505355b2405 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h @@ -5,10 +5,7 @@ #if !defined(__wasm__) -#include - -#include "core/providers/webgpu/math/matmul.h" -#include "core/providers/webgpu/math/subgroup_matrix_matmul.h" +#include "core/providers/webgpu/subgroup_matrix_common.h" namespace onnxruntime { namespace webgpu { diff --git a/onnxruntime/test/providers/webgpu/conv_large_test.cc b/onnxruntime/test/providers/webgpu/conv_large_test.cc new file mode 100644 index 0000000000000..7e7d71ba8c2de --- /dev/null +++ b/onnxruntime/test/providers/webgpu/conv_large_test.cc @@ -0,0 +1,190 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" + +#include "test/providers/provider_test_utils.h" +#include "test/common/tensor_op_test_utils.h" +#include "default_providers.h" + +namespace onnxruntime { +namespace test { + +// Compute Conv2D output height/width given input, kernel, pad, stride, dilation. +static int64_t ConvOutputSize(int64_t input_size, int64_t kernel_size, int64_t pad_begin, int64_t pad_end, + int64_t stride, int64_t dilation) { + int64_t effective_kernel = (kernel_size - 1) * dilation + 1; + return (input_size + pad_begin + pad_end - effective_kernel) / stride + 1; +} + +// Reference implementation for Conv2D (NCHW layout) with group, stride, pad, dilation, and optional bias. +static void ComputeExpectedConv2D(const std::vector& x_vals, const std::vector& w_vals, + const std::vector& b_vals, std::vector& out_vals, + int64_t N, int64_t C, int64_t H, int64_t W, + int64_t M, int64_t kH, int64_t kW, + int64_t oH, int64_t oW, + int64_t pad_h_begin, int64_t pad_w_begin, + int64_t stride_h, int64_t stride_w, + int64_t dilation_h, int64_t dilation_w, + int64_t group) { + const int64_t C_per_group = C / group; + const int64_t M_per_group = M / group; + + for (int64_t n = 0; n < N; ++n) { + for (int64_t m = 0; m < M; ++m) { + const int64_t g = m / M_per_group; + const int64_t m_in_group = m % M_per_group; + for (int64_t oh = 0; oh < oH; ++oh) { + for (int64_t ow = 0; ow < oW; ++ow) { + float sum = b_vals.empty() ? 0.0f : b_vals[m]; + for (int64_t c = 0; c < C_per_group; ++c) { + for (int64_t kh = 0; kh < kH; ++kh) { + const int64_t ih = oh * stride_h + kh * dilation_h - pad_h_begin; + if (ih < 0 || ih >= H) continue; + for (int64_t kw = 0; kw < kW; ++kw) { + const int64_t iw = ow * stride_w + kw * dilation_w - pad_w_begin; + if (iw < 0 || iw >= W) continue; + const size_t x_idx = static_cast(((n * C + g * C_per_group + c) * H + ih) * W + iw); + const size_t w_idx = static_cast(((m_in_group + g * M_per_group) * C_per_group + c) * kH * kW + + kh * kW + kw); + sum += x_vals[x_idx] * w_vals[w_idx]; + } + } + } + const size_t o_idx = static_cast(((n * M + m) * oH + oh) * oW + ow); + out_vals[o_idx] = sum; + } + } + } + } +} + +struct ConvTestParams { + // Input: [N, C, H, W] + int64_t N, C, H, W; + // Weight: [M, C/group, kH, kW] + int64_t M, kH, kW; + // Conv attributes. pad_h/pad_w are the leading (begin) pads. + int64_t pad_h, pad_w; + int64_t stride_h, stride_w; + int64_t dilation_h, dilation_w; + int64_t group; + bool has_bias; + // Trailing (end) pads. Negative means "same as the leading pad", i.e. symmetric + // padding; set them explicitly to build an asymmetrically padded Conv. + int64_t pad_h_end = -1; + int64_t pad_w_end = -1; +}; + +template +void RunConvTest(const ConvTestParams& p) { + static_assert(std::is_same_v || std::is_same_v, "unexpected type for T"); + + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + + const int64_t C_per_group = p.C / p.group; + const int64_t pad_h_end = p.pad_h_end < 0 ? p.pad_h : p.pad_h_end; + const int64_t pad_w_end = p.pad_w_end < 0 ? p.pad_w : p.pad_w_end; + const int64_t oH = ConvOutputSize(p.H, p.kH, p.pad_h, pad_h_end, p.stride_h, p.dilation_h); + const int64_t oW = ConvOutputSize(p.W, p.kW, p.pad_w, pad_w_end, p.stride_w, p.dilation_w); + + std::vector x_dims = {p.N, p.C, p.H, p.W}; + std::vector w_dims = {p.M, C_per_group, p.kH, p.kW}; + std::vector b_dims = p.has_bias ? std::vector{p.M} : std::vector{}; + std::vector y_dims = {p.N, p.M, oH, oW}; + + RandomValueGenerator random{1234}; + std::vector x_vals(random.Gaussian(x_dims, 0.0f, 0.25f)); + std::vector w_vals(random.Gaussian(w_dims, 0.0f, 0.25f)); + std::vector b_vals; + if (p.has_bias) { + b_vals = random.Gaussian(b_dims, 0.0f, 0.1f); + } + + std::vector expected_vals(static_cast(p.N * p.M * oH * oW), 0.0f); + ComputeExpectedConv2D(x_vals, w_vals, b_vals, expected_vals, + p.N, p.C, p.H, p.W, p.M, p.kH, p.kW, oH, oW, + p.pad_h, p.pad_w, p.stride_h, p.stride_w, + p.dilation_h, p.dilation_w, p.group); + + OpTester test("Conv", version); + test.AddAttribute("group", p.group); + test.AddAttribute("kernel_shape", std::vector{p.kH, p.kW}); + test.AddAttribute("pads", std::vector{p.pad_h, p.pad_w, pad_h_end, pad_w_end}); + test.AddAttribute("strides", std::vector{p.stride_h, p.stride_w}); + test.AddAttribute("dilations", std::vector{p.dilation_h, p.dilation_w}); + + if constexpr (std::is_same_v) { + test.AddInput("X", x_dims, x_vals); + test.AddInput("W", w_dims, w_vals, true /*is_initializer*/); + if (p.has_bias) + test.AddInput("B", b_dims, b_vals); + test.AddOutput("Y", y_dims, expected_vals); + } else { + test.AddInput("X", x_dims, FloatsToMLFloat16s(x_vals)); + test.AddInput("W", w_dims, FloatsToMLFloat16s(w_vals), true /*is_initializer*/); + if (p.has_bias) + test.AddInput("B", b_dims, FloatsToMLFloat16s(b_vals)); + test.AddOutput("Y", y_dims, FloatsToMLFloat16s(expected_vals)); + test.SetOutputAbsErr("Y", 0.06f); + test.SetOutputRelErr("Y", 0.02f); + } + + // WebGPU only: the reference above is computed on the host, and an MLFloat16 Conv has + // no CPU kernel on most platforms, so running the other EPs would only add noise. + test.ConfigEp(std::move(webgpu_ep)).RunWithConfig(); +} + +// Small enabled coverage for the Conv MatMul routes, including the subgroup-matrix +// 1x1 Conv path. It only engages when the device reports the f16 8x16x16 +// subgroup-matrix config; elsewhere these shapes take the generic path and the same +// assertions still apply. f16 with a constant weight, which is what the subgroup path +// requires (a float weight or a non-constant weight falls back). +// +// Input channels are the matmul K and output channels are the matmul N, so: +// - C=32: K is a multiple of the subgroup-matrix K (16), the path is eligible. +// - M=16: even N, the plain path. +// - M=17: odd N, which the f16 subgroup-matrix load cannot stride over. The constant +// weight is padded to an even stride once and cached, and the output is still +// written at the real, odd N. +// - C=24: K is not a multiple of 16 and the kernel has no K tail handling, so the +// tiling selector must decline and let the generic path compute it. +TEST(Conv_Large, Conv1x1_MatMulRoute_F16) { + RunConvTest({1, 32, 8, 8, 16, 1, 1, 0, 0, 1, 1, 1, 1, 1, true}); + RunConvTest({1, 32, 8, 8, 17, 1, 1, 0, 0, 1, 1, 1, 1, 1, true}); + RunConvTest({1, 24, 8, 8, 16, 1, 1, 0, 0, 1, 1, 1, 1, 1, true}); + // Batch > 1: the 1x1 reshape folds N,H,W into independent matmul rows. + RunConvTest({2, 32, 8, 8, 17, 1, 1, 0, 0, 1, 1, 1, 1, 1, true}); +} + +// A kernel that covers the whole input (the "same size" reshape): the window folds into +// a single matmul row per batch element, so the output is 1x1 spatially and the matmul K +// is H*W*C (256 here, a multiple of 16). Odd output channels again exercise the padded +// constant weight, and the 2-row matmul exercises the kernel's partial-M stores. +TEST(Conv_Large, ConvSameSize_MatMulRoute_F16) { + RunConvTest({2, 16, 4, 4, 16, 4, 4, 0, 0, 1, 1, 1, 1, 1, true}); + RunConvTest({2, 16, 4, 4, 17, 4, 4, 0, 0, 1, 1, 1, 1, 1, true}); +} + +// Trailing-only padding. Neither reshape is equivalent to the Conv here: the output grid +// is larger than what the matmul computes (1x1: (H+1) x (W+1) instead of H x W; +// same-size: 2x2 instead of 1x1), so both routes must decline and fall back to the +// general Conv path. A gate that only looks at the leading pads would accept these and +// leave part of the output unwritten. +TEST(Conv_Large, ConvMatMulRoute_TrailingPadDeclined_F16) { + ConvTestParams conv_1x1{1, 32, 8, 8, 16, 1, 1, 0, 0, 1, 1, 1, 1, 1, true}; + conv_1x1.pad_h_end = 1; + conv_1x1.pad_w_end = 1; + RunConvTest(conv_1x1); + + ConvTestParams conv_same_size{1, 16, 4, 4, 16, 4, 4, 0, 0, 1, 1, 1, 1, 1, true}; + conv_same_size.pad_h_end = 1; + conv_same_size.pad_w_end = 1; + RunConvTest(conv_same_size); +} + +} // namespace test +} // namespace onnxruntime