From 3702f0f77081356c9ebe0394a6de8d05e271d0f8 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Fri, 11 Sep 2026 09:46:14 +0800 Subject: [PATCH 01/14] webgpu: using subgroup matrix to optimize Conv operator --- onnxruntime/core/providers/webgpu/nn/conv.cc | 57 + onnxruntime/core/providers/webgpu/nn/conv.h | 63 + .../webgpu/nn/subgroup_matrix_conv.cc | 541 +++++ .../webgpu/nn/subgroup_matrix_conv.h | 123 ++ ...subgroup_matrix_conv_8x16x16.wgsl.template | 606 ++++++ .../math/subgroup_matrix_tiling_selector.cc | 162 +- .../math/subgroup_matrix_tiling_selector.h | 8 + .../intel/nn/subgroup_matrix_conv_tuned.inc | 1871 +++++++++++++++++ .../webgpu/subgroup_matrix_conv_test.cc | 351 ++++ 9 files changed, 3776 insertions(+), 6 deletions(-) create mode 100644 onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc create mode 100644 onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h create mode 100644 onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template create mode 100644 onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc create mode 100644 onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index 4b7937e02f240..a73388503980e 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -1,9 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. + #include "core/providers/webgpu/nn/conv.h" + +#include + #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" @@ -27,8 +32,31 @@ Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const Tens return Transpose::DoTranspose(context, perm, reshaped_kernel, *transposed_kernel); } +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 the subgroup-matrix config: implicit-GEMM + // Conv (fused im2col) via subgroupMatrixMultiplyAccumulate. Created once on first + // use (it needs a device query), then it decides inside Compute whether the + // kernel/attrs qualify. It declines (handled=false, allocating nothing) otherwise + // and we fall through to the normal path order below. + 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); @@ -322,6 +350,34 @@ Status Conv::PrePackInternal(ComputeContextBase& con return Status::OK(); } + // Subgroup-matrix Conv path: it reads the weight in its original OIHW layout (it + // applies its own OIHW -> OHWI transpose at runtime to build the GEMM right + // operand), so the weight must not be prepacked at all. + // + // Must come before the Im2ColMatMul branch below: ComputeInternal tries the + // subgroup-matrix path first, and on the devices this path targets + // CanApplyIm2ColMatMulProgram() accepts the same problems. If im2col prepacked + // first, SubgroupMatrixConvImpl::Compute would see a non-null prepacked_kernel_ and + // decline on every constant-weight Conv -- i.e. the path would silently never run. + // + // Placed before the auto_pad check below for the same reason the im2col branch is: + // CanApplySubgroupMatrixConv() only looks at the adapter, dtype, layout, group and + // the weight shape -- never at pads -- so its answer cannot change at runtime, and + // ComputeInternal tests it before every pads-dependent branch. + // + // Skipping the prepack is always safe: is_packed stays false, so ORT keeps input 1 + // in OIHW and every fallback path can still transpose it at runtime if the + // subgroup-matrix path ends up declining. Once CanApplySubgroupMatrixConv has said + // yes, the only things Compute can still decline on are the ones the weight alone + // cannot show: a degenerate empty batch/M/N, or fewer than two spatial entries in + // dilations/pads/strides. Its other two decline paths are unreachable from here -- + // the vendor selector's K and N conditions are exactly the ones checked above, and + // its tile_n > N guard cannot fire because the selector clamps tile_n to N. + if (CanApplySubgroupMatrixConv(context, is_channels_last, kernel_shape, + onnxruntime::narrow(conv_attrs_.group), tensor.DataType())) { + return Status::OK(); + } + // Im2ColMatMul path: transpose OIHW -> OHWI once here instead of on every inference. // // Placed before the auto_pad check below on purpose: @@ -367,6 +423,7 @@ Status Conv::PrePackInternal(ComputeContextBase& con // Analyze execution paths in ComputeInternal to determine if kernel transpose is needed: // + // 0. Subgroup-matrix Conv path: handled above (not prepacked -- stays OIHW) // 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 a64a206f1ef34..896d0bb805418 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -3,6 +3,9 @@ #pragma once +#include +#include + #include "core/common/optional.h" #include "core/providers/webgpu/webgpu_kernel.h" #include "core/providers/cpu/nn/conv_attributes.h" @@ -27,10 +30,35 @@ enum class KernelLayout { template class Conv : public WebGpuKernel { public: + // Abstract base class for alternative optimized Conv implementations (currently the + // subgroup-matrix path). Implementations hold the parent Conv and read + // conv_attrs/activation/... from it, mirroring MatMul::MatMulOptImpl. + class ConvOptImpl { + public: + explicit ConvOptImpl(const Conv& parent) : parent_(parent) {} + virtual ~ConvOptImpl() = default; + + // Attempts the optimized path, reading the Conv operands from `context` and the + // Conv attributes from the parent. Sets handled=true when it ran; leaves + // handled=false (allocating nothing) so the caller falls back to the normal Conv + // path. Called before ComputeInternal has resolved auto_pad or allocated the + // output, so an implementation that needs those runs its own shape inference. + 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) { ORT_ENFORCE(GetFusedActivationAttr(info, activation_).IsOK()); } + // Whether the weight input (index 1) is a constant initializer. An optimized + // implementation uses this to decide whether it may cache a derived form of the + // weight (e.g. a transposed copy) across Runs; a non-constant weight can change + // between Runs, so anything derived from it must be rebuilt every time. + const Tensor* w = nullptr; + w_is_constant_ = info.TryGetConstantInput(1, &w); } Status ComputeInternal(ComputeContext& context) const override; @@ -40,6 +68,15 @@ class Conv : public WebGpuKernel { AllocatorPtr alloc, /*out*/ bool& is_packed) override; + // State an optimized implementation (ConvOptImpl) reads from its parent Conv. + const ConvAttributes& ConvAttrs() const { return conv_attrs_; } + const Activation& ConvActivation() const { return activation_; } + // True when input 1 (the weight) is a constant initializer. See w_is_constant_. + bool IsWeightConstant() const { return w_is_constant_; } + // The prepacked weight, or null when it is not prepacked. Its layout is HWIO or OHWI, + // as reported by kernel_layout_. + const Tensor* PrepackedKernel() const { return prepacked_kernel_.get(); } + protected: ConvAttributes conv_attrs_; Activation activation_; @@ -50,6 +87,32 @@ class Conv : public WebGpuKernel { // is set, otherwise input 1. Stays `OIHW` while `prepacked_kernel_` is null. KernelLayout kernel_layout_{KernelLayout::OIHW}; mutable MatMulOptImplCache matmul_compute_cache_; + + bool w_is_constant_ = false; // whether the weight input (index 1) is a constant initializer + + private: + // Owns the alternative optimized implementation and builds it on first use, + // mirroring MatMulOptImplCache. The impl cannot be built in the ctor because it + // needs the device capabilities, which are only reachable through a compute + // context; and the same Conv kernel can see concurrent Compute calls, so creation + // is guarded by call_once. A null return after initialization means this device + // has no optimized path. Nested rather than a free class like MatMulOptImplCache + // only because ConvOptImpl is itself nested in this template. + class ConvOptImplCache { + public: + ConvOptImplCache() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ConvOptImplCache); + + // Defined in conv.cc: the factory it calls is declared in subgroup_matrix_conv.h, + // which includes this header. + 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); 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..483e7fd215dfb --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -0,0 +1,541 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/nn/subgroup_matrix_conv.h" + +#include +#include +#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/nn/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 { + +// Lanes per subgroup assumed by the subgroup-matrix kernel. The workgroup runs +// split_k subgroups, so its size is kSubgroupMatrixSubgroupSize * split_k. +constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; + +// Combined workgroup-memory (SLM) budget in f16 elements for the two scratch +// buffers the conv kernel uses: the split-K output tiles +// (split_k * tile_m * tile_n) plus the per-subgroup im2col A staging +// (split_k * tile_m * sg_mat_k). 16384 f16 == 32 KB, the common WebGPU +// maxComputeWorkgroupStorageSize. split_k is clamped to fit this. +constexpr uint32_t kMaxConvScratchElems = 16384; + +// Largest per-tile subgroup-matrix counts the WGSL template unrolls: it generates the +// accumulators sg_mat_c_ for mi < 8 and ni < 4. A tiling above these would be +// accepted by the selector but silently mis-computed -- the kernel would accumulate +// only the blocks it generated while the write-out still covers the whole tile. The +// tile candidates the selector picks from stay inside this, but the Conv pretuned tree +// is a generated file dropped into the build after a sweep, so enforce rather than +// assume. +constexpr uint32_t kMaxSgMatCountM = 8; +constexpr uint32_t kMaxSgMatCountN = 4; + +// Returns the device config index for the fixed 8x16x16 F16 subgroup-matrix +// config this kernel is implemented for, or -1 when the device does not report +// it. Also requires subgroup-size control and the vendor tiling policy (Intel). +// +// Single source of truth for device eligibility: CanApplySubgroupMatrixConv (which +// Conv::PrePackInternal consults to skip prepacking) and CreateSubgroupMatrixConvImpl +// both go through here, so the prepack decision and the runtime path can never +// disagree about whether this device is served. +int32_t GetSupportedConfigIndex(const ComputeContextBase& context) { + int32_t config_index = 0; + if (!IsSubgroupMatrixConfigSupported(context, /*is_fp16=*/true, config_index) || + !supported_subgroup_matrix_configs[config_index].Is(8, 16, 16)) { + return -1; + } + // The adapters reporting that config expose a 16-32 subgroup size range, while the + // kernel hard-codes kSubgroupMatrixSubgroupSize lanes per subgroup (sg_index = + // local_idx / 32). A driver-chosen 16 still computes the right answer -- the two + // hardware subgroups covered by one logical subgroup redundantly load, multiply and + // store the identical tile -- but it wastes about half the subgroup-matrix + // throughput, so the size is pinned with subgroup-size control instead. + if (!context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + return -1; + } + if (context.AdapterInfo().vendor != std::string_view{"intel"}) { + return -1; + } + return config_index; +} + +// The bias add and the fused activation are not template parameters: the template +// calls the `write_output` helper that GenerateShaderCode builds from fuse_utils. +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) { + return WGSL_TEMPLATE_APPLY(shader, "nn/subgroup_matrix_conv_8x16x16.wgsl.template", + 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)); +} + +// Subgroup-matrix Conv implementation. Computes a channels-last 2D convolution as +// an implicit GEMM during Compute: A is the im2col of the activation, gathered on +// the fly into workgroup memory (never materialized in global memory), B is the +// transposed weight, and the product is accumulated with +// subgroupMatrixMultiplyAccumulate. The class is intended to support all +// subgroup-matrix configs; for now only 8x16x16 is implemented. The per-problem +// output tiling is supplied by a vendor-specific selector kept internal to this +// impl, which declines problems its policy does not cover. +// It runs its own Conv shape inference to decide eligibility, mirroring +// Conv::ComputeInternal: ConvOptImpl::Compute is called before ComputeInternal has +// resolved auto_pad or allocated the output. +template +class SubgroupMatrixConvImpl final : public Conv::ConvOptImpl { + using ConvT = Conv; + using Base = typename ConvT::ConvOptImpl; + + public: + SubgroupMatrixConvImpl(const ConvT& parent, int32_t config_index, + SubgroupMatrixTilingSelector tiling_selector) + : Base(parent), config_index_(config_index), tiling_selector_(std::move(tiling_selector)) {} + + Status Compute(ComputeContext& context, /*out*/ bool& handled) override { + handled = false; + + const ConvT& parent = this->parent_; + const ConvAttributes& conv_attrs = parent.ConvAttrs(); + const Activation& activation = parent.ConvActivation(); + + // The kernel reads the weight in its original OIHW layout (it applies its own + // OIHW -> OHWI transpose below), so any prepacked layout (HWIO or OHWI) is + // unusable here. Conv::PrePackInternal tests CanApplySubgroupMatrixConv before + // every prepacking branch and skips prepacking for exactly the problems accepted + // here, so this normally never declines. + if (parent.PrepackedKernel() != nullptr) { + return Status::OK(); + } + const auto* src = context.Input(0); + const auto* weight = context.Input(1); + if (src == nullptr || weight == nullptr) { + return Status::OK(); + } + const TensorShape& weight_shape = weight->Shape(); // OIHW: [Cout, Cin, kh, kw] + if (!CanApplySubgroupMatrixConv(context, is_channels_last, weight_shape, + narrow(conv_attrs.group), weight->DataType())) { + return Status::OK(); + } + + // Shape inference mirrors Conv::ComputeInternal. channels-last and a 4D weight are + // guaranteed by CanApplySubgroupMatrixConv; only a rank-4 NHWC activation is served + // (a reshaped Conv1D never gets here -- its kernel becomes 1xW, which + // CanApplySubgroupMatrixConv rejects). + const TensorShape& src_shape = src->Shape(); // NHWC: [batch, src_h, src_w, Cin] + if (src_shape.NumDimensions() != 4) { + return Status::OK(); + } + 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(weight_shape, kernel_spacial_shape_vector, false)); + if (local_pads.empty()) { + local_pads.resize(kernel_spacial_shape_vector.size() * 2, 0); + } + if (local_dilations.empty()) { + local_dilations.resize(kernel_spacial_shape_vector.size(), 1); + } + if (local_strides.empty()) { + local_strides.resize(kernel_spacial_shape_vector.size(), 1); + } + TensorShapeVector src_shape_vector = src_shape.AsShapeVector(); + TensorShapeVector output_shape_vector = {src_shape[0]}; + TensorShape src_spacial_shape = TensorShape( + TensorShapeVector(std::next(src_shape_vector.begin()), std::prev(src_shape_vector.end()))); + ORT_RETURN_IF_ERROR(conv_attrs.InferPadsAndOutputShape(src_spacial_shape, kernel_spacial_shape_vector, + local_strides, local_dilations, local_pads, + output_shape_vector)); + output_shape_vector.push_back(weight_shape[0]); // Cout + const TensorShape output_shape(output_shape_vector); // NHWC: [batch, out_h, out_w, Cout] + + // Spatial attributes as the kernel's uniforms want them: pads is [begin..., end...] + // in ONNX order, strides/dilations one entry per spatial dim. static_cast, not + // narrow, so an input the normal Conv path accepts is not turned into an abort here + // -- ComputeInternal casts these the same way. + std::vector pads; + std::vector strides; + std::vector dilations; + auto transform_dim = [](int64_t dim) { return static_cast(dim); }; + 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); + if (dilations.size() < 2 || pads.size() < 2 || strides.size() < 2) { + return Status::OK(); + } + + const uint32_t batch = narrow(src_shape[0]); + const uint32_t src_h = narrow(src_shape[1]); + const uint32_t src_w = narrow(src_shape[2]); + const uint32_t channel_i = 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 * channel_i; // contraction dim + const uint32_t N = narrow(weight_shape[0]); // Cout + + // An empty batch, M or N yields an empty output (ONNX permits zero-length dims). + // Let the generic Conv paths handle it rather than dispatch a degenerate + // (zero-workgroup) kernel. + if (batch == 0 || M == 0 || N == 0) { + return Status::OK(); + } + + // Pick the per-problem output tiling from the vendor policy (Intel). The GEMM + // dims map directly: A [M, K] x B [K, N] with one z-slice per conv batch. The + // Conv selector uses the Conv scratch budget and Conv-specific pretuned data. + const std::optional tiling = tiling_selector_(context, M, N, K, batch); + if (!tiling) { + return Status::OK(); + } + + // The kernel keeps the unchecked weight loads of a trailing partial N tile in bounds + // by shifting that tile's origin back to N - tile_n, which only exists when the tile + // fits within N. The Conv selector clamps tile_n for exactly this reason + // (ClampTileNToProblem), so this normally never triggers -- it mirrors the same guard + // in the subgroup-matrix MatMul/Gemm impls, keeping the precondition enforced on this + // side of the vendor boundary too. M needs no such guard: the A operand is staged in + // workgroup memory sized by the tile, so tile_m > M only wastes lanes. + // + // Last thing that can decline, so it comes before the output allocation and the + // weight transpose below: an impl that declines must not have allocated or + // dispatched anything. + if (N < tiling->tile_n) { + return Status::OK(); + } + + // Past this point the path is committed, so the output can be allocated: an impl + // that declines must allocate nothing. + Tensor* output = context.Output(0, output_shape); + + // Transpose the weight OIHW -> OHWI = [Cout, kh, kw, Cin], which flattens to the + // GEMM right operand B stored column-major: element B[k, n] lives at n*K + k, so + // the kernel loads B right tiles as col_major with stride K. (This matches the + // im2col-matmul weight layout; the alternative HWIO -> row-major stride-N load is + // the other option.) This allocates a GPU tensor and dispatches a Transpose + // program, so the path is already committed at this point. + // + // A constant weight is transposed once and reused; a non-constant one can change + // between Runs, so it has to be redone every time. Declining the non-constant case + // instead would cost more than the transpose it saves. + // + // TODO: prepack the constant case, which would drop the cache and the first-Run + // dispatch too. The layout is not the obstacle -- KernelLayout::OHWI is exactly + // what PrePackIm2ColMatMulWeight already produces, with the same {0, 2, 3, 1} + // perm. What blocks it is the fallback: PrePackInternal has to commit from the + // weight alone, while Compute can still decline at runtime (empty batch/M/N, or + // fewer than two spatial entries in dilations/pads/strides), and an OHWI weight is + // only readable by the im2col path that ComputeInternal ORT_ENFORCEs on. Since + // CanApplyIm2ColMatMulProgram is not implied by CanApplySubgroupMatrixConv (it + // gates on an arch-string whitelist and a fixed list of activation kinds, this one + // on the reported subgroup-matrix config and all of fuse_utils), a decline could + // land on a path that cannot read OHWI. Prepacking therefore means dropping the + // early return in Conv::PrePackInternal so the im2col branch below it prepacks + // when *both* predicates hold, and accepting a prepacked OHWI weight here instead + // of declining on it above. + Tensor run_local_ohwi_weight; + const Tensor* ohwi_weight = nullptr; + if (parent.IsWeightConstant()) { + ORT_RETURN_IF_ERROR(EnsureTransposedKernel(context, weight, weight_shape)); + ohwi_weight = cached_ohwi_weight_.get(); + } else { + ORT_RETURN_IF_ERROR( + TransposeKernel(context, weight, weight_shape, &run_local_ohwi_weight, {0, 2, 3, 1})); + ohwi_weight = &run_local_ohwi_weight; + } + + const auto& config = supported_subgroup_matrix_configs[config_index_]; + const uint32_t tile_m = tiling->tile_m; + const uint32_t tile_n = tiling->tile_n; + 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); + + // Safety net: the Conv selector already enforces the Conv scratch budget + // (output tiles + im2col A staging), so this clamp normally does nothing; it + // guards against a stray oversized split-K (e.g. a hand-edited pretuned entry). + // Only split_k is clamped, not the tile area: the largest candidate tile needs + // 64 * (64 + 16) = 5120 elements at split_k 1, well inside the budget. + uint32_t split_k = tiling->split_k; + while (split_k > 1 && + split_k * tile_m * (tile_n + config.K) > kMaxConvScratchElems) { + split_k /= 2; + } + + 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 the input channel (cin), which is the + // fastest-varying, contiguous NHWC dimension. channel_i % vec_size == 0 keeps + // each vec chunk inside a single kernel position, so all lanes of a chunk share + // one bounds check and map to contiguous input. + const uint32_t vec_size = channel_i % 4 == 0 ? 4 : (channel_i % 2 == 0 ? 2 : 1); + + const Tensor* bias = context.InputCount() > 2 ? context.Input(2) : nullptr; + const bool has_bias = bias != nullptr; + SubgroupMatrixConvProgram program{has_bias, config_index_, sg_mat_count_m, sg_mat_count_n, split_k, vec_size, + activation}; + program.SetWorkgroupSize(kSubgroupMatrixSubgroupSize * split_k); + // Pin the lanes per subgroup the kernel assumes; GetSupportedConfigIndex has + // already established that the device supports subgroup-size control. + program.SetSubgroupSize(kSubgroupMatrixSubgroupSize); + program.SetDispatchGroupSize(dispatch_x, dispatch_y, batch); + program.CacheHint(has_bias, config_index_, sg_mat_count_m, sg_mat_count_n, split_k, vec_size, + activation.CacheKey()) + .AddInputs({{src, ProgramTensorMetadataDependency::TypeAndRank, static_cast(vec_size)}, + {ohwi_weight, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddOutput({output, ProgramTensorMetadataDependency::Rank, output_shape, 1}) + .AddUniformVariables({{M}, + {N}, + {K}, + {dispatch_x}, + {batch}, + {src_h}, + {src_w}, + {channel_i}, + {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}); + } + ORT_RETURN_IF_ERROR(context.RunProgram(program)); + + handled = true; + return Status::OK(); + } + + private: + // Builds the OIHW -> OHWI transposed weight on first use and caches it. Only valid + // because the caller checked the weight is a constant initializer. impl_ is shared + // across concurrent Compute calls on the same Conv kernel, so the build is guarded + // by call_once, which also synchronizes the read of cached_ohwi_weight_ below. + // The transpose is attempted exactly once: if it fails, every later Run on this + // Conv node reports the same error rather than retrying. + // const, with a mutable cache, like SubgroupMatrixMatMulImpl::EnsurePaddedB: the + // cached tensor is a memo of the constant weight, not observable state of the 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."); + } + + const int32_t config_index_; + 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_; +}; + +} // 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 once per output element: bias add, then the + // fused activation. The activation body comes from fuse_utils + // (GetActivationDeclaration emits whatever module-scope helper the snippet needs, + // e.g. fused_act_erf), so every ActivationKind those cover is supported without the + // template having to enumerate them -- and any kind added there is picked up here + // for free. Mirrors SubgroupMatrixMatMulProgram::GenerateShaderCode. + 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 = supported_subgroup_matrix_configs[config_index_]; + if (config.Is(8, 16, 16)) { + return GenerateShaderCode8x16x16(shader, input_src, sg_mat_count_m_, sg_mat_count_n_, split_k_, vec_size_); + } + 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) { + // The subgroup-matrix kernel is F16-only and computes a channels-last, + // non-grouped 2D convolution. A fused activation is applied in the write-out + // epilogue, so fused convs are accepted. + 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]); + + // Mirror the im2col-matmul path: skip 1xW / Hx1 (and 1x1) kernels, which the + // plain MatMul / naive conv paths handle better. This also excludes reshaped + // Conv1D (kernel becomes 1xW). + 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; + // The subgroup-matrix K must be a whole number of 16-wide blocks: the kernel walks + // K in sg_mat_k blocks with no tail handling. + // + // N must cover at least one whole subgroup-matrix N block (16). The kernel loads B + // straight from the weight buffer with subgroupMatrixLoad, which is not + // bounds-checked, and handles a trailing partial N tile by shifting the tile origin + // back to N - tile_n; below one block there is no tile_n that fits, so such a Cout + // cannot be served at all (it goes to the generic conv paths, which handle a tiny + // Cout better anyway). Must stay in agreement with the N guard in + // CreateSubgroupMatrixConvTilingSelector, or PrePackInternal would skip prepacking + // for a problem this path then declines. + // + // N has no parity or block-multiple requirement beyond that. The f16 right-operand + // load needs a 4-byte-aligned offset and row stride, but B is OHWI, i.e. + // column-major with stride K, so the offset n*K + kb*sg_mat_k is a multiple of + // sg_mat_k for any n -- alignment follows from K, not from N. (The MatMul kernel + // does need an even N: its B is [K, N] row-major, so N *is* the row stride, which is + // why it pads odd N. Do not copy that constraint here.) + 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 int32_t config_index = GetSupportedConfigIndex(context); + if (config_index < 0) { + 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 ACCEPTED: 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) { + // Only channels-last Conv is served by this kernel; the caller falls back to the + // normal Conv path for channels-first. + if constexpr (!is_channels_last) { + ORT_UNUSED_PARAMETER(parent); + ORT_UNUSED_PARAMETER(context); + return nullptr; + } else { + // Only run on devices that report the fixed 8x16x16 F16 subgroup-matrix config + // this kernel is implemented for, and for which a vendor tiling policy exists + // (Intel today - GetSupportedConfigIndex enforces both). + const int32_t config_index = GetSupportedConfigIndex(context); + if (config_index < 0) { + return nullptr; + } + // The Conv selector shares the autotuner hooks (tile override, device-arch + // capture) with MatMul, but budgets the extra im2col staging the Conv kernel + // keeps in workgroup memory and consults the Conv-specific pretuned data. + SubgroupMatrixTilingSelector tiling_selector = intel::CreateSubgroupMatrixConvTilingSelector(context); + if (!tiling_selector) { + return nullptr; + } + return std::make_unique>(parent, config_index, + std::move(tiling_selector)); + } +} + +// Explicit instantiation for every Conv instantiation (see conv.cc). +#define WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(CHANNELS_LAST, FUSED) \ + template std::unique_ptr::ConvOptImpl> \ + CreateSubgroupMatrixConvImpl(const Conv& parent, \ + const ComputeContextBase& context); + +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(false, false) +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(false, true) +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(true, false) +WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(true, true) + +#undef WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL + +} // namespace webgpu +} // namespace onnxruntime 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..3158d52173575 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -0,0 +1,123 @@ +// 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/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 { + +// Computes a channels-last (NHWC) 2D convolution as an implicit-GEMM using +// subgroupMatrixMultiplyAccumulate. The activation operand is the im2col of the +// input, gathered on the fly into workgroup memory (never materialized in global +// memory); the weight is pre-transposed to OHWI, i.e. B stored column-major. +// config_index selects the device subgroup-matrix config (into +// supported_subgroup_matrix_configs); sg_mat_count_m/n select how many subgroup +// matrices the output tile spans along M/N; split_k is the number of subgroups that +// cooperatively reduce K. +// +// The bias add and the fused activation live in a generated `write_output` helper +// rather than in the WGSL template: the activation body comes from +// GetActivationSnippet / GetActivationDeclaration, so every ActivationKind those +// support works here without the template knowing the kinds. +class SubgroupMatrixConvProgram final : public Program { + public: + SubgroupMatrixConvProgram(bool has_bias, int32_t config_index, + uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k, + uint32_t vec_size, const Activation& activation) + : Program{"SubgroupMatrixConv"}, + has_bias_(has_bias), + config_index_(config_index), + sg_mat_count_m_(sg_mat_count_m), + sg_mat_count_n_(sg_mat_count_n), + split_k_(split_k), + vec_size_(vec_size), + activation_(activation) {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, + {"num_n_tile", ProgramUniformVariableDataType::Uint32}, + // Number of batch slices, i.e. the logical z extent of the tile + // grid. The kernel needs it to drop the extra workgroups a + // normalized dispatch can add (see the guard in the template). + {"batch", ProgramUniformVariableDataType::Uint32}, + {"src_h", ProgramUniformVariableDataType::Uint32}, + {"src_w", ProgramUniformVariableDataType::Uint32}, + {"channel_i", 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 int32_t config_index_; + const uint32_t sg_mat_count_m_; + const uint32_t sg_mat_count_n_; + const uint32_t split_k_; + const uint32_t vec_size_; + // Held by value, like the other subgroup-matrix/MatMul programs (and unlike the + // neighbouring Conv programs, which hold a reference into the kernel). The program + // outlives the call that builds it -- CacheKey() and the uniform values are read + // during Run -- and Activation is two floats plus an enum, so a copy is cheaper + // than having to reason about the referent's lifetime. + const Activation activation_; +}; + +// Returns true when the subgroup-matrix Conv fast path can run for this problem: +// fp16, channels-last, group == 1, a 4D non-1x1 kernel, and a device that reports +// the 8x16x16 F16 subgroup-matrix config, supports subgroup-size control (the kernel +// pins its 32 lanes per subgroup) and has a vendor tiling policy. The GEMM +// contraction dim K = kernel_h * kernel_w * Cin must be a multiple of the +// subgroup-matrix K (16), and N = Cout must be at least one subgroup-matrix N block +// (16) wide so that a trailing partial N tile can be shifted back far enough for its +// unchecked weight loads to stay in bounds. Above that, N needs no parity or +// block-multiple (the OHWI weight is column-major with stride K, so load alignment +// comes from K); M = out_h * out_w is unconstrained. Callers fall back to the default +// conv paths when this returns false. +// +// The fused activation is deliberately not part of the test: the write-out epilogue +// is generated from GetActivationSnippet, so every ActivationKind is handled and +// there is nothing to reject. Unlike CanApplyIm2ColMatMulProgram, this does not need +// an IsActivationSupported() companion. +// +// This is the shape/device-level eligibility test, decided from the weight shape +// alone: SubgroupMatrixConvImpl calls it from Compute, and Conv::PrePackInternal +// calls it to skip prepacking for a weight this path will consume in its original +// OIHW layout. +bool CanApplySubgroupMatrixConv(ComputeContextBase& context, + bool is_channels_last, + const TensorShape& kernel_shape, + uint32_t group, + MLDataType data_type); + +// Creates a subgroup-matrix Conv implementation on devices whose vendor policy +// supports it; returns nullptr otherwise, so the caller falls back to the normal Conv +// paths. The impl reads the Conv attributes from `parent`, and the per-problem output +// tiling comes from a vendor-specific selector chosen internally from the device +// context. +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..75f218b48bb70 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template @@ -0,0 +1,606 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Subgroup-matrix Conv kernel (F16) for the 8x16x16 (sg_mat_m x sg_mat_n x +// sg_mat_k) subgroup-matrix shape. It computes a channels-last (NHWC) 2D +// convolution as an implicit GEMM Y[MxN] = A[MxK] @ B[KxN] (+ optional bias), +// where the activation operand A is the im2col of the input and never +// materialized in global memory. +// +// M = out_h * out_w (output pixels per batch slice) +// K = kernel_h * kernel_w * Cin +// N = Cout +// +// A: im2col of the NHWC input. subgroupMatrixLoad requires a contiguous strided +// buffer, but im2col rows are gathered from scattered input pixels, so each +// K-block A-tile is gathered into workgroup scratch (a_scratch) first and then +// loaded as a left operand (subgroup_matrix_left) with +// stride kSgMatK. K index order is (kh, kw, cin) so it matches the OHWI weight. +// Because cin is the fastest-varying, contiguous NHWC dimension, the gather +// reads vec_size channels at a time (host guarantees channel_i % vec_size == 0). +// B: weight pre-transposed OIHW -> OHWI = [Cout, kh, kw, Cin], i.e. B stored +// column-major ([N, K] row-major; element B[k, n] at n*K + k). Loaded directly +// from global memory as a right operand (subgroup_matrix_right) with col_major and stride K. Shared across all batch slices. +// +// Workgroup: split_k subgroups x 32 lanes (split_k in {1,2,4,8}). +// Tile: kTileM x kTileN, chosen adaptively by the host tiling selector. Each +// subgroup computes the whole tile as kSgMatCountM x kSgMatCountN subgroup +// matrices. K is iterated in sg_mat_k blocks. When split_k > 1 the K blocks are +// distributed round-robin across the subgroups, each accumulates its own partial +// tile in shared memory, and the partials are summed at write-out. +// +// Unlike the plain subgroup-matrix MatMul kernel, A must be staged in workgroup +// memory before it can be loaded as a subgroup matrix. To keep the workgroup +// barriers in uniform control flow when split_k > 1, the K reduction runs as a +// per-step loop: every subgroup gathers its K-block A-tile into its own +// a_scratch slot, a workgroup barrier makes those tiles visible, then every +// subgroup multiply-accumulates its tile. This keeps all subgroups in lockstep +// at each barrier regardless of how many K-blocks remain. +// +// The bias add and the fused activation are not implemented here: the write-out +// pass calls write_output(), which the host builds in AdditionalImplementation from +// GetActivationSnippet / GetActivationDeclaration. That keeps every ActivationKind +// available without this template enumerating them. +// +// Preconditions enforced by the host: K % sg_mat_k == 0, and kTileN <= N (the tiling +// selector clamps tile_n and the impl re-checks it) so a trailing partial N tile can be +// shifted back far enough for its weight loads to stay in bounds. N itself may be odd -- +// B is column-major with stride K, so the load alignment comes from K, not N. M may be +// any size; partial tiles are handled by the shift, by bounds-checked stores and by the +// im2col gather returning zero for out-of-range pixels. + +#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 + +#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 (channel_i % vec_size == 0) + +// Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each +// split-K subgroup accumulates its slice of K into its own slot; the slots are +// summed during write-back. Storing here also lets us write to global memory +// with M/N bounds checks for dimensions that are not multiples of the tile size. +var scratch: array; + +// Per-subgroup staging for the current K-block A-tile: kSplitK slots, each a +// [kTileM, kSgMatK] row-major (stride kSgMatK) tile gathered by im2col. +var a_scratch: array; + +// Gathers a contiguous vec_size-wide chunk of im2col elements +// A[m, k_base .. k_base + vec_size - 1] for the given batch slice, where +// k_base = k_vec * vec_size. The chunk always lies within a single (kh, kw) +// kernel position: the K order is (kh, kw, cin) with cin fastest and the host +// guarantees channel_i % vec_size == 0, so all vec lanes share one (ih, iw) and a +// single bounds check and map to contiguous NHWC input. Out-of-range input pixels +// (padding) contribute zero. +// +// m past M (a partial M tile stages whole rows) is index-safe but not zero: oh >= out_h +// there, yet ih = oh * stride + kh * dilation - pad can still land inside [0, src_h), so +// the gather reads real input. That is harmless -- the (ih, iw) bounds check is what +// keeps every read inside the input buffer, and those rows are never written back. +// m -> output pixel (oh, ow); k_vec -> (kh_idx, kw_idx, cin_vec) with cin_vec fastest. +fn load_im2col_vec(m: u32, k_vec: u32, src_vec_batch_offset: u32) -> input_src_value_t { + let channel_i_vec = uniforms.channel_i / kVecSize; + let oh = m / uniforms.output_w; + let ow = m % uniforms.output_w; + let cin_vec = k_vec % channel_i_vec; + let khw = k_vec / channel_i_vec; + let kw_idx = khw % uniforms.kernel_w; + let kh_idx = khw / 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.src_h) || iw < 0 || iw >= i32(uniforms.src_w)) { + return input_src_value_t(); + } + let src_vec_idx = src_vec_batch_offset + + (u32(ih) * uniforms.src_w + u32(iw)) * channel_i_vec + cin_vec; + return input_src.getByOffset(src_vec_idx); +} + +$MAIN { + // Recover the batch slice and the logical (n_tile, m_tile) from the flattened + // workgroup index. The host dispatches num_n_tile x num_m_tile tiles per batch + // slice on (x, y, z=batch), but that grid may be normalized (reshaped across + // x/y/z to stay within per-dimension limits), so workgroup_id.x/y/z cannot be + // used directly. workgroup_idx is the true linear index, invariant under + // normalization; batch slices are contiguous blocks of tiles_per_slice tiles. + let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; + let tiles_per_slice = uniforms.num_n_tile * num_m_tile; + let batch_id = workgroup_idx / tiles_per_slice; // which batch slice + // Normalization rounds the workgroup count up (see + // ProgramManager::NormalizeDispatchGroupSize), so more workgroups can be launched + // than there are tiles. Those extras land past the last slice; without this guard + // they would derive batch_id >= batch and write past the end of the output (the + // in-tile M/N checks below cannot catch them, since their m_tile/n_tile look + // perfectly valid). workgroup_idx is workgroup-uniform, so the whole workgroup + // returns together and the barriers below stay in uniform control flow. + if (batch_id >= uniforms.batch) { + return; + } + let slice_idx = workgroup_idx % tiles_per_slice; // tile index within the slice + let n_tile = slice_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 the B right-operand loads stay inside + // the weight buffer. subgroupMatrixLoad has no bounds check of its own and the + // N-block j of this tile reads weight columns global_base_n + j*kSgMatN .. + // +kSgMatN-1, so an unshifted last tile would read up to kTileN-1 columns past + // Cout whenever N is not a multiple of kTileN. + // + // Unlike the MatMul kernel this needs no alignment fixup: B is column-major with + // stride K and K % kSgMatK == 0, so the load offset n*K + kb*kSgMatK is a + // multiple of kSgMatK for any n. The max() only guards the u32 subtraction; the + // host clamps kTileN to at most N, so the shifted tile always fits. + let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); + // Columns before the unshifted origin belong to the previous tile and are skipped + // at write-out, so a shifted tile never writes the same output element twice. + let n_skip = n_tile_start - global_base_n; + // M needs no such shift: the A left operand is gathered into a_scratch, which is + // sized by the tile rather than by M, so a partial M tile only stages unused rows + // (and the im2col gather itself is index-safe for them). + let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile + let m_base = m_tile * kTileM; + // Flat-element offset into the NHWC input for this batch slice. The weight is + // shared across batch slices, so B has no batch offset. The output is laid out + // as [batch, M, N] row-major. + let src_vec_batch_offset = batch_id * uniforms.src_h * uniforms.src_w * (uniforms.channel_i / kVecSize); + let out_batch_offset = 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. Each subgroup processes K-blocks kb = sg_index, sg_index + + // kSplitK, ... The loop bound num_steps is workgroup-uniform, so the barriers + // below stay in uniform control flow even though a subgroup may run out of + // K-blocks before the last step (valid == false). + 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; + + // Stage this subgroup's [kTileM, kSgMatK] A-tile for K-block kb. The 32 + // lanes cooperatively gather the tile row-major (stride kSgMatK), each + // lane fetching a contiguous vec_size-wide chunk of cin per step. + if (valid) { + let k_vec_base = kb * kSgMatK / kVecSize; + let k_vec_per_row = kSgMatK / kVecSize; + for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_row; vidx = vidx + kSubgroupSize) { + let row = vidx / k_vec_per_row; + let kc_vec = vidx % k_vec_per_row; + let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, src_vec_batch_offset); + 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. The weight is OHWI, i.e. B is + // stored column-major ([N, K] row-major, element B[k, n] at n*K + k), so + // the tile for N-block j (columns global_base_n + j*sg_mat_n) starts at + // (global_base_n + j*sg_mat_n)*K + kb*sg_mat_k and is loaded col-major + // with 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 + + // 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 this subgroup's partial results into its split-K slot of scratch, + // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) + // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. + 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(); + + // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] + // tiles into slot 0 of scratch. Each lane owns a strided set of element + // positions and sums that position across all kSplitK slots. Writing slot 0 + // is safe: a position is only overwritten by the same lane that already read + // it as the slot-0 term of its own sum. + 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 pass: every subgroup writes out whole M-rows from the summed tile + // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows + // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a + // row's N elements (strided by the subgroup size). Columns below n_skip belong to + // the previous (unshifted) tile and M is bounds-checked, so non-multiple + // dimensions write each valid output element exactly once. write_output applies + // the bias and the fused activation before storing. + 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_offset + 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..602839c19001c 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 @@ -3,6 +3,7 @@ #include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" +#include #include #include #include @@ -23,6 +24,30 @@ namespace onnxruntime { namespace webgpu { namespace intel { +// Optional Conv-specific pretuned decision tree, emitted by +// tools/python/gen_sgmm_tree.py --mode conv from an on-device Conv sweep. When the +// .inc is present it defines sgconv_tree::::{kArch, PredictTilingTree} and is +// consulted before the Conv heuristic (see LookupConvTree / SelectConvTiling); when +// absent the heuristic is used. The file is dropped in after a sweep, so the build +// must not depend on it existing. +// +// One such tree is checked in, so the guard looks redundant at first glance. It is +// not: the tree is measured data, not source, and the include has to keep working +// both ways -- a build that deletes it (a fork with no interest in shipping Intel +// tuning data, or a size-constrained target) must still compile, and a new arch's +// tree must be droppable in without touching the build files. The heuristic is a +// complete fallback, so the only cost of absence is tiling quality. +// +// Included here, inside the namespace, because the generated code names +// SubgroupMatrixTiling unqualified: it is an oracle over that type and carries no +// includes or namespace boilerplate of its own. +#if defined(__has_include) +#if __has_include("core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc") +#include "core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc" +#define ORT_WEBGPU_HAS_SGCONV_TREE 1 +#endif +#endif + namespace { // Subgroup matrix configuration used by this implementation (Intel Xe2/Xe3, F16). @@ -36,9 +61,23 @@ 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 +// Extra per-tile N a kernel stages in workgroup memory beyond the split-K output +// tiles. The MatMul kernel stages nothing else (0); the Conv kernel also holds the +// im2col A tile (tile_m x sg_mat_k), so its scratch bound is +// split_k * tile_m * (tile_n + sg_mat_k) -- up to 2x stricter than MatMul's. This is +// the only difference between the two kernels' tiling constraints, so every helper +// below takes it as a parameter rather than being duplicated per kernel. +constexpr uint32_t kMatMulScratchNPad = 0; +constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK; + +// Workgroup-memory footprint of a tiling, in f16 elements. +uint64_t ScratchElems(const SubgroupMatrixTiling& t, uint32_t scratch_n_pad) { + return static_cast(t.split_k) * t.tile_m * (t.tile_n + scratch_n_pad); +} + // 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 scratch_n_pad) { if (t.tile_m == 0 || t.tile_n == 0 || t.tile_m % kSubgroupMatrixM != 0 || t.tile_n % kSubgroupMatrixN != 0) { return false; @@ -47,7 +86,7 @@ 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, scratch_n_pad) <= kMaxScratchElems; } // HwSubgroups returns 0 for an unrecognized arch; fall back to a conservative @@ -66,7 +105,8 @@ uint32_t EffectiveHwSubgroups(std::string_view arch) { // 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. -SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { +SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch, + uint32_t scratch_n_pad) { 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; }; @@ -119,13 +159,46 @@ 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}, scratch_n_pad) <= kMaxScratchElems) { split_k = cand; } } return {tile_m, tile_n, split_k}; } +// Repairs a pretuned tiling to this problem by halving split_k until it fits both +// K's block count and the scratch budget. The tuned trees are performance oracles +// that never see K's block count (they are trained on the problem's log dims only), +// so they can propose split_k > K / kSubgroupMatrixK on a problem outside the swept +// envelope. Clamping keeps the tuned tile shape -- the part that carries most of the +// win -- instead of discarding the prediction and falling back to the heuristic. +void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) { + const uint32_t k_blocks = K / kSubgroupMatrixK; + while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) { + t.split_k /= 2; + } +} + +// Caps tile_n at the largest whole number of subgroup-matrix N blocks that fits in N. +// +// Correctness, not performance: the Conv kernel loads its B right operand straight +// from the weight buffer, and subgroupMatrixLoad is not bounds-checked, so a trailing +// partial N tile is shifted back to N - tile_n. That shift only exists when +// tile_n <= N. Neither source guarantees it -- HeuristicTiling only skips candidates +// larger than N for all but the smallest one, and the tuned trees are unconstrained +// oracles -- so the cap is applied to both. Shrinking tile_n keeps it a multiple of +// kSubgroupMatrixN and only lowers ScratchElems, so it cannot invalidate a tiling +// that was already valid. +// +// Callers must have rejected N < kSubgroupMatrixN, which cannot be served at all +// (there is no non-zero multiple of kSubgroupMatrixN that fits). +void ClampTileNToProblem(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 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 @@ -160,11 +233,34 @@ std::optional LookupPretunedTiling(std::string_view arch, return std::nullopt; } +// Looks up the Conv-specific pretuned decision tree for this arch. One `if` per +// tuned arch; each generated tree carries its own kArch, so the arch spelling lives +// with the data instead of being duplicated here. Returns nullopt when the arch has +// not been swept (or no tree was dropped into the build), in which case the caller +// uses the Conv heuristic. +std::optional LookupConvTree([[maybe_unused]] std::string_view arch, + [[maybe_unused]] uint32_t M, + [[maybe_unused]] uint32_t N, + [[maybe_unused]] uint32_t K, + [[maybe_unused]] uint32_t batch) { +#ifdef ORT_WEBGPU_HAS_SGCONV_TREE + if (arch == sgconv_tree::xe3lpg::kArch) { + return sgconv_tree::xe3lpg::PredictTilingTree(M, N, K, batch); + } +#endif + 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. +// +// Only for batch-blind sources. The Conv tree takes batch as a training feature and +// was fitted on measurements at batch 1/2/4/8, so its split-K already prices batch in; +// applying this on top of it double-counts occupancy. Measured over the xe-3lpg sweep, +// doing so overrode 41% of the tuned picks and raised mean regret from 3.1% to 7.2%. void ClampSplitKForBatch(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); @@ -178,7 +274,7 @@ void ClampSplitKForBatch(SubgroupMatrixTiling& t, uint32_t M, uint32_t N, uint32 // batch is the number of z-dispatched slices; it scales occupancy (see // ClampSplitKForBatch) but not the per-slice tile shape. 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, kMatMulScratchNPad)) { SubgroupMatrixTiling tiling = *tuned; // The table is tuned per (M, N, K) at batch 1, so only revisit split-K when // batch adds occupancy. @@ -187,7 +283,33 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, } return tiling; } - return HeuristicTiling(arch, M, N, K, batch); + return HeuristicTiling(arch, M, N, K, batch, kMatMulScratchNPad); +} + +// Conv counterpart of SelectTiling. Consults the Conv-specific pretuned tree (when +// one was dropped in for this arch) before falling back to the Conv heuristic; both +// use the Conv scratch budget. There is no MatMul-table lookup here: the Conv +// kernel's cost profile (the im2col gather) differs, and its scratch budget is +// stricter, so it must not reuse the MatMul-tuned entries. +SubgroupMatrixTiling SelectConvTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { + if (std::optional tuned = LookupConvTree(arch, M, N, K, batch)) { + // Unlike the MatMul table this prediction already accounts for batch (it is a + // model input), so it is taken as-is apart from the correctness clamp -- no + // ClampSplitKForBatch, which would second-guess the measured choice. + // + // The tree only ever violates split_k (every leaf fits the Conv scratch budget by + // construction), so clamp rather than discard, then re-check. On the swept + // envelope the clamp is a no-op; it only bites for a K smaller than anything + // swept. + ClampSplitKToProblem(*tuned, K, kConvScratchNPad); + if (IsTilingValid(*tuned, K, kConvScratchNPad)) { + ClampTileNToProblem(*tuned, N); + return *tuned; + } + } + SubgroupMatrixTiling tiling = HeuristicTiling(arch, M, N, K, batch, kConvScratchNPad); + ClampTileNToProblem(tiling, N); + return tiling; } } // namespace @@ -212,6 +334,34 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( }; } +SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( + const ComputeContextBase& context) { + if (context.AdapterInfo().vendor != std::string_view{"intel"}) { + return nullptr; + } + // Conv tiling policy for the 8x16x16 subgroup-matrix Conv kernel: same tile + // candidates as MatMul, but selection uses the Conv scratch budget (which also + // stages the im2col A tile) and the Conv-specific pretuned tree / heuristic. + return [](const ComputeContext& context, uint32_t M, uint32_t N, + uint32_t K, uint32_t batch) -> std::optional { + const std::string_view arch = std::string_view{context.AdapterInfo().architecture}; + // K must align to the subgroup tiling, and N must hold at least one whole + // subgroup-matrix N block: the Conv kernel loads B unchecked from the weight + // buffer and shifts a trailing partial N tile back to N - tile_n, which is + // impossible below one block. A partial N tile above that is served by the shift + // (see ClampTileNToProblem); M partial tiles need no guarantee at all, since the + // A operand is staged in workgroup memory sized by the tile. + // + // Declining here rather than clamping keeps this in agreement with + // CanApplySubgroupMatrixConv, which rejects the same N so Conv::PrePackInternal + // does not skip prepacking for a problem this path will not serve. + if (K % kSubgroupMatrixK != 0 || N < kSubgroupMatrixN) { + return std::nullopt; + } + return SelectConvTiling(arch, M, N, K, batch); + }; +} + } // 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..46566d8ee1a2e 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 @@ -20,6 +20,14 @@ namespace intel { SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( const ComputeContextBase& context); +// Like CreateSubgroupMatrixTilingSelector, but for the subgroup-matrix Conv path. +// The tile candidates and autotuner hooks (override / arch capture) are shared, +// but selection uses the Conv workgroup-memory budget (which also stages the +// im2col A tile) and consults the Conv-specific pretuned data instead of the +// MatMul table. Returns an empty selector on non-Intel devices. +SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( + const ComputeContextBase& context); + } // namespace intel } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc b/onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc new file mode 100644 index 0000000000000..bfc3a52c3e051 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc @@ -0,0 +1,1871 @@ +// Auto-generated by tools/python/gen_sgmm_tree.py --mode conv. Do not edit. +// Distilled decision tree (depth 10, 460 leaves, 30 tilings) +// approximating the offline GBM cost model for arch: xe-3lpg. +// +// Arch-specific: dispatch on kArch. One .inc per arch. +// +// Every leaf satisfies the conv workgroup-memory budget +// (split_k * tile_m * (tile_n + 16) <= 16384), so a prediction can only be +// invalid through split_k: this is a performance oracle only, like HeuristicTiling in +// subgroup_matrix_tiling_selector.cc, and it does not see K's block count, so it can +// propose split_k > K / 16. The caller MUST clamp split-K to K's block count and +// re-check validity. +// +// The caller must NOT apply the batch-occupancy clamp (ClampSplitKForBatch) on top of +// a prediction from this tree: batch is a model feature here (f3), so the predicted +// split_k already prices occupancy in, and clamping again double-counts it. That clamp +// exists for the batch-blind pretuned table only. +// +// The including TU must provide , , and +// SubgroupMatrixTiling. + +namespace sgconv_tree::xe3lpg { + +// The device architecture string this tree was tuned on, as reported by the adapter. +inline constexpr std::string_view kArch = "xe-3lpg"; + +inline SubgroupMatrixTiling PredictTilingTree(uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { + [[maybe_unused]] const float f0 = std::log2(static_cast(M < 1u ? 1u : M)); + [[maybe_unused]] const float f1 = std::log2(static_cast(N < 1u ? 1u : N)); + [[maybe_unused]] const float f2 = std::log2(static_cast(K < 1u ? 1u : K)); + [[maybe_unused]] const float f3 = std::log2(static_cast(batch < 1u ? 1u : batch)); + if (f0 <= 7.792481f) { + if (f1 <= 5.792481f) { + if (f0 <= 6.292481f) { + if (f3 <= 2.500000f) { + if (f1 <= 4.292481f) { + if (f3 <= 0.500000f) { + if (f2 <= 8.406890f) { + return {8, 16, 4}; + } else { + if (f2 <= 9.921175f) { + return {16, 16, 8}; + } else { + return {8, 16, 8}; + } + } + } else { + return {8, 16, 8}; + } + } else { + if (f0 <= 5.792481f) { + return {8, 16, 8}; + } else { + if (f3 <= 1.500000f) { + return {8, 16, 8}; + } else { + if (f1 <= 5.292481f) { + return {8, 16, 8}; + } else { + return {16, 16, 8}; + } + } + } + } + } else { + if (f0 <= 5.292481f) { + if (f2 <= 8.406890f) { + if (f1 <= 5.292481f) { + return {8, 16, 8}; + } else { + return {8, 16, 4}; + } + } else { + return {8, 16, 8}; + } + } else { + if (f1 <= 4.292481f) { + return {8, 16, 8}; + } else { + if (f2 <= 8.906890f) { + if (f1 <= 5.292481f) { + if (f1 <= 4.792481f) { + return {8, 16, 4}; + } else { + return {8, 32, 8}; + } + } else { + return {8, 16, 4}; + } + } else { + if (f1 <= 4.792481f) { + return {16, 16, 8}; + } else { + if (f1 <= 5.292481f) { + return {8, 32, 8}; + } else { + return {16, 16, 8}; + } + } + } + } + } + } + } else { + if (f3 <= 1.500000f) { + if (f2 <= 10.421175f) { + if (f2 <= 8.406890f) { + if (f0 <= 6.792481f) { + return {8, 16, 8}; + } else { + if (f2 <= 7.435183f) { + return {8, 16, 4}; + } else { + if (f1 <= 4.792481f) { + if (f0 <= 7.292481f) { + return {8, 16, 8}; + } else { + return {8, 16, 4}; + } + } else { + return {8, 16, 8}; + } + } + } + } else { + if (f1 <= 5.292481f) { + if (f1 <= 4.292481f) { + if (f3 <= 0.500000f) { + return {16, 16, 8}; + } else { + return {8, 16, 8}; + } + } else { + if (f0 <= 7.292481f) { + return {8, 16, 8}; + } else { + if (f3 <= 0.500000f) { + return {8, 16, 8}; + } else { + return {8, 32, 8}; + } + } + } + } else { + if (f3 <= 0.500000f) { + return {8, 16, 8}; + } else { + return {16, 16, 8}; + } + } + } + } else { + if (f0 <= 7.292481f) { + return {8, 16, 8}; + } else { + if (f3 <= 0.500000f) { + return {8, 16, 8}; + } else { + if (f1 <= 4.792481f) { + return {8, 16, 8}; + } else { + return {8, 32, 8}; + } + } + } + } + } else { + if (f2 <= 9.421175f) { + if (f1 <= 4.792481f) { + if (f3 <= 2.500000f) { + if (f1 <= 4.292481f) { + return {8, 16, 8}; + } else { + return {8, 16, 4}; + } + } else { + if (f1 <= 4.292481f) { + return {8, 16, 4}; + } else { + return {16, 16, 4}; + } + } + } else { + if (f1 <= 5.292481f) { + if (f3 <= 2.500000f) { + return {8, 32, 8}; + } else { + return {8, 32, 4}; + } + } else { + if (f3 <= 2.500000f) { + return {8, 16, 4}; + } else { + return {16, 16, 4}; + } + } + } + } else { + if (f1 <= 4.792481f) { + if (f3 <= 2.500000f) { + if (f1 <= 4.292481f) { + return {8, 16, 8}; + } else { + return {16, 16, 8}; + } + } else { + return {16, 16, 8}; + } + } else { + if (f1 <= 5.292481f) { + if (f3 <= 2.500000f) { + return {8, 32, 8}; + } else { + if (f2 <= 10.921175f) { + return {8, 32, 4}; + } else { + return {16, 32, 8}; + } + } + } else { + if (f0 <= 7.292481f) { + return {16, 16, 8}; + } else { + return {16, 32, 4}; + } + } + } + } + } + } + } else { + if (f0 <= 5.292481f) { + if (f3 <= 1.500000f) { + if (f2 <= 7.435183f) { + if (f3 <= 0.500000f) { + return {8, 16, 4}; + } else { + if (f1 <= 7.292481f) { + return {8, 16, 8}; + } else { + return {8, 16, 4}; + } + } + } else { + if (f1 <= 8.292481f) { + return {8, 32, 8}; + } else { + if (f3 <= 0.500000f) { + return {8, 32, 8}; + } else { + if (f0 <= 4.292481f) { + return {8, 32, 8}; + } else { + if (f2 <= 8.406890f) { + return {8, 32, 4}; + } else { + if (f2 <= 11.421175f) { + return {8, 64, 8}; + } else { + return {8, 32, 8}; + } + } + } + } + } + } + } else { + if (f1 <= 7.292481f) { + if (f2 <= 7.435183f) { + if (f3 <= 2.500000f) { + return {8, 32, 8}; + } else { + return {8, 16, 4}; + } + } else { + if (f3 <= 2.500000f) { + return {8, 32, 8}; + } else { + if (f1 <= 6.792481f) { + if (f1 <= 6.292481f) { + return {8, 32, 8}; + } else { + if (f2 <= 9.421175f) { + return {8, 32, 4}; + } else { + return {8, 32, 8}; + } + } + } else { + if (f0 <= 4.292481f) { + return {8, 32, 8}; + } else { + if (f2 <= 9.421175f) { + return {8, 32, 4}; + } else { + return {8, 64, 8}; + } + } + } + } + } + } else { + if (f2 <= 10.421175f) { + if (f3 <= 2.500000f) { + if (f2 <= 8.906890f) { + return {8, 32, 4}; + } else { + if (f1 <= 8.292481f) { + return {8, 64, 8}; + } else { + return {8, 32, 4}; + } + } + } else { + if (f1 <= 8.292481f) { + if (f0 <= 4.792481f) { + return {8, 32, 4}; + } else { + return {8, 64, 4}; + } + } else { + if (f0 <= 4.292481f) { + return {8, 32, 4}; + } else { + return {8, 64, 4}; + } + } + } + } else { + if (f1 <= 8.792481f) { + if (f2 <= 11.421175f) { + return {8, 64, 8}; + } else { + if (f3 <= 2.500000f) { + return {8, 32, 8}; + } else { + return {8, 64, 8}; + } + } + } else { + if (f3 <= 2.500000f) { + return {8, 64, 8}; + } else { + return {8, 64, 4}; + } + } + } + } + } + } else { + if (f1 <= 7.292481f) { + if (f3 <= 1.500000f) { + if (f1 <= 6.292481f) { + if (f0 <= 6.292481f) { + if (f3 <= 0.500000f) { + return {8, 16, 8}; + } else { + if (f0 <= 5.792481f) { + return {8, 16, 8}; + } else { + return {8, 32, 8}; + } + } + } else { + if (f3 <= 0.500000f) { + if (f0 <= 7.292481f) { + return {8, 16, 8}; + } else { + return {8, 32, 8}; + } + } else { + if (f0 <= 7.292481f) { + return {8, 32, 8}; + } else { + return {8, 64, 8}; + } + } + } + } else { + if (f0 <= 6.792481f) { + if (f3 <= 0.500000f) { + if (f0 <= 6.292481f) { + if (f1 <= 6.792481f) { + return {8, 16, 8}; + } else { + return {8, 32, 8}; + } + } else { + return {8, 32, 8}; + } + } else { + if (f0 <= 6.292481f) { + return {8, 32, 8}; + } else { + if (f2 <= 8.906890f) { + return {8, 32, 4}; + } else { + return {8, 32, 8}; + } + } + } + } else { + if (f3 <= 0.500000f) { + if (f2 <= 8.406890f) { + return {8, 32, 4}; + } else { + return {8, 32, 8}; + } + } else { + if (f2 <= 10.421175f) { + return {8, 32, 4}; + } else { + if (f1 <= 6.792481f) { + return {8, 32, 8}; + } else { + return {8, 64, 8}; + } + } + } + } + } + } else { + if (f2 <= 9.421175f) { + if (f0 <= 6.792481f) { + if (f3 <= 2.500000f) { + if (f1 <= 6.292481f) { + return {8, 32, 8}; + } else { + return {8, 32, 4}; + } + } else { + if (f1 <= 6.792481f) { + return {8, 32, 4}; + } else { + return {8, 64, 4}; + } + } + } else { + if (f3 <= 2.500000f) { + if (f1 <= 6.792481f) { + if (f0 <= 7.292481f) { + return {8, 32, 4}; + } else { + return {8, 32, 2}; + } + } else { + return {8, 64, 4}; + } + } else { + if (f1 <= 6.292481f) { + return {8, 64, 4}; + } else { + if (f1 <= 6.792481f) { + return {8, 32, 1}; + } else { + return {8, 64, 2}; + } + } + } + } + } else { + if (f1 <= 6.792481f) { + if (f1 <= 6.292481f) { + if (f0 <= 6.792481f) { + if (f3 <= 2.500000f) { + return {8, 32, 8}; + } else { + return {8, 64, 8}; + } + } else { + if (f3 <= 2.500000f) { + return {8, 64, 8}; + } else { + return {8, 64, 4}; + } + } + } else { + if (f2 <= 11.421175f) { + if (f3 <= 2.500000f) { + return {16, 32, 8}; + } else { + return {8, 32, 4}; + } + } else { + if (f3 <= 2.500000f) { + return {8, 32, 8}; + } else { + return {16, 32, 8}; + } + } + } + } else { + if (f0 <= 5.792481f) { + return {8, 32, 8}; + } else { + if (f0 <= 6.792481f) { + if (f3 <= 2.500000f) { + return {8, 64, 8}; + } else { + return {8, 64, 4}; + } + } else { + if (f3 <= 2.500000f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } + } + } + } + } + } else { + if (f3 <= 0.500000f) { + if (f0 <= 6.292481f) { + if (f2 <= 8.906890f) { + if (f1 <= 8.292481f) { + if (f2 <= 7.935182f) { + return {8, 32, 4}; + } else { + return {8, 32, 8}; + } + } else { + return {8, 32, 4}; + } + } else { + return {8, 32, 8}; + } + } else { + if (f2 <= 10.421175f) { + if (f1 <= 8.792481f) { + if (f0 <= 7.292481f) { + return {8, 32, 4}; + } else { + if (f2 <= 8.406890f) { + return {8, 32, 2}; + } else { + return {8, 32, 4}; + } + } + } else { + return {8, 64, 4}; + } + } else { + if (f1 <= 8.792481f) { + if (f1 <= 7.792481f) { + return {8, 32, 8}; + } else { + if (f0 <= 6.792481f) { + return {8, 32, 8}; + } else { + return {8, 64, 8}; + } + } + } else { + return {8, 64, 4}; + } + } + } + } else { + if (f2 <= 9.421175f) { + if (f3 <= 1.500000f) { + if (f0 <= 6.792481f) { + if (f1 <= 8.792481f) { + return {8, 32, 4}; + } else { + return {8, 64, 4}; + } + } else { + if (f1 <= 8.792481f) { + if (f0 <= 7.292481f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } else { + return {8, 64, 2}; + } + } + } else { + if (f2 <= 7.435183f) { + if (f0 <= 7.292481f) { + if (f0 <= 6.292481f) { + return {8, 32, 1}; + } else { + return {8, 64, 1}; + } + } else { + return {8, 64, 1}; + } + } else { + if (f0 <= 6.792481f) { + if (f1 <= 8.792481f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } else { + if (f2 <= 8.406890f) { + return {8, 64, 1}; + } else { + return {8, 64, 2}; + } + } + } + } + } else { + if (f2 <= 11.421175f) { + if (f3 <= 1.500000f) { + if (f0 <= 6.792481f) { + if (f2 <= 10.421175f) { + return {8, 32, 4}; + } else { + return {8, 64, 8}; + } + } else { + return {8, 64, 4}; + } + } else { + return {8, 64, 4}; + } + } else { + if (f3 <= 2.500000f) { + if (f1 <= 7.792481f) { + return {8, 32, 8}; + } else { + if (f0 <= 5.792481f) { + return {8, 64, 8}; + } else { + return {8, 64, 4}; + } + } + } else { + if (f0 <= 6.292481f) { + if (f1 <= 8.292481f) { + return {8, 64, 4}; + } else { + return {16, 32, 4}; + } + } else { + if (f1 <= 8.292481f) { + return {16, 32, 4}; + } else { + return {16, 64, 4}; + } + } + } + } + } + } + } + } + } + } else { + if (f1 <= 6.792481f) { + if (f1 <= 4.792481f) { + if (f2 <= 8.906890f) { + if (f0 <= 12.292481f) { + if (f2 <= 7.435183f) { + if (f0 <= 9.292481f) { + if (f3 <= 1.500000f) { + return {8, 16, 4}; + } else { + return {8, 16, 1}; + } + } else { + return {8, 16, 1}; + } + } else { + if (f0 <= 10.792481f) { + if (f3 <= 2.500000f) { + if (f0 <= 9.792481f) { + if (f2 <= 8.406890f) { + return {8, 16, 4}; + } else { + return {16, 16, 8}; + } + } else { + if (f3 <= 0.500000f) { + return {8, 16, 4}; + } else { + return {8, 16, 1}; + } + } + } else { + if (f1 <= 4.292481f) { + return {16, 16, 4}; + } else { + return {16, 16, 2}; + } + } + } else { + if (f3 <= 2.500000f) { + if (f1 <= 4.292481f) { + if (f2 <= 8.406890f) { + return {8, 16, 2}; + } else { + return {16, 16, 4}; + } + } else { + return {16, 16, 2}; + } + } else { + if (f1 <= 4.292481f) { + return {16, 16, 2}; + } else { + return {32, 16, 1}; + } + } + } + } + } else { + if (f2 <= 7.435183f) { + if (f0 <= 13.292481f) { + if (f3 <= 1.500000f) { + return {8, 16, 1}; + } else { + return {16, 16, 1}; + } + } else { + return {16, 16, 1}; + } + } else { + if (f0 <= 14.292481f) { + if (f3 <= 1.500000f) { + if (f2 <= 8.406890f) { + if (f0 <= 12.792481f) { + return {32, 16, 2}; + } else { + return {16, 16, 2}; + } + } else { + if (f0 <= 13.292481f) { + return {16, 16, 4}; + } else { + return {16, 16, 2}; + } + } + } else { + if (f1 <= 4.292481f) { + return {32, 16, 2}; + } else { + if (f3 <= 2.500000f) { + return {32, 16, 2}; + } else { + return {16, 16, 1}; + } + } + } + } else { + if (f3 <= 2.500000f) { + if (f1 <= 4.292481f) { + if (f0 <= 15.792481f) { + return {32, 16, 2}; + } else { + return {16, 16, 1}; + } + } else { + return {16, 16, 1}; + } + } else { + if (f1 <= 4.292481f) { + return {16, 16, 1}; + } else { + return {32, 16, 1}; + } + } + } + } + } + } else { + if (f0 <= 11.292481f) { + if (f2 <= 10.921175f) { + if (f3 <= 1.500000f) { + if (f0 <= 9.792481f) { + if (f0 <= 8.292481f) { + return {8, 16, 8}; + } else { + return {16, 16, 8}; + } + } else { + if (f2 <= 9.421175f) { + return {16, 16, 4}; + } else { + if (f3 <= 0.500000f) { + return {16, 16, 8}; + } else { + return {16, 16, 4}; + } + } + } + } else { + if (f0 <= 9.292481f) { + if (f2 <= 9.421175f) { + return {16, 16, 4}; + } else { + return {16, 16, 8}; + } + } else { + if (f0 <= 10.792481f) { + return {16, 16, 4}; + } else { + if (f2 <= 9.921175f) { + return {32, 16, 4}; + } else { + return {16, 16, 4}; + } + } + } + } + } else { + if (f3 <= 0.500000f) { + if (f0 <= 8.792481f) { + return {8, 16, 8}; + } else { + return {16, 16, 8}; + } + } else { + if (f0 <= 9.792481f) { + if (f1 <= 4.292481f) { + if (f3 <= 1.500000f) { + return {8, 16, 8}; + } else { + return {16, 16, 8}; + } + } else { + if (f3 <= 2.500000f) { + return {16, 16, 8}; + } else { + return {32, 16, 8}; + } + } + } else { + if (f2 <= 11.421175f) { + if (f1 <= 4.292481f) { + return {16, 16, 8}; + } else { + return {32, 16, 8}; + } + } else { + return {16, 16, 8}; + } + } + } + } + } else { + if (f2 <= 10.921175f) { + if (f0 <= 14.292481f) { + if (f3 <= 2.500000f) { + return {16, 16, 4}; + } else { + if (f2 <= 10.421175f) { + if (f0 <= 13.292481f) { + return {32, 16, 2}; + } else { + return {16, 16, 2}; + } + } else { + return {16, 16, 2}; + } + } + } else { + if (f2 <= 9.921175f) { + if (f1 <= 4.292481f) { + if (f0 <= 14.792481f) { + return {16, 16, 4}; + } else { + return {32, 16, 2}; + } + } else { + if (f3 <= 1.500000f) { + return {32, 16, 2}; + } else { + return {16, 16, 1}; + } + } + } else { + if (f0 <= 14.792481f) { + if (f3 <= 1.500000f) { + return {32, 16, 4}; + } else { + return {16, 16, 2}; + } + } else { + return {16, 16, 2}; + } + } + } + } else { + if (f0 <= 12.792481f) { + if (f2 <= 11.421175f) { + if (f3 <= 1.500000f) { + return {16, 16, 4}; + } else { + return {32, 16, 8}; + } + } else { + if (f3 <= 2.500000f) { + if (f0 <= 12.292481f) { + return {16, 16, 8}; + } else { + return {32, 16, 8}; + } + } else { + return {32, 16, 8}; + } + } + } else { + if (f1 <= 4.292481f) { + if (f0 <= 14.792481f) { + if (f3 <= 1.500000f) { + return {16, 16, 8}; + } else { + return {16, 16, 4}; + } + } else { + if (f3 <= 2.500000f) { + return {16, 16, 4}; + } else { + return {16, 16, 2}; + } + } + } else { + if (f2 <= 11.421175f) { + if (f0 <= 14.292481f) { + return {16, 16, 4}; + } else { + return {32, 16, 4}; + } + } else { + return {16, 16, 4}; + } + } + } + } + } + } + } else { + if (f2 <= 9.421175f) { + if (f0 <= 10.292481f) { + if (f3 <= 1.500000f) { + if (f1 <= 5.792481f) { + if (f1 <= 5.292481f) { + if (f0 <= 9.292481f) { + if (f3 <= 0.500000f) { + return {8, 32, 8}; + } else { + return {8, 32, 4}; + } + } else { + if (f3 <= 0.500000f) { + return {8, 32, 4}; + } else { + return {8, 32, 2}; + } + } + } else { + if (f0 <= 9.292481f) { + if (f3 <= 0.500000f) { + return {8, 16, 4}; + } else { + return {16, 16, 4}; + } + } else { + if (f3 <= 0.500000f) { + return {8, 16, 2}; + } else { + return {16, 16, 2}; + } + } + } + } else { + if (f0 <= 8.792481f) { + if (f2 <= 7.435183f) { + return {8, 32, 2}; + } else { + return {8, 32, 4}; + } + } else { + if (f3 <= 0.500000f) { + if (f0 <= 9.292481f) { + return {8, 32, 4}; + } else { + return {8, 32, 2}; + } + } else { + if (f0 <= 9.292481f) { + return {8, 32, 2}; + } else { + return {8, 32, 1}; + } + } + } + } + } else { + if (f1 <= 5.292481f) { + if (f0 <= 9.292481f) { + if (f3 <= 2.500000f) { + return {8, 32, 4}; + } else { + return {8, 32, 2}; + } + } else { + if (f2 <= 8.406890f) { + return {8, 32, 1}; + } else { + return {8, 32, 4}; + } + } + } else { + if (f1 <= 6.292481f) { + if (f1 <= 5.792481f) { + if (f2 <= 7.435183f) { + return {16, 32, 1}; + } else { + return {16, 16, 4}; + } + } else { + if (f2 <= 7.435183f) { + return {8, 64, 1}; + } else { + return {8, 64, 2}; + } + } + } else { + if (f3 <= 2.500000f) { + if (f2 <= 8.406890f) { + return {8, 32, 1}; + } else { + return {16, 32, 4}; + } + } else { + if (f0 <= 8.792481f) { + return {8, 32, 1}; + } else { + return {16, 32, 1}; + } + } + } + } + } + } else { + if (f1 <= 6.292481f) { + if (f1 <= 5.792481f) { + if (f1 <= 5.292481f) { + if (f0 <= 13.792481f) { + if (f3 <= 2.500000f) { + return {8, 32, 1}; + } else { + return {16, 32, 1}; + } + } else { + return {16, 32, 1}; + } + } else { + if (f0 <= 14.292481f) { + return {16, 32, 1}; + } else { + if (f2 <= 7.435183f) { + return {16, 32, 1}; + } else { + return {32, 32, 1}; + } + } + } + } else { + if (f2 <= 7.435183f) { + return {8, 64, 1}; + } else { + if (f3 <= 1.500000f) { + if (f0 <= 11.792481f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } else { + if (f0 <= 15.292481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } + } + } + } else { + if (f0 <= 12.292481f) { + if (f3 <= 0.500000f) { + if (f2 <= 8.406890f) { + return {8, 32, 1}; + } else { + return {16, 32, 2}; + } + } else { + if (f3 <= 2.500000f) { + if (f2 <= 8.406890f) { + return {8, 32, 1}; + } else { + return {16, 32, 1}; + } + } else { + return {16, 32, 1}; + } + } + } else { + return {16, 32, 1}; + } + } + } + } else { + if (f0 <= 11.292481f) { + if (f1 <= 6.292481f) { + if (f1 <= 5.792481f) { + if (f1 <= 5.292481f) { + if (f3 <= 1.500000f) { + if (f0 <= 8.792481f) { + return {8, 32, 8}; + } else { + return {16, 32, 8}; + } + } else { + if (f0 <= 10.292481f) { + return {16, 32, 8}; + } else { + return {16, 32, 4}; + } + } + } else { + if (f3 <= 1.500000f) { + if (f0 <= 9.292481f) { + return {16, 16, 8}; + } else { + return {16, 32, 4}; + } + } else { + if (f0 <= 8.792481f) { + return {16, 32, 4}; + } else { + return {32, 32, 2}; + } + } + } + } else { + if (f3 <= 2.500000f) { + if (f0 <= 8.792481f) { + if (f3 <= 1.500000f) { + return {8, 64, 8}; + } else { + return {8, 64, 4}; + } + } else { + if (f0 <= 10.792481f) { + return {8, 64, 4}; + } else { + return {16, 64, 2}; + } + } + } else { + if (f0 <= 8.792481f) { + return {8, 64, 2}; + } else { + if (f0 <= 10.292481f) { + return {16, 64, 2}; + } else { + return {8, 64, 2}; + } + } + } + } + } else { + if (f2 <= 10.921175f) { + if (f3 <= 1.500000f) { + if (f0 <= 9.292481f) { + return {16, 32, 8}; + } else { + return {16, 32, 4}; + } + } else { + if (f0 <= 10.292481f) { + return {16, 32, 4}; + } else { + return {16, 32, 2}; + } + } + } else { + if (f0 <= 10.292481f) { + return {16, 32, 8}; + } else { + if (f3 <= 1.500000f) { + return {16, 32, 8}; + } else { + return {16, 32, 4}; + } + } + } + } + } else { + if (f1 <= 5.292481f) { + if (f0 <= 14.792481f) { + if (f2 <= 10.421175f) { + if (f3 <= 1.500000f) { + if (f0 <= 13.292481f) { + return {16, 32, 4}; + } else { + return {16, 32, 2}; + } + } else { + if (f0 <= 12.292481f) { + return {16, 32, 4}; + } else { + return {16, 32, 2}; + } + } + } else { + if (f3 <= 2.500000f) { + return {16, 32, 4}; + } else { + if (f0 <= 13.792481f) { + return {16, 32, 4}; + } else { + return {16, 32, 2}; + } + } + } + } else { + if (f2 <= 11.421175f) { + return {16, 32, 2}; + } else { + if (f3 <= 1.500000f) { + return {16, 32, 4}; + } else { + return {16, 32, 2}; + } + } + } + } else { + if (f1 <= 6.292481f) { + if (f1 <= 5.792481f) { + if (f2 <= 10.421175f) { + if (f0 <= 13.292481f) { + return {16, 32, 1}; + } else { + return {32, 32, 1}; + } + } else { + return {32, 32, 2}; + } + } else { + if (f2 <= 10.421175f) { + if (f2 <= 9.921175f) { + return {8, 64, 1}; + } else { + return {8, 64, 2}; + } + } else { + return {16, 64, 2}; + } + } + } else { + if (f2 <= 11.421175f) { + if (f2 <= 9.921175f) { + if (f3 <= 1.500000f) { + return {16, 32, 2}; + } else { + return {16, 32, 1}; + } + } else { + if (f3 <= 2.500000f) { + return {16, 32, 2}; + } else { + return {16, 32, 1}; + } + } + } else { + if (f0 <= 13.292481f) { + if (f3 <= 0.500000f) { + return {16, 32, 8}; + } else { + return {16, 32, 4}; + } + } else { + if (f3 <= 1.500000f) { + return {16, 32, 4}; + } else { + return {16, 32, 2}; + } + } + } + } + } + } + } + } + } else { + if (f2 <= 9.921175f) { + if (f2 <= 8.906890f) { + if (f0 <= 10.792481f) { + if (f2 <= 7.435183f) { + if (f3 <= 2.500000f) { + if (f0 <= 9.292481f) { + return {8, 64, 1}; + } else { + if (f1 <= 8.292481f) { + return {8, 64, 1}; + } else { + if (f1 <= 8.792481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } + } + } else { + if (f1 <= 7.792481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } + } else { + if (f3 <= 2.500000f) { + if (f2 <= 8.406890f) { + if (f0 <= 9.792481f) { + return {8, 64, 2}; + } else { + return {8, 64, 1}; + } + } else { + if (f0 <= 8.792481f) { + if (f3 <= 0.500000f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } else { + if (f3 <= 1.500000f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } + } + } else { + if (f1 <= 8.292481f) { + if (f0 <= 8.292481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } else { + if (f0 <= 9.292481f) { + if (f1 <= 8.792481f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } else { + return {16, 32, 1}; + } + } + } + } + } else { + if (f0 <= 14.792481f) { + if (f2 <= 8.406890f) { + if (f3 <= 2.500000f) { + return {8, 64, 1}; + } else { + if (f0 <= 12.792481f) { + return {8, 64, 1}; + } else { + if (f2 <= 7.435183f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } + } + } else { + if (f3 <= 0.500000f) { + if (f0 <= 13.292481f) { + if (f1 <= 8.292481f) { + return {16, 64, 1}; + } else { + return {8, 64, 1}; + } + } else { + return {8, 64, 1}; + } + } else { + if (f0 <= 14.292481f) { + if (f0 <= 11.792481f) { + return {8, 64, 1}; + } else { + return {16, 32, 1}; + } + } else { + return {8, 64, 1}; + } + } + } + } else { + if (f2 <= 7.435183f) { + if (f3 <= 1.500000f) { + return {8, 64, 1}; + } else { + if (f1 <= 8.292481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } + } else { + if (f2 <= 8.406890f) { + if (f3 <= 0.500000f) { + if (f1 <= 8.292481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } else { + return {16, 64, 1}; + } + } else { + if (f1 <= 8.292481f) { + return {8, 64, 1}; + } else { + if (f0 <= 15.792481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } + } + } + } + } + } else { + if (f0 <= 12.792481f) { + if (f0 <= 8.792481f) { + if (f3 <= 2.500000f) { + if (f3 <= 1.500000f) { + if (f1 <= 7.792481f) { + if (f3 <= 0.500000f) { + return {8, 32, 4}; + } else { + return {8, 64, 4}; + } + } else { + if (f3 <= 0.500000f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } + } else { + return {8, 64, 2}; + } + } else { + if (f1 <= 7.792481f) { + return {16, 64, 1}; + } else { + return {8, 64, 2}; + } + } + } else { + if (f2 <= 9.421175f) { + if (f3 <= 2.500000f) { + return {8, 64, 2}; + } else { + if (f0 <= 10.292481f) { + return {8, 64, 2}; + } else { + if (f0 <= 11.292481f) { + return {8, 64, 1}; + } else { + return {16, 32, 1}; + } + } + } + } else { + if (f0 <= 9.792481f) { + if (f3 <= 0.500000f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } else { + if (f3 <= 2.500000f) { + return {8, 64, 2}; + } else { + if (f0 <= 11.292481f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } + } + } + } + } else { + if (f2 <= 9.421175f) { + if (f1 <= 8.292481f) { + if (f0 <= 14.792481f) { + return {8, 64, 1}; + } else { + if (f3 <= 2.500000f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } + } else { + if (f3 <= 1.500000f) { + if (f0 <= 14.792481f) { + return {8, 64, 1}; + } else { + return {16, 64, 1}; + } + } else { + return {16, 64, 1}; + } + } + } else { + if (f3 <= 1.500000f) { + if (f1 <= 8.292481f) { + if (f0 <= 14.792481f) { + if (f0 <= 13.792481f) { + return {8, 64, 2}; + } else { + return {8, 64, 1}; + } + } else { + return {8, 64, 2}; + } + } else { + if (f0 <= 15.292481f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } + } else { + if (f0 <= 13.292481f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } + } + } + } + } else { + if (f2 <= 11.421175f) { + if (f0 <= 9.792481f) { + if (f3 <= 1.500000f) { + if (f3 <= 0.500000f) { + if (f1 <= 7.792481f) { + if (f0 <= 8.792481f) { + return {8, 64, 8}; + } else { + return {8, 64, 4}; + } + } else { + return {8, 64, 4}; + } + } else { + if (f1 <= 7.292481f) { + return {8, 64, 4}; + } else { + if (f0 <= 8.792481f) { + return {8, 64, 4}; + } else { + if (f1 <= 8.292481f) { + return {8, 64, 2}; + } else { + return {16, 32, 4}; + } + } + } + } + } else { + if (f2 <= 10.421175f) { + if (f1 <= 8.792481f) { + if (f3 <= 2.500000f) { + if (f0 <= 8.792481f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } else { + if (f0 <= 8.792481f) { + return {16, 32, 4}; + } else { + return {8, 64, 2}; + } + } + } else { + return {8, 64, 2}; + } + } else { + if (f1 <= 7.792481f) { + if (f3 <= 2.500000f) { + if (f0 <= 8.792481f) { + return {8, 64, 2}; + } else { + return {8, 64, 4}; + } + } else { + if (f0 <= 8.792481f) { + return {8, 64, 4}; + } else { + return {16, 64, 2}; + } + } + } else { + if (f0 <= 8.792481f) { + return {16, 64, 2}; + } else { + if (f3 <= 2.500000f) { + return {16, 64, 2}; + } else { + return {16, 64, 4}; + } + } + } + } + } + } else { + if (f2 <= 10.421175f) { + if (f0 <= 14.792481f) { + if (f3 <= 2.500000f) { + if (f0 <= 11.792481f) { + if (f3 <= 0.500000f) { + return {8, 64, 4}; + } else { + return {8, 64, 2}; + } + } else { + return {8, 64, 2}; + } + } else { + if (f1 <= 7.792481f) { + if (f0 <= 12.792481f) { + return {16, 64, 1}; + } else { + return {8, 64, 2}; + } + } else { + if (f0 <= 11.792481f) { + return {16, 64, 1}; + } else { + return {8, 64, 2}; + } + } + } + } else { + if (f3 <= 0.500000f) { + return {8, 64, 2}; + } else { + if (f3 <= 2.500000f) { + if (f1 <= 7.792481f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } else { + return {16, 64, 1}; + } + } + } + } else { + if (f0 <= 14.792481f) { + if (f2 <= 10.921175f) { + if (f3 <= 2.500000f) { + return {16, 64, 2}; + } else { + if (f0 <= 13.292481f) { + return {32, 64, 4}; + } else { + return {8, 64, 2}; + } + } + } else { + if (f0 <= 11.292481f) { + if (f3 <= 2.500000f) { + return {16, 64, 2}; + } else { + return {32, 64, 4}; + } + } else { + if (f0 <= 12.292481f) { + return {32, 64, 4}; + } else { + return {16, 32, 2}; + } + } + } + } else { + if (f3 <= 1.500000f) { + if (f2 <= 10.921175f) { + return {16, 64, 1}; + } else { + if (f1 <= 7.792481f) { + return {32, 64, 4}; + } else { + return {16, 32, 2}; + } + } + } else { + if (f2 <= 10.921175f) { + return {8, 64, 2}; + } else { + if (f1 <= 7.792481f) { + return {16, 32, 2}; + } else { + return {16, 64, 1}; + } + } + } + } + } + } + } else { + if (f0 <= 10.792481f) { + if (f3 <= 1.500000f) { + if (f0 <= 8.792481f) { + if (f1 <= 7.292481f) { + return {8, 64, 4}; + } else { + if (f3 <= 0.500000f) { + if (f0 <= 8.292481f) { + return {8, 64, 4}; + } else { + return {8, 32, 8}; + } + } else { + if (f1 <= 8.292481f) { + return {8, 32, 8}; + } else { + return {8, 64, 2}; + } + } + } + } else { + if (f1 <= 7.792481f) { + if (f0 <= 9.792481f) { + return {8, 64, 4}; + } else { + if (f3 <= 0.500000f) { + return {8, 64, 2}; + } else { + return {16, 64, 2}; + } + } + } else { + if (f0 <= 9.792481f) { + if (f3 <= 0.500000f) { + return {8, 64, 4}; + } else { + return {16, 64, 2}; + } + } else { + return {16, 64, 2}; + } + } + } + } else { + if (f1 <= 7.292481f) { + if (f0 <= 8.792481f) { + return {8, 32, 8}; + } else { + return {16, 64, 2}; + } + } else { + if (f0 <= 8.792481f) { + if (f1 <= 7.792481f) { + return {8, 64, 2}; + } else { + if (f3 <= 2.500000f) { + return {16, 64, 2}; + } else { + return {16, 32, 4}; + } + } + } else { + if (f3 <= 2.500000f) { + if (f0 <= 9.792481f) { + return {16, 64, 2}; + } else { + return {16, 32, 4}; + } + } else { + if (f0 <= 9.292481f) { + return {16, 32, 4}; + } else { + return {32, 64, 4}; + } + } + } + } + } + } else { + if (f0 <= 14.292481f) { + if (f3 <= 2.500000f) { + if (f0 <= 12.792481f) { + return {16, 64, 2}; + } else { + if (f3 <= 1.500000f) { + if (f1 <= 7.792481f) { + return {16, 32, 4}; + } else { + return {32, 64, 4}; + } + } else { + if (f1 <= 7.792481f) { + return {32, 64, 4}; + } else { + return {16, 64, 2}; + } + } + } + } else { + if (f0 <= 11.792481f) { + if (f1 <= 7.792481f) { + return {32, 64, 4}; + } else { + return {16, 64, 2}; + } + } else { + return {16, 64, 2}; + } + } + } else { + if (f1 <= 8.292481f) { + if (f3 <= 0.500000f) { + if (f0 <= 15.292481f) { + return {32, 64, 4}; + } else { + return {16, 64, 2}; + } + } else { + return {16, 64, 2}; + } + } else { + if (f3 <= 2.500000f) { + if (f3 <= 1.500000f) { + return {16, 64, 2}; + } else { + if (f0 <= 15.292481f) { + return {16, 64, 2}; + } else { + return {8, 64, 2}; + } + } + } else { + if (f0 <= 15.292481f) { + return {8, 64, 2}; + } else { + return {16, 64, 1}; + } + } + } + } + } + } + } + } + } +} + +} // namespace sgconv_tree::xe3lpg 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..3fb4386cca4e0 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc @@ -0,0 +1,351 @@ +// 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). Mirrors the ONNX Conv semantics the +// WebGPU EP exercises after its internal NCHW->NHWC layout transform, so the +// subgroup-matrix implicit-GEMM conv can be validated against it. +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, one case per branch of GetActivationSnippet in +// core/providers/webgpu/nn/fuse_utils.cc. The subgroup-matrix conv builds its epilogue +// from that same snippet, so these have to stay in step with it: if a branch there +// changes shape, the matching case here must change too. +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. +// Inputs/outputs are standard ONNX NCHW; the WebGPU EP transforms to its +// channels-last kernels internally, which is where the subgroup-matrix conv runs. +// +// A non-empty `activation` switches the node to com.microsoft::FusedConv, which the +// layout transform rewrites to kMSInternalNHWCDomain::Conv (Conv) with the +// attribute preserved -- the same kernel, now with a fused epilogue. +// +// These cases are device-independent by construction, so they are enabled rather than +// DISABLED_ like the *_large tests in this directory. The subgroup-matrix path needs an +// adapter that reports the 8x16x16 F16 subgroup-matrix config plus a vendor tiling +// policy; where that is missing, CanApplySubgroupMatrixConv() declines and the EP falls +// back to its regular conv kernels, which must produce the same result against the same +// reference. So a failure here is a real defect on either path, not a capability gap -- +// only the absence of the WebGPU EP itself is skipped. The problem sizes are small +// (<= 32x20x20 in, <= 130 out channels) to keep them cheap enough for that. +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) { +#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}); + test->AddAttribute("pads", std::vector{pad, pad, pad, pad}); + test->AddAttribute("strides", std::vector{stride, stride}); + test->AddAttribute("dilations", std::vector{dilation, dilation}); + + test->AddInput("X", x_dims, FloatsToMLFloat16s(x_vals)); + // A constant weight (the default) takes the branch that matters most: PrePackInternal + // must decline to prepack so the weight stays OIHW, and the impl caches its own OHWI + // transpose across runs. A non-constant weight instead redoes the transpose every + // run; RuntimeWeight covers that 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, which is what keeps the +// unchecked subgroupMatrixLoad of the weight inside the buffer. tile_n is clamped to +// the largest multiple of 16 that fits in Cout, so with Cout in (16, 32) the shifted +// tile overlaps the previous one by all but (Cout - 16) columns -- n_skip near kTileN, +// the narrowest write-out this kernel ever 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 neither the candidate list nor the +// tuned tree 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); +} + +#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); + } +} + +// 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 From 76a1a263a4b8ae65890a324bec7d102571de7d62 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Tue, 13 Oct 2026 10:56:37 +0800 Subject: [PATCH 02/14] Add a ConvHeuristicTiling function and rebase code --- .../webgpu/nn/subgroup_matrix_conv.cc | 54 +++++++------- .../webgpu/nn/subgroup_matrix_conv.h | 14 ++-- .../math/subgroup_matrix_tiling_selector.cc | 73 ++++++++++++++++++- 3 files changed, 108 insertions(+), 33 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index 483e7fd215dfb..5ab526af8e497 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -47,19 +47,20 @@ constexpr uint32_t kMaxConvScratchElems = 16384; constexpr uint32_t kMaxSgMatCountM = 8; constexpr uint32_t kMaxSgMatCountN = 4; -// Returns the device config index for the fixed 8x16x16 F16 subgroup-matrix -// config this kernel is implemented for, or -1 when the device does not report -// it. Also requires subgroup-size control and the vendor tiling policy (Intel). +// Returns the device's subgroup-matrix config for the fixed 8x16x16 F16 config +// this kernel is implemented for, or nullopt when the device does not report it. +// Also requires subgroup-size control and the vendor tiling policy (Intel). // // Single source of truth for device eligibility: CanApplySubgroupMatrixConv (which // Conv::PrePackInternal consults to skip prepacking) and CreateSubgroupMatrixConvImpl // both go through here, so the prepack decision and the runtime path can never // disagree about whether this device is served. -int32_t GetSupportedConfigIndex(const ComputeContextBase& context) { - int32_t config_index = 0; - if (!IsSubgroupMatrixConfigSupported(context, /*is_fp16=*/true, config_index) || - !supported_subgroup_matrix_configs[config_index].Is(8, 16, 16)) { - return -1; +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; } // The adapters reporting that config expose a 16-32 subgroup size range, while the // kernel hard-codes kSubgroupMatrixSubgroupSize lanes per subgroup (sg_index = @@ -68,12 +69,12 @@ int32_t GetSupportedConfigIndex(const ComputeContextBase& context) { // store the identical tile -- but it wastes about half the subgroup-matrix // throughput, so the size is pinned with subgroup-size control instead. if (!context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { - return -1; + return std::nullopt; } if (context.AdapterInfo().vendor != std::string_view{"intel"}) { - return -1; + return std::nullopt; } - return config_index; + return config; } // The bias add and the fused activation are not template parameters: the template @@ -109,9 +110,9 @@ class SubgroupMatrixConvImpl final : public Conv::Co using Base = typename ConvT::ConvOptImpl; public: - SubgroupMatrixConvImpl(const ConvT& parent, int32_t config_index, + SubgroupMatrixConvImpl(const ConvT& parent, SubgroupMatrixConfig config, SubgroupMatrixTilingSelector tiling_selector) - : Base(parent), config_index_(config_index), tiling_selector_(std::move(tiling_selector)) {} + : Base(parent), config_(config), tiling_selector_(std::move(tiling_selector)) {} Status Compute(ComputeContext& context, /*out*/ bool& handled) override { handled = false; @@ -269,7 +270,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co ohwi_weight = &run_local_ohwi_weight; } - const auto& config = supported_subgroup_matrix_configs[config_index_]; + 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, @@ -304,14 +305,17 @@ class SubgroupMatrixConvImpl final : public Conv::Co const Tensor* bias = context.InputCount() > 2 ? context.Input(2) : nullptr; const bool has_bias = bias != nullptr; - SubgroupMatrixConvProgram program{has_bias, config_index_, sg_mat_count_m, sg_mat_count_n, split_k, vec_size, + SubgroupMatrixConvProgram program{has_bias, config, sg_mat_count_m, sg_mat_count_n, split_k, vec_size, activation}; program.SetWorkgroupSize(kSubgroupMatrixSubgroupSize * split_k); - // Pin the lanes per subgroup the kernel assumes; GetSupportedConfigIndex has + // Pin the lanes per subgroup the kernel assumes; GetSupportedConfig has // already established that the device supports subgroup-size control. program.SetSubgroupSize(kSubgroupMatrixSubgroupSize); program.SetDispatchGroupSize(dispatch_x, dispatch_y, batch); - program.CacheHint(has_bias, config_index_, sg_mat_count_m, sg_mat_count_n, split_k, vec_size, + 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, activation.CacheKey()) .AddInputs({{src, ProgramTensorMetadataDependency::TypeAndRank, static_cast(vec_size)}, {ohwi_weight, ProgramTensorMetadataDependency::TypeAndRank, 1}}) @@ -363,7 +367,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co "Failed to transpose the subgroup-matrix Conv weight to OHWI."); } - const int32_t config_index_; + const SubgroupMatrixConfig config_; SubgroupMatrixTilingSelector tiling_selector_; // OIHW -> OHWI transposed constant weight, built once by EnsureTransposedKernel. @@ -401,7 +405,7 @@ Status SubgroupMatrixConvProgram::GenerateShaderCode(ShaderHelper& shader) const << " " << output.SetByOffset("output_offset", "value") << "\n" << "}\n"; - const auto& config = supported_subgroup_matrix_configs[config_index_]; + 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_); } @@ -475,8 +479,8 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, return false; } - const int32_t config_index = GetSupportedConfigIndex(context); - if (config_index < 0) { + 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} @@ -507,9 +511,9 @@ std::unique_ptr::ConvOptImpl> CreateSu } else { // Only run on devices that report the fixed 8x16x16 F16 subgroup-matrix config // this kernel is implemented for, and for which a vendor tiling policy exists - // (Intel today - GetSupportedConfigIndex enforces both). - const int32_t config_index = GetSupportedConfigIndex(context); - if (config_index < 0) { + // (Intel today - GetSupportedConfig enforces both). + const std::optional config = GetSupportedConfig(context); + if (!config) { return nullptr; } // The Conv selector shares the autotuner hooks (tile override, device-arch @@ -519,7 +523,7 @@ std::unique_ptr::ConvOptImpl> CreateSu if (!tiling_selector) { return nullptr; } - return std::make_unique>(parent, config_index, + return std::make_unique>(parent, *config, std::move(tiling_selector)); } } diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h index 3158d52173575..cf4dc9b388999 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -9,6 +9,7 @@ #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" @@ -21,10 +22,9 @@ namespace webgpu { // subgroupMatrixMultiplyAccumulate. The activation operand is the im2col of the // input, gathered on the fly into workgroup memory (never materialized in global // memory); the weight is pre-transposed to OHWI, i.e. B stored column-major. -// config_index selects the device subgroup-matrix config (into -// supported_subgroup_matrix_configs); sg_mat_count_m/n select how many subgroup -// matrices the output tile spans along M/N; split_k is the number of subgroups that -// cooperatively reduce K. +// config is the device subgroup-matrix config selected by SelectSubgroupMatrixConfig; +// sg_mat_count_m/n select how many subgroup matrices the output tile spans along +// M/N; split_k is the number of subgroups that cooperatively reduce K. // // The bias add and the fused activation live in a generated `write_output` helper // rather than in the WGSL template: the activation body comes from @@ -32,12 +32,12 @@ namespace webgpu { // support works here without the template knowing the kinds. class SubgroupMatrixConvProgram final : public Program { public: - SubgroupMatrixConvProgram(bool has_bias, int32_t config_index, + 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, const Activation& activation) : Program{"SubgroupMatrixConv"}, has_bias_(has_bias), - config_index_(config_index), + config_(config), sg_mat_count_m_(sg_mat_count_m), sg_mat_count_n_(sg_mat_count_n), split_k_(split_k), @@ -69,7 +69,7 @@ class SubgroupMatrixConvProgram final : public Program Date: Wed, 14 Oct 2026 10:33:50 +0800 Subject: [PATCH 03/14] modify the ConvHeuristicTiling --- .../webgpu/nn/subgroup_matrix_conv.cc | 17 +- .../math/subgroup_matrix_tiling_selector.cc | 191 +- .../math/subgroup_matrix_tiling_selector.h | 7 +- .../intel/nn/subgroup_matrix_conv_tuned.inc | 1871 ----------------- .../webgpu/subgroup_matrix_conv_test.cc | 4 +- 5 files changed, 81 insertions(+), 2009 deletions(-) delete mode 100644 onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index 5ab526af8e497..3bed70902d097 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -41,9 +41,8 @@ constexpr uint32_t kMaxConvScratchElems = 16384; // accumulators sg_mat_c_ for mi < 8 and ni < 4. A tiling above these would be // accepted by the selector but silently mis-computed -- the kernel would accumulate // only the blocks it generated while the write-out still covers the whole tile. The -// tile candidates the selector picks from stay inside this, but the Conv pretuned tree -// is a generated file dropped into the build after a sweep, so enforce rather than -// assume. +// tile candidates the selector picks from stay inside this, so this is enforced as a +// safety net rather than something expected to fire. constexpr uint32_t kMaxSgMatCountM = 8; constexpr uint32_t kMaxSgMatCountN = 4; @@ -209,7 +208,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co // Pick the per-problem output tiling from the vendor policy (Intel). The GEMM // dims map directly: A [M, K] x B [K, N] with one z-slice per conv batch. The - // Conv selector uses the Conv scratch budget and Conv-specific pretuned data. + // Conv selector uses the Conv scratch budget and the Conv-specific heuristic. const std::optional tiling = tiling_selector_(context, M, N, K, batch); if (!tiling) { return Status::OK(); @@ -285,8 +284,8 @@ class SubgroupMatrixConvImpl final : public Conv::Co // Safety net: the Conv selector already enforces the Conv scratch budget // (output tiles + im2col A staging), so this clamp normally does nothing; it - // guards against a stray oversized split-K (e.g. a hand-edited pretuned entry). - // Only split_k is clamped, not the tile area: the largest candidate tile needs + // guards against a stray oversized split-K from the selector. Only split_k is + // clamped, not the tile area: the largest candidate tile needs // 64 * (64 + 16) = 5120 elements at split_k 1, well inside the budget. uint32_t split_k = tiling->split_k; while (split_k > 1 && @@ -516,9 +515,9 @@ std::unique_ptr::ConvOptImpl> CreateSu if (!config) { return nullptr; } - // The Conv selector shares the autotuner hooks (tile override, device-arch - // capture) with MatMul, but budgets the extra im2col staging the Conv kernel - // keeps in workgroup memory and consults the Conv-specific pretuned data. + // The Conv selector shares the tile candidates with MatMul, but budgets the + // extra im2col staging the Conv kernel keeps in workgroup memory and uses the + // Conv-specific heuristic. SubgroupMatrixTilingSelector tiling_selector = intel::CreateSubgroupMatrixConvTilingSelector(context); if (!tiling_selector) { return nullptr; 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 6d5d60f9ec633..536fce9675171 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 @@ -24,30 +24,6 @@ namespace onnxruntime { namespace webgpu { namespace intel { -// Optional Conv-specific pretuned decision tree, emitted by -// tools/python/gen_sgmm_tree.py --mode conv from an on-device Conv sweep. When the -// .inc is present it defines sgconv_tree::::{kArch, PredictTilingTree} and is -// consulted before the Conv heuristic (see LookupConvTree / SelectConvTiling); when -// absent the heuristic is used. The file is dropped in after a sweep, so the build -// must not depend on it existing. -// -// One such tree is checked in, so the guard looks redundant at first glance. It is -// not: the tree is measured data, not source, and the include has to keep working -// both ways -- a build that deletes it (a fork with no interest in shipping Intel -// tuning data, or a size-constrained target) must still compile, and a new arch's -// tree must be droppable in without touching the build files. The heuristic is a -// complete fallback, so the only cost of absence is tiling quality. -// -// Included here, inside the namespace, because the generated code names -// SubgroupMatrixTiling unqualified: it is an oracle over that type and carries no -// includes or namespace boilerplate of its own. -#if defined(__has_include) -#if __has_include("core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc") -#include "core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc" -#define ORT_WEBGPU_HAS_SGCONV_TREE 1 -#endif -#endif - namespace { // Subgroup matrix configuration used by this implementation (Intel Xe2/Xe3, F16). @@ -166,78 +142,82 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t return {tile_m, tile_n, split_k}; } -// Arch the sweep behind ConvHeuristicTiling (below) was measured on. Kept as a -// plain literal, independent of ORT_WEBGPU_HAS_SGCONV_TREE / subgroup_matrix_conv_tuned.inc, -// because this heuristic ships as checked-in source and must stay usable even in a -// build that drops the generated tree. +// Arch the sweep behind ConvHeuristicTiling (below) was measured on. constexpr std::string_view kConvHeuristicArch = "xe-3lpg"; -// Restricted-candidate approximation of the generated Conv tree -// (subgroup_matrix_conv_tuned.inc), fit to the same sweep but for kConvHeuristicArch -// only -- SelectConvTiling uses it as the fallback when the tree is absent or its -// prediction fails validation, never for another arch with no sweep data of its own. -// -// Picks among 6 (tile_m, tile_n, split_k) combos, each within 20% regret of the true -// best tiling for its region; deliberately narrower than the tree, so some shapes get -// a "good enough" pick here instead of the true optimum. -// -// batch is unused (not a useful splitting feature on the swept range). split_k is not -// guaranteed to fit K's block count -- callers must run the result through -// ClampSplitKToProblem, like a tuned-table/tree result. -SubgroupMatrixTiling ConvHeuristicTiling(uint32_t M, uint32_t N, uint32_t K, uint32_t /*batch*/) { - if (M <= 512u) { - if (M <= 160u) { - if (N <= 96u) { - return {8, 32, 8}; - } else { - if (M <= 32u) { - return {8, 32, 8}; +// Same decisions as ConvHeuristicTilingV3, verified against all 16,535 shapes to +// produce identical predictions (0 mismatches): every M/N/K threshold slid to the +// nearest multiple of 16 (or 8 where no multiple of 16 fits) that stays strictly +// between the same two real data values the original threshold separated, so no +// shape crosses leaves. batch's threshold (6) has no such multiple in its gap +// (4, 8) -- rounding it to 8 would pull real batch=8 shapes into the wrong leaf -- +// so it is left as-is. Used by SelectConvTiling for kConvHeuristicArch. +SubgroupMatrixTiling ConvHeuristicTiling(uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { + if (K <= 2016u) { + if (N <= 112u) { + if (N <= 24u) { + if (M <= 224u) { + return {8, 16, 8}; } else { - return {8, 64, 8}; + if (K <= 416u) { + return {16, 16, 2}; + } else { + if (batch <= 6u) { + if (M <= 896u) { + return {8, 16, 8}; + } else { + return {16, 16, 4}; + } + } else { + return {16, 16, 4}; + } + } } - } - } else { - if (N <= 384u) { - if (K <= 1024u) { - return {8, 64, 4}; + } else { + if (M <= 512u) { + if (batch <= 6u) { + return {8, 32, 8}; + } else { + return {8, 32, 4}; + } } else { - return {8, 64, 8}; + if (K <= 512u) { + return {8, 32, 1}; + } else { + return {16, 32, 4}; + } } - } else { - return {16, 64, 2}; } + } else { + return {8, 64, 2}; } } else { - if (K <= 512u) { - if (N <= 48u) { - return {16, 32, 2}; - } else { - return {16, 64, 2}; - } - } else { - if (N <= 48u) { - if (M <= 2048u) { - return {8, 32, 8}; + if (N <= 112u) { + if (N <= 24u) { + if (M <= 224u) { + return {8, 16, 8}; } else { - return {16, 32, 2}; + return {16, 16, 8}; } } else { - if (M <= 896u) { - return {8, 64, 8}; + if (M <= 512u) { + return {8, 32, 8}; } else { - return {16, 64, 2}; + return {16, 32, 4}; } } + } else { + return {8, 64, 8}; } } } -// Repairs a pretuned tiling to this problem by halving split_k until it fits both -// K's block count and the scratch budget. The tuned trees are performance oracles -// that never see K's block count (they are trained on the problem's log dims only), -// so they can propose split_k > K / kSubgroupMatrixK on a problem outside the swept -// envelope. Clamping keeps the tuned tile shape -- the part that carries most of the -// win -- instead of discarding the prediction and falling back to the heuristic. +// Repairs a predicted tiling to this problem by halving split_k until it fits both +// K's block count and the scratch budget. ConvHeuristicTiling is a performance oracle +// that never sees K's block count (it is trained on the problem's log dims only), so it +// can propose split_k > K / kSubgroupMatrixK on a problem outside its swept envelope. +// Clamping keeps the predicted tile shape -- the part that carries most of the win -- +// instead of discarding the prediction and falling back to the arch-agnostic heuristic. void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) { const uint32_t k_blocks = K / kSubgroupMatrixK; while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) { @@ -251,8 +231,8 @@ void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_ // from the weight buffer, and subgroupMatrixLoad is not bounds-checked, so a trailing // partial N tile is shifted back to N - tile_n. That shift only exists when // tile_n <= N. Neither source guarantees it -- HeuristicTiling only skips candidates -// larger than N for all but the smallest one, and the tuned trees are unconstrained -// oracles -- so the cap is applied to both. Shrinking tile_n keeps it a multiple of +// larger than N for all but the smallest one, and ConvHeuristicTiling is an +// unconstrained oracle -- so the cap is applied to both. Shrinking tile_n keeps it a multiple of // kSubgroupMatrixN and only lowers ScratchElems, so it cannot invalidate a tiling // that was already valid. // @@ -299,34 +279,15 @@ std::optional LookupPretunedTiling(std::string_view arch, return std::nullopt; } -// Looks up the Conv-specific pretuned decision tree for this arch. One `if` per -// tuned arch; each generated tree carries its own kArch, so the arch spelling lives -// with the data instead of being duplicated here. Returns nullopt when the arch has -// not been swept (or no tree was dropped into the build), in which case the caller -// uses the Conv heuristic. -std::optional LookupConvTree([[maybe_unused]] std::string_view arch, - [[maybe_unused]] uint32_t M, - [[maybe_unused]] uint32_t N, - [[maybe_unused]] uint32_t K, - [[maybe_unused]] uint32_t batch) { -#ifdef ORT_WEBGPU_HAS_SGCONV_TREE - if (arch == sgconv_tree::xe3lpg::kArch) { - return sgconv_tree::xe3lpg::PredictTilingTree(M, N, K, batch); - } -#endif - 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. // -// Only for batch-blind sources. The Conv tree takes batch as a training feature and -// was fitted on measurements at batch 1/2/4/8, so its split-K already prices batch in; -// applying this on top of it double-counts occupancy. Measured over the xe-3lpg sweep, -// doing so overrode 41% of the tuned picks and raised mean regret from 3.1% to 7.2%. +// Only for the batch-blind MatMul pretuned table: SelectConvTiling's ConvHeuristicTiling +// already takes batch into account when picking split-K, so applying this on top of it +// would double-count occupancy. void ClampSplitKForBatch(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); @@ -352,27 +313,11 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, return HeuristicTiling(arch, M, N, K, batch, kMatMulScratchNPad); } -// Conv counterpart of SelectTiling. Consults the Conv-specific pretuned tree (when -// one was dropped in for this arch) before falling back to the Conv heuristic; both -// use the Conv scratch budget. There is no MatMul-table lookup here: the Conv -// kernel's cost profile (the im2col gather) differs, and its scratch budget is -// stricter, so it must not reuse the MatMul-tuned entries. +// Conv counterpart of SelectTiling: uses the Conv heuristic (kConvHeuristicArch) or the +// arch-agnostic occupancy heuristic, both under the Conv scratch budget. There is no +// MatMul-table lookup here: the Conv kernel's cost profile (the im2col gather) differs, +// and its scratch budget is stricter, so it must not reuse the MatMul-tuned entries. SubgroupMatrixTiling SelectConvTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - if (std::optional tuned = LookupConvTree(arch, M, N, K, batch)) { - // Unlike the MatMul table this prediction already accounts for batch (it is a - // model input), so it is taken as-is apart from the correctness clamp -- no - // ClampSplitKForBatch, which would second-guess the measured choice. - // - // The tree only ever violates split_k (every leaf fits the Conv scratch budget by - // construction), so clamp rather than discard, then re-check. On the swept - // envelope the clamp is a no-op; it only bites for a K smaller than anything - // swept. - ClampSplitKToProblem(*tuned, K, kConvScratchNPad); - if (IsTilingValid(*tuned, K, kConvScratchNPad)) { - ClampTileNToProblem(*tuned, N); - return *tuned; - } - } // ConvHeuristicTiling is fit to kConvHeuristicArch's sweep only; any other arch // (no Conv sweep data) keeps using the arch-agnostic occupancy heuristic. SubgroupMatrixTiling tiling = (arch == kConvHeuristicArch) @@ -412,10 +357,9 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( } // Conv tiling policy for the 8x16x16 subgroup-matrix Conv kernel: same tile // candidates as MatMul, but selection uses the Conv scratch budget (which also - // stages the im2col A tile) and the Conv-specific pretuned tree / heuristic. + // stages the im2col A tile) and the Conv-specific heuristic. return [](const ComputeContext& context, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) -> std::optional { - const std::string_view arch = std::string_view{context.AdapterInfo().architecture}; // K must align to the subgroup tiling, and N must hold at least one whole // subgroup-matrix N block: the Conv kernel loads B unchecked from the weight // buffer and shifts a trailing partial N tile back to N - tile_n, which is @@ -429,6 +373,7 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( if (K % kSubgroupMatrixK != 0 || N < kSubgroupMatrixN) { return std::nullopt; } + const std::string_view arch = std::string_view{context.AdapterInfo().architecture}; return SelectConvTiling(arch, M, N, K, batch); }; } 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 46566d8ee1a2e..9a010c79b70b2 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 @@ -21,10 +21,9 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( const ComputeContextBase& context); // Like CreateSubgroupMatrixTilingSelector, but for the subgroup-matrix Conv path. -// The tile candidates and autotuner hooks (override / arch capture) are shared, -// but selection uses the Conv workgroup-memory budget (which also stages the -// im2col A tile) and consults the Conv-specific pretuned data instead of the -// MatMul table. Returns an empty selector on non-Intel devices. +// The tile candidates are shared, but selection uses the Conv workgroup-memory +// budget (which also stages the im2col A tile) and the Conv-specific heuristic +// instead of the MatMul table. Returns an empty selector on non-Intel devices. SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( const ComputeContextBase& context); diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc b/onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc deleted file mode 100644 index bfc3a52c3e051..0000000000000 --- a/onnxruntime/core/providers/webgpu/vendor/intel/nn/subgroup_matrix_conv_tuned.inc +++ /dev/null @@ -1,1871 +0,0 @@ -// Auto-generated by tools/python/gen_sgmm_tree.py --mode conv. Do not edit. -// Distilled decision tree (depth 10, 460 leaves, 30 tilings) -// approximating the offline GBM cost model for arch: xe-3lpg. -// -// Arch-specific: dispatch on kArch. One .inc per arch. -// -// Every leaf satisfies the conv workgroup-memory budget -// (split_k * tile_m * (tile_n + 16) <= 16384), so a prediction can only be -// invalid through split_k: this is a performance oracle only, like HeuristicTiling in -// subgroup_matrix_tiling_selector.cc, and it does not see K's block count, so it can -// propose split_k > K / 16. The caller MUST clamp split-K to K's block count and -// re-check validity. -// -// The caller must NOT apply the batch-occupancy clamp (ClampSplitKForBatch) on top of -// a prediction from this tree: batch is a model feature here (f3), so the predicted -// split_k already prices occupancy in, and clamping again double-counts it. That clamp -// exists for the batch-blind pretuned table only. -// -// The including TU must provide , , and -// SubgroupMatrixTiling. - -namespace sgconv_tree::xe3lpg { - -// The device architecture string this tree was tuned on, as reported by the adapter. -inline constexpr std::string_view kArch = "xe-3lpg"; - -inline SubgroupMatrixTiling PredictTilingTree(uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - [[maybe_unused]] const float f0 = std::log2(static_cast(M < 1u ? 1u : M)); - [[maybe_unused]] const float f1 = std::log2(static_cast(N < 1u ? 1u : N)); - [[maybe_unused]] const float f2 = std::log2(static_cast(K < 1u ? 1u : K)); - [[maybe_unused]] const float f3 = std::log2(static_cast(batch < 1u ? 1u : batch)); - if (f0 <= 7.792481f) { - if (f1 <= 5.792481f) { - if (f0 <= 6.292481f) { - if (f3 <= 2.500000f) { - if (f1 <= 4.292481f) { - if (f3 <= 0.500000f) { - if (f2 <= 8.406890f) { - return {8, 16, 4}; - } else { - if (f2 <= 9.921175f) { - return {16, 16, 8}; - } else { - return {8, 16, 8}; - } - } - } else { - return {8, 16, 8}; - } - } else { - if (f0 <= 5.792481f) { - return {8, 16, 8}; - } else { - if (f3 <= 1.500000f) { - return {8, 16, 8}; - } else { - if (f1 <= 5.292481f) { - return {8, 16, 8}; - } else { - return {16, 16, 8}; - } - } - } - } - } else { - if (f0 <= 5.292481f) { - if (f2 <= 8.406890f) { - if (f1 <= 5.292481f) { - return {8, 16, 8}; - } else { - return {8, 16, 4}; - } - } else { - return {8, 16, 8}; - } - } else { - if (f1 <= 4.292481f) { - return {8, 16, 8}; - } else { - if (f2 <= 8.906890f) { - if (f1 <= 5.292481f) { - if (f1 <= 4.792481f) { - return {8, 16, 4}; - } else { - return {8, 32, 8}; - } - } else { - return {8, 16, 4}; - } - } else { - if (f1 <= 4.792481f) { - return {16, 16, 8}; - } else { - if (f1 <= 5.292481f) { - return {8, 32, 8}; - } else { - return {16, 16, 8}; - } - } - } - } - } - } - } else { - if (f3 <= 1.500000f) { - if (f2 <= 10.421175f) { - if (f2 <= 8.406890f) { - if (f0 <= 6.792481f) { - return {8, 16, 8}; - } else { - if (f2 <= 7.435183f) { - return {8, 16, 4}; - } else { - if (f1 <= 4.792481f) { - if (f0 <= 7.292481f) { - return {8, 16, 8}; - } else { - return {8, 16, 4}; - } - } else { - return {8, 16, 8}; - } - } - } - } else { - if (f1 <= 5.292481f) { - if (f1 <= 4.292481f) { - if (f3 <= 0.500000f) { - return {16, 16, 8}; - } else { - return {8, 16, 8}; - } - } else { - if (f0 <= 7.292481f) { - return {8, 16, 8}; - } else { - if (f3 <= 0.500000f) { - return {8, 16, 8}; - } else { - return {8, 32, 8}; - } - } - } - } else { - if (f3 <= 0.500000f) { - return {8, 16, 8}; - } else { - return {16, 16, 8}; - } - } - } - } else { - if (f0 <= 7.292481f) { - return {8, 16, 8}; - } else { - if (f3 <= 0.500000f) { - return {8, 16, 8}; - } else { - if (f1 <= 4.792481f) { - return {8, 16, 8}; - } else { - return {8, 32, 8}; - } - } - } - } - } else { - if (f2 <= 9.421175f) { - if (f1 <= 4.792481f) { - if (f3 <= 2.500000f) { - if (f1 <= 4.292481f) { - return {8, 16, 8}; - } else { - return {8, 16, 4}; - } - } else { - if (f1 <= 4.292481f) { - return {8, 16, 4}; - } else { - return {16, 16, 4}; - } - } - } else { - if (f1 <= 5.292481f) { - if (f3 <= 2.500000f) { - return {8, 32, 8}; - } else { - return {8, 32, 4}; - } - } else { - if (f3 <= 2.500000f) { - return {8, 16, 4}; - } else { - return {16, 16, 4}; - } - } - } - } else { - if (f1 <= 4.792481f) { - if (f3 <= 2.500000f) { - if (f1 <= 4.292481f) { - return {8, 16, 8}; - } else { - return {16, 16, 8}; - } - } else { - return {16, 16, 8}; - } - } else { - if (f1 <= 5.292481f) { - if (f3 <= 2.500000f) { - return {8, 32, 8}; - } else { - if (f2 <= 10.921175f) { - return {8, 32, 4}; - } else { - return {16, 32, 8}; - } - } - } else { - if (f0 <= 7.292481f) { - return {16, 16, 8}; - } else { - return {16, 32, 4}; - } - } - } - } - } - } - } else { - if (f0 <= 5.292481f) { - if (f3 <= 1.500000f) { - if (f2 <= 7.435183f) { - if (f3 <= 0.500000f) { - return {8, 16, 4}; - } else { - if (f1 <= 7.292481f) { - return {8, 16, 8}; - } else { - return {8, 16, 4}; - } - } - } else { - if (f1 <= 8.292481f) { - return {8, 32, 8}; - } else { - if (f3 <= 0.500000f) { - return {8, 32, 8}; - } else { - if (f0 <= 4.292481f) { - return {8, 32, 8}; - } else { - if (f2 <= 8.406890f) { - return {8, 32, 4}; - } else { - if (f2 <= 11.421175f) { - return {8, 64, 8}; - } else { - return {8, 32, 8}; - } - } - } - } - } - } - } else { - if (f1 <= 7.292481f) { - if (f2 <= 7.435183f) { - if (f3 <= 2.500000f) { - return {8, 32, 8}; - } else { - return {8, 16, 4}; - } - } else { - if (f3 <= 2.500000f) { - return {8, 32, 8}; - } else { - if (f1 <= 6.792481f) { - if (f1 <= 6.292481f) { - return {8, 32, 8}; - } else { - if (f2 <= 9.421175f) { - return {8, 32, 4}; - } else { - return {8, 32, 8}; - } - } - } else { - if (f0 <= 4.292481f) { - return {8, 32, 8}; - } else { - if (f2 <= 9.421175f) { - return {8, 32, 4}; - } else { - return {8, 64, 8}; - } - } - } - } - } - } else { - if (f2 <= 10.421175f) { - if (f3 <= 2.500000f) { - if (f2 <= 8.906890f) { - return {8, 32, 4}; - } else { - if (f1 <= 8.292481f) { - return {8, 64, 8}; - } else { - return {8, 32, 4}; - } - } - } else { - if (f1 <= 8.292481f) { - if (f0 <= 4.792481f) { - return {8, 32, 4}; - } else { - return {8, 64, 4}; - } - } else { - if (f0 <= 4.292481f) { - return {8, 32, 4}; - } else { - return {8, 64, 4}; - } - } - } - } else { - if (f1 <= 8.792481f) { - if (f2 <= 11.421175f) { - return {8, 64, 8}; - } else { - if (f3 <= 2.500000f) { - return {8, 32, 8}; - } else { - return {8, 64, 8}; - } - } - } else { - if (f3 <= 2.500000f) { - return {8, 64, 8}; - } else { - return {8, 64, 4}; - } - } - } - } - } - } else { - if (f1 <= 7.292481f) { - if (f3 <= 1.500000f) { - if (f1 <= 6.292481f) { - if (f0 <= 6.292481f) { - if (f3 <= 0.500000f) { - return {8, 16, 8}; - } else { - if (f0 <= 5.792481f) { - return {8, 16, 8}; - } else { - return {8, 32, 8}; - } - } - } else { - if (f3 <= 0.500000f) { - if (f0 <= 7.292481f) { - return {8, 16, 8}; - } else { - return {8, 32, 8}; - } - } else { - if (f0 <= 7.292481f) { - return {8, 32, 8}; - } else { - return {8, 64, 8}; - } - } - } - } else { - if (f0 <= 6.792481f) { - if (f3 <= 0.500000f) { - if (f0 <= 6.292481f) { - if (f1 <= 6.792481f) { - return {8, 16, 8}; - } else { - return {8, 32, 8}; - } - } else { - return {8, 32, 8}; - } - } else { - if (f0 <= 6.292481f) { - return {8, 32, 8}; - } else { - if (f2 <= 8.906890f) { - return {8, 32, 4}; - } else { - return {8, 32, 8}; - } - } - } - } else { - if (f3 <= 0.500000f) { - if (f2 <= 8.406890f) { - return {8, 32, 4}; - } else { - return {8, 32, 8}; - } - } else { - if (f2 <= 10.421175f) { - return {8, 32, 4}; - } else { - if (f1 <= 6.792481f) { - return {8, 32, 8}; - } else { - return {8, 64, 8}; - } - } - } - } - } - } else { - if (f2 <= 9.421175f) { - if (f0 <= 6.792481f) { - if (f3 <= 2.500000f) { - if (f1 <= 6.292481f) { - return {8, 32, 8}; - } else { - return {8, 32, 4}; - } - } else { - if (f1 <= 6.792481f) { - return {8, 32, 4}; - } else { - return {8, 64, 4}; - } - } - } else { - if (f3 <= 2.500000f) { - if (f1 <= 6.792481f) { - if (f0 <= 7.292481f) { - return {8, 32, 4}; - } else { - return {8, 32, 2}; - } - } else { - return {8, 64, 4}; - } - } else { - if (f1 <= 6.292481f) { - return {8, 64, 4}; - } else { - if (f1 <= 6.792481f) { - return {8, 32, 1}; - } else { - return {8, 64, 2}; - } - } - } - } - } else { - if (f1 <= 6.792481f) { - if (f1 <= 6.292481f) { - if (f0 <= 6.792481f) { - if (f3 <= 2.500000f) { - return {8, 32, 8}; - } else { - return {8, 64, 8}; - } - } else { - if (f3 <= 2.500000f) { - return {8, 64, 8}; - } else { - return {8, 64, 4}; - } - } - } else { - if (f2 <= 11.421175f) { - if (f3 <= 2.500000f) { - return {16, 32, 8}; - } else { - return {8, 32, 4}; - } - } else { - if (f3 <= 2.500000f) { - return {8, 32, 8}; - } else { - return {16, 32, 8}; - } - } - } - } else { - if (f0 <= 5.792481f) { - return {8, 32, 8}; - } else { - if (f0 <= 6.792481f) { - if (f3 <= 2.500000f) { - return {8, 64, 8}; - } else { - return {8, 64, 4}; - } - } else { - if (f3 <= 2.500000f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } - } - } - } - } - } else { - if (f3 <= 0.500000f) { - if (f0 <= 6.292481f) { - if (f2 <= 8.906890f) { - if (f1 <= 8.292481f) { - if (f2 <= 7.935182f) { - return {8, 32, 4}; - } else { - return {8, 32, 8}; - } - } else { - return {8, 32, 4}; - } - } else { - return {8, 32, 8}; - } - } else { - if (f2 <= 10.421175f) { - if (f1 <= 8.792481f) { - if (f0 <= 7.292481f) { - return {8, 32, 4}; - } else { - if (f2 <= 8.406890f) { - return {8, 32, 2}; - } else { - return {8, 32, 4}; - } - } - } else { - return {8, 64, 4}; - } - } else { - if (f1 <= 8.792481f) { - if (f1 <= 7.792481f) { - return {8, 32, 8}; - } else { - if (f0 <= 6.792481f) { - return {8, 32, 8}; - } else { - return {8, 64, 8}; - } - } - } else { - return {8, 64, 4}; - } - } - } - } else { - if (f2 <= 9.421175f) { - if (f3 <= 1.500000f) { - if (f0 <= 6.792481f) { - if (f1 <= 8.792481f) { - return {8, 32, 4}; - } else { - return {8, 64, 4}; - } - } else { - if (f1 <= 8.792481f) { - if (f0 <= 7.292481f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } else { - return {8, 64, 2}; - } - } - } else { - if (f2 <= 7.435183f) { - if (f0 <= 7.292481f) { - if (f0 <= 6.292481f) { - return {8, 32, 1}; - } else { - return {8, 64, 1}; - } - } else { - return {8, 64, 1}; - } - } else { - if (f0 <= 6.792481f) { - if (f1 <= 8.792481f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } else { - if (f2 <= 8.406890f) { - return {8, 64, 1}; - } else { - return {8, 64, 2}; - } - } - } - } - } else { - if (f2 <= 11.421175f) { - if (f3 <= 1.500000f) { - if (f0 <= 6.792481f) { - if (f2 <= 10.421175f) { - return {8, 32, 4}; - } else { - return {8, 64, 8}; - } - } else { - return {8, 64, 4}; - } - } else { - return {8, 64, 4}; - } - } else { - if (f3 <= 2.500000f) { - if (f1 <= 7.792481f) { - return {8, 32, 8}; - } else { - if (f0 <= 5.792481f) { - return {8, 64, 8}; - } else { - return {8, 64, 4}; - } - } - } else { - if (f0 <= 6.292481f) { - if (f1 <= 8.292481f) { - return {8, 64, 4}; - } else { - return {16, 32, 4}; - } - } else { - if (f1 <= 8.292481f) { - return {16, 32, 4}; - } else { - return {16, 64, 4}; - } - } - } - } - } - } - } - } - } - } else { - if (f1 <= 6.792481f) { - if (f1 <= 4.792481f) { - if (f2 <= 8.906890f) { - if (f0 <= 12.292481f) { - if (f2 <= 7.435183f) { - if (f0 <= 9.292481f) { - if (f3 <= 1.500000f) { - return {8, 16, 4}; - } else { - return {8, 16, 1}; - } - } else { - return {8, 16, 1}; - } - } else { - if (f0 <= 10.792481f) { - if (f3 <= 2.500000f) { - if (f0 <= 9.792481f) { - if (f2 <= 8.406890f) { - return {8, 16, 4}; - } else { - return {16, 16, 8}; - } - } else { - if (f3 <= 0.500000f) { - return {8, 16, 4}; - } else { - return {8, 16, 1}; - } - } - } else { - if (f1 <= 4.292481f) { - return {16, 16, 4}; - } else { - return {16, 16, 2}; - } - } - } else { - if (f3 <= 2.500000f) { - if (f1 <= 4.292481f) { - if (f2 <= 8.406890f) { - return {8, 16, 2}; - } else { - return {16, 16, 4}; - } - } else { - return {16, 16, 2}; - } - } else { - if (f1 <= 4.292481f) { - return {16, 16, 2}; - } else { - return {32, 16, 1}; - } - } - } - } - } else { - if (f2 <= 7.435183f) { - if (f0 <= 13.292481f) { - if (f3 <= 1.500000f) { - return {8, 16, 1}; - } else { - return {16, 16, 1}; - } - } else { - return {16, 16, 1}; - } - } else { - if (f0 <= 14.292481f) { - if (f3 <= 1.500000f) { - if (f2 <= 8.406890f) { - if (f0 <= 12.792481f) { - return {32, 16, 2}; - } else { - return {16, 16, 2}; - } - } else { - if (f0 <= 13.292481f) { - return {16, 16, 4}; - } else { - return {16, 16, 2}; - } - } - } else { - if (f1 <= 4.292481f) { - return {32, 16, 2}; - } else { - if (f3 <= 2.500000f) { - return {32, 16, 2}; - } else { - return {16, 16, 1}; - } - } - } - } else { - if (f3 <= 2.500000f) { - if (f1 <= 4.292481f) { - if (f0 <= 15.792481f) { - return {32, 16, 2}; - } else { - return {16, 16, 1}; - } - } else { - return {16, 16, 1}; - } - } else { - if (f1 <= 4.292481f) { - return {16, 16, 1}; - } else { - return {32, 16, 1}; - } - } - } - } - } - } else { - if (f0 <= 11.292481f) { - if (f2 <= 10.921175f) { - if (f3 <= 1.500000f) { - if (f0 <= 9.792481f) { - if (f0 <= 8.292481f) { - return {8, 16, 8}; - } else { - return {16, 16, 8}; - } - } else { - if (f2 <= 9.421175f) { - return {16, 16, 4}; - } else { - if (f3 <= 0.500000f) { - return {16, 16, 8}; - } else { - return {16, 16, 4}; - } - } - } - } else { - if (f0 <= 9.292481f) { - if (f2 <= 9.421175f) { - return {16, 16, 4}; - } else { - return {16, 16, 8}; - } - } else { - if (f0 <= 10.792481f) { - return {16, 16, 4}; - } else { - if (f2 <= 9.921175f) { - return {32, 16, 4}; - } else { - return {16, 16, 4}; - } - } - } - } - } else { - if (f3 <= 0.500000f) { - if (f0 <= 8.792481f) { - return {8, 16, 8}; - } else { - return {16, 16, 8}; - } - } else { - if (f0 <= 9.792481f) { - if (f1 <= 4.292481f) { - if (f3 <= 1.500000f) { - return {8, 16, 8}; - } else { - return {16, 16, 8}; - } - } else { - if (f3 <= 2.500000f) { - return {16, 16, 8}; - } else { - return {32, 16, 8}; - } - } - } else { - if (f2 <= 11.421175f) { - if (f1 <= 4.292481f) { - return {16, 16, 8}; - } else { - return {32, 16, 8}; - } - } else { - return {16, 16, 8}; - } - } - } - } - } else { - if (f2 <= 10.921175f) { - if (f0 <= 14.292481f) { - if (f3 <= 2.500000f) { - return {16, 16, 4}; - } else { - if (f2 <= 10.421175f) { - if (f0 <= 13.292481f) { - return {32, 16, 2}; - } else { - return {16, 16, 2}; - } - } else { - return {16, 16, 2}; - } - } - } else { - if (f2 <= 9.921175f) { - if (f1 <= 4.292481f) { - if (f0 <= 14.792481f) { - return {16, 16, 4}; - } else { - return {32, 16, 2}; - } - } else { - if (f3 <= 1.500000f) { - return {32, 16, 2}; - } else { - return {16, 16, 1}; - } - } - } else { - if (f0 <= 14.792481f) { - if (f3 <= 1.500000f) { - return {32, 16, 4}; - } else { - return {16, 16, 2}; - } - } else { - return {16, 16, 2}; - } - } - } - } else { - if (f0 <= 12.792481f) { - if (f2 <= 11.421175f) { - if (f3 <= 1.500000f) { - return {16, 16, 4}; - } else { - return {32, 16, 8}; - } - } else { - if (f3 <= 2.500000f) { - if (f0 <= 12.292481f) { - return {16, 16, 8}; - } else { - return {32, 16, 8}; - } - } else { - return {32, 16, 8}; - } - } - } else { - if (f1 <= 4.292481f) { - if (f0 <= 14.792481f) { - if (f3 <= 1.500000f) { - return {16, 16, 8}; - } else { - return {16, 16, 4}; - } - } else { - if (f3 <= 2.500000f) { - return {16, 16, 4}; - } else { - return {16, 16, 2}; - } - } - } else { - if (f2 <= 11.421175f) { - if (f0 <= 14.292481f) { - return {16, 16, 4}; - } else { - return {32, 16, 4}; - } - } else { - return {16, 16, 4}; - } - } - } - } - } - } - } else { - if (f2 <= 9.421175f) { - if (f0 <= 10.292481f) { - if (f3 <= 1.500000f) { - if (f1 <= 5.792481f) { - if (f1 <= 5.292481f) { - if (f0 <= 9.292481f) { - if (f3 <= 0.500000f) { - return {8, 32, 8}; - } else { - return {8, 32, 4}; - } - } else { - if (f3 <= 0.500000f) { - return {8, 32, 4}; - } else { - return {8, 32, 2}; - } - } - } else { - if (f0 <= 9.292481f) { - if (f3 <= 0.500000f) { - return {8, 16, 4}; - } else { - return {16, 16, 4}; - } - } else { - if (f3 <= 0.500000f) { - return {8, 16, 2}; - } else { - return {16, 16, 2}; - } - } - } - } else { - if (f0 <= 8.792481f) { - if (f2 <= 7.435183f) { - return {8, 32, 2}; - } else { - return {8, 32, 4}; - } - } else { - if (f3 <= 0.500000f) { - if (f0 <= 9.292481f) { - return {8, 32, 4}; - } else { - return {8, 32, 2}; - } - } else { - if (f0 <= 9.292481f) { - return {8, 32, 2}; - } else { - return {8, 32, 1}; - } - } - } - } - } else { - if (f1 <= 5.292481f) { - if (f0 <= 9.292481f) { - if (f3 <= 2.500000f) { - return {8, 32, 4}; - } else { - return {8, 32, 2}; - } - } else { - if (f2 <= 8.406890f) { - return {8, 32, 1}; - } else { - return {8, 32, 4}; - } - } - } else { - if (f1 <= 6.292481f) { - if (f1 <= 5.792481f) { - if (f2 <= 7.435183f) { - return {16, 32, 1}; - } else { - return {16, 16, 4}; - } - } else { - if (f2 <= 7.435183f) { - return {8, 64, 1}; - } else { - return {8, 64, 2}; - } - } - } else { - if (f3 <= 2.500000f) { - if (f2 <= 8.406890f) { - return {8, 32, 1}; - } else { - return {16, 32, 4}; - } - } else { - if (f0 <= 8.792481f) { - return {8, 32, 1}; - } else { - return {16, 32, 1}; - } - } - } - } - } - } else { - if (f1 <= 6.292481f) { - if (f1 <= 5.792481f) { - if (f1 <= 5.292481f) { - if (f0 <= 13.792481f) { - if (f3 <= 2.500000f) { - return {8, 32, 1}; - } else { - return {16, 32, 1}; - } - } else { - return {16, 32, 1}; - } - } else { - if (f0 <= 14.292481f) { - return {16, 32, 1}; - } else { - if (f2 <= 7.435183f) { - return {16, 32, 1}; - } else { - return {32, 32, 1}; - } - } - } - } else { - if (f2 <= 7.435183f) { - return {8, 64, 1}; - } else { - if (f3 <= 1.500000f) { - if (f0 <= 11.792481f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } else { - if (f0 <= 15.292481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } - } - } - } else { - if (f0 <= 12.292481f) { - if (f3 <= 0.500000f) { - if (f2 <= 8.406890f) { - return {8, 32, 1}; - } else { - return {16, 32, 2}; - } - } else { - if (f3 <= 2.500000f) { - if (f2 <= 8.406890f) { - return {8, 32, 1}; - } else { - return {16, 32, 1}; - } - } else { - return {16, 32, 1}; - } - } - } else { - return {16, 32, 1}; - } - } - } - } else { - if (f0 <= 11.292481f) { - if (f1 <= 6.292481f) { - if (f1 <= 5.792481f) { - if (f1 <= 5.292481f) { - if (f3 <= 1.500000f) { - if (f0 <= 8.792481f) { - return {8, 32, 8}; - } else { - return {16, 32, 8}; - } - } else { - if (f0 <= 10.292481f) { - return {16, 32, 8}; - } else { - return {16, 32, 4}; - } - } - } else { - if (f3 <= 1.500000f) { - if (f0 <= 9.292481f) { - return {16, 16, 8}; - } else { - return {16, 32, 4}; - } - } else { - if (f0 <= 8.792481f) { - return {16, 32, 4}; - } else { - return {32, 32, 2}; - } - } - } - } else { - if (f3 <= 2.500000f) { - if (f0 <= 8.792481f) { - if (f3 <= 1.500000f) { - return {8, 64, 8}; - } else { - return {8, 64, 4}; - } - } else { - if (f0 <= 10.792481f) { - return {8, 64, 4}; - } else { - return {16, 64, 2}; - } - } - } else { - if (f0 <= 8.792481f) { - return {8, 64, 2}; - } else { - if (f0 <= 10.292481f) { - return {16, 64, 2}; - } else { - return {8, 64, 2}; - } - } - } - } - } else { - if (f2 <= 10.921175f) { - if (f3 <= 1.500000f) { - if (f0 <= 9.292481f) { - return {16, 32, 8}; - } else { - return {16, 32, 4}; - } - } else { - if (f0 <= 10.292481f) { - return {16, 32, 4}; - } else { - return {16, 32, 2}; - } - } - } else { - if (f0 <= 10.292481f) { - return {16, 32, 8}; - } else { - if (f3 <= 1.500000f) { - return {16, 32, 8}; - } else { - return {16, 32, 4}; - } - } - } - } - } else { - if (f1 <= 5.292481f) { - if (f0 <= 14.792481f) { - if (f2 <= 10.421175f) { - if (f3 <= 1.500000f) { - if (f0 <= 13.292481f) { - return {16, 32, 4}; - } else { - return {16, 32, 2}; - } - } else { - if (f0 <= 12.292481f) { - return {16, 32, 4}; - } else { - return {16, 32, 2}; - } - } - } else { - if (f3 <= 2.500000f) { - return {16, 32, 4}; - } else { - if (f0 <= 13.792481f) { - return {16, 32, 4}; - } else { - return {16, 32, 2}; - } - } - } - } else { - if (f2 <= 11.421175f) { - return {16, 32, 2}; - } else { - if (f3 <= 1.500000f) { - return {16, 32, 4}; - } else { - return {16, 32, 2}; - } - } - } - } else { - if (f1 <= 6.292481f) { - if (f1 <= 5.792481f) { - if (f2 <= 10.421175f) { - if (f0 <= 13.292481f) { - return {16, 32, 1}; - } else { - return {32, 32, 1}; - } - } else { - return {32, 32, 2}; - } - } else { - if (f2 <= 10.421175f) { - if (f2 <= 9.921175f) { - return {8, 64, 1}; - } else { - return {8, 64, 2}; - } - } else { - return {16, 64, 2}; - } - } - } else { - if (f2 <= 11.421175f) { - if (f2 <= 9.921175f) { - if (f3 <= 1.500000f) { - return {16, 32, 2}; - } else { - return {16, 32, 1}; - } - } else { - if (f3 <= 2.500000f) { - return {16, 32, 2}; - } else { - return {16, 32, 1}; - } - } - } else { - if (f0 <= 13.292481f) { - if (f3 <= 0.500000f) { - return {16, 32, 8}; - } else { - return {16, 32, 4}; - } - } else { - if (f3 <= 1.500000f) { - return {16, 32, 4}; - } else { - return {16, 32, 2}; - } - } - } - } - } - } - } - } - } else { - if (f2 <= 9.921175f) { - if (f2 <= 8.906890f) { - if (f0 <= 10.792481f) { - if (f2 <= 7.435183f) { - if (f3 <= 2.500000f) { - if (f0 <= 9.292481f) { - return {8, 64, 1}; - } else { - if (f1 <= 8.292481f) { - return {8, 64, 1}; - } else { - if (f1 <= 8.792481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } - } - } else { - if (f1 <= 7.792481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } - } else { - if (f3 <= 2.500000f) { - if (f2 <= 8.406890f) { - if (f0 <= 9.792481f) { - return {8, 64, 2}; - } else { - return {8, 64, 1}; - } - } else { - if (f0 <= 8.792481f) { - if (f3 <= 0.500000f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } else { - if (f3 <= 1.500000f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } - } - } else { - if (f1 <= 8.292481f) { - if (f0 <= 8.292481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } else { - if (f0 <= 9.292481f) { - if (f1 <= 8.792481f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } else { - return {16, 32, 1}; - } - } - } - } - } else { - if (f0 <= 14.792481f) { - if (f2 <= 8.406890f) { - if (f3 <= 2.500000f) { - return {8, 64, 1}; - } else { - if (f0 <= 12.792481f) { - return {8, 64, 1}; - } else { - if (f2 <= 7.435183f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } - } - } else { - if (f3 <= 0.500000f) { - if (f0 <= 13.292481f) { - if (f1 <= 8.292481f) { - return {16, 64, 1}; - } else { - return {8, 64, 1}; - } - } else { - return {8, 64, 1}; - } - } else { - if (f0 <= 14.292481f) { - if (f0 <= 11.792481f) { - return {8, 64, 1}; - } else { - return {16, 32, 1}; - } - } else { - return {8, 64, 1}; - } - } - } - } else { - if (f2 <= 7.435183f) { - if (f3 <= 1.500000f) { - return {8, 64, 1}; - } else { - if (f1 <= 8.292481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } - } else { - if (f2 <= 8.406890f) { - if (f3 <= 0.500000f) { - if (f1 <= 8.292481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } else { - return {16, 64, 1}; - } - } else { - if (f1 <= 8.292481f) { - return {8, 64, 1}; - } else { - if (f0 <= 15.792481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } - } - } - } - } - } else { - if (f0 <= 12.792481f) { - if (f0 <= 8.792481f) { - if (f3 <= 2.500000f) { - if (f3 <= 1.500000f) { - if (f1 <= 7.792481f) { - if (f3 <= 0.500000f) { - return {8, 32, 4}; - } else { - return {8, 64, 4}; - } - } else { - if (f3 <= 0.500000f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } - } else { - return {8, 64, 2}; - } - } else { - if (f1 <= 7.792481f) { - return {16, 64, 1}; - } else { - return {8, 64, 2}; - } - } - } else { - if (f2 <= 9.421175f) { - if (f3 <= 2.500000f) { - return {8, 64, 2}; - } else { - if (f0 <= 10.292481f) { - return {8, 64, 2}; - } else { - if (f0 <= 11.292481f) { - return {8, 64, 1}; - } else { - return {16, 32, 1}; - } - } - } - } else { - if (f0 <= 9.792481f) { - if (f3 <= 0.500000f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } else { - if (f3 <= 2.500000f) { - return {8, 64, 2}; - } else { - if (f0 <= 11.292481f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } - } - } - } - } else { - if (f2 <= 9.421175f) { - if (f1 <= 8.292481f) { - if (f0 <= 14.792481f) { - return {8, 64, 1}; - } else { - if (f3 <= 2.500000f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } - } else { - if (f3 <= 1.500000f) { - if (f0 <= 14.792481f) { - return {8, 64, 1}; - } else { - return {16, 64, 1}; - } - } else { - return {16, 64, 1}; - } - } - } else { - if (f3 <= 1.500000f) { - if (f1 <= 8.292481f) { - if (f0 <= 14.792481f) { - if (f0 <= 13.792481f) { - return {8, 64, 2}; - } else { - return {8, 64, 1}; - } - } else { - return {8, 64, 2}; - } - } else { - if (f0 <= 15.292481f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } - } else { - if (f0 <= 13.292481f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } - } - } - } - } else { - if (f2 <= 11.421175f) { - if (f0 <= 9.792481f) { - if (f3 <= 1.500000f) { - if (f3 <= 0.500000f) { - if (f1 <= 7.792481f) { - if (f0 <= 8.792481f) { - return {8, 64, 8}; - } else { - return {8, 64, 4}; - } - } else { - return {8, 64, 4}; - } - } else { - if (f1 <= 7.292481f) { - return {8, 64, 4}; - } else { - if (f0 <= 8.792481f) { - return {8, 64, 4}; - } else { - if (f1 <= 8.292481f) { - return {8, 64, 2}; - } else { - return {16, 32, 4}; - } - } - } - } - } else { - if (f2 <= 10.421175f) { - if (f1 <= 8.792481f) { - if (f3 <= 2.500000f) { - if (f0 <= 8.792481f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } else { - if (f0 <= 8.792481f) { - return {16, 32, 4}; - } else { - return {8, 64, 2}; - } - } - } else { - return {8, 64, 2}; - } - } else { - if (f1 <= 7.792481f) { - if (f3 <= 2.500000f) { - if (f0 <= 8.792481f) { - return {8, 64, 2}; - } else { - return {8, 64, 4}; - } - } else { - if (f0 <= 8.792481f) { - return {8, 64, 4}; - } else { - return {16, 64, 2}; - } - } - } else { - if (f0 <= 8.792481f) { - return {16, 64, 2}; - } else { - if (f3 <= 2.500000f) { - return {16, 64, 2}; - } else { - return {16, 64, 4}; - } - } - } - } - } - } else { - if (f2 <= 10.421175f) { - if (f0 <= 14.792481f) { - if (f3 <= 2.500000f) { - if (f0 <= 11.792481f) { - if (f3 <= 0.500000f) { - return {8, 64, 4}; - } else { - return {8, 64, 2}; - } - } else { - return {8, 64, 2}; - } - } else { - if (f1 <= 7.792481f) { - if (f0 <= 12.792481f) { - return {16, 64, 1}; - } else { - return {8, 64, 2}; - } - } else { - if (f0 <= 11.792481f) { - return {16, 64, 1}; - } else { - return {8, 64, 2}; - } - } - } - } else { - if (f3 <= 0.500000f) { - return {8, 64, 2}; - } else { - if (f3 <= 2.500000f) { - if (f1 <= 7.792481f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } else { - return {16, 64, 1}; - } - } - } - } else { - if (f0 <= 14.792481f) { - if (f2 <= 10.921175f) { - if (f3 <= 2.500000f) { - return {16, 64, 2}; - } else { - if (f0 <= 13.292481f) { - return {32, 64, 4}; - } else { - return {8, 64, 2}; - } - } - } else { - if (f0 <= 11.292481f) { - if (f3 <= 2.500000f) { - return {16, 64, 2}; - } else { - return {32, 64, 4}; - } - } else { - if (f0 <= 12.292481f) { - return {32, 64, 4}; - } else { - return {16, 32, 2}; - } - } - } - } else { - if (f3 <= 1.500000f) { - if (f2 <= 10.921175f) { - return {16, 64, 1}; - } else { - if (f1 <= 7.792481f) { - return {32, 64, 4}; - } else { - return {16, 32, 2}; - } - } - } else { - if (f2 <= 10.921175f) { - return {8, 64, 2}; - } else { - if (f1 <= 7.792481f) { - return {16, 32, 2}; - } else { - return {16, 64, 1}; - } - } - } - } - } - } - } else { - if (f0 <= 10.792481f) { - if (f3 <= 1.500000f) { - if (f0 <= 8.792481f) { - if (f1 <= 7.292481f) { - return {8, 64, 4}; - } else { - if (f3 <= 0.500000f) { - if (f0 <= 8.292481f) { - return {8, 64, 4}; - } else { - return {8, 32, 8}; - } - } else { - if (f1 <= 8.292481f) { - return {8, 32, 8}; - } else { - return {8, 64, 2}; - } - } - } - } else { - if (f1 <= 7.792481f) { - if (f0 <= 9.792481f) { - return {8, 64, 4}; - } else { - if (f3 <= 0.500000f) { - return {8, 64, 2}; - } else { - return {16, 64, 2}; - } - } - } else { - if (f0 <= 9.792481f) { - if (f3 <= 0.500000f) { - return {8, 64, 4}; - } else { - return {16, 64, 2}; - } - } else { - return {16, 64, 2}; - } - } - } - } else { - if (f1 <= 7.292481f) { - if (f0 <= 8.792481f) { - return {8, 32, 8}; - } else { - return {16, 64, 2}; - } - } else { - if (f0 <= 8.792481f) { - if (f1 <= 7.792481f) { - return {8, 64, 2}; - } else { - if (f3 <= 2.500000f) { - return {16, 64, 2}; - } else { - return {16, 32, 4}; - } - } - } else { - if (f3 <= 2.500000f) { - if (f0 <= 9.792481f) { - return {16, 64, 2}; - } else { - return {16, 32, 4}; - } - } else { - if (f0 <= 9.292481f) { - return {16, 32, 4}; - } else { - return {32, 64, 4}; - } - } - } - } - } - } else { - if (f0 <= 14.292481f) { - if (f3 <= 2.500000f) { - if (f0 <= 12.792481f) { - return {16, 64, 2}; - } else { - if (f3 <= 1.500000f) { - if (f1 <= 7.792481f) { - return {16, 32, 4}; - } else { - return {32, 64, 4}; - } - } else { - if (f1 <= 7.792481f) { - return {32, 64, 4}; - } else { - return {16, 64, 2}; - } - } - } - } else { - if (f0 <= 11.792481f) { - if (f1 <= 7.792481f) { - return {32, 64, 4}; - } else { - return {16, 64, 2}; - } - } else { - return {16, 64, 2}; - } - } - } else { - if (f1 <= 8.292481f) { - if (f3 <= 0.500000f) { - if (f0 <= 15.292481f) { - return {32, 64, 4}; - } else { - return {16, 64, 2}; - } - } else { - return {16, 64, 2}; - } - } else { - if (f3 <= 2.500000f) { - if (f3 <= 1.500000f) { - return {16, 64, 2}; - } else { - if (f0 <= 15.292481f) { - return {16, 64, 2}; - } else { - return {8, 64, 2}; - } - } - } else { - if (f0 <= 15.292481f) { - return {8, 64, 2}; - } else { - return {16, 64, 1}; - } - } - } - } - } - } - } - } - } -} - -} // namespace sgconv_tree::xe3lpg diff --git a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc index 3fb4386cca4e0..ffe3c8671e26d 100644 --- a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc +++ b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc @@ -215,8 +215,8 @@ TEST(SubgroupMatrixConv, PartialTiles) { // 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 neither the candidate list nor the -// tuned tree produces on its own +// 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 From 57ed167d74e0bb1afaa7468221be5b3e7daed4ba Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Thu, 15 Oct 2026 16:57:31 +0800 Subject: [PATCH 04/14] Optimize some code --- onnxruntime/core/providers/webgpu/nn/conv.cc | 31 ++++-------- onnxruntime/core/providers/webgpu/nn/conv.h | 19 ++++---- .../webgpu/nn/subgroup_matrix_conv.cc | 40 +++++++--------- .../webgpu/nn/subgroup_matrix_conv.h | 16 +++---- ...subgroup_matrix_conv_8x16x16.wgsl.template | 5 +- .../math/subgroup_matrix_tiling_selector.cc | 47 ++++++++++--------- .../math/subgroup_matrix_tiling_selector.h | 10 +++- 7 files changed, 79 insertions(+), 89 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index 72b56447d51a1..92c028346fcdf 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -2,9 +2,6 @@ // Licensed under the MIT License. #include "core/providers/webgpu/nn/conv.h" - -#include - #include "core/providers/webgpu/nn/conv2d_mm.h" #include "core/providers/webgpu/nn/conv3d_naive.h" #include "core/providers/webgpu/nn/im2col_matmul.h" @@ -43,11 +40,11 @@ Conv::ConvOptImplCache::GetOrCreate( template Status Conv::ComputeInternal(ComputeContext& context) const { - // Preferred fast path on devices with the subgroup-matrix config: implicit-GEMM - // Conv (fused im2col) via subgroupMatrixMultiplyAccumulate. Created once on first - // use (it needs a device query), then it decides inside Compute whether the - // kernel/attrs qualify. It declines (handled=false, allocating nothing) otherwise - // and we fall through to the normal path order below. + // Preferred fast path on devices with a vendor subgroup-matrix tiling policy (see + // subgroup_matrix_conv.h). Created once on first use (it needs a device query), then + // it decides inside Compute whether the kernel/attrs qualify. It declines + // (handled=false, allocating nothing) otherwise and we fall through to the normal + // path order below. 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)); @@ -360,19 +357,11 @@ Status Conv::PrePackInternal(ComputeContextBase& con // first, SubgroupMatrixConvImpl::Compute would see a non-null prepacked_kernel_ and // decline on every constant-weight Conv -- i.e. the path would silently never run. // - // Placed before the auto_pad check below for the same reason the im2col branch is: - // CanApplySubgroupMatrixConv() only looks at the adapter, dtype, layout, group and - // the weight shape -- never at pads -- so its answer cannot change at runtime, and - // ComputeInternal tests it before every pads-dependent branch. - // - // Skipping the prepack is always safe: is_packed stays false, so ORT keeps input 1 - // in OIHW and every fallback path can still transpose it at runtime if the - // subgroup-matrix path ends up declining. Once CanApplySubgroupMatrixConv has said - // yes, the only things Compute can still decline on are the ones the weight alone - // cannot show: a degenerate empty batch/M/N, or fewer than two spatial entries in - // dilations/pads/strides. Its other two decline paths are unreachable from here -- - // the vendor selector's K and N conditions are exactly the ones checked above, and - // its tile_n > N guard cannot fire because the selector clamps tile_n to N. + // Placed before the auto_pad check below for the same pads-independence reason as + // the im2col branch (see the comment there). Skipping the prepack here is always + // safe even if Compute() later declines for another reason: is_packed stays false, + // so input 1 stays in OIHW and every fallback path can still transpose it at + // runtime. if (CanApplySubgroupMatrixConv(context, is_channels_last, kernel_shape, onnxruntime::narrow(conv_attrs_.group), tensor.DataType())) { return Status::OK(); diff --git a/onnxruntime/core/providers/webgpu/nn/conv.h b/onnxruntime/core/providers/webgpu/nn/conv.h index 2e3bdac43bc41..600fc27ae7adf 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -32,7 +32,7 @@ class Conv : public WebGpuKernel { public: // Abstract base class for alternative optimized Conv implementations (currently the // subgroup-matrix path). Implementations hold the parent Conv and read - // conv_attrs/activation/... from it, mirroring MatMul::MatMulOptImpl. + // conv_attrs/activation/... from it, mirroring Gemm::GemmOptImpl. class ConvOptImpl { public: explicit ConvOptImpl(const Conv& parent) : parent_(parent) {} @@ -91,20 +91,19 @@ class Conv : public WebGpuKernel { bool w_is_constant_ = false; // whether the weight input (index 1) is a constant initializer private: - // Owns the alternative optimized implementation and builds it on first use, - // mirroring MatMulOptImplCache. The impl cannot be built in the ctor because it - // needs the device capabilities, which are only reachable through a compute - // context; and the same Conv kernel can see concurrent Compute calls, so creation - // is guarded by call_once. A null return after initialization means this device - // has no optimized path. Nested rather than a free class like MatMulOptImplCache - // only because ConvOptImpl is itself nested in this template. + // Owns the alternative optimized implementation and builds it lazily: the impl + // needs device capabilities only reachable through a compute context (not the + // ctor), and the same Conv kernel can see concurrent Compute calls, so creation is + // guarded by call_once -- mirroring MatMulComputeDispatcher, except nested here + // because ConvOptImpl is itself nested in this template. A null impl after + // initialization means this device has no optimized path. class ConvOptImplCache { public: ConvOptImplCache() = default; ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ConvOptImplCache); - // Defined in conv.cc: the factory it calls is declared in subgroup_matrix_conv.h, - // which includes this header. + // Defined in conv.cc, not here: 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: diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index 3bed70902d097..bbe9501876b85 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -1,8 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/providers/webgpu/nn/subgroup_matrix_conv.h" - #include #include #include @@ -17,6 +15,7 @@ #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/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" @@ -33,9 +32,9 @@ constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; // Combined workgroup-memory (SLM) budget in f16 elements for the two scratch // buffers the conv kernel uses: the split-K output tiles // (split_k * tile_m * tile_n) plus the per-subgroup im2col A staging -// (split_k * tile_m * sg_mat_k). 16384 f16 == 32 KB, the common WebGPU -// maxComputeWorkgroupStorageSize. split_k is clamped to fit this. -constexpr uint32_t kMaxConvScratchElems = 16384; +// (split_k * tile_m * sg_mat_k). split_k is clamped to fit this. Shares +// intel::kMaxScratchElems rather than hardcoding a second copy of the same +// budget the selector already sizes tilings against. // Largest per-tile subgroup-matrix counts the WGSL template unrolls: it generates the // accumulators sg_mat_c_ for mi < 8 and ni < 4. A tiling above these would be @@ -246,18 +245,16 @@ class SubgroupMatrixConvImpl final : public Conv::Co // // TODO: prepack the constant case, which would drop the cache and the first-Run // dispatch too. The layout is not the obstacle -- KernelLayout::OHWI is exactly - // what PrePackIm2ColMatMulWeight already produces, with the same {0, 2, 3, 1} - // perm. What blocks it is the fallback: PrePackInternal has to commit from the - // weight alone, while Compute can still decline at runtime (empty batch/M/N, or - // fewer than two spatial entries in dilations/pads/strides), and an OHWI weight is - // only readable by the im2col path that ComputeInternal ORT_ENFORCEs on. Since - // CanApplyIm2ColMatMulProgram is not implied by CanApplySubgroupMatrixConv (it - // gates on an arch-string whitelist and a fixed list of activation kinds, this one - // on the reported subgroup-matrix config and all of fuse_utils), a decline could - // land on a path that cannot read OHWI. Prepacking therefore means dropping the - // early return in Conv::PrePackInternal so the im2col branch below it prepacks - // when *both* predicates hold, and accepting a prepacked OHWI weight here instead - // of declining on it above. + // what PrePackIm2ColMatMulWeight already produces, with the same {0, 2, 3, 1} perm. + // The blocker is the fallback: PrePackInternal must commit from the weight alone, + // while Compute can still decline at runtime (empty batch/M/N, or fewer than two + // spatial entries in dilations/pads/strides) onto a path that only reads OHWI via + // im2col -- and CanApplyIm2ColMatMulProgram is not implied by + // CanApplySubgroupMatrixConv (different gates: arch/activation whitelist vs + // subgroup-matrix config/fuse_utils), so a decline could land where OHWI can't be + // read. Prepacking therefore means dropping PrePackInternal's early return so the + // im2col branch prepacks when *both* predicates hold, and accepting a prepacked + // OHWI weight here instead of declining on it. Tensor run_local_ohwi_weight; const Tensor* ohwi_weight = nullptr; if (parent.IsWeightConstant()) { @@ -289,7 +286,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co // 64 * (64 + 16) = 5120 elements at split_k 1, well inside the budget. uint32_t split_k = tiling->split_k; while (split_k > 1 && - split_k * tile_m * (tile_n + config.K) > kMaxConvScratchElems) { + split_k * tile_m * (tile_n + config.K) > intel::kMaxScratchElems) { split_k /= 2; } @@ -386,11 +383,8 @@ Status SubgroupMatrixConvProgram::GenerateShaderCode(ShaderHelper& shader) const "output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); // Write-out epilogue the template calls once per output element: bias add, then the - // fused activation. The activation body comes from fuse_utils - // (GetActivationDeclaration emits whatever module-scope helper the snippet needs, - // e.g. fused_act_erf), so every ActivationKind those cover is supported without the - // template having to enumerate them -- and any kind added there is picked up here - // for free. Mirrors SubgroupMatrixMatMulProgram::GenerateShaderCode. + // fused activation (see the class doc in subgroup_matrix_conv.h for why this needs no + // per-ActivationKind handling). Mirrors SubgroupMatrixMatMulProgram::GenerateShaderCode. auto& additional_implementation = shader.AdditionalImplementation(); additional_implementation << GetActivationDeclaration(activation_, "output_value_t", "output_element_t") diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h index cf4dc9b388999..e04b89e8b7148 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -91,18 +91,16 @@ class SubgroupMatrixConvProgram final : public Program #include #include @@ -15,6 +13,7 @@ #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. @@ -35,14 +34,19 @@ 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 +// Declared in subgroup_matrix_tiling_selector.h (as intel::kMaxScratchElems) so the +// Conv kernel's own split-K safety clamp can reuse the same value instead of +// hardcoding a second copy that has to be kept in sync by hand. // Extra per-tile N a kernel stages in workgroup memory beyond the split-K output // tiles. The MatMul kernel stages nothing else (0); the Conv kernel also holds the // im2col A tile (tile_m x sg_mat_k), so its scratch bound is // split_k * tile_m * (tile_n + sg_mat_k) -- up to 2x stricter than MatMul's. This is -// the only difference between the two kernels' tiling constraints, so every helper -// below takes it as a parameter rather than being duplicated per kernel. +// the only difference in the scratch-budget math shared by ScratchElems, +// IsTilingValid, HeuristicTiling and ClampSplitKToProblem below, so those helpers +// take it as a parameter instead of being duplicated per kernel. (The two kernels' +// tiling constraints differ elsewhere too -- e.g. ClampTileNToProblem and the +// pretuned-table-vs-heuristic-tree split -- those are handled separately.) constexpr uint32_t kMatMulScratchNPad = 0; constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK; @@ -142,17 +146,16 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t return {tile_m, tile_n, split_k}; } -// Arch the sweep behind ConvHeuristicTiling (below) was measured on. -constexpr std::string_view kConvHeuristicArch = "xe-3lpg"; +// Arch the sweep behind ConvTunedTiling (below) was measured on. +constexpr std::string_view kConvTunedArch = "xe-3lpg"; -// Same decisions as ConvHeuristicTilingV3, verified against all 16,535 shapes to -// produce identical predictions (0 mismatches): every M/N/K threshold slid to the -// nearest multiple of 16 (or 8 where no multiple of 16 fits) that stays strictly -// between the same two real data values the original threshold separated, so no -// shape crosses leaves. batch's threshold (6) has no such multiple in its gap -// (4, 8) -- rounding it to 8 would pull real batch=8 shapes into the wrong leaf -- -// so it is left as-is. Used by SelectConvTiling for kConvHeuristicArch. -SubgroupMatrixTiling ConvHeuristicTiling(uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { +// Decision tree fit to a shape sweep on kConvTunedArch, with every M/N/K +// threshold slid to the nearest multiple of 16 (or 8 where no multiple of 16 fits) +// that stays strictly between the same two real data values the original threshold +// separated, so no shape crosses leaves. batch's threshold (6) has no such multiple +// in its gap (4, 8) -- rounding it to 8 would pull real batch=8 shapes into the +// wrong leaf -- so it is left as-is. Used by SelectConvTiling for kConvTunedArch. +SubgroupMatrixTiling ConvTunedTiling(uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { if (K <= 2016u) { if (N <= 112u) { if (N <= 24u) { @@ -213,7 +216,7 @@ SubgroupMatrixTiling ConvHeuristicTiling(uint32_t M, uint32_t N, uint32_t K, uin } // Repairs a predicted tiling to this problem by halving split_k until it fits both -// K's block count and the scratch budget. ConvHeuristicTiling is a performance oracle +// K's block count and the scratch budget. ConvTunedTiling is a performance oracle // that never sees K's block count (it is trained on the problem's log dims only), so it // can propose split_k > K / kSubgroupMatrixK on a problem outside its swept envelope. // Clamping keeps the predicted tile shape -- the part that carries most of the win -- @@ -231,7 +234,7 @@ void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_ // from the weight buffer, and subgroupMatrixLoad is not bounds-checked, so a trailing // partial N tile is shifted back to N - tile_n. That shift only exists when // tile_n <= N. Neither source guarantees it -- HeuristicTiling only skips candidates -// larger than N for all but the smallest one, and ConvHeuristicTiling is an +// larger than N for all but the smallest one, and ConvTunedTiling is an // unconstrained oracle -- so the cap is applied to both. Shrinking tile_n keeps it a multiple of // kSubgroupMatrixN and only lowers ScratchElems, so it cannot invalidate a tiling // that was already valid. @@ -285,7 +288,7 @@ std::optional LookupPretunedTiling(std::string_view arch, // cooperative subgroups are redundant. Retire split-K factors whose batch-scaled // occupancy would exceed ~2x hardware. // -// Only for the batch-blind MatMul pretuned table: SelectConvTiling's ConvHeuristicTiling +// Only for the batch-blind MatMul pretuned table: SelectConvTiling's ConvTunedTiling // already takes batch into account when picking split-K, so applying this on top of it // would double-count occupancy. void ClampSplitKForBatch(SubgroupMatrixTiling& t, uint32_t M, uint32_t N, uint32_t batch, uint32_t hw) { @@ -313,15 +316,15 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, return HeuristicTiling(arch, M, N, K, batch, kMatMulScratchNPad); } -// Conv counterpart of SelectTiling: uses the Conv heuristic (kConvHeuristicArch) or the +// Conv counterpart of SelectTiling: uses the Conv-tuned tree (kConvTunedArch) or the // arch-agnostic occupancy heuristic, both under the Conv scratch budget. There is no // MatMul-table lookup here: the Conv kernel's cost profile (the im2col gather) differs, // and its scratch budget is stricter, so it must not reuse the MatMul-tuned entries. SubgroupMatrixTiling SelectConvTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - // ConvHeuristicTiling is fit to kConvHeuristicArch's sweep only; any other arch + // ConvTunedTiling is fit to kConvTunedArch's sweep only; any other arch // (no Conv sweep data) keeps using the arch-agnostic occupancy heuristic. - SubgroupMatrixTiling tiling = (arch == kConvHeuristicArch) - ? ConvHeuristicTiling(M, N, K, batch) + SubgroupMatrixTiling tiling = (arch == kConvTunedArch) + ? ConvTunedTiling(M, N, K, batch) : HeuristicTiling(arch, M, N, K, batch, kConvScratchNPad); ClampSplitKToProblem(tiling, K, kConvScratchNPad); ClampTileNToProblem(tiling, N); 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 9a010c79b70b2..8af19111c8bbe 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,6 +12,12 @@ namespace onnxruntime { namespace webgpu { namespace intel { +// Workgroup-memory (SLM) budget, in f16 elements, that the selector below fits +// tilings into. 16384 f16 == 32 KB, the common WebGPU maxComputeWorkgroupStorageSize. +// Shared with (not duplicated by) the Conv kernel's own split-K safety clamp, since +// both bound the same physical scratch allocation the selector sizes. +constexpr uint32_t kMaxScratchElems = 16384; + // 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 @@ -22,8 +28,8 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( // Like CreateSubgroupMatrixTilingSelector, but for the subgroup-matrix Conv path. // The tile candidates are shared, but selection uses the Conv workgroup-memory -// budget (which also stages the im2col A tile) and the Conv-specific heuristic -// instead of the MatMul table. Returns an empty selector on non-Intel devices. +// budget (which also stages the im2col A tile) and the Conv-tuned tree instead +// of the MatMul table. Returns an empty selector on non-Intel devices. SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( const ComputeContextBase& context); From 9c98929cd355545729857b56b6fed9440b4d2631 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Mon, 19 Oct 2026 14:52:41 +0800 Subject: [PATCH 05/14] Reject weights exceeding maxStorageBufferBindingSize --- .../webgpu/nn/subgroup_matrix_conv.cc | 18 +++ .../math/subgroup_matrix_tiling_selector.cc | 133 ++++++++++++++---- 2 files changed, 121 insertions(+), 30 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index bbe9501876b85..b355a8ca38533 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -472,6 +472,24 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, return false; } + // The weight buffer is bound as a single WGSL binding (input_b), and the shader's + // subgroupMatrixLoad calls address it with offsets derived from the full [N, K] + // weight -- they know nothing about ProgramManager's binding-segment split, which + // kicks in once a tensor exceeds maxStorageBufferBindingSize. A weight past that + // limit would still pass every check above, get split into multiple bindings, and + // then have its later output channels load past the first binding. Reject it here, + // before PrePackInternal or any output allocation/transpose, so it falls back to + // the normal (prepacked) conv paths instead. + const uint64_t max_weight_elems = context.DeviceLimits().maxStorageBufferBindingSize / sizeof(MLFloat16); + if (static_cast(kernel_shape.Size()) > max_weight_elems) { + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: weight exceeds maxStorageBufferBindingSize" + << " (kernel_shape[OIHW]=" << kernel_shape.ToString() << ", elems=" + << kernel_shape.Size() << ", limit_elems=" << max_weight_elems + << ", maxStorageBufferBindingSize=" << context.DeviceLimits().maxStorageBufferBindingSize + << " bytes)."; + return false; + } + const std::optional config = GetSupportedConfig(context); if (!config) { LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: device does not report the required 8x16x16 F16" 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 17326356aa164..c030d2f97963d 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,15 +1,11 @@ // 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_tiling.h" #include "core/providers/webgpu/vendor/intel/intel_device_info.h" @@ -43,10 +39,10 @@ constexpr uint32_t kSplitKCandidates[] = {1, 2, 4, 8}; // im2col A tile (tile_m x sg_mat_k), so its scratch bound is // split_k * tile_m * (tile_n + sg_mat_k) -- up to 2x stricter than MatMul's. This is // the only difference in the scratch-budget math shared by ScratchElems, -// IsTilingValid, HeuristicTiling and ClampSplitKToProblem below, so those helpers -// take it as a parameter instead of being duplicated per kernel. (The two kernels' -// tiling constraints differ elsewhere too -- e.g. ClampTileNToProblem and the -// pretuned-table-vs-heuristic-tree split -- those are handled separately.) +// IsTilingValid and ClampSplitKToProblem below, so those helpers take it as a +// parameter instead of being duplicated per kernel. +// (The two kernels' tiling constraints differ elsewhere too -- e.g. ClampTileNToProblem +// and the pretuned-table-vs-heuristic-tree split -- those are handled separately.) constexpr uint32_t kMatMulScratchNPad = 0; constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK; @@ -76,17 +72,24 @@ uint32_t EffectiveHwSubgroups(std::string_view 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 +// Fallback MatMul 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. -SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch, - uint32_t scratch_n_pad) { +// 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. +// +// Deliberately not shared with ConvHeuristicTiling below, even where the two look +// similar: shrinking tile_n to manufacture more occupancy tiles is free here (for +// MatMul, tile_n only trades off against tile_m for a fixed amount of +// global-memory traffic) but not for Conv -- see that function's comment -- so the +// two kernels' tile and split-K policies are expected to keep diverging and tend +// to get tuned/debugged independently. A shared helper would just mean every +// future change has to re-check it does not perturb the other kernel. +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; }; @@ -139,13 +142,82 @@ 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 && - ScratchElems({tile_m, tile_n, cand}, scratch_n_pad) <= kMaxScratchElems) { + ScratchElems({tile_m, tile_n, cand}, kMatMulScratchNPad) <= kMaxScratchElems) { split_k = cand; } } return {tile_m, tile_n, split_k}; } +// Once N needs at least two full-width (kTileNCandidates' max) N tiles, ConvTunedTiling's +// own sweep (see its N > 112 branches) picks split-K from K alone, completely +// ignoring M -- e.g. {8, 64, 8} at M = 64 and at M = 4096 alike. An occupancy-based +// search cannot reproduce that for any fill target: a large M already +// saturates occupancy at split_k = 1 on its own, yet the tuned table still splits, +// because the win there is not filling idle subgroups but shortening each +// workgroup's serial K-reduction chain (gather-latency hiding), which does not care +// how many other workgroups are already in flight. kConvWideN and +// kConvWideNSplitKBlocks mirror that branch's shape -- wide N, two K tiers -- as the +// closest available proxy for architectures with no sweep of their own; ClampSplitKToProblem +// (applied by SelectConvTiling after this returns) shrinks the pick back down if K or the +// scratch budget cannot actually support it. +constexpr uint32_t kConvWideN = 2 * kTileNCandidates[2]; // 128: needs 2+ tile_n=64 tiles +constexpr uint32_t kConvWideNSplitKBlocks = 2016 / kSubgroupMatrixK; // 126, ConvTunedTiling's own cut + +// Fallback Conv tiling on arches with no Conv-tuned tree (anything but +// kConvTunedArch). +// +// Does *not* pick tile_m/tile_n the way HeuristicTiling does -- shrinking tile_n to +// manufacture more independent occupancy tiles is free for MatMul (tile_n only +// trades off against tile_m for a fixed amount of global-memory traffic) but not +// for Conv: the im2col A tile (tile_m x K) is staged per-workgroup in workgroup +// memory, not shared, so it is re-gathered once per N tile that covers it. Total +// gather volume is M * N * K / tile_n -- inversely proportional to tile_n and +// independent of tile_m or split_k. Measured against ConvTunedTiling's real +// hardware sweep, HeuristicTiling's occupancy-driven tile choice was shrinking +// tile_n to 16-32 on shapes where the tuned table keeps it at 64, multiplying +// gather volume 2x-4x and dominating the kernel's cost regardless of how split-K +// was tuned -- which is why an earlier fix that only tightened split-K's fill +// target did not help. +// +// So tile_n is fixed at the largest candidate that fits N (minimizing gather +// volume) and tile_m at the smallest candidate (splitting M adds independent tiles +// without touching gather volume). split-K then takes one of two policies: below +// kConvWideN, an occupancy fill nearly identical in shape to HeuristicTiling's (not +// shared with it -- see that function's comment); at or above it, the +// M-independent K-tiered policy from kConvWideN's comment, since occupancy no +// longer predicts ConvTunedTiling's choice there. +SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { + 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]; + + uint32_t split_k = 1; + if (N >= kConvWideN) { + // Occupancy (hw) is deliberately not consulted here -- see kConvWideN's comment. + split_k = k_blocks > kConvWideNSplitKBlocks ? 8u : 2u; + } else { + const uint32_t hw = EffectiveHwSubgroups(arch); + auto tile_count = [](uint32_t dim, uint32_t tile) { return (dim + tile - 1) / tile; }; + const uint32_t num_tiles = batch * tile_count(M, tile_m) * tile_count(N, tile_n); + constexpr uint32_t kMinBlocksPerSplit = 2; + for (uint32_t cand : kSplitKCandidates) { + if (k_blocks >= cand * kMinBlocksPerSplit && + num_tiles * cand <= 2 * hw && + ScratchElems({tile_m, tile_n, cand}, kConvScratchNPad) <= kMaxScratchElems) { + split_k = cand; + } + } + } + return {tile_m, tile_n, split_k}; +} + // Arch the sweep behind ConvTunedTiling (below) was measured on. constexpr std::string_view kConvTunedArch = "xe-3lpg"; @@ -233,9 +305,10 @@ void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_ // Correctness, not performance: the Conv kernel loads its B right operand straight // from the weight buffer, and subgroupMatrixLoad is not bounds-checked, so a trailing // partial N tile is shifted back to N - tile_n. That shift only exists when -// tile_n <= N. Neither source guarantees it -- HeuristicTiling only skips candidates -// larger than N for all but the smallest one, and ConvTunedTiling is an -// unconstrained oracle -- so the cap is applied to both. Shrinking tile_n keeps it a multiple of +// tile_n <= N. Neither source guarantees it -- ConvHeuristicTiling already picks the +// largest candidate that fits N but N itself need not be a multiple of +// kSubgroupMatrixN, and ConvTunedTiling is an unconstrained oracle -- so the cap is +// applied to both. Shrinking tile_n keeps it a multiple of // kSubgroupMatrixN and only lowers ScratchElems, so it cannot invalidate a tiling // that was already valid. // @@ -313,19 +386,19 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, } return tiling; } - return HeuristicTiling(arch, M, N, K, batch, kMatMulScratchNPad); + return HeuristicTiling(arch, M, N, K, batch); } -// Conv counterpart of SelectTiling: uses the Conv-tuned tree (kConvTunedArch) or the -// arch-agnostic occupancy heuristic, both under the Conv scratch budget. There is no -// MatMul-table lookup here: the Conv kernel's cost profile (the im2col gather) differs, -// and its scratch budget is stricter, so it must not reuse the MatMul-tuned entries. +// Conv counterpart of SelectTiling: uses the Conv-tuned tree (kConvTunedArch) or +// ConvHeuristicTiling, both under the Conv scratch budget. There is no MatMul-table +// lookup here: the Conv kernel's cost profile (the im2col gather) differs, and its +// scratch budget is stricter, so it must not reuse the MatMul-tuned entries. SubgroupMatrixTiling SelectConvTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - // ConvTunedTiling is fit to kConvTunedArch's sweep only; any other arch - // (no Conv sweep data) keeps using the arch-agnostic occupancy heuristic. + // ConvTunedTiling is fit to kConvTunedArch's sweep only; any other arch (no Conv + // sweep data) falls back to ConvHeuristicTiling. SubgroupMatrixTiling tiling = (arch == kConvTunedArch) ? ConvTunedTiling(M, N, K, batch) - : HeuristicTiling(arch, M, N, K, batch, kConvScratchNPad); + : ConvHeuristicTiling(arch, M, N, K, batch); ClampSplitKToProblem(tiling, K, kConvScratchNPad); ClampTileNToProblem(tiling, N); return tiling; @@ -360,7 +433,7 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( } // Conv tiling policy for the 8x16x16 subgroup-matrix Conv kernel: same tile // candidates as MatMul, but selection uses the Conv scratch budget (which also - // stages the im2col A tile) and the Conv-specific heuristic. + // stages the im2col A tile) and SelectConvTiling's Conv-tuned tree / ConvHeuristicTiling. return [](const ComputeContext& context, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) -> std::optional { // K must align to the subgroup tiling, and N must hold at least one whole From 2b3fabbbb0c36ad31b2b05c29f79cd0bae05c622 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Mon, 19 Oct 2026 16:06:11 +0800 Subject: [PATCH 06/14] Remove conv tuned tree --- .../webgpu/nn/subgroup_matrix_conv.cc | 6 +- .../math/subgroup_matrix_tiling_selector.cc | 167 ++++++------------ .../math/subgroup_matrix_tiling_selector.h | 6 +- 3 files changed, 59 insertions(+), 120 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index b355a8ca38533..10e0774014eb3 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -282,8 +282,10 @@ class SubgroupMatrixConvImpl final : public Conv::Co // Safety net: the Conv selector already enforces the Conv scratch budget // (output tiles + im2col A staging), so this clamp normally does nothing; it // guards against a stray oversized split-K from the selector. Only split_k is - // clamped, not the tile area: the largest candidate tile needs - // 64 * (64 + 16) = 5120 elements at split_k 1, well inside the budget. + // clamped, not the tile area: ConvHeuristicTiling always picks the smallest + // tile_m (8), so the largest candidate tile needs 8 * (64 + 16) = 640 elements + // at split_k 1, or 8 * 8 * 80 = 5120 at the largest split_k (8), well inside + // the budget. uint32_t split_k = tiling->split_k; while (split_k > 1 && split_k * tile_m * (tile_n + config.K) > intel::kMaxScratchElems) { 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 c030d2f97963d..c550af604710e 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 @@ -41,8 +41,9 @@ constexpr uint32_t kSplitKCandidates[] = {1, 2, 4, 8}; // the only difference in the scratch-budget math shared by ScratchElems, // IsTilingValid and ClampSplitKToProblem below, so those helpers take it as a // parameter instead of being duplicated per kernel. -// (The two kernels' tiling constraints differ elsewhere too -- e.g. ClampTileNToProblem -// and the pretuned-table-vs-heuristic-tree split -- those are handled separately.) +// (The two kernels' tiling constraints differ elsewhere too -- e.g. ClampTileNToProblem, +// and MatMul alone has a pretuned table to consult before its heuristic -- those are +// handled separately.) constexpr uint32_t kMatMulScratchNPad = 0; constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK; @@ -149,23 +150,29 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t return {tile_m, tile_n, split_k}; } -// Once N needs at least two full-width (kTileNCandidates' max) N tiles, ConvTunedTiling's -// own sweep (see its N > 112 branches) picks split-K from K alone, completely -// ignoring M -- e.g. {8, 64, 8} at M = 64 and at M = 4096 alike. An occupancy-based -// search cannot reproduce that for any fill target: a large M already -// saturates occupancy at split_k = 1 on its own, yet the tuned table still splits, -// because the win there is not filling idle subgroups but shortening each -// workgroup's serial K-reduction chain (gather-latency hiding), which does not care -// how many other workgroups are already in flight. kConvWideN and -// kConvWideNSplitKBlocks mirror that branch's shape -- wide N, two K tiers -- as the -// closest available proxy for architectures with no sweep of their own; ClampSplitKToProblem -// (applied by SelectConvTiling after this returns) shrinks the pick back down if K or the -// scratch budget cannot actually support it. +// Once N needs at least two full-width (kTileNCandidates' max) N tiles, real +// hardware profiling (an xe-3lpg shape sweep) showed split-K should be picked from +// K alone, completely ignoring M -- e.g. 8-way splitting pays off equally at M = 64 +// and at M = 4096. An occupancy-based search cannot reproduce that for any fill +// target: a large M already saturates occupancy at split_k = 1 on its own, yet the +// profiled-good choice still splits, because the win there is not filling idle +// subgroups but shortening each workgroup's serial K-reduction chain +// (gather-latency hiding), which does not care how many other workgroups are +// already in flight. kConvWideN and kConvWideNSplitKBlocks mirror that one sweep's +// shape -- wide N, two K tiers -- as the best available estimate on hardware with +// no sweep of its own; ClampSplitKToProblem (applied by SelectConvTiling after this +// returns) shrinks the pick back down if K or the scratch budget cannot actually +// support it. constexpr uint32_t kConvWideN = 2 * kTileNCandidates[2]; // 128: needs 2+ tile_n=64 tiles -constexpr uint32_t kConvWideNSplitKBlocks = 2016 / kSubgroupMatrixK; // 126, ConvTunedTiling's own cut +constexpr uint32_t kConvWideNSplitKBlocks = 2016 / kSubgroupMatrixK; // 126 K-blocks -// Fallback Conv tiling on arches with no Conv-tuned tree (anything but -// kConvTunedArch). +// Conv tiling heuristic: the sole tiling strategy SelectConvTiling uses, on every +// architecture. (SelectConvTiling previously also consulted a separate, hardcoded +// decision tree fit to a shape sweep collected on one specific GPU architecture, +// used only on that architecture; a per-architecture sweep like that does not +// generalize to hardware it was never collected on, so that tree was dropped in +// favor of this heuristic alone, pending a more broadly-informed retuning in a +// follow-up.) // // Does *not* pick tile_m/tile_n the way HeuristicTiling does -- shrinking tile_n to // manufacture more independent occupancy tiles is free for MatMul (tile_n only @@ -173,20 +180,19 @@ constexpr uint32_t kConvWideNSplitKBlocks = 2016 / kSubgroupMatrixK; // 126, Co // for Conv: the im2col A tile (tile_m x K) is staged per-workgroup in workgroup // memory, not shared, so it is re-gathered once per N tile that covers it. Total // gather volume is M * N * K / tile_n -- inversely proportional to tile_n and -// independent of tile_m or split_k. Measured against ConvTunedTiling's real -// hardware sweep, HeuristicTiling's occupancy-driven tile choice was shrinking -// tile_n to 16-32 on shapes where the tuned table keeps it at 64, multiplying -// gather volume 2x-4x and dominating the kernel's cost regardless of how split-K -// was tuned -- which is why an earlier fix that only tightened split-K's fill -// target did not help. +// independent of tile_m or split_k. Measured against a real hardware sweep, +// HeuristicTiling's occupancy-driven tile choice was shrinking tile_n to 16-32 on +// shapes that profiled best at 64, multiplying gather volume 2x-4x and dominating +// the kernel's cost regardless of how split-K was tuned -- which is why an earlier +// fix that only tightened split-K's fill target did not help. // // So tile_n is fixed at the largest candidate that fits N (minimizing gather // volume) and tile_m at the smallest candidate (splitting M adds independent tiles // without touching gather volume). split-K then takes one of two policies: below // kConvWideN, an occupancy fill nearly identical in shape to HeuristicTiling's (not // shared with it -- see that function's comment); at or above it, the -// M-independent K-tiered policy from kConvWideN's comment, since occupancy no -// longer predicts ConvTunedTiling's choice there. +// M-independent K-tiered policy from kConvWideN's comment, since occupancy does not +// predict the profiled-good choice there. SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { const uint32_t k_blocks = K / kSubgroupMatrixK; @@ -218,81 +224,13 @@ SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint return {tile_m, tile_n, split_k}; } -// Arch the sweep behind ConvTunedTiling (below) was measured on. -constexpr std::string_view kConvTunedArch = "xe-3lpg"; - -// Decision tree fit to a shape sweep on kConvTunedArch, with every M/N/K -// threshold slid to the nearest multiple of 16 (or 8 where no multiple of 16 fits) -// that stays strictly between the same two real data values the original threshold -// separated, so no shape crosses leaves. batch's threshold (6) has no such multiple -// in its gap (4, 8) -- rounding it to 8 would pull real batch=8 shapes into the -// wrong leaf -- so it is left as-is. Used by SelectConvTiling for kConvTunedArch. -SubgroupMatrixTiling ConvTunedTiling(uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - if (K <= 2016u) { - if (N <= 112u) { - if (N <= 24u) { - if (M <= 224u) { - return {8, 16, 8}; - } else { - if (K <= 416u) { - return {16, 16, 2}; - } else { - if (batch <= 6u) { - if (M <= 896u) { - return {8, 16, 8}; - } else { - return {16, 16, 4}; - } - } else { - return {16, 16, 4}; - } - } - } - } else { - if (M <= 512u) { - if (batch <= 6u) { - return {8, 32, 8}; - } else { - return {8, 32, 4}; - } - } else { - if (K <= 512u) { - return {8, 32, 1}; - } else { - return {16, 32, 4}; - } - } - } - } else { - return {8, 64, 2}; - } - } else { - if (N <= 112u) { - if (N <= 24u) { - if (M <= 224u) { - return {8, 16, 8}; - } else { - return {16, 16, 8}; - } - } else { - if (M <= 512u) { - return {8, 32, 8}; - } else { - return {16, 32, 4}; - } - } - } else { - return {8, 64, 8}; - } - } -} - // Repairs a predicted tiling to this problem by halving split_k until it fits both -// K's block count and the scratch budget. ConvTunedTiling is a performance oracle -// that never sees K's block count (it is trained on the problem's log dims only), so it -// can propose split_k > K / kSubgroupMatrixK on a problem outside its swept envelope. -// Clamping keeps the predicted tile shape -- the part that carries most of the win -- -// instead of discarding the prediction and falling back to the arch-agnostic heuristic. +// K's block count and the scratch budget. ConvHeuristicTiling's wide-N branch picks +// split_k from a K-blocks threshold without checking whether K actually has enough +// blocks to divide that many ways (e.g. K with only 1 block still gets a tentative +// split_k of 2 there), so this clamp is relied on as the actual correctness +// guarantee, not just a safety net. Clamping keeps the predicted tile shape -- the +// part that carries most of the win -- instead of discarding the prediction. void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) { const uint32_t k_blocks = K / kSubgroupMatrixK; while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) { @@ -305,10 +243,9 @@ void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_ // Correctness, not performance: the Conv kernel loads its B right operand straight // from the weight buffer, and subgroupMatrixLoad is not bounds-checked, so a trailing // partial N tile is shifted back to N - tile_n. That shift only exists when -// tile_n <= N. Neither source guarantees it -- ConvHeuristicTiling already picks the -// largest candidate that fits N but N itself need not be a multiple of -// kSubgroupMatrixN, and ConvTunedTiling is an unconstrained oracle -- so the cap is -// applied to both. Shrinking tile_n keeps it a multiple of +// tile_n <= N. ConvHeuristicTiling already picks the largest candidate that fits N, +// but N itself need not be a multiple of kSubgroupMatrixN, so this cap is still +// needed. Shrinking tile_n keeps it a multiple of // kSubgroupMatrixN and only lowers ScratchElems, so it cannot invalidate a tiling // that was already valid. // @@ -361,9 +298,9 @@ std::optional LookupPretunedTiling(std::string_view arch, // cooperative subgroups are redundant. Retire split-K factors whose batch-scaled // occupancy would exceed ~2x hardware. // -// Only for the batch-blind MatMul pretuned table: SelectConvTiling's ConvTunedTiling -// already takes batch into account when picking split-K, so applying this on top of it -// would double-count occupancy. +// Only for the batch-blind MatMul pretuned table: SelectConvTiling's +// ConvHeuristicTiling already takes batch into account when picking split-K (see its +// narrow-N branch), so applying this on top of it would double-count occupancy. void ClampSplitKForBatch(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); @@ -389,16 +326,14 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, return HeuristicTiling(arch, M, N, K, batch); } -// Conv counterpart of SelectTiling: uses the Conv-tuned tree (kConvTunedArch) or -// ConvHeuristicTiling, both under the Conv scratch budget. There is no MatMul-table -// lookup here: the Conv kernel's cost profile (the im2col gather) differs, and its -// scratch budget is stricter, so it must not reuse the MatMul-tuned entries. +// Conv counterpart of SelectTiling: always goes through ConvHeuristicTiling, under +// the Conv scratch budget. There is no MatMul-table lookup here: the Conv kernel's +// cost profile (the im2col gather) differs, and its scratch budget is stricter, so +// it must not reuse the MatMul-tuned entries; and unlike SelectTiling, there is no +// pretuned Conv entry point at all to try first (see ConvHeuristicTiling's comment +// for why). SubgroupMatrixTiling SelectConvTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - // ConvTunedTiling is fit to kConvTunedArch's sweep only; any other arch (no Conv - // sweep data) falls back to ConvHeuristicTiling. - SubgroupMatrixTiling tiling = (arch == kConvTunedArch) - ? ConvTunedTiling(M, N, K, batch) - : ConvHeuristicTiling(arch, M, N, K, batch); + SubgroupMatrixTiling tiling = ConvHeuristicTiling(arch, M, N, K, batch); ClampSplitKToProblem(tiling, K, kConvScratchNPad); ClampTileNToProblem(tiling, N); return tiling; @@ -433,7 +368,7 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( } // Conv tiling policy for the 8x16x16 subgroup-matrix Conv kernel: same tile // candidates as MatMul, but selection uses the Conv scratch budget (which also - // stages the im2col A tile) and SelectConvTiling's Conv-tuned tree / ConvHeuristicTiling. + // stages the im2col A tile) and SelectConvTiling's ConvHeuristicTiling. return [](const ComputeContext& context, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) -> std::optional { // K must align to the subgroup tiling, and N must hold at least one whole 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 8af19111c8bbe..7447def75aa9f 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 @@ -28,8 +28,10 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixTilingSelector( // Like CreateSubgroupMatrixTilingSelector, but for the subgroup-matrix Conv path. // The tile candidates are shared, but selection uses the Conv workgroup-memory -// budget (which also stages the im2col A tile) and the Conv-tuned tree instead -// of the MatMul table. Returns an empty selector on non-Intel devices. +// budget (which also stages the im2col A tile) and a Conv-specific heuristic +// instead of the MatMul pretuned table. Conv has no pretuned table of its own: a +// per-architecture sweep does not generalize across hardware, so this always goes +// through the heuristic. Returns an empty selector on non-Intel devices. SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( const ComputeContextBase& context); From fd5af6637d18009d2eaaaf906160dbfdfd3114d4 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Tue, 20 Oct 2026 15:15:09 +0800 Subject: [PATCH 07/14] Remove kConvWideNSplitKBlocks --- .../math/subgroup_matrix_tiling_selector.cc | 105 ++++++++---------- 1 file changed, 46 insertions(+), 59 deletions(-) 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 c550af604710e..c97d3be7b947d 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 @@ -150,49 +150,37 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t return {tile_m, tile_n, split_k}; } -// Once N needs at least two full-width (kTileNCandidates' max) N tiles, real -// hardware profiling (an xe-3lpg shape sweep) showed split-K should be picked from -// K alone, completely ignoring M -- e.g. 8-way splitting pays off equally at M = 64 -// and at M = 4096. An occupancy-based search cannot reproduce that for any fill -// target: a large M already saturates occupancy at split_k = 1 on its own, yet the -// profiled-good choice still splits, because the win there is not filling idle -// subgroups but shortening each workgroup's serial K-reduction chain -// (gather-latency hiding), which does not care how many other workgroups are -// already in flight. kConvWideN and kConvWideNSplitKBlocks mirror that one sweep's -// shape -- wide N, two K tiers -- as the best available estimate on hardware with -// no sweep of its own; ClampSplitKToProblem (applied by SelectConvTiling after this -// returns) shrinks the pick back down if K or the scratch budget cannot actually -// support it. -constexpr uint32_t kConvWideN = 2 * kTileNCandidates[2]; // 128: needs 2+ tile_n=64 tiles -constexpr uint32_t kConvWideNSplitKBlocks = 2016 / kSubgroupMatrixK; // 126 K-blocks +// Once N needs at least two full-width (kTileNCandidates' max) N tiles, gather +// volume is already at its floor (see ConvHeuristicTiling's comment), and real +// hardware profiling (an xe-3lpg shape sweep) showed split-K keeps paying off past +// the point where occupancy alone would call the machine full: the win there is +// shortening each workgroup's serial K-reduction chain (gather-latency hiding), not +// filling idle subgroups, so a large M does not make split-K less useful the way it +// would under an occupancy model. kConvWideN marks where the search below drops its +// occupancy cap accordingly. +constexpr uint32_t kConvWideN = 2 * kTileNCandidates[2]; // 128: needs 2+ tile_n=64 tiles // Conv tiling heuristic: the sole tiling strategy SelectConvTiling uses, on every -// architecture. (SelectConvTiling previously also consulted a separate, hardcoded -// decision tree fit to a shape sweep collected on one specific GPU architecture, -// used only on that architecture; a per-architecture sweep like that does not -// generalize to hardware it was never collected on, so that tree was dropped in -// favor of this heuristic alone, pending a more broadly-informed retuning in a -// follow-up.) +// architecture. (A separate decision tree fit to a shape sweep on one GPU used to +// cover that one architecture; a sweep like that does not generalize to untested +// hardware, so it was dropped in favor of this heuristic alone, pending a more +// broadly-informed retuning in a follow-up.) // -// Does *not* pick tile_m/tile_n the way HeuristicTiling does -- shrinking tile_n to -// manufacture more independent occupancy tiles is free for MatMul (tile_n only -// trades off against tile_m for a fixed amount of global-memory traffic) but not -// for Conv: the im2col A tile (tile_m x K) is staged per-workgroup in workgroup -// memory, not shared, so it is re-gathered once per N tile that covers it. Total -// gather volume is M * N * K / tile_n -- inversely proportional to tile_n and -// independent of tile_m or split_k. Measured against a real hardware sweep, -// HeuristicTiling's occupancy-driven tile choice was shrinking tile_n to 16-32 on -// shapes that profiled best at 64, multiplying gather volume 2x-4x and dominating -// the kernel's cost regardless of how split-K was tuned -- which is why an earlier -// fix that only tightened split-K's fill target did not help. +// Does *not* pick tile_m/tile_n the way HeuristicTiling does. Shrinking tile_n to +// manufacture more occupancy tiles is free for MatMul (tile_n only trades off +// against tile_m for a fixed amount of global-memory traffic), but not for Conv: +// the im2col A tile is staged per-workgroup in workgroup memory, not shared, so it +// is re-gathered once per N tile that covers it -- gather volume is +// M * N * K / tile_n, independent of tile_m or split_k. A real hardware sweep +// showed HeuristicTiling's occupancy-driven choice shrinking tile_n to 16-32 on +// shapes that profiled best at 64, multiplying gather volume 2x-4x regardless of +// how split-K was tuned (why an earlier split-K-only fix did not help). // -// So tile_n is fixed at the largest candidate that fits N (minimizing gather -// volume) and tile_m at the smallest candidate (splitting M adds independent tiles -// without touching gather volume). split-K then takes one of two policies: below -// kConvWideN, an occupancy fill nearly identical in shape to HeuristicTiling's (not -// shared with it -- see that function's comment); at or above it, the -// M-independent K-tiered policy from kConvWideN's comment, since occupancy does not -// predict the profiled-good choice there. +// So tile_n is fixed at the largest candidate that fits N, and tile_m at the +// smallest (splitting M adds tiles without touching gather volume). split-K then +// takes the largest candidate with enough K-blocks per split and scratch headroom, +// subject to an occupancy cap below kConvWideN (see its comment for why that cap is +// dropped above it). SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { const uint32_t k_blocks = K / kSubgroupMatrixK; @@ -204,33 +192,32 @@ SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint } const uint32_t tile_m = kTileMCandidates[0]; + const bool wide_n = N >= kConvWideN; + const uint32_t hw = EffectiveHwSubgroups(arch); + auto tile_count = [](uint32_t dim, uint32_t tile) { return (dim + tile - 1) / tile; }; + const uint32_t num_tiles = batch * tile_count(M, tile_m) * tile_count(N, tile_n); + + constexpr uint32_t kMinBlocksPerSplit = 2; uint32_t split_k = 1; - if (N >= kConvWideN) { - // Occupancy (hw) is deliberately not consulted here -- see kConvWideN's comment. - split_k = k_blocks > kConvWideNSplitKBlocks ? 8u : 2u; - } else { - const uint32_t hw = EffectiveHwSubgroups(arch); - auto tile_count = [](uint32_t dim, uint32_t tile) { return (dim + tile - 1) / tile; }; - const uint32_t num_tiles = batch * tile_count(M, tile_m) * tile_count(N, tile_n); - constexpr uint32_t kMinBlocksPerSplit = 2; - for (uint32_t cand : kSplitKCandidates) { - if (k_blocks >= cand * kMinBlocksPerSplit && - num_tiles * cand <= 2 * hw && - ScratchElems({tile_m, tile_n, cand}, kConvScratchNPad) <= kMaxScratchElems) { - split_k = cand; - } + for (uint32_t cand : kSplitKCandidates) { + if (k_blocks >= cand * kMinBlocksPerSplit && + (wide_n || num_tiles * cand <= 2 * hw) && + ScratchElems({tile_m, tile_n, cand}, kConvScratchNPad) <= kMaxScratchElems) { + split_k = cand; } } return {tile_m, tile_n, split_k}; } // Repairs a predicted tiling to this problem by halving split_k until it fits both -// K's block count and the scratch budget. ConvHeuristicTiling's wide-N branch picks -// split_k from a K-blocks threshold without checking whether K actually has enough -// blocks to divide that many ways (e.g. K with only 1 block still gets a tentative -// split_k of 2 there), so this clamp is relied on as the actual correctness -// guarantee, not just a safety net. Clamping keeps the predicted tile shape -- the -// part that carries most of the win -- instead of discarding the prediction. +// K's block count and the scratch budget. ConvHeuristicTiling's own split-K search +// already checks both of these for the tiling it returns, so in practice this is a +// no-op safety net (mirroring the Conv kernel's own split-K clamp in +// subgroup_matrix_conv.cc) rather than something relied on to fire -- kept in case a +// future change to that search (or a future non-heuristic tiling source) ever +// predicts a split_k outside what K or the scratch budget can actually support. +// Clamping keeps the predicted tile shape -- the part that carries most of the win +// -- instead of discarding the prediction. void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) { const uint32_t k_blocks = K / kSubgroupMatrixK; while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) { From 0e3f41a47f7538052541aee2f060906f15876112 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Wed, 21 Oct 2026 15:05:38 +0800 Subject: [PATCH 08/14] Address comments and optimize ConvHeuristicTiling --- onnxruntime/core/providers/webgpu/nn/conv.cc | 27 +- onnxruntime/core/providers/webgpu/nn/conv.h | 34 +-- .../webgpu/nn/subgroup_matrix_conv.cc | 235 +++++------------- .../webgpu/nn/subgroup_matrix_conv.h | 56 +---- ...subgroup_matrix_conv_8x16x16.wgsl.template | 194 +++++---------- .../math/subgroup_matrix_tiling_selector.cc | 197 +++------------ .../math/subgroup_matrix_tiling_selector.h | 21 +- .../webgpu/subgroup_matrix_conv_test.cc | 47 ++-- 8 files changed, 202 insertions(+), 609 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index 92c028346fcdf..e655fe92a2671 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -40,11 +40,8 @@ Conv::ConvOptImplCache::GetOrCreate( template Status Conv::ComputeInternal(ComputeContext& context) const { - // Preferred fast path on devices with a vendor subgroup-matrix tiling policy (see - // subgroup_matrix_conv.h). Created once on first use (it needs a device query), then - // it decides inside Compute whether the kernel/attrs qualify. It declines - // (handled=false, allocating nothing) otherwise and we fall through to the normal - // path order below. + // 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)); @@ -347,21 +344,11 @@ Status Conv::PrePackInternal(ComputeContextBase& con return Status::OK(); } - // Subgroup-matrix Conv path: it reads the weight in its original OIHW layout (it - // applies its own OIHW -> OHWI transpose at runtime to build the GEMM right - // operand), so the weight must not be prepacked at all. - // - // Must come before the Im2ColMatMul branch below: ComputeInternal tries the - // subgroup-matrix path first, and on the devices this path targets - // CanApplyIm2ColMatMulProgram() accepts the same problems. If im2col prepacked - // first, SubgroupMatrixConvImpl::Compute would see a non-null prepacked_kernel_ and - // decline on every constant-weight Conv -- i.e. the path would silently never run. - // - // Placed before the auto_pad check below for the same pads-independence reason as - // the im2col branch (see the comment there). Skipping the prepack here is always - // safe even if Compute() later declines for another reason: is_packed stays false, - // so input 1 stays in OIHW and every fallback path can still transpose it at - // runtime. + // Subgroup-matrix Conv needs the weight in OIHW, so skip prepacking for it. Must + // come before Im2ColMatMul below: prepacking there first would leave + // prepacked_kernel_ non-null, making SubgroupMatrixConvImpl::Compute silently + // decline on every constant-weight Conv. Safe even if Compute() later declines + // for another reason: is_packed stays false, so input 1 stays OIHW. if (CanApplySubgroupMatrixConv(context, is_channels_last, kernel_shape, onnxruntime::narrow(conv_attrs_.group), tensor.DataType())) { return Status::OK(); diff --git a/onnxruntime/core/providers/webgpu/nn/conv.h b/onnxruntime/core/providers/webgpu/nn/conv.h index cf6459c25db33..41593c77ac43d 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -30,19 +30,15 @@ enum class KernelLayout { template class Conv : public WebGpuKernel { public: - // Abstract base class for alternative optimized Conv implementations (currently the - // subgroup-matrix path). Implementations hold the parent Conv and read - // conv_attrs/activation/... from it, mirroring Gemm::GemmOptImpl. + // 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; - // Attempts the optimized path, reading the Conv operands from `context` and the - // Conv attributes from the parent. Sets handled=true when it ran; leaves - // handled=false (allocating nothing) so the caller falls back to the normal Conv - // path. Called before ComputeInternal has resolved auto_pad or allocated the - // output, so an implementation that needs those runs its own shape inference. + // 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: @@ -54,10 +50,6 @@ class Conv : public WebGpuKernel { // Throw the Status, not the expression: the message names the activation ORT_ENFORCE would drop. ORT_THROW_IF_ERROR(GetFusedActivationAttr(info, activation_)); } - // Whether the weight input (index 1) is a constant initializer. An optimized - // implementation uses this to decide whether it may cache a derived form of the - // weight (e.g. a transposed copy) across Runs; a non-constant weight can change - // between Runs, so anything derived from it must be rebuilt every time. const Tensor* w = nullptr; w_is_constant_ = info.TryGetConstantInput(1, &w); } @@ -69,13 +61,10 @@ class Conv : public WebGpuKernel { AllocatorPtr alloc, /*out*/ bool& is_packed) override; - // State an optimized implementation (ConvOptImpl) reads from its parent Conv. const ConvAttributes& ConvAttrs() const { return conv_attrs_; } const Activation& ConvActivation() const { return activation_; } - // True when input 1 (the weight) is a constant initializer. See w_is_constant_. bool IsWeightConstant() const { return w_is_constant_; } - // The prepacked weight, or null when it is not prepacked. Its layout is HWIO or OHWI, - // as reported by kernel_layout_. + // Prepacked weight, or null if not prepacked (layout given by kernel_layout_). const Tensor* PrepackedKernel() const { return prepacked_kernel_.get(); } protected: @@ -92,19 +81,16 @@ class Conv : public WebGpuKernel { bool w_is_constant_ = false; // whether the weight input (index 1) is a constant initializer private: - // Owns the alternative optimized implementation and builds it lazily: the impl - // needs device capabilities only reachable through a compute context (not the - // ctor), and the same Conv kernel can see concurrent Compute calls, so creation is - // guarded by call_once -- mirroring MatMulComputeDispatcher, except nested here - // because ConvOptImpl is itself nested in this template. A null impl after - // initialization means this device has no optimized path. + // 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, not here: it calls into subgroup_matrix_conv.h, which - // includes this header, so the include can't go the other way. + // 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: diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index 10e0774014eb3..fff032851026b 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -25,34 +25,17 @@ namespace webgpu { namespace { -// Lanes per subgroup assumed by the subgroup-matrix kernel. The workgroup runs -// split_k subgroups, so its size is kSubgroupMatrixSubgroupSize * split_k. +// Workgroup size is kSubgroupMatrixSubgroupSize * split_k. constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; -// Combined workgroup-memory (SLM) budget in f16 elements for the two scratch -// buffers the conv kernel uses: the split-K output tiles -// (split_k * tile_m * tile_n) plus the per-subgroup im2col A staging -// (split_k * tile_m * sg_mat_k). split_k is clamped to fit this. Shares -// intel::kMaxScratchElems rather than hardcoding a second copy of the same -// budget the selector already sizes tilings against. - -// Largest per-tile subgroup-matrix counts the WGSL template unrolls: it generates the -// accumulators sg_mat_c_ for mi < 8 and ni < 4. A tiling above these would be -// accepted by the selector but silently mis-computed -- the kernel would accumulate -// only the blocks it generated while the write-out still covers the whole tile. The -// tile candidates the selector picks from stay inside this, so this is enforced as a -// safety net rather than something expected to fire. +// 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; -// Returns the device's subgroup-matrix config for the fixed 8x16x16 F16 config -// this kernel is implemented for, or nullopt when the device does not report it. -// Also requires subgroup-size control and the vendor tiling policy (Intel). -// -// Single source of truth for device eligibility: CanApplySubgroupMatrixConv (which -// Conv::PrePackInternal consults to skip prepacking) and CreateSubgroupMatrixConvImpl -// both go through here, so the prepack decision and the runtime path can never -// disagree about whether this device is served. +// 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 = @@ -60,12 +43,8 @@ std::optional GetSupportedConfig(const ComputeContextBase& if (!config || !config->Is(8, 16, 16)) { return std::nullopt; } - // The adapters reporting that config expose a 16-32 subgroup size range, while the - // kernel hard-codes kSubgroupMatrixSubgroupSize lanes per subgroup (sg_index = - // local_idx / 32). A driver-chosen 16 still computes the right answer -- the two - // hardware subgroups covered by one logical subgroup redundantly load, multiply and - // store the identical tile -- but it wastes about half the subgroup-matrix - // throughput, so the size is pinned with subgroup-size control instead. + // 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; } @@ -75,8 +54,6 @@ std::optional GetSupportedConfig(const ComputeContextBase& return config; } -// The bias add and the fused activation are not template parameters: the template -// calls the `write_output` helper that GenerateShaderCode builds from fuse_utils. 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) { @@ -91,17 +68,9 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelpe WGSL_TEMPLATE_VARIABLE(input_src, input_src)); } -// Subgroup-matrix Conv implementation. Computes a channels-last 2D convolution as -// an implicit GEMM during Compute: A is the im2col of the activation, gathered on -// the fly into workgroup memory (never materialized in global memory), B is the -// transposed weight, and the product is accumulated with -// subgroupMatrixMultiplyAccumulate. The class is intended to support all -// subgroup-matrix configs; for now only 8x16x16 is implemented. The per-problem -// output tiling is supplied by a vendor-specific selector kept internal to this -// impl, which declines problems its policy does not cover. -// It runs its own Conv shape inference to decide eligibility, mirroring -// Conv::ComputeInternal: ConvOptImpl::Compute is called before ComputeInternal has -// resolved auto_pad or allocated the output. +// 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; @@ -119,17 +88,17 @@ class SubgroupMatrixConvImpl final : public Conv::Co const ConvAttributes& conv_attrs = parent.ConvAttrs(); const Activation& activation = parent.ConvActivation(); - // The kernel reads the weight in its original OIHW layout (it applies its own - // OIHW -> OHWI transpose below), so any prepacked layout (HWIO or OHWI) is - // unusable here. Conv::PrePackInternal tests CanApplySubgroupMatrixConv before - // every prepacking branch and skips prepacking for exactly the problems accepted - // here, so this normally never declines. + // Fast decline for any Conv this path doesn't accept: PrePackInternal only skips + // prepacking for problems CanApplySubgroupMatrixConv accepts, so a non-null + // prepacked_kernel_ here (grouped/1x1/etc. convs prepacked for another path) + // means this one would decline below anyway -- this just avoids the shape + // inference to get there. The kernel needs OIHW regardless. if (parent.PrepackedKernel() != nullptr) { return Status::OK(); } - const auto* src = context.Input(0); + const auto* input = context.Input(0); const auto* weight = context.Input(1); - if (src == nullptr || weight == nullptr) { + if (input == nullptr || weight == nullptr) { return Status::OK(); } const TensorShape& weight_shape = weight->Shape(); // OIHW: [Cout, Cin, kh, kw] @@ -138,12 +107,10 @@ class SubgroupMatrixConvImpl final : public Conv::Co return Status::OK(); } - // Shape inference mirrors Conv::ComputeInternal. channels-last and a 4D weight are - // guaranteed by CanApplySubgroupMatrixConv; only a rank-4 NHWC activation is served - // (a reshaped Conv1D never gets here -- its kernel becomes 1xW, which - // CanApplySubgroupMatrixConv rejects). - const TensorShape& src_shape = src->Shape(); // NHWC: [batch, src_h, src_w, Cin] - if (src_shape.NumDimensions() != 4) { + // 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(); } ConvAttributes::ConvPadVector local_pads(conv_attrs.pads.begin(), conv_attrs.pads.end()); @@ -160,20 +127,17 @@ class SubgroupMatrixConvImpl final : public Conv::Co if (local_strides.empty()) { local_strides.resize(kernel_spacial_shape_vector.size(), 1); } - TensorShapeVector src_shape_vector = src_shape.AsShapeVector(); - TensorShapeVector output_shape_vector = {src_shape[0]}; + TensorShapeVector input_shape_vector = input_shape.AsShapeVector(); + TensorShapeVector output_shape_vector = {input_shape[0]}; TensorShape src_spacial_shape = TensorShape( - TensorShapeVector(std::next(src_shape_vector.begin()), std::prev(src_shape_vector.end()))); + TensorShapeVector(std::next(input_shape_vector.begin()), std::prev(input_shape_vector.end()))); ORT_RETURN_IF_ERROR(conv_attrs.InferPadsAndOutputShape(src_spacial_shape, kernel_spacial_shape_vector, local_strides, local_dilations, local_pads, output_shape_vector)); output_shape_vector.push_back(weight_shape[0]); // Cout const TensorShape output_shape(output_shape_vector); // NHWC: [batch, out_h, out_w, Cout] - // Spatial attributes as the kernel's uniforms want them: pads is [begin..., end...] - // in ONNX order, strides/dilations one entry per spatial dim. static_cast, not - // narrow, so an input the normal Conv path accepts is not turned into an abort here - // -- ComputeInternal casts these the same way. + // static_cast, not narrow: an input the normal Conv path accepts must not abort here. std::vector pads; std::vector strides; std::vector dilations; @@ -185,9 +149,9 @@ class SubgroupMatrixConvImpl final : public Conv::Co return Status::OK(); } - const uint32_t batch = narrow(src_shape[0]); - const uint32_t src_h = narrow(src_shape[1]); - const uint32_t src_w = narrow(src_shape[2]); + 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 channel_i = narrow(weight_shape[1]); const uint32_t kernel_h = narrow(weight_shape[2]); const uint32_t kernel_w = narrow(weight_shape[3]); @@ -198,63 +162,31 @@ class SubgroupMatrixConvImpl final : public Conv::Co const uint32_t K = kernel_h * kernel_w * channel_i; // contraction dim const uint32_t N = narrow(weight_shape[0]); // Cout - // An empty batch, M or N yields an empty output (ONNX permits zero-length dims). - // Let the generic Conv paths handle it rather than dispatch a degenerate - // (zero-workgroup) kernel. + // 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(); } - // Pick the per-problem output tiling from the vendor policy (Intel). The GEMM - // dims map directly: A [M, K] x B [K, N] with one z-slice per conv batch. The - // Conv selector uses the Conv scratch budget and the Conv-specific heuristic. const std::optional tiling = tiling_selector_(context, M, N, K, batch); if (!tiling) { return Status::OK(); } - // The kernel keeps the unchecked weight loads of a trailing partial N tile in bounds - // by shifting that tile's origin back to N - tile_n, which only exists when the tile - // fits within N. The Conv selector clamps tile_n for exactly this reason - // (ClampTileNToProblem), so this normally never triggers -- it mirrors the same guard - // in the subgroup-matrix MatMul/Gemm impls, keeping the precondition enforced on this - // side of the vendor boundary too. M needs no such guard: the A operand is staged in - // workgroup memory sized by the tile, so tile_m > M only wastes lanes. - // - // Last thing that can decline, so it comes before the output allocation and the - // weight transpose below: an impl that declines must not have allocated or - // dispatched anything. + // ClampTileNToProblem 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(); } - // Past this point the path is committed, so the output can be allocated: an impl - // that declines must allocate nothing. Tensor* output = context.Output(0, output_shape); - // Transpose the weight OIHW -> OHWI = [Cout, kh, kw, Cin], which flattens to the - // GEMM right operand B stored column-major: element B[k, n] lives at n*K + k, so - // the kernel loads B right tiles as col_major with stride K. (This matches the - // im2col-matmul weight layout; the alternative HWIO -> row-major stride-N load is - // the other option.) This allocates a GPU tensor and dispatches a Transpose - // program, so the path is already committed at this point. - // - // A constant weight is transposed once and reused; a non-constant one can change - // between Runs, so it has to be redone every time. Declining the non-constant case - // instead would cost more than the transpose it saves. + // Transpose weight OIHW -> OHWI = [Cout, kh, kw, Cin] (B column-major, B[k,n] at + // n*K+k). Constant weights are cached; non-constant ones are redone every Run. // - // TODO: prepack the constant case, which would drop the cache and the first-Run - // dispatch too. The layout is not the obstacle -- KernelLayout::OHWI is exactly - // what PrePackIm2ColMatMulWeight already produces, with the same {0, 2, 3, 1} perm. - // The blocker is the fallback: PrePackInternal must commit from the weight alone, - // while Compute can still decline at runtime (empty batch/M/N, or fewer than two - // spatial entries in dilations/pads/strides) onto a path that only reads OHWI via - // im2col -- and CanApplyIm2ColMatMulProgram is not implied by - // CanApplySubgroupMatrixConv (different gates: arch/activation whitelist vs - // subgroup-matrix config/fuse_utils), so a decline could land where OHWI can't be - // read. Prepacking therefore means dropping PrePackInternal's early return so the - // im2col branch prepacks when *both* predicates hold, and accepting a prepacked - // OHWI weight here instead of declining on it. + // TODO: prepack the constant case too (PrePackIm2ColMatMulWeight already produces + // this layout) -- blocked on PrePackInternal committing from the weight alone + // while Compute can still decline at runtime onto the im2col path. Tensor run_local_ohwi_weight; const Tensor* ohwi_weight = nullptr; if (parent.IsWeightConstant()) { @@ -279,35 +211,18 @@ class SubgroupMatrixConvImpl final : public Conv::Co tile_m, "x", tile_n, " spans ", sg_mat_count_m, "x", sg_mat_count_n, " subgroup matrices, at most ", kMaxSgMatCountM, "x", kMaxSgMatCountN); - // Safety net: the Conv selector already enforces the Conv scratch budget - // (output tiles + im2col A staging), so this clamp normally does nothing; it - // guards against a stray oversized split-K from the selector. Only split_k is - // clamped, not the tile area: ConvHeuristicTiling always picks the smallest - // tile_m (8), so the largest candidate tile needs 8 * (64 + 16) = 640 elements - // at split_k 1, or 8 * 8 * 80 = 5120 at the largest split_k (8), well inside - // the budget. - uint32_t split_k = tiling->split_k; - while (split_k > 1 && - split_k * tile_m * (tile_n + config.K) > intel::kMaxScratchElems) { - split_k /= 2; - } - 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 the input channel (cin), which is the - // fastest-varying, contiguous NHWC dimension. channel_i % vec_size == 0 keeps - // each vec chunk inside a single kernel position, so all lanes of a chunk share - // one bounds check and map to contiguous input. + // Vectorize the im2col gather along cin (contiguous in NHWC). const uint32_t vec_size = channel_i % 4 == 0 ? 4 : (channel_i % 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, activation}; program.SetWorkgroupSize(kSubgroupMatrixSubgroupSize * split_k); - // Pin the lanes per subgroup the kernel assumes; GetSupportedConfig has - // already established that the device supports subgroup-size control. program.SetSubgroupSize(kSubgroupMatrixSubgroupSize); program.SetDispatchGroupSize(dispatch_x, dispatch_y, batch); program.CacheHint(has_bias, static_cast(config.componentType), @@ -315,7 +230,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co config.M, config.N, config.K, config.subgroupSize, sg_mat_count_m, sg_mat_count_n, split_k, vec_size, activation.CacheKey()) - .AddInputs({{src, ProgramTensorMetadataDependency::TypeAndRank, static_cast(vec_size)}, + .AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank, static_cast(vec_size)}, {ohwi_weight, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddOutput({output, ProgramTensorMetadataDependency::Rank, output_shape, 1}) .AddUniformVariables({{M}, @@ -323,8 +238,8 @@ class SubgroupMatrixConvImpl final : public Conv::Co {K}, {dispatch_x}, {batch}, - {src_h}, - {src_w}, + {input_h}, + {input_w}, {channel_i}, {kernel_h}, {kernel_w}, @@ -344,14 +259,8 @@ class SubgroupMatrixConvImpl final : public Conv::Co } private: - // Builds the OIHW -> OHWI transposed weight on first use and caches it. Only valid - // because the caller checked the weight is a constant initializer. impl_ is shared - // across concurrent Compute calls on the same Conv kernel, so the build is guarded - // by call_once, which also synchronizes the read of cached_ohwi_weight_ below. - // The transpose is attempted exactly once: if it fails, every later Run on this - // Conv node reports the same error rather than retrying. - // const, with a mutable cache, like SubgroupMatrixMatMulImpl::EnsurePaddedB: the - // cached tensor is a memo of the constant weight, not observable state of the impl. + // 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_, [&]() { @@ -384,9 +293,8 @@ Status SubgroupMatrixConvProgram::GenerateShaderCode(ShaderHelper& shader) const const auto& output = shader.AddOutput( "output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); - // Write-out epilogue the template calls once per output element: bias add, then the - // fused activation (see the class doc in subgroup_matrix_conv.h for why this needs no - // per-ActivationKind handling). Mirrors SubgroupMatrixMatMulProgram::GenerateShaderCode. + // 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") @@ -413,9 +321,6 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, const TensorShape& kernel_shape, uint32_t group, MLDataType data_type) { - // The subgroup-matrix kernel is F16-only and computes a channels-last, - // non-grouped 2D convolution. A fused activation is applied in the write-out - // epilogue, so fused convs are accepted. if (data_type != DataTypeImpl::GetType()) { LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: data_type is not F16."; return false; @@ -437,9 +342,8 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, const uint32_t kernel_height = narrow(kernel_shape[2]); const uint32_t kernel_width = narrow(kernel_shape[3]); - // Mirror the im2col-matmul path: skip 1xW / Hx1 (and 1x1) kernels, which the - // plain MatMul / naive conv paths handle better. This also excludes reshaped - // Conv1D (kernel becomes 1xW). + // 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 << ")."; @@ -448,24 +352,10 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, const uint32_t K = kernel_height * kernel_width * channel_input; const uint32_t N = channel_output; - // The subgroup-matrix K must be a whole number of 16-wide blocks: the kernel walks - // K in sg_mat_k blocks with no tail handling. - // - // N must cover at least one whole subgroup-matrix N block (16). The kernel loads B - // straight from the weight buffer with subgroupMatrixLoad, which is not - // bounds-checked, and handles a trailing partial N tile by shifting the tile origin - // back to N - tile_n; below one block there is no tile_n that fits, so such a Cout - // cannot be served at all (it goes to the generic conv paths, which handle a tiny - // Cout better anyway). Must stay in agreement with the N guard in - // CreateSubgroupMatrixConvTilingSelector, or PrePackInternal would skip prepacking - // for a problem this path then declines. - // - // N has no parity or block-multiple requirement beyond that. The f16 right-operand - // load needs a 4-byte-aligned offset and row stride, but B is OHWI, i.e. - // column-major with stride K, so the offset n*K + kb*sg_mat_k is a multiple of - // sg_mat_k for any n -- alignment follows from K, not from N. (The MatMul kernel - // does need an even N: its B is [K, N] row-major, so N *is* the row stride, which is - // why it pads odd N. Do not copy that constraint here.) + // 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 @@ -474,14 +364,9 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, return false; } - // The weight buffer is bound as a single WGSL binding (input_b), and the shader's - // subgroupMatrixLoad calls address it with offsets derived from the full [N, K] - // weight -- they know nothing about ProgramManager's binding-segment split, which - // kicks in once a tensor exceeds maxStorageBufferBindingSize. A weight past that - // limit would still pass every check above, get split into multiple bindings, and - // then have its later output channels load past the first binding. Reject it here, - // before PrePackInternal or any output allocation/transpose, so it falls back to - // the normal (prepacked) conv paths instead. + // subgroupMatrixLoad addresses the weight as a single binding; a weight past + // maxStorageBufferBindingSize gets split by ProgramManager, which this offset + // arithmetic doesn't know about. const uint64_t max_weight_elems = context.DeviceLimits().maxStorageBufferBindingSize / sizeof(MLFloat16); if (static_cast(kernel_shape.Size()) > max_weight_elems) { LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: weight exceeds maxStorageBufferBindingSize" @@ -515,23 +400,15 @@ template std::unique_ptr::ConvOptImpl> CreateSubgroupMatrixConvImpl( const Conv& parent, const ComputeContextBase& context) { - // Only channels-last Conv is served by this kernel; the caller falls back to the - // normal Conv path for channels-first. if constexpr (!is_channels_last) { ORT_UNUSED_PARAMETER(parent); ORT_UNUSED_PARAMETER(context); return nullptr; } else { - // Only run on devices that report the fixed 8x16x16 F16 subgroup-matrix config - // this kernel is implemented for, and for which a vendor tiling policy exists - // (Intel today - GetSupportedConfig enforces both). const std::optional config = GetSupportedConfig(context); if (!config) { return nullptr; } - // The Conv selector shares the tile candidates with MatMul, but budgets the - // extra im2col staging the Conv kernel keeps in workgroup memory and uses the - // Conv-specific heuristic. SubgroupMatrixTilingSelector tiling_selector = intel::CreateSubgroupMatrixConvTilingSelector(context); if (!tiling_selector) { return nullptr; diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h index e04b89e8b7148..481ad8353a4b0 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -18,18 +18,9 @@ namespace onnxruntime { namespace webgpu { -// Computes a channels-last (NHWC) 2D convolution as an implicit-GEMM using -// subgroupMatrixMultiplyAccumulate. The activation operand is the im2col of the -// input, gathered on the fly into workgroup memory (never materialized in global -// memory); the weight is pre-transposed to OHWI, i.e. B stored column-major. -// config is the device subgroup-matrix config selected by SelectSubgroupMatrixConfig; -// sg_mat_count_m/n select how many subgroup matrices the output tile spans along -// M/N; split_k is the number of subgroups that cooperatively reduce K. -// -// The bias add and the fused activation live in a generated `write_output` helper -// rather than in the WGSL template: the activation body comes from -// GetActivationSnippet / GetActivationDeclaration, so every ActivationKind those -// support works here without the template knowing the kinds. +// 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, @@ -50,12 +41,9 @@ class SubgroupMatrixConvProgram final : public Program= 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); -// Creates a subgroup-matrix Conv implementation on devices whose vendor policy -// supports it; returns nullptr otherwise, so the caller falls back to the normal Conv -// paths. The impl reads the Conv attributes from `parent`, and the per-problem output -// tiling comes from a vendor-specific selector chosen internally from the device -// context. +// Nullptr on devices with no vendor tiling policy; callers fall back to normal Conv. template std::unique_ptr::ConvOptImpl> CreateSubgroupMatrixConvImpl( const Conv& parent, 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 index 5b5545655aaf5..6291d6f24dfdf 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template @@ -1,55 +1,25 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -// Subgroup-matrix Conv kernel (F16) for the 8x16x16 (sg_mat_m x sg_mat_n x -// sg_mat_k) subgroup-matrix shape. It computes a channels-last (NHWC) 2D -// convolution as an implicit GEMM Y[MxN] = A[MxK] @ B[KxN] (+ optional bias), -// where the activation operand A is the im2col of the input and never -// materialized in global memory. +// 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 // -// M = out_h * out_w (output pixels per batch slice) -// 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 +// the OHWI weight. +// B: weight pre-transposed to OHWI, column-major (B[k,n] at n*K+k). // -// A: im2col of the NHWC input. subgroupMatrixLoad requires a contiguous strided -// buffer, but im2col rows are gathered from scattered input pixels, so each -// K-block A-tile is gathered into workgroup scratch (a_scratch) first and then -// loaded as a left operand (subgroup_matrix_left) with -// stride kSgMatK. K index order is (kh, kw, cin) so it matches the OHWI weight. -// Because cin is the fastest-varying, contiguous NHWC dimension, the gather -// reads vec_size channels at a time (host guarantees channel_i % vec_size == 0). -// B: weight pre-transposed OIHW -> OHWI = [Cout, kh, kw, Cin], i.e. B stored -// column-major ([N, K] row-major; element B[k, n] at n*K + k). Loaded directly -// from global memory as a right operand (subgroup_matrix_right) with col_major and stride K. Shared across all batch slices. +// 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. // -// Workgroup: split_k subgroups x 32 lanes (split_k in {1,2,4,8}). -// Tile: kTileM x kTileN, chosen adaptively by the host tiling selector. Each -// subgroup computes the whole tile as kSgMatCountM x kSgMatCountN subgroup -// matrices. K is iterated in sg_mat_k blocks. When split_k > 1 the K blocks are -// distributed round-robin across the subgroups, each accumulates its own partial -// tile in shared memory, and the partials are summed at write-out. +// Bias/activation: write_output(), built by the host from GetActivationSnippet. // -// Unlike the plain subgroup-matrix MatMul kernel, A must be staged in workgroup -// memory before it can be loaded as a subgroup matrix. To keep the workgroup -// barriers in uniform control flow when split_k > 1, the K reduction runs as a -// per-step loop: every subgroup gathers its K-block A-tile into its own -// a_scratch slot, a workgroup barrier makes those tiles visible, then every -// subgroup multiply-accumulates its tile. This keeps all subgroups in lockstep -// at each barrier regardless of how many K-blocks remain. -// -// The bias add and the fused activation are not implemented here: the write-out -// pass calls write_output(), which the host builds in AdditionalImplementation from -// GetActivationSnippet / GetActivationDeclaration. That keeps every ActivationKind -// available without this template enumerating them. -// -// Preconditions enforced by the host: K % sg_mat_k == 0, and kTileN <= N (the tiling -// selector clamps tile_n and the impl re-checks it) so a trailing partial N tile can be -// shifted back far enough for its weight loads to stay in bounds. N itself may be odd -- -// B is column-major with stride K, so the load alignment comes from K, not N. M may be -// any size: a partial M tile just wastes a few a_scratch rows (its im2col reads stay -// in-bounds even past out_h/out_w, see load_im2col_vec), and the write-out bounds-checks -// M so those rows are never stored. +// Host guarantees: K % sg_mat_k == 0; kTileN <= N 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 @@ -72,30 +42,18 @@ const kSplitK: u32 = u32(split_k); // Number of subgroups cooperat const kSubgroupSize: u32 = 32; // Lanes per subgroup const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (channel_i % vec_size == 0) -// Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each -// split-K subgroup accumulates its slice of K into its own slot; the slots are -// summed during write-back. Storing here also lets us write to global memory -// with M/N bounds checks for dimensions that are not multiples of the tile size. +// kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// subgroup, summed at write-back. var scratch: array; -// Per-subgroup staging for the current K-block A-tile: kSplitK slots, each a -// [kTileM, kSgMatK] row-major (stride kSgMatK) tile gathered by im2col. +// Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. var a_scratch: array; -// Gathers a contiguous vec_size-wide chunk of im2col elements -// A[m, k_base .. k_base + vec_size - 1] for the given batch slice, where -// k_base = k_vec * vec_size. The chunk always lies within a single (kh, kw) -// kernel position: the K order is (kh, kw, cin) with cin fastest and the host -// guarantees channel_i % vec_size == 0, so all vec lanes share one (ih, iw) and a -// single bounds check and map to contiguous NHWC input. Out-of-range input pixels -// (padding) contribute zero. -// -// m past M (a partial M tile stages whole rows) is index-safe but not zero: oh >= out_h -// there, yet ih = oh * stride + kh * dilation - pad can still land inside [0, src_h), so -// the gather reads real input. That is harmless -- the (ih, iw) bounds check is what -// keeps every read inside the input buffer, and those rows are never written back. -// m -> output pixel (oh, ow); k_vec -> (kh_idx, kw_idx, cin_vec) with cin_vec fastest. -fn load_im2col_vec(m: u32, k_vec: u32, src_vec_batch_offset: u32) -> input_src_value_t { +// 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, cin_vec), cin_vec fastest. +fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { let channel_i_vec = uniforms.channel_i / kVecSize; let oh = m / uniforms.output_w; let ow = m % uniforms.output_w; @@ -105,61 +63,36 @@ fn load_im2col_vec(m: u32, k_vec: u32, src_vec_batch_offset: u32) -> input_src_v let kh_idx = khw / 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.src_h) || iw < 0 || iw >= i32(uniforms.src_w)) { + if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { return input_src_value_t(); } - let src_vec_idx = src_vec_batch_offset + - (u32(ih) * uniforms.src_w + u32(iw)) * channel_i_vec + cin_vec; - return input_src.getByOffset(src_vec_idx); + let input_vec_idx = input_batch_base_vec + + (u32(ih) * uniforms.input_w + u32(iw)) * channel_i_vec + cin_vec; + return input_src.getByOffset(input_vec_idx); } $MAIN { - // Recover the batch slice and the logical (n_tile, m_tile) from the flattened - // workgroup index. The host dispatches num_n_tile x num_m_tile tiles per batch - // slice on (x, y, z=batch), but that grid may be normalized (reshaped across - // x/y/z to stay within per-dimension limits), so workgroup_id.x/y/z cannot be - // used directly. workgroup_idx is the true linear index, invariant under - // normalization; batch slices are contiguous blocks of tiles_per_slice tiles. + // 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_slice = uniforms.num_n_tile * num_m_tile; - let batch_id = workgroup_idx / tiles_per_slice; // which batch slice - // Normalization rounds the workgroup count up (see - // ProgramManager::NormalizeDispatchGroupSize), so more workgroups can be launched - // than there are tiles. Those extras land past the last slice; without this guard - // they would derive batch_id >= batch and write past the end of the output (the - // in-tile M/N checks below cannot catch them, since their m_tile/n_tile look - // perfectly valid). workgroup_idx is workgroup-uniform, so the whole workgroup - // returns together and the barriers below stay in uniform control flow. + 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 slice_idx = workgroup_idx % tiles_per_slice; // tile index within the slice - let n_tile = slice_idx % uniforms.num_n_tile; // which kTileN-wide column tile + 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 the B right-operand loads stay inside - // the weight buffer. subgroupMatrixLoad has no bounds check of its own and the - // N-block j of this tile reads weight columns global_base_n + j*kSgMatN .. - // +kSgMatN-1, so an unshifted last tile would read up to kTileN-1 columns past - // Cout whenever N is not a multiple of kTileN. - // - // Unlike the MatMul kernel this needs no alignment fixup: B is column-major with - // stride K and K % kSgMatK == 0, so the load offset n*K + kb*kSgMatK is a - // multiple of kSgMatK for any n. The max() only guards the u32 subtraction; the - // host clamps kTileN to at most N, so the shifted tile always fits. + // 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). let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); - // Columns before the unshifted origin belong to the previous tile and are skipped - // at write-out, so a shifted tile never writes the same output element twice. - let n_skip = n_tile_start - global_base_n; - // M needs no such shift: the A left operand is gathered into a_scratch, which is - // sized by the tile rather than by M, so a partial M tile only stages unused rows - // (and the im2col gather itself is index-safe for them). - let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile + 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; - // Flat-element offset into the NHWC input for this batch slice. The weight is - // shared across batch slices, so B has no batch offset. The output is laid out - // as [batch, M, N] row-major. - let src_vec_batch_offset = batch_id * uniforms.src_h * uniforms.src_w * (uniforms.channel_i / kVecSize); - let out_batch_offset = batch_id * uniforms.M * uniforms.N; + let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.channel_i / 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) @@ -262,25 +195,21 @@ $MAIN { #endif #endif - // K reduction. Each subgroup processes K-blocks kb = sg_index, sg_index + - // kSplitK, ... The loop bound num_steps is workgroup-uniform, so the barriers - // below stay in uniform control flow even though a subgroup may run out of - // K-blocks before the last step (valid == false). + // 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; - // Stage this subgroup's [kTileM, kSgMatK] A-tile for K-block kb. The 32 - // lanes cooperatively gather the tile row-major (stride kSgMatK), each - // lane fetching a contiguous vec_size-wide chunk of cin per step. if (valid) { - let k_vec_base = kb * kSgMatK / kVecSize; - let k_vec_per_row = kSgMatK / kVecSize; - for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_row; vidx = vidx + kSubgroupSize) { - let row = vidx / k_vec_per_row; - let kc_vec = vidx % k_vec_per_row; - let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, src_vec_batch_offset); + 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; @@ -465,9 +394,8 @@ $MAIN { workgroupBarrier(); } - // Store this subgroup's partial results into its split-K slot of scratch, - // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) - // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. + // 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); @@ -567,11 +495,8 @@ $MAIN { // Make all subgroups' partial tiles visible before the reduction. workgroupBarrier(); - // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] - // tiles into slot 0 of scratch. Each lane owns a strided set of element - // positions and sums that position across all kSplitK slots. Writing slot 0 - // is safe: a position is only overwritten by the same lane that already read - // it as the slot-0 term of its own sum. + // 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); @@ -585,19 +510,14 @@ $MAIN { workgroupBarrier(); #endif - // Write-out pass: every subgroup writes out whole M-rows from the summed tile - // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows - // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a - // row's N elements (strided by the subgroup size). Columns below n_skip belong to - // the previous (unshifted) tile and M is bounds-checked, so non-multiple - // dimensions write each valid output element exactly once. write_output applies - // the bias and the fused activation before storing. + // 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_offset + global_m * uniforms.N + global_base_n; + 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])); 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 c97d3be7b947d..045a0799815f4 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 @@ -29,21 +29,8 @@ 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. -// Declared in subgroup_matrix_tiling_selector.h (as intel::kMaxScratchElems) so the -// Conv kernel's own split-K safety clamp can reuse the same value instead of -// hardcoding a second copy that has to be kept in sync by hand. - -// Extra per-tile N a kernel stages in workgroup memory beyond the split-K output -// tiles. The MatMul kernel stages nothing else (0); the Conv kernel also holds the -// im2col A tile (tile_m x sg_mat_k), so its scratch bound is -// split_k * tile_m * (tile_n + sg_mat_k) -- up to 2x stricter than MatMul's. This is -// the only difference in the scratch-budget math shared by ScratchElems, -// IsTilingValid and ClampSplitKToProblem below, so those helpers take it as a -// parameter instead of being duplicated per kernel. -// (The two kernels' tiling constraints differ elsewhere too -- e.g. ClampTileNToProblem, -// and MatMul alone has a pretuned table to consult before its heuristic -- those are -// handled separately.) +// Extra per-tile N staged beyond split-K output tiles: 0 for MatMul; kSubgroupMatrixK +// for Conv (which also stages an im2col A tile). constexpr uint32_t kMatMulScratchNPad = 0; constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK; @@ -52,8 +39,6 @@ uint64_t ScratchElems(const SubgroupMatrixTiling& t, uint32_t scratch_n_pad) { return static_cast(t.split_k) * t.tile_m * (t.tile_n + scratch_n_pad); } -// 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, uint32_t scratch_n_pad) { if (t.tile_m == 0 || t.tile_n == 0 || t.tile_m % kSubgroupMatrixM != 0 || t.tile_n % kSubgroupMatrixN != 0) { @@ -66,39 +51,23 @@ bool IsTilingValid(const SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n return ScratchElems(t, scratch_n_pad) <= 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 MatMul 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. -// -// Deliberately not shared with ConvHeuristicTiling below, even where the two look -// similar: shrinking tile_n to manufacture more occupancy tiles is free here (for -// MatMul, tile_n only trades off against tile_m for a fixed amount of -// global-memory traffic) but not for Conv -- see that function's comment -- so the -// two kernels' tile and split-K policies are expected to keep diverging and tend -// to get tuned/debugged independently. A shared helper would just mean every -// future change has to re-check it does not perturb the other kernel. +// 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) @@ -115,9 +84,6 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t const uint32_t tiles = batch * tile_count(M, tm) * tile_count(N, tn); const uint32_t area = tm * tn; if (tiles >= hw) { - // This tile fills the machine on its own. Among all such tiles, keep the - // largest-area one (best operand reuse); this also supersedes any - // under-filling pick made earlier. if (!filled || area > best_area) { filled = true; best_area = area; @@ -125,8 +91,6 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t tile_n = tn; } } else if (!filled && (tiles > best_tiles || (tiles == best_tiles && area > best_area))) { - // No tile has filled the machine yet: chase occupancy by maximizing the - // number of independent tiles, breaking ties toward larger area. best_tiles = tiles; best_area = area; tile_m = tm; @@ -135,8 +99,6 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t } } - // Add split-K only while the independent tiles under-fill the machine; cap the - // total at 2x hardware and keep enough K work per subgroup and scratch budget. const uint32_t num_tiles = batch * tile_count(M, tile_m) * tile_count(N, tile_n); constexpr uint32_t kMinBlocksPerSplit = 2; uint32_t split_k = 1; @@ -150,38 +112,14 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t return {tile_m, tile_n, split_k}; } -// Once N needs at least two full-width (kTileNCandidates' max) N tiles, gather -// volume is already at its floor (see ConvHeuristicTiling's comment), and real -// hardware profiling (an xe-3lpg shape sweep) showed split-K keeps paying off past -// the point where occupancy alone would call the machine full: the win there is -// shortening each workgroup's serial K-reduction chain (gather-latency hiding), not -// filling idle subgroups, so a large M does not make split-K less useful the way it -// would under an occupancy model. kConvWideN marks where the search below drops its -// occupancy cap accordingly. -constexpr uint32_t kConvWideN = 2 * kTileNCandidates[2]; // 128: needs 2+ tile_n=64 tiles - -// Conv tiling heuristic: the sole tiling strategy SelectConvTiling uses, on every -// architecture. (A separate decision tree fit to a shape sweep on one GPU used to -// cover that one architecture; a sweep like that does not generalize to untested -// hardware, so it was dropped in favor of this heuristic alone, pending a more -// broadly-informed retuning in a follow-up.) -// -// Does *not* pick tile_m/tile_n the way HeuristicTiling does. Shrinking tile_n to -// manufacture more occupancy tiles is free for MatMul (tile_n only trades off -// against tile_m for a fixed amount of global-memory traffic), but not for Conv: -// the im2col A tile is staged per-workgroup in workgroup memory, not shared, so it -// is re-gathered once per N tile that covers it -- gather volume is -// M * N * K / tile_n, independent of tile_m or split_k. A real hardware sweep -// showed HeuristicTiling's occupancy-driven choice shrinking tile_n to 16-32 on -// shapes that profiled best at 64, multiplying gather volume 2x-4x regardless of -// how split-K was tuned (why an earlier split-K-only fix did not help). -// -// So tile_n is fixed at the largest candidate that fits N, and tile_m at the -// smallest (splitting M adds tiles without touching gather volume). split-K then -// takes the largest candidate with enough K-blocks per split and scratch headroom, -// subject to an occupancy cap below kConvWideN (see its comment for why that cap is -// dropped above it). -SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { +// 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]; @@ -192,16 +130,10 @@ SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint } const uint32_t tile_m = kTileMCandidates[0]; - const bool wide_n = N >= kConvWideN; - const uint32_t hw = EffectiveHwSubgroups(arch); - auto tile_count = [](uint32_t dim, uint32_t tile) { return (dim + tile - 1) / tile; }; - const uint32_t num_tiles = batch * tile_count(M, tile_m) * tile_count(N, tile_n); - constexpr uint32_t kMinBlocksPerSplit = 2; uint32_t split_k = 1; for (uint32_t cand : kSplitKCandidates) { if (k_blocks >= cand * kMinBlocksPerSplit && - (wide_n || num_tiles * cand <= 2 * hw) && ScratchElems({tile_m, tile_n, cand}, kConvScratchNPad) <= kMaxScratchElems) { split_k = cand; } @@ -209,15 +141,8 @@ SubgroupMatrixTiling ConvHeuristicTiling(std::string_view arch, uint32_t M, uint return {tile_m, tile_n, split_k}; } -// Repairs a predicted tiling to this problem by halving split_k until it fits both -// K's block count and the scratch budget. ConvHeuristicTiling's own split-K search -// already checks both of these for the tiling it returns, so in practice this is a -// no-op safety net (mirroring the Conv kernel's own split-K clamp in -// subgroup_matrix_conv.cc) rather than something relied on to fire -- kept in case a -// future change to that search (or a future non-heuristic tiling source) ever -// predicts a split_k outside what K or the scratch budget can actually support. -// Clamping keeps the predicted tile shape -- the part that carries most of the win -// -- instead of discarding the prediction. +// Halves split_k until it fits K's block count and the scratch budget. No-op today +// (ConvHeuristicTiling already guarantees both); a safety net for future callers. void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) { const uint32_t k_blocks = K / kSubgroupMatrixK; while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) { @@ -225,19 +150,11 @@ void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_ } } -// Caps tile_n at the largest whole number of subgroup-matrix N blocks that fits in N. -// -// Correctness, not performance: the Conv kernel loads its B right operand straight -// from the weight buffer, and subgroupMatrixLoad is not bounds-checked, so a trailing -// partial N tile is shifted back to N - tile_n. That shift only exists when -// tile_n <= N. ConvHeuristicTiling already picks the largest candidate that fits N, -// but N itself need not be a multiple of kSubgroupMatrixN, so this cap is still -// needed. Shrinking tile_n keeps it a multiple of -// kSubgroupMatrixN and only lowers ScratchElems, so it cannot invalidate a tiling -// that was already valid. -// -// Callers must have rejected N < kSubgroupMatrixN, which cannot be served at all -// (there is no non-zero multiple of kSubgroupMatrixN that fits). +// 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 ClampTileNToProblem(SubgroupMatrixTiling& t, uint32_t N) { const uint32_t max_tile_n = (N / kSubgroupMatrixN) * kSubgroupMatrixN; if (t.tile_n > max_tile_n) { @@ -245,11 +162,8 @@ void ClampTileNToProblem(SubgroupMatrixTiling& t, uint32_t N) { } } -// 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. +// 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) { @@ -279,15 +193,9 @@ 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. -// -// Only for the batch-blind MatMul pretuned table: SelectConvTiling's -// ConvHeuristicTiling already takes batch into account when picking split-K (see its -// narrow-N branch), so applying this on top of it would double-count occupancy. +// 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 ClampSplitKForBatch(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); @@ -296,15 +204,11 @@ 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, kMatMulScratchNPad)) { 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)); } @@ -313,14 +217,9 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, return HeuristicTiling(arch, M, N, K, batch); } -// Conv counterpart of SelectTiling: always goes through ConvHeuristicTiling, under -// the Conv scratch budget. There is no MatMul-table lookup here: the Conv kernel's -// cost profile (the im2col gather) differs, and its scratch budget is stricter, so -// it must not reuse the MatMul-tuned entries; and unlike SelectTiling, there is no -// pretuned Conv entry point at all to try first (see ConvHeuristicTiling's comment -// for why). -SubgroupMatrixTiling SelectConvTiling(std::string_view arch, uint32_t M, uint32_t N, uint32_t K, uint32_t batch) { - SubgroupMatrixTiling tiling = ConvHeuristicTiling(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); ClampSplitKToProblem(tiling, K, kConvScratchNPad); ClampTileNToProblem(tiling, N); return tiling; @@ -333,13 +232,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; } @@ -353,26 +248,14 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( if (context.AdapterInfo().vendor != std::string_view{"intel"}) { return nullptr; } - // Conv tiling policy for the 8x16x16 subgroup-matrix Conv kernel: same tile - // candidates as MatMul, but selection uses the Conv scratch budget (which also - // stages the im2col A tile) and SelectConvTiling's ConvHeuristicTiling. - return [](const ComputeContext& context, uint32_t M, uint32_t N, - uint32_t K, uint32_t batch) -> std::optional { - // K must align to the subgroup tiling, and N must hold at least one whole - // subgroup-matrix N block: the Conv kernel loads B unchecked from the weight - // buffer and shifts a trailing partial N tile back to N - tile_n, which is - // impossible below one block. A partial N tile above that is served by the shift - // (see ClampTileNToProblem); M partial tiles need no guarantee at all, since the - // A operand is staged in workgroup memory sized by the tile. - // - // Declining here rather than clamping keeps this in agreement with - // CanApplySubgroupMatrixConv, which rejects the same N so Conv::PrePackInternal - // does not skip prepacking for a problem this path will not serve. + 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 (ClampTileNToProblem can't + // help below that); must agree with CanApplySubgroupMatrixConv's own check. if (K % kSubgroupMatrixK != 0 || N < kSubgroupMatrixN) { return std::nullopt; } - const std::string_view arch = std::string_view{context.AdapterInfo().architecture}; - return SelectConvTiling(arch, M, N, K, batch); + return SelectConvTiling(N, K); }; } 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 7447def75aa9f..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,26 +12,17 @@ namespace onnxruntime { namespace webgpu { namespace intel { -// Workgroup-memory (SLM) budget, in f16 elements, that the selector below fits -// tilings into. 16384 f16 == 32 KB, the common WebGPU maxComputeWorkgroupStorageSize. -// Shared with (not duplicated by) the Conv kernel's own split-K safety clamp, since -// both bound the same physical scratch allocation the selector sizes. +// 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; -// 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. +// 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); -// Like CreateSubgroupMatrixTilingSelector, but for the subgroup-matrix Conv path. -// The tile candidates are shared, but selection uses the Conv workgroup-memory -// budget (which also stages the im2col A tile) and a Conv-specific heuristic -// instead of the MatMul pretuned table. Conv has no pretuned table of its own: a -// per-architecture sweep does not generalize across hardware, so this always goes -// through the heuristic. Returns an empty selector on non-Intel devices. +// Conv counterpart: Conv scratch budget (also stages the im2col A tile) and a +// Conv-only heuristic, no pretuned table. SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( const ComputeContextBase& context); diff --git a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc index ffe3c8671e26d..570f4974fc41d 100644 --- a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc +++ b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc @@ -16,9 +16,7 @@ namespace onnxruntime { namespace test { -// Reference NCHW 2D convolution (group == 1). Mirrors the ONNX Conv semantics the -// WebGPU EP exercises after its internal NCHW->NHWC layout transform, so the -// subgroup-matrix implicit-GEMM conv can be validated against it. +// 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, @@ -54,10 +52,8 @@ static void ComputeExpectedConv(const std::vector& x, const std::vector& params, std::vector& y) { const float alpha = params.size() > 0 ? params[0] : 0.0f; @@ -98,22 +94,13 @@ static void ApplyExpectedActivation(const std::string& kind, const std::vector) with the -// attribute preserved -- the same kernel, now with a fused epilogue. -// -// These cases are device-independent by construction, so they are enabled rather than -// DISABLED_ like the *_large tests in this directory. The subgroup-matrix path needs an -// adapter that reports the 8x16x16 F16 subgroup-matrix config plus a vendor tiling -// policy; where that is missing, CanApplySubgroupMatrixConv() declines and the EP falls -// back to its regular conv kernels, which must produce the same result against the same -// reference. So a failure here is a real defect on either path, not a capability gap -- -// only the absence of the WebGPU EP itself is skipped. The problem sizes are small -// (<= 32x20x20 in, <= 130 out channels) to keep them cheap enough for that. +// 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, @@ -174,10 +161,8 @@ static void RunConvTest(int64_t N, int64_t Cin, int64_t H, int64_t W, test->AddAttribute("dilations", std::vector{dilation, dilation}); test->AddInput("X", x_dims, FloatsToMLFloat16s(x_vals)); - // A constant weight (the default) takes the branch that matters most: PrePackInternal - // must decline to prepack so the weight stays OIHW, and the impl caches its own OHWI - // transpose across runs. A non-constant weight instead redoes the transpose every - // run; RuntimeWeight covers that one. + // 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)); @@ -203,11 +188,11 @@ TEST(SubgroupMatrixConv, PartialTiles) { 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, which is what keeps the -// unchecked subgroupMatrixLoad of the weight inside the buffer. tile_n is clamped to -// the largest multiple of 16 that fits in Cout, so with Cout in (16, 32) the shifted -// tile overlaps the previous one by all but (Cout - 16) columns -- n_skip near kTileN, -// the narrowest write-out this kernel ever does. +// 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 From 1110cd1178c18910fe15f90f22546c0087dcdfad Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Thu, 22 Oct 2026 13:07:18 +0800 Subject: [PATCH 09/14] Address Jie's comments --- .../webgpu/nn/subgroup_matrix_conv.cc | 8 ++-- .../webgpu/nn/subgroup_matrix_conv.h | 2 +- ...subgroup_matrix_conv_8x16x16.wgsl.template | 18 ++++----- .../math/subgroup_matrix_tiling_selector.cc | 40 +++++++++---------- 4 files changed, 33 insertions(+), 35 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index fff032851026b..8f6d4ccc5b4a2 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -152,14 +152,14 @@ class SubgroupMatrixConvImpl final : public Conv::Co 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 channel_i = narrow(weight_shape[1]); + 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 * channel_i; // contraction dim + 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 @@ -215,7 +215,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co 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 = channel_i % 4 == 0 ? 4 : (channel_i % 2 == 0 ? 2 : 1); + 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; @@ -240,7 +240,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co {batch}, {input_h}, {input_w}, - {channel_i}, + {cin}, {kernel_h}, {kernel_w}, {output_w}, diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h index 481ad8353a4b0..5cb530f6289c4 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -44,7 +44,7 @@ class SubgroupMatrixConvProgram final : public Program 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, cin_vec), cin_vec fastest. +// 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 channel_i_vec = uniforms.channel_i / kVecSize; + let cin_vec = uniforms.cin / kVecSize; let oh = m / uniforms.output_w; let ow = m % uniforms.output_w; - let cin_vec = k_vec % channel_i_vec; - let khw = k_vec / channel_i_vec; - let kw_idx = khw % uniforms.kernel_w; - let kh_idx = khw / uniforms.kernel_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)) * channel_i_vec + cin_vec; + (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; return input_src.getByOffset(input_vec_idx); } @@ -91,7 +91,7 @@ $MAIN { 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.channel_i / kVecSize); + 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) 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 045a0799815f4..750c9fa38de3c 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 @@ -29,17 +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}; -// Extra per-tile N staged beyond split-K output tiles: 0 for MatMul; kSubgroupMatrixK -// for Conv (which also stages an im2col A tile). -constexpr uint32_t kMatMulScratchNPad = 0; -constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK; +// 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 scratch_n_pad) { - return static_cast(t.split_k) * t.tile_m * (t.tile_n + scratch_n_pad); +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); } -bool IsTilingValid(const SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) { +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; @@ -48,7 +49,7 @@ bool IsTilingValid(const SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n if (t.split_k == 0 || t.split_k > k_blocks) { return false; } - return ScratchElems(t, scratch_n_pad) <= kMaxScratchElems; + return ScratchElems(t, a_scratch_cols) <= kMaxScratchElems; } // Falls back to a conservative default when HwSubgroups doesn't recognize arch. @@ -84,6 +85,9 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t const uint32_t tiles = batch * tile_count(M, tm) * tile_count(N, tn); const uint32_t area = tm * tn; if (tiles >= hw) { + // This tile fills the machine on its own. Among all such tiles, keep the + // largest-area one (best operand reuse); this also supersedes any + // under-filling pick made earlier. if (!filled || area > best_area) { filled = true; best_area = area; @@ -91,6 +95,8 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t tile_n = tn; } } else if (!filled && (tiles > best_tiles || (tiles == best_tiles && area > best_area))) { + // No tile has filled the machine yet: chase occupancy by maximizing the + // number of independent tiles, breaking ties toward larger area. best_tiles = tiles; best_area = area; tile_m = tm; @@ -99,13 +105,15 @@ SubgroupMatrixTiling HeuristicTiling(std::string_view arch, uint32_t M, uint32_t } } + // Add split-K only while the independent tiles under-fill the machine; cap the + // total at 2x hardware and keep enough K work per subgroup and scratch budget. const uint32_t num_tiles = batch * tile_count(M, tile_m) * tile_count(N, tile_n); constexpr uint32_t kMinBlocksPerSplit = 2; uint32_t split_k = 1; for (uint32_t cand : kSplitKCandidates) { if (k_blocks >= cand * kMinBlocksPerSplit && num_tiles * cand <= 2 * hw && - ScratchElems({tile_m, tile_n, cand}, kMatMulScratchNPad) <= kMaxScratchElems) { + ScratchElems({tile_m, tile_n, cand}, kMatMulAScratchCols) <= kMaxScratchElems) { split_k = cand; } } @@ -134,22 +142,13 @@ SubgroupMatrixTiling ConvHeuristicTiling(uint32_t N, uint32_t K) { uint32_t split_k = 1; for (uint32_t cand : kSplitKCandidates) { if (k_blocks >= cand * kMinBlocksPerSplit && - ScratchElems({tile_m, tile_n, cand}, kConvScratchNPad) <= kMaxScratchElems) { + ScratchElems({tile_m, tile_n, cand}, kConvAScratchCols) <= kMaxScratchElems) { split_k = cand; } } return {tile_m, tile_n, split_k}; } -// Halves split_k until it fits K's block count and the scratch budget. No-op today -// (ConvHeuristicTiling already guarantees both); a safety net for future callers. -void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) { - const uint32_t k_blocks = K / kSubgroupMatrixK; - while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) { - t.split_k /= 2; - } -} - // 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 @@ -206,7 +205,7 @@ void ClampSplitKForBatch(SubgroupMatrixTiling& t, uint32_t M, uint32_t N, uint32 // 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, kMatMulScratchNPad)) { + if (const auto tuned = LookupPretunedTiling(arch, M, N, K); tuned && IsTilingValid(*tuned, K, kMatMulAScratchCols)) { SubgroupMatrixTiling tiling = *tuned; // The table is tuned at batch 1, so only revisit split-K when batch > 1. if (batch > 1) { @@ -220,7 +219,6 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, // Conv counterpart of SelectTiling: no pretuned table, always ConvHeuristicTiling. SubgroupMatrixTiling SelectConvTiling(uint32_t N, uint32_t K) { SubgroupMatrixTiling tiling = ConvHeuristicTiling(N, K); - ClampSplitKToProblem(tiling, K, kConvScratchNPad); ClampTileNToProblem(tiling, N); return tiling; } From 9ac6a2fdb800a293df5bd995d0e6f59f4b17a5cf Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Mon, 26 Oct 2026 10:14:19 +0800 Subject: [PATCH 10/14] Fix bots' errors --- .../webgpu/nn/subgroup_matrix_conv.cc | 6 +- .../nn/subgroup_matrix_conv_8x16x16.h | 1257 +++++++++++++++++ .../in_tree_golden/static-cpp-literal/index.h | 21 + .../static-cpp-literal/index_impl.h | 1 + .../nn/subgroup_matrix_conv_8x16x16.h | 1257 +++++++++++++++++ .../generated/tensor/oihw_to_ohwi.h | 54 +- .../static-cpp/generated/tensor/pad.h | 84 +- .../test/in_tree_golden/static-cpp/index.h | 21 + .../in_tree_golden/static-cpp/index_impl.h | 7 +- .../in_tree_golden/static-cpp/string_table.h | 235 ++- 10 files changed, 2804 insertions(+), 139 deletions(-) create mode 100644 tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/subgroup_matrix_conv_8x16x16.h create mode 100644 tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/subgroup_matrix_conv_8x16x16.h diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index 8f6d4ccc5b4a2..86b01379da956 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -134,7 +134,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co ORT_RETURN_IF_ERROR(conv_attrs.InferPadsAndOutputShape(src_spacial_shape, kernel_spacial_shape_vector, local_strides, local_dilations, local_pads, output_shape_vector)); - output_shape_vector.push_back(weight_shape[0]); // Cout + output_shape_vector.push_back(weight_shape[0]); // Cout const TensorShape output_shape(output_shape_vector); // NHWC: [batch, out_h, out_w, Cout] // static_cast, not narrow: an input the normal Conv path accepts must not abort here. @@ -414,12 +414,12 @@ std::unique_ptr::ConvOptImpl> CreateSu return nullptr; } return std::make_unique>(parent, *config, - std::move(tiling_selector)); + std::move(tiling_selector)); } } // Explicit instantiation for every Conv instantiation (see conv.cc). -#define WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(CHANNELS_LAST, FUSED) \ +#define WEBGPU_INSTANTIATE_CREATE_SUBGROUP_MATRIX_CONV_IMPL(CHANNELS_LAST, FUSED) \ template std::unique_ptr::ConvOptImpl> \ CreateSubgroupMatrixConvImpl(const Conv& parent, \ const ComputeContextBase& context); 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..d0eeaeaff7a1d --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/subgroup_matrix_conv_8x16x16.h @@ -0,0 +1,1257 @@ +// 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_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 | // the OHWI weight. +// 12 | // B: weight pre-transposed to OHWI, column-major (B[k,n] at n*K+k). +// 13 | // +// 14 | // Workgroup: split_k subgroups x 32 lanes, tile kTileM x kTileN (host-selected). K +// 15 | // blocks round-robin across subgroups when split_k > 1; partials summed at +// 16 | // write-out. Staging A forces the K reduction into a per-step +// 17 | // gather/barrier/multiply-accumulate loop in uniform control flow. +// 18 | // +// 19 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. +// 20 | // +// 21 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N so a trailing partial N tile can +// 22 | // shift back in bounds. M is unconstrained (see load_im2col_vec). +// 23 | +// 24 | #param sg_mat_k +// 25 | #param sg_mat_m +// 26 | #param sg_mat_n +// 27 | #param sg_mat_count_m +// 28 | #param sg_mat_count_n +// 29 | #param split_k +// 30 | #param vec_size +// 31 | +// 32 | #use .getByOffset .setByOffset +// 33 | +// 34 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +ss << "const kSgMatM: u32 = u32("; +ss << __param_sg_mat_m; +ss << ");\n"; +// 35 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +ss << "const kSgMatN: u32 = u32("; +ss << __param_sg_mat_n; +ss << ");\n"; +// 36 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +ss << "const kSgMatK: u32 = u32("; +ss << __param_sg_mat_k; +ss << ");\n"; +// 37 | 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"; +// 38 | 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"; +// 39 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +ss << "const kTileM: u32 = kSgMatM * kSgMatCountM;\n"; +// 40 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +ss << "const kTileN: u32 = kSgMatN * kSgMatCountN;\n"; +// 41 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +ss << "const kSplitK: u32 = u32("; +ss << __param_split_k; +ss << ");\n"; +// 42 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +ss << "const kSubgroupSize: u32 = 32;\n"; +// 43 | 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"; +// 44 | +ss << "\n"; +// 45 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// 46 | // subgroup, summed at write-back. +// 47 | var scratch: array; +ss << "var scratch: array;\n"; +// 48 | +ss << "\n"; +// 49 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 50 | var a_scratch: array; +ss << "var a_scratch: array;\n"; +// 51 | +ss << "\n"; +// 52 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always +// 53 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding +// 54 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. +// 55 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. +// 56 | 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"; +// 57 | let cin_vec = uniforms.cin / kVecSize; +ss << " let cin_vec = uniforms.cin / kVecSize;\n"; +// 58 | let oh = m / uniforms.output_w; +ss << " let oh = m / uniforms.output_w;\n"; +// 59 | let ow = m % uniforms.output_w; +ss << " let ow = m % uniforms.output_w;\n"; +// 60 | let k_cin = k_vec % cin_vec; +ss << " let k_cin = k_vec % cin_vec;\n"; +// 61 | let k_hw = k_vec / cin_vec; +ss << " let k_hw = k_vec / cin_vec;\n"; +// 62 | let kw_idx = k_hw % uniforms.kernel_w; +ss << " let kw_idx = k_hw % uniforms.kernel_w;\n"; +// 63 | let kh_idx = k_hw / uniforms.kernel_w; +ss << " let kh_idx = k_hw / uniforms.kernel_w;\n"; +// 64 | 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"; +// 65 | 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"; +// 66 | 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"; +// 67 | return input_src_value_t(); +ss << " return input_src_value_t();\n"; +// 68 | } +ss << " }\n"; +// 69 | let input_vec_idx = input_batch_base_vec + +ss << " let input_vec_idx = input_batch_base_vec +\n"; +// 70 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +ss << " (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin;\n"; +// 71 | return input_src.getByOffset(input_vec_idx); +ss << " return "; +ss << __var_input_src->GetByOffset("input_vec_idx"); +ss << ";\n"; +// 72 | } +ss << "}\n"; +// 73 | +ss << "\n"; +// 74 | $MAIN { +MainFunctionStart(); +ss << "\n"; +// 75 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +ss << "\n"; +// 76 | // the true linear index used to recover batch_id and (n_tile, m_tile). +// 77 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +ss << " let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM;\n"; +// 78 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +ss << " let tiles_per_batch = uniforms.num_n_tile * num_m_tile;\n"; +// 79 | let batch_id = workgroup_idx / tiles_per_batch; +ss << " let batch_id = workgroup_idx / tiles_per_batch;\n"; +// 80 | // Dispatch normalization can add extra workgroups past the last batch element; +ss << "\n"; +// 81 | // drop them here (workgroup-uniform, so the barriers below stay uniform). +// 82 | if (batch_id >= uniforms.batch) { +ss << " if (batch_id >= uniforms.batch) {\n"; +// 83 | return; +ss << " return;\n"; +// 84 | } +ss << " }\n"; +// 85 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +ss << " let tile_idx = workgroup_idx % tiles_per_batch;\n"; +// 86 | 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"; +// 87 | let n_tile_start = n_tile * kTileN; +ss << " let n_tile_start = n_tile * kTileN;\n"; +// 88 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +ss << "\n"; +// 89 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). +// 90 | let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); +ss << " let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN);\n"; +// 91 | 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"; +// 92 | 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"; +// 93 | let m_base = m_tile * kTileM; +ss << " let m_base = m_tile * kTileM;\n"; +// 94 | 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"; +// 95 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +ss << " let out_batch_base = batch_id * uniforms.M * uniforms.N;\n"; +// 96 | let k_blocks = uniforms.K / kSgMatK; +ss << " let k_blocks = uniforms.K / kSgMatK;\n"; +// 97 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +ss << " let sg_index = local_idx / kSubgroupSize;\n"; +// 98 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +ss << " let sg_lane = local_idx % kSubgroupSize;\n"; +// 99 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +ss << " let sg_scratch_base = sg_index * kTileM * kTileN;\n"; +// 100 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +ss << " let a_scratch_base = sg_index * kTileM * kSgMatK;\n"; +// 101 | +ss << "\n"; +// 102 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 103 | var sg_mat_c0_0: subgroup_matrix_result; +ss << " var sg_mat_c0_0: subgroup_matrix_result;\n"; +// 104 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 105 | var sg_mat_c0_1: subgroup_matrix_result; +ss << " var sg_mat_c0_1: subgroup_matrix_result;\n"; +// 106 | #endif +} +// 107 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 108 | var sg_mat_c0_2: subgroup_matrix_result; +ss << " var sg_mat_c0_2: subgroup_matrix_result;\n"; +// 109 | #endif +} +// 110 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 111 | var sg_mat_c0_3: subgroup_matrix_result; +ss << " var sg_mat_c0_3: subgroup_matrix_result;\n"; +// 112 | #endif +} +// 113 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 114 | var sg_mat_c1_0: subgroup_matrix_result; +ss << " var sg_mat_c1_0: subgroup_matrix_result;\n"; +// 115 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 116 | var sg_mat_c1_1: subgroup_matrix_result; +ss << " var sg_mat_c1_1: subgroup_matrix_result;\n"; +// 117 | #endif +} +// 118 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 119 | var sg_mat_c1_2: subgroup_matrix_result; +ss << " var sg_mat_c1_2: subgroup_matrix_result;\n"; +// 120 | #endif +} +// 121 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 122 | var sg_mat_c1_3: subgroup_matrix_result; +ss << " var sg_mat_c1_3: subgroup_matrix_result;\n"; +// 123 | #endif +} +// 124 | #endif +} +// 125 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 126 | var sg_mat_c2_0: subgroup_matrix_result; +ss << " var sg_mat_c2_0: subgroup_matrix_result;\n"; +// 127 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 128 | var sg_mat_c2_1: subgroup_matrix_result; +ss << " var sg_mat_c2_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_c2_2: subgroup_matrix_result; +ss << " var sg_mat_c2_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_c2_3: subgroup_matrix_result; +ss << " var sg_mat_c2_3: subgroup_matrix_result;\n"; +// 135 | #endif +} +// 136 | #endif +} +// 137 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 138 | var sg_mat_c3_0: subgroup_matrix_result; +ss << " var sg_mat_c3_0: subgroup_matrix_result;\n"; +// 139 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 140 | var sg_mat_c3_1: subgroup_matrix_result; +ss << " var sg_mat_c3_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_c3_2: subgroup_matrix_result; +ss << " var sg_mat_c3_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_c3_3: subgroup_matrix_result; +ss << " var sg_mat_c3_3: subgroup_matrix_result;\n"; +// 147 | #endif +} +// 148 | #endif +} +// 149 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 150 | var sg_mat_c4_0: subgroup_matrix_result; +ss << " var sg_mat_c4_0: subgroup_matrix_result;\n"; +// 151 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 152 | var sg_mat_c4_1: subgroup_matrix_result; +ss << " var sg_mat_c4_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_c4_2: subgroup_matrix_result; +ss << " var sg_mat_c4_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_c4_3: subgroup_matrix_result; +ss << " var sg_mat_c4_3: subgroup_matrix_result;\n"; +// 159 | #endif +} +// 160 | #endif +} +// 161 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 162 | var sg_mat_c5_0: subgroup_matrix_result; +ss << " var sg_mat_c5_0: subgroup_matrix_result;\n"; +// 163 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 164 | var sg_mat_c5_1: subgroup_matrix_result; +ss << " var sg_mat_c5_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_c5_2: subgroup_matrix_result; +ss << " var sg_mat_c5_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_c5_3: subgroup_matrix_result; +ss << " var sg_mat_c5_3: subgroup_matrix_result;\n"; +// 171 | #endif +} +// 172 | #endif +} +// 173 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 174 | var sg_mat_c6_0: subgroup_matrix_result; +ss << " var sg_mat_c6_0: subgroup_matrix_result;\n"; +// 175 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 176 | var sg_mat_c6_1: subgroup_matrix_result; +ss << " var sg_mat_c6_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_c6_2: subgroup_matrix_result; +ss << " var sg_mat_c6_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_c6_3: subgroup_matrix_result; +ss << " var sg_mat_c6_3: subgroup_matrix_result;\n"; +// 183 | #endif +} +// 184 | #endif +} +// 185 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 186 | var sg_mat_c7_0: subgroup_matrix_result; +ss << " var sg_mat_c7_0: subgroup_matrix_result;\n"; +// 187 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 188 | var sg_mat_c7_1: subgroup_matrix_result; +ss << " var sg_mat_c7_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_c7_2: subgroup_matrix_result; +ss << " var sg_mat_c7_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_c7_3: subgroup_matrix_result; +ss << " var sg_mat_c7_3: subgroup_matrix_result;\n"; +// 195 | #endif +} +// 196 | #endif +} +// 197 | +// 198 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, +// 199 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks +// 200 | // early (valid == false), keeping the barriers below in uniform control flow. +// 201 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +ss << " let num_steps = (k_blocks + kSplitK - 1u) / kSplitK;\n"; +// 202 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +ss << " for (var step: u32 = 0u; step < num_steps; step = step + 1u) {\n"; +// 203 | let kb = sg_index + step * kSplitK; +ss << " let kb = sg_index + step * kSplitK;\n"; +// 204 | let valid = kb < k_blocks; +ss << " let valid = kb < k_blocks;\n"; +// 205 | +ss << "\n"; +// 206 | if (valid) { +ss << " if (valid) {\n"; +// 207 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +ss << " let k_vec_per_sgmat_row = kSgMatK / kVecSize;\n"; +// 208 | let k_vec_base = kb * k_vec_per_sgmat_row; +ss << " let k_vec_base = kb * k_vec_per_sgmat_row;\n"; +// 209 | 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"; +// 210 | let row = vidx / k_vec_per_sgmat_row; +ss << " let row = vidx / k_vec_per_sgmat_row;\n"; +// 211 | let kc_vec = vidx % k_vec_per_sgmat_row; +ss << " let kc_vec = vidx % k_vec_per_sgmat_row;\n"; +// 212 | 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"; +// 213 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +ss << " let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize;\n"; +// 214 | #if vec_size == 1 +if (__param_vec_size == 1) { +// 215 | a_scratch[dst] = chunk; +ss << " a_scratch[dst] = chunk;\n"; +// 216 | #else +} else { +// 217 | a_scratch[dst + 0u] = chunk[0]; +ss << " a_scratch[dst + 0u] = chunk[0];\n"; +// 218 | a_scratch[dst + 1u] = chunk[1]; +ss << " a_scratch[dst + 1u] = chunk[1];\n"; +// 219 | #if vec_size >= 4 +if (__param_vec_size >= 4) { +// 220 | a_scratch[dst + 2u] = chunk[2]; +ss << " a_scratch[dst + 2u] = chunk[2];\n"; +// 221 | a_scratch[dst + 3u] = chunk[3]; +ss << " a_scratch[dst + 3u] = chunk[3];\n"; +// 222 | #endif +} +// 223 | #endif +} +// 224 | } +ss << " }\n"; +// 225 | } +ss << " }\n"; +// 226 | // Make the staged A-tiles visible before loading them as subgroup matrices. +ss << "\n"; +// 227 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 228 | +ss << "\n"; +// 229 | if (valid) { +ss << " if (valid) {\n"; +// 230 | // Load the B right tiles for this K block. The weight is OHWI, i.e. B is +ss << "\n"; +// 231 | // stored column-major ([N, K] row-major, element B[k, n] at n*K + k), so +// 232 | // the tile for N-block j (columns global_base_n + j*sg_mat_n) starts at +// 233 | // (global_base_n + j*sg_mat_n)*K + kb*sg_mat_k and is loaded col-major +// 234 | // with stride K. +// 235 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +ss << " let b_base = global_base_n * uniforms.K + kb * kSgMatK;\n"; +// 236 | var sg_mat_b0: subgroup_matrix_right = +ss << " var sg_mat_b0: subgroup_matrix_right =\n"; +// 237 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 238 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 239 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 240 | var sg_mat_b1: subgroup_matrix_right = +ss << " var sg_mat_b1: subgroup_matrix_right =\n"; +// 241 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 242 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 243 | #endif +} +// 244 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 245 | var sg_mat_b2: subgroup_matrix_right = +ss << " var sg_mat_b2: subgroup_matrix_right =\n"; +// 246 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 247 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 248 | #endif +} +// 249 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 250 | var sg_mat_b3: subgroup_matrix_right = +ss << " var sg_mat_b3: subgroup_matrix_right =\n"; +// 251 | subgroupMatrixLoad, col_major>( +ss << " subgroupMatrixLoad, col_major>(\n"; +// 252 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +ss << " &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K);\n"; +// 253 | #endif +} +// 254 | +// 255 | // Load the A left tiles (one per M block) from workgroup staging, +// 256 | // row-major with stride kSgMatK. +// 257 | var sg_mat_a0: subgroup_matrix_left = +ss << " var sg_mat_a0: subgroup_matrix_left =\n"; +// 258 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 259 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 260 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 261 | var sg_mat_a1: subgroup_matrix_left = +ss << " var sg_mat_a1: subgroup_matrix_left =\n"; +// 262 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 263 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 264 | #endif +} +// 265 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 266 | var sg_mat_a2: subgroup_matrix_left = +ss << " var sg_mat_a2: subgroup_matrix_left =\n"; +// 267 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 268 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 269 | #endif +} +// 270 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 271 | var sg_mat_a3: subgroup_matrix_left = +ss << " var sg_mat_a3: subgroup_matrix_left =\n"; +// 272 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 273 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 274 | #endif +} +// 275 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 276 | var sg_mat_a4: subgroup_matrix_left = +ss << " var sg_mat_a4: subgroup_matrix_left =\n"; +// 277 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 278 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 279 | #endif +} +// 280 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 281 | var sg_mat_a5: subgroup_matrix_left = +ss << " var sg_mat_a5: subgroup_matrix_left =\n"; +// 282 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 283 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 284 | #endif +} +// 285 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 286 | var sg_mat_a6: subgroup_matrix_left = +ss << " var sg_mat_a6: subgroup_matrix_left =\n"; +// 287 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 288 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 289 | #endif +} +// 290 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 291 | var sg_mat_a7: subgroup_matrix_left = +ss << " var sg_mat_a7: subgroup_matrix_left =\n"; +// 292 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 293 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +ss << " &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK);\n"; +// 294 | #endif +} +// 295 | +// 296 | // Accumulate every (M block, N block) pair. +// 297 | 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"; +// 298 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 299 | 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"; +// 300 | #endif +} +// 301 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 302 | 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"; +// 303 | #endif +} +// 304 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 305 | 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"; +// 306 | #endif +} +// 307 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 308 | 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"; +// 309 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 310 | 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"; +// 311 | #endif +} +// 312 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 313 | 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"; +// 314 | #endif +} +// 315 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 316 | 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"; +// 317 | #endif +} +// 318 | #endif +} +// 319 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 320 | 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"; +// 321 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 322 | 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"; +// 323 | #endif +} +// 324 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 325 | 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"; +// 326 | #endif +} +// 327 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 328 | 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"; +// 329 | #endif +} +// 330 | #endif +} +// 331 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 332 | 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"; +// 333 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 334 | 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"; +// 335 | #endif +} +// 336 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 337 | 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"; +// 338 | #endif +} +// 339 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 340 | 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"; +// 341 | #endif +} +// 342 | #endif +} +// 343 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 344 | 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"; +// 345 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 346 | 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"; +// 347 | #endif +} +// 348 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 349 | 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"; +// 350 | #endif +} +// 351 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 352 | 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"; +// 353 | #endif +} +// 354 | #endif +} +// 355 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 356 | 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"; +// 357 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 358 | 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"; +// 359 | #endif +} +// 360 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 361 | 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"; +// 362 | #endif +} +// 363 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 364 | 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"; +// 365 | #endif +} +// 366 | #endif +} +// 367 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 368 | 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"; +// 369 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 370 | 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"; +// 371 | #endif +} +// 372 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 373 | 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"; +// 374 | #endif +} +// 375 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 376 | 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"; +// 377 | #endif +} +// 378 | #endif +} +// 379 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 380 | 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"; +// 381 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 382 | 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"; +// 383 | #endif +} +// 384 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 385 | 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"; +// 386 | #endif +} +// 387 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 388 | 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"; +// 389 | #endif +} +// 390 | #endif +} +// 391 | } +ss << " }\n"; +// 392 | // Ensure every subgroup is done reading a_scratch before it is reused for +ss << "\n"; +// 393 | // the next K-block step. +// 394 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 395 | } +ss << " }\n"; +// 396 | +ss << "\n"; +// 397 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; +// 398 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). +// 399 | 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"; +// 400 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 401 | 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"; +// 402 | #endif +} +// 403 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 404 | 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"; +// 405 | #endif +} +// 406 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 407 | 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"; +// 408 | #endif +} +// 409 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 410 | 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"; +// 411 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 412 | 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"; +// 413 | #endif +} +// 414 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 415 | 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"; +// 416 | #endif +} +// 417 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 418 | 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"; +// 419 | #endif +} +// 420 | #endif +} +// 421 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 422 | 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"; +// 423 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 424 | 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"; +// 425 | #endif +} +// 426 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 427 | 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"; +// 428 | #endif +} +// 429 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 430 | 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"; +// 431 | #endif +} +// 432 | #endif +} +// 433 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 434 | 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"; +// 435 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 436 | 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"; +// 437 | #endif +} +// 438 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 439 | 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"; +// 440 | #endif +} +// 441 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 442 | 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"; +// 443 | #endif +} +// 444 | #endif +} +// 445 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 446 | 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"; +// 447 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 448 | 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"; +// 449 | #endif +} +// 450 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 451 | 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"; +// 452 | #endif +} +// 453 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 454 | 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"; +// 455 | #endif +} +// 456 | #endif +} +// 457 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 458 | 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"; +// 459 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 460 | 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"; +// 461 | #endif +} +// 462 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 463 | 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"; +// 464 | #endif +} +// 465 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 466 | 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"; +// 467 | #endif +} +// 468 | #endif +} +// 469 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 470 | 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"; +// 471 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 472 | 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"; +// 473 | #endif +} +// 474 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 475 | 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"; +// 476 | #endif +} +// 477 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 478 | 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"; +// 479 | #endif +} +// 480 | #endif +} +// 481 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 482 | 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"; +// 483 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 484 | 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"; +// 485 | #endif +} +// 486 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 487 | 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"; +// 488 | #endif +} +// 489 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 490 | 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"; +// 491 | #endif +} +// 492 | #endif +} +// 493 | +// 494 | #if split_k >= 2 +if (__param_split_k >= 2) { +// 495 | // Make all subgroups' partial tiles visible before the reduction. +// 496 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 497 | +ss << "\n"; +// 498 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each +// 499 | // position is only overwritten by the lane that already read it. +// 500 | if (sg_index == 0u) { +ss << " if (sg_index == 0u) {\n"; +// 501 | 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"; +// 502 | var acc: f16 = f16(0); +ss << " var acc: f16 = f16(0);\n"; +// 503 | for (var s: u32 = 0; s < kSplitK; s++) { +ss << " for (var s: u32 = 0; s < kSplitK; s++) {\n"; +// 504 | acc += scratch[s * kTileM * kTileN + idx]; +ss << " acc += scratch[s * kTileM * kTileN + idx];\n"; +// 505 | } +ss << " }\n"; +// 506 | scratch[idx] = acc; +ss << " scratch[idx] = acc;\n"; +// 507 | } +ss << " }\n"; +// 508 | } +ss << " }\n"; +// 509 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +ss << "\n"; +// 510 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 511 | #endif +} +// 512 | +// 513 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip +// 514 | // and M's bounds check make non-multiple dimensions write each element once. +// 515 | let n_count = min(kTileN, uniforms.N - global_base_n); +ss << " let n_count = min(kTileN, uniforms.N - global_base_n);\n"; +// 516 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +ss << " for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) {\n"; +// 517 | let global_m = m_base + r; +ss << " let global_m = m_base + r;\n"; +// 518 | if (global_m < uniforms.M) { +ss << " if (global_m < uniforms.M) {\n"; +// 519 | let scratch_base = r * kTileN; +ss << " let scratch_base = r * kTileN;\n"; +// 520 | 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"; +// 521 | 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"; +// 522 | write_output(out_base + i, global_base_n + i, +ss << " write_output(out_base + i, global_base_n + i,\n"; +// 523 | output_value_t(scratch[scratch_base + i])); +ss << " output_value_t(scratch[scratch_base + i]));\n"; +// 524 | } +ss << " }\n"; +// 525 | } +ss << " }\n"; +// 526 | } +ss << " }\n"; +// 527 | } // MAIN +MainFunctionEnd(); +ss << "\n"; +// 528 | + + + 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..31b02837e4c3a 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,27 @@ 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_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..7e76b49326049 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 @@ -40,6 +40,7 @@ std::string pass_as_string(T&& v) { #include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // c01fae062d8059604670635900f3dd036b692cc3d16ebb8a1e54491e83439e66 #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" // 8ea2bc37b92b7f0b615005b3856710d09f1a37f48e3680252bdc239f8f0ba2d5 #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/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..8be84a7ab6a25 --- /dev/null +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/subgroup_matrix_conv_8x16x16.h @@ -0,0 +1,1257 @@ +// 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_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 | // the OHWI weight. +// 12 | // B: weight pre-transposed to OHWI, column-major (B[k,n] at n*K+k). +// 13 | // +// 14 | // Workgroup: split_k subgroups x 32 lanes, tile kTileM x kTileN (host-selected). K +// 15 | // blocks round-robin across subgroups when split_k > 1; partials summed at +// 16 | // write-out. Staging A forces the K reduction into a per-step +// 17 | // gather/barrier/multiply-accumulate loop in uniform control flow. +// 18 | // +// 19 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. +// 20 | // +// 21 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N so a trailing partial N tile can +// 22 | // shift back in bounds. M is unconstrained (see load_im2col_vec). +// 23 | +// 24 | #param sg_mat_k +// 25 | #param sg_mat_m +// 26 | #param sg_mat_n +// 27 | #param sg_mat_count_m +// 28 | #param sg_mat_count_n +// 29 | #param split_k +// 30 | #param vec_size +// 31 | +// 32 | #use .getByOffset .setByOffset +// 33 | +// 34 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +ss << __str_32; +ss << __param_sg_mat_m; +ss << __str_17; +// 35 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +ss << __str_33; +ss << __param_sg_mat_n; +ss << __str_17; +// 36 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +ss << __str_34; +ss << __param_sg_mat_k; +ss << __str_17; +// 37 | 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; +// 38 | 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; +// 39 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +ss << __str_37; +// 40 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +ss << __str_38; +// 41 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +ss << __str_39; +ss << __param_split_k; +ss << __str_17; +// 42 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +ss << __str_40; +// 43 | 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; +// 44 | +ss << __str_5; +// 45 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// 46 | // subgroup, summed at write-back. +// 47 | var scratch: array; +ss << __str_41; +// 48 | +ss << __str_5; +// 49 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 50 | var a_scratch: array; +ss << __str_408; +// 51 | +ss << __str_5; +// 52 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always +// 53 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding +// 54 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. +// 55 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. +// 56 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { +ss << __str_409; +// 57 | let cin_vec = uniforms.cin / kVecSize; +ss << __str_410; +// 58 | let oh = m / uniforms.output_w; +ss << __str_411; +// 59 | let ow = m % uniforms.output_w; +ss << __str_412; +// 60 | let k_cin = k_vec % cin_vec; +ss << __str_413; +// 61 | let k_hw = k_vec / cin_vec; +ss << __str_414; +// 62 | let kw_idx = k_hw % uniforms.kernel_w; +ss << __str_415; +// 63 | let kh_idx = k_hw / uniforms.kernel_w; +ss << __str_416; +// 64 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); +ss << __str_417; +// 65 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); +ss << __str_418; +// 66 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { +ss << __str_419; +// 67 | return input_src_value_t(); +ss << __str_420; +// 68 | } +ss << __str_20; +// 69 | let input_vec_idx = input_batch_base_vec + +ss << __str_421; +// 70 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +ss << __str_422; +// 71 | return input_src.getByOffset(input_vec_idx); +ss << __str_351; +ss << __var_input_src->GetByOffset(__str_406); +ss << __str_29; +// 72 | } +ss << __str_345; +// 73 | +ss << __str_5; +// 74 | $MAIN { +MainFunctionStart(); +ss << __str_5; +// 75 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +ss << __str_5; +// 76 | // the true linear index used to recover batch_id and (n_tile, m_tile). +// 77 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +ss << __str_218; +// 78 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +ss << __str_423; +// 79 | let batch_id = workgroup_idx / tiles_per_batch; +ss << __str_424; +// 80 | // Dispatch normalization can add extra workgroups past the last batch element; +ss << __str_5; +// 81 | // drop them here (workgroup-uniform, so the barriers below stay uniform). +// 82 | if (batch_id >= uniforms.batch) { +ss << __str_425; +// 83 | return; +ss << __str_426; +// 84 | } +ss << __str_20; +// 85 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +ss << __str_427; +// 86 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile +ss << __str_428; +// 87 | let n_tile_start = n_tile * kTileN; +ss << __str_45; +// 88 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +ss << __str_5; +// 89 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). +// 90 | let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); +ss << __str_47; +// 91 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile +ss << __str_205; +// 92 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile +ss << __str_429; +// 93 | let m_base = m_tile * kTileM; +ss << __str_430; +// 94 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); +ss << __str_431; +// 95 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +ss << __str_432; +// 96 | let k_blocks = uniforms.K / kSgMatK; +ss << __str_49; +// 97 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +ss << __str_50; +// 98 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +ss << __str_51; +// 99 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +ss << __str_52; +// 100 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +ss << __str_433; +// 101 | +ss << __str_5; +// 102 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 103 | 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; +// 104 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 105 | 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; +// 106 | #endif +} +// 107 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 108 | 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; +// 109 | #endif +} +// 110 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 111 | 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; +// 112 | #endif +} +// 113 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 114 | 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; +// 115 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 116 | 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; +// 117 | #endif +} +// 118 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 119 | 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; +// 120 | #endif +} +// 121 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 122 | 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; +// 123 | #endif +} +// 124 | #endif +} +// 125 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 126 | 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; +// 127 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 128 | 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; +// 129 | #endif +} +// 130 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 131 | 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; +// 132 | #endif +} +// 133 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 134 | 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; +// 135 | #endif +} +// 136 | #endif +} +// 137 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 138 | 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; +// 139 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 140 | 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; +// 141 | #endif +} +// 142 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 143 | 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; +// 144 | #endif +} +// 145 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 146 | 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; +// 147 | #endif +} +// 148 | #endif +} +// 149 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 150 | 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; +// 151 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 152 | 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; +// 153 | #endif +} +// 154 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 155 | 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; +// 156 | #endif +} +// 157 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 158 | 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; +// 159 | #endif +} +// 160 | #endif +} +// 161 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 162 | 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; +// 163 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 164 | 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; +// 165 | #endif +} +// 166 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 167 | 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; +// 168 | #endif +} +// 169 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 170 | 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; +// 171 | #endif +} +// 172 | #endif +} +// 173 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 174 | 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; +// 175 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 176 | 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; +// 177 | #endif +} +// 178 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 179 | 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; +// 180 | #endif +} +// 181 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 182 | 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; +// 183 | #endif +} +// 184 | #endif +} +// 185 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 186 | 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; +// 187 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 188 | 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; +// 189 | #endif +} +// 190 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 191 | 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; +// 192 | #endif +} +// 193 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 194 | 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; +// 195 | #endif +} +// 196 | #endif +} +// 197 | +// 198 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, +// 199 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks +// 200 | // early (valid == false), keeping the barriers below in uniform control flow. +// 201 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +ss << __str_434; +// 202 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +ss << __str_435; +// 203 | let kb = sg_index + step * kSplitK; +ss << __str_436; +// 204 | let valid = kb < k_blocks; +ss << __str_437; +// 205 | +ss << __str_5; +// 206 | if (valid) { +ss << __str_438; +// 207 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +ss << __str_439; +// 208 | let k_vec_base = kb * k_vec_per_sgmat_row; +ss << __str_440; +// 209 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { +ss << __str_441; +// 210 | let row = vidx / k_vec_per_sgmat_row; +ss << __str_442; +// 211 | let kc_vec = vidx % k_vec_per_sgmat_row; +ss << __str_443; +// 212 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); +ss << __str_444; +// 213 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +ss << __str_445; +// 214 | #if vec_size == 1 +if (__param_vec_size == 1) { +// 215 | a_scratch[dst] = chunk; +ss << __str_446; +// 216 | #else +} else { +// 217 | a_scratch[dst + 0u] = chunk[0]; +ss << __str_447; +// 218 | a_scratch[dst + 1u] = chunk[1]; +ss << __str_448; +// 219 | #if vec_size >= 4 +if (__param_vec_size >= 4) { +// 220 | a_scratch[dst + 2u] = chunk[2]; +ss << __str_449; +// 221 | a_scratch[dst + 3u] = chunk[3]; +ss << __str_450; +// 222 | #endif +} +// 223 | #endif +} +// 224 | } +ss << __str_202; +// 225 | } +ss << __str_204; +// 226 | // Make the staged A-tiles visible before loading them as subgroup matrices. +ss << __str_5; +// 227 | workgroupBarrier(); +ss << __str_451; +// 228 | +ss << __str_5; +// 229 | if (valid) { +ss << __str_438; +// 230 | // Load the B right tiles for this K block. The weight is OHWI, i.e. B is +ss << __str_5; +// 231 | // stored column-major ([N, K] row-major, element B[k, n] at n*K + k), so +// 232 | // the tile for N-block j (columns global_base_n + j*sg_mat_n) starts at +// 233 | // (global_base_n + j*sg_mat_n)*K + kb*sg_mat_k and is loaded col-major +// 234 | // with stride K. +// 235 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +ss << __str_452; +// 236 | 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; +// 237 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 238 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_455; +// 239 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 240 | 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; +// 241 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 242 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_457; +// 243 | #endif +} +// 244 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 245 | 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; +// 246 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 247 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_459; +// 248 | #endif +} +// 249 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 250 | 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; +// 251 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 252 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_461; +// 253 | #endif +} +// 254 | +// 255 | // Load the A left tiles (one per M block) from workgroup staging, +// 256 | // row-major with stride kSgMatK. +// 257 | var sg_mat_a0: subgroup_matrix_left = +ss << __str_462; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 258 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 259 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_464; +// 260 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 261 | var sg_mat_a1: subgroup_matrix_left = +ss << __str_465; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 262 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 263 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_466; +// 264 | #endif +} +// 265 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 266 | var sg_mat_a2: subgroup_matrix_left = +ss << __str_467; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 267 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 268 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_468; +// 269 | #endif +} +// 270 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 271 | var sg_mat_a3: subgroup_matrix_left = +ss << __str_469; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 272 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 273 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_470; +// 274 | #endif +} +// 275 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 276 | var sg_mat_a4: subgroup_matrix_left = +ss << __str_471; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 277 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 278 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_472; +// 279 | #endif +} +// 280 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 281 | var sg_mat_a5: subgroup_matrix_left = +ss << __str_473; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 282 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 283 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_474; +// 284 | #endif +} +// 285 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 286 | var sg_mat_a6: subgroup_matrix_left = +ss << __str_475; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 287 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 288 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_476; +// 289 | #endif +} +// 290 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 291 | var sg_mat_a7: subgroup_matrix_left = +ss << __str_477; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 292 | subgroupMatrixLoad, row_major>( +ss << __str_463; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 293 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_478; +// 294 | #endif +} +// 295 | +// 296 | // Accumulate every (M block, N block) pair. +// 297 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +ss << __str_479; +// 298 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 299 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +ss << __str_480; +// 300 | #endif +} +// 301 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 302 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +ss << __str_481; +// 303 | #endif +} +// 304 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 305 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +ss << __str_482; +// 306 | #endif +} +// 307 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 308 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +ss << __str_483; +// 309 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 310 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +ss << __str_484; +// 311 | #endif +} +// 312 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 313 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +ss << __str_485; +// 314 | #endif +} +// 315 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 316 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +ss << __str_486; +// 317 | #endif +} +// 318 | #endif +} +// 319 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 320 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +ss << __str_487; +// 321 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 322 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +ss << __str_488; +// 323 | #endif +} +// 324 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 325 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +ss << __str_489; +// 326 | #endif +} +// 327 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 328 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +ss << __str_490; +// 329 | #endif +} +// 330 | #endif +} +// 331 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 332 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +ss << __str_491; +// 333 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 334 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +ss << __str_492; +// 335 | #endif +} +// 336 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 337 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +ss << __str_493; +// 338 | #endif +} +// 339 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 340 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +ss << __str_494; +// 341 | #endif +} +// 342 | #endif +} +// 343 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 344 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +ss << __str_495; +// 345 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 346 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +ss << __str_496; +// 347 | #endif +} +// 348 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 349 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +ss << __str_497; +// 350 | #endif +} +// 351 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 352 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +ss << __str_498; +// 353 | #endif +} +// 354 | #endif +} +// 355 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 356 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +ss << __str_499; +// 357 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 358 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +ss << __str_500; +// 359 | #endif +} +// 360 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 361 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +ss << __str_501; +// 362 | #endif +} +// 363 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 364 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +ss << __str_502; +// 365 | #endif +} +// 366 | #endif +} +// 367 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 368 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +ss << __str_503; +// 369 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 370 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +ss << __str_504; +// 371 | #endif +} +// 372 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 373 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +ss << __str_505; +// 374 | #endif +} +// 375 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 376 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +ss << __str_506; +// 377 | #endif +} +// 378 | #endif +} +// 379 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 380 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +ss << __str_507; +// 381 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 382 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +ss << __str_508; +// 383 | #endif +} +// 384 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 385 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +ss << __str_509; +// 386 | #endif +} +// 387 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 388 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +ss << __str_510; +// 389 | #endif +} +// 390 | #endif +} +// 391 | } +ss << __str_204; +// 392 | // Ensure every subgroup is done reading a_scratch before it is reused for +ss << __str_5; +// 393 | // the next K-block step. +// 394 | workgroupBarrier(); +ss << __str_451; +// 395 | } +ss << __str_20; +// 396 | +ss << __str_5; +// 397 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; +// 398 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). +// 399 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +ss << __str_165; +// 400 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 401 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +ss << __str_166; +// 402 | #endif +} +// 403 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 404 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +ss << __str_167; +// 405 | #endif +} +// 406 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 407 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +ss << __str_168; +// 408 | #endif +} +// 409 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 410 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +ss << __str_169; +// 411 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 412 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +ss << __str_170; +// 413 | #endif +} +// 414 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 415 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +ss << __str_171; +// 416 | #endif +} +// 417 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 418 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +ss << __str_172; +// 419 | #endif +} +// 420 | #endif +} +// 421 | #if sg_mat_count_m >= 3 +if (__param_sg_mat_count_m >= 3) { +// 422 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +ss << __str_173; +// 423 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 424 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +ss << __str_174; +// 425 | #endif +} +// 426 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 427 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +ss << __str_175; +// 428 | #endif +} +// 429 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 430 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +ss << __str_176; +// 431 | #endif +} +// 432 | #endif +} +// 433 | #if sg_mat_count_m >= 4 +if (__param_sg_mat_count_m >= 4) { +// 434 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +ss << __str_177; +// 435 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 436 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +ss << __str_178; +// 437 | #endif +} +// 438 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 439 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +ss << __str_179; +// 440 | #endif +} +// 441 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 442 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +ss << __str_180; +// 443 | #endif +} +// 444 | #endif +} +// 445 | #if sg_mat_count_m >= 5 +if (__param_sg_mat_count_m >= 5) { +// 446 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +ss << __str_181; +// 447 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 448 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +ss << __str_182; +// 449 | #endif +} +// 450 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 451 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +ss << __str_183; +// 452 | #endif +} +// 453 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 454 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +ss << __str_184; +// 455 | #endif +} +// 456 | #endif +} +// 457 | #if sg_mat_count_m >= 6 +if (__param_sg_mat_count_m >= 6) { +// 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +ss << __str_185; +// 459 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 460 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +ss << __str_186; +// 461 | #endif +} +// 462 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 463 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +ss << __str_187; +// 464 | #endif +} +// 465 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 466 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +ss << __str_188; +// 467 | #endif +} +// 468 | #endif +} +// 469 | #if sg_mat_count_m >= 7 +if (__param_sg_mat_count_m >= 7) { +// 470 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +ss << __str_189; +// 471 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 472 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +ss << __str_190; +// 473 | #endif +} +// 474 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 475 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +ss << __str_191; +// 476 | #endif +} +// 477 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 478 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +ss << __str_192; +// 479 | #endif +} +// 480 | #endif +} +// 481 | #if sg_mat_count_m >= 8 +if (__param_sg_mat_count_m >= 8) { +// 482 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +ss << __str_193; +// 483 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 484 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +ss << __str_194; +// 485 | #endif +} +// 486 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 487 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +ss << __str_195; +// 488 | #endif +} +// 489 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 490 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +ss << __str_196; +// 491 | #endif +} +// 492 | #endif +} +// 493 | +// 494 | #if split_k >= 2 +if (__param_split_k >= 2) { +// 495 | // Make all subgroups' partial tiles visible before the reduction. +// 496 | workgroupBarrier(); +ss << __str_26; +// 497 | +ss << __str_5; +// 498 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each +// 499 | // position is only overwritten by the lane that already read it. +// 500 | if (sg_index == 0u) { +ss << __str_197; +// 501 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +ss << __str_198; +// 502 | var acc: f16 = f16(0); +ss << __str_199; +// 503 | for (var s: u32 = 0; s < kSplitK; s++) { +ss << __str_200; +// 504 | acc += scratch[s * kTileM * kTileN + idx]; +ss << __str_201; +// 505 | } +ss << __str_202; +// 506 | scratch[idx] = acc; +ss << __str_203; +// 507 | } +ss << __str_204; +// 508 | } +ss << __str_20; +// 509 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +ss << __str_5; +// 510 | workgroupBarrier(); +ss << __str_26; +// 511 | #endif +} +// 512 | +// 513 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip +// 514 | // and M's bounds check make non-multiple dimensions write each element once. +// 515 | let n_count = min(kTileN, uniforms.N - global_base_n); +ss << __str_206; +// 516 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +ss << __str_207; +// 517 | let global_m = m_base + r; +ss << __str_208; +// 518 | if (global_m < uniforms.M) { +ss << __str_511; +// 519 | let scratch_base = r * kTileN; +ss << __str_210; +// 520 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; +ss << __str_512; +// 521 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +ss << __str_212; +// 522 | write_output(out_base + i, global_base_n + i, +ss << __str_235; +// 523 | output_value_t(scratch[scratch_base + i])); +ss << __str_236; +// 524 | } +ss << __str_202; +// 525 | } +ss << __str_204; +// 526 | } +ss << __str_20; +// 527 | } // MAIN +MainFunctionEnd(); +ss << __str_5; +// 528 | + + + 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..2b2638ef80532 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_514; // 7 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_408; +ss << __str_515; // 8 | let offset = co * uniforms.I * uniforms.H * uniforms.W + -ss << __str_409; +ss << __str_516; // 9 | ci * uniforms.H * uniforms.W + -ss << __str_410; +ss << __str_517; // 10 | h_w; -ss << __str_411; +ss << __str_518; // 11 | return src.getByOffset(offset); ss << __str_351; -ss << __var_src->GetByOffset(__str_406); +ss << __var_src->GetByOffset(__str_513); ss << __str_29; // 12 | } ss << __str_11; // 13 | return src_element_t(); -ss << __str_412; +ss << __str_519; // 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_520; // 17 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_408; +ss << __str_515; // 18 | let offset = co * uniforms.H * uniforms.W * uniforms.I + -ss << __str_414; +ss << __str_521; // 19 | h_w * uniforms.I + -ss << __str_415; +ss << __str_522; // 20 | ci; -ss << __str_416; +ss << __str_523; // 21 | output.setByOffset(offset, value); ss << __str_28; -ss << __var_output->SetByOffset(__str_406, __str_322); +ss << __var_output->SetByOffset(__str_513, __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_524; // 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_525; // 29 | let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64; -ss << __str_419; +ss << __str_526; // 30 | ss << __str_5; // 31 | if (group_co >= uniforms.O || group_ci >= uniforms.I) { -ss << __str_420; +ss << __str_527; // 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_528; // 36 | // load ss << __str_5; // 37 | for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) { -ss << __str_422; +ss << __str_529; // 38 | let load_ci_idx = ci_idx + local_idx / 4; -ss << __str_423; +ss << __str_530; // 39 | let load_h_w_idx = local_idx % 4; -ss << __str_424; +ss << __str_531; // 40 | ss << __str_5; // 41 | data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co, -ss << __str_425; +ss << __str_532; // 42 | group_ci + load_ci_idx, -ss << __str_426; +ss << __str_533; // 43 | h_w_idx * 4 + load_h_w_idx); -ss << __str_427; +ss << __str_534; // 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_535; // 49 | let output_data = data_cache[local_h_w_idx][local_idx]; -ss << __str_429; +ss << __str_536; // 50 | write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data); -ss << __str_430; +ss << __str_537; // 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..52118960dc647 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_543; // 18 | #if is_float16 if (__param_is_float16) { // 19 | bitcast>(uniforms.constant_value)[0]; -ss << __str_437; +ss << __str_544; // 20 | #else } else { // 21 | bitcast(uniforms.constant_value); -ss << __str_438; +ss << __str_545; // 22 | #endif } // 23 | // 24 | #if dim_value_zero if (__param_dim_value_zero) { // 25 | output[global_idx] = constant_value; -ss << __str_439; +ss << __str_546; // 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_547; // 29 | var use_pad_value = false; -ss << __str_441; +ss << __str_548; // 30 | var in_coord = i32(0); -ss << __str_442; +ss << __str_549; // 31 | ss << __str_5; // 32 | for (var dim = 0; dim < output.rank && !use_pad_value; dim++) { -ss << __str_443; +ss << __str_550; ss << __var_output->Rank(); -ss << __str_444; +ss << __str_551; // 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_552; +ss << GetElementAt(__str_247, __str_538, __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_553; +ss << GetElementAt(__str_539, __str_538, __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_554; +ss << GetElementAt(__str_540, __str_538, __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_555; // 38 | use_pad_value = true; -ss << __str_449; +ss << __str_556; // 39 | #elif pad_mode == PAD_MODE_EDGE } else if (__param_pad_mode == 2) { // 40 | if (output_index < lower_pads) { -ss << __str_450; +ss << __str_557; // 41 | in_coord = 0; -ss << __str_451; +ss << __str_558; // 42 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_452; +ss << __str_559; // 43 | in_coord = data_shape - 1; -ss << __str_453; +ss << __str_560; // 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_555; // 46 | in_coord = output_index - lower_pads; -ss << __str_454; +ss << __str_561; // 47 | if (in_coord < 0) { -ss << __str_455; +ss << __str_562; // 48 | in_coord = -in_coord; -ss << __str_456; +ss << __str_563; // 49 | } ss << __str_292; // 50 | let _2n_1 = 2 * (data_shape - 1); -ss << __str_457; +ss << __str_564; // 51 | in_coord = in_coord % _2n_1; -ss << __str_458; +ss << __str_565; // 52 | if (in_coord >= data_shape) { -ss << __str_459; +ss << __str_566; // 53 | in_coord = _2n_1 - in_coord; -ss << __str_460; +ss << __str_567; // 54 | } ss << __str_292; // 55 | #else // PAD_MODE_WRAP } else { // 56 | if (output_index < lower_pads) { -ss << __str_450; +ss << __str_557; // 57 | in_coord = data_shape + output_index - lower_pads; -ss << __str_461; +ss << __str_568; // 58 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_452; +ss << __str_559; // 59 | in_coord = output_index - data_shape - lower_pads; -ss << __str_462; +ss << __str_569; // 60 | #endif // pad_mode } // 61 | } else { -ss << __str_463; +ss << __str_570; // 62 | in_coord = output_index - lower_pads; -ss << __str_464; +ss << __str_571; // 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_572; // 67 | #endif } // 68 | // 69 | input_index += select(u32(in_coord) -ss << __str_466; +ss << __str_573; // 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_574; +ss << GetElementAt(__str_541, __str_538, __var_output->Rank() - 1); ss << __str_5; // 72 | #endif } // 73 | , u32(in_coord), dim == output.rank - 1); -ss << __str_468; +ss << __str_575; ss << __var_output->Rank(); -ss << __str_469; +ss << __str_576; // 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_542); 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..31b02837e4c3a 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,27 @@ 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_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..0aaaec60f7447 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,7 +31,7 @@ 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" // 678f18362f8d5afb8555278608d06770e427839a2727213dd26bcc53d8b0f486 // Include template implementations @@ -41,8 +41,9 @@ std::string pass_as_string(T&& v) { #include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // 4b28d0cbe9ca50020bbcb75121e787982c268992efdac401b821dee7c59c604e #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" // 1805ea3004b0ac1e3048cd71876a56470086f6eac0c24dcca1434fed14dd6ea8 +#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // 166e2a2db54c8477de7e164973ee31f79b40a2f1049b22770b9a1c7c96ec1907 +#include "wgsl_template_gen/generated/tensor/pad.h" // 7f732549c9b2533acacbd683d9625b55b3ce6fbe0159c509bfae46cadadb6e1c #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..814db34704fce 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,174 @@ 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 = global_base_n * uniforms.K + kb * kSgMatK;\n"; +constexpr const char* __str_453 = " var sg_mat_b0: subgroup_matrix_right Date: Tue, 27 Oct 2026 13:07:09 +0800 Subject: [PATCH 11/14] Address Jiajia's comments --- onnxruntime/core/providers/webgpu/nn/conv.cc | 108 ++++++++++-------- onnxruntime/core/providers/webgpu/nn/conv.h | 10 ++ .../webgpu/nn/subgroup_matrix_conv.cc | 102 +++++++---------- .../webgpu/nn/subgroup_matrix_conv.h | 5 +- ...subgroup_matrix_conv_8x16x16.wgsl.template | 36 ++++-- .../math/subgroup_matrix_tiling_selector.cc | 10 +- .../webgpu/subgroup_matrix_conv_test.cc | 37 +++++- 7 files changed, 188 insertions(+), 120 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index e655fe92a2671..e7d0b9837f8f8 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -28,6 +28,51 @@ Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const Tens return Transpose::DoTranspose(context, perm, reshaped_kernel, *transposed_kernel); } +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)); + if (local_pads.empty()) { + local_pads.resize(kernel_spacial_shape_vector.size() * 2, 0); + } + if (local_dilations.empty()) { + local_dilations.resize(kernel_spacial_shape_vector.size(), 1); + } + if (local_strides.empty()) { + local_strides.resize(kernel_spacial_shape_vector.size(), 1); + } + TensorShapeVector input_shape_vector = input_shape.AsShapeVector(); + 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)); + 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); + } + 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( @@ -68,40 +113,16 @@ Status Conv::ComputeInternal(ComputeContext& context 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()); - TensorShapeVector kernel_spacial_shape_vector; - ORT_RETURN_IF_ERROR(conv_attrs_.ComputeKernelShape(kernel_shape, kernel_spacial_shape_vector, false)); - if (local_pads.empty()) { - local_pads.resize(kernel_spacial_shape_vector.size() * 2, 0); - } - if (local_dilations.empty()) { - local_dilations.resize(kernel_spacial_shape_vector.size(), 1); - } - if (local_strides.empty()) { - local_strides.resize(kernel_spacial_shape_vector.size(), 1); - } - TensorShapeVector input_shape_vector = input_shape.AsShapeVector(); - 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)); - 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 strides; std::vector dilations; - auto transform_dim = [](int64_t dim) { return static_cast(dim); }; - 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); + 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) { @@ -147,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); @@ -344,17 +366,8 @@ Status Conv::PrePackInternal(ComputeContextBase& con return Status::OK(); } - // Subgroup-matrix Conv needs the weight in OIHW, so skip prepacking for it. Must - // come before Im2ColMatMul below: prepacking there first would leave - // prepacked_kernel_ non-null, making SubgroupMatrixConvImpl::Compute silently - // decline on every constant-weight Conv. Safe even if Compute() later declines - // for another reason: is_packed stays false, so input 1 stays OIHW. - if (CanApplySubgroupMatrixConv(context, is_channels_last, kernel_shape, - onnxruntime::narrow(conv_attrs_.group), tensor.DataType())) { - return Status::OK(); - } - // 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, @@ -364,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), @@ -399,7 +412,8 @@ Status Conv::PrePackInternal(ComputeContextBase& con // Analyze execution paths in ComputeInternal to determine if kernel transpose is needed: // - // 0. Subgroup-matrix Conv path: handled above (not prepacked -- stays OIHW) + // 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 41593c77ac43d..713661cc072f6 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -5,6 +5,7 @@ #include #include +#include #include "core/common/optional.h" #include "core/providers/webgpu/webgpu_kernel.h" @@ -66,6 +67,8 @@ class Conv : public WebGpuKernel { 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_; @@ -102,5 +105,12 @@ class Conv : public WebGpuKernel { 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 index 86b01379da956..d6a0ea418d0a5 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -1,9 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include #include -#include #include #include #include @@ -56,8 +54,9 @@ std::optional GetSupportedConfig(const ComputeContextBase& 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) { + 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), @@ -88,20 +87,27 @@ class SubgroupMatrixConvImpl final : public Conv::Co const ConvAttributes& conv_attrs = parent.ConvAttrs(); const Activation& activation = parent.ConvActivation(); - // Fast decline for any Conv this path doesn't accept: PrePackInternal only skips - // prepacking for problems CanApplySubgroupMatrixConv accepts, so a non-null - // prepacked_kernel_ here (grouped/1x1/etc. convs prepacked for another path) - // means this one would decline below anyway -- this just avoids the shape - // inference to get there. The kernel needs OIHW regardless. - if (parent.PrepackedKernel() != nullptr) { - return Status::OK(); - } + // 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); - const auto* weight = context.Input(1); if (input == nullptr || weight == nullptr) { return Status::OK(); } - const TensorShape& weight_shape = weight->Shape(); // OIHW: [Cout, Cin, kh, kw] + 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(); @@ -113,38 +119,12 @@ class SubgroupMatrixConvImpl final : public Conv::Co if (input_shape.NumDimensions() != 4) { return Status::OK(); } - 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(weight_shape, kernel_spacial_shape_vector, false)); - if (local_pads.empty()) { - local_pads.resize(kernel_spacial_shape_vector.size() * 2, 0); - } - if (local_dilations.empty()) { - local_dilations.resize(kernel_spacial_shape_vector.size(), 1); - } - if (local_strides.empty()) { - local_strides.resize(kernel_spacial_shape_vector.size(), 1); - } - TensorShapeVector input_shape_vector = input_shape.AsShapeVector(); - TensorShapeVector output_shape_vector = {input_shape[0]}; - TensorShape src_spacial_shape = TensorShape( - TensorShapeVector(std::next(input_shape_vector.begin()), std::prev(input_shape_vector.end()))); - ORT_RETURN_IF_ERROR(conv_attrs.InferPadsAndOutputShape(src_spacial_shape, kernel_spacial_shape_vector, - local_strides, local_dilations, local_pads, - output_shape_vector)); - output_shape_vector.push_back(weight_shape[0]); // Cout - const TensorShape output_shape(output_shape_vector); // NHWC: [batch, out_h, out_w, Cout] - - // static_cast, not narrow: an input the normal Conv path accepts must not abort here. std::vector pads; std::vector strides; std::vector dilations; - auto transform_dim = [](int64_t dim) { return static_cast(dim); }; - 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); + 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(); } @@ -173,7 +153,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co return Status::OK(); } - // ClampTileNToProblem already guarantees tile_n <= N; mirrored here as the last + // 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(); @@ -181,21 +161,22 @@ class SubgroupMatrixConvImpl final : public Conv::Co Tensor* output = context.Output(0, output_shape); - // Transpose weight OIHW -> OHWI = [Cout, kh, kw, Cin] (B column-major, B[k,n] at - // n*K+k). Constant weights are cached; non-constant ones are redone every Run. - // - // TODO: prepack the constant case too (PrePackIm2ColMatMulWeight already produces - // this layout) -- blocked on PrePackInternal committing from the weight alone - // while Compute can still decline at runtime onto the im2col path. + // B operand: OHWI or HWIO (see b_is_hwio and the WGSL template). An already + // prepacked weight is used as-is; otherwise this transposes to OHWI itself, + // caching the result for constant weights. Tensor run_local_ohwi_weight; - const Tensor* ohwi_weight = nullptr; - if (parent.IsWeightConstant()) { + 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); + } else if (parent.IsWeightConstant()) { ORT_RETURN_IF_ERROR(EnsureTransposedKernel(context, weight, weight_shape)); - ohwi_weight = cached_ohwi_weight_.get(); + b_weight = cached_ohwi_weight_.get(); } else { ORT_RETURN_IF_ERROR( TransposeKernel(context, weight, weight_shape, &run_local_ohwi_weight, {0, 2, 3, 1})); - ohwi_weight = &run_local_ohwi_weight; + b_weight = &run_local_ohwi_weight; } const auto& config = config_; @@ -221,17 +202,17 @@ class SubgroupMatrixConvImpl final : public Conv::Co 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, - activation}; + 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, + 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)}, - {ohwi_weight, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + {b_weight, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddOutput({output, ProgramTensorMetadataDependency::Rank, output_shape, 1}) .AddUniformVariables({{M}, {N}, @@ -252,6 +233,9 @@ class SubgroupMatrixConvImpl final : public Conv::Co if (has_bias) { program.AddInput({bias, ProgramTensorMetadataDependency::None}); } + LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv dispatching: M=" << M << " N=" << N << " 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; @@ -310,7 +294,8 @@ Status SubgroupMatrixConvProgram::GenerateShaderCode(ShaderHelper& shader) const 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_); + 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."); @@ -390,7 +375,8 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, return false; } - LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv ACCEPTED: kernel_shape[OIHW]=" << kernel_shape.ToString() + 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; diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h index 5cb530f6289c4..6a38f20b606da 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -25,7 +25,7 @@ class SubgroupMatrixConvProgram final : public Program 1; partials summed at @@ -28,6 +29,7 @@ #param sg_mat_count_n #param split_k #param vec_size +#param b_is_hwio #use .getByOffset .setByOffset @@ -227,11 +229,30 @@ $MAIN { workgroupBarrier(); if (valid) { - // Load the B right tiles for this K block. The weight is OHWI, i.e. B is - // stored column-major ([N, K] row-major, element B[k, n] at n*K + k), so - // the tile for N-block j (columns global_base_n + j*sg_mat_n) starts at - // (global_base_n + j*sg_mat_n)*K + kb*sg_mat_k and is loaded col-major - // with stride K. + // Load the B right tiles for this K block. +#if b_is_hwio + // HWIO: B is row-major ([K, N], B[k, n] at k*N + n); load row-major, stride N. + let b_base = kb * kSgMatK * uniforms.N + global_base_n; + var sg_mat_b0: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &input_b, b_base + 0 * kSgMatN, uniforms.N); +#if sg_mat_count_n >= 2 + var sg_mat_b1: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &input_b, b_base + 1 * kSgMatN, uniforms.N); +#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); +#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); +#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>( @@ -250,6 +271,7 @@ $MAIN { 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, 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 750c9fa38de3c..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 @@ -154,7 +154,7 @@ SubgroupMatrixTiling ConvHeuristicTiling(uint32_t N, uint32_t K) { // 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 ClampTileNToProblem(SubgroupMatrixTiling& t, uint32_t N) { +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; @@ -195,7 +195,7 @@ std::optional LookupPretunedTiling(std::string_view arch, // 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 ClampSplitKForBatch(SubgroupMatrixTiling& t, uint32_t M, uint32_t N, uint32_t batch, uint32_t hw) { +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) { @@ -209,7 +209,7 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, SubgroupMatrixTiling tiling = *tuned; // 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; } @@ -219,7 +219,7 @@ SubgroupMatrixTiling SelectTiling(std::string_view arch, uint32_t M, uint32_t N, // Conv counterpart of SelectTiling: no pretuned table, always ConvHeuristicTiling. SubgroupMatrixTiling SelectConvTiling(uint32_t N, uint32_t K) { SubgroupMatrixTiling tiling = ConvHeuristicTiling(N, K); - ClampTileNToProblem(tiling, N); + CapTileNToAlignedN(tiling, N); return tiling; } @@ -248,7 +248,7 @@ SubgroupMatrixTilingSelector CreateSubgroupMatrixConvTilingSelector( } 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 (ClampTileNToProblem can't + // 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; diff --git a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc index 570f4974fc41d..db873c53ec68b 100644 --- a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc +++ b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc @@ -108,7 +108,8 @@ static void RunConvTest(int64_t N, int64_t Cin, int64_t H, int64_t W, const std::string& activation = "", const std::vector& activation_params = {}, float abs_err = 0.06f, float rel_err = 0.02f, - bool weight_is_initializer = true) { + 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."; @@ -156,7 +157,13 @@ static void RunConvTest(int64_t N, int64_t Cin, int64_t H, int64_t W, } test->AddAttribute("group", static_cast(1)); test->AddAttribute("kernel_shape", std::vector{kh, kw}); - test->AddAttribute("pads", std::vector{pad, pad, pad, pad}); + 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}); @@ -268,6 +275,24 @@ TEST(SubgroupMatrixConv, Batched) { /*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 @@ -320,6 +345,14 @@ TEST(SubgroupMatrixConv, FusedActivationWithDeclaration) { } } +// HardSwish isn't supported by Im2ColMatMul, so PrePackInternal falls through to the +// general HWIO transpose, and Compute() must read that layout directly too. +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"); +} + // 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) { From 5ce227f3a704fda518d8c07a3a3aeee53faf4c1c Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Wed, 28 Oct 2026 13:24:16 +0800 Subject: [PATCH 12/14] Pad b for hwio layout --- .../webgpu/math/subgroup_matrix_matmul.cc | 59 +++--------------- ...subgroup_matrix_matmul_pad_b.wgsl.template | 3 +- .../webgpu/math/subgroup_matrix_pad_b.cc | 62 +++++++++++++++++++ .../webgpu/math/subgroup_matrix_pad_b.h | 27 ++++++++ .../webgpu/nn/subgroup_matrix_conv.cc | 44 +++++++++++-- .../webgpu/nn/subgroup_matrix_conv.h | 1 + ...subgroup_matrix_conv_8x16x16.wgsl.template | 30 +++++---- .../webgpu/subgroup_matrix_conv_test.cc | 12 +++- 8 files changed, 169 insertions(+), 69 deletions(-) create mode 100644 onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.cc create mode 100644 onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.h 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/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index d6a0ea418d0a5..cff2eac9d1119 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -12,6 +12,7 @@ #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" @@ -159,17 +160,33 @@ class SubgroupMatrixConvImpl final : public Conv::Co 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; + 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; otherwise this transposes to OHWI itself, - // caching the result for constant weights. + // 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(); @@ -216,6 +233,7 @@ class SubgroupMatrixConvImpl final : public Conv::Co .AddOutput({output, ProgramTensorMetadataDependency::Rank, output_shape, 1}) .AddUniformVariables({{M}, {N}, + {n_b}, {K}, {dispatch_x}, {batch}, @@ -233,8 +251,8 @@ class SubgroupMatrixConvImpl final : public Conv::Co if (has_bias) { program.AddInput({bias, ProgramTensorMetadataDependency::None}); } - LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv dispatching: M=" << M << " N=" << N << " K=" << K - << " batch=" << batch << " tile_m=" << tile_m << " tile_n=" << tile_n + 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)); @@ -258,12 +276,30 @@ class SubgroupMatrixConvImpl final : public Conv::Co "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 diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h index 6a38f20b606da..2d05c270a7804 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h @@ -40,6 +40,7 @@ class SubgroupMatrixConvProgram final : public Program 1; partials summed at @@ -19,8 +22,8 @@ // // Bias/activation: write_output(), built by the host from GetActivationSnippet. // -// Host guarantees: K % sg_mat_k == 0; kTileN <= N so a trailing partial N tile can -// shift back in bounds. M is unconstrained (see load_im2col_vec). +// 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 @@ -89,7 +92,10 @@ $MAIN { 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). - let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); + // 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; @@ -231,25 +237,27 @@ $MAIN { if (valid) { // Load the B right tiles for this K block. #if b_is_hwio - // HWIO: B is row-major ([K, N], B[k, n] at k*N + n); load row-major, stride N. - let b_base = kb * kSgMatK * uniforms.N + global_base_n; + // 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); + &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); + &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); + &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); + &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. diff --git a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc index db873c53ec68b..e2909c55cd25e 100644 --- a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc +++ b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc @@ -346,11 +346,21 @@ TEST(SubgroupMatrixConv, FusedActivationWithDeclaration) { } // HardSwish isn't supported by Im2ColMatMul, so PrePackInternal falls through to the -// general HWIO transpose, and Compute() must read that layout directly too. +// 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"); + 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"); + 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 From e5fffbade6f7e4f363397286baa381c831b1cbdb Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Wed, 28 Oct 2026 13:33:09 +0800 Subject: [PATCH 13/14] Regenerate WGSL template goldens for the HWIO B-padding change --- .../math/subgroup_matrix_matmul_pad_b.h | 33 +- .../nn/subgroup_matrix_conv_8x16x16.h | 1123 ++++++++------- .../in_tree_golden/static-cpp-literal/index.h | 1 + .../static-cpp-literal/index_impl.h | 4 +- .../math/subgroup_matrix_matmul_pad_b.h | 33 +- .../nn/subgroup_matrix_conv_8x16x16.h | 1245 +++++++++-------- .../generated/tensor/oihw_to_ohwi.h | 54 +- .../static-cpp/generated/tensor/pad.h | 84 +- .../test/in_tree_golden/static-cpp/index.h | 1 + .../in_tree_golden/static-cpp/index_impl.h | 10 +- .../in_tree_golden/static-cpp/string_table.h | 245 ++-- 11 files changed, 1506 insertions(+), 1327 deletions(-) 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 index d0eeaeaff7a1d..e6b773c7bf389 100644 --- 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 @@ -10,6 +10,7 @@ Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelp [[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; @@ -31,1226 +32,1310 @@ Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelp // 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 | // the OHWI weight. -// 12 | // B: weight pre-transposed to OHWI, column-major (B[k,n] at n*K+k). -// 13 | // -// 14 | // Workgroup: split_k subgroups x 32 lanes, tile kTileM x kTileN (host-selected). K -// 15 | // blocks round-robin across subgroups when split_k > 1; partials summed at -// 16 | // write-out. Staging A forces the K reduction into a per-step -// 17 | // gather/barrier/multiply-accumulate loop in uniform control flow. -// 18 | // -// 19 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. -// 20 | // -// 21 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N so a trailing partial N tile can -// 22 | // shift back in bounds. M is unconstrained (see load_im2col_vec). -// 23 | -// 24 | #param sg_mat_k -// 25 | #param sg_mat_m -// 26 | #param sg_mat_n -// 27 | #param sg_mat_count_m -// 28 | #param sg_mat_count_n -// 29 | #param split_k -// 30 | #param vec_size -// 31 | -// 32 | #use .getByOffset .setByOffset -// 33 | -// 34 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +// 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 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N_b so a trailing partial N tile +// 26 | // can shift back in bounds. M is unconstrained (see load_im2col_vec). +// 27 | +// 28 | #param sg_mat_k +// 29 | #param sg_mat_m +// 30 | #param sg_mat_n +// 31 | #param sg_mat_count_m +// 32 | #param sg_mat_count_n +// 33 | #param split_k +// 34 | #param vec_size +// 35 | #param b_is_hwio +// 36 | +// 37 | #use .getByOffset .setByOffset +// 38 | +// 39 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) ss << "const kSgMatM: u32 = u32("; ss << __param_sg_mat_m; ss << ");\n"; -// 35 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +// 40 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) ss << "const kSgMatN: u32 = u32("; ss << __param_sg_mat_n; ss << ");\n"; -// 36 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +// 41 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) ss << "const kSgMatK: u32 = u32("; ss << __param_sg_mat_k; ss << ");\n"; -// 37 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +// 42 | 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"; -// 38 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +// 43 | 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"; -// 39 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +// 44 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) ss << "const kTileM: u32 = kSgMatM * kSgMatCountM;\n"; -// 40 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +// 45 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) ss << "const kTileN: u32 = kSgMatN * kSgMatCountN;\n"; -// 41 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +// 46 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K ss << "const kSplitK: u32 = u32("; ss << __param_split_k; ss << ");\n"; -// 42 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +// 47 | const kSubgroupSize: u32 = 32; // Lanes per subgroup ss << "const kSubgroupSize: u32 = 32;\n"; -// 43 | const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (cin % vec_size == 0) +// 48 | 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"; -// 44 | +// 49 | ss << "\n"; -// 45 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K -// 46 | // subgroup, summed at write-back. -// 47 | var scratch: array; +// 50 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// 51 | // subgroup, summed at write-back. +// 52 | var scratch: array; ss << "var scratch: array;\n"; -// 48 | +// 53 | ss << "\n"; -// 49 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. -// 50 | var a_scratch: array; +// 54 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 55 | var a_scratch: array; ss << "var a_scratch: array;\n"; -// 51 | +// 56 | ss << "\n"; -// 52 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always -// 53 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding -// 54 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. -// 55 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. -// 56 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { +// 57 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always +// 58 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding +// 59 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. +// 60 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. +// 61 | 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"; -// 57 | let cin_vec = uniforms.cin / kVecSize; +// 62 | let cin_vec = uniforms.cin / kVecSize; ss << " let cin_vec = uniforms.cin / kVecSize;\n"; -// 58 | let oh = m / uniforms.output_w; +// 63 | let oh = m / uniforms.output_w; ss << " let oh = m / uniforms.output_w;\n"; -// 59 | let ow = m % uniforms.output_w; +// 64 | let ow = m % uniforms.output_w; ss << " let ow = m % uniforms.output_w;\n"; -// 60 | let k_cin = k_vec % cin_vec; +// 65 | let k_cin = k_vec % cin_vec; ss << " let k_cin = k_vec % cin_vec;\n"; -// 61 | let k_hw = k_vec / cin_vec; +// 66 | let k_hw = k_vec / cin_vec; ss << " let k_hw = k_vec / cin_vec;\n"; -// 62 | let kw_idx = k_hw % uniforms.kernel_w; +// 67 | let kw_idx = k_hw % uniforms.kernel_w; ss << " let kw_idx = k_hw % uniforms.kernel_w;\n"; -// 63 | let kh_idx = k_hw / uniforms.kernel_w; +// 68 | let kh_idx = k_hw / uniforms.kernel_w; ss << " let kh_idx = k_hw / uniforms.kernel_w;\n"; -// 64 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); +// 69 | 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"; -// 65 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); +// 70 | 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"; -// 66 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { +// 71 | 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"; -// 67 | return input_src_value_t(); +// 72 | return input_src_value_t(); ss << " return input_src_value_t();\n"; -// 68 | } +// 73 | } ss << " }\n"; -// 69 | let input_vec_idx = input_batch_base_vec + +// 74 | let input_vec_idx = input_batch_base_vec + ss << " let input_vec_idx = input_batch_base_vec +\n"; -// 70 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +// 75 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; ss << " (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin;\n"; -// 71 | return input_src.getByOffset(input_vec_idx); +// 76 | return input_src.getByOffset(input_vec_idx); ss << " return "; ss << __var_input_src->GetByOffset("input_vec_idx"); ss << ";\n"; -// 72 | } +// 77 | } ss << "}\n"; -// 73 | +// 78 | ss << "\n"; -// 74 | $MAIN { +// 79 | $MAIN { MainFunctionStart(); ss << "\n"; -// 75 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +// 80 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is ss << "\n"; -// 76 | // the true linear index used to recover batch_id and (n_tile, m_tile). -// 77 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +// 81 | // the true linear index used to recover batch_id and (n_tile, m_tile). +// 82 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; ss << " let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM;\n"; -// 78 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +// 83 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; ss << " let tiles_per_batch = uniforms.num_n_tile * num_m_tile;\n"; -// 79 | let batch_id = workgroup_idx / tiles_per_batch; +// 84 | let batch_id = workgroup_idx / tiles_per_batch; ss << " let batch_id = workgroup_idx / tiles_per_batch;\n"; -// 80 | // Dispatch normalization can add extra workgroups past the last batch element; +// 85 | // Dispatch normalization can add extra workgroups past the last batch element; ss << "\n"; -// 81 | // drop them here (workgroup-uniform, so the barriers below stay uniform). -// 82 | if (batch_id >= uniforms.batch) { +// 86 | // drop them here (workgroup-uniform, so the barriers below stay uniform). +// 87 | if (batch_id >= uniforms.batch) { ss << " if (batch_id >= uniforms.batch) {\n"; -// 83 | return; +// 88 | return; ss << " return;\n"; -// 84 | } +// 89 | } ss << " }\n"; -// 85 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +// 90 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element ss << " let tile_idx = workgroup_idx % tiles_per_batch;\n"; -// 86 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile +// 91 | 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"; -// 87 | let n_tile_start = n_tile * kTileN; +// 92 | let n_tile_start = n_tile * kTileN; ss << " let n_tile_start = n_tile * kTileN;\n"; -// 88 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +// 93 | // Shift a trailing partial N tile back so unchecked B loads stay inside the ss << "\n"; -// 89 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). -// 90 | let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); -ss << " let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN);\n"; -// 91 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile +// 94 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). +// 95 | // Uses N_b (== N unless B was padded for row-major alignment) so the shifted +// 96 | // origin -- itself n_align-even since N_b, kTileN and n_tile_start all are -- +// 97 | // stays alignment-safe for the HWIO load below. +// 98 | 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"; +// 99 | 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"; -// 92 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile +// 100 | 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"; -// 93 | let m_base = m_tile * kTileM; +// 101 | let m_base = m_tile * kTileM; ss << " let m_base = m_tile * kTileM;\n"; -// 94 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); +// 102 | 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"; -// 95 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +// 103 | let out_batch_base = batch_id * uniforms.M * uniforms.N; ss << " let out_batch_base = batch_id * uniforms.M * uniforms.N;\n"; -// 96 | let k_blocks = uniforms.K / kSgMatK; +// 104 | let k_blocks = uniforms.K / kSgMatK; ss << " let k_blocks = uniforms.K / kSgMatK;\n"; -// 97 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +// 105 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) ss << " let sg_index = local_idx / kSubgroupSize;\n"; -// 98 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +// 106 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) ss << " let sg_lane = local_idx % kSubgroupSize;\n"; -// 99 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +// 107 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot ss << " let sg_scratch_base = sg_index * kTileM * kTileN;\n"; -// 100 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +// 108 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot ss << " let a_scratch_base = sg_index * kTileM * kSgMatK;\n"; -// 101 | +// 109 | ss << "\n"; -// 102 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). -// 103 | var sg_mat_c0_0: subgroup_matrix_result; +// 110 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 111 | var sg_mat_c0_0: subgroup_matrix_result; ss << " var sg_mat_c0_0: subgroup_matrix_result;\n"; -// 104 | #if sg_mat_count_n >= 2 +// 112 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 105 | var sg_mat_c0_1: subgroup_matrix_result; +// 113 | var sg_mat_c0_1: subgroup_matrix_result; ss << " var sg_mat_c0_1: subgroup_matrix_result;\n"; -// 106 | #endif +// 114 | #endif } -// 107 | #if sg_mat_count_n >= 3 +// 115 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 108 | var sg_mat_c0_2: subgroup_matrix_result; +// 116 | var sg_mat_c0_2: subgroup_matrix_result; ss << " var sg_mat_c0_2: subgroup_matrix_result;\n"; -// 109 | #endif +// 117 | #endif } -// 110 | #if sg_mat_count_n >= 4 +// 118 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 111 | var sg_mat_c0_3: subgroup_matrix_result; +// 119 | var sg_mat_c0_3: subgroup_matrix_result; ss << " var sg_mat_c0_3: subgroup_matrix_result;\n"; -// 112 | #endif +// 120 | #endif } -// 113 | #if sg_mat_count_m >= 2 +// 121 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 114 | var sg_mat_c1_0: subgroup_matrix_result; +// 122 | var sg_mat_c1_0: subgroup_matrix_result; ss << " var sg_mat_c1_0: subgroup_matrix_result;\n"; -// 115 | #if sg_mat_count_n >= 2 +// 123 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 116 | var sg_mat_c1_1: subgroup_matrix_result; +// 124 | var sg_mat_c1_1: subgroup_matrix_result; ss << " var sg_mat_c1_1: subgroup_matrix_result;\n"; -// 117 | #endif +// 125 | #endif } -// 118 | #if sg_mat_count_n >= 3 +// 126 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 119 | var sg_mat_c1_2: subgroup_matrix_result; +// 127 | var sg_mat_c1_2: subgroup_matrix_result; ss << " var sg_mat_c1_2: subgroup_matrix_result;\n"; -// 120 | #endif +// 128 | #endif } -// 121 | #if sg_mat_count_n >= 4 +// 129 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 122 | var sg_mat_c1_3: subgroup_matrix_result; +// 130 | var sg_mat_c1_3: subgroup_matrix_result; ss << " var sg_mat_c1_3: subgroup_matrix_result;\n"; -// 123 | #endif +// 131 | #endif } -// 124 | #endif +// 132 | #endif } -// 125 | #if sg_mat_count_m >= 3 +// 133 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 126 | var sg_mat_c2_0: subgroup_matrix_result; +// 134 | var sg_mat_c2_0: subgroup_matrix_result; ss << " var sg_mat_c2_0: subgroup_matrix_result;\n"; -// 127 | #if sg_mat_count_n >= 2 +// 135 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 128 | var sg_mat_c2_1: subgroup_matrix_result; +// 136 | var sg_mat_c2_1: subgroup_matrix_result; ss << " var sg_mat_c2_1: subgroup_matrix_result;\n"; -// 129 | #endif +// 137 | #endif } -// 130 | #if sg_mat_count_n >= 3 +// 138 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 131 | var sg_mat_c2_2: subgroup_matrix_result; +// 139 | var sg_mat_c2_2: subgroup_matrix_result; ss << " var sg_mat_c2_2: subgroup_matrix_result;\n"; -// 132 | #endif +// 140 | #endif } -// 133 | #if sg_mat_count_n >= 4 +// 141 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 134 | var sg_mat_c2_3: subgroup_matrix_result; +// 142 | var sg_mat_c2_3: subgroup_matrix_result; ss << " var sg_mat_c2_3: subgroup_matrix_result;\n"; -// 135 | #endif +// 143 | #endif } -// 136 | #endif +// 144 | #endif } -// 137 | #if sg_mat_count_m >= 4 +// 145 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 138 | var sg_mat_c3_0: subgroup_matrix_result; +// 146 | var sg_mat_c3_0: subgroup_matrix_result; ss << " var sg_mat_c3_0: subgroup_matrix_result;\n"; -// 139 | #if sg_mat_count_n >= 2 +// 147 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 140 | var sg_mat_c3_1: subgroup_matrix_result; +// 148 | var sg_mat_c3_1: subgroup_matrix_result; ss << " var sg_mat_c3_1: subgroup_matrix_result;\n"; -// 141 | #endif +// 149 | #endif } -// 142 | #if sg_mat_count_n >= 3 +// 150 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 143 | var sg_mat_c3_2: subgroup_matrix_result; +// 151 | var sg_mat_c3_2: subgroup_matrix_result; ss << " var sg_mat_c3_2: subgroup_matrix_result;\n"; -// 144 | #endif +// 152 | #endif } -// 145 | #if sg_mat_count_n >= 4 +// 153 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 146 | var sg_mat_c3_3: subgroup_matrix_result; +// 154 | var sg_mat_c3_3: subgroup_matrix_result; ss << " var sg_mat_c3_3: subgroup_matrix_result;\n"; -// 147 | #endif +// 155 | #endif } -// 148 | #endif +// 156 | #endif } -// 149 | #if sg_mat_count_m >= 5 +// 157 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 150 | var sg_mat_c4_0: subgroup_matrix_result; +// 158 | var sg_mat_c4_0: subgroup_matrix_result; ss << " var sg_mat_c4_0: subgroup_matrix_result;\n"; -// 151 | #if sg_mat_count_n >= 2 +// 159 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 152 | var sg_mat_c4_1: subgroup_matrix_result; +// 160 | var sg_mat_c4_1: subgroup_matrix_result; ss << " var sg_mat_c4_1: subgroup_matrix_result;\n"; -// 153 | #endif +// 161 | #endif } -// 154 | #if sg_mat_count_n >= 3 +// 162 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 155 | var sg_mat_c4_2: subgroup_matrix_result; +// 163 | var sg_mat_c4_2: subgroup_matrix_result; ss << " var sg_mat_c4_2: subgroup_matrix_result;\n"; -// 156 | #endif +// 164 | #endif } -// 157 | #if sg_mat_count_n >= 4 +// 165 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 158 | var sg_mat_c4_3: subgroup_matrix_result; +// 166 | var sg_mat_c4_3: subgroup_matrix_result; ss << " var sg_mat_c4_3: subgroup_matrix_result;\n"; -// 159 | #endif +// 167 | #endif } -// 160 | #endif +// 168 | #endif } -// 161 | #if sg_mat_count_m >= 6 +// 169 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 162 | var sg_mat_c5_0: subgroup_matrix_result; +// 170 | var sg_mat_c5_0: subgroup_matrix_result; ss << " var sg_mat_c5_0: subgroup_matrix_result;\n"; -// 163 | #if sg_mat_count_n >= 2 +// 171 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 164 | var sg_mat_c5_1: subgroup_matrix_result; +// 172 | var sg_mat_c5_1: subgroup_matrix_result; ss << " var sg_mat_c5_1: subgroup_matrix_result;\n"; -// 165 | #endif +// 173 | #endif } -// 166 | #if sg_mat_count_n >= 3 +// 174 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 167 | var sg_mat_c5_2: subgroup_matrix_result; +// 175 | var sg_mat_c5_2: subgroup_matrix_result; ss << " var sg_mat_c5_2: subgroup_matrix_result;\n"; -// 168 | #endif +// 176 | #endif } -// 169 | #if sg_mat_count_n >= 4 +// 177 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 170 | var sg_mat_c5_3: subgroup_matrix_result; +// 178 | var sg_mat_c5_3: subgroup_matrix_result; ss << " var sg_mat_c5_3: subgroup_matrix_result;\n"; -// 171 | #endif +// 179 | #endif } -// 172 | #endif +// 180 | #endif } -// 173 | #if sg_mat_count_m >= 7 +// 181 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 174 | var sg_mat_c6_0: subgroup_matrix_result; +// 182 | var sg_mat_c6_0: subgroup_matrix_result; ss << " var sg_mat_c6_0: subgroup_matrix_result;\n"; -// 175 | #if sg_mat_count_n >= 2 +// 183 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 176 | var sg_mat_c6_1: subgroup_matrix_result; +// 184 | var sg_mat_c6_1: subgroup_matrix_result; ss << " var sg_mat_c6_1: subgroup_matrix_result;\n"; -// 177 | #endif +// 185 | #endif } -// 178 | #if sg_mat_count_n >= 3 +// 186 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 179 | var sg_mat_c6_2: subgroup_matrix_result; +// 187 | var sg_mat_c6_2: subgroup_matrix_result; ss << " var sg_mat_c6_2: subgroup_matrix_result;\n"; -// 180 | #endif +// 188 | #endif } -// 181 | #if sg_mat_count_n >= 4 +// 189 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 182 | var sg_mat_c6_3: subgroup_matrix_result; +// 190 | var sg_mat_c6_3: subgroup_matrix_result; ss << " var sg_mat_c6_3: subgroup_matrix_result;\n"; -// 183 | #endif +// 191 | #endif } -// 184 | #endif +// 192 | #endif } -// 185 | #if sg_mat_count_m >= 8 +// 193 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 186 | var sg_mat_c7_0: subgroup_matrix_result; +// 194 | var sg_mat_c7_0: subgroup_matrix_result; ss << " var sg_mat_c7_0: subgroup_matrix_result;\n"; -// 187 | #if sg_mat_count_n >= 2 +// 195 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 188 | var sg_mat_c7_1: subgroup_matrix_result; +// 196 | var sg_mat_c7_1: subgroup_matrix_result; ss << " var sg_mat_c7_1: subgroup_matrix_result;\n"; -// 189 | #endif +// 197 | #endif } -// 190 | #if sg_mat_count_n >= 3 +// 198 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 191 | var sg_mat_c7_2: subgroup_matrix_result; +// 199 | var sg_mat_c7_2: subgroup_matrix_result; ss << " var sg_mat_c7_2: subgroup_matrix_result;\n"; -// 192 | #endif +// 200 | #endif } -// 193 | #if sg_mat_count_n >= 4 +// 201 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 194 | var sg_mat_c7_3: subgroup_matrix_result; +// 202 | var sg_mat_c7_3: subgroup_matrix_result; ss << " var sg_mat_c7_3: subgroup_matrix_result;\n"; -// 195 | #endif +// 203 | #endif } -// 196 | #endif +// 204 | #endif } -// 197 | -// 198 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, -// 199 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks -// 200 | // early (valid == false), keeping the barriers below in uniform control flow. -// 201 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +// 205 | +// 206 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, +// 207 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks +// 208 | // early (valid == false), keeping the barriers below in uniform control flow. +// 209 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; ss << " let num_steps = (k_blocks + kSplitK - 1u) / kSplitK;\n"; -// 202 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +// 210 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { ss << " for (var step: u32 = 0u; step < num_steps; step = step + 1u) {\n"; -// 203 | let kb = sg_index + step * kSplitK; +// 211 | let kb = sg_index + step * kSplitK; ss << " let kb = sg_index + step * kSplitK;\n"; -// 204 | let valid = kb < k_blocks; +// 212 | let valid = kb < k_blocks; ss << " let valid = kb < k_blocks;\n"; -// 205 | +// 213 | ss << "\n"; -// 206 | if (valid) { +// 214 | if (valid) { ss << " if (valid) {\n"; -// 207 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +// 215 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; ss << " let k_vec_per_sgmat_row = kSgMatK / kVecSize;\n"; -// 208 | let k_vec_base = kb * k_vec_per_sgmat_row; +// 216 | let k_vec_base = kb * k_vec_per_sgmat_row; ss << " let k_vec_base = kb * k_vec_per_sgmat_row;\n"; -// 209 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { +// 217 | 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"; -// 210 | let row = vidx / k_vec_per_sgmat_row; +// 218 | let row = vidx / k_vec_per_sgmat_row; ss << " let row = vidx / k_vec_per_sgmat_row;\n"; -// 211 | let kc_vec = vidx % k_vec_per_sgmat_row; +// 219 | let kc_vec = vidx % k_vec_per_sgmat_row; ss << " let kc_vec = vidx % k_vec_per_sgmat_row;\n"; -// 212 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); +// 220 | 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"; -// 213 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +// 221 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; ss << " let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize;\n"; -// 214 | #if vec_size == 1 +// 222 | #if vec_size == 1 if (__param_vec_size == 1) { -// 215 | a_scratch[dst] = chunk; +// 223 | a_scratch[dst] = chunk; ss << " a_scratch[dst] = chunk;\n"; -// 216 | #else +// 224 | #else } else { -// 217 | a_scratch[dst + 0u] = chunk[0]; +// 225 | a_scratch[dst + 0u] = chunk[0]; ss << " a_scratch[dst + 0u] = chunk[0];\n"; -// 218 | a_scratch[dst + 1u] = chunk[1]; +// 226 | a_scratch[dst + 1u] = chunk[1]; ss << " a_scratch[dst + 1u] = chunk[1];\n"; -// 219 | #if vec_size >= 4 +// 227 | #if vec_size >= 4 if (__param_vec_size >= 4) { -// 220 | a_scratch[dst + 2u] = chunk[2]; +// 228 | a_scratch[dst + 2u] = chunk[2]; ss << " a_scratch[dst + 2u] = chunk[2];\n"; -// 221 | a_scratch[dst + 3u] = chunk[3]; +// 229 | a_scratch[dst + 3u] = chunk[3]; ss << " a_scratch[dst + 3u] = chunk[3];\n"; -// 222 | #endif +// 230 | #endif } -// 223 | #endif +// 231 | #endif } -// 224 | } +// 232 | } ss << " }\n"; -// 225 | } +// 233 | } ss << " }\n"; -// 226 | // Make the staged A-tiles visible before loading them as subgroup matrices. +// 234 | // Make the staged A-tiles visible before loading them as subgroup matrices. ss << "\n"; -// 227 | workgroupBarrier(); +// 235 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 228 | +// 236 | ss << "\n"; -// 229 | if (valid) { +// 237 | if (valid) { ss << " if (valid) {\n"; -// 230 | // Load the B right tiles for this K block. The weight is OHWI, i.e. B is +// 238 | // Load the B right tiles for this K block. ss << "\n"; -// 231 | // stored column-major ([N, K] row-major, element B[k, n] at n*K + k), so -// 232 | // the tile for N-block j (columns global_base_n + j*sg_mat_n) starts at -// 233 | // (global_base_n + j*sg_mat_n)*K + kb*sg_mat_k and is loaded col-major -// 234 | // with stride K. -// 235 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +// 239 | #if b_is_hwio +if (__param_b_is_hwio) { +// 240 | // HWIO: B is row-major ([K, N_b], B[k, n] at k*N_b + n); load row-major, +// 241 | // stride N_b. N_b (== N unless padded) keeps the row stride, and via +// 242 | // global_base_n above, the column offset, 4-byte aligned. +// 243 | let b_base = kb * kSgMatK * uniforms.N_b + global_base_n; +ss << " let b_base = kb * kSgMatK * uniforms.N_b + global_base_n;\n"; +// 244 | var sg_mat_b0: subgroup_matrix_right = +ss << " var sg_mat_b0: subgroup_matrix_right =\n"; +// 245 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 246 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 0 * kSgMatN, uniforms.N_b);\n"; +// 247 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 248 | var sg_mat_b1: subgroup_matrix_right = +ss << " var sg_mat_b1: subgroup_matrix_right =\n"; +// 249 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 250 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 1 * kSgMatN, uniforms.N_b);\n"; +// 251 | #endif +} +// 252 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 253 | var sg_mat_b2: subgroup_matrix_right = +ss << " var sg_mat_b2: subgroup_matrix_right =\n"; +// 254 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 255 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 2 * kSgMatN, uniforms.N_b);\n"; +// 256 | #endif +} +// 257 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 258 | var sg_mat_b3: subgroup_matrix_right = +ss << " var sg_mat_b3: subgroup_matrix_right =\n"; +// 259 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 260 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +ss << " &input_b, b_base + 3 * kSgMatN, uniforms.N_b);\n"; +// 261 | #endif +} +// 262 | #else +} else { +// 263 | // OHWI: B is column-major ([N, K], B[k, n] at n*K + k); load col-major, stride K. +// 264 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; ss << " let b_base = global_base_n * uniforms.K + kb * kSgMatK;\n"; -// 236 | var sg_mat_b0: subgroup_matrix_right = +// 265 | var sg_mat_b0: subgroup_matrix_right = ss << " var sg_mat_b0: subgroup_matrix_right =\n"; -// 237 | subgroupMatrixLoad, col_major>( +// 266 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 238 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +// 267 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 239 | #if sg_mat_count_n >= 2 +// 268 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 240 | var sg_mat_b1: subgroup_matrix_right = +// 269 | var sg_mat_b1: subgroup_matrix_right = ss << " var sg_mat_b1: subgroup_matrix_right =\n"; -// 241 | subgroupMatrixLoad, col_major>( +// 270 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 242 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +// 271 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 243 | #endif +// 272 | #endif } -// 244 | #if sg_mat_count_n >= 3 +// 273 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 245 | var sg_mat_b2: subgroup_matrix_right = +// 274 | var sg_mat_b2: subgroup_matrix_right = ss << " var sg_mat_b2: subgroup_matrix_right =\n"; -// 246 | subgroupMatrixLoad, col_major>( +// 275 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 247 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +// 276 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 248 | #endif +// 277 | #endif } -// 249 | #if sg_mat_count_n >= 4 +// 278 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 250 | var sg_mat_b3: subgroup_matrix_right = +// 279 | var sg_mat_b3: subgroup_matrix_right = ss << " var sg_mat_b3: subgroup_matrix_right =\n"; -// 251 | subgroupMatrixLoad, col_major>( +// 280 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 252 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +// 281 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 253 | #endif +// 282 | #endif +} +// 283 | #endif } -// 254 | -// 255 | // Load the A left tiles (one per M block) from workgroup staging, -// 256 | // row-major with stride kSgMatK. -// 257 | var sg_mat_a0: subgroup_matrix_left = +// 284 | +// 285 | // Load the A left tiles (one per M block) from workgroup staging, +// 286 | // row-major with stride kSgMatK. +// 287 | var sg_mat_a0: subgroup_matrix_left = ss << " var sg_mat_a0: subgroup_matrix_left =\n"; -// 258 | subgroupMatrixLoad, row_major>( +// 288 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 259 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +// 289 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 260 | #if sg_mat_count_m >= 2 +// 290 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 261 | var sg_mat_a1: subgroup_matrix_left = +// 291 | var sg_mat_a1: subgroup_matrix_left = ss << " var sg_mat_a1: subgroup_matrix_left =\n"; -// 262 | subgroupMatrixLoad, row_major>( +// 292 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 263 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +// 293 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 264 | #endif +// 294 | #endif } -// 265 | #if sg_mat_count_m >= 3 +// 295 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 266 | var sg_mat_a2: subgroup_matrix_left = +// 296 | var sg_mat_a2: subgroup_matrix_left = ss << " var sg_mat_a2: subgroup_matrix_left =\n"; -// 267 | subgroupMatrixLoad, row_major>( +// 297 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 268 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +// 298 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 269 | #endif +// 299 | #endif } -// 270 | #if sg_mat_count_m >= 4 +// 300 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 271 | var sg_mat_a3: subgroup_matrix_left = +// 301 | var sg_mat_a3: subgroup_matrix_left = ss << " var sg_mat_a3: subgroup_matrix_left =\n"; -// 272 | subgroupMatrixLoad, row_major>( +// 302 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 273 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +// 303 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 274 | #endif +// 304 | #endif } -// 275 | #if sg_mat_count_m >= 5 +// 305 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 276 | var sg_mat_a4: subgroup_matrix_left = +// 306 | var sg_mat_a4: subgroup_matrix_left = ss << " var sg_mat_a4: subgroup_matrix_left =\n"; -// 277 | subgroupMatrixLoad, row_major>( +// 307 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 278 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +// 308 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 279 | #endif +// 309 | #endif } -// 280 | #if sg_mat_count_m >= 6 +// 310 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 281 | var sg_mat_a5: subgroup_matrix_left = +// 311 | var sg_mat_a5: subgroup_matrix_left = ss << " var sg_mat_a5: subgroup_matrix_left =\n"; -// 282 | subgroupMatrixLoad, row_major>( +// 312 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 283 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +// 313 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 284 | #endif +// 314 | #endif } -// 285 | #if sg_mat_count_m >= 7 +// 315 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 286 | var sg_mat_a6: subgroup_matrix_left = +// 316 | var sg_mat_a6: subgroup_matrix_left = ss << " var sg_mat_a6: subgroup_matrix_left =\n"; -// 287 | subgroupMatrixLoad, row_major>( +// 317 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 288 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +// 318 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 289 | #endif +// 319 | #endif } -// 290 | #if sg_mat_count_m >= 8 +// 320 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 291 | var sg_mat_a7: subgroup_matrix_left = +// 321 | var sg_mat_a7: subgroup_matrix_left = ss << " var sg_mat_a7: subgroup_matrix_left =\n"; -// 292 | subgroupMatrixLoad, row_major>( +// 322 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 293 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +// 323 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 294 | #endif +// 324 | #endif } -// 295 | -// 296 | // Accumulate every (M block, N block) pair. -// 297 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +// 325 | +// 326 | // Accumulate every (M block, N block) pair. +// 327 | 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"; -// 298 | #if sg_mat_count_n >= 2 +// 328 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 299 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +// 329 | 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"; -// 300 | #endif +// 330 | #endif } -// 301 | #if sg_mat_count_n >= 3 +// 331 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 302 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +// 332 | 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"; -// 303 | #endif +// 333 | #endif } -// 304 | #if sg_mat_count_n >= 4 +// 334 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 305 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +// 335 | 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"; -// 306 | #endif +// 336 | #endif } -// 307 | #if sg_mat_count_m >= 2 +// 337 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 308 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +// 338 | 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"; -// 309 | #if sg_mat_count_n >= 2 +// 339 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 310 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +// 340 | 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"; -// 311 | #endif +// 341 | #endif } -// 312 | #if sg_mat_count_n >= 3 +// 342 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 313 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +// 343 | 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"; -// 314 | #endif +// 344 | #endif } -// 315 | #if sg_mat_count_n >= 4 +// 345 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 316 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +// 346 | 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"; -// 317 | #endif +// 347 | #endif } -// 318 | #endif +// 348 | #endif } -// 319 | #if sg_mat_count_m >= 3 +// 349 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 320 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +// 350 | 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"; -// 321 | #if sg_mat_count_n >= 2 +// 351 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 322 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +// 352 | 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"; -// 323 | #endif +// 353 | #endif } -// 324 | #if sg_mat_count_n >= 3 +// 354 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 325 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +// 355 | 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"; -// 326 | #endif +// 356 | #endif } -// 327 | #if sg_mat_count_n >= 4 +// 357 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 328 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +// 358 | 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"; -// 329 | #endif +// 359 | #endif } -// 330 | #endif +// 360 | #endif } -// 331 | #if sg_mat_count_m >= 4 +// 361 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 332 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +// 362 | 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"; -// 333 | #if sg_mat_count_n >= 2 +// 363 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 334 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +// 364 | 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"; -// 335 | #endif +// 365 | #endif } -// 336 | #if sg_mat_count_n >= 3 +// 366 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 337 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +// 367 | 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"; -// 338 | #endif +// 368 | #endif } -// 339 | #if sg_mat_count_n >= 4 +// 369 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 340 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +// 370 | 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"; -// 341 | #endif +// 371 | #endif } -// 342 | #endif +// 372 | #endif } -// 343 | #if sg_mat_count_m >= 5 +// 373 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 344 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +// 374 | 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"; -// 345 | #if sg_mat_count_n >= 2 +// 375 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 346 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +// 376 | 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"; -// 347 | #endif +// 377 | #endif } -// 348 | #if sg_mat_count_n >= 3 +// 378 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 349 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +// 379 | 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"; -// 350 | #endif +// 380 | #endif } -// 351 | #if sg_mat_count_n >= 4 +// 381 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 352 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +// 382 | 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"; -// 353 | #endif +// 383 | #endif } -// 354 | #endif +// 384 | #endif } -// 355 | #if sg_mat_count_m >= 6 +// 385 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 356 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +// 386 | 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"; -// 357 | #if sg_mat_count_n >= 2 +// 387 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 358 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +// 388 | 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"; -// 359 | #endif +// 389 | #endif } -// 360 | #if sg_mat_count_n >= 3 +// 390 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 361 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +// 391 | 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"; -// 362 | #endif +// 392 | #endif } -// 363 | #if sg_mat_count_n >= 4 +// 393 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 364 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +// 394 | 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"; -// 365 | #endif +// 395 | #endif } -// 366 | #endif +// 396 | #endif } -// 367 | #if sg_mat_count_m >= 7 +// 397 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 368 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +// 398 | 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"; -// 369 | #if sg_mat_count_n >= 2 +// 399 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 370 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +// 400 | 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"; -// 371 | #endif +// 401 | #endif } -// 372 | #if sg_mat_count_n >= 3 +// 402 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 373 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +// 403 | 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"; -// 374 | #endif +// 404 | #endif } -// 375 | #if sg_mat_count_n >= 4 +// 405 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 376 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +// 406 | 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"; -// 377 | #endif +// 407 | #endif } -// 378 | #endif +// 408 | #endif } -// 379 | #if sg_mat_count_m >= 8 +// 409 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 380 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +// 410 | 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"; -// 381 | #if sg_mat_count_n >= 2 +// 411 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 382 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +// 412 | 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"; -// 383 | #endif +// 413 | #endif } -// 384 | #if sg_mat_count_n >= 3 +// 414 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 385 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +// 415 | 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"; -// 386 | #endif +// 416 | #endif } -// 387 | #if sg_mat_count_n >= 4 +// 417 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 388 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +// 418 | 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"; -// 389 | #endif +// 419 | #endif } -// 390 | #endif +// 420 | #endif } -// 391 | } +// 421 | } ss << " }\n"; -// 392 | // Ensure every subgroup is done reading a_scratch before it is reused for +// 422 | // Ensure every subgroup is done reading a_scratch before it is reused for ss << "\n"; -// 393 | // the next K-block step. -// 394 | workgroupBarrier(); +// 423 | // the next K-block step. +// 424 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 395 | } +// 425 | } ss << " }\n"; -// 396 | +// 426 | ss << "\n"; -// 397 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; -// 398 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). -// 399 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +// 427 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; +// 428 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). +// 429 | 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"; -// 400 | #if sg_mat_count_n >= 2 +// 430 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 401 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +// 431 | 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"; -// 402 | #endif +// 432 | #endif } -// 403 | #if sg_mat_count_n >= 3 +// 433 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 404 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +// 434 | 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"; -// 405 | #endif +// 435 | #endif } -// 406 | #if sg_mat_count_n >= 4 +// 436 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 407 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +// 437 | 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"; -// 408 | #endif +// 438 | #endif } -// 409 | #if sg_mat_count_m >= 2 +// 439 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 410 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +// 440 | 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"; -// 411 | #if sg_mat_count_n >= 2 +// 441 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 412 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +// 442 | 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"; -// 413 | #endif +// 443 | #endif } -// 414 | #if sg_mat_count_n >= 3 +// 444 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 415 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +// 445 | 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"; -// 416 | #endif +// 446 | #endif } -// 417 | #if sg_mat_count_n >= 4 +// 447 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 418 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +// 448 | 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"; -// 419 | #endif +// 449 | #endif } -// 420 | #endif +// 450 | #endif } -// 421 | #if sg_mat_count_m >= 3 +// 451 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 422 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +// 452 | 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"; -// 423 | #if sg_mat_count_n >= 2 +// 453 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 424 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +// 454 | 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"; -// 425 | #endif +// 455 | #endif } -// 426 | #if sg_mat_count_n >= 3 +// 456 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 427 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +// 457 | 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"; -// 428 | #endif +// 458 | #endif } -// 429 | #if sg_mat_count_n >= 4 +// 459 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 430 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +// 460 | 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"; -// 431 | #endif +// 461 | #endif } -// 432 | #endif +// 462 | #endif } -// 433 | #if sg_mat_count_m >= 4 +// 463 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 434 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +// 464 | 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"; -// 435 | #if sg_mat_count_n >= 2 +// 465 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 436 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +// 466 | 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"; -// 437 | #endif +// 467 | #endif } -// 438 | #if sg_mat_count_n >= 3 +// 468 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 439 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +// 469 | 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"; -// 440 | #endif +// 470 | #endif } -// 441 | #if sg_mat_count_n >= 4 +// 471 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 442 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +// 472 | 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"; -// 443 | #endif +// 473 | #endif } -// 444 | #endif +// 474 | #endif } -// 445 | #if sg_mat_count_m >= 5 +// 475 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 446 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +// 476 | 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"; -// 447 | #if sg_mat_count_n >= 2 +// 477 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 448 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +// 478 | 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"; -// 449 | #endif +// 479 | #endif } -// 450 | #if sg_mat_count_n >= 3 +// 480 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 451 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +// 481 | 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"; -// 452 | #endif +// 482 | #endif } -// 453 | #if sg_mat_count_n >= 4 +// 483 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 454 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +// 484 | 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"; -// 455 | #endif +// 485 | #endif } -// 456 | #endif +// 486 | #endif } -// 457 | #if sg_mat_count_m >= 6 +// 487 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +// 488 | 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"; -// 459 | #if sg_mat_count_n >= 2 +// 489 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 460 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +// 490 | 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"; -// 461 | #endif +// 491 | #endif } -// 462 | #if sg_mat_count_n >= 3 +// 492 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 463 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +// 493 | 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"; -// 464 | #endif +// 494 | #endif } -// 465 | #if sg_mat_count_n >= 4 +// 495 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 466 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +// 496 | 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"; -// 467 | #endif +// 497 | #endif } -// 468 | #endif +// 498 | #endif } -// 469 | #if sg_mat_count_m >= 7 +// 499 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 470 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +// 500 | 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"; -// 471 | #if sg_mat_count_n >= 2 +// 501 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 472 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +// 502 | 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"; -// 473 | #endif +// 503 | #endif } -// 474 | #if sg_mat_count_n >= 3 +// 504 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 475 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +// 505 | 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"; -// 476 | #endif +// 506 | #endif } -// 477 | #if sg_mat_count_n >= 4 +// 507 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 478 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +// 508 | 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"; -// 479 | #endif +// 509 | #endif } -// 480 | #endif +// 510 | #endif } -// 481 | #if sg_mat_count_m >= 8 +// 511 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 482 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +// 512 | 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"; -// 483 | #if sg_mat_count_n >= 2 +// 513 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 484 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +// 514 | 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"; -// 485 | #endif +// 515 | #endif } -// 486 | #if sg_mat_count_n >= 3 +// 516 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 487 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +// 517 | 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"; -// 488 | #endif +// 518 | #endif } -// 489 | #if sg_mat_count_n >= 4 +// 519 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 490 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +// 520 | 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"; -// 491 | #endif +// 521 | #endif } -// 492 | #endif +// 522 | #endif } -// 493 | -// 494 | #if split_k >= 2 +// 523 | +// 524 | #if split_k >= 2 if (__param_split_k >= 2) { -// 495 | // Make all subgroups' partial tiles visible before the reduction. -// 496 | workgroupBarrier(); +// 525 | // Make all subgroups' partial tiles visible before the reduction. +// 526 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 497 | +// 527 | ss << "\n"; -// 498 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each -// 499 | // position is only overwritten by the lane that already read it. -// 500 | if (sg_index == 0u) { +// 528 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each +// 529 | // position is only overwritten by the lane that already read it. +// 530 | if (sg_index == 0u) { ss << " if (sg_index == 0u) {\n"; -// 501 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +// 531 | 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"; -// 502 | var acc: f16 = f16(0); +// 532 | var acc: f16 = f16(0); ss << " var acc: f16 = f16(0);\n"; -// 503 | for (var s: u32 = 0; s < kSplitK; s++) { +// 533 | for (var s: u32 = 0; s < kSplitK; s++) { ss << " for (var s: u32 = 0; s < kSplitK; s++) {\n"; -// 504 | acc += scratch[s * kTileM * kTileN + idx]; +// 534 | acc += scratch[s * kTileM * kTileN + idx]; ss << " acc += scratch[s * kTileM * kTileN + idx];\n"; -// 505 | } +// 535 | } ss << " }\n"; -// 506 | scratch[idx] = acc; +// 536 | scratch[idx] = acc; ss << " scratch[idx] = acc;\n"; -// 507 | } +// 537 | } ss << " }\n"; -// 508 | } +// 538 | } ss << " }\n"; -// 509 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +// 539 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << "\n"; -// 510 | workgroupBarrier(); +// 540 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 511 | #endif +// 541 | #endif } -// 512 | -// 513 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip -// 514 | // and M's bounds check make non-multiple dimensions write each element once. -// 515 | let n_count = min(kTileN, uniforms.N - global_base_n); +// 542 | +// 543 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip +// 544 | // and M's bounds check make non-multiple dimensions write each element once. +// 545 | let n_count = min(kTileN, uniforms.N - global_base_n); ss << " let n_count = min(kTileN, uniforms.N - global_base_n);\n"; -// 516 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +// 546 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { ss << " for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) {\n"; -// 517 | let global_m = m_base + r; +// 547 | let global_m = m_base + r; ss << " let global_m = m_base + r;\n"; -// 518 | if (global_m < uniforms.M) { +// 548 | if (global_m < uniforms.M) { ss << " if (global_m < uniforms.M) {\n"; -// 519 | let scratch_base = r * kTileN; +// 549 | let scratch_base = r * kTileN; ss << " let scratch_base = r * kTileN;\n"; -// 520 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; +// 550 | 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"; -// 521 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +// 551 | 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"; -// 522 | write_output(out_base + i, global_base_n + i, +// 552 | write_output(out_base + i, global_base_n + i, ss << " write_output(out_base + i, global_base_n + i,\n"; -// 523 | output_value_t(scratch[scratch_base + i])); +// 553 | output_value_t(scratch[scratch_base + i])); ss << " output_value_t(scratch[scratch_base + i]));\n"; -// 524 | } +// 554 | } ss << " }\n"; -// 525 | } +// 555 | } ss << " }\n"; -// 526 | } +// 556 | } ss << " }\n"; -// 527 | } // MAIN +// 557 | } // MAIN MainFunctionEnd(); ss << "\n"; -// 528 | +// 558 | return Status::OK(); 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 31b02837e4c3a..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 @@ -128,6 +128,7 @@ Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_help 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; 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 7e76b49326049..b66530c7507b4 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,10 +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" // 8ea2bc37b92b7f0b615005b3856710d09f1a37f48e3680252bdc239f8f0ba2d5 +#include "wgsl_template_gen/generated/nn/subgroup_matrix_conv_8x16x16.h" // 8e5c73e9849494d81da110da2fb808dddc4cdfae6eb96d054b205c8f8c505e51 #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 index 8be84a7ab6a25..9e47cc73911b2 100644 --- 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 @@ -10,6 +10,7 @@ Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelp [[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; @@ -31,1226 +32,1310 @@ Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelp // 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 | // the OHWI weight. -// 12 | // B: weight pre-transposed to OHWI, column-major (B[k,n] at n*K+k). -// 13 | // -// 14 | // Workgroup: split_k subgroups x 32 lanes, tile kTileM x kTileN (host-selected). K -// 15 | // blocks round-robin across subgroups when split_k > 1; partials summed at -// 16 | // write-out. Staging A forces the K reduction into a per-step -// 17 | // gather/barrier/multiply-accumulate loop in uniform control flow. -// 18 | // -// 19 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. -// 20 | // -// 21 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N so a trailing partial N tile can -// 22 | // shift back in bounds. M is unconstrained (see load_im2col_vec). -// 23 | -// 24 | #param sg_mat_k -// 25 | #param sg_mat_m -// 26 | #param sg_mat_n -// 27 | #param sg_mat_count_m -// 28 | #param sg_mat_count_n -// 29 | #param split_k -// 30 | #param vec_size -// 31 | -// 32 | #use .getByOffset .setByOffset -// 33 | -// 34 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +// 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 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N_b so a trailing partial N tile +// 26 | // can shift back in bounds. M is unconstrained (see load_im2col_vec). +// 27 | +// 28 | #param sg_mat_k +// 29 | #param sg_mat_m +// 30 | #param sg_mat_n +// 31 | #param sg_mat_count_m +// 32 | #param sg_mat_count_n +// 33 | #param split_k +// 34 | #param vec_size +// 35 | #param b_is_hwio +// 36 | +// 37 | #use .getByOffset .setByOffset +// 38 | +// 39 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) ss << __str_32; ss << __param_sg_mat_m; ss << __str_17; -// 35 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +// 40 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) ss << __str_33; ss << __param_sg_mat_n; ss << __str_17; -// 36 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +// 41 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) ss << __str_34; ss << __param_sg_mat_k; ss << __str_17; -// 37 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +// 42 | 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; -// 38 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +// 43 | 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; -// 39 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +// 44 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) ss << __str_37; -// 40 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +// 45 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) ss << __str_38; -// 41 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +// 46 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K ss << __str_39; ss << __param_split_k; ss << __str_17; -// 42 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +// 47 | const kSubgroupSize: u32 = 32; // Lanes per subgroup ss << __str_40; -// 43 | const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (cin % vec_size == 0) +// 48 | 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; -// 44 | +// 49 | ss << __str_5; -// 45 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K -// 46 | // subgroup, summed at write-back. -// 47 | var scratch: array; +// 50 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K +// 51 | // subgroup, summed at write-back. +// 52 | var scratch: array; ss << __str_41; -// 48 | +// 53 | ss << __str_5; -// 49 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. -// 50 | var a_scratch: array; +// 54 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 55 | var a_scratch: array; ss << __str_408; -// 51 | +// 56 | ss << __str_5; -// 52 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always -// 53 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding -// 54 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. -// 55 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. -// 56 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { +// 57 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always +// 58 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding +// 59 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. +// 60 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. +// 61 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { ss << __str_409; -// 57 | let cin_vec = uniforms.cin / kVecSize; +// 62 | let cin_vec = uniforms.cin / kVecSize; ss << __str_410; -// 58 | let oh = m / uniforms.output_w; +// 63 | let oh = m / uniforms.output_w; ss << __str_411; -// 59 | let ow = m % uniforms.output_w; +// 64 | let ow = m % uniforms.output_w; ss << __str_412; -// 60 | let k_cin = k_vec % cin_vec; +// 65 | let k_cin = k_vec % cin_vec; ss << __str_413; -// 61 | let k_hw = k_vec / cin_vec; +// 66 | let k_hw = k_vec / cin_vec; ss << __str_414; -// 62 | let kw_idx = k_hw % uniforms.kernel_w; +// 67 | let kw_idx = k_hw % uniforms.kernel_w; ss << __str_415; -// 63 | let kh_idx = k_hw / uniforms.kernel_w; +// 68 | let kh_idx = k_hw / uniforms.kernel_w; ss << __str_416; -// 64 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); +// 69 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); ss << __str_417; -// 65 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); +// 70 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); ss << __str_418; -// 66 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { +// 71 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { ss << __str_419; -// 67 | return input_src_value_t(); +// 72 | return input_src_value_t(); ss << __str_420; -// 68 | } +// 73 | } ss << __str_20; -// 69 | let input_vec_idx = input_batch_base_vec + +// 74 | let input_vec_idx = input_batch_base_vec + ss << __str_421; -// 70 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +// 75 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; ss << __str_422; -// 71 | return input_src.getByOffset(input_vec_idx); +// 76 | return input_src.getByOffset(input_vec_idx); ss << __str_351; ss << __var_input_src->GetByOffset(__str_406); ss << __str_29; -// 72 | } +// 77 | } ss << __str_345; -// 73 | +// 78 | ss << __str_5; -// 74 | $MAIN { +// 79 | $MAIN { MainFunctionStart(); ss << __str_5; -// 75 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +// 80 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is ss << __str_5; -// 76 | // the true linear index used to recover batch_id and (n_tile, m_tile). -// 77 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +// 81 | // the true linear index used to recover batch_id and (n_tile, m_tile). +// 82 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; ss << __str_218; -// 78 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +// 83 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; ss << __str_423; -// 79 | let batch_id = workgroup_idx / tiles_per_batch; +// 84 | let batch_id = workgroup_idx / tiles_per_batch; ss << __str_424; -// 80 | // Dispatch normalization can add extra workgroups past the last batch element; +// 85 | // Dispatch normalization can add extra workgroups past the last batch element; ss << __str_5; -// 81 | // drop them here (workgroup-uniform, so the barriers below stay uniform). -// 82 | if (batch_id >= uniforms.batch) { +// 86 | // drop them here (workgroup-uniform, so the barriers below stay uniform). +// 87 | if (batch_id >= uniforms.batch) { ss << __str_425; -// 83 | return; +// 88 | return; ss << __str_426; -// 84 | } +// 89 | } ss << __str_20; -// 85 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +// 90 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element ss << __str_427; -// 86 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile +// 91 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile ss << __str_428; -// 87 | let n_tile_start = n_tile * kTileN; +// 92 | let n_tile_start = n_tile * kTileN; ss << __str_45; -// 88 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +// 93 | // Shift a trailing partial N tile back so unchecked B loads stay inside the ss << __str_5; -// 89 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). -// 90 | let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); -ss << __str_47; -// 91 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile +// 94 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). +// 95 | // Uses N_b (== N unless B was padded for row-major alignment) so the shifted +// 96 | // origin -- itself n_align-even since N_b, kTileN and n_tile_start all are -- +// 97 | // stays alignment-safe for the HWIO load below. +// 98 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); +ss << __str_224; +// 99 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile ss << __str_205; -// 92 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile +// 100 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile ss << __str_429; -// 93 | let m_base = m_tile * kTileM; +// 101 | let m_base = m_tile * kTileM; ss << __str_430; -// 94 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); +// 102 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); ss << __str_431; -// 95 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +// 103 | let out_batch_base = batch_id * uniforms.M * uniforms.N; ss << __str_432; -// 96 | let k_blocks = uniforms.K / kSgMatK; +// 104 | let k_blocks = uniforms.K / kSgMatK; ss << __str_49; -// 97 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +// 105 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) ss << __str_50; -// 98 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +// 106 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) ss << __str_51; -// 99 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +// 107 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot ss << __str_52; -// 100 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +// 108 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot ss << __str_433; -// 101 | +// 109 | ss << __str_5; -// 102 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). -// 103 | var sg_mat_c0_0: subgroup_matrix_result; +// 110 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 111 | 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; -// 104 | #if sg_mat_count_n >= 2 +// 112 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 105 | var sg_mat_c0_1: subgroup_matrix_result; +// 113 | 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; -// 106 | #endif +// 114 | #endif } -// 107 | #if sg_mat_count_n >= 3 +// 115 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 108 | var sg_mat_c0_2: subgroup_matrix_result; +// 116 | 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; -// 109 | #endif +// 117 | #endif } -// 110 | #if sg_mat_count_n >= 4 +// 118 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 111 | var sg_mat_c0_3: subgroup_matrix_result; +// 119 | 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; -// 112 | #endif +// 120 | #endif } -// 113 | #if sg_mat_count_m >= 2 +// 121 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 114 | var sg_mat_c1_0: subgroup_matrix_result; +// 122 | 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; -// 115 | #if sg_mat_count_n >= 2 +// 123 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 116 | var sg_mat_c1_1: subgroup_matrix_result; +// 124 | 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; -// 117 | #endif +// 125 | #endif } -// 118 | #if sg_mat_count_n >= 3 +// 126 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 119 | var sg_mat_c1_2: subgroup_matrix_result; +// 127 | 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; -// 120 | #endif +// 128 | #endif } -// 121 | #if sg_mat_count_n >= 4 +// 129 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 122 | var sg_mat_c1_3: subgroup_matrix_result; +// 130 | 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; -// 123 | #endif +// 131 | #endif } -// 124 | #endif +// 132 | #endif } -// 125 | #if sg_mat_count_m >= 3 +// 133 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 126 | var sg_mat_c2_0: subgroup_matrix_result; +// 134 | 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; -// 127 | #if sg_mat_count_n >= 2 +// 135 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 128 | var sg_mat_c2_1: subgroup_matrix_result; +// 136 | 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; -// 129 | #endif +// 137 | #endif } -// 130 | #if sg_mat_count_n >= 3 +// 138 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 131 | var sg_mat_c2_2: subgroup_matrix_result; +// 139 | 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; -// 132 | #endif +// 140 | #endif } -// 133 | #if sg_mat_count_n >= 4 +// 141 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 134 | var sg_mat_c2_3: subgroup_matrix_result; +// 142 | 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; -// 135 | #endif +// 143 | #endif } -// 136 | #endif +// 144 | #endif } -// 137 | #if sg_mat_count_m >= 4 +// 145 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 138 | var sg_mat_c3_0: subgroup_matrix_result; +// 146 | 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; -// 139 | #if sg_mat_count_n >= 2 +// 147 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 140 | var sg_mat_c3_1: subgroup_matrix_result; +// 148 | 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; -// 141 | #endif +// 149 | #endif } -// 142 | #if sg_mat_count_n >= 3 +// 150 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 143 | var sg_mat_c3_2: subgroup_matrix_result; +// 151 | 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; -// 144 | #endif +// 152 | #endif } -// 145 | #if sg_mat_count_n >= 4 +// 153 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 146 | var sg_mat_c3_3: subgroup_matrix_result; +// 154 | 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; -// 147 | #endif +// 155 | #endif } -// 148 | #endif +// 156 | #endif } -// 149 | #if sg_mat_count_m >= 5 +// 157 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 150 | var sg_mat_c4_0: subgroup_matrix_result; +// 158 | 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; -// 151 | #if sg_mat_count_n >= 2 +// 159 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 152 | var sg_mat_c4_1: subgroup_matrix_result; +// 160 | 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; -// 153 | #endif +// 161 | #endif } -// 154 | #if sg_mat_count_n >= 3 +// 162 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 155 | var sg_mat_c4_2: subgroup_matrix_result; +// 163 | 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; -// 156 | #endif +// 164 | #endif } -// 157 | #if sg_mat_count_n >= 4 +// 165 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 158 | var sg_mat_c4_3: subgroup_matrix_result; +// 166 | 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; -// 159 | #endif +// 167 | #endif } -// 160 | #endif +// 168 | #endif } -// 161 | #if sg_mat_count_m >= 6 +// 169 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 162 | var sg_mat_c5_0: subgroup_matrix_result; +// 170 | 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; -// 163 | #if sg_mat_count_n >= 2 +// 171 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 164 | var sg_mat_c5_1: subgroup_matrix_result; +// 172 | 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; -// 165 | #endif +// 173 | #endif } -// 166 | #if sg_mat_count_n >= 3 +// 174 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 167 | var sg_mat_c5_2: subgroup_matrix_result; +// 175 | 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; -// 168 | #endif +// 176 | #endif } -// 169 | #if sg_mat_count_n >= 4 +// 177 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 170 | var sg_mat_c5_3: subgroup_matrix_result; +// 178 | 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; -// 171 | #endif +// 179 | #endif } -// 172 | #endif +// 180 | #endif } -// 173 | #if sg_mat_count_m >= 7 +// 181 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 174 | var sg_mat_c6_0: subgroup_matrix_result; +// 182 | 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; -// 175 | #if sg_mat_count_n >= 2 +// 183 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 176 | var sg_mat_c6_1: subgroup_matrix_result; +// 184 | 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; -// 177 | #endif +// 185 | #endif } -// 178 | #if sg_mat_count_n >= 3 +// 186 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 179 | var sg_mat_c6_2: subgroup_matrix_result; +// 187 | 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; -// 180 | #endif +// 188 | #endif } -// 181 | #if sg_mat_count_n >= 4 +// 189 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 182 | var sg_mat_c6_3: subgroup_matrix_result; +// 190 | 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; -// 183 | #endif +// 191 | #endif } -// 184 | #endif +// 192 | #endif } -// 185 | #if sg_mat_count_m >= 8 +// 193 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 186 | var sg_mat_c7_0: subgroup_matrix_result; +// 194 | 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; -// 187 | #if sg_mat_count_n >= 2 +// 195 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 188 | var sg_mat_c7_1: subgroup_matrix_result; +// 196 | 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; -// 189 | #endif +// 197 | #endif } -// 190 | #if sg_mat_count_n >= 3 +// 198 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 191 | var sg_mat_c7_2: subgroup_matrix_result; +// 199 | 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; -// 192 | #endif +// 200 | #endif } -// 193 | #if sg_mat_count_n >= 4 +// 201 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 194 | var sg_mat_c7_3: subgroup_matrix_result; +// 202 | 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; -// 195 | #endif +// 203 | #endif } -// 196 | #endif +// 204 | #endif } -// 197 | -// 198 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, -// 199 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks -// 200 | // early (valid == false), keeping the barriers below in uniform control flow. -// 201 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +// 205 | +// 206 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, +// 207 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks +// 208 | // early (valid == false), keeping the barriers below in uniform control flow. +// 209 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; ss << __str_434; -// 202 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +// 210 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { ss << __str_435; -// 203 | let kb = sg_index + step * kSplitK; +// 211 | let kb = sg_index + step * kSplitK; ss << __str_436; -// 204 | let valid = kb < k_blocks; +// 212 | let valid = kb < k_blocks; ss << __str_437; -// 205 | +// 213 | ss << __str_5; -// 206 | if (valid) { +// 214 | if (valid) { ss << __str_438; -// 207 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +// 215 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; ss << __str_439; -// 208 | let k_vec_base = kb * k_vec_per_sgmat_row; +// 216 | let k_vec_base = kb * k_vec_per_sgmat_row; ss << __str_440; -// 209 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { +// 217 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { ss << __str_441; -// 210 | let row = vidx / k_vec_per_sgmat_row; +// 218 | let row = vidx / k_vec_per_sgmat_row; ss << __str_442; -// 211 | let kc_vec = vidx % k_vec_per_sgmat_row; +// 219 | let kc_vec = vidx % k_vec_per_sgmat_row; ss << __str_443; -// 212 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); +// 220 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); ss << __str_444; -// 213 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +// 221 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; ss << __str_445; -// 214 | #if vec_size == 1 +// 222 | #if vec_size == 1 if (__param_vec_size == 1) { -// 215 | a_scratch[dst] = chunk; +// 223 | a_scratch[dst] = chunk; ss << __str_446; -// 216 | #else +// 224 | #else } else { -// 217 | a_scratch[dst + 0u] = chunk[0]; +// 225 | a_scratch[dst + 0u] = chunk[0]; ss << __str_447; -// 218 | a_scratch[dst + 1u] = chunk[1]; +// 226 | a_scratch[dst + 1u] = chunk[1]; ss << __str_448; -// 219 | #if vec_size >= 4 +// 227 | #if vec_size >= 4 if (__param_vec_size >= 4) { -// 220 | a_scratch[dst + 2u] = chunk[2]; +// 228 | a_scratch[dst + 2u] = chunk[2]; ss << __str_449; -// 221 | a_scratch[dst + 3u] = chunk[3]; +// 229 | a_scratch[dst + 3u] = chunk[3]; ss << __str_450; -// 222 | #endif +// 230 | #endif } -// 223 | #endif +// 231 | #endif } -// 224 | } +// 232 | } ss << __str_202; -// 225 | } +// 233 | } ss << __str_204; -// 226 | // Make the staged A-tiles visible before loading them as subgroup matrices. +// 234 | // Make the staged A-tiles visible before loading them as subgroup matrices. ss << __str_5; -// 227 | workgroupBarrier(); +// 235 | workgroupBarrier(); ss << __str_451; -// 228 | +// 236 | ss << __str_5; -// 229 | if (valid) { +// 237 | if (valid) { ss << __str_438; -// 230 | // Load the B right tiles for this K block. The weight is OHWI, i.e. B is +// 238 | // Load the B right tiles for this K block. ss << __str_5; -// 231 | // stored column-major ([N, K] row-major, element B[k, n] at n*K + k), so -// 232 | // the tile for N-block j (columns global_base_n + j*sg_mat_n) starts at -// 233 | // (global_base_n + j*sg_mat_n)*K + kb*sg_mat_k and is loaded col-major -// 234 | // with stride K. -// 235 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +// 239 | #if b_is_hwio +if (__param_b_is_hwio) { +// 240 | // HWIO: B is row-major ([K, N_b], B[k, n] at k*N_b + n); load row-major, +// 241 | // stride N_b. N_b (== N unless padded) keeps the row stride, and via +// 242 | // global_base_n above, the column offset, 4-byte aligned. +// 243 | let b_base = kb * kSgMatK * uniforms.N_b + global_base_n; ss << __str_452; -// 236 | var sg_mat_b0: subgroup_matrix_right = +// 244 | 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; -// 237 | subgroupMatrixLoad, col_major>( +// 245 | subgroupMatrixLoad, row_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; -ss << __str_92; -// 238 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_101; +// 246 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); ss << __str_455; -// 239 | #if sg_mat_count_n >= 2 +// 247 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 240 | var sg_mat_b1: subgroup_matrix_right = +// 248 | 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; -// 241 | subgroupMatrixLoad, col_major>( +// 249 | subgroupMatrixLoad, row_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; -ss << __str_92; -// 242 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_101; +// 250 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); ss << __str_457; -// 243 | #endif +// 251 | #endif } -// 244 | #if sg_mat_count_n >= 3 +// 252 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 245 | var sg_mat_b2: subgroup_matrix_right = +// 253 | 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; -// 246 | subgroupMatrixLoad, col_major>( +// 254 | subgroupMatrixLoad, row_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; -ss << __str_92; -// 247 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_101; +// 255 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); ss << __str_459; -// 248 | #endif +// 256 | #endif } -// 249 | #if sg_mat_count_n >= 4 +// 257 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 250 | var sg_mat_b3: subgroup_matrix_right = +// 258 | 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; -// 251 | subgroupMatrixLoad, col_major>( +// 259 | subgroupMatrixLoad, row_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; -ss << __str_92; -// 252 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_101; +// 260 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); ss << __str_461; -// 253 | #endif +// 261 | #endif } -// 254 | -// 255 | // Load the A left tiles (one per M block) from workgroup staging, -// 256 | // row-major with stride kSgMatK. -// 257 | var sg_mat_a0: subgroup_matrix_left = +// 262 | #else +} else { +// 263 | // OHWI: B is column-major ([N, K], B[k, n] at n*K + k); load col-major, stride K. +// 264 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; ss << __str_462; +// 265 | 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; +// 266 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 267 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_463; +// 268 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 269 | 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; +// 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 + 1 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_464; +// 272 | #endif +} +// 273 | #if sg_mat_count_n >= 3 +if (__param_sg_mat_count_n >= 3) { +// 274 | 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; +// 275 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 276 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_465; +// 277 | #endif +} +// 278 | #if sg_mat_count_n >= 4 +if (__param_sg_mat_count_n >= 4) { +// 279 | 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; +// 280 | subgroupMatrixLoad, col_major>( +ss << __str_454; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_92; +// 281 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +ss << __str_466; +// 282 | #endif +} +// 283 | #endif +} +// 284 | +// 285 | // Load the A left tiles (one per M block) from workgroup staging, +// 286 | // row-major with stride kSgMatK. +// 287 | 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; -// 258 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 288 | subgroupMatrixLoad, row_major>( +ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 259 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_464; -// 260 | #if sg_mat_count_m >= 2 +// 289 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_469; +// 290 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 261 | var sg_mat_a1: subgroup_matrix_left = -ss << __str_465; +// 291 | 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; -// 262 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 292 | subgroupMatrixLoad, row_major>( +ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 263 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_466; -// 264 | #endif +// 293 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_471; +// 294 | #endif } -// 265 | #if sg_mat_count_m >= 3 +// 295 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 266 | var sg_mat_a2: subgroup_matrix_left = -ss << __str_467; +// 296 | 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; -// 267 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 297 | subgroupMatrixLoad, row_major>( +ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 268 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_468; -// 269 | #endif +// 298 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_473; +// 299 | #endif } -// 270 | #if sg_mat_count_m >= 4 +// 300 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 271 | var sg_mat_a3: subgroup_matrix_left = -ss << __str_469; +// 301 | 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; -// 272 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 302 | subgroupMatrixLoad, row_major>( +ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 273 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_470; -// 274 | #endif +// 303 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_475; +// 304 | #endif } -// 275 | #if sg_mat_count_m >= 5 +// 305 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 276 | var sg_mat_a4: subgroup_matrix_left = -ss << __str_471; +// 306 | 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; -// 277 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 307 | subgroupMatrixLoad, row_major>( +ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 278 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_472; -// 279 | #endif +// 308 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_477; +// 309 | #endif } -// 280 | #if sg_mat_count_m >= 6 +// 310 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 281 | var sg_mat_a5: subgroup_matrix_left = -ss << __str_473; +// 311 | 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; -// 282 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 312 | subgroupMatrixLoad, row_major>( +ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 283 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_474; -// 284 | #endif +// 313 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_479; +// 314 | #endif } -// 285 | #if sg_mat_count_m >= 7 +// 315 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 286 | var sg_mat_a6: subgroup_matrix_left = -ss << __str_475; +// 316 | 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; -// 287 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 317 | subgroupMatrixLoad, row_major>( +ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 288 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_476; -// 289 | #endif +// 318 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_481; +// 319 | #endif } -// 290 | #if sg_mat_count_m >= 8 +// 320 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 291 | var sg_mat_a7: subgroup_matrix_left = -ss << __str_477; +// 321 | 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; -// 292 | subgroupMatrixLoad, row_major>( -ss << __str_463; +// 322 | 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 + 7 * kSgMatM * kSgMatK, kSgMatK); -ss << __str_478; -// 294 | #endif +// 323 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +ss << __str_483; +// 324 | #endif } -// 295 | -// 296 | // Accumulate every (M block, N block) pair. -// 297 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); -ss << __str_479; -// 298 | #if sg_mat_count_n >= 2 +// 325 | +// 326 | // Accumulate every (M block, N block) pair. +// 327 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +ss << __str_484; +// 328 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 299 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); -ss << __str_480; -// 300 | #endif +// 329 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +ss << __str_485; +// 330 | #endif } -// 301 | #if sg_mat_count_n >= 3 +// 331 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 302 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); -ss << __str_481; -// 303 | #endif +// 332 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +ss << __str_486; +// 333 | #endif } -// 304 | #if sg_mat_count_n >= 4 +// 334 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 305 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); -ss << __str_482; -// 306 | #endif +// 335 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +ss << __str_487; +// 336 | #endif } -// 307 | #if sg_mat_count_m >= 2 +// 337 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 308 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); -ss << __str_483; -// 309 | #if sg_mat_count_n >= 2 +// 338 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +ss << __str_488; +// 339 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 310 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); -ss << __str_484; -// 311 | #endif +// 340 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +ss << __str_489; +// 341 | #endif } -// 312 | #if sg_mat_count_n >= 3 +// 342 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 313 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); -ss << __str_485; -// 314 | #endif +// 343 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +ss << __str_490; +// 344 | #endif } -// 315 | #if sg_mat_count_n >= 4 +// 345 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 316 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); -ss << __str_486; -// 317 | #endif +// 346 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +ss << __str_491; +// 347 | #endif } -// 318 | #endif +// 348 | #endif } -// 319 | #if sg_mat_count_m >= 3 +// 349 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 320 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); -ss << __str_487; -// 321 | #if sg_mat_count_n >= 2 +// 350 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +ss << __str_492; +// 351 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 322 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); -ss << __str_488; -// 323 | #endif +// 352 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +ss << __str_493; +// 353 | #endif } -// 324 | #if sg_mat_count_n >= 3 +// 354 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 325 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); -ss << __str_489; -// 326 | #endif +// 355 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +ss << __str_494; +// 356 | #endif } -// 327 | #if sg_mat_count_n >= 4 +// 357 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 328 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); -ss << __str_490; -// 329 | #endif +// 358 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +ss << __str_495; +// 359 | #endif } -// 330 | #endif +// 360 | #endif } -// 331 | #if sg_mat_count_m >= 4 +// 361 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 332 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); -ss << __str_491; -// 333 | #if sg_mat_count_n >= 2 +// 362 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +ss << __str_496; +// 363 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 334 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); -ss << __str_492; -// 335 | #endif +// 364 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +ss << __str_497; +// 365 | #endif } -// 336 | #if sg_mat_count_n >= 3 +// 366 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 337 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); -ss << __str_493; -// 338 | #endif +// 367 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +ss << __str_498; +// 368 | #endif } -// 339 | #if sg_mat_count_n >= 4 +// 369 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 340 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); -ss << __str_494; -// 341 | #endif +// 370 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +ss << __str_499; +// 371 | #endif } -// 342 | #endif +// 372 | #endif } -// 343 | #if sg_mat_count_m >= 5 +// 373 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 344 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); -ss << __str_495; -// 345 | #if sg_mat_count_n >= 2 +// 374 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +ss << __str_500; +// 375 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 346 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); -ss << __str_496; -// 347 | #endif +// 376 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +ss << __str_501; +// 377 | #endif } -// 348 | #if sg_mat_count_n >= 3 +// 378 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 349 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); -ss << __str_497; -// 350 | #endif +// 379 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +ss << __str_502; +// 380 | #endif } -// 351 | #if sg_mat_count_n >= 4 +// 381 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 352 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); -ss << __str_498; -// 353 | #endif +// 382 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +ss << __str_503; +// 383 | #endif } -// 354 | #endif +// 384 | #endif } -// 355 | #if sg_mat_count_m >= 6 +// 385 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 356 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); -ss << __str_499; -// 357 | #if sg_mat_count_n >= 2 +// 386 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +ss << __str_504; +// 387 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 358 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); -ss << __str_500; -// 359 | #endif +// 388 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +ss << __str_505; +// 389 | #endif } -// 360 | #if sg_mat_count_n >= 3 +// 390 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 361 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); -ss << __str_501; -// 362 | #endif +// 391 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +ss << __str_506; +// 392 | #endif } -// 363 | #if sg_mat_count_n >= 4 +// 393 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 364 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); -ss << __str_502; -// 365 | #endif +// 394 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +ss << __str_507; +// 395 | #endif } -// 366 | #endif +// 396 | #endif } -// 367 | #if sg_mat_count_m >= 7 +// 397 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 368 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); -ss << __str_503; -// 369 | #if sg_mat_count_n >= 2 +// 398 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +ss << __str_508; +// 399 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 370 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); -ss << __str_504; -// 371 | #endif +// 400 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +ss << __str_509; +// 401 | #endif } -// 372 | #if sg_mat_count_n >= 3 +// 402 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 373 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); -ss << __str_505; -// 374 | #endif +// 403 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +ss << __str_510; +// 404 | #endif } -// 375 | #if sg_mat_count_n >= 4 +// 405 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 376 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); -ss << __str_506; -// 377 | #endif +// 406 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +ss << __str_511; +// 407 | #endif } -// 378 | #endif +// 408 | #endif } -// 379 | #if sg_mat_count_m >= 8 +// 409 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 380 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); -ss << __str_507; -// 381 | #if sg_mat_count_n >= 2 +// 410 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +ss << __str_512; +// 411 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 382 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); -ss << __str_508; -// 383 | #endif +// 412 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +ss << __str_513; +// 413 | #endif } -// 384 | #if sg_mat_count_n >= 3 +// 414 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 385 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); -ss << __str_509; -// 386 | #endif +// 415 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +ss << __str_514; +// 416 | #endif } -// 387 | #if sg_mat_count_n >= 4 +// 417 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 388 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); -ss << __str_510; -// 389 | #endif +// 418 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +ss << __str_515; +// 419 | #endif } -// 390 | #endif +// 420 | #endif } -// 391 | } +// 421 | } ss << __str_204; -// 392 | // Ensure every subgroup is done reading a_scratch before it is reused for +// 422 | // Ensure every subgroup is done reading a_scratch before it is reused for ss << __str_5; -// 393 | // the next K-block step. -// 394 | workgroupBarrier(); +// 423 | // the next K-block step. +// 424 | workgroupBarrier(); ss << __str_451; -// 395 | } +// 425 | } ss << __str_20; -// 396 | +// 426 | ss << __str_5; -// 397 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; -// 398 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). -// 399 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +// 427 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; +// 428 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). +// 429 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); ss << __str_165; -// 400 | #if sg_mat_count_n >= 2 +// 430 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 401 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +// 431 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); ss << __str_166; -// 402 | #endif +// 432 | #endif } -// 403 | #if sg_mat_count_n >= 3 +// 433 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 404 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +// 434 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); ss << __str_167; -// 405 | #endif +// 435 | #endif } -// 406 | #if sg_mat_count_n >= 4 +// 436 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 407 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +// 437 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); ss << __str_168; -// 408 | #endif +// 438 | #endif } -// 409 | #if sg_mat_count_m >= 2 +// 439 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 410 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +// 440 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); ss << __str_169; -// 411 | #if sg_mat_count_n >= 2 +// 441 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 412 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +// 442 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); ss << __str_170; -// 413 | #endif +// 443 | #endif } -// 414 | #if sg_mat_count_n >= 3 +// 444 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 415 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +// 445 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); ss << __str_171; -// 416 | #endif +// 446 | #endif } -// 417 | #if sg_mat_count_n >= 4 +// 447 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 418 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +// 448 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); ss << __str_172; -// 419 | #endif +// 449 | #endif } -// 420 | #endif +// 450 | #endif } -// 421 | #if sg_mat_count_m >= 3 +// 451 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 422 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); ss << __str_173; -// 423 | #if sg_mat_count_n >= 2 +// 453 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 424 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +// 454 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); ss << __str_174; -// 425 | #endif +// 455 | #endif } -// 426 | #if sg_mat_count_n >= 3 +// 456 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 427 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +// 457 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); ss << __str_175; -// 428 | #endif +// 458 | #endif } -// 429 | #if sg_mat_count_n >= 4 +// 459 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 430 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +// 460 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); ss << __str_176; -// 431 | #endif +// 461 | #endif } -// 432 | #endif +// 462 | #endif } -// 433 | #if sg_mat_count_m >= 4 +// 463 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 434 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +// 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); ss << __str_177; -// 435 | #if sg_mat_count_n >= 2 +// 465 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 436 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +// 466 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); ss << __str_178; -// 437 | #endif +// 467 | #endif } -// 438 | #if sg_mat_count_n >= 3 +// 468 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 439 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +// 469 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); ss << __str_179; -// 440 | #endif +// 470 | #endif } -// 441 | #if sg_mat_count_n >= 4 +// 471 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 442 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +// 472 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); ss << __str_180; -// 443 | #endif +// 473 | #endif } -// 444 | #endif +// 474 | #endif } -// 445 | #if sg_mat_count_m >= 5 +// 475 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 446 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +// 476 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); ss << __str_181; -// 447 | #if sg_mat_count_n >= 2 +// 477 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 448 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +// 478 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); ss << __str_182; -// 449 | #endif +// 479 | #endif } -// 450 | #if sg_mat_count_n >= 3 +// 480 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 451 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +// 481 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); ss << __str_183; -// 452 | #endif +// 482 | #endif } -// 453 | #if sg_mat_count_n >= 4 +// 483 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 454 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +// 484 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); ss << __str_184; -// 455 | #endif +// 485 | #endif } -// 456 | #endif +// 486 | #endif } -// 457 | #if sg_mat_count_m >= 6 +// 487 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +// 488 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); ss << __str_185; -// 459 | #if sg_mat_count_n >= 2 +// 489 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 460 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +// 490 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); ss << __str_186; -// 461 | #endif +// 491 | #endif } -// 462 | #if sg_mat_count_n >= 3 +// 492 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 463 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +// 493 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); ss << __str_187; -// 464 | #endif +// 494 | #endif } -// 465 | #if sg_mat_count_n >= 4 +// 495 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 466 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +// 496 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); ss << __str_188; -// 467 | #endif +// 497 | #endif } -// 468 | #endif +// 498 | #endif } -// 469 | #if sg_mat_count_m >= 7 +// 499 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 470 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +// 500 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); ss << __str_189; -// 471 | #if sg_mat_count_n >= 2 +// 501 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 472 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +// 502 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); ss << __str_190; -// 473 | #endif +// 503 | #endif } -// 474 | #if sg_mat_count_n >= 3 +// 504 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 475 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +// 505 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); ss << __str_191; -// 476 | #endif +// 506 | #endif } -// 477 | #if sg_mat_count_n >= 4 +// 507 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 478 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +// 508 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); ss << __str_192; -// 479 | #endif +// 509 | #endif } -// 480 | #endif +// 510 | #endif } -// 481 | #if sg_mat_count_m >= 8 +// 511 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 482 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +// 512 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); ss << __str_193; -// 483 | #if sg_mat_count_n >= 2 +// 513 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 484 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +// 514 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); ss << __str_194; -// 485 | #endif +// 515 | #endif } -// 486 | #if sg_mat_count_n >= 3 +// 516 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 487 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +// 517 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); ss << __str_195; -// 488 | #endif +// 518 | #endif } -// 489 | #if sg_mat_count_n >= 4 +// 519 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 490 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +// 520 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); ss << __str_196; -// 491 | #endif +// 521 | #endif } -// 492 | #endif +// 522 | #endif } -// 493 | -// 494 | #if split_k >= 2 +// 523 | +// 524 | #if split_k >= 2 if (__param_split_k >= 2) { -// 495 | // Make all subgroups' partial tiles visible before the reduction. -// 496 | workgroupBarrier(); +// 525 | // Make all subgroups' partial tiles visible before the reduction. +// 526 | workgroupBarrier(); ss << __str_26; -// 497 | +// 527 | ss << __str_5; -// 498 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each -// 499 | // position is only overwritten by the lane that already read it. -// 500 | if (sg_index == 0u) { +// 528 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each +// 529 | // position is only overwritten by the lane that already read it. +// 530 | if (sg_index == 0u) { ss << __str_197; -// 501 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +// 531 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { ss << __str_198; -// 502 | var acc: f16 = f16(0); +// 532 | var acc: f16 = f16(0); ss << __str_199; -// 503 | for (var s: u32 = 0; s < kSplitK; s++) { +// 533 | for (var s: u32 = 0; s < kSplitK; s++) { ss << __str_200; -// 504 | acc += scratch[s * kTileM * kTileN + idx]; +// 534 | acc += scratch[s * kTileM * kTileN + idx]; ss << __str_201; -// 505 | } +// 535 | } ss << __str_202; -// 506 | scratch[idx] = acc; +// 536 | scratch[idx] = acc; ss << __str_203; -// 507 | } +// 537 | } ss << __str_204; -// 508 | } +// 538 | } ss << __str_20; -// 509 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +// 539 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << __str_5; -// 510 | workgroupBarrier(); +// 540 | workgroupBarrier(); ss << __str_26; -// 511 | #endif +// 541 | #endif } -// 512 | -// 513 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip -// 514 | // and M's bounds check make non-multiple dimensions write each element once. -// 515 | let n_count = min(kTileN, uniforms.N - global_base_n); +// 542 | +// 543 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip +// 544 | // and M's bounds check make non-multiple dimensions write each element once. +// 545 | let n_count = min(kTileN, uniforms.N - global_base_n); ss << __str_206; -// 516 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +// 546 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { ss << __str_207; -// 517 | let global_m = m_base + r; +// 547 | let global_m = m_base + r; ss << __str_208; -// 518 | if (global_m < uniforms.M) { -ss << __str_511; -// 519 | let scratch_base = r * kTileN; +// 548 | if (global_m < uniforms.M) { +ss << __str_516; +// 549 | let scratch_base = r * kTileN; ss << __str_210; -// 520 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; -ss << __str_512; -// 521 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +// 550 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; +ss << __str_517; +// 551 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { ss << __str_212; -// 522 | write_output(out_base + i, global_base_n + i, +// 552 | write_output(out_base + i, global_base_n + i, ss << __str_235; -// 523 | output_value_t(scratch[scratch_base + i])); +// 553 | output_value_t(scratch[scratch_base + i])); ss << __str_236; -// 524 | } +// 554 | } ss << __str_202; -// 525 | } +// 555 | } ss << __str_204; -// 526 | } +// 556 | } ss << __str_20; -// 527 | } // MAIN +// 557 | } // MAIN MainFunctionEnd(); ss << __str_5; -// 528 | +// 558 | return Status::OK(); 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 2b2638ef80532..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_514; +ss << __str_519; // 7 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_515; +ss << __str_520; // 8 | let offset = co * uniforms.I * uniforms.H * uniforms.W + -ss << __str_516; +ss << __str_521; // 9 | ci * uniforms.H * uniforms.W + -ss << __str_517; +ss << __str_522; // 10 | h_w; -ss << __str_518; +ss << __str_523; // 11 | return src.getByOffset(offset); ss << __str_351; -ss << __var_src->GetByOffset(__str_513); +ss << __var_src->GetByOffset(__str_518); ss << __str_29; // 12 | } ss << __str_11; // 13 | return src_element_t(); -ss << __str_519; +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_520; +ss << __str_525; // 17 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_515; +ss << __str_520; // 18 | let offset = co * uniforms.H * uniforms.W * uniforms.I + -ss << __str_521; +ss << __str_526; // 19 | h_w * uniforms.I + -ss << __str_522; +ss << __str_527; // 20 | ci; -ss << __str_523; +ss << __str_528; // 21 | output.setByOffset(offset, value); ss << __str_28; -ss << __var_output->SetByOffset(__str_513, __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_524; +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_525; +ss << __str_530; // 29 | let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64; -ss << __str_526; +ss << __str_531; // 30 | ss << __str_5; // 31 | if (group_co >= uniforms.O || group_ci >= uniforms.I) { -ss << __str_527; +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_528; +ss << __str_533; // 36 | // load ss << __str_5; // 37 | for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) { -ss << __str_529; +ss << __str_534; // 38 | let load_ci_idx = ci_idx + local_idx / 4; -ss << __str_530; +ss << __str_535; // 39 | let load_h_w_idx = local_idx % 4; -ss << __str_531; +ss << __str_536; // 40 | ss << __str_5; // 41 | data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co, -ss << __str_532; +ss << __str_537; // 42 | group_ci + load_ci_idx, -ss << __str_533; +ss << __str_538; // 43 | h_w_idx * 4 + load_h_w_idx); -ss << __str_534; +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_535; +ss << __str_540; // 49 | let output_data = data_cache[local_h_w_idx][local_idx]; -ss << __str_536; +ss << __str_541; // 50 | write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data); -ss << __str_537; +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 52118960dc647..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_543; +ss << __str_548; // 18 | #if is_float16 if (__param_is_float16) { // 19 | bitcast>(uniforms.constant_value)[0]; -ss << __str_544; +ss << __str_549; // 20 | #else } else { // 21 | bitcast(uniforms.constant_value); -ss << __str_545; +ss << __str_550; // 22 | #endif } // 23 | // 24 | #if dim_value_zero if (__param_dim_value_zero) { // 25 | output[global_idx] = constant_value; -ss << __str_546; +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_547; +ss << __str_552; // 29 | var use_pad_value = false; -ss << __str_548; +ss << __str_553; // 30 | var in_coord = i32(0); -ss << __str_549; +ss << __str_554; // 31 | ss << __str_5; // 32 | for (var dim = 0; dim < output.rank && !use_pad_value; dim++) { -ss << __str_550; +ss << __str_555; ss << __var_output->Rank(); -ss << __str_551; +ss << __str_556; // 33 | let output_index = i32(getElementAt(output_indices, dim, output.rank)); -ss << __str_552; -ss << GetElementAt(__str_247, __str_538, __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_553; -ss << GetElementAt(__str_539, __str_538, __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_554; -ss << GetElementAt(__str_540, __str_538, __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_555; +ss << __str_560; // 38 | use_pad_value = true; -ss << __str_556; +ss << __str_561; // 39 | #elif pad_mode == PAD_MODE_EDGE } else if (__param_pad_mode == 2) { // 40 | if (output_index < lower_pads) { -ss << __str_557; +ss << __str_562; // 41 | in_coord = 0; -ss << __str_558; +ss << __str_563; // 42 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_559; +ss << __str_564; // 43 | in_coord = data_shape - 1; -ss << __str_560; +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_555; +ss << __str_560; // 46 | in_coord = output_index - lower_pads; -ss << __str_561; +ss << __str_566; // 47 | if (in_coord < 0) { -ss << __str_562; +ss << __str_567; // 48 | in_coord = -in_coord; -ss << __str_563; +ss << __str_568; // 49 | } ss << __str_292; // 50 | let _2n_1 = 2 * (data_shape - 1); -ss << __str_564; +ss << __str_569; // 51 | in_coord = in_coord % _2n_1; -ss << __str_565; +ss << __str_570; // 52 | if (in_coord >= data_shape) { -ss << __str_566; +ss << __str_571; // 53 | in_coord = _2n_1 - in_coord; -ss << __str_567; +ss << __str_572; // 54 | } ss << __str_292; // 55 | #else // PAD_MODE_WRAP } else { // 56 | if (output_index < lower_pads) { -ss << __str_557; +ss << __str_562; // 57 | in_coord = data_shape + output_index - lower_pads; -ss << __str_568; +ss << __str_573; // 58 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_559; +ss << __str_564; // 59 | in_coord = output_index - data_shape - lower_pads; -ss << __str_569; +ss << __str_574; // 60 | #endif // pad_mode } // 61 | } else { -ss << __str_570; +ss << __str_575; // 62 | in_coord = output_index - lower_pads; -ss << __str_571; +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_572; +ss << __str_577; // 67 | #endif } // 68 | // 69 | input_index += select(u32(in_coord) -ss << __str_573; +ss << __str_578; // 70 | #if output.rank > 1 if (__var_output->Rank() > 1) { // 71 | * getElementAt(uniforms.data_stride, dim, output.rank - 1) -ss << __str_574; -ss << GetElementAt(__str_541, __str_538, __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_575; +ss << __str_580; ss << __var_output->Rank(); -ss << __str_576; +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_542); +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 31b02837e4c3a..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 @@ -128,6 +128,7 @@ Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_help 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; 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 0aaaec60f7447..49274aa9fcdae 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,19 +31,19 @@ std::string pass_as_string(T&& v) { } } // namespace wgsl_detail -#include "wgsl_template_gen/string_table.h" // 678f18362f8d5afb8555278608d06770e427839a2727213dd26bcc53d8b0f486 +#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/nn/subgroup_matrix_conv_8x16x16.h" // 1805ea3004b0ac1e3048cd71876a56470086f6eac0c24dcca1434fed14dd6ea8 -#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // 166e2a2db54c8477de7e164973ee31f79b40a2f1049b22770b9a1c7c96ec1907 -#include "wgsl_template_gen/generated/tensor/pad.h" // 7f732549c9b2533acacbd683d9625b55b3ce6fbe0159c509bfae46cadadb6e1c +#include "wgsl_template_gen/generated/nn/subgroup_matrix_conv_8x16x16.h" // b0abd8ed2662209255943ac112b323db06a470d2c62c2fb7430f1925085172f9 +#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 814db34704fce..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 @@ -458,128 +458,133 @@ constexpr const char* __str_448 = " a_scratch[dst + 1u] = chunk[1 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 = global_base_n * uniforms.K + kb * kSgMatK;\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 Date: Wed, 28 Oct 2026 16:52:07 +0800 Subject: [PATCH 14/14] Check maxStorageBufferBindingSize for padB --- .../webgpu/nn/subgroup_matrix_conv.cc | 23 +- ...subgroup_matrix_conv_8x16x16.wgsl.template | 4 + .../webgpu/subgroup_matrix_conv_test.cc | 10 +- .../nn/subgroup_matrix_conv_8x16x16.h | 1072 +++++++++-------- .../static-cpp-literal/index_impl.h | 2 +- .../nn/subgroup_matrix_conv_8x16x16.h | 1072 +++++++++-------- .../in_tree_golden/static-cpp/index_impl.h | 2 +- 7 files changed, 1100 insertions(+), 1085 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc index cff2eac9d1119..c8100ebdf19bd 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc @@ -171,6 +171,16 @@ class SubgroupMatrixConvImpl final : public Conv::Co } 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 @@ -385,19 +395,6 @@ bool CanApplySubgroupMatrixConv(ComputeContextBase& context, return false; } - // subgroupMatrixLoad addresses the weight as a single binding; a weight past - // maxStorageBufferBindingSize gets split by ProgramManager, which this offset - // arithmetic doesn't know about. - const uint64_t max_weight_elems = context.DeviceLimits().maxStorageBufferBindingSize / sizeof(MLFloat16); - if (static_cast(kernel_shape.Size()) > max_weight_elems) { - LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: weight exceeds maxStorageBufferBindingSize" - << " (kernel_shape[OIHW]=" << kernel_shape.ToString() << ", elems=" - << kernel_shape.Size() << ", limit_elems=" << max_weight_elems - << ", maxStorageBufferBindingSize=" << context.DeviceLimits().maxStorageBufferBindingSize - << " bytes)."; - return false; - } - const std::optional config = GetSupportedConfig(context); if (!config) { LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv rejected: device does not report the required 8x16x16 F16" 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 index 3fe26066adb51..2e07af49c9558 100644 --- a/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template +++ b/onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template @@ -22,6 +22,10 @@ // // 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). diff --git a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc index e2909c55cd25e..690724b7c9ed0 100644 --- a/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc +++ b/onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc @@ -354,8 +354,14 @@ 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"); - RunConvTest(/*N=*/1, /*Cin=*/16, /*H=*/10, /*W=*/10, /*Cout=*/17, /*kh=*/3, /*kw=*/3, - /*stride=*/1, /*pad=*/1, /*dilation=*/1, /*has_bias=*/false, + 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, 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 index e6b773c7bf389..0ceaec859f0a9 100644 --- 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 @@ -46,1296 +46,1300 @@ Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelp // 22 | // // 23 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. // 24 | // -// 25 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N_b so a trailing partial N tile -// 26 | // can shift back in bounds. M is unconstrained (see load_im2col_vec). -// 27 | -// 28 | #param sg_mat_k -// 29 | #param sg_mat_m -// 30 | #param sg_mat_n -// 31 | #param sg_mat_count_m -// 32 | #param sg_mat_count_n -// 33 | #param split_k -// 34 | #param vec_size -// 35 | #param b_is_hwio -// 36 | -// 37 | #use .getByOffset .setByOffset -// 38 | -// 39 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +// 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"; -// 40 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +// 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"; -// 41 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +// 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"; -// 42 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +// 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"; -// 43 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along 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"; -// 44 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +// 48 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) ss << "const kTileM: u32 = kSgMatM * kSgMatCountM;\n"; -// 45 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +// 49 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) ss << "const kTileN: u32 = kSgMatN * kSgMatCountN;\n"; -// 46 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +// 50 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K ss << "const kSplitK: u32 = u32("; ss << __param_split_k; ss << ");\n"; -// 47 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +// 51 | const kSubgroupSize: u32 = 32; // Lanes per subgroup ss << "const kSubgroupSize: u32 = 32;\n"; -// 48 | const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (cin % vec_size == 0) +// 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"; -// 49 | +// 53 | ss << "\n"; -// 50 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K -// 51 | // subgroup, summed at write-back. -// 52 | var scratch: array; +// 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"; -// 53 | +// 57 | ss << "\n"; -// 54 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. -// 55 | var a_scratch: array; +// 58 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 59 | var a_scratch: array; ss << "var a_scratch: array;\n"; -// 56 | +// 60 | ss << "\n"; -// 57 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always -// 58 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding -// 59 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. -// 60 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. -// 61 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { +// 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"; -// 62 | let cin_vec = uniforms.cin / kVecSize; +// 66 | let cin_vec = uniforms.cin / kVecSize; ss << " let cin_vec = uniforms.cin / kVecSize;\n"; -// 63 | let oh = m / uniforms.output_w; +// 67 | let oh = m / uniforms.output_w; ss << " let oh = m / uniforms.output_w;\n"; -// 64 | let ow = m % uniforms.output_w; +// 68 | let ow = m % uniforms.output_w; ss << " let ow = m % uniforms.output_w;\n"; -// 65 | let k_cin = k_vec % cin_vec; +// 69 | let k_cin = k_vec % cin_vec; ss << " let k_cin = k_vec % cin_vec;\n"; -// 66 | let k_hw = k_vec / cin_vec; +// 70 | let k_hw = k_vec / cin_vec; ss << " let k_hw = k_vec / cin_vec;\n"; -// 67 | let kw_idx = k_hw % uniforms.kernel_w; +// 71 | let kw_idx = k_hw % uniforms.kernel_w; ss << " let kw_idx = k_hw % uniforms.kernel_w;\n"; -// 68 | let kh_idx = k_hw / uniforms.kernel_w; +// 72 | let kh_idx = k_hw / uniforms.kernel_w; ss << " let kh_idx = k_hw / uniforms.kernel_w;\n"; -// 69 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); +// 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"; -// 70 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); +// 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"; -// 71 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { +// 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"; -// 72 | return input_src_value_t(); +// 76 | return input_src_value_t(); ss << " return input_src_value_t();\n"; -// 73 | } +// 77 | } ss << " }\n"; -// 74 | let input_vec_idx = input_batch_base_vec + +// 78 | let input_vec_idx = input_batch_base_vec + ss << " let input_vec_idx = input_batch_base_vec +\n"; -// 75 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +// 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"; -// 76 | return input_src.getByOffset(input_vec_idx); +// 80 | return input_src.getByOffset(input_vec_idx); ss << " return "; ss << __var_input_src->GetByOffset("input_vec_idx"); ss << ";\n"; -// 77 | } +// 81 | } ss << "}\n"; -// 78 | +// 82 | ss << "\n"; -// 79 | $MAIN { +// 83 | $MAIN { MainFunctionStart(); ss << "\n"; -// 80 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +// 84 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is ss << "\n"; -// 81 | // the true linear index used to recover batch_id and (n_tile, m_tile). -// 82 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +// 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"; -// 83 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +// 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"; -// 84 | let batch_id = workgroup_idx / tiles_per_batch; +// 88 | let batch_id = workgroup_idx / tiles_per_batch; ss << " let batch_id = workgroup_idx / tiles_per_batch;\n"; -// 85 | // Dispatch normalization can add extra workgroups past the last batch element; +// 89 | // Dispatch normalization can add extra workgroups past the last batch element; ss << "\n"; -// 86 | // drop them here (workgroup-uniform, so the barriers below stay uniform). -// 87 | if (batch_id >= uniforms.batch) { +// 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"; -// 88 | return; +// 92 | return; ss << " return;\n"; -// 89 | } +// 93 | } ss << " }\n"; -// 90 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +// 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"; -// 91 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile +// 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"; -// 92 | let n_tile_start = n_tile * kTileN; +// 96 | let n_tile_start = n_tile * kTileN; ss << " let n_tile_start = n_tile * kTileN;\n"; -// 93 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +// 97 | // Shift a trailing partial N tile back so unchecked B loads stay inside the ss << "\n"; -// 94 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). -// 95 | // Uses N_b (== N unless B was padded for row-major alignment) so the shifted -// 96 | // origin -- itself n_align-even since N_b, kTileN and n_tile_start all are -- -// 97 | // stays alignment-safe for the HWIO load below. -// 98 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); +// 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"; -// 99 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile +// 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"; -// 100 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile +// 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"; -// 101 | let m_base = m_tile * kTileM; +// 105 | let m_base = m_tile * kTileM; ss << " let m_base = m_tile * kTileM;\n"; -// 102 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); +// 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"; -// 103 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +// 107 | let out_batch_base = batch_id * uniforms.M * uniforms.N; ss << " let out_batch_base = batch_id * uniforms.M * uniforms.N;\n"; -// 104 | let k_blocks = uniforms.K / kSgMatK; +// 108 | let k_blocks = uniforms.K / kSgMatK; ss << " let k_blocks = uniforms.K / kSgMatK;\n"; -// 105 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +// 109 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) ss << " let sg_index = local_idx / kSubgroupSize;\n"; -// 106 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +// 110 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) ss << " let sg_lane = local_idx % kSubgroupSize;\n"; -// 107 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +// 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"; -// 108 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +// 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"; -// 109 | +// 113 | ss << "\n"; -// 110 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). -// 111 | var sg_mat_c0_0: subgroup_matrix_result; +// 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"; -// 112 | #if sg_mat_count_n >= 2 +// 116 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 113 | var sg_mat_c0_1: subgroup_matrix_result; +// 117 | var sg_mat_c0_1: subgroup_matrix_result; ss << " var sg_mat_c0_1: subgroup_matrix_result;\n"; -// 114 | #endif +// 118 | #endif } -// 115 | #if sg_mat_count_n >= 3 +// 119 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 116 | var sg_mat_c0_2: subgroup_matrix_result; +// 120 | var sg_mat_c0_2: subgroup_matrix_result; ss << " var sg_mat_c0_2: subgroup_matrix_result;\n"; -// 117 | #endif +// 121 | #endif } -// 118 | #if sg_mat_count_n >= 4 +// 122 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 119 | var sg_mat_c0_3: subgroup_matrix_result; +// 123 | var sg_mat_c0_3: subgroup_matrix_result; ss << " var sg_mat_c0_3: subgroup_matrix_result;\n"; -// 120 | #endif +// 124 | #endif } -// 121 | #if sg_mat_count_m >= 2 +// 125 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 122 | var sg_mat_c1_0: subgroup_matrix_result; +// 126 | var sg_mat_c1_0: subgroup_matrix_result; ss << " var sg_mat_c1_0: subgroup_matrix_result;\n"; -// 123 | #if sg_mat_count_n >= 2 +// 127 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 124 | var sg_mat_c1_1: subgroup_matrix_result; +// 128 | var sg_mat_c1_1: subgroup_matrix_result; ss << " var sg_mat_c1_1: subgroup_matrix_result;\n"; -// 125 | #endif +// 129 | #endif } -// 126 | #if sg_mat_count_n >= 3 +// 130 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 127 | var sg_mat_c1_2: subgroup_matrix_result; +// 131 | var sg_mat_c1_2: subgroup_matrix_result; ss << " var sg_mat_c1_2: subgroup_matrix_result;\n"; -// 128 | #endif +// 132 | #endif } -// 129 | #if sg_mat_count_n >= 4 +// 133 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 130 | var sg_mat_c1_3: subgroup_matrix_result; +// 134 | var sg_mat_c1_3: subgroup_matrix_result; ss << " var sg_mat_c1_3: subgroup_matrix_result;\n"; -// 131 | #endif +// 135 | #endif } -// 132 | #endif +// 136 | #endif } -// 133 | #if sg_mat_count_m >= 3 +// 137 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 134 | var sg_mat_c2_0: subgroup_matrix_result; +// 138 | var sg_mat_c2_0: subgroup_matrix_result; ss << " var sg_mat_c2_0: subgroup_matrix_result;\n"; -// 135 | #if sg_mat_count_n >= 2 +// 139 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 136 | var sg_mat_c2_1: subgroup_matrix_result; +// 140 | var sg_mat_c2_1: subgroup_matrix_result; ss << " var sg_mat_c2_1: subgroup_matrix_result;\n"; -// 137 | #endif +// 141 | #endif } -// 138 | #if sg_mat_count_n >= 3 +// 142 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 139 | var sg_mat_c2_2: subgroup_matrix_result; +// 143 | var sg_mat_c2_2: subgroup_matrix_result; ss << " var sg_mat_c2_2: subgroup_matrix_result;\n"; -// 140 | #endif +// 144 | #endif } -// 141 | #if sg_mat_count_n >= 4 +// 145 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 142 | var sg_mat_c2_3: subgroup_matrix_result; +// 146 | var sg_mat_c2_3: subgroup_matrix_result; ss << " var sg_mat_c2_3: subgroup_matrix_result;\n"; -// 143 | #endif +// 147 | #endif } -// 144 | #endif +// 148 | #endif } -// 145 | #if sg_mat_count_m >= 4 +// 149 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 146 | var sg_mat_c3_0: subgroup_matrix_result; +// 150 | var sg_mat_c3_0: subgroup_matrix_result; ss << " var sg_mat_c3_0: subgroup_matrix_result;\n"; -// 147 | #if sg_mat_count_n >= 2 +// 151 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 148 | var sg_mat_c3_1: subgroup_matrix_result; +// 152 | var sg_mat_c3_1: subgroup_matrix_result; ss << " var sg_mat_c3_1: subgroup_matrix_result;\n"; -// 149 | #endif +// 153 | #endif } -// 150 | #if sg_mat_count_n >= 3 +// 154 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 151 | var sg_mat_c3_2: subgroup_matrix_result; +// 155 | var sg_mat_c3_2: subgroup_matrix_result; ss << " var sg_mat_c3_2: subgroup_matrix_result;\n"; -// 152 | #endif +// 156 | #endif } -// 153 | #if sg_mat_count_n >= 4 +// 157 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 154 | var sg_mat_c3_3: subgroup_matrix_result; +// 158 | var sg_mat_c3_3: subgroup_matrix_result; ss << " var sg_mat_c3_3: subgroup_matrix_result;\n"; -// 155 | #endif +// 159 | #endif } -// 156 | #endif +// 160 | #endif } -// 157 | #if sg_mat_count_m >= 5 +// 161 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 158 | var sg_mat_c4_0: subgroup_matrix_result; +// 162 | var sg_mat_c4_0: subgroup_matrix_result; ss << " var sg_mat_c4_0: subgroup_matrix_result;\n"; -// 159 | #if sg_mat_count_n >= 2 +// 163 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 160 | var sg_mat_c4_1: subgroup_matrix_result; +// 164 | var sg_mat_c4_1: subgroup_matrix_result; ss << " var sg_mat_c4_1: subgroup_matrix_result;\n"; -// 161 | #endif +// 165 | #endif } -// 162 | #if sg_mat_count_n >= 3 +// 166 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 163 | var sg_mat_c4_2: subgroup_matrix_result; +// 167 | var sg_mat_c4_2: subgroup_matrix_result; ss << " var sg_mat_c4_2: subgroup_matrix_result;\n"; -// 164 | #endif +// 168 | #endif } -// 165 | #if sg_mat_count_n >= 4 +// 169 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 166 | var sg_mat_c4_3: subgroup_matrix_result; +// 170 | var sg_mat_c4_3: subgroup_matrix_result; ss << " var sg_mat_c4_3: subgroup_matrix_result;\n"; -// 167 | #endif +// 171 | #endif } -// 168 | #endif +// 172 | #endif } -// 169 | #if sg_mat_count_m >= 6 +// 173 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 170 | var sg_mat_c5_0: subgroup_matrix_result; +// 174 | var sg_mat_c5_0: subgroup_matrix_result; ss << " var sg_mat_c5_0: subgroup_matrix_result;\n"; -// 171 | #if sg_mat_count_n >= 2 +// 175 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 172 | var sg_mat_c5_1: subgroup_matrix_result; +// 176 | var sg_mat_c5_1: subgroup_matrix_result; ss << " var sg_mat_c5_1: subgroup_matrix_result;\n"; -// 173 | #endif +// 177 | #endif } -// 174 | #if sg_mat_count_n >= 3 +// 178 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 175 | var sg_mat_c5_2: subgroup_matrix_result; +// 179 | var sg_mat_c5_2: subgroup_matrix_result; ss << " var sg_mat_c5_2: subgroup_matrix_result;\n"; -// 176 | #endif +// 180 | #endif } -// 177 | #if sg_mat_count_n >= 4 +// 181 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 178 | var sg_mat_c5_3: subgroup_matrix_result; +// 182 | var sg_mat_c5_3: subgroup_matrix_result; ss << " var sg_mat_c5_3: subgroup_matrix_result;\n"; -// 179 | #endif +// 183 | #endif } -// 180 | #endif +// 184 | #endif } -// 181 | #if sg_mat_count_m >= 7 +// 185 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 182 | var sg_mat_c6_0: subgroup_matrix_result; +// 186 | var sg_mat_c6_0: subgroup_matrix_result; ss << " var sg_mat_c6_0: subgroup_matrix_result;\n"; -// 183 | #if sg_mat_count_n >= 2 +// 187 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 184 | var sg_mat_c6_1: subgroup_matrix_result; +// 188 | var sg_mat_c6_1: subgroup_matrix_result; ss << " var sg_mat_c6_1: subgroup_matrix_result;\n"; -// 185 | #endif +// 189 | #endif } -// 186 | #if sg_mat_count_n >= 3 +// 190 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 187 | var sg_mat_c6_2: subgroup_matrix_result; +// 191 | var sg_mat_c6_2: subgroup_matrix_result; ss << " var sg_mat_c6_2: subgroup_matrix_result;\n"; -// 188 | #endif +// 192 | #endif } -// 189 | #if sg_mat_count_n >= 4 +// 193 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 190 | var sg_mat_c6_3: subgroup_matrix_result; +// 194 | var sg_mat_c6_3: subgroup_matrix_result; ss << " var sg_mat_c6_3: subgroup_matrix_result;\n"; -// 191 | #endif +// 195 | #endif } -// 192 | #endif +// 196 | #endif } -// 193 | #if sg_mat_count_m >= 8 +// 197 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 194 | var sg_mat_c7_0: subgroup_matrix_result; +// 198 | var sg_mat_c7_0: subgroup_matrix_result; ss << " var sg_mat_c7_0: subgroup_matrix_result;\n"; -// 195 | #if sg_mat_count_n >= 2 +// 199 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 196 | var sg_mat_c7_1: subgroup_matrix_result; +// 200 | var sg_mat_c7_1: subgroup_matrix_result; ss << " var sg_mat_c7_1: subgroup_matrix_result;\n"; -// 197 | #endif +// 201 | #endif } -// 198 | #if sg_mat_count_n >= 3 +// 202 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 199 | var sg_mat_c7_2: subgroup_matrix_result; +// 203 | var sg_mat_c7_2: subgroup_matrix_result; ss << " var sg_mat_c7_2: subgroup_matrix_result;\n"; -// 200 | #endif +// 204 | #endif } -// 201 | #if sg_mat_count_n >= 4 +// 205 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 202 | var sg_mat_c7_3: subgroup_matrix_result; +// 206 | var sg_mat_c7_3: subgroup_matrix_result; ss << " var sg_mat_c7_3: subgroup_matrix_result;\n"; -// 203 | #endif +// 207 | #endif } -// 204 | #endif +// 208 | #endif } -// 205 | -// 206 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, -// 207 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks -// 208 | // early (valid == false), keeping the barriers below in uniform control flow. -// 209 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +// 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"; -// 210 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +// 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"; -// 211 | let kb = sg_index + step * kSplitK; +// 215 | let kb = sg_index + step * kSplitK; ss << " let kb = sg_index + step * kSplitK;\n"; -// 212 | let valid = kb < k_blocks; +// 216 | let valid = kb < k_blocks; ss << " let valid = kb < k_blocks;\n"; -// 213 | +// 217 | ss << "\n"; -// 214 | if (valid) { +// 218 | if (valid) { ss << " if (valid) {\n"; -// 215 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +// 219 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; ss << " let k_vec_per_sgmat_row = kSgMatK / kVecSize;\n"; -// 216 | let k_vec_base = kb * k_vec_per_sgmat_row; +// 220 | let k_vec_base = kb * k_vec_per_sgmat_row; ss << " let k_vec_base = kb * k_vec_per_sgmat_row;\n"; -// 217 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { +// 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"; -// 218 | let row = vidx / k_vec_per_sgmat_row; +// 222 | let row = vidx / k_vec_per_sgmat_row; ss << " let row = vidx / k_vec_per_sgmat_row;\n"; -// 219 | let kc_vec = vidx % k_vec_per_sgmat_row; +// 223 | let kc_vec = vidx % k_vec_per_sgmat_row; ss << " let kc_vec = vidx % k_vec_per_sgmat_row;\n"; -// 220 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); +// 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"; -// 221 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +// 225 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; ss << " let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize;\n"; -// 222 | #if vec_size == 1 +// 226 | #if vec_size == 1 if (__param_vec_size == 1) { -// 223 | a_scratch[dst] = chunk; +// 227 | a_scratch[dst] = chunk; ss << " a_scratch[dst] = chunk;\n"; -// 224 | #else +// 228 | #else } else { -// 225 | a_scratch[dst + 0u] = chunk[0]; +// 229 | a_scratch[dst + 0u] = chunk[0]; ss << " a_scratch[dst + 0u] = chunk[0];\n"; -// 226 | a_scratch[dst + 1u] = chunk[1]; +// 230 | a_scratch[dst + 1u] = chunk[1]; ss << " a_scratch[dst + 1u] = chunk[1];\n"; -// 227 | #if vec_size >= 4 +// 231 | #if vec_size >= 4 if (__param_vec_size >= 4) { -// 228 | a_scratch[dst + 2u] = chunk[2]; +// 232 | a_scratch[dst + 2u] = chunk[2]; ss << " a_scratch[dst + 2u] = chunk[2];\n"; -// 229 | a_scratch[dst + 3u] = chunk[3]; +// 233 | a_scratch[dst + 3u] = chunk[3]; ss << " a_scratch[dst + 3u] = chunk[3];\n"; -// 230 | #endif +// 234 | #endif } -// 231 | #endif +// 235 | #endif } -// 232 | } +// 236 | } ss << " }\n"; -// 233 | } +// 237 | } ss << " }\n"; -// 234 | // Make the staged A-tiles visible before loading them as subgroup matrices. +// 238 | // Make the staged A-tiles visible before loading them as subgroup matrices. ss << "\n"; -// 235 | workgroupBarrier(); +// 239 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 236 | +// 240 | ss << "\n"; -// 237 | if (valid) { +// 241 | if (valid) { ss << " if (valid) {\n"; -// 238 | // Load the B right tiles for this K block. +// 242 | // Load the B right tiles for this K block. ss << "\n"; -// 239 | #if b_is_hwio +// 243 | #if b_is_hwio if (__param_b_is_hwio) { -// 240 | // HWIO: B is row-major ([K, N_b], B[k, n] at k*N_b + n); load row-major, -// 241 | // stride N_b. N_b (== N unless padded) keeps the row stride, and via -// 242 | // global_base_n above, the column offset, 4-byte aligned. -// 243 | let b_base = kb * kSgMatK * uniforms.N_b + global_base_n; +// 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"; -// 244 | var sg_mat_b0: subgroup_matrix_right = +// 248 | var sg_mat_b0: subgroup_matrix_right = ss << " var sg_mat_b0: subgroup_matrix_right =\n"; -// 245 | subgroupMatrixLoad, row_major>( +// 249 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 246 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +// 250 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 0 * kSgMatN, uniforms.N_b);\n"; -// 247 | #if sg_mat_count_n >= 2 +// 251 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 248 | var sg_mat_b1: subgroup_matrix_right = +// 252 | var sg_mat_b1: subgroup_matrix_right = ss << " var sg_mat_b1: subgroup_matrix_right =\n"; -// 249 | subgroupMatrixLoad, row_major>( +// 253 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 250 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); +// 254 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 1 * kSgMatN, uniforms.N_b);\n"; -// 251 | #endif +// 255 | #endif } -// 252 | #if sg_mat_count_n >= 3 +// 256 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 253 | var sg_mat_b2: subgroup_matrix_right = +// 257 | var sg_mat_b2: subgroup_matrix_right = ss << " var sg_mat_b2: subgroup_matrix_right =\n"; -// 254 | subgroupMatrixLoad, row_major>( +// 258 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 255 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +// 259 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 2 * kSgMatN, uniforms.N_b);\n"; -// 256 | #endif +// 260 | #endif } -// 257 | #if sg_mat_count_n >= 4 +// 261 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 258 | var sg_mat_b3: subgroup_matrix_right = +// 262 | var sg_mat_b3: subgroup_matrix_right = ss << " var sg_mat_b3: subgroup_matrix_right =\n"; -// 259 | subgroupMatrixLoad, row_major>( +// 263 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 260 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +// 264 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 3 * kSgMatN, uniforms.N_b);\n"; -// 261 | #endif +// 265 | #endif } -// 262 | #else +// 266 | #else } else { -// 263 | // OHWI: B is column-major ([N, K], B[k, n] at n*K + k); load col-major, stride K. -// 264 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +// 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"; -// 265 | var sg_mat_b0: subgroup_matrix_right = +// 269 | var sg_mat_b0: subgroup_matrix_right = ss << " var sg_mat_b0: subgroup_matrix_right =\n"; -// 266 | subgroupMatrixLoad, col_major>( +// 270 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 267 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +// 271 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 268 | #if sg_mat_count_n >= 2 +// 272 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 269 | var sg_mat_b1: subgroup_matrix_right = +// 273 | var sg_mat_b1: subgroup_matrix_right = ss << " var sg_mat_b1: subgroup_matrix_right =\n"; -// 270 | subgroupMatrixLoad, col_major>( +// 274 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 271 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); +// 275 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 272 | #endif +// 276 | #endif } -// 273 | #if sg_mat_count_n >= 3 +// 277 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 274 | var sg_mat_b2: subgroup_matrix_right = +// 278 | var sg_mat_b2: subgroup_matrix_right = ss << " var sg_mat_b2: subgroup_matrix_right =\n"; -// 275 | subgroupMatrixLoad, col_major>( +// 279 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 276 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +// 280 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 277 | #endif +// 281 | #endif } -// 278 | #if sg_mat_count_n >= 4 +// 282 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 279 | var sg_mat_b3: subgroup_matrix_right = +// 283 | var sg_mat_b3: subgroup_matrix_right = ss << " var sg_mat_b3: subgroup_matrix_right =\n"; -// 280 | subgroupMatrixLoad, col_major>( +// 284 | subgroupMatrixLoad, col_major>( ss << " subgroupMatrixLoad, col_major>(\n"; -// 281 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +// 285 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); ss << " &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K);\n"; -// 282 | #endif +// 286 | #endif } -// 283 | #endif +// 287 | #endif } -// 284 | -// 285 | // Load the A left tiles (one per M block) from workgroup staging, -// 286 | // row-major with stride kSgMatK. -// 287 | var sg_mat_a0: subgroup_matrix_left = +// 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"; -// 288 | subgroupMatrixLoad, row_major>( +// 292 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 289 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +// 293 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 290 | #if sg_mat_count_m >= 2 +// 294 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 291 | var sg_mat_a1: subgroup_matrix_left = +// 295 | var sg_mat_a1: subgroup_matrix_left = ss << " var sg_mat_a1: subgroup_matrix_left =\n"; -// 292 | subgroupMatrixLoad, row_major>( +// 296 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 293 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); +// 297 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 294 | #endif +// 298 | #endif } -// 295 | #if sg_mat_count_m >= 3 +// 299 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 296 | var sg_mat_a2: subgroup_matrix_left = +// 300 | var sg_mat_a2: subgroup_matrix_left = ss << " var sg_mat_a2: subgroup_matrix_left =\n"; -// 297 | subgroupMatrixLoad, row_major>( +// 301 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 298 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +// 302 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 299 | #endif +// 303 | #endif } -// 300 | #if sg_mat_count_m >= 4 +// 304 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 301 | var sg_mat_a3: subgroup_matrix_left = +// 305 | var sg_mat_a3: subgroup_matrix_left = ss << " var sg_mat_a3: subgroup_matrix_left =\n"; -// 302 | subgroupMatrixLoad, row_major>( +// 306 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 303 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +// 307 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 304 | #endif +// 308 | #endif } -// 305 | #if sg_mat_count_m >= 5 +// 309 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 306 | var sg_mat_a4: subgroup_matrix_left = +// 310 | var sg_mat_a4: subgroup_matrix_left = ss << " var sg_mat_a4: subgroup_matrix_left =\n"; -// 307 | subgroupMatrixLoad, row_major>( +// 311 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 308 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +// 312 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 309 | #endif +// 313 | #endif } -// 310 | #if sg_mat_count_m >= 6 +// 314 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 311 | var sg_mat_a5: subgroup_matrix_left = +// 315 | var sg_mat_a5: subgroup_matrix_left = ss << " var sg_mat_a5: subgroup_matrix_left =\n"; -// 312 | subgroupMatrixLoad, row_major>( +// 316 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 313 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +// 317 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 314 | #endif +// 318 | #endif } -// 315 | #if sg_mat_count_m >= 7 +// 319 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 316 | var sg_mat_a6: subgroup_matrix_left = +// 320 | var sg_mat_a6: subgroup_matrix_left = ss << " var sg_mat_a6: subgroup_matrix_left =\n"; -// 317 | subgroupMatrixLoad, row_major>( +// 321 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 318 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +// 322 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 319 | #endif +// 323 | #endif } -// 320 | #if sg_mat_count_m >= 8 +// 324 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 321 | var sg_mat_a7: subgroup_matrix_left = +// 325 | var sg_mat_a7: subgroup_matrix_left = ss << " var sg_mat_a7: subgroup_matrix_left =\n"; -// 322 | subgroupMatrixLoad, row_major>( +// 326 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 323 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +// 327 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); ss << " &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK);\n"; -// 324 | #endif +// 328 | #endif } -// 325 | -// 326 | // Accumulate every (M block, N block) pair. -// 327 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +// 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"; -// 328 | #if sg_mat_count_n >= 2 +// 332 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 329 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +// 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"; -// 330 | #endif +// 334 | #endif } -// 331 | #if sg_mat_count_n >= 3 +// 335 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 332 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +// 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"; -// 333 | #endif +// 337 | #endif } -// 334 | #if sg_mat_count_n >= 4 +// 338 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 335 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +// 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"; -// 336 | #endif +// 340 | #endif } -// 337 | #if sg_mat_count_m >= 2 +// 341 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 338 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +// 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"; -// 339 | #if sg_mat_count_n >= 2 +// 343 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 340 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +// 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"; -// 341 | #endif +// 345 | #endif } -// 342 | #if sg_mat_count_n >= 3 +// 346 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 343 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +// 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"; -// 344 | #endif +// 348 | #endif } -// 345 | #if sg_mat_count_n >= 4 +// 349 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 346 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +// 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"; -// 347 | #endif +// 351 | #endif } -// 348 | #endif +// 352 | #endif } -// 349 | #if sg_mat_count_m >= 3 +// 353 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 350 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +// 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"; -// 351 | #if sg_mat_count_n >= 2 +// 355 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 352 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +// 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"; -// 353 | #endif +// 357 | #endif } -// 354 | #if sg_mat_count_n >= 3 +// 358 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 355 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +// 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"; -// 356 | #endif +// 360 | #endif } -// 357 | #if sg_mat_count_n >= 4 +// 361 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 358 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +// 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"; -// 359 | #endif +// 363 | #endif } -// 360 | #endif +// 364 | #endif } -// 361 | #if sg_mat_count_m >= 4 +// 365 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 362 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +// 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"; -// 363 | #if sg_mat_count_n >= 2 +// 367 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 364 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +// 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"; -// 365 | #endif +// 369 | #endif } -// 366 | #if sg_mat_count_n >= 3 +// 370 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 367 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +// 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"; -// 368 | #endif +// 372 | #endif } -// 369 | #if sg_mat_count_n >= 4 +// 373 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 370 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +// 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"; -// 371 | #endif +// 375 | #endif } -// 372 | #endif +// 376 | #endif } -// 373 | #if sg_mat_count_m >= 5 +// 377 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 374 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +// 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"; -// 375 | #if sg_mat_count_n >= 2 +// 379 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 376 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +// 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"; -// 377 | #endif +// 381 | #endif } -// 378 | #if sg_mat_count_n >= 3 +// 382 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 379 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +// 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"; -// 380 | #endif +// 384 | #endif } -// 381 | #if sg_mat_count_n >= 4 +// 385 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 382 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +// 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"; -// 383 | #endif +// 387 | #endif } -// 384 | #endif +// 388 | #endif } -// 385 | #if sg_mat_count_m >= 6 +// 389 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 386 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +// 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"; -// 387 | #if sg_mat_count_n >= 2 +// 391 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 388 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +// 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"; -// 389 | #endif +// 393 | #endif } -// 390 | #if sg_mat_count_n >= 3 +// 394 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 391 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +// 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"; -// 392 | #endif +// 396 | #endif } -// 393 | #if sg_mat_count_n >= 4 +// 397 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 394 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +// 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"; -// 395 | #endif +// 399 | #endif } -// 396 | #endif +// 400 | #endif } -// 397 | #if sg_mat_count_m >= 7 +// 401 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 398 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +// 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"; -// 399 | #if sg_mat_count_n >= 2 +// 403 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 400 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +// 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"; -// 401 | #endif +// 405 | #endif } -// 402 | #if sg_mat_count_n >= 3 +// 406 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 403 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +// 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"; -// 404 | #endif +// 408 | #endif } -// 405 | #if sg_mat_count_n >= 4 +// 409 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 406 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +// 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"; -// 407 | #endif +// 411 | #endif } -// 408 | #endif +// 412 | #endif } -// 409 | #if sg_mat_count_m >= 8 +// 413 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 410 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +// 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"; -// 411 | #if sg_mat_count_n >= 2 +// 415 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 412 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +// 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"; -// 413 | #endif +// 417 | #endif } -// 414 | #if sg_mat_count_n >= 3 +// 418 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 415 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +// 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"; -// 416 | #endif +// 420 | #endif } -// 417 | #if sg_mat_count_n >= 4 +// 421 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 418 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +// 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"; -// 419 | #endif +// 423 | #endif } -// 420 | #endif +// 424 | #endif } -// 421 | } +// 425 | } ss << " }\n"; -// 422 | // Ensure every subgroup is done reading a_scratch before it is reused for +// 426 | // Ensure every subgroup is done reading a_scratch before it is reused for ss << "\n"; -// 423 | // the next K-block step. -// 424 | workgroupBarrier(); +// 427 | // the next K-block step. +// 428 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 425 | } +// 429 | } ss << " }\n"; -// 426 | +// 430 | ss << "\n"; -// 427 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; -// 428 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). -// 429 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +// 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"; -// 430 | #if sg_mat_count_n >= 2 +// 434 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 431 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +// 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"; -// 432 | #endif +// 436 | #endif } -// 433 | #if sg_mat_count_n >= 3 +// 437 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 434 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +// 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"; -// 435 | #endif +// 439 | #endif } -// 436 | #if sg_mat_count_n >= 4 +// 440 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 437 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +// 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"; -// 438 | #endif +// 442 | #endif } -// 439 | #if sg_mat_count_m >= 2 +// 443 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 440 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +// 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"; -// 441 | #if sg_mat_count_n >= 2 +// 445 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 442 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +// 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"; -// 443 | #endif +// 447 | #endif } -// 444 | #if sg_mat_count_n >= 3 +// 448 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 445 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +// 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"; -// 446 | #endif +// 450 | #endif } -// 447 | #if sg_mat_count_n >= 4 +// 451 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 448 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +// 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"; -// 449 | #endif +// 453 | #endif } -// 450 | #endif +// 454 | #endif } -// 451 | #if sg_mat_count_m >= 3 +// 455 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +// 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"; -// 453 | #if sg_mat_count_n >= 2 +// 457 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 454 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +// 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"; -// 455 | #endif +// 459 | #endif } -// 456 | #if sg_mat_count_n >= 3 +// 460 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 457 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +// 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"; -// 458 | #endif +// 462 | #endif } -// 459 | #if sg_mat_count_n >= 4 +// 463 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 460 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +// 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"; -// 461 | #endif +// 465 | #endif } -// 462 | #endif +// 466 | #endif } -// 463 | #if sg_mat_count_m >= 4 +// 467 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +// 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"; -// 465 | #if sg_mat_count_n >= 2 +// 469 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 466 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +// 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"; -// 467 | #endif +// 471 | #endif } -// 468 | #if sg_mat_count_n >= 3 +// 472 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 469 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +// 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"; -// 470 | #endif +// 474 | #endif } -// 471 | #if sg_mat_count_n >= 4 +// 475 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 472 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +// 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"; -// 473 | #endif +// 477 | #endif } -// 474 | #endif +// 478 | #endif } -// 475 | #if sg_mat_count_m >= 5 +// 479 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 476 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +// 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"; -// 477 | #if sg_mat_count_n >= 2 +// 481 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 478 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +// 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"; -// 479 | #endif +// 483 | #endif } -// 480 | #if sg_mat_count_n >= 3 +// 484 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 481 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +// 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"; -// 482 | #endif +// 486 | #endif } -// 483 | #if sg_mat_count_n >= 4 +// 487 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 484 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +// 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"; -// 485 | #endif +// 489 | #endif } -// 486 | #endif +// 490 | #endif } -// 487 | #if sg_mat_count_m >= 6 +// 491 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 488 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +// 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"; -// 489 | #if sg_mat_count_n >= 2 +// 493 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 490 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +// 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"; -// 491 | #endif +// 495 | #endif } -// 492 | #if sg_mat_count_n >= 3 +// 496 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 493 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +// 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"; -// 494 | #endif +// 498 | #endif } -// 495 | #if sg_mat_count_n >= 4 +// 499 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 496 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +// 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"; -// 497 | #endif +// 501 | #endif } -// 498 | #endif +// 502 | #endif } -// 499 | #if sg_mat_count_m >= 7 +// 503 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 500 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +// 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"; -// 501 | #if sg_mat_count_n >= 2 +// 505 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 502 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +// 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"; -// 503 | #endif +// 507 | #endif } -// 504 | #if sg_mat_count_n >= 3 +// 508 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 505 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +// 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"; -// 506 | #endif +// 510 | #endif } -// 507 | #if sg_mat_count_n >= 4 +// 511 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 508 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +// 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"; -// 509 | #endif +// 513 | #endif } -// 510 | #endif +// 514 | #endif } -// 511 | #if sg_mat_count_m >= 8 +// 515 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 512 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +// 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"; -// 513 | #if sg_mat_count_n >= 2 +// 517 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 514 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +// 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"; -// 515 | #endif +// 519 | #endif } -// 516 | #if sg_mat_count_n >= 3 +// 520 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 517 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +// 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"; -// 518 | #endif +// 522 | #endif } -// 519 | #if sg_mat_count_n >= 4 +// 523 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 520 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +// 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"; -// 521 | #endif +// 525 | #endif } -// 522 | #endif +// 526 | #endif } -// 523 | -// 524 | #if split_k >= 2 +// 527 | +// 528 | #if split_k >= 2 if (__param_split_k >= 2) { -// 525 | // Make all subgroups' partial tiles visible before the reduction. -// 526 | workgroupBarrier(); +// 529 | // Make all subgroups' partial tiles visible before the reduction. +// 530 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 527 | +// 531 | ss << "\n"; -// 528 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each -// 529 | // position is only overwritten by the lane that already read it. -// 530 | if (sg_index == 0u) { +// 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"; -// 531 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +// 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"; -// 532 | var acc: f16 = f16(0); +// 536 | var acc: f16 = f16(0); ss << " var acc: f16 = f16(0);\n"; -// 533 | for (var s: u32 = 0; s < kSplitK; s++) { +// 537 | for (var s: u32 = 0; s < kSplitK; s++) { ss << " for (var s: u32 = 0; s < kSplitK; s++) {\n"; -// 534 | acc += scratch[s * kTileM * kTileN + idx]; +// 538 | acc += scratch[s * kTileM * kTileN + idx]; ss << " acc += scratch[s * kTileM * kTileN + idx];\n"; -// 535 | } +// 539 | } ss << " }\n"; -// 536 | scratch[idx] = acc; +// 540 | scratch[idx] = acc; ss << " scratch[idx] = acc;\n"; -// 537 | } +// 541 | } ss << " }\n"; -// 538 | } +// 542 | } ss << " }\n"; -// 539 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +// 543 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << "\n"; -// 540 | workgroupBarrier(); +// 544 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 541 | #endif +// 545 | #endif } -// 542 | -// 543 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip -// 544 | // and M's bounds check make non-multiple dimensions write each element once. -// 545 | let n_count = min(kTileN, uniforms.N - global_base_n); +// 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"; -// 546 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +// 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"; -// 547 | let global_m = m_base + r; +// 551 | let global_m = m_base + r; ss << " let global_m = m_base + r;\n"; -// 548 | if (global_m < uniforms.M) { +// 552 | if (global_m < uniforms.M) { ss << " if (global_m < uniforms.M) {\n"; -// 549 | let scratch_base = r * kTileN; +// 553 | let scratch_base = r * kTileN; ss << " let scratch_base = r * kTileN;\n"; -// 550 | let out_base = out_batch_base + global_m * uniforms.N + global_base_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"; -// 551 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +// 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"; -// 552 | write_output(out_base + i, global_base_n + i, +// 556 | write_output(out_base + i, global_base_n + i, ss << " write_output(out_base + i, global_base_n + i,\n"; -// 553 | output_value_t(scratch[scratch_base + i])); +// 557 | output_value_t(scratch[scratch_base + i])); ss << " output_value_t(scratch[scratch_base + i]));\n"; -// 554 | } +// 558 | } ss << " }\n"; -// 555 | } +// 559 | } ss << " }\n"; -// 556 | } +// 560 | } ss << " }\n"; -// 557 | } // MAIN +// 561 | } // MAIN MainFunctionEnd(); ss << "\n"; -// 558 | +// 562 | return Status::OK(); 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 b66530c7507b4..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 @@ -40,7 +40,7 @@ std::string pass_as_string(T&& v) { #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" // 8e5c73e9849494d81da110da2fb808dddc4cdfae6eb96d054b205c8f8c505e51 +#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/nn/subgroup_matrix_conv_8x16x16.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/subgroup_matrix_conv_8x16x16.h index 9e47cc73911b2..79282ca894ed7 100644 --- 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 @@ -46,1296 +46,1300 @@ Status ApplyTemplate<"nn/subgroup_matrix_conv_8x16x16.wgsl.template">(ShaderHelp // 22 | // // 23 | // Bias/activation: write_output(), built by the host from GetActivationSnippet. // 24 | // -// 25 | // Host guarantees: K % sg_mat_k == 0; kTileN <= N_b so a trailing partial N tile -// 26 | // can shift back in bounds. M is unconstrained (see load_im2col_vec). -// 27 | -// 28 | #param sg_mat_k -// 29 | #param sg_mat_m -// 30 | #param sg_mat_n -// 31 | #param sg_mat_count_m -// 32 | #param sg_mat_count_n -// 33 | #param split_k -// 34 | #param vec_size -// 35 | #param b_is_hwio -// 36 | -// 37 | #use .getByOffset .setByOffset -// 38 | -// 39 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +// 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; -// 40 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +// 44 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) ss << __str_33; ss << __param_sg_mat_n; ss << __str_17; -// 41 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +// 45 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) ss << __str_34; ss << __param_sg_mat_k; ss << __str_17; -// 42 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +// 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; -// 43 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +// 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; -// 44 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) +// 48 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) ss << __str_37; -// 45 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) +// 49 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) ss << __str_38; -// 46 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +// 50 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K ss << __str_39; ss << __param_split_k; ss << __str_17; -// 47 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +// 51 | const kSubgroupSize: u32 = 32; // Lanes per subgroup ss << __str_40; -// 48 | const kVecSize: u32 = u32(vec_size); // im2col gather width along cin (cin % vec_size == 0) +// 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; -// 49 | +// 53 | ss << __str_5; -// 50 | // kSplitK partial [kTileM, kTileN] tiles, row-major; one slot per split-K -// 51 | // subgroup, summed at write-back. -// 52 | var scratch: array; +// 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; -// 53 | +// 57 | ss << __str_5; -// 54 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. -// 55 | var a_scratch: array; +// 58 | // Per-subgroup [kTileM, kSgMatK] staging for the current K-block A-tile. +// 59 | var a_scratch: array; ss << __str_408; -// 56 | +// 60 | ss << __str_5; -// 57 | // Gathers A[m, k_base..k_base+vec_size-1] (k_base = k_vec * vec_size), always -// 58 | // within one (kh, kw) position, so all vec lanes share one bounds check. Padding -// 59 | // reads as zero; m past M is index-safe (not necessarily zero) but never stored. -// 60 | // m -> (oh, ow); k_vec -> (kh_idx, kw_idx, k_cin), k_cin fastest. -// 61 | fn load_im2col_vec(m: u32, k_vec: u32, input_batch_base_vec: u32) -> input_src_value_t { +// 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; -// 62 | let cin_vec = uniforms.cin / kVecSize; +// 66 | let cin_vec = uniforms.cin / kVecSize; ss << __str_410; -// 63 | let oh = m / uniforms.output_w; +// 67 | let oh = m / uniforms.output_w; ss << __str_411; -// 64 | let ow = m % uniforms.output_w; +// 68 | let ow = m % uniforms.output_w; ss << __str_412; -// 65 | let k_cin = k_vec % cin_vec; +// 69 | let k_cin = k_vec % cin_vec; ss << __str_413; -// 66 | let k_hw = k_vec / cin_vec; +// 70 | let k_hw = k_vec / cin_vec; ss << __str_414; -// 67 | let kw_idx = k_hw % uniforms.kernel_w; +// 71 | let kw_idx = k_hw % uniforms.kernel_w; ss << __str_415; -// 68 | let kh_idx = k_hw / uniforms.kernel_w; +// 72 | let kh_idx = k_hw / uniforms.kernel_w; ss << __str_416; -// 69 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); +// 73 | let ih = i32(oh * uniforms.strides.x + kh_idx * uniforms.dilations.x) - i32(uniforms.pads.x); ss << __str_417; -// 70 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); +// 74 | let iw = i32(ow * uniforms.strides.y + kw_idx * uniforms.dilations.y) - i32(uniforms.pads.y); ss << __str_418; -// 71 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { +// 75 | if (ih < 0 || ih >= i32(uniforms.input_h) || iw < 0 || iw >= i32(uniforms.input_w)) { ss << __str_419; -// 72 | return input_src_value_t(); +// 76 | return input_src_value_t(); ss << __str_420; -// 73 | } +// 77 | } ss << __str_20; -// 74 | let input_vec_idx = input_batch_base_vec + +// 78 | let input_vec_idx = input_batch_base_vec + ss << __str_421; -// 75 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; +// 79 | (u32(ih) * uniforms.input_w + u32(iw)) * cin_vec + k_cin; ss << __str_422; -// 76 | return input_src.getByOffset(input_vec_idx); +// 80 | return input_src.getByOffset(input_vec_idx); ss << __str_351; ss << __var_input_src->GetByOffset(__str_406); ss << __str_29; -// 77 | } +// 81 | } ss << __str_345; -// 78 | +// 82 | ss << __str_5; -// 79 | $MAIN { +// 83 | $MAIN { MainFunctionStart(); ss << __str_5; -// 80 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is +// 84 | // workgroup_idx (not workgroup_id.x/y/z, which may be normalized/reshaped) is ss << __str_5; -// 81 | // the true linear index used to recover batch_id and (n_tile, m_tile). -// 82 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +// 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; -// 83 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; +// 87 | let tiles_per_batch = uniforms.num_n_tile * num_m_tile; ss << __str_423; -// 84 | let batch_id = workgroup_idx / tiles_per_batch; +// 88 | let batch_id = workgroup_idx / tiles_per_batch; ss << __str_424; -// 85 | // Dispatch normalization can add extra workgroups past the last batch element; +// 89 | // Dispatch normalization can add extra workgroups past the last batch element; ss << __str_5; -// 86 | // drop them here (workgroup-uniform, so the barriers below stay uniform). -// 87 | if (batch_id >= uniforms.batch) { +// 90 | // drop them here (workgroup-uniform, so the barriers below stay uniform). +// 91 | if (batch_id >= uniforms.batch) { ss << __str_425; -// 88 | return; +// 92 | return; ss << __str_426; -// 89 | } +// 93 | } ss << __str_20; -// 90 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element +// 94 | let tile_idx = workgroup_idx % tiles_per_batch; // tile index within this batch element ss << __str_427; -// 91 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile +// 95 | let n_tile = tile_idx % uniforms.num_n_tile; // which kTileN-wide column tile ss << __str_428; -// 92 | let n_tile_start = n_tile * kTileN; +// 96 | let n_tile_start = n_tile * kTileN; ss << __str_45; -// 93 | // Shift a trailing partial N tile back so unchecked B loads stay inside the +// 97 | // Shift a trailing partial N tile back so unchecked B loads stay inside the ss << __str_5; -// 94 | // weight buffer (max() just guards the u32 subtraction; host clamps kTileN <= N). -// 95 | // Uses N_b (== N unless B was padded for row-major alignment) so the shifted -// 96 | // origin -- itself n_align-even since N_b, kTileN and n_tile_start all are -- -// 97 | // stays alignment-safe for the HWIO load below. -// 98 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); +// 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; -// 99 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile +// 103 | let n_skip = n_tile_start - global_base_n; // columns already written by the previous tile ss << __str_205; -// 100 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile +// 104 | let m_tile = tile_idx / uniforms.num_n_tile; // which kTileM-tall row tile ss << __str_429; -// 101 | let m_base = m_tile * kTileM; +// 105 | let m_base = m_tile * kTileM; ss << __str_430; -// 102 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); +// 106 | let input_batch_base_vec = batch_id * uniforms.input_h * uniforms.input_w * (uniforms.cin / kVecSize); ss << __str_431; -// 103 | let out_batch_base = batch_id * uniforms.M * uniforms.N; +// 107 | let out_batch_base = batch_id * uniforms.M * uniforms.N; ss << __str_432; -// 104 | let k_blocks = uniforms.K / kSgMatK; +// 108 | let k_blocks = uniforms.K / kSgMatK; ss << __str_49; -// 105 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +// 109 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) ss << __str_50; -// 106 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +// 110 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) ss << __str_51; -// 107 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot +// 111 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's output scratch slot ss << __str_52; -// 108 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot +// 112 | let a_scratch_base = sg_index * kTileM * kSgMatK; // this subgroup's A staging slot ss << __str_433; -// 109 | +// 113 | ss << __str_5; -// 110 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). -// 111 | var sg_mat_c0_0: subgroup_matrix_result; +// 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; -// 112 | #if sg_mat_count_n >= 2 +// 116 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 113 | var sg_mat_c0_1: subgroup_matrix_result; +// 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; -// 114 | #endif +// 118 | #endif } -// 115 | #if sg_mat_count_n >= 3 +// 119 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 116 | var sg_mat_c0_2: subgroup_matrix_result; +// 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; -// 117 | #endif +// 121 | #endif } -// 118 | #if sg_mat_count_n >= 4 +// 122 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 119 | var sg_mat_c0_3: subgroup_matrix_result; +// 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; -// 120 | #endif +// 124 | #endif } -// 121 | #if sg_mat_count_m >= 2 +// 125 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 122 | var sg_mat_c1_0: subgroup_matrix_result; +// 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; -// 123 | #if sg_mat_count_n >= 2 +// 127 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 124 | var sg_mat_c1_1: subgroup_matrix_result; +// 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; -// 125 | #endif +// 129 | #endif } -// 126 | #if sg_mat_count_n >= 3 +// 130 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 127 | var sg_mat_c1_2: subgroup_matrix_result; +// 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; -// 128 | #endif +// 132 | #endif } -// 129 | #if sg_mat_count_n >= 4 +// 133 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 130 | var sg_mat_c1_3: subgroup_matrix_result; +// 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; -// 131 | #endif +// 135 | #endif } -// 132 | #endif +// 136 | #endif } -// 133 | #if sg_mat_count_m >= 3 +// 137 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 134 | var sg_mat_c2_0: subgroup_matrix_result; +// 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; -// 135 | #if sg_mat_count_n >= 2 +// 139 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 136 | var sg_mat_c2_1: subgroup_matrix_result; +// 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; -// 137 | #endif +// 141 | #endif } -// 138 | #if sg_mat_count_n >= 3 +// 142 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 139 | var sg_mat_c2_2: subgroup_matrix_result; +// 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; -// 140 | #endif +// 144 | #endif } -// 141 | #if sg_mat_count_n >= 4 +// 145 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 142 | var sg_mat_c2_3: subgroup_matrix_result; +// 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; -// 143 | #endif +// 147 | #endif } -// 144 | #endif +// 148 | #endif } -// 145 | #if sg_mat_count_m >= 4 +// 149 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 146 | var sg_mat_c3_0: subgroup_matrix_result; +// 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; -// 147 | #if sg_mat_count_n >= 2 +// 151 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 148 | var sg_mat_c3_1: subgroup_matrix_result; +// 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; -// 149 | #endif +// 153 | #endif } -// 150 | #if sg_mat_count_n >= 3 +// 154 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 151 | var sg_mat_c3_2: subgroup_matrix_result; +// 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; -// 152 | #endif +// 156 | #endif } -// 153 | #if sg_mat_count_n >= 4 +// 157 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 154 | var sg_mat_c3_3: subgroup_matrix_result; +// 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; -// 155 | #endif +// 159 | #endif } -// 156 | #endif +// 160 | #endif } -// 157 | #if sg_mat_count_m >= 5 +// 161 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 158 | var sg_mat_c4_0: subgroup_matrix_result; +// 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; -// 159 | #if sg_mat_count_n >= 2 +// 163 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 160 | var sg_mat_c4_1: subgroup_matrix_result; +// 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; -// 161 | #endif +// 165 | #endif } -// 162 | #if sg_mat_count_n >= 3 +// 166 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 163 | var sg_mat_c4_2: subgroup_matrix_result; +// 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; -// 164 | #endif +// 168 | #endif } -// 165 | #if sg_mat_count_n >= 4 +// 169 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 166 | var sg_mat_c4_3: subgroup_matrix_result; +// 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; -// 167 | #endif +// 171 | #endif } -// 168 | #endif +// 172 | #endif } -// 169 | #if sg_mat_count_m >= 6 +// 173 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 170 | var sg_mat_c5_0: subgroup_matrix_result; +// 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; -// 171 | #if sg_mat_count_n >= 2 +// 175 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 172 | var sg_mat_c5_1: subgroup_matrix_result; +// 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; -// 173 | #endif +// 177 | #endif } -// 174 | #if sg_mat_count_n >= 3 +// 178 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 175 | var sg_mat_c5_2: subgroup_matrix_result; +// 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; -// 176 | #endif +// 180 | #endif } -// 177 | #if sg_mat_count_n >= 4 +// 181 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 178 | var sg_mat_c5_3: subgroup_matrix_result; +// 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; -// 179 | #endif +// 183 | #endif } -// 180 | #endif +// 184 | #endif } -// 181 | #if sg_mat_count_m >= 7 +// 185 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 182 | var sg_mat_c6_0: subgroup_matrix_result; +// 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; -// 183 | #if sg_mat_count_n >= 2 +// 187 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 184 | var sg_mat_c6_1: subgroup_matrix_result; +// 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; -// 185 | #endif +// 189 | #endif } -// 186 | #if sg_mat_count_n >= 3 +// 190 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 187 | var sg_mat_c6_2: subgroup_matrix_result; +// 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; -// 188 | #endif +// 192 | #endif } -// 189 | #if sg_mat_count_n >= 4 +// 193 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 190 | var sg_mat_c6_3: subgroup_matrix_result; +// 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; -// 191 | #endif +// 195 | #endif } -// 192 | #endif +// 196 | #endif } -// 193 | #if sg_mat_count_m >= 8 +// 197 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 194 | var sg_mat_c7_0: subgroup_matrix_result; +// 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; -// 195 | #if sg_mat_count_n >= 2 +// 199 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 196 | var sg_mat_c7_1: subgroup_matrix_result; +// 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; -// 197 | #endif +// 201 | #endif } -// 198 | #if sg_mat_count_n >= 3 +// 202 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 199 | var sg_mat_c7_2: subgroup_matrix_result; +// 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; -// 200 | #endif +// 204 | #endif } -// 201 | #if sg_mat_count_n >= 4 +// 205 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 202 | var sg_mat_c7_3: subgroup_matrix_result; +// 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; -// 203 | #endif +// 207 | #endif } -// 204 | #endif +// 208 | #endif } -// 205 | -// 206 | // K reduction: subgroup sg_index processes kb = sg_index, sg_index + kSplitK, -// 207 | // ... num_steps is workgroup-uniform even if a subgroup runs out of K-blocks -// 208 | // early (valid == false), keeping the barriers below in uniform control flow. -// 209 | let num_steps = (k_blocks + kSplitK - 1u) / kSplitK; +// 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; -// 210 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { +// 214 | for (var step: u32 = 0u; step < num_steps; step = step + 1u) { ss << __str_435; -// 211 | let kb = sg_index + step * kSplitK; +// 215 | let kb = sg_index + step * kSplitK; ss << __str_436; -// 212 | let valid = kb < k_blocks; +// 216 | let valid = kb < k_blocks; ss << __str_437; -// 213 | +// 217 | ss << __str_5; -// 214 | if (valid) { +// 218 | if (valid) { ss << __str_438; -// 215 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; +// 219 | let k_vec_per_sgmat_row = kSgMatK / kVecSize; ss << __str_439; -// 216 | let k_vec_base = kb * k_vec_per_sgmat_row; +// 220 | let k_vec_base = kb * k_vec_per_sgmat_row; ss << __str_440; -// 217 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { +// 221 | for (var vidx: u32 = sg_lane; vidx < kTileM * k_vec_per_sgmat_row; vidx = vidx + kSubgroupSize) { ss << __str_441; -// 218 | let row = vidx / k_vec_per_sgmat_row; +// 222 | let row = vidx / k_vec_per_sgmat_row; ss << __str_442; -// 219 | let kc_vec = vidx % k_vec_per_sgmat_row; +// 223 | let kc_vec = vidx % k_vec_per_sgmat_row; ss << __str_443; -// 220 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); +// 224 | let chunk = load_im2col_vec(m_base + row, k_vec_base + kc_vec, input_batch_base_vec); ss << __str_444; -// 221 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; +// 225 | let dst = a_scratch_base + row * kSgMatK + kc_vec * kVecSize; ss << __str_445; -// 222 | #if vec_size == 1 +// 226 | #if vec_size == 1 if (__param_vec_size == 1) { -// 223 | a_scratch[dst] = chunk; +// 227 | a_scratch[dst] = chunk; ss << __str_446; -// 224 | #else +// 228 | #else } else { -// 225 | a_scratch[dst + 0u] = chunk[0]; +// 229 | a_scratch[dst + 0u] = chunk[0]; ss << __str_447; -// 226 | a_scratch[dst + 1u] = chunk[1]; +// 230 | a_scratch[dst + 1u] = chunk[1]; ss << __str_448; -// 227 | #if vec_size >= 4 +// 231 | #if vec_size >= 4 if (__param_vec_size >= 4) { -// 228 | a_scratch[dst + 2u] = chunk[2]; +// 232 | a_scratch[dst + 2u] = chunk[2]; ss << __str_449; -// 229 | a_scratch[dst + 3u] = chunk[3]; +// 233 | a_scratch[dst + 3u] = chunk[3]; ss << __str_450; -// 230 | #endif +// 234 | #endif } -// 231 | #endif +// 235 | #endif } -// 232 | } +// 236 | } ss << __str_202; -// 233 | } +// 237 | } ss << __str_204; -// 234 | // Make the staged A-tiles visible before loading them as subgroup matrices. +// 238 | // Make the staged A-tiles visible before loading them as subgroup matrices. ss << __str_5; -// 235 | workgroupBarrier(); +// 239 | workgroupBarrier(); ss << __str_451; -// 236 | +// 240 | ss << __str_5; -// 237 | if (valid) { +// 241 | if (valid) { ss << __str_438; -// 238 | // Load the B right tiles for this K block. +// 242 | // Load the B right tiles for this K block. ss << __str_5; -// 239 | #if b_is_hwio +// 243 | #if b_is_hwio if (__param_b_is_hwio) { -// 240 | // HWIO: B is row-major ([K, N_b], B[k, n] at k*N_b + n); load row-major, -// 241 | // stride N_b. N_b (== N unless padded) keeps the row stride, and via -// 242 | // global_base_n above, the column offset, 4-byte aligned. -// 243 | let b_base = kb * kSgMatK * uniforms.N_b + global_base_n; +// 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; -// 244 | var sg_mat_b0: subgroup_matrix_right = +// 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; -// 245 | subgroupMatrixLoad, row_major>( +// 249 | subgroupMatrixLoad, row_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_101; -// 246 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +// 250 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); ss << __str_455; -// 247 | #if sg_mat_count_n >= 2 +// 251 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 248 | var sg_mat_b1: subgroup_matrix_right = +// 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; -// 249 | subgroupMatrixLoad, row_major>( +// 253 | 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 + 1 * kSgMatN, uniforms.N_b); +// 254 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); ss << __str_457; -// 251 | #endif +// 255 | #endif } -// 252 | #if sg_mat_count_n >= 3 +// 256 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 253 | var sg_mat_b2: subgroup_matrix_right = +// 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; -// 254 | subgroupMatrixLoad, row_major>( +// 258 | subgroupMatrixLoad, row_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_101; -// 255 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +// 259 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); ss << __str_459; -// 256 | #endif +// 260 | #endif } -// 257 | #if sg_mat_count_n >= 4 +// 261 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 258 | var sg_mat_b3: subgroup_matrix_right = +// 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; -// 259 | subgroupMatrixLoad, row_major>( +// 263 | subgroupMatrixLoad, row_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_101; -// 260 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +// 264 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); ss << __str_461; -// 261 | #endif +// 265 | #endif } -// 262 | #else +// 266 | #else } else { -// 263 | // OHWI: B is column-major ([N, K], B[k, n] at n*K + k); load col-major, stride K. -// 264 | let b_base = global_base_n * uniforms.K + kb * kSgMatK; +// 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; -// 265 | var sg_mat_b0: subgroup_matrix_right = +// 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; -// 266 | subgroupMatrixLoad, col_major>( +// 270 | subgroupMatrixLoad, col_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_92; -// 267 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); +// 271 | &input_b, b_base + 0 * kSgMatN * uniforms.K, uniforms.K); ss << __str_463; -// 268 | #if sg_mat_count_n >= 2 +// 272 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 269 | var sg_mat_b1: subgroup_matrix_right = +// 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; -// 270 | subgroupMatrixLoad, col_major>( +// 274 | 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 + 1 * kSgMatN * uniforms.K, uniforms.K); +// 275 | &input_b, b_base + 1 * kSgMatN * uniforms.K, uniforms.K); ss << __str_464; -// 272 | #endif +// 276 | #endif } -// 273 | #if sg_mat_count_n >= 3 +// 277 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 274 | var sg_mat_b2: subgroup_matrix_right = +// 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; -// 275 | subgroupMatrixLoad, col_major>( +// 279 | subgroupMatrixLoad, col_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_92; -// 276 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); +// 280 | &input_b, b_base + 2 * kSgMatN * uniforms.K, uniforms.K); ss << __str_465; -// 277 | #endif +// 281 | #endif } -// 278 | #if sg_mat_count_n >= 4 +// 282 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 279 | var sg_mat_b3: subgroup_matrix_right = +// 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; -// 280 | subgroupMatrixLoad, col_major>( +// 284 | subgroupMatrixLoad, col_major>( ss << __str_454; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_92; -// 281 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); +// 285 | &input_b, b_base + 3 * kSgMatN * uniforms.K, uniforms.K); ss << __str_466; -// 282 | #endif +// 286 | #endif } -// 283 | #endif +// 287 | #endif } -// 284 | -// 285 | // Load the A left tiles (one per M block) from workgroup staging, -// 286 | // row-major with stride kSgMatK. -// 287 | var sg_mat_a0: subgroup_matrix_left = +// 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; -// 288 | subgroupMatrixLoad, row_major>( +// 292 | subgroupMatrixLoad, row_major>( ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 289 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); +// 293 | &a_scratch, a_scratch_base + 0 * kSgMatM * kSgMatK, kSgMatK); ss << __str_469; -// 290 | #if sg_mat_count_m >= 2 +// 294 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 291 | var sg_mat_a1: subgroup_matrix_left = +// 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; -// 292 | subgroupMatrixLoad, row_major>( +// 296 | 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 + 1 * kSgMatM * kSgMatK, kSgMatK); +// 297 | &a_scratch, a_scratch_base + 1 * kSgMatM * kSgMatK, kSgMatK); ss << __str_471; -// 294 | #endif +// 298 | #endif } -// 295 | #if sg_mat_count_m >= 3 +// 299 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 296 | var sg_mat_a2: subgroup_matrix_left = +// 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; -// 297 | subgroupMatrixLoad, row_major>( +// 301 | subgroupMatrixLoad, row_major>( ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 298 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); +// 302 | &a_scratch, a_scratch_base + 2 * kSgMatM * kSgMatK, kSgMatK); ss << __str_473; -// 299 | #endif +// 303 | #endif } -// 300 | #if sg_mat_count_m >= 4 +// 304 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 301 | var sg_mat_a3: subgroup_matrix_left = +// 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; -// 302 | subgroupMatrixLoad, row_major>( +// 306 | subgroupMatrixLoad, row_major>( ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 303 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); +// 307 | &a_scratch, a_scratch_base + 3 * kSgMatM * kSgMatK, kSgMatK); ss << __str_475; -// 304 | #endif +// 308 | #endif } -// 305 | #if sg_mat_count_m >= 5 +// 309 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 306 | var sg_mat_a4: subgroup_matrix_left = +// 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; -// 307 | subgroupMatrixLoad, row_major>( +// 311 | subgroupMatrixLoad, row_major>( ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 308 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); +// 312 | &a_scratch, a_scratch_base + 4 * kSgMatM * kSgMatK, kSgMatK); ss << __str_477; -// 309 | #endif +// 313 | #endif } -// 310 | #if sg_mat_count_m >= 6 +// 314 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 311 | var sg_mat_a5: subgroup_matrix_left = +// 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; -// 312 | subgroupMatrixLoad, row_major>( +// 316 | subgroupMatrixLoad, row_major>( ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 313 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); +// 317 | &a_scratch, a_scratch_base + 5 * kSgMatM * kSgMatK, kSgMatK); ss << __str_479; -// 314 | #endif +// 318 | #endif } -// 315 | #if sg_mat_count_m >= 7 +// 319 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 316 | var sg_mat_a6: subgroup_matrix_left = +// 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; -// 317 | subgroupMatrixLoad, row_major>( +// 321 | subgroupMatrixLoad, row_major>( ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 318 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); +// 322 | &a_scratch, a_scratch_base + 6 * kSgMatM * kSgMatK, kSgMatK); ss << __str_481; -// 319 | #endif +// 323 | #endif } -// 320 | #if sg_mat_count_m >= 8 +// 324 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 321 | var sg_mat_a7: subgroup_matrix_left = +// 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; -// 322 | subgroupMatrixLoad, row_major>( +// 326 | subgroupMatrixLoad, row_major>( ss << __str_468; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 323 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); +// 327 | &a_scratch, a_scratch_base + 7 * kSgMatM * kSgMatK, kSgMatK); ss << __str_483; -// 324 | #endif +// 328 | #endif } -// 325 | -// 326 | // Accumulate every (M block, N block) pair. -// 327 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +// 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; -// 328 | #if sg_mat_count_n >= 2 +// 332 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 329 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +// 333 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); ss << __str_485; -// 330 | #endif +// 334 | #endif } -// 331 | #if sg_mat_count_n >= 3 +// 335 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 332 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +// 336 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); ss << __str_486; -// 333 | #endif +// 337 | #endif } -// 334 | #if sg_mat_count_n >= 4 +// 338 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 335 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +// 339 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); ss << __str_487; -// 336 | #endif +// 340 | #endif } -// 337 | #if sg_mat_count_m >= 2 +// 341 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 338 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +// 342 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); ss << __str_488; -// 339 | #if sg_mat_count_n >= 2 +// 343 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 340 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +// 344 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); ss << __str_489; -// 341 | #endif +// 345 | #endif } -// 342 | #if sg_mat_count_n >= 3 +// 346 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 343 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +// 347 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); ss << __str_490; -// 344 | #endif +// 348 | #endif } -// 345 | #if sg_mat_count_n >= 4 +// 349 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 346 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +// 350 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); ss << __str_491; -// 347 | #endif +// 351 | #endif } -// 348 | #endif +// 352 | #endif } -// 349 | #if sg_mat_count_m >= 3 +// 353 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 350 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +// 354 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); ss << __str_492; -// 351 | #if sg_mat_count_n >= 2 +// 355 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 352 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +// 356 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); ss << __str_493; -// 353 | #endif +// 357 | #endif } -// 354 | #if sg_mat_count_n >= 3 +// 358 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 355 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +// 359 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); ss << __str_494; -// 356 | #endif +// 360 | #endif } -// 357 | #if sg_mat_count_n >= 4 +// 361 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 358 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +// 362 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); ss << __str_495; -// 359 | #endif +// 363 | #endif } -// 360 | #endif +// 364 | #endif } -// 361 | #if sg_mat_count_m >= 4 +// 365 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 362 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +// 366 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); ss << __str_496; -// 363 | #if sg_mat_count_n >= 2 +// 367 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 364 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +// 368 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); ss << __str_497; -// 365 | #endif +// 369 | #endif } -// 366 | #if sg_mat_count_n >= 3 +// 370 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 367 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +// 371 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); ss << __str_498; -// 368 | #endif +// 372 | #endif } -// 369 | #if sg_mat_count_n >= 4 +// 373 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 370 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +// 374 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); ss << __str_499; -// 371 | #endif +// 375 | #endif } -// 372 | #endif +// 376 | #endif } -// 373 | #if sg_mat_count_m >= 5 +// 377 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 374 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +// 378 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); ss << __str_500; -// 375 | #if sg_mat_count_n >= 2 +// 379 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 376 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +// 380 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); ss << __str_501; -// 377 | #endif +// 381 | #endif } -// 378 | #if sg_mat_count_n >= 3 +// 382 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 379 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +// 383 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); ss << __str_502; -// 380 | #endif +// 384 | #endif } -// 381 | #if sg_mat_count_n >= 4 +// 385 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 382 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +// 386 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); ss << __str_503; -// 383 | #endif +// 387 | #endif } -// 384 | #endif +// 388 | #endif } -// 385 | #if sg_mat_count_m >= 6 +// 389 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 386 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +// 390 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); ss << __str_504; -// 387 | #if sg_mat_count_n >= 2 +// 391 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 388 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +// 392 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); ss << __str_505; -// 389 | #endif +// 393 | #endif } -// 390 | #if sg_mat_count_n >= 3 +// 394 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 391 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +// 395 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); ss << __str_506; -// 392 | #endif +// 396 | #endif } -// 393 | #if sg_mat_count_n >= 4 +// 397 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 394 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +// 398 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); ss << __str_507; -// 395 | #endif +// 399 | #endif } -// 396 | #endif +// 400 | #endif } -// 397 | #if sg_mat_count_m >= 7 +// 401 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 398 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +// 402 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); ss << __str_508; -// 399 | #if sg_mat_count_n >= 2 +// 403 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 400 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +// 404 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); ss << __str_509; -// 401 | #endif +// 405 | #endif } -// 402 | #if sg_mat_count_n >= 3 +// 406 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 403 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +// 407 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); ss << __str_510; -// 404 | #endif +// 408 | #endif } -// 405 | #if sg_mat_count_n >= 4 +// 409 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 406 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +// 410 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); ss << __str_511; -// 407 | #endif +// 411 | #endif } -// 408 | #endif +// 412 | #endif } -// 409 | #if sg_mat_count_m >= 8 +// 413 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 410 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +// 414 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); ss << __str_512; -// 411 | #if sg_mat_count_n >= 2 +// 415 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 412 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +// 416 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); ss << __str_513; -// 413 | #endif +// 417 | #endif } -// 414 | #if sg_mat_count_n >= 3 +// 418 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 415 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +// 419 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); ss << __str_514; -// 416 | #endif +// 420 | #endif } -// 417 | #if sg_mat_count_n >= 4 +// 421 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 418 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +// 422 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); ss << __str_515; -// 419 | #endif +// 423 | #endif } -// 420 | #endif +// 424 | #endif } -// 421 | } +// 425 | } ss << __str_204; -// 422 | // Ensure every subgroup is done reading a_scratch before it is reused for +// 426 | // Ensure every subgroup is done reading a_scratch before it is reused for ss << __str_5; -// 423 | // the next K-block step. -// 424 | workgroupBarrier(); +// 427 | // the next K-block step. +// 428 | workgroupBarrier(); ss << __str_451; -// 425 | } +// 429 | } ss << __str_20; -// 426 | +// 430 | ss << __str_5; -// 427 | // Store into this subgroup's split-K slot (kSplitK x [kTileM, kTileN] row-major; -// 428 | // block (mi, ni) at mi*kSgMatM*kTileN + ni*kSgMatN). -// 429 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +// 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; -// 430 | #if sg_mat_count_n >= 2 +// 434 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 431 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +// 435 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); ss << __str_166; -// 432 | #endif +// 436 | #endif } -// 433 | #if sg_mat_count_n >= 3 +// 437 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 434 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +// 438 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); ss << __str_167; -// 435 | #endif +// 439 | #endif } -// 436 | #if sg_mat_count_n >= 4 +// 440 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 437 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +// 441 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); ss << __str_168; -// 438 | #endif +// 442 | #endif } -// 439 | #if sg_mat_count_m >= 2 +// 443 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 440 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +// 444 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); ss << __str_169; -// 441 | #if sg_mat_count_n >= 2 +// 445 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 442 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +// 446 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); ss << __str_170; -// 443 | #endif +// 447 | #endif } -// 444 | #if sg_mat_count_n >= 3 +// 448 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 445 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +// 449 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); ss << __str_171; -// 446 | #endif +// 450 | #endif } -// 447 | #if sg_mat_count_n >= 4 +// 451 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 448 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); ss << __str_172; -// 449 | #endif +// 453 | #endif } -// 450 | #endif +// 454 | #endif } -// 451 | #if sg_mat_count_m >= 3 +// 455 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +// 456 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); ss << __str_173; -// 453 | #if sg_mat_count_n >= 2 +// 457 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 454 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +// 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); ss << __str_174; -// 455 | #endif +// 459 | #endif } -// 456 | #if sg_mat_count_n >= 3 +// 460 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 457 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +// 461 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); ss << __str_175; -// 458 | #endif +// 462 | #endif } -// 459 | #if sg_mat_count_n >= 4 +// 463 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 460 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +// 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); ss << __str_176; -// 461 | #endif +// 465 | #endif } -// 462 | #endif +// 466 | #endif } -// 463 | #if sg_mat_count_m >= 4 +// 467 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +// 468 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); ss << __str_177; -// 465 | #if sg_mat_count_n >= 2 +// 469 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 466 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +// 470 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); ss << __str_178; -// 467 | #endif +// 471 | #endif } -// 468 | #if sg_mat_count_n >= 3 +// 472 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 469 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +// 473 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); ss << __str_179; -// 470 | #endif +// 474 | #endif } -// 471 | #if sg_mat_count_n >= 4 +// 475 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 472 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +// 476 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); ss << __str_180; -// 473 | #endif +// 477 | #endif } -// 474 | #endif +// 478 | #endif } -// 475 | #if sg_mat_count_m >= 5 +// 479 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 476 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +// 480 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); ss << __str_181; -// 477 | #if sg_mat_count_n >= 2 +// 481 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 478 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +// 482 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); ss << __str_182; -// 479 | #endif +// 483 | #endif } -// 480 | #if sg_mat_count_n >= 3 +// 484 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 481 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +// 485 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); ss << __str_183; -// 482 | #endif +// 486 | #endif } -// 483 | #if sg_mat_count_n >= 4 +// 487 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 484 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +// 488 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); ss << __str_184; -// 485 | #endif +// 489 | #endif } -// 486 | #endif +// 490 | #endif } -// 487 | #if sg_mat_count_m >= 6 +// 491 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 488 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +// 492 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); ss << __str_185; -// 489 | #if sg_mat_count_n >= 2 +// 493 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 490 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +// 494 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); ss << __str_186; -// 491 | #endif +// 495 | #endif } -// 492 | #if sg_mat_count_n >= 3 +// 496 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 493 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +// 497 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); ss << __str_187; -// 494 | #endif +// 498 | #endif } -// 495 | #if sg_mat_count_n >= 4 +// 499 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 496 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +// 500 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); ss << __str_188; -// 497 | #endif +// 501 | #endif } -// 498 | #endif +// 502 | #endif } -// 499 | #if sg_mat_count_m >= 7 +// 503 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 500 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +// 504 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); ss << __str_189; -// 501 | #if sg_mat_count_n >= 2 +// 505 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 502 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +// 506 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); ss << __str_190; -// 503 | #endif +// 507 | #endif } -// 504 | #if sg_mat_count_n >= 3 +// 508 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 505 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +// 509 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); ss << __str_191; -// 506 | #endif +// 510 | #endif } -// 507 | #if sg_mat_count_n >= 4 +// 511 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 508 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +// 512 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); ss << __str_192; -// 509 | #endif +// 513 | #endif } -// 510 | #endif +// 514 | #endif } -// 511 | #if sg_mat_count_m >= 8 +// 515 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 512 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +// 516 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); ss << __str_193; -// 513 | #if sg_mat_count_n >= 2 +// 517 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 514 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +// 518 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); ss << __str_194; -// 515 | #endif +// 519 | #endif } -// 516 | #if sg_mat_count_n >= 3 +// 520 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 517 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +// 521 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); ss << __str_195; -// 518 | #endif +// 522 | #endif } -// 519 | #if sg_mat_count_n >= 4 +// 523 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 520 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +// 524 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); ss << __str_196; -// 521 | #endif +// 525 | #endif } -// 522 | #endif +// 526 | #endif } -// 523 | -// 524 | #if split_k >= 2 +// 527 | +// 528 | #if split_k >= 2 if (__param_split_k >= 2) { -// 525 | // Make all subgroups' partial tiles visible before the reduction. -// 526 | workgroupBarrier(); +// 529 | // Make all subgroups' partial tiles visible before the reduction. +// 530 | workgroupBarrier(); ss << __str_26; -// 527 | +// 531 | ss << __str_5; -// 528 | // Subgroup 0 reduces the kSplitK partial tiles into slot 0; safe since each -// 529 | // position is only overwritten by the lane that already read it. -// 530 | if (sg_index == 0u) { +// 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; -// 531 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +// 535 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { ss << __str_198; -// 532 | var acc: f16 = f16(0); +// 536 | var acc: f16 = f16(0); ss << __str_199; -// 533 | for (var s: u32 = 0; s < kSplitK; s++) { +// 537 | for (var s: u32 = 0; s < kSplitK; s++) { ss << __str_200; -// 534 | acc += scratch[s * kTileM * kTileN + idx]; +// 538 | acc += scratch[s * kTileM * kTileN + idx]; ss << __str_201; -// 535 | } +// 539 | } ss << __str_202; -// 536 | scratch[idx] = acc; +// 540 | scratch[idx] = acc; ss << __str_203; -// 537 | } +// 541 | } ss << __str_204; -// 538 | } +// 542 | } ss << __str_20; -// 539 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +// 543 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << __str_5; -// 540 | workgroupBarrier(); +// 544 | workgroupBarrier(); ss << __str_26; -// 541 | #endif +// 545 | #endif } -// 542 | -// 543 | // Write-out: subgroups stride M-rows by kSplitK, lanes cooperate per row. n_skip -// 544 | // and M's bounds check make non-multiple dimensions write each element once. -// 545 | let n_count = min(kTileN, uniforms.N - global_base_n); +// 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; -// 546 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +// 550 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { ss << __str_207; -// 547 | let global_m = m_base + r; +// 551 | let global_m = m_base + r; ss << __str_208; -// 548 | if (global_m < uniforms.M) { +// 552 | if (global_m < uniforms.M) { ss << __str_516; -// 549 | let scratch_base = r * kTileN; +// 553 | let scratch_base = r * kTileN; ss << __str_210; -// 550 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; +// 554 | let out_base = out_batch_base + global_m * uniforms.N + global_base_n; ss << __str_517; -// 551 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +// 555 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { ss << __str_212; -// 552 | write_output(out_base + i, global_base_n + i, +// 556 | write_output(out_base + i, global_base_n + i, ss << __str_235; -// 553 | output_value_t(scratch[scratch_base + i])); +// 557 | output_value_t(scratch[scratch_base + i])); ss << __str_236; -// 554 | } +// 558 | } ss << __str_202; -// 555 | } +// 559 | } ss << __str_204; -// 556 | } +// 560 | } ss << __str_20; -// 557 | } // MAIN +// 561 | } // MAIN MainFunctionEnd(); ss << __str_5; -// 558 | +// 562 | return Status::OK(); 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 49274aa9fcdae..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 @@ -41,7 +41,7 @@ std::string pass_as_string(T&& v) { #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/nn/subgroup_matrix_conv_8x16x16.h" // b0abd8ed2662209255943ac112b323db06a470d2c62c2fb7430f1925085172f9 +#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