diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc index aa17baeb2bae1..3b9e86fa75416 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc @@ -15,10 +15,10 @@ #include "core/providers/webgpu/compute_context.h" #include "core/providers/webgpu/math/matmul.h" #include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/math/subgroup_matrix_pad_b.h" #include "core/providers/webgpu/program.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 { @@ -53,24 +53,6 @@ class SubgroupMatrixMatMulProgram final : public Program= N), zero-filling columns [N, N_b). Gives B an aligned row stride so the -// subgroup-matrix f16 load's row-start alignment holds for any 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 8x16x16 and 16x16x16 @@ -264,45 +246,18 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl { } private: - // Lazily builds an aligned copy of a constant weight B [..., K, N] with unaligned N - // by widening its last dim to an aligned N_b (zero-filling the extra columns) 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. + // Lazily builds an aligned copy of a constant weight B with unaligned N, caching + // it so the per-run pad cost is paid once. 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, const TensorShape& b_shape, uint32_t N, uint32_t n_b) const { - 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 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, b_shape, 1}) - .AddOutput({padded.get(), ProgramTensorMetadataDependency::TypeAndRank, padded->Shape(), 1}) - .AddUniformVariables({{output_size}, {N}, {n_b}}); - s = context.RunProgram(program); - } - if (s.IsOK()) { + std::unique_ptr padded; + if (PadSubgroupMatrixWeightColumns(context, b, b_shape, N, n_b, padded).IsOK()) { padded_b_ = std::move(padded); } }); diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template index 8932768357e2b..09c2df0d12949 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template @@ -4,7 +4,8 @@ // 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). One thread per padded-output element. // Gives B an aligned row stride so the subgroup-matrix f16 load's row-start -// alignment holds for any N. See EnsurePaddedB in subgroup_matrix_matmul.cc. +// alignment holds for any N. Shared by MatMul and Conv; see +// PadSubgroupMatrixWeightColumns in subgroup_matrix_pad_b.cc. #use guardAgainstOutOfBoundsWorkgroupSizes #use .getByOffset .setByOffset diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.cc b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.cc new file mode 100644 index 0000000000000..216a91ee8fcd9 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.cc @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/math/subgroup_matrix_pad_b.h" + +#include + +#include "core/common/narrow.h" +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_utils.h" + +namespace onnxruntime { +namespace webgpu { + +namespace { + +class SubgroupMatrixPadBProgram final : public Program { + public: + SubgroupMatrixPadBProgram() : Program{"SubgroupMatrixPadB"} {} + 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}); +}; + +} // namespace + +Status PadSubgroupMatrixWeightColumns(ComputeContext& context, const Tensor& b, const TensorShape& b_shape, + uint32_t N, uint32_t n_b, std::unique_ptr& padded) { + 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 the subgroup-matrix weight because the padded tensor has ", output_size_i64, + " elements, exceeding uint32_t shader indexing range."); + const uint32_t output_size = narrow(output_size_i64); + + padded = std::make_unique(context.CreateGPUTensor(b.DataType(), padded_shape)); + // A zero-element padded tensor needs no pad pass -- dispatching 0 workgroups is + // pointless and some drivers reject it. + if (output_size == 0) { + return Status::OK(); + } + SubgroupMatrixPadBProgram program; + program.SetWorkgroupSize(WORKGROUP_SIZE) + .SetDispatchGroupSize(CeilDiv(output_size, WORKGROUP_SIZE)) + .AddInput({&b, ProgramTensorMetadataDependency::TypeAndRank, b_shape, 1}) + .AddOutput({padded.get(), ProgramTensorMetadataDependency::TypeAndRank, padded->Shape(), 1}) + .AddUniformVariables({{output_size}, {N}, {n_b}}); + return context.RunProgram(program); +} + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.h b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.h new file mode 100644 index 0000000000000..f7fa895d47246 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.h @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/framework/tensor.h" +#include "core/framework/tensor_shape.h" +#include "core/providers/webgpu/compute_context.h" + +namespace onnxruntime { +namespace webgpu { + +// Pads a row-major f16 weight B [..., K, N] along its last dim to N_b (>= N), +// zero-filling the new columns, giving B an aligned row stride/offset for a +// subgroup-matrix row-major load. Treats B as a flat [rows, N] -> [rows, N_b] copy +// with rows = numel / N, so a 2D [K, N] weight and a batched [batch, K, N] one (or, +// for Conv, an HWIO [kh, kw, Cin, Cout] one) all work. Always runs the GPU pad +// pass; callers needing to pay this once for a constant weight should cache the +// result themselves (e.g. via std::call_once), as subgroup_matrix_matmul.cc does. +Status PadSubgroupMatrixWeightColumns(ComputeContext& context, const Tensor& b, const TensorShape& b_shape, + uint32_t N, uint32_t n_b, std::unique_ptr& padded); + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index 6e7fb6f80a4d1..e7d0b9837f8f8 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -1,9 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. + #include "core/providers/webgpu/nn/conv.h" #include "core/providers/webgpu/nn/conv2d_mm.h" #include "core/providers/webgpu/nn/conv3d_naive.h" #include "core/providers/webgpu/nn/im2col_matmul.h" +#include "core/providers/webgpu/nn/subgroup_matrix_conv.h" #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" #include "core/providers/webgpu/tensor/transpose.h" @@ -26,31 +28,15 @@ Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const Tens return Transpose::DoTranspose(context, perm, reshaped_kernel, *transposed_kernel); } -template -Status Conv::ComputeInternal(ComputeContext& context) const { - bool has_bias = context.InputCount() > 2; - const auto* input = context.Input(0); - const Tensor* kernel = prepacked_kernel_ ? prepacked_kernel_.get() : context.Input(1); - const auto* bias = has_bias ? context.Input(2) : nullptr; - TensorShape input_shape = input->Shape(); - ORT_ENFORCE(kernel != nullptr, "Conv kernel tensor is required."); - // Prepacked kernels are stored permuted; recover the logical OIHW shape. - TensorShape kernel_shape = kernel->Shape(); - switch (kernel_layout_) { - case KernelLayout::OIHW: - break; - case KernelLayout::HWIO: - kernel_shape = TensorShape(TensorShapeVector{kernel_shape[3], kernel_shape[2], kernel_shape[0], kernel_shape[1]}); - break; - case KernelLayout::OHWI: - kernel_shape = TensorShape(TensorShapeVector{kernel_shape[0], kernel_shape[3], kernel_shape[1], kernel_shape[2]}); - break; - } - 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()); +Status ComputeConvGeometry(const ConvAttributes& conv_attrs, const TensorShape& input_shape, + const TensorShape& kernel_shape, bool is_channels_last, + std::vector& pads, std::vector& strides, + std::vector& dilations, TensorShape& output_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)); + 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); } @@ -63,23 +49,80 @@ Status Conv::ComputeInternal(ComputeContext& context TensorShapeVector input_shape_vector = input_shape.AsShapeVector(); auto batch = input_shape[0]; TensorShapeVector output_shape_vector = {batch}; - TensorShape input_spacial_shape = is_channels_last ? TensorShape(TensorShapeVector(std::next(input_shape_vector.begin()), std::prev(input_shape_vector.end()))) : input_shape.Slice(2); - ORT_RETURN_IF_ERROR(conv_attrs_.InferPadsAndOutputShape(input_spacial_shape, kernel_spacial_shape_vector, local_strides, local_dilations, local_pads, output_shape_vector)); + TensorShape input_spacial_shape = is_channels_last + ? TensorShape(TensorShapeVector(std::next(input_shape_vector.begin()), std::prev(input_shape_vector.end()))) + : input_shape.Slice(2); + ORT_RETURN_IF_ERROR(conv_attrs.InferPadsAndOutputShape(input_spacial_shape, kernel_spacial_shape_vector, + local_strides, local_dilations, local_pads, + output_shape_vector)); auto output_channels = kernel_shape[0]; if (is_channels_last) { output_shape_vector.push_back(output_channels); } else { output_shape_vector.insert(output_shape_vector.begin() + 1, output_channels); } - auto output_shape = TensorShape(output_shape_vector); - auto* output = context.Output(0, output_shape); - std::vector strides; - std::vector pads; - std::vector dilations; + output_shape = TensorShape(output_shape_vector); + auto transform_dim = [](int64_t dim) { return static_cast(dim); }; + pads.clear(); + strides.clear(); + dilations.clear(); std::transform(local_pads.begin(), local_pads.end(), std::back_inserter(pads), transform_dim); std::transform(local_strides.begin(), local_strides.end(), std::back_inserter(strides), transform_dim); std::transform(local_dilations.begin(), local_dilations.end(), std::back_inserter(dilations), transform_dim); + return Status::OK(); +} + +template +typename Conv::ConvOptImpl* +Conv::ConvOptImplCache::GetOrCreate( + const Conv& parent, const ComputeContextBase& context) { + std::call_once(init_flag_, [&]() { + impl_ = CreateSubgroupMatrixConvImpl(parent, context); + }); + return impl_.get(); +} + +template +Status Conv::ComputeInternal(ComputeContext& context) const { + // Preferred fast path on devices with a vendor subgroup-matrix tiling policy; + // declines (handled=false) when the kernel/attrs don't qualify. + if (ConvOptImpl* opt_impl = opt_impl_cache_.GetOrCreate(*this, context); opt_impl != nullptr) { + bool handled = false; + ORT_RETURN_IF_ERROR(opt_impl->Compute(context, handled)); + if (handled) { + return Status::OK(); + } + } + + bool has_bias = context.InputCount() > 2; + const auto* input = context.Input(0); + const Tensor* kernel = prepacked_kernel_ ? prepacked_kernel_.get() : context.Input(1); + const auto* bias = has_bias ? context.Input(2) : nullptr; + TensorShape input_shape = input->Shape(); + ORT_ENFORCE(kernel != nullptr, "Conv kernel tensor is required."); + // Prepacked kernels are stored permuted; recover the logical OIHW shape. + TensorShape kernel_shape = kernel->Shape(); + switch (kernel_layout_) { + case KernelLayout::OIHW: + break; + case KernelLayout::HWIO: + kernel_shape = TensorShape(TensorShapeVector{kernel_shape[3], kernel_shape[2], kernel_shape[0], kernel_shape[1]}); + break; + case KernelLayout::OHWI: + kernel_shape = TensorShape(TensorShapeVector{kernel_shape[0], kernel_shape[3], kernel_shape[1], kernel_shape[2]}); + break; + } + std::vector pads; + std::vector strides; + std::vector dilations; + TensorShape output_shape; + ORT_RETURN_IF_ERROR(ComputeConvGeometry(conv_attrs_, input_shape, kernel_shape, is_channels_last, + pads, strides, dilations, output_shape)); + auto* output = context.Output(0, output_shape); + auto batch = input_shape[0]; + auto output_channels = kernel_shape[0]; + TensorShapeVector output_shape_vector = output_shape.AsShapeVector(); auto rank = input_shape.NumDimensions(); const InlinedVector perm = {2, 3, 1, 0}; if (rank > 5) { @@ -125,6 +168,7 @@ Status Conv::ComputeInternal(ComputeContext& context } else if (rank == 3) { // Conv1D TensorShapeVector kernel_shape_vector = kernel_shape.AsShapeVector(); + TensorShapeVector input_shape_vector = input_shape.AsShapeVector(); input_shape_vector.insert(input_shape_vector.begin() + (is_channels_last ? 1 : 2), 1, 1); output_shape_vector.insert(output_shape_vector.begin() + (is_channels_last ? 1 : 2), 1, 1); kernel_shape_vector.insert(kernel_shape_vector.begin() + 2, 1); @@ -323,6 +367,7 @@ Status Conv::PrePackInternal(ComputeContextBase& con } // Im2ColMatMul path: transpose OIHW -> OHWI once here instead of on every inference. + // Subgroup-matrix Conv (tried first in ComputeInternal) shares this same prepack. // // Placed before the auto_pad check below on purpose: // - Safe: CanApplyIm2ColMatMulProgram() only looks at the adapter, dtype, layout, @@ -332,10 +377,10 @@ Status Conv::PrePackInternal(ComputeContextBase& con // - Necessary: otherwise every auto_pad != NOTSET model would bail out below and // keep paying for the transpose on every inference. // - // This call and the one in ComputeInternal must stay in agreement: only the im2col - // path can read the OHWI layout, so a decision made here that ComputeInternal later - // reverses would corrupt the weights. If CanApplyIm2ColMatMulProgram() ever gains a - // condition that is not known at prepack time (pads, strides, input shape), this + // This call and the one in ComputeInternal must stay in agreement: only im2col and + // subgroup-matrix Conv can read OHWI, so a decision made here that ComputeInternal + // later reverses would corrupt the weights. If CanApplyIm2ColMatMulProgram() ever + // gains a condition not known at prepack time (pads, strides, input shape), this // shortcut must go away. ComputeInternal ORT_ENFORCEs the invariant. if (CanApplyIm2ColMatMulProgram(context, is_channels_last, activation_, kernel_shape, onnxruntime::narrow(conv_attrs_.group), @@ -367,6 +412,8 @@ Status Conv::PrePackInternal(ComputeContextBase& con // Analyze execution paths in ComputeInternal to determine if kernel transpose is needed: // + // 0. Subgroup-matrix Conv path: no dedicated branch -- reads whatever layout falls + // out below (OIHW, OHWI from Im2ColMatMul, or HWIO from the general transpose). // 1. Im2ColMatMul path: handled above (prepacked as OHWI) // 2. Grouped conv (group > 1): handled above (skip if !is_channels_last) // 3. MatMul optimization (same_size || is_1x1_conv): diff --git a/onnxruntime/core/providers/webgpu/nn/conv.h b/onnxruntime/core/providers/webgpu/nn/conv.h index af98fa76e30ca..713661cc072f6 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -3,6 +3,10 @@ #pragma once +#include +#include +#include + #include "core/common/optional.h" #include "core/providers/webgpu/webgpu_kernel.h" #include "core/providers/cpu/nn/conv_attributes.h" @@ -27,11 +31,28 @@ enum class KernelLayout { template class Conv : public WebGpuKernel { public: + // Abstract base for alternative optimized Conv implementations (mirrors + // Gemm::GemmOptImpl); currently just the subgroup-matrix path. + class ConvOptImpl { + public: + explicit ConvOptImpl(const Conv& parent) : parent_(parent) {} + virtual ~ConvOptImpl() = default; + + // Sets handled=true if it ran, else leaves handled=false (allocating nothing). + // Runs before ComputeInternal resolves auto_pad/output. + virtual Status Compute(ComputeContext& context, /*out*/ bool& handled) = 0; + + protected: + const Conv& parent_; + }; + Conv(const OpKernelInfo& info) : WebGpuKernel(info), conv_attrs_(info) { if (is_fused) { // Throw the Status, not the expression: the message names the activation ORT_ENFORCE would drop. ORT_THROW_IF_ERROR(GetFusedActivationAttr(info, activation_)); } + const Tensor* w = nullptr; + w_is_constant_ = info.TryGetConstantInput(1, &w); } Status ComputeInternal(ComputeContext& context) const override; @@ -41,6 +62,14 @@ class Conv : public WebGpuKernel { AllocatorPtr alloc, /*out*/ bool& is_packed) override; + const ConvAttributes& ConvAttrs() const { return conv_attrs_; } + const Activation& ConvActivation() const { return activation_; } + bool IsWeightConstant() const { return w_is_constant_; } + // Prepacked weight, or null if not prepacked (layout given by kernel_layout_). + const Tensor* PrepackedKernel() const { return prepacked_kernel_.get(); } + // Layout of PrepackedKernel(), or OIHW while it is null (input 1 is read directly). + KernelLayout GetKernelLayout() const { return kernel_layout_; } + protected: ConvAttributes conv_attrs_; Activation activation_; @@ -51,9 +80,37 @@ class Conv : public WebGpuKernel { // is set, otherwise input 1. Stays `OIHW` while `prepacked_kernel_` is null. KernelLayout kernel_layout_{KernelLayout::OIHW}; mutable MatMulComputeDispatcher matmul_compute_dispatcher_; + + bool w_is_constant_ = false; // whether the weight input (index 1) is a constant initializer + + private: + // Owns the optimized implementation, built lazily (needs a compute context) and + // guarded by call_once against concurrent Compute calls. Null after init means no + // optimized path on this device. + class ConvOptImplCache { + public: + ConvOptImplCache() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ConvOptImplCache); + + // Defined in conv.cc: it calls into subgroup_matrix_conv.h, which includes this + // header, so the include can't go the other way. + ConvOptImpl* GetOrCreate(const Conv& parent, const ComputeContextBase& context); + + private: + std::once_flag init_flag_; + std::unique_ptr impl_; + }; + mutable ConvOptImplCache opt_impl_cache_; }; Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const TensorShape& kernel_shape, Tensor* transposed_kernel, const InlinedVector& perm); +// Resolves pads/strides/dilations and the output shape. Shared by ComputeInternal and +// the subgroup-matrix opt path so the two can't diverge on auto_pad handling. +Status ComputeConvGeometry(const ConvAttributes& conv_attrs, const TensorShape& input_shape, + const TensorShape& kernel_shape, bool is_channels_last, + /*out*/ std::vector& pads, std::vector& strides, + std::vector& dilations, TensorShape& output_shape); + } // namespace webgpu } // namespace onnxruntime 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..c8100ebdf19bd --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -0,0 +1,454 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#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/math/subgroup_matrix_pad_b.h" +#include "core/providers/webgpu/nn/conv.h" +#include "core/providers/webgpu/nn/subgroup_matrix_conv.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" + +namespace onnxruntime { +namespace webgpu { + +namespace { + +// Workgroup size is kSubgroupMatrixSubgroupSize * split_k. +constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; + +// Max sg_mat_c_ the WGSL template unrolls (mi < 8, ni < 4); a larger +// tiling would be accepted by the selector but silently mis-computed. +constexpr uint32_t kMaxSgMatCountM = 8; +constexpr uint32_t kMaxSgMatCountN = 4; + +// Device's config for the fixed 8x16x16 F16 config this kernel implements, or +// nullopt if unsupported. Single source of truth for eligibility: both +// CanApplySubgroupMatrixConv and CreateSubgroupMatrixConvImpl call this. +std::optional GetSupportedConfig(const ComputeContextBase& context) { + constexpr auto kF16 = wgpu::SubgroupMatrixComponentType::F16; + const std::optional config = + SelectSubgroupMatrixConfig(context, {{kF16, kF16, 8, 16, 16, 32, false}}); + if (!config || !config->Is(8, 16, 16)) { + return std::nullopt; + } + // Adapters here may expose a 16-32 subgroup size range, but the kernel hard-codes + // 32 (sg_index = local_idx / 32); pin it rather than risk a wasteful driver choice. + if (!context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + return std::nullopt; + } + if (context.AdapterInfo().vendor != std::string_view{"intel"}) { + return std::nullopt; + } + return config; +} + +Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelper& input_src, + uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, + uint32_t split_k, uint32_t vec_size, bool b_is_hwio) { + return WGSL_TEMPLATE_APPLY(shader, "nn/subgroup_matrix_conv_8x16x16.wgsl.template", + WGSL_TEMPLATE_PARAMETER(b_is_hwio, b_is_hwio), + WGSL_TEMPLATE_PARAMETER(sg_mat_count_m, sg_mat_count_m), + WGSL_TEMPLATE_PARAMETER(sg_mat_count_n, sg_mat_count_n), + WGSL_TEMPLATE_PARAMETER(sg_mat_k, 16), + WGSL_TEMPLATE_PARAMETER(sg_mat_m, 8), + WGSL_TEMPLATE_PARAMETER(sg_mat_n, 16), + WGSL_TEMPLATE_PARAMETER(split_k, split_k), + WGSL_TEMPLATE_PARAMETER(vec_size, vec_size), + WGSL_TEMPLATE_VARIABLE(input_src, input_src)); +} + +// Computes a channels-last 2D convolution as an implicit GEMM via +// subgroupMatrixMultiplyAccumulate (only 8x16x16 implemented so far). Runs its own +// Conv shape inference, since Compute runs before auto_pad/output are resolved. +template +class SubgroupMatrixConvImpl final : public Conv::ConvOptImpl { + using ConvT = Conv; + using Base = typename ConvT::ConvOptImpl; + + public: + SubgroupMatrixConvImpl(const ConvT& parent, SubgroupMatrixConfig config, + SubgroupMatrixTilingSelector tiling_selector) + : Base(parent), config_(config), 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(); + + // Accept whatever layout PrePackInternal left the weight in (OIHW, or OHWI/HWIO + // shared with another path) and recover the logical OIHW shape from it. + const KernelLayout kernel_layout = parent.GetKernelLayout(); + const Tensor* weight = parent.PrepackedKernel() != nullptr ? parent.PrepackedKernel() : context.Input(1); + const auto* input = context.Input(0); + if (input == nullptr || weight == nullptr) { + return Status::OK(); + } + TensorShape weight_shape = weight->Shape(); + switch (kernel_layout) { + case KernelLayout::OIHW: + break; + case KernelLayout::HWIO: + weight_shape = TensorShape( + TensorShapeVector{weight_shape[3], weight_shape[2], weight_shape[0], weight_shape[1]}); + break; + case KernelLayout::OHWI: + weight_shape = TensorShape( + TensorShapeVector{weight_shape[0], weight_shape[3], weight_shape[1], weight_shape[2]}); + break; + } + if (!CanApplySubgroupMatrixConv(context, is_channels_last, weight_shape, + narrow(conv_attrs.group), weight->DataType())) { + return Status::OK(); + } + + // Only a rank-4 NHWC activation is served (reshaped Conv1D's 1xW kernel is + // already rejected by CanApplySubgroupMatrixConv). + const TensorShape& input_shape = input->Shape(); // NHWC: [batch, input_h, input_w, Cin] + if (input_shape.NumDimensions() != 4) { + return Status::OK(); + } + std::vector pads; + std::vector strides; + std::vector dilations; + TensorShape output_shape; // NHWC: [batch, out_h, out_w, Cout] + ORT_RETURN_IF_ERROR(ComputeConvGeometry(conv_attrs, input_shape, weight_shape, is_channels_last, + pads, strides, dilations, output_shape)); + if (dilations.size() < 2 || pads.size() < 2 || strides.size() < 2) { + return Status::OK(); + } + + const uint32_t batch = narrow(input_shape[0]); + const uint32_t input_h = narrow(input_shape[1]); + const uint32_t input_w = narrow(input_shape[2]); + const uint32_t cin = narrow(weight_shape[1]); + const uint32_t kernel_h = narrow(weight_shape[2]); + const uint32_t kernel_w = narrow(weight_shape[3]); + const uint32_t output_h = narrow(output_shape[1]); + const uint32_t output_w = narrow(output_shape[2]); + + const uint32_t M = output_h * output_w; // output pixels per batch slice + const uint32_t K = kernel_h * kernel_w * cin; // contraction dim + const uint32_t N = narrow(weight_shape[0]); // Cout + + // Empty batch/M/N (ONNX permits zero-length dims): let the generic Conv paths + // handle it instead of dispatching a zero-workgroup kernel. + if (batch == 0 || M == 0 || N == 0) { + return Status::OK(); + } + + const std::optional tiling = tiling_selector_(context, M, N, K, batch); + if (!tiling) { + return Status::OK(); + } + + // CapTileNToAlignedN already guarantees tile_n <= N; mirrored here as the last + // possible decline point before output allocation/weight transpose commit below. + if (N < tiling->tile_n) { + return Status::OK(); + } + + // HWIO's row-major B has row stride N, with no alignment guarantee (unlike OHWI's + // stride K, always a multiple of 16); an odd N breaks the subgroup-matrix load's + // row-start alignment. A constant weight can pad N up to N_b and cache the + // result; a non-constant one can't be cached, so decline instead. + constexpr uint32_t kHwioNAlign = 2; + const bool needs_padded_b = kernel_layout == KernelLayout::HWIO && N % kHwioNAlign != 0; + if (needs_padded_b && !parent.IsWeightConstant()) { + return Status::OK(); + } + const uint32_t n_b = needs_padded_b ? ((N / kHwioNAlign) + 1) * kHwioNAlign : N; + + // subgroupMatrixLoad addresses the weight as a single binding; a weight past + // maxStorageBufferBindingSize gets split by ProgramManager, which this offset + // arithmetic doesn't know about. Checked against the bound size (K x N_b). + const uint64_t max_weight_elems = context.DeviceLimits().maxStorageBufferBindingSize / sizeof(MLFloat16); + if (static_cast(K) * n_b > max_weight_elems) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv declined: weight exceeds maxStorageBufferBindingSize (K=" << K + << " N_b=" << n_b << ", limit_elems=" << max_weight_elems << ")."; + return Status::OK(); + } + + Tensor* output = context.Output(0, output_shape); + + // B operand: OHWI or HWIO (see b_is_hwio and the WGSL template). An already + // prepacked weight is used as-is (padded to n_b first if its N is unaligned); + // otherwise this transposes to OHWI itself, caching the result for constant + // weights. + Tensor run_local_ohwi_weight; + const Tensor* b_weight = nullptr; + bool b_is_hwio = false; + if (parent.PrepackedKernel() != nullptr) { + b_weight = parent.PrepackedKernel(); + b_is_hwio = (kernel_layout == KernelLayout::HWIO); + if (needs_padded_b) { + ORT_RETURN_IF_ERROR(EnsurePaddedWeight(context, *b_weight, N, n_b)); + b_weight = padded_b_weight_.get(); + } + } else if (parent.IsWeightConstant()) { + ORT_RETURN_IF_ERROR(EnsureTransposedKernel(context, weight, weight_shape)); + b_weight = cached_ohwi_weight_.get(); + } else { + ORT_RETURN_IF_ERROR( + TransposeKernel(context, weight, weight_shape, &run_local_ohwi_weight, {0, 2, 3, 1})); + b_weight = &run_local_ohwi_weight; + } + + const auto& config = config_; + const uint32_t tile_m = tiling->tile_m; + const uint32_t tile_n = tiling->tile_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 sg_mat_count_m = tile_m / config.M; + const uint32_t sg_mat_count_n = tile_n / config.N; + ORT_ENFORCE(sg_mat_count_m <= kMaxSgMatCountM && sg_mat_count_n <= kMaxSgMatCountN, + "Tiling exceeds the subgroup-matrix counts the kernel unrolls: ", + tile_m, "x", tile_n, " spans ", sg_mat_count_m, "x", sg_mat_count_n, + " subgroup matrices, at most ", kMaxSgMatCountM, "x", kMaxSgMatCountN); + + const uint32_t dispatch_x = (N + tile_n - 1) / tile_n; + const uint32_t dispatch_y = (M + tile_m - 1) / tile_m; + + // Vectorize the im2col gather along cin (contiguous in NHWC). + const uint32_t vec_size = cin % 4 == 0 ? 4 : (cin % 2 == 0 ? 2 : 1); + + const Tensor* bias = context.InputCount() > 2 ? context.Input(2) : nullptr; + const bool has_bias = bias != nullptr; + const uint32_t split_k = tiling->split_k; + SubgroupMatrixConvProgram program{has_bias, config, sg_mat_count_m, sg_mat_count_n, split_k, vec_size, + b_is_hwio, activation}; + program.SetWorkgroupSize(kSubgroupMatrixSubgroupSize * split_k); + program.SetSubgroupSize(kSubgroupMatrixSubgroupSize); + program.SetDispatchGroupSize(dispatch_x, dispatch_y, batch); + program.CacheHint(has_bias, static_cast(config.componentType), + static_cast(config.resultComponentType), + config.M, config.N, config.K, config.subgroupSize, + sg_mat_count_m, sg_mat_count_n, split_k, vec_size, b_is_hwio, + activation.CacheKey()) + .AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank, static_cast(vec_size)}, + {b_weight, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddOutput({output, ProgramTensorMetadataDependency::Rank, output_shape, 1}) + .AddUniformVariables({{M}, + {N}, + {n_b}, + {K}, + {dispatch_x}, + {batch}, + {input_h}, + {input_w}, + {cin}, + {kernel_h}, + {kernel_w}, + {output_w}, + {dilations}, + {pads}, + {strides}}); + // Activation uniforms must remain last because definitions and values are matched by index. + AppendActivationUniformsData(activation, program); + if (has_bias) { + program.AddInput({bias, ProgramTensorMetadataDependency::None}); + } + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv dispatching: M=" << M << " N=" << N << " N_b=" << n_b + << " K=" << K << " batch=" << batch << " tile_m=" << tile_m << " tile_n=" << tile_n + << " split_k=" << split_k << " vec_size=" << vec_size << "."; + ORT_RETURN_IF_ERROR(context.RunProgram(program)); + + handled = true; + return Status::OK(); + } + + private: + // Builds the OIHW -> OHWI transpose on first use and caches it; call_once guards + // against concurrent Compute calls on the same (shared) impl. + Status EnsureTransposedKernel(ComputeContext& context, const Tensor* weight, + const TensorShape& weight_shape) const { + std::call_once(transpose_once_, [&]() { + auto transposed = std::make_unique(); + if (TransposeKernel(context, weight, weight_shape, transposed.get(), {0, 2, 3, 1}).IsOK()) { + cached_ohwi_weight_ = std::move(transposed); + } + }); + return cached_ohwi_weight_ ? Status::OK() + : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Failed to transpose the subgroup-matrix Conv weight to OHWI."); + } + + // Pads a prepacked HWIO weight's unaligned N up to n_b on first use and caches it; + // only reached for a constant weight (see needs_padded_b above). + Status EnsurePaddedWeight(ComputeContext& context, const Tensor& b_weight, uint32_t N, uint32_t n_b) const { + std::call_once(pad_once_, [&]() { + std::unique_ptr padded; + if (PadSubgroupMatrixWeightColumns(context, b_weight, b_weight.Shape(), N, n_b, padded).IsOK()) { + padded_b_weight_ = std::move(padded); + } + }); + return padded_b_weight_ ? Status::OK() + : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Failed to pad the subgroup-matrix Conv weight."); + } + + const SubgroupMatrixConfig config_; + SubgroupMatrixTilingSelector tiling_selector_; + + // OIHW -> OHWI transposed constant weight, built once by EnsureTransposedKernel. + mutable std::once_flag transpose_once_; + mutable std::unique_ptr cached_ohwi_weight_; + + // N-padded prepacked HWIO weight, built once by EnsurePaddedWeight. + mutable std::once_flag pad_once_; + mutable std::unique_ptr padded_b_weight_; +}; + +} // namespace + +Status SubgroupMatrixConvProgram::GenerateShaderCode(ShaderHelper& shader) const { + const auto& input_src = shader.AddInput("input_src", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + shader.AddInput("input_b", ShaderUsage::UseUniform); + if (has_bias_) { + shader.AddInput("bias", ShaderUsage::UseUniform); + } + const auto& output = shader.AddOutput( + "output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + + // Write-out epilogue the template calls per output element: bias add, then the + // fused activation. + auto& additional_implementation = shader.AdditionalImplementation(); + additional_implementation + << GetActivationDeclaration(activation_, "output_value_t", "output_element_t") + << "fn write_output(output_offset: u32, bias_offset: u32, value_in: output_value_t) {\n" + " var value = value_in;\n"; + if (has_bias_) { + additional_implementation << " value += output_value_t(bias[bias_offset]);\n"; + } + additional_implementation + << " " << GetActivationSnippet(activation_, "output_value_t", "output_element_t") << "\n" + << " " << output.SetByOffset("output_offset", "value") << "\n" + << "}\n"; + + const auto& config = config_; + if (config.Is(8, 16, 16)) { + return GenerateShaderCode8x16x16(shader, input_src, sg_mat_count_m_, sg_mat_count_n_, split_k_, vec_size_, + b_is_hwio_); + } + return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, + "Unsupported subgroup matrix config dimensions."); +} + +bool CanApplySubgroupMatrixConv(ComputeContextBase& context, + bool is_channels_last, + const TensorShape& kernel_shape, + uint32_t group, + MLDataType data_type) { + if (data_type != DataTypeImpl::GetType()) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: data_type is not F16."; + return false; + } + if (!is_channels_last || group != 1) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: requires channels_last && group==1, but got" + << " is_channels_last=" << is_channels_last + << " group=" << group << "."; + return false; + } + if (kernel_shape.NumDimensions() != 4) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: kernel is not 4D (rank=" + << kernel_shape.NumDimensions() << ", shape=" << kernel_shape.ToString() << ")."; + return false; + } + + const uint32_t channel_output = narrow(kernel_shape[0]); // Cout (GEMM N) + const uint32_t channel_input = narrow(kernel_shape[1]); // Cin + const uint32_t kernel_height = narrow(kernel_shape[2]); + const uint32_t kernel_width = narrow(kernel_shape[3]); + + // Mirrors the im2col-matmul path: 1xW/Hx1/1x1 kernels (incl. reshaped Conv1D) are + // handled better by the plain MatMul/naive conv paths. + if (kernel_height == 1 || kernel_width == 1) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: 1xW/Hx1/1x1 kernel not handled (kernel_shape[OIHW]=" + << kernel_shape.ToString() << ", kh=" << kernel_height << " kw=" << kernel_width << ")."; + return false; + } + + const uint32_t K = kernel_height * kernel_width * channel_input; + const uint32_t N = channel_output; + // K: the kernel walks it in sg_mat_k blocks with no tail handling, so it must be a + // multiple of 16. N: B is loaded unchecked, shifting a trailing partial tile back + // to N - tile_n, which needs one full 16-wide block to exist; must agree with the + // N guard in CreateSubgroupMatrixConvTilingSelector. + if (K == 0 || K % 16 != 0 || N < 16) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: K/N constraints not met (kernel_shape[OIHW]=" + << kernel_shape.ToString() << ", Cout(N)=" << N << " Cin=" << channel_input + << " kh=" << kernel_height << " kw=" << kernel_width << " K=kh*kw*Cin=" << K + << "). Need K%16==0 (K%16=" << (K % 16) << ") and N>=16."; + return false; + } + + const std::optional config = GetSupportedConfig(context); + if (!config) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: device does not report the required 8x16x16 F16" + << " subgroup-matrix config, lacks subgroup-size control, or vendor != intel." + << " vendor=\"" << std::string_view{context.AdapterInfo().vendor} + << "\" architecture=\"" << std::string_view{context.AdapterInfo().architecture} + << "\" subgroup_size_control=" + << context.HasFeature(wgpu::FeatureName::SubgroupSizeControl) + << " subgroup_size_range=[" << context.AdapterInfo().subgroupMinSize << ", " + << context.AdapterInfo().subgroupMaxSize << "]."; + return false; + } + + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv initial eligibility checks passed (Compute() may still" + << " decline): kernel_shape[OIHW]=" << kernel_shape.ToString() + << " Cout(N)=" << N << " Cin=" << channel_input << " kh=" << kernel_height + << " kw=" << kernel_width << " K=" << K << "."; + return true; +} + +template +std::unique_ptr::ConvOptImpl> CreateSubgroupMatrixConvImpl( + const Conv& parent, + const ComputeContextBase& context) { + if constexpr (!is_channels_last) { + ORT_UNUSED_PARAMETER(parent); + ORT_UNUSED_PARAMETER(context); + return nullptr; + } else { + const std::optional config = GetSupportedConfig(context); + if (!config) { + return nullptr; + } + SubgroupMatrixTilingSelector tiling_selector = intel::CreateSubgroupMatrixConvTilingSelector(context); + if (!tiling_selector) { + return nullptr; + } + return std::make_unique>(parent, *config, + 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 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..2d05c270a7804 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/framework/tensor.h" +#include "core/framework/tensor_shape.h" +#include "core/providers/webgpu/compute_context.h" +#include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/nn/conv.h" +#include "core/providers/webgpu/nn/fuse_utils.h" +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/shader_helper.h" + +namespace onnxruntime { +namespace webgpu { + +// Implicit-GEMM channels-last Conv via subgroupMatrixMultiplyAccumulate. +// sg_mat_count_m/n: subgroup matrices per output tile along M/N. split_k: +// subgroups cooperatively reducing K. +class SubgroupMatrixConvProgram final : public Program { + public: + SubgroupMatrixConvProgram(bool has_bias, SubgroupMatrixConfig config, + uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k, + uint32_t vec_size, bool b_is_hwio, const Activation& activation) + : Program{"SubgroupMatrixConv"}, + has_bias_(has_bias), + config_(config), + sg_mat_count_m_(sg_mat_count_m), + sg_mat_count_n_(sg_mat_count_n), + split_k_(split_k), + vec_size_(vec_size), + b_is_hwio_(b_is_hwio), + activation_(activation) {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"N_b", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, + {"num_n_tile", ProgramUniformVariableDataType::Uint32}, + {"batch", ProgramUniformVariableDataType::Uint32}, + {"input_h", ProgramUniformVariableDataType::Uint32}, + {"input_w", ProgramUniformVariableDataType::Uint32}, + {"cin", ProgramUniformVariableDataType::Uint32}, + {"kernel_h", ProgramUniformVariableDataType::Uint32}, + {"kernel_w", ProgramUniformVariableDataType::Uint32}, + {"output_w", ProgramUniformVariableDataType::Uint32}, + {"dilations", ProgramUniformVariableDataType::Uint32}, + {"pads", ProgramUniformVariableDataType::Uint32}, + {"strides", ProgramUniformVariableDataType::Uint32}, + // Activation uniforms must remain last because definitions and values + // are matched by index (see AppendActivationUniformsData). + WEBGPU_PROGRAM_ACTIVATION_UNIFORM_VARIABLES); + + private: + const bool has_bias_; + const SubgroupMatrixConfig config_; + const uint32_t sg_mat_count_m_; + const uint32_t sg_mat_count_n_; + const uint32_t split_k_; + const uint32_t vec_size_; + // Whether input_b is HWIO rather than OHWI. + const bool b_is_hwio_; + const Activation activation_; +}; + +// Eligibility: fp16, channels-last, group == 1, a 4D non-1x1 kernel, a device +// reporting the 8x16x16 F16 subgroup-matrix config, and K = kernel_h * kernel_w * +// Cin a multiple of 16, N = Cout >= 16. Callers fall back to the default conv +// paths when this returns false. +bool CanApplySubgroupMatrixConv(ComputeContextBase& context, + bool is_channels_last, + const TensorShape& kernel_shape, + uint32_t group, + MLDataType data_type); + +// Nullptr on devices with no vendor tiling policy; callers fall back to normal Conv. +template +std::unique_ptr::ConvOptImpl> CreateSubgroupMatrixConvImpl( + const Conv& parent, + const ComputeContextBase& context); + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template new file mode 100644 index 0000000000000..2e07af49c9558 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template @@ -0,0 +1,561 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Subgroup-matrix Conv kernel (F16), 8x16x16 shape: channels-last 2D convolution +// as an implicit GEMM Y[MxN] = A[MxK] @ B[KxN] (+ optional bias), A the im2col of +// the input (never materialized in global memory). +// M = out_h * out_w, K = kernel_h * kernel_w * Cin, N = Cout +// +// A: gathered per K-block into a_scratch (im2col rows are scattered, but +// subgroupMatrixLoad needs a contiguous buffer); K order (kh, kw, cin) matches +// both weight layouts below. +// B: OHWI, column-major (B[k,n] at n*K+k), or HWIO, row-major, row stride N_b +// (B[k,n] at k*N_b+n) -- see b_is_hwio, set by the host to whichever the weight +// arrived in. N_b pads HWIO's N up to a 4-byte-aligned width when N is odd +// (never needed for OHWI, whose stride K is already a multiple of 16); padded +// columns are read but never written (output always uses real N). +// +// Workgroup: split_k subgroups x 32 lanes, tile kTileM x kTileN (host-selected). K +// blocks round-robin across subgroups when split_k > 1; partials summed at +// write-out. Staging A forces the K reduction into a per-step +// gather/barrier/multiply-accumulate loop in uniform control flow. +// +// Bias/activation: write_output(), built by the host from GetActivationSnippet. +// +// TODO: Honor EnableMatmulFp32Accumulation(): when it is set, accumulate in +// f32 (an F16 x F16 -> F32 subgroup-matrix config, f32 accumulators and scratch) +// instead of the f16 used today. +// +// Host guarantees: K % sg_mat_k == 0; kTileN <= N_b so a trailing partial N tile +// can shift back in bounds. M is unconstrained (see load_im2col_vec). + +#param sg_mat_k +#param sg_mat_m +#param sg_mat_n +#param sg_mat_count_m +#param sg_mat_count_n +#param split_k +#param vec_size +#param b_is_hwio + +#use .getByOffset .setByOffset + +const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +const kSubgroupSize: u32 = 32; // Lanes per subgroup +const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (cin % vec_size == 0) + +// kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// subgroup, summed at write-back. +var scratch: array; + +// Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +var a_scratch: array; + +// Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always +// within one (kh, kw) position, so all vec lanes share one bounds check. Padding +// reads as zero; m past M is index-safe (not necessarily zero) but never stored. +// m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. +fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { + let cin_vec = uniforms.cin / kVecSize; + let oh = m / uniforms.output_w; + let ow = m % uniforms.output_w; + let k_cin = k_vec % cin_vec; + let k_hw = k_vec / cin_vec; + let kw_idx = k_hw % uniforms.kernel_w; + let kh_idx = k_hw / uniforms.kernel_w; + let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); + let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); + if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { + return input_src_value_t(); + } + let input_vec_idx = input_batch_base_vec + + (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; + return input_src.getByOffset(input_vec_idx); +} + +$MAIN { + // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is + // the true linear index used to recover batch_id and (n_tile, m_tile). + let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; + let tiles_per_batch = uniforms.num_n_tile * num_m_tile; + let batch_id = workgroup_idx / tiles_per_batch; + // Dispatch normalization can add extra workgroups past the last batch element; + // drop them here (workgroup-uniform, so the barriers below stay uniform). + if (batch_id >= uniforms.batch) { + return; + } + let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element + let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile + let n_tile_start = n_tile * kTileN; + // Shift a trailing partial N tile back so unchecked B loads stay inside the + // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). + // Uses N_b (== N unless B was padded for row-major alignment) so the shifted + // origin -- itself n_align-even since N_b, kTileN and n_tile_start all are -- + // stays alignment-safe for the HWIO load below. + let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); + let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile + let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile + let m_base = m_tile * kTileM; + let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); + let out_batch_base = batch_id * uniforms.M * uniforms.N; + let k_blocks = uniforms.K / kSgMatK; + let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) + let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) + let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot + let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot + + // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). + var sg_mat_c0_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c0_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c0_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c0_3: subgroup_matrix_result; +#endif +#if sg_mat_count_m >= 2 + var sg_mat_c1_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c1_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c1_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c1_3: subgroup_matrix_result; +#endif +#endif +#if sg_mat_count_m >= 3 + var sg_mat_c2_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c2_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c2_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c2_3: subgroup_matrix_result; +#endif +#endif +#if sg_mat_count_m >= 4 + var sg_mat_c3_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c3_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c3_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c3_3: subgroup_matrix_result; +#endif +#endif +#if sg_mat_count_m >= 5 + var sg_mat_c4_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c4_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c4_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c4_3: subgroup_matrix_result; +#endif +#endif +#if sg_mat_count_m >= 6 + var sg_mat_c5_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c5_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c5_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c5_3: subgroup_matrix_result; +#endif +#endif +#if sg_mat_count_m >= 7 + var sg_mat_c6_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c6_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c6_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c6_3: subgroup_matrix_result; +#endif +#endif +#if sg_mat_count_m >= 8 + var sg_mat_c7_0: subgroup_matrix_result; +#if sg_mat_count_n >= 2 + var sg_mat_c7_1: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 3 + var sg_mat_c7_2: subgroup_matrix_result; +#endif +#if sg_mat_count_n >= 4 + var sg_mat_c7_3: subgroup_matrix_result; +#endif +#endif + + // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, + // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks + // early (valid == false), keeping the barriers below in uniform control flow. + let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; + for (var step: u32 = 0u; step < num_steps; step = step + 1u) { + let kb = sg_index + step * kSplitK; + let valid = kb < k_blocks; + + if (valid) { + let k_vec_per_sgmat_row = kSgMatK / kVecSize; + let k_vec_base = kb * k_vec_per_sgmat_row; + for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { + let row = vidx / k_vec_per_sgmat_row; + let kc_vec = vidx % k_vec_per_sgmat_row; + let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); + let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +#if vec_size == 1 + a_scratch[dst] = chunk; +#else + a_scratch[dst + 0u] = chunk[0]; + a_scratch[dst + 1u] = chunk[1]; +#if vec_size >= 4 + a_scratch[dst + 2u] = chunk[2]; + a_scratch[dst + 3u] = chunk[3]; +#endif +#endif + } + } + // Make the staged A-tiles visible before loading them as subgroup matrices. + workgroupBarrier(); + + if (valid) { + // Load the B right tiles for this K block. +#if b_is_hwio + // HWIO: B is row-major ([K, N_b], B[k, n] at k*N_b + n); load row-major, + // stride N_b. N_b (== N unless padded) keeps the row stride, and via + // global_base_n above, the column offset, 4-byte aligned. + let b_base = kb * kSgMatK * uniforms.N_b + global_base_n; + var sg_mat_b0: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +#if sg_mat_count_n >= 2 + var sg_mat_b1: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &input_b, b_base + 1 * kSgMatN, uniforms.N_b); +#endif +#if sg_mat_count_n >= 3 + var sg_mat_b2: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +#endif +#if sg_mat_count_n >= 4 + var sg_mat_b3: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +#endif +#else + // OHWI: B is column-major ([N, K], B[k, n] at n*K + k); load col-major, stride K. + let b_base = global_base_n * uniforms.K + kb * kSgMatK; + var sg_mat_b0: subgroup_matrix_right = + subgroupMatrixLoad, col_major>( + &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +#if sg_mat_count_n >= 2 + var sg_mat_b1: subgroup_matrix_right = + subgroupMatrixLoad, col_major>( + &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +#endif +#if sg_mat_count_n >= 3 + var sg_mat_b2: subgroup_matrix_right = + subgroupMatrixLoad, col_major>( + &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +#endif +#if sg_mat_count_n >= 4 + var sg_mat_b3: subgroup_matrix_right = + subgroupMatrixLoad, col_major>( + &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +#endif +#endif + + // Load the A left tiles (one per M block) from workgroup staging, + // row-major with stride kSgMatK. + var sg_mat_a0: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +#if sg_mat_count_m >= 2 + var sg_mat_a1: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +#endif +#if sg_mat_count_m >= 3 + var sg_mat_a2: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +#endif +#if sg_mat_count_m >= 4 + var sg_mat_a3: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +#endif +#if sg_mat_count_m >= 5 + var sg_mat_a4: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +#endif +#if sg_mat_count_m >= 6 + var sg_mat_a5: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +#endif +#if sg_mat_count_m >= 7 + var sg_mat_a6: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +#endif +#if sg_mat_count_m >= 8 + var sg_mat_a7: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +#endif + + // Accumulate every (M block, N block) pair. + sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +#if sg_mat_count_n >= 2 + sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +#endif +#if sg_mat_count_m >= 2 + sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +#if sg_mat_count_n >= 2 + sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +#endif +#endif +#if sg_mat_count_m >= 3 + sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +#if sg_mat_count_n >= 2 + sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +#endif +#endif +#if sg_mat_count_m >= 4 + sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +#if sg_mat_count_n >= 2 + sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +#endif +#endif +#if sg_mat_count_m >= 5 + sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +#if sg_mat_count_n >= 2 + sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +#endif +#endif +#if sg_mat_count_m >= 6 + sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +#if sg_mat_count_n >= 2 + sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +#endif +#endif +#if sg_mat_count_m >= 7 + sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +#if sg_mat_count_n >= 2 + sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +#endif +#endif +#if sg_mat_count_m >= 8 + sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +#if sg_mat_count_n >= 2 + sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +#endif +#if sg_mat_count_n >= 3 + sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +#endif +#if sg_mat_count_n >= 4 + sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +#endif +#endif + } + // Ensure every subgroup is done reading a_scratch before it is reused for + // the next K-block step. + workgroupBarrier(); + } + + // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; + // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). + subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +#endif +#if sg_mat_count_m >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +#endif +#endif +#if sg_mat_count_m >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +#endif +#endif +#if sg_mat_count_m >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +#endif +#endif +#if sg_mat_count_m >= 5 + subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +#endif +#endif +#if sg_mat_count_m >= 6 + subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +#endif +#endif +#if sg_mat_count_m >= 7 + subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +#endif +#endif +#if sg_mat_count_m >= 8 + subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +#if sg_mat_count_n >= 2 + subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +#endif +#if sg_mat_count_n >= 3 + subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +#endif +#if sg_mat_count_n >= 4 + subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +#endif +#endif + +#if split_k >= 2 + // Make all subgroups' partial tiles visible before the reduction. + workgroupBarrier(); + + // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each + // position is only overwritten by the lane that already read it. + if (sg_index == 0u) { + for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { + var acc: f16 = f16(0); + for (var s: u32 = 0; s < kSplitK; s++) { + acc += scratch[s * kTileM * kTileN + idx]; + } + scratch[idx] = acc; + } + } + // Publish the summed tile (slot 0) to every subgroup for the write-out pass. + workgroupBarrier(); +#endif + + // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip + // and M's bounds check make non-multiple dimensions write each element once. + let n_count = min(kTileN, uniforms.N - global_base_n); + for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { + let global_m = m_base + r; + if (global_m < uniforms.M) { + let scratch_base = r * kTileN; + let out_base = out_batch_base + global_m * uniforms.N + global_base_n; + for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { + write_output(out_base + i, global_base_n + i, + output_value_t(scratch[scratch_base + i])); + } + } + } +} // MAIN 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 959b54b508b89..e4f7493beb5ea 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 @@ -1,19 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" - #include #include -#include #include #include -#include -#include "core/common/narrow.h" #include "core/providers/webgpu/compute_context.h" #include "core/providers/webgpu/math/subgroup_matrix_tiling.h" #include "core/providers/webgpu/vendor/intel/intel_device_info.h" +#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" // Pretuned tile + split-K table baked into the build; consulted before the // heuristic in SelectTiling. @@ -33,12 +29,18 @@ constexpr uint32_t kSubgroupMatrixK = 16; constexpr uint32_t kTileMCandidates[] = {8, 16, 32, 64}; constexpr uint32_t kTileNCandidates[] = {16, 32, 64}; constexpr uint32_t kSplitKCandidates[] = {1, 2, 4, 8}; -// Scratch holds split_k partial f16 tiles; cap keeps it within the SLM budget. -constexpr uint32_t kMaxScratchElems = 16384; // 32 KB +// Per-tile-row a_scratch footprint, in elements, beyond the output scratch tile: +// 0 for MatMul (no im2col A-tile staging); kSubgroupMatrixK for Conv (width of +// its im2col A-tile staging buffer). +constexpr uint32_t kMatMulAScratchCols = 0; +constexpr uint32_t kConvAScratchCols = kSubgroupMatrixK; + +// Workgroup-memory footprint of a tiling, in f16 elements. +uint64_t ScratchElems(const SubgroupMatrixTiling& t, uint32_t a_scratch_cols) { + return static_cast(t.split_k) * t.tile_m * (t.tile_n + a_scratch_cols); +} -// Hard constraints a tiling must satisfy to run correctly for this problem. -// Used to reject a pretuned entry whose bucket does not fit the actual K. -bool IsTilingValid(const SubgroupMatrixTiling& t, uint32_t K) { +bool IsTilingValid(const SubgroupMatrixTiling& t, uint32_t K, uint32_t a_scratch_cols) { if (t.tile_m == 0 || t.tile_n == 0 || t.tile_m % kSubgroupMatrixM != 0 || t.tile_n % kSubgroupMatrixN != 0) { return false; @@ -47,34 +49,26 @@ bool IsTilingValid(const SubgroupMatrixTiling& t, uint32_t K) { if (t.split_k == 0 || t.split_k > k_blocks) { return false; } - return t.tile_m * t.tile_n * t.split_k <= kMaxScratchElems; + return ScratchElems(t, a_scratch_cols) <= kMaxScratchElems; } -// HwSubgroups returns 0 for an unrecognized arch; fall back to a conservative -// default so the occupancy target is still reasonable. +// Falls back to a conservative default when HwSubgroups doesn't recognize arch. uint32_t EffectiveHwSubgroups(std::string_view arch) { const uint32_t hw = HwSubgroups(arch); return hw == 0 ? 256u : hw; } -// Fallback tiling selection when no pretuned entry applies. The goal is to keep -// the GPU busy without over-subscribing it: pick the tile whose independent -// output-tile grid just fills the resident subgroups, preferring larger tiles -// (more data reuse) among those that qualify, and only using smaller tiles when -// nothing else would fill the machine. Split-K then adds cooperative subgroups -// to reach up to 2x occupancy, subject to the per-subgroup K-work minimum and -// the scratch (SLM) budget. batch scales the independent-tile count: each z -// slice contributes its own output-tile grid, so a larger batch fills the -// machine with bigger tiles and needs less (or no) split-K. +// Fallback MatMul tiling: picks the tile whose independent output-tile grid just +// fills the resident subgroups, then adds split-K up to 2x occupancy. Not shared +// with ConvHeuristicTiling: shrinking tile_n for occupancy is free for MatMul but +// not Conv (see that function's comment). SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { const uint32_t hw = EffectiveHwSubgroups(arch); const uint32_t k_blocks = K / kSubgroupMatrixK; auto tile_count = [](uint32_t dim, uint32_t tile) { return (dim + tile - 1) / tile; }; - // Choose tile_m x tile_n. Skip tiles larger than the dimension (they only waste - // lanes). Prefer the largest-area tile whose grid already fills the hardware - // (num_tiles >= hw); if none does, take the tile that yields the most - // independent tiles (largest area breaks ties). + // Skip tiles larger than the dimension. Prefer the largest-area tile that + // already fills the hardware on its own; else the one with the most tiles. uint32_t tile_m = kTileMCandidates[0]; uint32_t tile_n = kTileNCandidates[0]; uint32_t best_tiles = 0; // most output tiles seen so far (used only while under-filling) @@ -119,18 +113,56 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t for (uint32_t cand : kSplitKCandidates) { if (k_blocks >= cand * kMinBlocksPerSplit && num_tiles * cand <= 2 * hw && - tile_m * tile_n * cand <= kMaxScratchElems) { + ScratchElems({tile_m, tile_n, cand}, kMatMulAScratchCols) <= kMaxScratchElems) { split_k = cand; } } return {tile_m, tile_n, split_k}; } -// Looks up a pretuned tiling for this problem in the baked-in table -// (subgroup_matrix_matmul_tuned.inc). The table holds one sub-table per GPU -// architecture; we match `arch`, then map each problem dimension to the smallest -// grid value >= it (clamped to the largest bucket). Returns nullopt when the -// arch is not covered or no entry matches. +// Conv's sole tiling strategy (a per-architecture tuned tree used to exist but +// didn't generalize). Unlike HeuristicTiling, never shrinks tile_n for occupancy: +// the im2col A tile is re-gathered per N tile, so gather volume (M*N*K/tile_n) +// would multiply -- tile_n is fixed at the largest candidate that fits N instead, +// tile_m at the smallest. split-K likewise has no occupancy cap (it hides gather +// latency independent of occupancy): just the largest candidate with enough work +// per split and scratch headroom. +SubgroupMatrixTiling ConvHeuristicTiling(uint32_t N, uint32_t K) { + const uint32_t k_blocks = K / kSubgroupMatrixK; + + uint32_t tile_n = kTileNCandidates[0]; + for (uint32_t tn : kTileNCandidates) { + if (tn <= N) { + tile_n = tn; + } + } + const uint32_t tile_m = kTileMCandidates[0]; + + constexpr uint32_t kMinBlocksPerSplit = 2; + uint32_t split_k = 1; + for (uint32_t cand : kSplitKCandidates) { + if (k_blocks >= cand * kMinBlocksPerSplit && + ScratchElems({tile_m, tile_n, cand}, kConvAScratchCols) <= kMaxScratchElems) { + split_k = cand; + } + } + return {tile_m, tile_n, split_k}; +} + +// Caps tile_n at the largest multiple of kSubgroupMatrixN that fits in N. +// Correctness, not performance: the Conv kernel's unchecked B load shifts a +// trailing partial N tile back to N - tile_n, which needs tile_n <= N, and N +// itself need not be a multiple of kSubgroupMatrixN. Callers must have already +// rejected N < kSubgroupMatrixN. +void CapTileNToAlignedN(SubgroupMatrixTiling& t, uint32_t N) { + const uint32_t max_tile_n = (N / kSubgroupMatrixN) * kSubgroupMatrixN; + if (t.tile_n > max_tile_n) { + t.tile_n = max_tile_n; + } +} + +// Looks up a pretuned tiling (subgroup_matrix_matmul_tuned.inc): match `arch`, map +// each dimension to the smallest grid value >= it. Nullopt if not covered. std::optional LookupPretunedTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K) { for (const auto& table : sgmm_tuned::kArchTables) { if (table.arch != arch) { @@ -160,12 +192,10 @@ std::optional LookupPretunedTiling(std::string_view arch, return std::nullopt; } -// Batch slices are dispatched on z as independent output-tile grids, so batch is -// a pure occupancy multiplier. The pretuned table and heuristic pick tile shape -// for a single (M, N, K) slice; once batch fills the machine, split-K's -// cooperative subgroups are redundant. Retire split-K factors whose batch-scaled -// occupancy would exceed ~2x hardware. -void ClampSplitKForBatch(SubgroupMatrixTiling& t, uint32_t M, uint32_t N, uint32_t batch, uint32_t hw) { +// batch is a pure occupancy multiplier (one output-tile grid per z slice); retires +// split-K factors whose batch-scaled occupancy would exceed ~2x hardware. +// MatMul-only: Conv's tiling ignores occupancy and batch entirely. +void CapSplitKForBatch(SubgroupMatrixTiling& t, uint32_t M, uint32_t N, uint32_t batch, uint32_t hw) { auto tile_count = [](uint32_t dim, uint32_t tile) { return (dim + tile - 1) / tile; }; const uint32_t eff_tiles = batch * tile_count(M, t.tile_m) * tile_count(N, t.tile_n); while (t.split_k > 1 && eff_tiles * t.split_k > 2 * hw) { @@ -173,23 +203,26 @@ void ClampSplitKForBatch(SubgroupMatrixTiling& t, uint32_t M, uint32_t N, uint32 } } -// Chooses the tile + split-K tiling for the given problem: use the pretuned -// table entry when one exists and fits, otherwise fall back to the heuristic. -// batch is the number of z-dispatched slices; it scales occupancy (see -// ClampSplitKForBatch) but not the per-slice tile shape. +// Uses the pretuned table entry when one exists and fits, else the heuristic. SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - if (const auto tuned = LookupPretunedTiling(arch, M, N, K); tuned && IsTilingValid(*tuned, K)) { + if (const auto tuned = LookupPretunedTiling(arch, M, N, K); tuned && IsTilingValid(*tuned, K, kMatMulAScratchCols)) { SubgroupMatrixTiling tiling = *tuned; - // The table is tuned per (M, N, K) at batch 1, so only revisit split-K when - // batch adds occupancy. + // The table is tuned at batch 1, so only revisit split-K when batch > 1. if (batch > 1) { - ClampSplitKForBatch(tiling, M, N, batch, EffectiveHwSubgroups(arch)); + CapSplitKForBatch(tiling, M, N, batch, EffectiveHwSubgroups(arch)); } return tiling; } return HeuristicTiling(arch, M, N, K, batch); } +// Conv counterpart of SelectTiling: no pretuned table, always ConvHeuristicTiling. +SubgroupMatrixTiling SelectConvTiling(uint32_t N, uint32_t K) { + SubgroupMatrixTiling tiling = ConvHeuristicTiling(N, K); + CapTileNToAlignedN(tiling, N); + return tiling; +} + } // namespace SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( @@ -197,13 +230,9 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( if (context.AdapterInfo().vendor != std::string_view{"intel"}) { return nullptr; } - // Intel tiling policy consumed by the common 8x16x16 subgroup-matrix kernel: - // the tile shape and split-K factor from the pretuned table or the heuristic. - // The subgroup-matrix shape itself is fixed by the kernel and not selected here. return [](const ComputeContext& context, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) -> std::optional { - // Only K needs to align to the subgroup tiling; M and N partial tiles are - // handled by bounds-checked stores in the kernel. + // M/N partial tiles are handled by bounds-checked stores in the kernel. if (K % kSubgroupMatrixK != 0) { return std::nullopt; } @@ -212,6 +241,22 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( }; } +SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( + const ComputeContextBase& context) { + if (context.AdapterInfo().vendor != std::string_view{"intel"}) { + return nullptr; + } + return [](const ComputeContext& /*context*/, uint32_t /*M*/, uint32_t N, + uint32_t K, uint32_t /*batch*/) -> std::optional { + // N must hold at least one subgroup-matrix N block (CapTileNToAlignedN can't + // help below that); must agree with CanApplySubgroupMatrixConv's own check. + if (K % kSubgroupMatrixK != 0 || N < kSubgroupMatrixN) { + return std::nullopt; + } + return SelectConvTiling(N, K); + }; +} + } // namespace intel } // namespace webgpu } // namespace onnxruntime 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 16914454600c4..816bb5aaaa960 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 @@ -12,14 +12,20 @@ namespace onnxruntime { namespace webgpu { namespace intel { -// Creates the Intel tile/split-K selection callback for the common subgroup-matrix -// kernel. Returns an empty selector on devices that are not Intel or that do not -// support the required 8x16x16 F16 subgroup matrix configuration, so the common -// factory yields no implementation and the caller falls back to the default -// MatMul path. +// MatMul selector's SLM budget, in f16 elements (16384 == 32 KB). Shared with the +// Conv kernel's own split-K safety clamp. +constexpr uint32_t kMaxScratchElems = 16384; + +// Tile/split-K selector for the common subgroup-matrix kernel. Empty on non-Intel +// devices or ones lacking the 8x16x16 F16 config. SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( const ComputeContextBase& context); +// Conv counterpart: Conv scratch budget (also stages the im2col A tile) and a +// Conv-only heuristic, no pretuned table. +SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( + const ComputeContextBase& context); + } // namespace intel } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc new file mode 100644 index 0000000000000..690724b7c9ed0 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc @@ -0,0 +1,385 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include + +#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 { + +// Reference NCHW 2D convolution (group == 1). +static void ComputeExpectedConv(const std::vector& x, const std::vector& w, + const std::vector& b, std::vector& y, + int64_t N, int64_t Cin, int64_t H, int64_t W, + int64_t Cout, int64_t kh, int64_t kw, + int64_t stride_h, int64_t stride_w, + int64_t pad_h, int64_t pad_w, + int64_t dil_h, int64_t dil_w, + int64_t out_h, int64_t out_w) { + for (int64_t n = 0; n < N; ++n) { + for (int64_t co = 0; co < Cout; ++co) { + for (int64_t oh = 0; oh < out_h; ++oh) { + for (int64_t ow = 0; ow < out_w; ++ow) { + float sum = b.empty() ? 0.0f : b[co]; + for (int64_t ci = 0; ci < Cin; ++ci) { + for (int64_t r = 0; r < kh; ++r) { + for (int64_t s = 0; s < kw; ++s) { + const int64_t ih = oh * stride_h - pad_h + r * dil_h; + const int64_t iw = ow * stride_w - pad_w + s * dil_w; + if (ih < 0 || ih >= H || iw < 0 || iw >= W) { + continue; + } + const float xv = x[((n * Cin + ci) * H + ih) * W + iw]; + const float wv = w[((co * Cin + ci) * kh + r) * kw + s]; + sum += xv * wv; + } + } + } + y[((n * Cout + co) * out_h + oh) * out_w + ow] = sum; + } + } + } + } +} + +#if !defined(DISABLE_CONTRIB_OPS) +// Reference for the fused activations; must stay in step with GetActivationSnippet +// (core/providers/webgpu/nn/fuse_utils.cc), which the kernel's epilogue builds from. +static void ApplyExpectedActivation(const std::string& kind, const std::vector& params, + std::vector& y) { + const float alpha = params.size() > 0 ? params[0] : 0.0f; + const float beta = params.size() > 1 ? params[1] : 0.0f; + for (float& v : y) { + if (kind == "Relu") { + v = std::max(v, 0.0f); + } else if (kind == "Sigmoid") { + v = 1.0f / (1.0f + std::exp(-v)); + } else if (kind == "Clip") { + v = std::clamp(v, alpha, beta); + } else if (kind == "HardSigmoid") { + v = std::clamp(alpha * v + beta, 0.0f, 1.0f); + } else if (kind == "LeakyRelu") { + v = v >= 0.0f ? v : alpha * v; + } else if (kind == "Tanh") { + v = std::tanh(v); + } else if (kind == "QuickGelu") { + v = v * (1.0f / (1.0f + std::exp(-(alpha * v)))); + } else if (kind == "HardSwish") { + v = v * std::clamp(v * (1.0f / 6.0f) + 0.5f, 0.0f, 1.0f); + } else if (kind == "Elu") { + v = v >= 0.0f ? v : alpha * (std::exp(v) - 1.0f); + } else if (kind == "ThresholdedRelu") { + v = v > alpha ? v : 0.0f; + } else if (kind == "Erf") { + v = std::erf(v); + } else if (kind == "Gelu") { + v = 0.5f * v * (1.0f + std::erf(v * 0.70710678118654752f)); + } else if (kind == "FastGelu") { + v = v * (0.5f + 0.5f * std::tanh(v * (0.035677408136300125f * v * v + 0.79788456080286535f))); + } else if (kind == "Softplus") { + v = std::max(v, 0.0f) + std::log(1.0f + std::exp(-std::abs(v))); + } else { + FAIL() << "unhandled activation in the test reference: " << kind; + } + } +} +#endif // !defined(DISABLE_CONTRIB_OPS) + +// Runs a Conv over the WebGPU EP and compares against the CPU reference above. A +// non-empty `activation` switches to FusedConv instead. +// +// Device-independent by construction (enabled, not DISABLED_ like the *_large tests +// here): where CanApplySubgroupMatrixConv() declines, the EP falls back to regular +// conv, which must match the same reference -- so a failure is a real defect +// either way, not a capability gap. +static void RunConvTest(int64_t N, int64_t Cin, int64_t H, int64_t W, + int64_t Cout, int64_t kh, int64_t kw, + int64_t stride = 1, int64_t pad = 0, int64_t dilation = 1, + bool has_bias = false, + const std::string& activation = "", + const std::vector& activation_params = {}, + float abs_err = 0.06f, float rel_err = 0.02f, + bool weight_is_initializer = true, + const std::string& auto_pad = "NOTSET") { +#if defined(DISABLE_CONTRIB_OPS) + if (!activation.empty()) { + GTEST_SKIP() << "FusedConv requires contrib ops."; + } +#endif + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + + const int64_t out_h = (H + 2 * pad - dilation * (kh - 1) - 1) / stride + 1; + const int64_t out_w = (W + 2 * pad - dilation * (kw - 1) - 1) / stride + 1; + ASSERT_GT(out_h, 0); + ASSERT_GT(out_w, 0); + + const std::vector x_dims{N, Cin, H, W}; + const std::vector w_dims{Cout, Cin, kh, kw}; + const std::vector b_dims{Cout}; + const std::vector y_dims{N, Cout, out_h, out_w}; + + RandomValueGenerator random{1234}; + std::vector x_vals(random.Gaussian(AsSpan(x_dims), 0.0f, 0.25f)); + std::vector w_vals(random.Gaussian(AsSpan(w_dims), 0.0f, 0.25f)); + std::vector b_vals = has_bias ? random.Gaussian(AsSpan(b_dims), 0.0f, 0.25f) + : std::vector{}; + + std::vector expected(static_cast(N * Cout * out_h * out_w)); + ComputeExpectedConv(x_vals, w_vals, b_vals, expected, N, Cin, H, W, Cout, kh, kw, + stride, stride, pad, pad, dilation, dilation, out_h, out_w); + + std::unique_ptr test; + if (activation.empty()) { + test = std::make_unique("Conv", 11); + } else { +#if !defined(DISABLE_CONTRIB_OPS) + test = std::make_unique("FusedConv", 1, onnxruntime::kMSDomain); + test->AddAttribute("activation", activation); + // GetFusedActivationAttr requires the attribute for the parameterized kinds and + // ignores it for the rest, so only set it when the case supplies parameters. + if (!activation_params.empty()) { + test->AddAttribute("activation_params", activation_params); + } + ApplyExpectedActivation(activation, activation_params, expected); +#endif + } + test->AddAttribute("group", static_cast(1)); + test->AddAttribute("kernel_shape", std::vector{kh, kw}); + if (auto_pad == "NOTSET") { + test->AddAttribute("pads", std::vector{pad, pad, pad, pad}); + } else { + // auto_pad and pads are mutually exclusive; pick pad/N/H/W to match (e.g. VALID + // with pad == 0). + test->AddAttribute("auto_pad", auto_pad); + } + test->AddAttribute("strides", std::vector{stride, stride}); + test->AddAttribute("dilations", std::vector{dilation, dilation}); + + test->AddInput("X", x_dims, FloatsToMLFloat16s(x_vals)); + // Default (constant) weight exercises the cached-transpose branch; RuntimeWeight + // covers the non-constant (rebuilt every run) one. + test->AddInput("W", w_dims, FloatsToMLFloat16s(w_vals), weight_is_initializer); + if (has_bias) { + test->AddInput("B", b_dims, FloatsToMLFloat16s(b_vals)); + } + test->AddOutput("Y", y_dims, FloatsToMLFloat16s(expected)); + test->SetOutputAbsErr("Y", abs_err); + test->SetOutputRelErr("Y", rel_err); + + test->ConfigEp(std::move(webgpu_ep)).RunWithConfig(); +} + +// K = kh*kw*Cin is a multiple of 16: the aligned fast path. out_h*out_w and Cout are +// tile multiples here (no partial tiles). +TEST(SubgroupMatrixConv, Aligned) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/16, /*W=*/16, /*Cout=*/64, /*kh=*/3, /*kw=*/3); + RunConvTest(/*N=*/1, /*Cin=*/32, /*H=*/8, /*W=*/8, /*Cout=*/128, /*kh=*/3, /*kw=*/3); +} + +// Partial output tiles: out_h*out_w (M) and/or Cout (N) are not tile multiples, +// exercising the kernel's bounds-checked M stores and its N tile shift. +TEST(SubgroupMatrixConv, PartialTiles) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/15, /*W=*/13, /*Cout=*/66, /*kh=*/3, /*kw=*/3); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/17, /*W=*/9, /*Cout=*/130, /*kh=*/3, /*kw=*/3); +} + +// Cout (GEMM N) values that stress the trailing-N-tile shift keeping the unchecked +// weight load inside the buffer. tile_n clamps to the largest multiple of 16 in +// Cout, so Cout in (16, 32) overlaps the shifted tile with the previous one by all +// but (Cout - 16) columns -- n_skip near kTileN, the narrowest write-out this +// kernel does. +// +// 16 -> exactly one N block, num_n_tile == 1, no shift at all +// 17 -> odd, minimal overhang: last tile shifts back by 15, writes 1 column +// 18 -> n_skip == 14, writes 2 columns +// 31 -> odd, largest overhang below the next block +// 33 -> odd, just past two blocks +// 50 -> if the selector picks tile_n 64 the clamp cuts it to 48, i.e. three +// subgroup matrices along N, a count the candidate list never produces +// on its own +// +// Odd values also cover the removal of the old N % 2 == 0 restriction: unlike the +// MatMul kernel, this one needs no even N because its weight is column-major with +// stride K, so load alignment follows from K % 16 == 0 rather than from Cout. +TEST(SubgroupMatrixConv, ShiftedAndOddNTiles) { + for (int64_t channels_out : {16, 17, 18, 31, 33, 50}) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/10, /*W=*/10, /*Cout=*/channels_out, /*kh=*/3, /*kw=*/3); + } + // Odd Cout together with a bias, which the epilogue indexes by the shifted column. + RunConvTest(/*N=*/2, /*Cin=*/16, /*H=*/11, /*W=*/7, /*Cout=*/17, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true); +} + +// Cin values that select each im2col gather width (vec_size 4 / 2 / 1). K must stay a +// multiple of 16, so the kernel size is chosen per case to make kh*kw*Cin align. +TEST(SubgroupMatrixConv, GatherVectorWidths) { + // Cin % 4 == 0 -> vec_size 4. + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3); + // Cin % 4 != 0, Cin % 2 == 0 -> vec_size 2. K = 2*4*6 = 48. + RunConvTest(/*N=*/1, /*Cin=*/6, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/2, /*kw=*/4); + // Cin odd -> vec_size 1. K = 2*8*3 = 48. + RunConvTest(/*N=*/1, /*Cin=*/3, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/2, /*kw=*/8); +} + +// Padding and strides: the im2col gather must return zero outside the input. +TEST(SubgroupMatrixConv, PadStride) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/16, /*W=*/16, /*Cout=*/64, /*kh=*/3, /*kw=*/3, + /*stride=*/2, /*pad=*/1); + RunConvTest(/*N=*/1, /*Cin=*/32, /*H=*/14, /*W=*/14, /*Cout=*/64, /*kh=*/5, /*kw=*/5, + /*stride=*/1, /*pad=*/2); +} + +// Dilated convolution. +TEST(SubgroupMatrixConv, Dilation) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/20, /*W=*/20, /*Cout=*/64, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/2, /*dilation=*/2); +} + +// Optional bias add. +TEST(SubgroupMatrixConv, Bias) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/16, /*W=*/16, /*Cout=*/64, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/0, /*dilation=*/1, /*has_bias=*/true); +} + +// Non-constant weight: PrePackInternal never runs, IsWeightConstant() is false and the +// impl rebuilds the OIHW -> OHWI transpose on every run instead of caching it. Same +// numerics as the constant-weight cases above, so any divergence points at the cache. +TEST(SubgroupMatrixConv, RuntimeWeight) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"", /*activation_params=*/{}, + /*abs_err=*/0.06f, /*rel_err=*/0.02f, /*weight_is_initializer=*/false); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/11, /*W=*/9, /*Cout=*/17, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/0, /*dilation=*/1, /*has_bias=*/false, + /*activation=*/"", /*activation_params=*/{}, + /*abs_err=*/0.06f, /*rel_err=*/0.02f, /*weight_is_initializer=*/false); +} + +// Batched conv: each batch slice is dispatched on z. Small per-slice M x N grids +// with a larger batch also stress the selector's split-K clamping. +TEST(SubgroupMatrixConv, Batched) { + RunConvTest(/*N=*/4, /*Cin=*/16, /*H=*/16, /*W=*/16, /*Cout=*/64, /*kh=*/3, /*kw=*/3); + RunConvTest(/*N=*/8, /*Cin=*/32, /*H=*/8, /*W=*/8, /*Cout=*/64, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1); +} + +// Plain Conv is also eligible for Im2ColMatMul, so PrePackInternal prepacks the +// constant weight to OHWI and Compute() reads that tensor directly (GetKernelLayout() +// == OHWI) instead of re-transposing it. +TEST(SubgroupMatrixConv, SharedPrepackOHWI) { + RunConvTest(/*N=*/1, /*Cin=*/32, /*H=*/10, /*W=*/10, /*Cout=*/48, /*kh=*/3, /*kw=*/3); +} + +// auto_pad != NOTSET: PrePackInternal can't predict the runtime pads, so the weight +// stays OIHW and Compute() builds its own cached OHWI transpose -- the one +// weight-handling branch the other tests don't exercise. +TEST(SubgroupMatrixConv, AutoPadSkipsPrepack) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/16, /*W=*/16, /*Cout=*/64, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/0, /*dilation=*/1, /*has_bias=*/false, + /*activation=*/"", /*activation_params=*/{}, + /*abs_err=*/0.06f, /*rel_err=*/0.02f, /*weight_is_initializer=*/true, + /*auto_pad=*/"VALID"); +} + +#if !defined(DISABLE_CONTRIB_OPS) +// Fused activations. The epilogue is generated from GetActivationSnippet, so every +// ActivationKind is expected to work -- including the ones a hand-written shader +// switch would have silently dropped. Each case below is a distinct branch of that +// snippet; the parameterized kinds also check that activation_params reach the +// uniforms in the right slots. +TEST(SubgroupMatrixConv, FusedActivationUnparameterized) { + for (const char* act : {"Relu", "Sigmoid", "Tanh", "HardSwish", "Softplus"}) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, /*activation=*/act); + } +} + +TEST(SubgroupMatrixConv, FusedActivationParameterized) { + // alpha, beta land in activation_param_0 / activation_param_1. + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"Clip", /*activation_params=*/{-0.5f, 0.5f}); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"HardSigmoid", /*activation_params=*/{0.2f, 0.5f}); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"LeakyRelu", /*activation_params=*/{0.1f}); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"Elu", /*activation_params=*/{1.0f}); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"ThresholdedRelu", /*activation_params=*/{0.25f}); + // QuickGelu has two shader variants: alpha == 1 drops the multiply and the uniform. + for (float alpha : {1.0f, 1.702f}) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"QuickGelu", /*activation_params=*/{alpha}); + } +} + +// The kinds that need GetActivationDeclaration to emit a module-scope helper +// (fused_act_erf for Erf/Gelu, fused_act_tanh for FastGelu) ahead of the epilogue -- +// the part of fuse_utils a hand-written shader switch cannot express at all. Erf and +// Gelu go through the Abramowitz & Stegun 7.1.26 approximation evaluated in f16, +// which costs roughly another 1e-2 of absolute accuracy against the CPU's std::erf, +// hence the looser bound. +TEST(SubgroupMatrixConv, FusedActivationWithDeclaration) { + for (const char* act : {"Erf", "Gelu", "FastGelu"}) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/12, /*W=*/12, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, /*activation=*/act, + /*activation_params=*/{}, /*abs_err=*/0.08f, /*rel_err=*/0.03f); + } +} + +// HardSwish isn't supported by Im2ColMatMul, so PrePackInternal falls through to the +// general HWIO transpose, and Compute() must read that layout directly too. Odd Cout +// (HWIO's row stride) must still produce the correct result: a constant weight pads +// N up to N_b (EnsurePaddedWeight), a non-constant one can't be cached so Compute() +// declines that unaligned load and falls back instead. +TEST(SubgroupMatrixConv, SharedPrepackHWIO) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/10, /*W=*/10, /*Cout=*/32, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/false, + /*activation=*/"HardSwish"); + for (int64_t channels_out : {17, 31, 33}) { + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/10, /*W=*/10, /*Cout=*/channels_out, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/false, + /*activation=*/"HardSwish"); + } + // Padded HWIO weight together with a bias and batch > 1. + RunConvTest(/*N=*/2, /*Cin=*/16, /*H=*/11, /*W=*/7, /*Cout=*/33, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"HardSwish"); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/10, /*W=*/10, /*Cout=*/17, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/false, + /*activation=*/"HardSwish", /*activation_params=*/{}, + /*abs_err=*/0.06f, /*rel_err=*/0.02f, /*weight_is_initializer=*/false); +} + +// A fused activation on a shifted/odd N tile: the epilogue, the bias index and the +// tile shift all interact in the same loop. +TEST(SubgroupMatrixConv, FusedActivationOddN) { + RunConvTest(/*N=*/2, /*Cin=*/16, /*H=*/11, /*W=*/9, /*Cout=*/17, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"Relu"); + RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/11, /*W=*/9, /*Cout=*/33, /*kh=*/3, /*kw=*/3, + /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/true, + /*activation=*/"LeakyRelu", /*activation_params=*/{0.1f}); +} +#endif // !defined(DISABLE_CONTRIB_OPS) + +} // namespace test +} // namespace onnxruntime diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul_pad_b.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul_pad_b.h index 1a7408e0588f2..2795d055853ed 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul_pad_b.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul_pad_b.h @@ -19,40 +19,41 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul_pad_b.wgsl.template">(ShaderHe // 4 | // Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer // 5 | // (N_b >= N), zero-filling columns [N, N_b). One thread per padded-output element. // 6 | // Gives B an aligned row stride so the subgroup-matrix f16 load's row-start -// 7 | // alignment holds for any N. See EnsurePaddedB in subgroup_matrix_matmul.cc. -// 8 | -// 9 | #use guardAgainstOutOfBoundsWorkgroupSizes -// 10 | #use .getByOffset .setByOffset -// 11 | -// 12 | $MAIN { +// 7 | // alignment holds for any N. Shared by MatMul and Conv; see +// 8 | // PadSubgroupMatrixWeightColumns in subgroup_matrix_pad_b.cc. +// 9 | +// 10 | #use guardAgainstOutOfBoundsWorkgroupSizes +// 11 | #use .getByOffset .setByOffset +// 12 | +// 13 | $MAIN { MainFunctionStart(); ss << "\n"; -// 13 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); +// 14 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); ss << " "; ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size"); ss << ";\n"; -// 14 | let r = global_idx / uniforms.N_b; // padded-output row +// 15 | let r = global_idx / uniforms.N_b; // padded-output row ss << " let r = global_idx / uniforms.N_b;\n"; -// 15 | let c = global_idx % uniforms.N_b; // padded-output column +// 16 | let c = global_idx % uniforms.N_b; // padded-output column ss << " let c = global_idx % uniforms.N_b;\n"; -// 16 | var v = output_value_t(0); +// 17 | var v = output_value_t(0); ss << " var v = output_value_t(0);\n"; -// 17 | if (c < uniforms.N) { // real column -> copy; else zero pad +// 18 | if (c < uniforms.N) { // real column -> copy; else zero pad ss << " if (c < uniforms.N) {\n"; -// 18 | v = output_value_t(input_b.getByOffset(r * uniforms.N + c)); +// 19 | v = output_value_t(input_b.getByOffset(r * uniforms.N + c)); ss << " v = output_value_t("; ss << __var_input_b->GetByOffset("r * uniforms.N + c"); ss << ");\n"; -// 19 | } +// 20 | } ss << " }\n"; -// 20 | output.setByOffset(global_idx, v); +// 21 | output.setByOffset(global_idx, v); ss << " "; ss << __var_output->SetByOffset("global_idx", "v"); ss << ";\n"; -// 21 | } // MAIN +// 22 | } // MAIN MainFunctionEnd(); ss << "\n"; -// 22 | +// 23 | return Status::OK(); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/subgroup_matrix_conv_8x16x16.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/subgroup_matrix_conv_8x16x16.h new file mode 100644 index 0000000000000..0ceaec859f0a9 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/subgroup_matrix_conv_8x16x16.h @@ -0,0 +1,1346 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: nn/subgroup_matrix_conv_8x16x16.wgsl.template + +template <> +Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract parameters + auto& __param_b_is_hwio = params.param_b_is_hwio; + auto& __param_sg_mat_count_m = params.param_sg_mat_count_m; + auto& __param_sg_mat_count_n = params.param_sg_mat_count_n; + auto& __param_sg_mat_k = params.param_sg_mat_k; + auto& __param_sg_mat_m = params.param_sg_mat_m; + auto& __param_sg_mat_n = params.param_sg_mat_n; + auto& __param_split_k = params.param_split_k; + auto& __param_vec_size = params.param_vec_size; + + // Extract variables + auto* __var_input_src = params.var_input_src; + +// 1 | // Copyright (c) Microsoft Corporation. All rights reserved. +// 2 | // Licensed under the MIT License. +// 3 | +// 4 | // Subgroup-matrix Conv kernel (F16), 8x16x16 shape: channels-last 2D convolution +// 5 | // as an implicit GEMM Y[MxN] = A[MxK] @ B[KxN] (+ optional bias), A the im2col of +// 6 | // the input (never materialized in global memory). +// 7 | // M = out_h * out_w, K = kernel_h * kernel_w * Cin, N = Cout +// 8 | // +// 9 | // A: gathered per K-block into a_scratch (im2col rows are scattered, but +// 10 | // subgroupMatrixLoad needs a contiguous buffer); K order (kh, kw, cin) matches +// 11 | // both weight layouts below. +// 12 | // B: OHWI, column-major (B[k,n] at n*K+k), or HWIO, row-major, row stride N_b +// 13 | // (B[k,n] at k*N_b+n) -- see b_is_hwio, set by the host to whichever the weight +// 14 | // arrived in. N_b pads HWIO's N up to a 4-byte-aligned width when N is odd +// 15 | // (never needed for OHWI, whose stride K is already a multiple of 16); padded +// 16 | // columns are read but never written (output always uses real N). +// 17 | // +// 18 | // Workgroup: split_k subgroups x 32 lanes, tile kTileM x kTileN (host-selected). K +// 19 | // blocks round-robin across subgroups when split_k > 1; partials summed at +// 20 | // write-out. Staging A forces the K reduction into a per-step +// 21 | // gather/barrier/multiply-accumulate loop in uniform control flow. +// 22 | // +// 23 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. +// 24 | // +// 25 | // TODO: Honor EnableMatmulFp32Accumulation() like matmul_packed.cc does: when it +// 26 | // is set, accumulate in f32 (an F16 x F16 -> F32 subgroup-matrix config, f32 +// 27 | // accumulators and scratch) instead of the f16 used today. +// 28 | // +// 29 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N_b so a trailing partial N tile +// 30 | // can shift back in bounds. M is unconstrained (see load_im2col_vec). +// 31 | +// 32 | #param sg_mat_k +// 33 | #param sg_mat_m +// 34 | #param sg_mat_n +// 35 | #param sg_mat_count_m +// 36 | #param sg_mat_count_n +// 37 | #param split_k +// 38 | #param vec_size +// 39 | #param b_is_hwio +// 40 | +// 41 | #use .getByOffset .setByOffset +// 42 | +// 43 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +ss << "const kSgMatM: u32 = u32("; +ss << __param_sg_mat_m; +ss << ");\n"; +// 44 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +ss << "const kSgMatN: u32 = u32("; +ss << __param_sg_mat_n; +ss << ");\n"; +// 45 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +ss << "const kSgMatK: u32 = u32("; +ss << __param_sg_mat_k; +ss << ");\n"; +// 46 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +ss << "const kSgMatCountM: u32 = u32("; +ss << __param_sg_mat_count_m; +ss << ");\n"; +// 47 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +ss << "const kSgMatCountN: u32 = u32("; +ss << __param_sg_mat_count_n; +ss << ");\n"; +// 48 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +ss << "const kTileM: u32 = kSgMatM * kSgMatCountM;\n"; +// 49 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +ss << "const kTileN: u32 = kSgMatN * kSgMatCountN;\n"; +// 50 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +ss << "const kSplitK: u32 = u32("; +ss << __param_split_k; +ss << ");\n"; +// 51 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +ss << "const kSubgroupSize: u32 = 32;\n"; +// 52 | const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (cin % vec_size == 0) +ss << "const kVecSize: u32 = u32("; +ss << __param_vec_size; +ss << ");\n"; +// 53 | +ss << "\n"; +// 54 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// 55 | // subgroup, summed at write-back. +// 56 | var scratch: array; +ss << "var scratch: array;\n"; +// 57 | +ss << "\n"; +// 58 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 59 | var a_scratch: array; +ss << "var a_scratch: array;\n"; +// 60 | +ss << "\n"; +// 61 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always +// 62 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding +// 63 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. +// 64 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. +// 65 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { +ss << "fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t {\n"; +// 66 | let cin_vec = uniforms.cin / kVecSize; +ss << " let cin_vec = uniforms.cin / kVecSize;\n"; +// 67 | let oh = m / uniforms.output_w; +ss << " let oh = m / uniforms.output_w;\n"; +// 68 | let ow = m % uniforms.output_w; +ss << " let ow = m % uniforms.output_w;\n"; +// 69 | let k_cin = k_vec % cin_vec; +ss << " let k_cin = k_vec % cin_vec;\n"; +// 70 | let k_hw = k_vec / cin_vec; +ss << " let k_hw = k_vec / cin_vec;\n"; +// 71 | let kw_idx = k_hw % uniforms.kernel_w; +ss << " let kw_idx = k_hw % uniforms.kernel_w;\n"; +// 72 | let kh_idx = k_hw / uniforms.kernel_w; +ss << " let kh_idx = k_hw / uniforms.kernel_w;\n"; +// 73 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); +ss << " let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x);\n"; +// 74 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); +ss << " let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y);\n"; +// 75 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { +ss << " if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) {\n"; +// 76 | return input_src_value_t(); +ss << " return input_src_value_t();\n"; +// 77 | } +ss << " }\n"; +// 78 | let input_vec_idx = input_batch_base_vec + +ss << " let input_vec_idx = input_batch_base_vec +\n"; +// 79 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +ss << " (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin;\n"; +// 80 | return input_src.getByOffset(input_vec_idx); +ss << " return "; +ss << __var_input_src->GetByOffset("input_vec_idx"); +ss << ";\n"; +// 81 | } +ss << "}\n"; +// 82 | +ss << "\n"; +// 83 | $MAIN { +MainFunctionStart(); +ss << "\n"; +// 84 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +ss << "\n"; +// 85 | // the true linear index used to recover batch_id and (n_tile, m_tile). +// 86 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +ss << " let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM;\n"; +// 87 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +ss << " let tiles_per_batch = uniforms.num_n_tile * num_m_tile;\n"; +// 88 | let batch_id = workgroup_idx / tiles_per_batch; +ss << " let batch_id = workgroup_idx / tiles_per_batch;\n"; +// 89 | // Dispatch normalization can add extra workgroups past the last batch element; +ss << "\n"; +// 90 | // drop them here (workgroup-uniform, so the barriers below stay uniform). +// 91 | if (batch_id >= uniforms.batch) { +ss << " if (batch_id >= uniforms.batch) {\n"; +// 92 | return; +ss << " return;\n"; +// 93 | } +ss << " }\n"; +// 94 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +ss << " let tile_idx = workgroup_idx % tiles_per_batch;\n"; +// 95 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile +ss << " let n_tile = tile_idx % uniforms.num_n_tile;\n"; +// 96 | let n_tile_start = n_tile * kTileN; +ss << " let n_tile_start = n_tile * kTileN;\n"; +// 97 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +ss << "\n"; +// 98 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). +// 99 | // Uses N_b (== N unless B was padded for row-major alignment) so the shifted +// 100 | // origin -- itself n_align-even since N_b, kTileN and n_tile_start all are -- +// 101 | // stays alignment-safe for the HWIO load below. +// 102 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); +ss << " let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN);\n"; +// 103 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile +ss << " let n_skip = n_tile_start - global_base_n;\n"; +// 104 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile +ss << " let m_tile = tile_idx / uniforms.num_n_tile;\n"; +// 105 | let m_base = m_tile * kTileM; +ss << " let m_base = m_tile * kTileM;\n"; +// 106 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); +ss << " let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize);\n"; +// 107 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +ss << " let out_batch_base = batch_id * uniforms.M * uniforms.N;\n"; +// 108 | let k_blocks = uniforms.K / kSgMatK; +ss << " let k_blocks = uniforms.K / kSgMatK;\n"; +// 109 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +ss << " let sg_index = local_idx / kSubgroupSize;\n"; +// 110 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +ss << " let sg_lane = local_idx % kSubgroupSize;\n"; +// 111 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +ss << " let sg_scratch_base = sg_index * kTileM * kTileN;\n"; +// 112 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +ss << " let a_scratch_base = sg_index * kTileM * kSgMatK;\n"; +// 113 | +ss << "\n"; +// 114 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 115 | var sg_mat_c0_0: subgroup_matrix_result; +ss << " var sg_mat_c0_0: subgroup_matrix_result;\n"; +// 116 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 117 | var sg_mat_c0_1: subgroup_matrix_result; +ss << " var sg_mat_c0_1: subgroup_matrix_result;\n"; +// 118 | #endif +} +// 119 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 120 | var sg_mat_c0_2: subgroup_matrix_result; +ss << " var sg_mat_c0_2: subgroup_matrix_result;\n"; +// 121 | #endif +} +// 122 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 123 | var sg_mat_c0_3: subgroup_matrix_result; +ss << " var sg_mat_c0_3: subgroup_matrix_result;\n"; +// 124 | #endif +} +// 125 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 126 | var sg_mat_c1_0: subgroup_matrix_result; +ss << " var sg_mat_c1_0: subgroup_matrix_result;\n"; +// 127 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 128 | var sg_mat_c1_1: subgroup_matrix_result; +ss << " var sg_mat_c1_1: subgroup_matrix_result;\n"; +// 129 | #endif +} +// 130 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 131 | var sg_mat_c1_2: subgroup_matrix_result; +ss << " var sg_mat_c1_2: subgroup_matrix_result;\n"; +// 132 | #endif +} +// 133 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 134 | var sg_mat_c1_3: subgroup_matrix_result; +ss << " var sg_mat_c1_3: subgroup_matrix_result;\n"; +// 135 | #endif +} +// 136 | #endif +} +// 137 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 138 | var sg_mat_c2_0: subgroup_matrix_result; +ss << " var sg_mat_c2_0: subgroup_matrix_result;\n"; +// 139 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 140 | var sg_mat_c2_1: subgroup_matrix_result; +ss << " var sg_mat_c2_1: subgroup_matrix_result;\n"; +// 141 | #endif +} +// 142 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 143 | var sg_mat_c2_2: subgroup_matrix_result; +ss << " var sg_mat_c2_2: subgroup_matrix_result;\n"; +// 144 | #endif +} +// 145 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 146 | var sg_mat_c2_3: subgroup_matrix_result; +ss << " var sg_mat_c2_3: subgroup_matrix_result;\n"; +// 147 | #endif +} +// 148 | #endif +} +// 149 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 150 | var sg_mat_c3_0: subgroup_matrix_result; +ss << " var sg_mat_c3_0: subgroup_matrix_result;\n"; +// 151 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 152 | var sg_mat_c3_1: subgroup_matrix_result; +ss << " var sg_mat_c3_1: subgroup_matrix_result;\n"; +// 153 | #endif +} +// 154 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 155 | var sg_mat_c3_2: subgroup_matrix_result; +ss << " var sg_mat_c3_2: subgroup_matrix_result;\n"; +// 156 | #endif +} +// 157 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 158 | var sg_mat_c3_3: subgroup_matrix_result; +ss << " var sg_mat_c3_3: subgroup_matrix_result;\n"; +// 159 | #endif +} +// 160 | #endif +} +// 161 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 162 | var sg_mat_c4_0: subgroup_matrix_result; +ss << " var sg_mat_c4_0: subgroup_matrix_result;\n"; +// 163 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 164 | var sg_mat_c4_1: subgroup_matrix_result; +ss << " var sg_mat_c4_1: subgroup_matrix_result;\n"; +// 165 | #endif +} +// 166 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 167 | var sg_mat_c4_2: subgroup_matrix_result; +ss << " var sg_mat_c4_2: subgroup_matrix_result;\n"; +// 168 | #endif +} +// 169 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 170 | var sg_mat_c4_3: subgroup_matrix_result; +ss << " var sg_mat_c4_3: subgroup_matrix_result;\n"; +// 171 | #endif +} +// 172 | #endif +} +// 173 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 174 | var sg_mat_c5_0: subgroup_matrix_result; +ss << " var sg_mat_c5_0: subgroup_matrix_result;\n"; +// 175 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 176 | var sg_mat_c5_1: subgroup_matrix_result; +ss << " var sg_mat_c5_1: subgroup_matrix_result;\n"; +// 177 | #endif +} +// 178 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 179 | var sg_mat_c5_2: subgroup_matrix_result; +ss << " var sg_mat_c5_2: subgroup_matrix_result;\n"; +// 180 | #endif +} +// 181 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 182 | var sg_mat_c5_3: subgroup_matrix_result; +ss << " var sg_mat_c5_3: subgroup_matrix_result;\n"; +// 183 | #endif +} +// 184 | #endif +} +// 185 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 186 | var sg_mat_c6_0: subgroup_matrix_result; +ss << " var sg_mat_c6_0: subgroup_matrix_result;\n"; +// 187 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 188 | var sg_mat_c6_1: subgroup_matrix_result; +ss << " var sg_mat_c6_1: subgroup_matrix_result;\n"; +// 189 | #endif +} +// 190 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 191 | var sg_mat_c6_2: subgroup_matrix_result; +ss << " var sg_mat_c6_2: subgroup_matrix_result;\n"; +// 192 | #endif +} +// 193 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 194 | var sg_mat_c6_3: subgroup_matrix_result; +ss << " var sg_mat_c6_3: subgroup_matrix_result;\n"; +// 195 | #endif +} +// 196 | #endif +} +// 197 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 198 | var sg_mat_c7_0: subgroup_matrix_result; +ss << " var sg_mat_c7_0: subgroup_matrix_result;\n"; +// 199 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 200 | var sg_mat_c7_1: subgroup_matrix_result; +ss << " var sg_mat_c7_1: subgroup_matrix_result;\n"; +// 201 | #endif +} +// 202 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 203 | var sg_mat_c7_2: subgroup_matrix_result; +ss << " var sg_mat_c7_2: subgroup_matrix_result;\n"; +// 204 | #endif +} +// 205 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 206 | var sg_mat_c7_3: subgroup_matrix_result; +ss << " var sg_mat_c7_3: subgroup_matrix_result;\n"; +// 207 | #endif +} +// 208 | #endif +} +// 209 | +// 210 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, +// 211 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks +// 212 | // early (valid == false), keeping the barriers below in uniform control flow. +// 213 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +ss << " let num_steps = (k_blocks + kSplitK - 1u) / kSplitK;\n"; +// 214 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +ss << " for (var step: u32 = 0u; step < num_steps; step = step + 1u) {\n"; +// 215 | let kb = sg_index + step * kSplitK; +ss << " let kb = sg_index + step * kSplitK;\n"; +// 216 | let valid = kb < k_blocks; +ss << " let valid = kb < k_blocks;\n"; +// 217 | +ss << "\n"; +// 218 | if (valid) { +ss << " if (valid) {\n"; +// 219 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +ss << " let k_vec_per_sgmat_row = kSgMatK / kVecSize;\n"; +// 220 | let k_vec_base = kb * k_vec_per_sgmat_row; +ss << " let k_vec_base = kb * k_vec_per_sgmat_row;\n"; +// 221 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { +ss << " for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) {\n"; +// 222 | let row = vidx / k_vec_per_sgmat_row; +ss << " let row = vidx / k_vec_per_sgmat_row;\n"; +// 223 | let kc_vec = vidx % k_vec_per_sgmat_row; +ss << " let kc_vec = vidx % k_vec_per_sgmat_row;\n"; +// 224 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); +ss << " let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec);\n"; +// 225 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +ss << " let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize;\n"; +// 226 | #if vec_size == 1 +if (__param_vec_size == 1) { +// 227 | a_scratch[dst] = chunk; +ss << " a_scratch[dst] = chunk;\n"; +// 228 | #else +} else { +// 229 | a_scratch[dst + 0u] = chunk[0]; +ss << " a_scratch[dst + 0u] = chunk[0];\n"; +// 230 | a_scratch[dst + 1u] = chunk[1]; +ss << " a_scratch[dst + 1u] = chunk[1];\n"; +// 231 | #if vec_size >= 4 +if (__param_vec_size >= 4) { +// 232 | a_scratch[dst + 2u] = chunk[2]; +ss << " a_scratch[dst + 2u] = chunk[2];\n"; +// 233 | a_scratch[dst + 3u] = chunk[3]; +ss << " a_scratch[dst + 3u] = chunk[3];\n"; +// 234 | #endif +} +// 235 | #endif +} +// 236 | } +ss << " }\n"; +// 237 | } +ss << " }\n"; +// 238 | // Make the staged A-tiles visible before loading them as subgroup matrices. +ss << "\n"; +// 239 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 240 | +ss << "\n"; +// 241 | if (valid) { +ss << " if (valid) {\n"; +// 242 | // Load the B right tiles for this K block. +ss << "\n"; +// 243 | #if b_is_hwio +if (__param_b_is_hwio) { +// 244 | // HWIO: B is row-major ([K, N_b], B[k, n] at k*N_b + n); load row-major, +// 245 | // stride N_b. N_b (== N unless padded) keeps the row stride, and via +// 246 | // global_base_n above, the column offset, 4-byte aligned. +// 247 | let b_base = kb * kSgMatK * uniforms.N_b + global_base_n; +ss << " let b_base = kb * kSgMatK * uniforms.N_b + global_base_n;\n"; +// 248 | var sg_mat_b0: subgroup_matrix_right = +ss << " var sg_mat_b0: subgroup_matrix_right =\n"; +// 249 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 250 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 0 * kSgMatN, uniforms.N_b);\n"; +// 251 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 252 | var sg_mat_b1: subgroup_matrix_right = +ss << " var sg_mat_b1: subgroup_matrix_right =\n"; +// 253 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 254 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 1 * kSgMatN, uniforms.N_b);\n"; +// 255 | #endif +} +// 256 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 257 | var sg_mat_b2: subgroup_matrix_right = +ss << " var sg_mat_b2: subgroup_matrix_right =\n"; +// 258 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 259 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 2 * kSgMatN, uniforms.N_b);\n"; +// 260 | #endif +} +// 261 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 262 | var sg_mat_b3: subgroup_matrix_right = +ss << " var sg_mat_b3: subgroup_matrix_right =\n"; +// 263 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 264 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 3 * kSgMatN, uniforms.N_b);\n"; +// 265 | #endif +} +// 266 | #else +} else { +// 267 | // OHWI: B is column-major ([N, K], B[k, n] at n*K + k); load col-major, stride K. +// 268 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +ss << " let b_base = global_base_n * uniforms.K + kb * kSgMatK;\n"; +// 269 | var sg_mat_b0: subgroup_matrix_right = +ss << " var sg_mat_b0: subgroup_matrix_right =\n"; +// 270 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 271 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 272 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 273 | var sg_mat_b1: subgroup_matrix_right = +ss << " var sg_mat_b1: subgroup_matrix_right =\n"; +// 274 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 275 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 276 | #endif +} +// 277 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 278 | var sg_mat_b2: subgroup_matrix_right = +ss << " var sg_mat_b2: subgroup_matrix_right =\n"; +// 279 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 280 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 281 | #endif +} +// 282 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 283 | var sg_mat_b3: subgroup_matrix_right = +ss << " var sg_mat_b3: subgroup_matrix_right =\n"; +// 284 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 285 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 286 | #endif +} +// 287 | #endif +} +// 288 | +// 289 | // Load the A left tiles (one per M block) from workgroup staging, +// 290 | // row-major with stride kSgMatK. +// 291 | var sg_mat_a0: subgroup_matrix_left = +ss << " var sg_mat_a0: subgroup_matrix_left =\n"; +// 292 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 293 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 294 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 295 | var sg_mat_a1: subgroup_matrix_left = +ss << " var sg_mat_a1: subgroup_matrix_left =\n"; +// 296 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 297 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 298 | #endif +} +// 299 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 300 | var sg_mat_a2: subgroup_matrix_left = +ss << " var sg_mat_a2: subgroup_matrix_left =\n"; +// 301 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 302 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 303 | #endif +} +// 304 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 305 | var sg_mat_a3: subgroup_matrix_left = +ss << " var sg_mat_a3: subgroup_matrix_left =\n"; +// 306 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 307 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 308 | #endif +} +// 309 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 310 | var sg_mat_a4: subgroup_matrix_left = +ss << " var sg_mat_a4: subgroup_matrix_left =\n"; +// 311 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 312 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 313 | #endif +} +// 314 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 315 | var sg_mat_a5: subgroup_matrix_left = +ss << " var sg_mat_a5: subgroup_matrix_left =\n"; +// 316 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 317 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 318 | #endif +} +// 319 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 320 | var sg_mat_a6: subgroup_matrix_left = +ss << " var sg_mat_a6: subgroup_matrix_left =\n"; +// 321 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 322 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 323 | #endif +} +// 324 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 325 | var sg_mat_a7: subgroup_matrix_left = +ss << " var sg_mat_a7: subgroup_matrix_left =\n"; +// 326 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 327 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 328 | #endif +} +// 329 | +// 330 | // Accumulate every (M block, N block) pair. +// 331 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +ss << " sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0);\n"; +// 332 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 333 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +ss << " sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1);\n"; +// 334 | #endif +} +// 335 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 336 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +ss << " sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2);\n"; +// 337 | #endif +} +// 338 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 339 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +ss << " sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3);\n"; +// 340 | #endif +} +// 341 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 342 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +ss << " sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0);\n"; +// 343 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 344 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +ss << " sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1);\n"; +// 345 | #endif +} +// 346 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 347 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +ss << " sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2);\n"; +// 348 | #endif +} +// 349 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 350 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +ss << " sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3);\n"; +// 351 | #endif +} +// 352 | #endif +} +// 353 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 354 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +ss << " sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0);\n"; +// 355 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 356 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +ss << " sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1);\n"; +// 357 | #endif +} +// 358 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 359 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +ss << " sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2);\n"; +// 360 | #endif +} +// 361 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 362 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +ss << " sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3);\n"; +// 363 | #endif +} +// 364 | #endif +} +// 365 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 366 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +ss << " sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0);\n"; +// 367 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 368 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +ss << " sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1);\n"; +// 369 | #endif +} +// 370 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 371 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +ss << " sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2);\n"; +// 372 | #endif +} +// 373 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 374 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +ss << " sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3);\n"; +// 375 | #endif +} +// 376 | #endif +} +// 377 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 378 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +ss << " sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0);\n"; +// 379 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 380 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +ss << " sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1);\n"; +// 381 | #endif +} +// 382 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 383 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +ss << " sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2);\n"; +// 384 | #endif +} +// 385 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 386 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +ss << " sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3);\n"; +// 387 | #endif +} +// 388 | #endif +} +// 389 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 390 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +ss << " sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0);\n"; +// 391 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 392 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +ss << " sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1);\n"; +// 393 | #endif +} +// 394 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 395 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +ss << " sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2);\n"; +// 396 | #endif +} +// 397 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 398 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +ss << " sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3);\n"; +// 399 | #endif +} +// 400 | #endif +} +// 401 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 402 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +ss << " sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0);\n"; +// 403 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 404 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +ss << " sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1);\n"; +// 405 | #endif +} +// 406 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 407 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +ss << " sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2);\n"; +// 408 | #endif +} +// 409 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 410 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +ss << " sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3);\n"; +// 411 | #endif +} +// 412 | #endif +} +// 413 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 414 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +ss << " sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0);\n"; +// 415 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 416 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +ss << " sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1);\n"; +// 417 | #endif +} +// 418 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 419 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +ss << " sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2);\n"; +// 420 | #endif +} +// 421 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 422 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +ss << " sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3);\n"; +// 423 | #endif +} +// 424 | #endif +} +// 425 | } +ss << " }\n"; +// 426 | // Ensure every subgroup is done reading a_scratch before it is reused for +ss << "\n"; +// 427 | // the next K-block step. +// 428 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 429 | } +ss << " }\n"; +// 430 | +ss << "\n"; +// 431 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; +// 432 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). +// 433 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN);\n"; +// 434 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 435 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN);\n"; +// 436 | #endif +} +// 437 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 438 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN);\n"; +// 439 | #endif +} +// 440 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 441 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN);\n"; +// 442 | #endif +} +// 443 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 444 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN);\n"; +// 445 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 446 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN);\n"; +// 447 | #endif +} +// 448 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 449 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN);\n"; +// 450 | #endif +} +// 451 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN);\n"; +// 453 | #endif +} +// 454 | #endif +} +// 455 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 456 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN);\n"; +// 457 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN);\n"; +// 459 | #endif +} +// 460 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 461 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN);\n"; +// 462 | #endif +} +// 463 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN);\n"; +// 465 | #endif +} +// 466 | #endif +} +// 467 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 468 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN);\n"; +// 469 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 470 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN);\n"; +// 471 | #endif +} +// 472 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 473 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN);\n"; +// 474 | #endif +} +// 475 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 476 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN);\n"; +// 477 | #endif +} +// 478 | #endif +} +// 479 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 480 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN);\n"; +// 481 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 482 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN);\n"; +// 483 | #endif +} +// 484 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 485 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN);\n"; +// 486 | #endif +} +// 487 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 488 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN);\n"; +// 489 | #endif +} +// 490 | #endif +} +// 491 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 492 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN);\n"; +// 493 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 494 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN);\n"; +// 495 | #endif +} +// 496 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 497 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN);\n"; +// 498 | #endif +} +// 499 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 500 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN);\n"; +// 501 | #endif +} +// 502 | #endif +} +// 503 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 504 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN);\n"; +// 505 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 506 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN);\n"; +// 507 | #endif +} +// 508 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 509 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN);\n"; +// 510 | #endif +} +// 511 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 512 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN);\n"; +// 513 | #endif +} +// 514 | #endif +} +// 515 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 516 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN);\n"; +// 517 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 518 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN);\n"; +// 519 | #endif +} +// 520 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 521 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN);\n"; +// 522 | #endif +} +// 523 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 524 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +ss << " subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN);\n"; +// 525 | #endif +} +// 526 | #endif +} +// 527 | +// 528 | #if split_k >= 2 +if (__param_split_k >= 2) { +// 529 | // Make all subgroups' partial tiles visible before the reduction. +// 530 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 531 | +ss << "\n"; +// 532 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each +// 533 | // position is only overwritten by the lane that already read it. +// 534 | if (sg_index == 0u) { +ss << " if (sg_index == 0u) {\n"; +// 535 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +ss << " for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) {\n"; +// 536 | var acc: f16 = f16(0); +ss << " var acc: f16 = f16(0);\n"; +// 537 | for (var s: u32 = 0; s < kSplitK; s++) { +ss << " for (var s: u32 = 0; s < kSplitK; s++) {\n"; +// 538 | acc += scratch[s * kTileM * kTileN + idx]; +ss << " acc += scratch[s * kTileM * kTileN + idx];\n"; +// 539 | } +ss << " }\n"; +// 540 | scratch[idx] = acc; +ss << " scratch[idx] = acc;\n"; +// 541 | } +ss << " }\n"; +// 542 | } +ss << " }\n"; +// 543 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +ss << "\n"; +// 544 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 545 | #endif +} +// 546 | +// 547 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip +// 548 | // and M's bounds check make non-multiple dimensions write each element once. +// 549 | let n_count = min(kTileN, uniforms.N - global_base_n); +ss << " let n_count = min(kTileN, uniforms.N - global_base_n);\n"; +// 550 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +ss << " for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) {\n"; +// 551 | let global_m = m_base + r; +ss << " let global_m = m_base + r;\n"; +// 552 | if (global_m < uniforms.M) { +ss << " if (global_m < uniforms.M) {\n"; +// 553 | let scratch_base = r * kTileN; +ss << " let scratch_base = r * kTileN;\n"; +// 554 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; +ss << " let out_base = out_batch_base + global_m * uniforms.N + global_base_n;\n"; +// 555 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +ss << " for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) {\n"; +// 556 | write_output(out_base + i, global_base_n + i, +ss << " write_output(out_base + i, global_base_n + i,\n"; +// 557 | output_value_t(scratch[scratch_base + i])); +ss << " output_value_t(scratch[scratch_base + i]));\n"; +// 558 | } +ss << " }\n"; +// 559 | } +ss << " }\n"; +// 560 | } +ss << " }\n"; +// 561 | } // MAIN +MainFunctionEnd(); +ss << "\n"; +// 562 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h index 3e837a94da226..fe0ae419c99f4 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h @@ -121,6 +121,28 @@ struct TemplateParameter<"nn/im2col_matmul.wgsl.template"> { template <> Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/im2col_matmul.wgsl.template">::type params); +// +// Template: nn/subgroup_matrix_conv_8x16x16.wgsl.template +// + +template <> +struct TemplateParameter<"nn/subgroup_matrix_conv_8x16x16.wgsl.template"> { + using type = struct { + int param_b_is_hwio; + int param_sg_mat_count_m; + int param_sg_mat_count_n; + int param_sg_mat_k; + int param_sg_mat_m; + int param_sg_mat_n; + int param_split_k; + int param_vec_size; + const ShaderVariableHelper* var_input_src; + }; +}; + +template <> +Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">::type params); + // // Template: tensor/oihw_to_ohwi.wgsl.template // diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h index b85499e73881b..131531a59778a 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h @@ -37,9 +37,10 @@ std::string pass_as_string(T&& v) { #include "wgsl_template_gen/generated/math/matmul_gemv.h" // ad945fab49f626209e741b76ac3900250cb43159d4ba8a3f66de12fad97f5c43 #include "wgsl_template_gen/generated/math/subgroup_matrix_gemm_8x16x16.h" // 1aaff98897bad5e6e46888a2948ca3ac245eda97bdbbe7463c7cec4b6da6d0a9 #include "wgsl_template_gen/generated/math/subgroup_matrix_matmul.h" // 5eb470fbe35d55b1ccfbf62f8d780a6f2798853a2628f81e19622f0b24a287f7 -#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // c01fae062d8059604670635900f3dd036b692cc3d16ebb8a1e54491e83439e66 +#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // e6c1647249a3190f14419e17b2cfc389cacb6d4b28e85dcc2827519e14c7cd67 #include "wgsl_template_gen/generated/nn/conv_transpose3d.h" // 1d4758c416a863e99d7a7b09368c8ef8f90272fcbb271d5020b842a437aa7198 #include "wgsl_template_gen/generated/nn/im2col_matmul.h" // 9842e2aba8c56fad8744906918b044a4a018ce7cfe5efe78d1d18742ac4d3b87 +#include "wgsl_template_gen/generated/nn/subgroup_matrix_conv_8x16x16.h" // b84b2d2788699e7088f71421393e8c1b9f68ada69037885c2cd85ce74326b9df #include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // ea6d6ea3ad39540d34a2153b291792f6e419401941c9f1fcbc910a9d0880df54 #include "wgsl_template_gen/generated/tensor/pad.h" // 2c39520bb179cbd18964c7364065e3de93d7f99b6e3eacee1e662ec8c4b88c04 diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h index 72c280617cd01..ab69dd39f0e93 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h @@ -19,40 +19,41 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul_pad_b.wgsl.template">(ShaderHe // 4 | // Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer // 5 | // (N_b >= N), zero-filling columns [N, N_b). One thread per padded-output element. // 6 | // Gives B an aligned row stride so the subgroup-matrix f16 load's row-start -// 7 | // alignment holds for any N. See EnsurePaddedB in subgroup_matrix_matmul.cc. -// 8 | -// 9 | #use guardAgainstOutOfBoundsWorkgroupSizes -// 10 | #use .getByOffset .setByOffset -// 11 | -// 12 | $MAIN { +// 7 | // alignment holds for any N. Shared by MatMul and Conv; see +// 8 | // PadSubgroupMatrixWeightColumns in subgroup_matrix_pad_b.cc. +// 9 | +// 10 | #use guardAgainstOutOfBoundsWorkgroupSizes +// 11 | #use .getByOffset .setByOffset +// 12 | +// 13 | $MAIN { MainFunctionStart(); ss << __str_5; -// 13 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); +// 14 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); ss << __str_241; ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_237); ss << __str_29; -// 14 | let r = global_idx / uniforms.N_b; // padded-output row +// 15 | let r = global_idx / uniforms.N_b; // padded-output row ss << __str_242; -// 15 | let c = global_idx % uniforms.N_b; // padded-output column +// 16 | let c = global_idx % uniforms.N_b; // padded-output column ss << __str_243; -// 16 | var v = output_value_t(0); +// 17 | var v = output_value_t(0); ss << __str_244; -// 17 | if (c < uniforms.N) { // real column -> copy; else zero pad +// 18 | if (c < uniforms.N) { // real column -> copy; else zero pad ss << __str_245; -// 18 | v = output_value_t(input_b.getByOffset(r * uniforms.N + c)); +// 19 | v = output_value_t(input_b.getByOffset(r * uniforms.N + c)); ss << __str_246; ss << __var_input_b->GetByOffset(__str_238); ss << __str_17; -// 19 | } +// 20 | } ss << __str_11; -// 20 | output.setByOffset(global_idx, v); +// 21 | output.setByOffset(global_idx, v); ss << __str_241; ss << __var_output->SetByOffset(__str_239, __str_240); ss << __str_29; -// 21 | } // MAIN +// 22 | } // MAIN MainFunctionEnd(); ss << __str_5; -// 22 | +// 23 | return Status::OK(); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/subgroup_matrix_conv_8x16x16.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/subgroup_matrix_conv_8x16x16.h new file mode 100644 index 0000000000000..79282ca894ed7 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/subgroup_matrix_conv_8x16x16.h @@ -0,0 +1,1346 @@ +// This file is auto-generated by wgsl-gen. Do not edit manually. + +#pragma once + +// Template implementation +// Source: nn/subgroup_matrix_conv_8x16x16.wgsl.template + +template <> +Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">::type params) { + [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); + + // Extract parameters + auto& __param_b_is_hwio = params.param_b_is_hwio; + auto& __param_sg_mat_count_m = params.param_sg_mat_count_m; + auto& __param_sg_mat_count_n = params.param_sg_mat_count_n; + auto& __param_sg_mat_k = params.param_sg_mat_k; + auto& __param_sg_mat_m = params.param_sg_mat_m; + auto& __param_sg_mat_n = params.param_sg_mat_n; + auto& __param_split_k = params.param_split_k; + auto& __param_vec_size = params.param_vec_size; + + // Extract variables + auto* __var_input_src = params.var_input_src; + +// 1 | // Copyright (c) Microsoft Corporation. All rights reserved. +// 2 | // Licensed under the MIT License. +// 3 | +// 4 | // Subgroup-matrix Conv kernel (F16), 8x16x16 shape: channels-last 2D convolution +// 5 | // as an implicit GEMM Y[MxN] = A[MxK] @ B[KxN] (+ optional bias), A the im2col of +// 6 | // the input (never materialized in global memory). +// 7 | // M = out_h * out_w, K = kernel_h * kernel_w * Cin, N = Cout +// 8 | // +// 9 | // A: gathered per K-block into a_scratch (im2col rows are scattered, but +// 10 | // subgroupMatrixLoad needs a contiguous buffer); K order (kh, kw, cin) matches +// 11 | // both weight layouts below. +// 12 | // B: OHWI, column-major (B[k,n] at n*K+k), or HWIO, row-major, row stride N_b +// 13 | // (B[k,n] at k*N_b+n) -- see b_is_hwio, set by the host to whichever the weight +// 14 | // arrived in. N_b pads HWIO's N up to a 4-byte-aligned width when N is odd +// 15 | // (never needed for OHWI, whose stride K is already a multiple of 16); padded +// 16 | // columns are read but never written (output always uses real N). +// 17 | // +// 18 | // Workgroup: split_k subgroups x 32 lanes, tile kTileM x kTileN (host-selected). K +// 19 | // blocks round-robin across subgroups when split_k > 1; partials summed at +// 20 | // write-out. Staging A forces the K reduction into a per-step +// 21 | // gather/barrier/multiply-accumulate loop in uniform control flow. +// 22 | // +// 23 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. +// 24 | // +// 25 | // TODO: Honor EnableMatmulFp32Accumulation() like matmul_packed.cc does: when it +// 26 | // is set, accumulate in f32 (an F16 x F16 -> F32 subgroup-matrix config, f32 +// 27 | // accumulators and scratch) instead of the f16 used today. +// 28 | // +// 29 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N_b so a trailing partial N tile +// 30 | // can shift back in bounds. M is unconstrained (see load_im2col_vec). +// 31 | +// 32 | #param sg_mat_k +// 33 | #param sg_mat_m +// 34 | #param sg_mat_n +// 35 | #param sg_mat_count_m +// 36 | #param sg_mat_count_n +// 37 | #param split_k +// 38 | #param vec_size +// 39 | #param b_is_hwio +// 40 | +// 41 | #use .getByOffset .setByOffset +// 42 | +// 43 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +ss << __str_32; +ss << __param_sg_mat_m; +ss << __str_17; +// 44 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +ss << __str_33; +ss << __param_sg_mat_n; +ss << __str_17; +// 45 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +ss << __str_34; +ss << __param_sg_mat_k; +ss << __str_17; +// 46 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +ss << __str_35; +ss << __param_sg_mat_count_m; +ss << __str_17; +// 47 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +ss << __str_36; +ss << __param_sg_mat_count_n; +ss << __str_17; +// 48 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +ss << __str_37; +// 49 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +ss << __str_38; +// 50 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +ss << __str_39; +ss << __param_split_k; +ss << __str_17; +// 51 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +ss << __str_40; +// 52 | const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (cin % vec_size == 0) +ss << __str_407; +ss << __param_vec_size; +ss << __str_17; +// 53 | +ss << __str_5; +// 54 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// 55 | // subgroup, summed at write-back. +// 56 | var scratch: array; +ss << __str_41; +// 57 | +ss << __str_5; +// 58 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 59 | var a_scratch: array; +ss << __str_408; +// 60 | +ss << __str_5; +// 61 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always +// 62 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding +// 63 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. +// 64 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. +// 65 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { +ss << __str_409; +// 66 | let cin_vec = uniforms.cin / kVecSize; +ss << __str_410; +// 67 | let oh = m / uniforms.output_w; +ss << __str_411; +// 68 | let ow = m % uniforms.output_w; +ss << __str_412; +// 69 | let k_cin = k_vec % cin_vec; +ss << __str_413; +// 70 | let k_hw = k_vec / cin_vec; +ss << __str_414; +// 71 | let kw_idx = k_hw % uniforms.kernel_w; +ss << __str_415; +// 72 | let kh_idx = k_hw / uniforms.kernel_w; +ss << __str_416; +// 73 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); +ss << __str_417; +// 74 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); +ss << __str_418; +// 75 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { +ss << __str_419; +// 76 | return input_src_value_t(); +ss << __str_420; +// 77 | } +ss << __str_20; +// 78 | let input_vec_idx = input_batch_base_vec + +ss << __str_421; +// 79 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +ss << __str_422; +// 80 | return input_src.getByOffset(input_vec_idx); +ss << __str_351; +ss << __var_input_src->GetByOffset(__str_406); +ss << __str_29; +// 81 | } +ss << __str_345; +// 82 | +ss << __str_5; +// 83 | $MAIN { +MainFunctionStart(); +ss << __str_5; +// 84 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +ss << __str_5; +// 85 | // the true linear index used to recover batch_id and (n_tile, m_tile). +// 86 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +ss << __str_218; +// 87 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +ss << __str_423; +// 88 | let batch_id = workgroup_idx / tiles_per_batch; +ss << __str_424; +// 89 | // Dispatch normalization can add extra workgroups past the last batch element; +ss << __str_5; +// 90 | // drop them here (workgroup-uniform, so the barriers below stay uniform). +// 91 | if (batch_id >= uniforms.batch) { +ss << __str_425; +// 92 | return; +ss << __str_426; +// 93 | } +ss << __str_20; +// 94 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +ss << __str_427; +// 95 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile +ss << __str_428; +// 96 | let n_tile_start = n_tile * kTileN; +ss << __str_45; +// 97 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +ss << __str_5; +// 98 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). +// 99 | // Uses N_b (== N unless B was padded for row-major alignment) so the shifted +// 100 | // origin -- itself n_align-even since N_b, kTileN and n_tile_start all are -- +// 101 | // stays alignment-safe for the HWIO load below. +// 102 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); +ss << __str_224; +// 103 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile +ss << __str_205; +// 104 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile +ss << __str_429; +// 105 | let m_base = m_tile * kTileM; +ss << __str_430; +// 106 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); +ss << __str_431; +// 107 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +ss << __str_432; +// 108 | let k_blocks = uniforms.K / kSgMatK; +ss << __str_49; +// 109 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +ss << __str_50; +// 110 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +ss << __str_51; +// 111 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +ss << __str_52; +// 112 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +ss << __str_433; +// 113 | +ss << __str_5; +// 114 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 115 | var sg_mat_c0_0: subgroup_matrix_result; +ss << __str_53; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 116 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 117 | var sg_mat_c0_1: subgroup_matrix_result; +ss << __str_56; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 118 | #endif +} +// 119 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 120 | var sg_mat_c0_2: subgroup_matrix_result; +ss << __str_57; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 121 | #endif +} +// 122 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 123 | var sg_mat_c0_3: subgroup_matrix_result; +ss << __str_58; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 124 | #endif +} +// 125 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 126 | var sg_mat_c1_0: subgroup_matrix_result; +ss << __str_59; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 127 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 128 | var sg_mat_c1_1: subgroup_matrix_result; +ss << __str_60; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 129 | #endif +} +// 130 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 131 | var sg_mat_c1_2: subgroup_matrix_result; +ss << __str_61; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 132 | #endif +} +// 133 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 134 | var sg_mat_c1_3: subgroup_matrix_result; +ss << __str_62; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 135 | #endif +} +// 136 | #endif +} +// 137 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 138 | var sg_mat_c2_0: subgroup_matrix_result; +ss << __str_63; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 139 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 140 | var sg_mat_c2_1: subgroup_matrix_result; +ss << __str_64; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 141 | #endif +} +// 142 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 143 | var sg_mat_c2_2: subgroup_matrix_result; +ss << __str_65; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 144 | #endif +} +// 145 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 146 | var sg_mat_c2_3: subgroup_matrix_result; +ss << __str_66; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 147 | #endif +} +// 148 | #endif +} +// 149 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 150 | var sg_mat_c3_0: subgroup_matrix_result; +ss << __str_67; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 151 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 152 | var sg_mat_c3_1: subgroup_matrix_result; +ss << __str_68; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 153 | #endif +} +// 154 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 155 | var sg_mat_c3_2: subgroup_matrix_result; +ss << __str_69; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 156 | #endif +} +// 157 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 158 | var sg_mat_c3_3: subgroup_matrix_result; +ss << __str_70; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 159 | #endif +} +// 160 | #endif +} +// 161 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 162 | var sg_mat_c4_0: subgroup_matrix_result; +ss << __str_71; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 163 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 164 | var sg_mat_c4_1: subgroup_matrix_result; +ss << __str_72; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 165 | #endif +} +// 166 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 167 | var sg_mat_c4_2: subgroup_matrix_result; +ss << __str_73; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 168 | #endif +} +// 169 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 170 | var sg_mat_c4_3: subgroup_matrix_result; +ss << __str_74; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 171 | #endif +} +// 172 | #endif +} +// 173 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 174 | var sg_mat_c5_0: subgroup_matrix_result; +ss << __str_75; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 175 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 176 | var sg_mat_c5_1: subgroup_matrix_result; +ss << __str_76; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 177 | #endif +} +// 178 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 179 | var sg_mat_c5_2: subgroup_matrix_result; +ss << __str_77; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 180 | #endif +} +// 181 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 182 | var sg_mat_c5_3: subgroup_matrix_result; +ss << __str_78; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 183 | #endif +} +// 184 | #endif +} +// 185 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 186 | var sg_mat_c6_0: subgroup_matrix_result; +ss << __str_79; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 187 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 188 | var sg_mat_c6_1: subgroup_matrix_result; +ss << __str_80; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 189 | #endif +} +// 190 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 191 | var sg_mat_c6_2: subgroup_matrix_result; +ss << __str_81; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 192 | #endif +} +// 193 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 194 | var sg_mat_c6_3: subgroup_matrix_result; +ss << __str_82; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 195 | #endif +} +// 196 | #endif +} +// 197 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 198 | var sg_mat_c7_0: subgroup_matrix_result; +ss << __str_83; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 199 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 200 | var sg_mat_c7_1: subgroup_matrix_result; +ss << __str_84; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 201 | #endif +} +// 202 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 203 | var sg_mat_c7_2: subgroup_matrix_result; +ss << __str_85; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 204 | #endif +} +// 205 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 206 | var sg_mat_c7_3: subgroup_matrix_result; +ss << __str_86; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_55; +// 207 | #endif +} +// 208 | #endif +} +// 209 | +// 210 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, +// 211 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks +// 212 | // early (valid == false), keeping the barriers below in uniform control flow. +// 213 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +ss << __str_434; +// 214 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +ss << __str_435; +// 215 | let kb = sg_index + step * kSplitK; +ss << __str_436; +// 216 | let valid = kb < k_blocks; +ss << __str_437; +// 217 | +ss << __str_5; +// 218 | if (valid) { +ss << __str_438; +// 219 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +ss << __str_439; +// 220 | let k_vec_base = kb * k_vec_per_sgmat_row; +ss << __str_440; +// 221 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { +ss << __str_441; +// 222 | let row = vidx / k_vec_per_sgmat_row; +ss << __str_442; +// 223 | let kc_vec = vidx % k_vec_per_sgmat_row; +ss << __str_443; +// 224 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); +ss << __str_444; +// 225 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +ss << __str_445; +// 226 | #if vec_size == 1 +if (__param_vec_size == 1) { +// 227 | a_scratch[dst] = chunk; +ss << __str_446; +// 228 | #else +} else { +// 229 | a_scratch[dst + 0u] = chunk[0]; +ss << __str_447; +// 230 | a_scratch[dst + 1u] = chunk[1]; +ss << __str_448; +// 231 | #if vec_size >= 4 +if (__param_vec_size >= 4) { +// 232 | a_scratch[dst + 2u] = chunk[2]; +ss << __str_449; +// 233 | a_scratch[dst + 3u] = chunk[3]; +ss << __str_450; +// 234 | #endif +} +// 235 | #endif +} +// 236 | } +ss << __str_202; +// 237 | } +ss << __str_204; +// 238 | // Make the staged A-tiles visible before loading them as subgroup matrices. +ss << __str_5; +// 239 | workgroupBarrier(); +ss << __str_451; +// 240 | +ss << __str_5; +// 241 | if (valid) { +ss << __str_438; +// 242 | // Load the B right tiles for this K block. +ss << __str_5; +// 243 | #if b_is_hwio +if (__param_b_is_hwio) { +// 244 | // HWIO: B is row-major ([K, N_b], B[k, n] at k*N_b + n); load row-major, +// 245 | // stride N_b. N_b (== N unless padded) keeps the row stride, and via +// 246 | // global_base_n above, the column offset, 4-byte aligned. +// 247 | let b_base = kb * kSgMatK * uniforms.N_b + global_base_n; +ss << __str_452; +// 248 | var sg_mat_b0: subgroup_matrix_right = +ss << __str_453; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 249 | subgroupMatrixLoad, row_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_101; +// 250 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +ss << __str_455; +// 251 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 252 | var sg_mat_b1: subgroup_matrix_right = +ss << __str_456; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 253 | subgroupMatrixLoad, row_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_101; +// 254 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); +ss << __str_457; +// 255 | #endif +} +// 256 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 257 | var sg_mat_b2: subgroup_matrix_right = +ss << __str_458; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 258 | subgroupMatrixLoad, row_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_101; +// 259 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +ss << __str_459; +// 260 | #endif +} +// 261 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 262 | var sg_mat_b3: subgroup_matrix_right = +ss << __str_460; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 263 | subgroupMatrixLoad, row_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_101; +// 264 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +ss << __str_461; +// 265 | #endif +} +// 266 | #else +} else { +// 267 | // OHWI: B is column-major ([N, K], B[k, n] at n*K + k); load col-major, stride K. +// 268 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +ss << __str_462; +// 269 | var sg_mat_b0: subgroup_matrix_right = +ss << __str_453; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 270 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 271 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_463; +// 272 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 273 | var sg_mat_b1: subgroup_matrix_right = +ss << __str_456; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 274 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 275 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_464; +// 276 | #endif +} +// 277 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 278 | var sg_mat_b2: subgroup_matrix_right = +ss << __str_458; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 279 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 280 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_465; +// 281 | #endif +} +// 282 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 283 | var sg_mat_b3: subgroup_matrix_right = +ss << __str_460; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 284 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 285 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_466; +// 286 | #endif +} +// 287 | #endif +} +// 288 | +// 289 | // Load the A left tiles (one per M block) from workgroup staging, +// 290 | // row-major with stride kSgMatK. +// 291 | var sg_mat_a0: subgroup_matrix_left = +ss << __str_467; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 292 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 293 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_469; +// 294 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 295 | var sg_mat_a1: subgroup_matrix_left = +ss << __str_470; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 296 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 297 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_471; +// 298 | #endif +} +// 299 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 300 | var sg_mat_a2: subgroup_matrix_left = +ss << __str_472; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 301 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 302 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_473; +// 303 | #endif +} +// 304 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 305 | var sg_mat_a3: subgroup_matrix_left = +ss << __str_474; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 306 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 307 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_475; +// 308 | #endif +} +// 309 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 310 | var sg_mat_a4: subgroup_matrix_left = +ss << __str_476; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 311 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 312 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_477; +// 313 | #endif +} +// 314 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 315 | var sg_mat_a5: subgroup_matrix_left = +ss << __str_478; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 316 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 317 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_479; +// 318 | #endif +} +// 319 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 320 | var sg_mat_a6: subgroup_matrix_left = +ss << __str_480; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 321 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 322 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_481; +// 323 | #endif +} +// 324 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 325 | var sg_mat_a7: subgroup_matrix_left = +ss << __str_482; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 326 | subgroupMatrixLoad, row_major>( +ss << __str_468; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 327 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_483; +// 328 | #endif +} +// 329 | +// 330 | // Accumulate every (M block, N block) pair. +// 331 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +ss << __str_484; +// 332 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 333 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +ss << __str_485; +// 334 | #endif +} +// 335 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 336 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +ss << __str_486; +// 337 | #endif +} +// 338 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 339 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +ss << __str_487; +// 340 | #endif +} +// 341 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 342 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +ss << __str_488; +// 343 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 344 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +ss << __str_489; +// 345 | #endif +} +// 346 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 347 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +ss << __str_490; +// 348 | #endif +} +// 349 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 350 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +ss << __str_491; +// 351 | #endif +} +// 352 | #endif +} +// 353 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 354 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +ss << __str_492; +// 355 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 356 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +ss << __str_493; +// 357 | #endif +} +// 358 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 359 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +ss << __str_494; +// 360 | #endif +} +// 361 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 362 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +ss << __str_495; +// 363 | #endif +} +// 364 | #endif +} +// 365 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 366 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +ss << __str_496; +// 367 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 368 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +ss << __str_497; +// 369 | #endif +} +// 370 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 371 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +ss << __str_498; +// 372 | #endif +} +// 373 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 374 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +ss << __str_499; +// 375 | #endif +} +// 376 | #endif +} +// 377 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 378 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +ss << __str_500; +// 379 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 380 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +ss << __str_501; +// 381 | #endif +} +// 382 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 383 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +ss << __str_502; +// 384 | #endif +} +// 385 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 386 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +ss << __str_503; +// 387 | #endif +} +// 388 | #endif +} +// 389 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 390 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +ss << __str_504; +// 391 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 392 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +ss << __str_505; +// 393 | #endif +} +// 394 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 395 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +ss << __str_506; +// 396 | #endif +} +// 397 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 398 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +ss << __str_507; +// 399 | #endif +} +// 400 | #endif +} +// 401 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 402 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +ss << __str_508; +// 403 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 404 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +ss << __str_509; +// 405 | #endif +} +// 406 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 407 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +ss << __str_510; +// 408 | #endif +} +// 409 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 410 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +ss << __str_511; +// 411 | #endif +} +// 412 | #endif +} +// 413 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 414 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +ss << __str_512; +// 415 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 416 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +ss << __str_513; +// 417 | #endif +} +// 418 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 419 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +ss << __str_514; +// 420 | #endif +} +// 421 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 422 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +ss << __str_515; +// 423 | #endif +} +// 424 | #endif +} +// 425 | } +ss << __str_204; +// 426 | // Ensure every subgroup is done reading a_scratch before it is reused for +ss << __str_5; +// 427 | // the next K-block step. +// 428 | workgroupBarrier(); +ss << __str_451; +// 429 | } +ss << __str_20; +// 430 | +ss << __str_5; +// 431 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; +// 432 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). +// 433 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +ss << __str_165; +// 434 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 435 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +ss << __str_166; +// 436 | #endif +} +// 437 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 438 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +ss << __str_167; +// 439 | #endif +} +// 440 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 441 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +ss << __str_168; +// 442 | #endif +} +// 443 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 444 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +ss << __str_169; +// 445 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 446 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +ss << __str_170; +// 447 | #endif +} +// 448 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 449 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +ss << __str_171; +// 450 | #endif +} +// 451 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +ss << __str_172; +// 453 | #endif +} +// 454 | #endif +} +// 455 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 456 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +ss << __str_173; +// 457 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +ss << __str_174; +// 459 | #endif +} +// 460 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 461 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +ss << __str_175; +// 462 | #endif +} +// 463 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +ss << __str_176; +// 465 | #endif +} +// 466 | #endif +} +// 467 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 468 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +ss << __str_177; +// 469 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 470 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +ss << __str_178; +// 471 | #endif +} +// 472 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 473 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +ss << __str_179; +// 474 | #endif +} +// 475 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 476 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +ss << __str_180; +// 477 | #endif +} +// 478 | #endif +} +// 479 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 480 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +ss << __str_181; +// 481 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 482 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +ss << __str_182; +// 483 | #endif +} +// 484 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 485 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +ss << __str_183; +// 486 | #endif +} +// 487 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 488 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +ss << __str_184; +// 489 | #endif +} +// 490 | #endif +} +// 491 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 492 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +ss << __str_185; +// 493 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 494 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +ss << __str_186; +// 495 | #endif +} +// 496 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 497 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +ss << __str_187; +// 498 | #endif +} +// 499 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 500 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +ss << __str_188; +// 501 | #endif +} +// 502 | #endif +} +// 503 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 504 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +ss << __str_189; +// 505 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 506 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +ss << __str_190; +// 507 | #endif +} +// 508 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 509 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +ss << __str_191; +// 510 | #endif +} +// 511 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 512 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +ss << __str_192; +// 513 | #endif +} +// 514 | #endif +} +// 515 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 516 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +ss << __str_193; +// 517 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 518 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +ss << __str_194; +// 519 | #endif +} +// 520 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 521 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +ss << __str_195; +// 522 | #endif +} +// 523 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 524 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +ss << __str_196; +// 525 | #endif +} +// 526 | #endif +} +// 527 | +// 528 | #if split_k >= 2 +if (__param_split_k >= 2) { +// 529 | // Make all subgroups' partial tiles visible before the reduction. +// 530 | workgroupBarrier(); +ss << __str_26; +// 531 | +ss << __str_5; +// 532 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each +// 533 | // position is only overwritten by the lane that already read it. +// 534 | if (sg_index == 0u) { +ss << __str_197; +// 535 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +ss << __str_198; +// 536 | var acc: f16 = f16(0); +ss << __str_199; +// 537 | for (var s: u32 = 0; s < kSplitK; s++) { +ss << __str_200; +// 538 | acc += scratch[s * kTileM * kTileN + idx]; +ss << __str_201; +// 539 | } +ss << __str_202; +// 540 | scratch[idx] = acc; +ss << __str_203; +// 541 | } +ss << __str_204; +// 542 | } +ss << __str_20; +// 543 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +ss << __str_5; +// 544 | workgroupBarrier(); +ss << __str_26; +// 545 | #endif +} +// 546 | +// 547 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip +// 548 | // and M's bounds check make non-multiple dimensions write each element once. +// 549 | let n_count = min(kTileN, uniforms.N - global_base_n); +ss << __str_206; +// 550 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +ss << __str_207; +// 551 | let global_m = m_base + r; +ss << __str_208; +// 552 | if (global_m < uniforms.M) { +ss << __str_516; +// 553 | let scratch_base = r * kTileN; +ss << __str_210; +// 554 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; +ss << __str_517; +// 555 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +ss << __str_212; +// 556 | write_output(out_base + i, global_base_n + i, +ss << __str_235; +// 557 | output_value_t(scratch[scratch_base + i])); +ss << __str_236; +// 558 | } +ss << __str_202; +// 559 | } +ss << __str_204; +// 560 | } +ss << __str_20; +// 561 | } // MAIN +MainFunctionEnd(); +ss << __str_5; +// 562 | + + + return Status::OK(); +} \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h index 06e403c405f07..e26fbe2247ed1 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h @@ -19,40 +19,40 @@ Status ApplyTemplate<"tensor/oihw_to_ohwi.wgsl.template">(ShaderHelper& shader_h // 4 | #use .getByOffset .setByOffset // 5 | // 6 | fn load_src(co : u32, ci : u32, h_w : u32) -> src_element_t { -ss << __str_407; +ss << __str_519; // 7 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_408; +ss << __str_520; // 8 | let offset = co * uniforms.I * uniforms.H * uniforms.W + -ss << __str_409; +ss << __str_521; // 9 | ci * uniforms.H * uniforms.W + -ss << __str_410; +ss << __str_522; // 10 | h_w; -ss << __str_411; +ss << __str_523; // 11 | return src.getByOffset(offset); ss << __str_351; -ss << __var_src->GetByOffset(__str_406); +ss << __var_src->GetByOffset(__str_518); ss << __str_29; // 12 | } ss << __str_11; // 13 | return src_element_t(); -ss << __str_412; +ss << __str_524; // 14 | } ss << __str_345; // 15 | ss << __str_5; // 16 | fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) { -ss << __str_413; +ss << __str_525; // 17 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_408; +ss << __str_520; // 18 | let offset = co * uniforms.H * uniforms.W * uniforms.I + -ss << __str_414; +ss << __str_526; // 19 | h_w * uniforms.I + -ss << __str_415; +ss << __str_527; // 20 | ci; -ss << __str_416; +ss << __str_528; // 21 | output.setByOffset(offset, value); ss << __str_28; -ss << __var_output->SetByOffset(__str_406, __str_322); +ss << __var_output->SetByOffset(__str_518, __str_322); ss << __str_29; // 22 | } ss << __str_11; @@ -61,20 +61,20 @@ ss << __str_345; // 24 | ss << __str_5; // 25 | var data_cache : array, 4>; -ss << __str_417; +ss << __str_529; // 26 | ss << __str_5; // 27 | $MAIN { MainFunctionStart(); ss << __str_5; // 28 | let group_co : u32 = workgroup_idx / uniforms.Ci_tiles; -ss << __str_418; +ss << __str_530; // 29 | let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64; -ss << __str_419; +ss << __str_531; // 30 | ss << __str_5; // 31 | if (group_co >= uniforms.O || group_ci >= uniforms.I) { -ss << __str_420; +ss << __str_532; // 32 | return; ss << __str_10; // 33 | } @@ -82,23 +82,23 @@ ss << __str_11; // 34 | ss << __str_5; // 35 | for (var h_w_idx = 0u; h_w_idx < uniforms.H_W_tiles; h_w_idx++) { -ss << __str_421; +ss << __str_533; // 36 | // load ss << __str_5; // 37 | for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) { -ss << __str_422; +ss << __str_534; // 38 | let load_ci_idx = ci_idx + local_idx / 4; -ss << __str_423; +ss << __str_535; // 39 | let load_h_w_idx = local_idx % 4; -ss << __str_424; +ss << __str_536; // 40 | ss << __str_5; // 41 | data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co, -ss << __str_425; +ss << __str_537; // 42 | group_ci + load_ci_idx, -ss << __str_426; +ss << __str_538; // 43 | h_w_idx * 4 + load_h_w_idx); -ss << __str_427; +ss << __str_539; // 44 | } ss << __str_20; // 45 | workgroupBarrier(); @@ -107,11 +107,11 @@ ss << __str_26; ss << __str_5; // 47 | // store // 48 | for (var local_h_w_idx = 0u; local_h_w_idx < 4u; local_h_w_idx++) { -ss << __str_428; +ss << __str_540; // 49 | let output_data = data_cache[local_h_w_idx][local_idx]; -ss << __str_429; +ss << __str_541; // 50 | write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data); -ss << __str_430; +ss << __str_542; // 51 | } ss << __str_20; // 52 | workgroupBarrier(); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h index a40daa675052e..4a2df66b174e6 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h @@ -40,22 +40,22 @@ ss << __str_29; // 16 | ss << __str_5; // 17 | let constant_value = -ss << __str_436; +ss << __str_548; // 18 | #if is_float16 if (__param_is_float16) { // 19 | bitcast>(uniforms.constant_value)[0]; -ss << __str_437; +ss << __str_549; // 20 | #else } else { // 21 | bitcast(uniforms.constant_value); -ss << __str_438; +ss << __str_550; // 22 | #endif } // 23 | // 24 | #if dim_value_zero if (__param_dim_value_zero) { // 25 | output[global_idx] = constant_value; -ss << __str_439; +ss << __str_551; // 26 | #else } else { // 27 | let output_indices = output.offsetToIndices(global_idx); @@ -63,83 +63,83 @@ ss << __str_265; ss << __var_output->OffsetToIndices(__str_239); ss << __str_29; // 28 | var input_index = u32(0); -ss << __str_440; +ss << __str_552; // 29 | var use_pad_value = false; -ss << __str_441; +ss << __str_553; // 30 | var in_coord = i32(0); -ss << __str_442; +ss << __str_554; // 31 | ss << __str_5; // 32 | for (var dim = 0; dim < output.rank && !use_pad_value; dim++) { -ss << __str_443; +ss << __str_555; ss << __var_output->Rank(); -ss << __str_444; +ss << __str_556; // 33 | let output_index = i32(getElementAt(output_indices, dim, output.rank)); -ss << __str_445; -ss << GetElementAt(__str_247, __str_431, __var_output->Rank()); +ss << __str_557; +ss << GetElementAt(__str_247, __str_543, __var_output->Rank()); ss << __str_17; // 34 | let lower_pads = getElementAt(uniforms.lower_pads, dim, output.rank); -ss << __str_446; -ss << GetElementAt(__str_432, __str_431, __var_output->Rank()); +ss << __str_558; +ss << GetElementAt(__str_544, __str_543, __var_output->Rank()); ss << __str_29; // 35 | let data_shape = i32(getElementAt(uniforms.data_shape, dim, output.rank)); -ss << __str_447; -ss << GetElementAt(__str_433, __str_431, __var_output->Rank()); +ss << __str_559; +ss << GetElementAt(__str_545, __str_543, __var_output->Rank()); ss << __str_17; // 36 | #if pad_mode == PAD_MODE_CONSTANT if (__param_pad_mode == 0) { // 37 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { -ss << __str_448; +ss << __str_560; // 38 | use_pad_value = true; -ss << __str_449; +ss << __str_561; // 39 | #elif pad_mode == PAD_MODE_EDGE } else if (__param_pad_mode == 2) { // 40 | if (output_index < lower_pads) { -ss << __str_450; +ss << __str_562; // 41 | in_coord = 0; -ss << __str_451; +ss << __str_563; // 42 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_452; +ss << __str_564; // 43 | in_coord = data_shape - 1; -ss << __str_453; +ss << __str_565; // 44 | #elif pad_mode == PAD_MODE_REFLECT } else if (__param_pad_mode == 1) { // 45 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { -ss << __str_448; +ss << __str_560; // 46 | in_coord = output_index - lower_pads; -ss << __str_454; +ss << __str_566; // 47 | if (in_coord < 0) { -ss << __str_455; +ss << __str_567; // 48 | in_coord = -in_coord; -ss << __str_456; +ss << __str_568; // 49 | } ss << __str_292; // 50 | let _2n_1 = 2 * (data_shape - 1); -ss << __str_457; +ss << __str_569; // 51 | in_coord = in_coord % _2n_1; -ss << __str_458; +ss << __str_570; // 52 | if (in_coord >= data_shape) { -ss << __str_459; +ss << __str_571; // 53 | in_coord = _2n_1 - in_coord; -ss << __str_460; +ss << __str_572; // 54 | } ss << __str_292; // 55 | #else // PAD_MODE_WRAP } else { // 56 | if (output_index < lower_pads) { -ss << __str_450; +ss << __str_562; // 57 | in_coord = data_shape + output_index - lower_pads; -ss << __str_461; +ss << __str_573; // 58 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_452; +ss << __str_564; // 59 | in_coord = output_index - data_shape - lower_pads; -ss << __str_462; +ss << __str_574; // 60 | #endif // pad_mode } // 61 | } else { -ss << __str_463; +ss << __str_575; // 62 | in_coord = output_index - lower_pads; -ss << __str_464; +ss << __str_576; // 63 | } ss << __str_20; // 64 | @@ -147,31 +147,31 @@ ss << __str_5; // 65 | #if pad_mode == PAD_MODE_WRAP if (__param_pad_mode == 3) { // 66 | in_coord = ((in_coord % data_shape) + data_shape) % data_shape; -ss << __str_465; +ss << __str_577; // 67 | #endif } // 68 | // 69 | input_index += select(u32(in_coord) -ss << __str_466; +ss << __str_578; // 70 | #if output.rank > 1 if (__var_output->Rank() > 1) { // 71 | * getElementAt(uniforms.data_stride, dim, output.rank - 1) -ss << __str_467; -ss << GetElementAt(__str_434, __str_431, __var_output->Rank() - 1); +ss << __str_579; +ss << GetElementAt(__str_546, __str_543, __var_output->Rank() - 1); ss << __str_5; // 72 | #endif } // 73 | , u32(in_coord), dim == output.rank - 1); -ss << __str_468; +ss << __str_580; ss << __var_output->Rank(); -ss << __str_469; +ss << __str_581; // 74 | } ss << __str_11; // 75 | ss << __str_5; // 76 | output.setByOffset(global_idx, select(data[input_index], constant_value, use_pad_value)); ss << __str_241; -ss << __var_output->SetByOffset(__str_239, __str_435); +ss << __var_output->SetByOffset(__str_239, __str_547); ss << __str_29; // 77 | #endif } diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h index 3e837a94da226..fe0ae419c99f4 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h @@ -121,6 +121,28 @@ struct TemplateParameter<"nn/im2col_matmul.wgsl.template"> { template <> Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/im2col_matmul.wgsl.template">::type params); +// +// Template: nn/subgroup_matrix_conv_8x16x16.wgsl.template +// + +template <> +struct TemplateParameter<"nn/subgroup_matrix_conv_8x16x16.wgsl.template"> { + using type = struct { + int param_b_is_hwio; + int param_sg_mat_count_m; + int param_sg_mat_count_n; + int param_sg_mat_k; + int param_sg_mat_m; + int param_sg_mat_n; + int param_split_k; + int param_vec_size; + const ShaderVariableHelper* var_input_src; + }; +}; + +template <> +Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelper& shader_helper, TemplateParameter<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">::type params); + // // Template: tensor/oihw_to_ohwi.wgsl.template // diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h index 3bb8396be1a50..c69062b7879c1 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h @@ -31,18 +31,19 @@ std::string pass_as_string(T&& v) { } } // namespace wgsl_detail -#include "wgsl_template_gen/string_table.h" // 0164a7040addea01a1f6e9ff048e816229c3b38c2d8eee7625c241df23ef3072 +#include "wgsl_template_gen/string_table.h" // 8bfc19bf0715588a4f8171d90b06809ccad9f7e99fd3cce7fba664bb204d9004 // Include template implementations #include "wgsl_template_gen/generated/math/matmul_gemv.h" // 512f5afa8421c4630912b0d87d46061d43e1b3226fbbc6c993241b285adf2fc7 #include "wgsl_template_gen/generated/math/subgroup_matrix_gemm_8x16x16.h" // 4251accac15bcce4e64a1ec7582de07ee5626d6cfe05527a9187586ade097870 #include "wgsl_template_gen/generated/math/subgroup_matrix_matmul.h" // 82baad53c5eb895073a12f45a2bb5ff6e073bbe089448e0d5ba52255c4ee7ce7 -#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // 4b28d0cbe9ca50020bbcb75121e787982c268992efdac401b821dee7c59c604e +#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // 8d25b9aab5a8f6d1d557f4c8efceb1bd44cdb2592a8483b396df8888c683a80c #include "wgsl_template_gen/generated/nn/conv_transpose3d.h" // 33bd52e16ae1e748ef45b492ae4730f520eeb9706cab055391d296ec54265bab #include "wgsl_template_gen/generated/nn/im2col_matmul.h" // 7185c8aeed8391c916822b35d16fcbfdd342e4c3cc5ea9a4d09667a2651a3c08 -#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // 473c95130b9ad534439ed746bf37df45be692792142d57c48db3cea0e864f352 -#include "wgsl_template_gen/generated/tensor/pad.h" // b60f1767797524bcf834535f8e4f8d214c0c8a6f6d7c4e12ceae44c39abec554 +#include "wgsl_template_gen/generated/nn/subgroup_matrix_conv_8x16x16.h" // 398f225f4b34778fc8826d92cd6952806fa231f57aeaf840bfa98cc47e483c1c +#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // f6adc81e9e253de771363ee54143064bffbefd468498f9e85e7a0f028d088678 +#include "wgsl_template_gen/generated/tensor/pad.h" // 598f8f41c0336b04fab5705643f439a307d059fd8971812c207a2273b6976518 #pragma pop_macro("MainFunctionStart") #pragma pop_macro("MainFunctionEnd") \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h index 9764cea3df588..4b64335feb014 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h @@ -412,67 +412,179 @@ constexpr const char* __str_402 = " output_data = tanh(output_data);\n"; constexpr const char* __str_403 = " output_data = output_data * (output_element_t(1) / (output_element_t(1) + exp(-output_data)));\n"; constexpr const char* __str_404 = " output_data = output_data * (output_element_t(1) / (output_element_t(1) + exp(-(output_element_t(uniforms.activation_param_0) * output_data))));\n"; constexpr const char* __str_405 = " write_output(batch, m_base + m_idx, n_base, output_data);\n"; -constexpr const char* __str_406 = "offset"; -constexpr const char* __str_407 = "fn load_src(co : u32, ci : u32, h_w : u32) -> src_element_t {\n"; -constexpr const char* __str_408 = " if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) {\n"; -constexpr const char* __str_409 = " let offset = co * uniforms.I * uniforms.H * uniforms.W +\n"; -constexpr const char* __str_410 = " ci * uniforms.H * uniforms.W +\n"; -constexpr const char* __str_411 = " h_w;\n"; -constexpr const char* __str_412 = " return src_element_t();\n"; -constexpr const char* __str_413 = "fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) {\n"; -constexpr const char* __str_414 = " let offset = co * uniforms.H * uniforms.W * uniforms.I +\n"; -constexpr const char* __str_415 = " h_w * uniforms.I +\n"; -constexpr const char* __str_416 = " ci;\n"; -constexpr const char* __str_417 = "var data_cache : array, 4>;\n"; -constexpr const char* __str_418 = " let group_co : u32 = workgroup_idx / uniforms.Ci_tiles;\n"; -constexpr const char* __str_419 = " let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64;\n"; -constexpr const char* __str_420 = " if (group_co >= uniforms.O || group_ci >= uniforms.I) {\n"; -constexpr const char* __str_421 = " for (var h_w_idx = 0u; h_w_idx < uniforms.H_W_tiles; h_w_idx++) {\n"; -constexpr const char* __str_422 = " for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) {\n"; -constexpr const char* __str_423 = " let load_ci_idx = ci_idx + local_idx / 4;\n"; -constexpr const char* __str_424 = " let load_h_w_idx = local_idx % 4;\n"; -constexpr const char* __str_425 = " data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co,\n"; -constexpr const char* __str_426 = " group_ci + load_ci_idx,\n"; -constexpr const char* __str_427 = " h_w_idx * 4 + load_h_w_idx);\n"; -constexpr const char* __str_428 = " for (var local_h_w_idx = 0u; local_h_w_idx < 4u; local_h_w_idx++) {\n"; -constexpr const char* __str_429 = " let output_data = data_cache[local_h_w_idx][local_idx];\n"; -constexpr const char* __str_430 = " write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data);\n"; -constexpr const char* __str_431 = "dim"; -constexpr const char* __str_432 = "uniforms.lower_pads"; -constexpr const char* __str_433 = "uniforms.data_shape"; -constexpr const char* __str_434 = "uniforms.data_stride"; -constexpr const char* __str_435 = "select(data[input_index], constant_value, use_pad_value)"; -constexpr const char* __str_436 = " let constant_value =\n"; -constexpr const char* __str_437 = " bitcast>(uniforms.constant_value)[0];\n"; -constexpr const char* __str_438 = " bitcast(uniforms.constant_value);\n"; -constexpr const char* __str_439 = " output[global_idx] = constant_value;\n"; -constexpr const char* __str_440 = " var input_index = u32(0);\n"; -constexpr const char* __str_441 = " var use_pad_value = false;\n"; -constexpr const char* __str_442 = " var in_coord = i32(0);\n"; -constexpr const char* __str_443 = " for (var dim = 0; dim < "; -constexpr const char* __str_444 = " && !use_pad_value; dim++) {\n"; -constexpr const char* __str_445 = " let output_index = i32("; -constexpr const char* __str_446 = " let lower_pads = "; -constexpr const char* __str_447 = " let data_shape = i32("; -constexpr const char* __str_448 = " if (output_index < lower_pads || output_index >= data_shape + lower_pads) {\n"; -constexpr const char* __str_449 = " use_pad_value = true;\n"; -constexpr const char* __str_450 = " if (output_index < lower_pads) {\n"; -constexpr const char* __str_451 = " in_coord = 0;\n"; -constexpr const char* __str_452 = " } else if (output_index >= data_shape + lower_pads) {\n"; -constexpr const char* __str_453 = " in_coord = data_shape - 1;\n"; -constexpr const char* __str_454 = " in_coord = output_index - lower_pads;\n"; -constexpr const char* __str_455 = " if (in_coord < 0) {\n"; -constexpr const char* __str_456 = " in_coord = -in_coord;\n"; -constexpr const char* __str_457 = " let _2n_1 = 2 * (data_shape - 1);\n"; -constexpr const char* __str_458 = " in_coord = in_coord % _2n_1;\n"; -constexpr const char* __str_459 = " if (in_coord >= data_shape) {\n"; -constexpr const char* __str_460 = " in_coord = _2n_1 - in_coord;\n"; -constexpr const char* __str_461 = " in_coord = data_shape + output_index - lower_pads;\n"; -constexpr const char* __str_462 = " in_coord = output_index - data_shape - lower_pads;\n"; -constexpr const char* __str_463 = " } else {\n"; -constexpr const char* __str_464 = " in_coord = output_index - lower_pads;\n"; -constexpr const char* __str_465 = " in_coord = ((in_coord % data_shape) + data_shape) % data_shape;\n"; -constexpr const char* __str_466 = " input_index += select(u32(in_coord)\n"; -constexpr const char* __str_467 = " * "; -constexpr const char* __str_468 = " , u32(in_coord), dim == "; -constexpr const char* __str_469 = " - 1);\n"; +constexpr const char* __str_406 = "input_vec_idx"; +constexpr const char* __str_407 = "const kVecSize: u32 = u32("; +constexpr const char* __str_408 = "var a_scratch: array;\n"; +constexpr const char* __str_409 = "fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t {\n"; +constexpr const char* __str_410 = " let cin_vec = uniforms.cin / kVecSize;\n"; +constexpr const char* __str_411 = " let oh = m / uniforms.output_w;\n"; +constexpr const char* __str_412 = " let ow = m % uniforms.output_w;\n"; +constexpr const char* __str_413 = " let k_cin = k_vec % cin_vec;\n"; +constexpr const char* __str_414 = " let k_hw = k_vec / cin_vec;\n"; +constexpr const char* __str_415 = " let kw_idx = k_hw % uniforms.kernel_w;\n"; +constexpr const char* __str_416 = " let kh_idx = k_hw / uniforms.kernel_w;\n"; +constexpr const char* __str_417 = " let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x);\n"; +constexpr const char* __str_418 = " let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y);\n"; +constexpr const char* __str_419 = " if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) {\n"; +constexpr const char* __str_420 = " return input_src_value_t();\n"; +constexpr const char* __str_421 = " let input_vec_idx = input_batch_base_vec +\n"; +constexpr const char* __str_422 = " (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin;\n"; +constexpr const char* __str_423 = " let tiles_per_batch = uniforms.num_n_tile * num_m_tile;\n"; +constexpr const char* __str_424 = " let batch_id = workgroup_idx / tiles_per_batch;\n"; +constexpr const char* __str_425 = " if (batch_id >= uniforms.batch) {\n"; +constexpr const char* __str_426 = " return;\n"; +constexpr const char* __str_427 = " let tile_idx = workgroup_idx % tiles_per_batch;\n"; +constexpr const char* __str_428 = " let n_tile = tile_idx % uniforms.num_n_tile;\n"; +constexpr const char* __str_429 = " let m_tile = tile_idx / uniforms.num_n_tile;\n"; +constexpr const char* __str_430 = " let m_base = m_tile * kTileM;\n"; +constexpr const char* __str_431 = " let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize);\n"; +constexpr const char* __str_432 = " let out_batch_base = batch_id * uniforms.M * uniforms.N;\n"; +constexpr const char* __str_433 = " let a_scratch_base = sg_index * kTileM * kSgMatK;\n"; +constexpr const char* __str_434 = " let num_steps = (k_blocks + kSplitK - 1u) / kSplitK;\n"; +constexpr const char* __str_435 = " for (var step: u32 = 0u; step < num_steps; step = step + 1u) {\n"; +constexpr const char* __str_436 = " let kb = sg_index + step * kSplitK;\n"; +constexpr const char* __str_437 = " let valid = kb < k_blocks;\n"; +constexpr const char* __str_438 = " if (valid) {\n"; +constexpr const char* __str_439 = " let k_vec_per_sgmat_row = kSgMatK / kVecSize;\n"; +constexpr const char* __str_440 = " let k_vec_base = kb * k_vec_per_sgmat_row;\n"; +constexpr const char* __str_441 = " for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) {\n"; +constexpr const char* __str_442 = " let row = vidx / k_vec_per_sgmat_row;\n"; +constexpr const char* __str_443 = " let kc_vec = vidx % k_vec_per_sgmat_row;\n"; +constexpr const char* __str_444 = " let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec);\n"; +constexpr const char* __str_445 = " let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize;\n"; +constexpr const char* __str_446 = " a_scratch[dst] = chunk;\n"; +constexpr const char* __str_447 = " a_scratch[dst + 0u] = chunk[0];\n"; +constexpr const char* __str_448 = " a_scratch[dst + 1u] = chunk[1];\n"; +constexpr const char* __str_449 = " a_scratch[dst + 2u] = chunk[2];\n"; +constexpr const char* __str_450 = " a_scratch[dst + 3u] = chunk[3];\n"; +constexpr const char* __str_451 = " workgroupBarrier();\n"; +constexpr const char* __str_452 = " let b_base = kb * kSgMatK * uniforms.N_b + global_base_n;\n"; +constexpr const char* __str_453 = " var sg_mat_b0: subgroup_matrix_right