diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc index aa17baeb2bae1..43582604bfff1 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc @@ -3,6 +3,7 @@ #include "core/providers/webgpu/math/subgroup_matrix_tiling.h" +#include #include #include #include @@ -26,14 +27,18 @@ namespace webgpu { class SubgroupMatrixMatMulProgram final : public Program { public: SubgroupMatrixMatMulProgram(const Activation& activation, bool has_bias, SubgroupMatrixConfig config, - uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k) + uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k, + uint32_t sg_grid_m, uint32_t sg_grid_n, uint32_t tile_k) : Program{"SubgroupMatrixMatMul"}, activation_(activation), has_bias_(has_bias), config_(config), sg_mat_count_m_(sg_mat_count_m), sg_mat_count_n_(sg_mat_count_n), - split_k_(split_k) {} + split_k_(split_k), + sg_grid_m_(sg_grid_m), + sg_grid_n_(sg_grid_n), + tile_k_(tile_k) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32}, {"N", ProgramUniformVariableDataType::Uint32}, @@ -49,6 +54,9 @@ class SubgroupMatrixMatMulProgram final : public Program 2 || tiling.tile_n / cell_n > 2) { + return false; + } + const uint64_t threads = uint64_t{config.subgroupSize} * tiling.sg_grid_m * tiling.sg_grid_n; + // Matches the scratch size in the template where each staged row is padded by 8 elements. + const uint64_t staged_elems = uint64_t{tiling.tile_m} * (tiling.tile_k + 8) + uint64_t{tiling.tile_k} * (tiling.tile_n + 8); + const uint64_t scratch_bytes = sizeof(uint16_t) * std::max(staged_elems, uint64_t{tiling.tile_m} * tiling.tile_n); + return threads <= limits.maxComputeInvocationsPerWorkgroup && threads <= limits.maxComputeWorkgroupSizeX && + scratch_bytes <= limits.maxComputeWorkgroupStorageSize; +} + // Subgroup-matrix MatMul implementation. Loads both A and B directly from global -// memory and runs the subgroup-matrix kernel during Compute. The class is +// memory (or stages them in shared memory when the tiling sets tile_k) and runs the +// subgroup-matrix kernel during Compute. The class is // intended to support all subgroup-matrix configs. For now 8x16x16 and 16x16x16 // are implemented. The per-problem output tiling is supplied by a vendor-specific // selector kept internal to this impl. @@ -135,6 +163,9 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl { } const auto& config = config_; + if (tiling->tile_k > 0) { + return config.K != 0 && K % config.K == 0 && IsStagedTilingSupported(context.DeviceLimits(), config, *tiling, N); + } const uint32_t n_align = config.Is(16, 16, 16) ? 8 : 2; const bool needs_padded_b = N % n_align != 0; if (needs_padded_b && (!b_is_constant || N > std::numeric_limits::max() - n_align)) { @@ -194,10 +225,15 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl { "MatMul algorithm subgroup_matrix has no tiling for these inputs."); const auto& config = config_; + // Staged tilings read B through shared memory so they never need the padded copy. + const bool staged = tiling->tile_k > 0; + ORT_RETURN_IF_NOT(staged ? IsStagedTilingSupported(context.DeviceLimits(), config, *tiling, N) + : tiling->sg_grid_m == 1 && tiling->sg_grid_n == 1, + "MatMul algorithm subgroup_matrix got a tiling this device cannot run."); // Vulkan cooperative matrix loads need 16 byte aligned rows so the // 16x16x16 config pads B to a multiple of 8 columns. const uint32_t n_align = config.Is(16, 16, 16) ? 8 : 2; - const bool needs_padded_b = N % n_align != 0; + const bool needs_padded_b = !staged && N % n_align != 0; // Require whole subgroup-matrix K blocks. An unaligned B must be constant // because its padded copy is cached by this implementation. ORT_RETURN_IF_NOT(config.K != 0 && K % config.K == 0 && @@ -215,7 +251,7 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl { // The kernel keeps its operand loads in bounds by shifting a trailing partial // tile back, which is only possible when the tile fits within M and N. - ORT_ENFORCE(M >= tiling->tile_m && N_b >= tiling->tile_n); + ORT_ENFORCE(staged || (M >= tiling->tile_m && N_b >= tiling->tile_n)); // The optimized path will run: now materialize the aligned B for unaligned N. const Tensor* b_used = b; @@ -229,17 +265,21 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl { const uint32_t tile_m = tiling->tile_m; const uint32_t tile_n = tiling->tile_n; const uint32_t split_k = tiling->split_k; - const uint32_t sg_mat_count_m = tile_m / config.M; - const uint32_t sg_mat_count_n = tile_n / config.N; - ORT_ENFORCE(tile_m % config.M == 0 && tile_n % config.N == 0, + const uint32_t sg_grid_m = tiling->sg_grid_m; + const uint32_t sg_grid_n = tiling->sg_grid_n; + const uint32_t tile_k = tiling->tile_k; + const uint32_t sg_mat_count_m = tile_m / (config.M * sg_grid_m); + const uint32_t sg_mat_count_n = tile_n / (config.N * sg_grid_n); + ORT_ENFORCE(tile_m % (config.M * sg_grid_m) == 0 && tile_n % (config.N * sg_grid_n) == 0, "Tiling must be a multiple of the subgroup-matrix shape: ", tile_m, "x", tile_n, " vs ", config.M, "x", config.N); const uint32_t dispatch_x = (N + tile_n - 1) / tile_n; const uint32_t dispatch_y = (M + tile_m - 1) / tile_m; SubgroupMatrixMatMulProgram program{activation, has_bias, config, - sg_mat_count_m, sg_mat_count_n, split_k}; - program.SetWorkgroupSize(config.subgroupSize * split_k); + sg_mat_count_m, sg_mat_count_n, split_k, + sg_grid_m, sg_grid_n, tile_k}; + program.SetWorkgroupSize(config.subgroupSize * split_k * sg_grid_m * sg_grid_n); if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { program.SetSubgroupSize(config.subgroupSize); } @@ -248,9 +288,9 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl { 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) - .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b_used, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + sg_mat_count_m, sg_mat_count_n, split_k, sg_grid_m, sg_grid_n, tile_k) + .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, staged ? 4 : 1}, + {b_used, ProgramTensorMetadataDependency::TypeAndRank, staged ? 4 : 1}}) .AddOutput({output, ProgramTensorMetadataDependency::Rank, 1}) .AddUniformVariables({{M}, {N}, {K}, {dispatch_x}, {N_b}}); // Activation uniforms must remain last because definitions and values are matched by index. @@ -322,14 +362,18 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl { Status GenerateSubgroupMatrixShaderCode(ShaderHelper& shader, const SubgroupMatrixConfig& config, uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, - uint32_t split_k) { + uint32_t split_k, uint32_t sg_grid_m, uint32_t sg_grid_n, + uint32_t tile_k) { return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul.wgsl.template", + WGSL_TEMPLATE_PARAMETER(sg_grid_m, sg_grid_m), + WGSL_TEMPLATE_PARAMETER(sg_grid_n, sg_grid_n), 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, config.K), WGSL_TEMPLATE_PARAMETER(sg_mat_m, config.M), WGSL_TEMPLATE_PARAMETER(sg_mat_n, config.N), - WGSL_TEMPLATE_PARAMETER(split_k, split_k)); + WGSL_TEMPLATE_PARAMETER(split_k, split_k), + WGSL_TEMPLATE_PARAMETER(tile_k, tile_k)); } // Default tiling used on any vendor without a specialized policy: a 64x64 output @@ -345,6 +389,30 @@ SubgroupMatrixTilingSelector MakeDefaultTilingSelector() { }; } +// A 4x4 grid of subgroups stages K slices 64 deep and each subgroup accumulates 2x2 subgroup +// matrices when the problem has at least 128 rows and 256 tiles of 64x64 and 1x1 otherwise. +// Staging starts at 64 rows and below that only for outputs at least 8192 wide because the +// default split K tiling fills the GPU better on small problems. +SubgroupMatrixTilingSelector MakeStagedTilingSelector(const SubgroupMatrixConfig& config) { + return [config, fallback = MakeDefaultTilingSelector()]( + const ComputeContext& context, uint32_t M, uint32_t N, + uint32_t K, uint32_t batch) -> std::optional { + if ((M >= 64 || N >= 8192) && K % config.K == 0) { + const auto& limits = context.DeviceLimits(); + const bool large = M >= 128 && uint64_t{(M + 63) / 64} * ((N + 63) / 64) >= 256; + const SubgroupMatrixTiling big{config.M * 8, config.N * 8, 1, 4, 4, 64}; + const SubgroupMatrixTiling small{config.M * 4, config.N * 4, 1, 4, 4, 64}; + if (large && IsStagedTilingSupported(limits, config, big, N)) { + return big; + } + if (IsStagedTilingSupported(limits, config, small, N)) { + return small; + } + } + return fallback(context, M, N, K, batch); + }; +} + } // namespace Status SubgroupMatrixMatMulProgram::GenerateShaderCode(ShaderHelper& shader) const { @@ -369,7 +437,8 @@ Status SubgroupMatrixMatMulProgram::GenerateShaderCode(ShaderHelper& shader) con << "}\n"; if (config_.Is(8, 16, 16) || config_.Is(16, 16, 16)) { - return GenerateSubgroupMatrixShaderCode(shader, config_, sg_mat_count_m_, sg_mat_count_n_, split_k_); + return GenerateSubgroupMatrixShaderCode(shader, config_, sg_mat_count_m_, sg_mat_count_n_, split_k_, + sg_grid_m_, sg_grid_n_, tile_k_); } return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, "Unsupported subgroup matrix config dimensions."); @@ -384,11 +453,13 @@ std::unique_ptr CreateSubgroupMatrixMatMulImpl(const C if (!config) { return nullptr; } - // Intel GPUs use a tuned/heuristic tiling policy; every other vendor falls back + // Intel and NVIDIA GPUs use a tuned/heuristic tiling policy; every other vendor falls back // to a default tiling. - const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; + const auto vendor = context.AdapterInfo().vendor; SubgroupMatrixTilingSelector tiling_selector = - is_intel ? intel::CreateSubgroupMatrixTilingSelector(context) : MakeDefaultTilingSelector(); + vendor == std::string_view{"intel"} ? intel::CreateSubgroupMatrixTilingSelector(context) + : vendor == std::string_view{"nvidia"} ? MakeStagedTilingSelector(*config) + : MakeDefaultTilingSelector(); if (!tiling_selector) { return nullptr; } diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template index a41c869cefd9a..db5de1001552b 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template @@ -15,10 +15,12 @@ // N_b. N_b is the padded row stride of the (possibly column-padded) B buffer and // equals N unless the host padded B to an aligned stride to satisfy the load's // row-start alignment (see the B padding in subgroup_matrix_matmul.cc). +// When tile_k > 0 the workgroup instead stages tile_k deep slices of A and B in +// shared memory and the subgroups load their operands from there. // -// Workgroup: split_k subgroups x 32 lanes (split_k in {1,2,4,8}). +// Workgroup: sg_grid_m x sg_grid_n x split_k subgroups x 32 lanes (split_k in {1,2,4,8}). // Tile: kTileM x kTileN, chosen adaptively by the host config provider. Each -// subgroup computes the whole tile as kSgMatCountM x kSgMatCountN subgroup +// subgroup computes its grid cell 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. @@ -34,7 +36,9 @@ // resets the whole load to offset 0 with the minimum stride, so it silently // returns a different, valid tile, which would corrupt the in-range rows of that // tile too. Output width/stride uses N (not N_b), so a padded B column in -// [N, N_b) may be read but is never written. +// [N, N_b) may be read but is never written. When tile_k > 0 none of this applies +// because the staging loads are bounds checked and fill zeros outside M N and K so +// the host only requires N % 4 == 0 and split_k == 1 and B is never padded. // // Batching: the host dispatches one num_n_tile x num_m_tile tile grid per batch // slice; the batch slice is recovered from the flattened workgroup_idx (the @@ -49,16 +53,27 @@ #param sg_mat_count_m #param sg_mat_count_n #param split_k +#param sg_grid_m +#param sg_grid_n +#param tile_k // #use .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 kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per subgroup along M +const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per subgroup along N +const kSgGridM: u32 = u32(sg_grid_m); // Subgroups splitting the tile along M +const kSgGridN: u32 = u32(sg_grid_n); // Subgroups splitting the tile along N +const kTileM: u32 = kSgMatM * kSgMatCountM * kSgGridM; // Tile size along M (output rows per workgroup) +const kTileN: u32 = kSgMatN * kSgMatCountN * kSgGridN; // Tile size along N (output columns per workgroup) +#if tile_k > 0 +const kTileK: u32 = u32(tile_k); // K depth of A and B staged in shared memory per step +const kWorkgroupSize: u32 = kSgGridM * kSgGridN * kSubgroupSize; // Threads per workgroup since split_k is 1 here +const kStrideA: u32 = kTileK + 8u; // Staged A row stride padded by 16 bytes to spread banks +const kStrideB: u32 = kTileN + 8u; // Staged B row stride padded by 16 bytes to spread banks +#endif const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K const kSubgroupSize: u32 = 32; // Lanes per subgroup // @@ -66,7 +81,13 @@ const kSubgroupSize: u32 = 32; // Lanes per subgroup // 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. +// When tile_k > 0 it holds the staged kTileM x kTileK slice of A and kTileK x kTileN +// slice of B during the K loop and the single output tile after it. +#if tile_k > 0 +var scratch: array; +#else var scratch: array; +#endif // $MAIN { // Recover the batch slice and the logical (n_tile, m_tile) from the flattened @@ -83,6 +104,15 @@ $MAIN { let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile let n_tile_start = n_tile * kTileN; let m_tile_start = m_tile * kTileM; +#if tile_k > 0 + let a_batch_offset = batch_id * uniforms.M * uniforms.K; + let b_batch_offset = batch_id * uniforms.K * uniforms.N; + let out_batch_offset = batch_id * uniforms.M * uniforms.N; + let sg_index = local_idx / kSubgroupSize; + let m_base = (sg_index % kSgGridM) * kSgMatM * kSgMatCountM; // first tile row of this subgroup's cell + let n_base = (sg_index / kSgGridM) * kSgMatN * kSgMatCountN; // first tile column of this subgroup's cell + let sg_scratch_base = m_base * kTileN + n_base; // this subgroup's cell in the output tile +#else // Shift a trailing partial tile back so its operand loads stay inside A and B. // The N shift uses B's padded stride N_b so the shifted origin stays aligned (the // load needs aligned row starts). The extra columns are clipped at write-out. @@ -99,6 +129,7 @@ $MAIN { 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 scratch slot +#endif // // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). var sg_mat_c0_0: subgroup_matrix_result; @@ -196,6 +227,59 @@ $MAIN { #endif #endif // +#if tile_k > 0 + for (var k0: u32 = 0u; k0 < uniforms.K; k0 = k0 + kTileK) { + // k_start equals k0. Deriving it through a division made this loop up to 40% + // faster on an RTX 5080. + let k_start = (k0 / kSgMatK) * kSgMatK; + // The barrier keeps the previous slice alive until every subgroup has read it. + workgroupBarrier(); + for (var i: u32 = local_idx; i < kTileM * kTileK / 4u; i = i + kWorkgroupSize) { + let r = i / (kTileK / 4u); + let c = (i % (kTileK / 4u)) * 4u; + var v = vec4(); + if (m_tile_start + r < uniforms.M && k_start + c < uniforms.K) { + v = input_a[(a_batch_offset + (m_tile_start + r) * uniforms.K + k_start + c) / 4u]; + } + let s = r * kStrideA + c; + scratch[s] = v.x; + scratch[s + 1u] = v.y; + scratch[s + 2u] = v.z; + scratch[s + 3u] = v.w; + } + for (var i: u32 = local_idx; i < kTileK * kTileN / 4u; i = i + kWorkgroupSize) { + let r = i / (kTileN / 4u); + let c = (i % (kTileN / 4u)) * 4u; + var v = vec4(); + if (k_start + r < uniforms.K && n_tile_start + c < uniforms.N) { + v = input_b[(b_batch_offset + (k_start + r) * uniforms.N + n_tile_start + c) / 4u]; + } + let s = kTileM * kStrideA + r * kStrideB + c; + scratch[s] = v.x; + scratch[s + 1u] = v.y; + scratch[s + 2u] = v.z; + scratch[s + 3u] = v.w; + } + workgroupBarrier(); + for (var kk: u32 = 0u; kk < kTileK; kk = kk + kSgMatK) { + let b_base = kTileM * kStrideA + kk * kStrideB + n_base; + var sg_mat_b0: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &scratch, b_base + 0 * kSgMatN, kStrideB); +#if sg_mat_count_n >= 2 + var sg_mat_b1: subgroup_matrix_right = + subgroupMatrixLoad, row_major>( + &scratch, b_base + 1 * kSgMatN, kStrideB); +#endif + var sg_mat_a0: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &scratch, (m_base + 0 * kSgMatM) * kStrideA + kk, kStrideA); +#if sg_mat_count_m >= 2 + var sg_mat_a1: subgroup_matrix_left = + subgroupMatrixLoad, row_major>( + &scratch, (m_base + 1 * kSgMatM) * kStrideA + kk, kStrideA); +#endif +#else for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { // Load the B right tiles for this K block (KxN_b row-major, stride N_b). let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n; @@ -259,6 +343,9 @@ $MAIN { var sg_mat_a7: subgroup_matrix_left = subgroupMatrixLoad, row_major>( &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K); +#endif + // This block stands in for the staged inner loop so both paths close the same braces. + { #endif // // Accumulate every (M block, N block) pair. @@ -356,8 +443,13 @@ $MAIN { sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); #endif #endif + } } // +#if tile_k > 0 + // Every subgroup must finish reading the staged slices before scratch holds the output. + workgroupBarrier(); +#endif // 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. @@ -456,6 +548,18 @@ $MAIN { #endif #endif // +#if tile_k > 0 + // The barrier makes every cell visible before all threads write out the tile. + workgroupBarrier(); + for (var i: u32 = local_idx; i < kTileM * kTileN; i = i + kWorkgroupSize) { + let global_m = m_tile_start + i / kTileN; + let global_n = n_tile_start + i % kTileN; + if (global_m < uniforms.M && global_n < uniforms.N) { + write_output(out_batch_offset + global_m * uniforms.N + global_n, global_n, + output_value_t(scratch[i])); + } + } +#else #if split_k >= 2 // Make all subgroups' partial tiles visible before the reduction. workgroupBarrier(); @@ -498,4 +602,5 @@ $MAIN { } } } +#endif } // MAIN diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_tiling.h b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_tiling.h index 8621750fe5d2e..426476578d0e9 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_tiling.h +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_tiling.h @@ -19,6 +19,10 @@ struct SubgroupMatrixTiling { uint32_t tile_m; // output rows per workgroup uint32_t tile_n; // output cols per workgroup uint32_t split_k; // subgroups cooperating along K (1 = no split) + + uint32_t sg_grid_m = 1; // subgroups splitting the tile along M + uint32_t sg_grid_n = 1; // subgroups splitting the tile along N + uint32_t tile_k = 0; // K depth staged in shared memory per step (0 loads A and B directly) }; // Vendor-supplied callback that selects the output tiling for a given problem. diff --git a/onnxruntime/test/providers/webgpu/math/matmul_test.cc b/onnxruntime/test/providers/webgpu/math/matmul_test.cc index 2ff6ad9efca94..c411daa05f744 100644 --- a/onnxruntime/test/providers/webgpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/webgpu/math/matmul_test.cc @@ -185,8 +185,8 @@ TEST(MathOpTest, MatMulPackedScalarFp16LongReductionUsesFloat32Accumulator) { } // f16 MatMul cases that exercise the 8x16x16 and 16x16x16 subgroup-matrix impl. -// The host picks the tile shape adaptively (TileM in {8,16,32,64}, TileN in -// {16,32,64}); M and N may be any size and K must be a multiple of 16. When the +// The host picks the tile shape adaptively (TileM in {8,16,32,64,128}, TileN in +// {16,32,64,128}); M and N may be any size and K must be a multiple of 16. When the // output has few tiles and K is large, the host also splits K across multiple // cooperating subgroups (split_k in {1,2,4,8}) that reduce in shared memory. B // is a constant initializer except in the last case. On hardware that reports @@ -296,6 +296,10 @@ TEST(MathOpTest, MatMulSubgroupMatrix) { {"UnalignedN44 PartialM (M=40,N=44)", {40, 128}, 128, 44}, // Odd N and partial M past 64 so the default tiling takes its 64x64 tile on padded B. {"OddN LargeTile (M=72,N=67)", {72, 96}, 96, 67}, + // Staged tilings cover partial tiles and a partial final K slice. + {"Staged PartialMN KTail (M=100,N=196,K=80)", {100, 80}, 80, 196}, + {"Staged PartialMN KTail (M=520,N=8196,K=48)", {520, 48}, 48, 8196}, + {"Staged SmallM (M=40,N=8192)", {40, 64}, 64, 8192}, }; for (const auto& c : cases) { @@ -303,8 +307,8 @@ TEST(MathOpTest, MatMulSubgroupMatrix) { RunSubgroupMatrixMatMulTest(c.a_dims, c.K, c.N); } - // B is a runtime input here so its padded copy could not be cached and the - // 16x16x16 config falls back to the default impl for this width. + // B is a runtime input here so its padded copy could not be cached and without + // a staged tiling the 16x16x16 config falls back to the default impl for this width. SCOPED_TRACE("RuntimeB UnalignedN36 (N=36)"); RunSubgroupMatrixMatMulTest({64, 128}, 128, 36, /*b_is_initializer=*/false); } diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul.h index d2acac261dd1c..86cece6f33d0b 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/subgroup_matrix_matmul.h @@ -10,12 +10,15 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul.wgsl.template">(ShaderHelper& [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); // Extract parameters + auto& __param_sg_grid_m = params.param_sg_grid_m; + auto& __param_sg_grid_n = params.param_sg_grid_n; 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_tile_k = params.param_tile_k; // 1 | // Copyright (c) Microsoft Corporation. All rights reserved. // 2 | // Licensed under the MIT License. @@ -34,1136 +37,1368 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul.wgsl.template">(ShaderHelper& // 15 | // N_b. N_b is the padded row stride of the (possibly column-padded) B buffer and // 16 | // equals N unless the host padded B to an aligned stride to satisfy the load's // 17 | // row-start alignment (see the B padding in subgroup_matrix_matmul.cc). -// 18 | // -// 19 | // Workgroup: split_k subgroups x 32 lanes (split_k in {1,2,4,8}). -// 20 | // Tile: kTileM x kTileN, chosen adaptively by the host config provider. Each -// 21 | // subgroup computes the whole tile as kSgMatCountM x kSgMatCountN subgroup -// 22 | // matrices. K is iterated in sg_mat_k blocks. When split_k > 1 the K blocks are -// 23 | // distributed round-robin across the subgroups, each accumulates its own partial -// 24 | // tile in shared memory, and the partials are summed at write-out. -// 25 | // -// 26 | // Preconditions enforced by the host: K % sg_mat_k == 0; B's row stride N_b is -// 27 | // padded to the row-start alignment the load requires and the tile fits the -// 28 | // matrix (kTileM <= M, kTileN <= N - the host falls back otherwise). M and N may -// 29 | // otherwise be any size: a trailing partial tile is shifted back to end exactly -// 30 | // at M / N_b so no operand load runs past the end of A or B, and the rows/columns -// 31 | // the shifted tile re-covers are skipped at write-out. The shift is required for -// 32 | // correctness, not just safety - a subgroupMatrixLoad whose footprint leaves the -// 33 | // buffer does not degrade to a per-element zero fill: Tint's robustness transform -// 34 | // resets the whole load to offset 0 with the minimum stride, so it silently -// 35 | // returns a different, valid tile, which would corrupt the in-range rows of that -// 36 | // tile too. Output width/stride uses N (not N_b), so a padded B column in -// 37 | // [N, N_b) may be read but is never written. -// 38 | // -// 39 | // Batching: the host dispatches one num_n_tile x num_m_tile tile grid per batch -// 40 | // slice; the batch slice is recovered from the flattened workgroup_idx (the -// 41 | // dispatch may be normalized across x/y/z). A and the output are laid out as -// 42 | // [batch, M, K] / [batch, M, N]; B is either a shared [K, N] weight (batch 1) -// 43 | // or batched as [batch, K, N]. Per-slice flat offsets are derived from M, N, K. -// 44 | // Each slice is an independent tile grid. -// 45 | // -// 46 | #param sg_mat_k -// 47 | #param sg_mat_m -// 48 | #param sg_mat_n -// 49 | #param sg_mat_count_m -// 50 | #param sg_mat_count_n -// 51 | #param split_k -// 52 | // -// 53 | #use .setByOffset -// 54 | // -// 55 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +// 18 | // When tile_k > 0 the workgroup instead stages tile_k deep slices of A and B in +// 19 | // shared memory and the subgroups load their operands from there. +// 20 | // +// 21 | // Workgroup: sg_grid_m x sg_grid_n x split_k subgroups x 32 lanes (split_k in {1,2,4,8}). +// 22 | // Tile: kTileM x kTileN, chosen adaptively by the host config provider. Each +// 23 | // subgroup computes its grid cell as kSgMatCountM x kSgMatCountN subgroup +// 24 | // matrices. K is iterated in sg_mat_k blocks. When split_k > 1 the K blocks are +// 25 | // distributed round-robin across the subgroups, each accumulates its own partial +// 26 | // tile in shared memory, and the partials are summed at write-out. +// 27 | // +// 28 | // Preconditions enforced by the host: K % sg_mat_k == 0; B's row stride N_b is +// 29 | // padded to the row-start alignment the load requires and the tile fits the +// 30 | // matrix (kTileM <= M, kTileN <= N - the host falls back otherwise). M and N may +// 31 | // otherwise be any size: a trailing partial tile is shifted back to end exactly +// 32 | // at M / N_b so no operand load runs past the end of A or B, and the rows/columns +// 33 | // the shifted tile re-covers are skipped at write-out. The shift is required for +// 34 | // correctness, not just safety - a subgroupMatrixLoad whose footprint leaves the +// 35 | // buffer does not degrade to a per-element zero fill: Tint's robustness transform +// 36 | // resets the whole load to offset 0 with the minimum stride, so it silently +// 37 | // returns a different, valid tile, which would corrupt the in-range rows of that +// 38 | // tile too. Output width/stride uses N (not N_b), so a padded B column in +// 39 | // [N, N_b) may be read but is never written. When tile_k > 0 none of this applies +// 40 | // because the staging loads are bounds checked and fill zeros outside M N and K so +// 41 | // the host only requires N % 4 == 0 and split_k == 1 and B is never padded. +// 42 | // +// 43 | // Batching: the host dispatches one num_n_tile x num_m_tile tile grid per batch +// 44 | // slice; the batch slice is recovered from the flattened workgroup_idx (the +// 45 | // dispatch may be normalized across x/y/z). A and the output are laid out as +// 46 | // [batch, M, K] / [batch, M, N]; B is either a shared [K, N] weight (batch 1) +// 47 | // or batched as [batch, K, N]. Per-slice flat offsets are derived from M, N, K. +// 48 | // Each slice is an independent tile grid. +// 49 | // +// 50 | #param sg_mat_k +// 51 | #param sg_mat_m +// 52 | #param sg_mat_n +// 53 | #param sg_mat_count_m +// 54 | #param sg_mat_count_n +// 55 | #param split_k +// 56 | #param sg_grid_m +// 57 | #param sg_grid_n +// 58 | #param tile_k +// 59 | // +// 60 | #use .setByOffset +// 61 | // +// 62 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) ss << "const kSgMatM: u32 = u32("; ss << __param_sg_mat_m; ss << ");\n"; -// 56 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +// 63 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) ss << "const kSgMatN: u32 = u32("; ss << __param_sg_mat_n; ss << ");\n"; -// 57 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +// 64 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) ss << "const kSgMatK: u32 = u32("; ss << __param_sg_mat_k; ss << ");\n"; -// 58 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +// 65 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per subgroup along M ss << "const kSgMatCountM: u32 = u32("; ss << __param_sg_mat_count_m; ss << ");\n"; -// 59 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +// 66 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per subgroup along N ss << "const kSgMatCountN: u32 = u32("; ss << __param_sg_mat_count_n; ss << ");\n"; -// 60 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) -ss << "const kTileM: u32 = kSgMatM * kSgMatCountM;\n"; -// 61 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) -ss << "const kTileN: u32 = kSgMatN * kSgMatCountN;\n"; -// 62 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +// 67 | const kSgGridM: u32 = u32(sg_grid_m); // Subgroups splitting the tile along M +ss << "const kSgGridM: u32 = u32("; +ss << __param_sg_grid_m; +ss << ");\n"; +// 68 | const kSgGridN: u32 = u32(sg_grid_n); // Subgroups splitting the tile along N +ss << "const kSgGridN: u32 = u32("; +ss << __param_sg_grid_n; +ss << ");\n"; +// 69 | const kTileM: u32 = kSgMatM * kSgMatCountM * kSgGridM; // Tile size along M (output rows per workgroup) +ss << "const kTileM: u32 = kSgMatM * kSgMatCountM * kSgGridM;\n"; +// 70 | const kTileN: u32 = kSgMatN * kSgMatCountN * kSgGridN; // Tile size along N (output columns per workgroup) +ss << "const kTileN: u32 = kSgMatN * kSgMatCountN * kSgGridN;\n"; +// 71 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 72 | const kTileK: u32 = u32(tile_k); // K depth of A and B staged in shared memory per step +ss << "const kTileK: u32 = u32("; +ss << __param_tile_k; +ss << ");\n"; +// 73 | const kWorkgroupSize: u32 = kSgGridM * kSgGridN * kSubgroupSize; // Threads per workgroup since split_k is 1 here +ss << "const kWorkgroupSize: u32 = kSgGridM * kSgGridN * kSubgroupSize;\n"; +// 74 | const kStrideA: u32 = kTileK + 8u; // Staged A row stride padded by 16 bytes to spread banks +ss << "const kStrideA: u32 = kTileK + 8u;\n"; +// 75 | const kStrideB: u32 = kTileN + 8u; // Staged B row stride padded by 16 bytes to spread banks +ss << "const kStrideB: u32 = kTileN + 8u;\n"; +// 76 | #endif +} +// 77 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K ss << "const kSplitK: u32 = u32("; ss << __param_split_k; ss << ");\n"; -// 63 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +// 78 | const kSubgroupSize: u32 = 32; // Lanes per subgroup ss << "const kSubgroupSize: u32 = 32;\n"; -// 64 | // +// 79 | // ss << "\n"; -// 65 | // Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each -// 66 | // split-K subgroup accumulates its slice of K into its own slot; the slots are -// 67 | // summed during write-back. Storing here also lets us write to global memory -// 68 | // with M/N bounds checks for dimensions that are not multiples of the tile size. -// 69 | var scratch: array; +// 80 | // Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each +// 81 | // split-K subgroup accumulates its slice of K into its own slot; the slots are +// 82 | // summed during write-back. Storing here also lets us write to global memory +// 83 | // with M/N bounds checks for dimensions that are not multiples of the tile size. +// 84 | // When tile_k > 0 it holds the staged kTileM x kTileK slice of A and kTileK x kTileN +// 85 | // slice of B during the K loop and the single output tile after it. +// 86 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 87 | var scratch: array; +ss << "var scratch: array;\n"; +// 88 | #else +} else { +// 89 | var scratch: array; ss << "var scratch: array;\n"; -// 70 | // -ss << "\n"; -// 71 | $MAIN { +// 90 | #endif +} +// 91 | // +// 92 | $MAIN { MainFunctionStart(); ss << "\n"; -// 72 | // Recover the batch slice and the logical (n_tile, m_tile) from the flattened +// 93 | // Recover the batch slice and the logical (n_tile, m_tile) from the flattened ss << "\n"; -// 73 | // workgroup index. The host dispatches num_n_tile x num_m_tile tiles per batch -// 74 | // slice on (x, y, z=batch), but that grid may be normalized (reshaped across -// 75 | // x/y/z to stay within per-dimension limits), so workgroup_id.x/y/z cannot be -// 76 | // used directly. workgroup_idx is the true linear index, invariant under -// 77 | // normalization; batch slices are contiguous blocks of tiles_per_slice tiles. -// 78 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +// 94 | // workgroup index. The host dispatches num_n_tile x num_m_tile tiles per batch +// 95 | // slice on (x, y, z=batch), but that grid may be normalized (reshaped across +// 96 | // x/y/z to stay within per-dimension limits), so workgroup_id.x/y/z cannot be +// 97 | // used directly. workgroup_idx is the true linear index, invariant under +// 98 | // normalization; batch slices are contiguous blocks of tiles_per_slice tiles. +// 99 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; ss << " let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM;\n"; -// 79 | let tiles_per_slice = uniforms.num_n_tile * num_m_tile; +// 100 | let tiles_per_slice = uniforms.num_n_tile * num_m_tile; ss << " let tiles_per_slice = uniforms.num_n_tile * num_m_tile;\n"; -// 80 | let batch_id = workgroup_idx / tiles_per_slice; // which batch slice (0 for a shared 2D weight) +// 101 | let batch_id = workgroup_idx / tiles_per_slice; // which batch slice (0 for a shared 2D weight) ss << " let batch_id = workgroup_idx / tiles_per_slice;\n"; -// 81 | let slice_idx = workgroup_idx % tiles_per_slice; // tile index within the slice +// 102 | let slice_idx = workgroup_idx % tiles_per_slice; // tile index within the slice ss << " let slice_idx = workgroup_idx % tiles_per_slice;\n"; -// 82 | let n_tile = slice_idx % uniforms.num_n_tile; // which kTileN-wide column tile +// 103 | let n_tile = slice_idx % uniforms.num_n_tile; // which kTileN-wide column tile ss << " let n_tile = slice_idx % uniforms.num_n_tile;\n"; -// 83 | let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile +// 104 | let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile ss << " let m_tile = slice_idx / uniforms.num_n_tile;\n"; -// 84 | let n_tile_start = n_tile * kTileN; +// 105 | let n_tile_start = n_tile * kTileN; ss << " let n_tile_start = n_tile * kTileN;\n"; -// 85 | let m_tile_start = m_tile * kTileM; +// 106 | let m_tile_start = m_tile * kTileM; ss << " let m_tile_start = m_tile * kTileM;\n"; -// 86 | // Shift a trailing partial tile back so its operand loads stay inside A and B. -ss << "\n"; -// 87 | // The N shift uses B's padded stride N_b so the shifted origin stays aligned (the -// 88 | // load needs aligned row starts). The extra columns are clipped at write-out. -// 89 | // The max() only guards the u32 subtraction; the host guarantees -// 90 | // the tile fits. -// 91 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); +// 107 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 108 | let a_batch_offset = batch_id * uniforms.M * uniforms.K; +ss << " let a_batch_offset = batch_id * uniforms.M * uniforms.K;\n"; +// 109 | let b_batch_offset = batch_id * uniforms.K * uniforms.N; +ss << " let b_batch_offset = batch_id * uniforms.K * uniforms.N;\n"; +// 110 | let out_batch_offset = batch_id * uniforms.M * uniforms.N; +ss << " let out_batch_offset = batch_id * uniforms.M * uniforms.N;\n"; +// 111 | let sg_index = local_idx / kSubgroupSize; +ss << " let sg_index = local_idx / kSubgroupSize;\n"; +// 112 | let m_base = (sg_index % kSgGridM) * kSgMatM * kSgMatCountM; // first tile row of this subgroup's cell +ss << " let m_base = (sg_index % kSgGridM) * kSgMatM * kSgMatCountM;\n"; +// 113 | let n_base = (sg_index / kSgGridM) * kSgMatN * kSgMatCountN; // first tile column of this subgroup's cell +ss << " let n_base = (sg_index / kSgGridM) * kSgMatN * kSgMatCountN;\n"; +// 114 | let sg_scratch_base = m_base * kTileN + n_base; // this subgroup's cell in the output tile +ss << " let sg_scratch_base = m_base * kTileN + n_base;\n"; +// 115 | #else +} else { +// 116 | // Shift a trailing partial tile back so its operand loads stay inside A and B. +// 117 | // The N shift uses B's padded stride N_b so the shifted origin stays aligned (the +// 118 | // load needs aligned row starts). The extra columns are clipped at write-out. +// 119 | // The max() only guards the u32 subtraction; the host guarantees +// 120 | // the tile fits. +// 121 | 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"; -// 92 | let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM); +// 122 | let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM); ss << " let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM);\n"; -// 93 | // Flat-element offsets into A/B/output for this batch slice, derived from +// 123 | // Flat-element offsets into A/B/output for this batch slice, derived from ss << "\n"; -// 94 | // M/N/K. For a shared 2D weight batch_id is 0, so B collapses to its base. -// 95 | let a_batch_offset = batch_id * uniforms.M * uniforms.K; +// 124 | // M/N/K. For a shared 2D weight batch_id is 0, so B collapses to its base. +// 125 | let a_batch_offset = batch_id * uniforms.M * uniforms.K; ss << " let a_batch_offset = batch_id * uniforms.M * uniforms.K;\n"; -// 96 | let b_batch_offset = batch_id * uniforms.K * uniforms.N_b; +// 126 | let b_batch_offset = batch_id * uniforms.K * uniforms.N_b; ss << " let b_batch_offset = batch_id * uniforms.K * uniforms.N_b;\n"; -// 97 | let out_batch_offset = batch_id * uniforms.M * uniforms.N; +// 127 | let out_batch_offset = batch_id * uniforms.M * uniforms.N; ss << " let out_batch_offset = batch_id * uniforms.M * uniforms.N;\n"; -// 98 | let k_blocks = uniforms.K / kSgMatK; +// 128 | let k_blocks = uniforms.K / kSgMatK; ss << " let k_blocks = uniforms.K / kSgMatK;\n"; -// 99 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +// 129 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) ss << " let sg_index = local_idx / kSubgroupSize;\n"; -// 100 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +// 130 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) ss << " let sg_lane = local_idx % kSubgroupSize;\n"; -// 101 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's scratch slot +// 131 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's scratch slot ss << " let sg_scratch_base = sg_index * kTileM * kTileN;\n"; -// 102 | // -ss << "\n"; -// 103 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). -// 104 | var sg_mat_c0_0: subgroup_matrix_result; +// 132 | #endif +} +// 133 | // +// 134 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 135 | var sg_mat_c0_0: subgroup_matrix_result; ss << " var sg_mat_c0_0: subgroup_matrix_result;\n"; -// 105 | #if sg_mat_count_n >= 2 +// 136 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 106 | var sg_mat_c0_1: subgroup_matrix_result; +// 137 | var sg_mat_c0_1: subgroup_matrix_result; ss << " var sg_mat_c0_1: subgroup_matrix_result;\n"; -// 107 | #endif +// 138 | #endif } -// 108 | #if sg_mat_count_n >= 3 +// 139 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 109 | var sg_mat_c0_2: subgroup_matrix_result; +// 140 | var sg_mat_c0_2: subgroup_matrix_result; ss << " var sg_mat_c0_2: subgroup_matrix_result;\n"; -// 110 | #endif +// 141 | #endif } -// 111 | #if sg_mat_count_n >= 4 +// 142 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 112 | var sg_mat_c0_3: subgroup_matrix_result; +// 143 | var sg_mat_c0_3: subgroup_matrix_result; ss << " var sg_mat_c0_3: subgroup_matrix_result;\n"; -// 113 | #endif +// 144 | #endif } -// 114 | #if sg_mat_count_m >= 2 +// 145 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 115 | var sg_mat_c1_0: subgroup_matrix_result; +// 146 | var sg_mat_c1_0: subgroup_matrix_result; ss << " var sg_mat_c1_0: subgroup_matrix_result;\n"; -// 116 | #if sg_mat_count_n >= 2 +// 147 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 117 | var sg_mat_c1_1: subgroup_matrix_result; +// 148 | var sg_mat_c1_1: subgroup_matrix_result; ss << " var sg_mat_c1_1: subgroup_matrix_result;\n"; -// 118 | #endif +// 149 | #endif } -// 119 | #if sg_mat_count_n >= 3 +// 150 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 120 | var sg_mat_c1_2: subgroup_matrix_result; +// 151 | var sg_mat_c1_2: subgroup_matrix_result; ss << " var sg_mat_c1_2: subgroup_matrix_result;\n"; -// 121 | #endif +// 152 | #endif } -// 122 | #if sg_mat_count_n >= 4 +// 153 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 123 | var sg_mat_c1_3: subgroup_matrix_result; +// 154 | var sg_mat_c1_3: subgroup_matrix_result; ss << " var sg_mat_c1_3: subgroup_matrix_result;\n"; -// 124 | #endif +// 155 | #endif } -// 125 | #endif +// 156 | #endif } -// 126 | #if sg_mat_count_m >= 3 +// 157 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 127 | var sg_mat_c2_0: subgroup_matrix_result; +// 158 | var sg_mat_c2_0: subgroup_matrix_result; ss << " var sg_mat_c2_0: subgroup_matrix_result;\n"; -// 128 | #if sg_mat_count_n >= 2 +// 159 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 129 | var sg_mat_c2_1: subgroup_matrix_result; +// 160 | var sg_mat_c2_1: subgroup_matrix_result; ss << " var sg_mat_c2_1: subgroup_matrix_result;\n"; -// 130 | #endif +// 161 | #endif } -// 131 | #if sg_mat_count_n >= 3 +// 162 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 132 | var sg_mat_c2_2: subgroup_matrix_result; +// 163 | var sg_mat_c2_2: subgroup_matrix_result; ss << " var sg_mat_c2_2: subgroup_matrix_result;\n"; -// 133 | #endif +// 164 | #endif } -// 134 | #if sg_mat_count_n >= 4 +// 165 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 135 | var sg_mat_c2_3: subgroup_matrix_result; +// 166 | var sg_mat_c2_3: subgroup_matrix_result; ss << " var sg_mat_c2_3: subgroup_matrix_result;\n"; -// 136 | #endif +// 167 | #endif } -// 137 | #endif +// 168 | #endif } -// 138 | #if sg_mat_count_m >= 4 +// 169 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 139 | var sg_mat_c3_0: subgroup_matrix_result; +// 170 | var sg_mat_c3_0: subgroup_matrix_result; ss << " var sg_mat_c3_0: subgroup_matrix_result;\n"; -// 140 | #if sg_mat_count_n >= 2 +// 171 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 141 | var sg_mat_c3_1: subgroup_matrix_result; +// 172 | var sg_mat_c3_1: subgroup_matrix_result; ss << " var sg_mat_c3_1: subgroup_matrix_result;\n"; -// 142 | #endif +// 173 | #endif } -// 143 | #if sg_mat_count_n >= 3 +// 174 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 144 | var sg_mat_c3_2: subgroup_matrix_result; +// 175 | var sg_mat_c3_2: subgroup_matrix_result; ss << " var sg_mat_c3_2: subgroup_matrix_result;\n"; -// 145 | #endif +// 176 | #endif } -// 146 | #if sg_mat_count_n >= 4 +// 177 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 147 | var sg_mat_c3_3: subgroup_matrix_result; +// 178 | var sg_mat_c3_3: subgroup_matrix_result; ss << " var sg_mat_c3_3: subgroup_matrix_result;\n"; -// 148 | #endif +// 179 | #endif } -// 149 | #endif +// 180 | #endif } -// 150 | #if sg_mat_count_m >= 5 +// 181 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 151 | var sg_mat_c4_0: subgroup_matrix_result; +// 182 | var sg_mat_c4_0: subgroup_matrix_result; ss << " var sg_mat_c4_0: subgroup_matrix_result;\n"; -// 152 | #if sg_mat_count_n >= 2 +// 183 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 153 | var sg_mat_c4_1: subgroup_matrix_result; +// 184 | var sg_mat_c4_1: subgroup_matrix_result; ss << " var sg_mat_c4_1: subgroup_matrix_result;\n"; -// 154 | #endif +// 185 | #endif } -// 155 | #if sg_mat_count_n >= 3 +// 186 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 156 | var sg_mat_c4_2: subgroup_matrix_result; +// 187 | var sg_mat_c4_2: subgroup_matrix_result; ss << " var sg_mat_c4_2: subgroup_matrix_result;\n"; -// 157 | #endif +// 188 | #endif } -// 158 | #if sg_mat_count_n >= 4 +// 189 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 159 | var sg_mat_c4_3: subgroup_matrix_result; +// 190 | var sg_mat_c4_3: subgroup_matrix_result; ss << " var sg_mat_c4_3: subgroup_matrix_result;\n"; -// 160 | #endif +// 191 | #endif } -// 161 | #endif +// 192 | #endif } -// 162 | #if sg_mat_count_m >= 6 +// 193 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 163 | var sg_mat_c5_0: subgroup_matrix_result; +// 194 | var sg_mat_c5_0: subgroup_matrix_result; ss << " var sg_mat_c5_0: subgroup_matrix_result;\n"; -// 164 | #if sg_mat_count_n >= 2 +// 195 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 165 | var sg_mat_c5_1: subgroup_matrix_result; +// 196 | var sg_mat_c5_1: subgroup_matrix_result; ss << " var sg_mat_c5_1: subgroup_matrix_result;\n"; -// 166 | #endif +// 197 | #endif } -// 167 | #if sg_mat_count_n >= 3 +// 198 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 168 | var sg_mat_c5_2: subgroup_matrix_result; +// 199 | var sg_mat_c5_2: subgroup_matrix_result; ss << " var sg_mat_c5_2: subgroup_matrix_result;\n"; -// 169 | #endif +// 200 | #endif } -// 170 | #if sg_mat_count_n >= 4 +// 201 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 171 | var sg_mat_c5_3: subgroup_matrix_result; +// 202 | var sg_mat_c5_3: subgroup_matrix_result; ss << " var sg_mat_c5_3: subgroup_matrix_result;\n"; -// 172 | #endif +// 203 | #endif } -// 173 | #endif +// 204 | #endif } -// 174 | #if sg_mat_count_m >= 7 +// 205 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 175 | var sg_mat_c6_0: subgroup_matrix_result; +// 206 | var sg_mat_c6_0: subgroup_matrix_result; ss << " var sg_mat_c6_0: subgroup_matrix_result;\n"; -// 176 | #if sg_mat_count_n >= 2 +// 207 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 177 | var sg_mat_c6_1: subgroup_matrix_result; +// 208 | var sg_mat_c6_1: subgroup_matrix_result; ss << " var sg_mat_c6_1: subgroup_matrix_result;\n"; -// 178 | #endif +// 209 | #endif } -// 179 | #if sg_mat_count_n >= 3 +// 210 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 180 | var sg_mat_c6_2: subgroup_matrix_result; +// 211 | var sg_mat_c6_2: subgroup_matrix_result; ss << " var sg_mat_c6_2: subgroup_matrix_result;\n"; -// 181 | #endif +// 212 | #endif } -// 182 | #if sg_mat_count_n >= 4 +// 213 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 183 | var sg_mat_c6_3: subgroup_matrix_result; +// 214 | var sg_mat_c6_3: subgroup_matrix_result; ss << " var sg_mat_c6_3: subgroup_matrix_result;\n"; -// 184 | #endif +// 215 | #endif } -// 185 | #endif +// 216 | #endif } -// 186 | #if sg_mat_count_m >= 8 +// 217 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 187 | var sg_mat_c7_0: subgroup_matrix_result; +// 218 | var sg_mat_c7_0: subgroup_matrix_result; ss << " var sg_mat_c7_0: subgroup_matrix_result;\n"; -// 188 | #if sg_mat_count_n >= 2 +// 219 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 189 | var sg_mat_c7_1: subgroup_matrix_result; +// 220 | var sg_mat_c7_1: subgroup_matrix_result; ss << " var sg_mat_c7_1: subgroup_matrix_result;\n"; -// 190 | #endif +// 221 | #endif } -// 191 | #if sg_mat_count_n >= 3 +// 222 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 192 | var sg_mat_c7_2: subgroup_matrix_result; +// 223 | var sg_mat_c7_2: subgroup_matrix_result; ss << " var sg_mat_c7_2: subgroup_matrix_result;\n"; -// 193 | #endif +// 224 | #endif } -// 194 | #if sg_mat_count_n >= 4 +// 225 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 195 | var sg_mat_c7_3: subgroup_matrix_result; +// 226 | var sg_mat_c7_3: subgroup_matrix_result; ss << " var sg_mat_c7_3: subgroup_matrix_result;\n"; -// 196 | #endif +// 227 | #endif } -// 197 | #endif +// 228 | #endif } -// 198 | // -// 199 | for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { +// 229 | // +// 230 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 231 | for (var k0: u32 = 0u; k0 < uniforms.K; k0 = k0 + kTileK) { +ss << " for (var k0: u32 = 0u; k0 < uniforms.K; k0 = k0 + kTileK) {\n"; +// 232 | // k_start equals k0. Deriving it through a division made this loop up to 40% +ss << "\n"; +// 233 | // faster on an RTX 5080. +// 234 | let k_start = (k0 / kSgMatK) * kSgMatK; +ss << " let k_start = (k0 / kSgMatK) * kSgMatK;\n"; +// 235 | // The barrier keeps the previous slice alive until every subgroup has read it. +ss << "\n"; +// 236 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 237 | for (var i: u32 = local_idx; i < kTileM * kTileK / 4u; i = i + kWorkgroupSize) { +ss << " for (var i: u32 = local_idx; i < kTileM * kTileK / 4u; i = i + kWorkgroupSize) {\n"; +// 238 | let r = i / (kTileK / 4u); +ss << " let r = i / (kTileK / 4u);\n"; +// 239 | let c = (i % (kTileK / 4u)) * 4u; +ss << " let c = (i % (kTileK / 4u)) * 4u;\n"; +// 240 | var v = vec4(); +ss << " var v = vec4();\n"; +// 241 | if (m_tile_start + r < uniforms.M && k_start + c < uniforms.K) { +ss << " if (m_tile_start + r < uniforms.M && k_start + c < uniforms.K) {\n"; +// 242 | v = input_a[(a_batch_offset + (m_tile_start + r) * uniforms.K + k_start + c) / 4u]; +ss << " v = input_a[(a_batch_offset + (m_tile_start + r) * uniforms.K + k_start + c) / 4u];\n"; +// 243 | } +ss << " }\n"; +// 244 | let s = r * kStrideA + c; +ss << " let s = r * kStrideA + c;\n"; +// 245 | scratch[s] = v.x; +ss << " scratch[s] = v.x;\n"; +// 246 | scratch[s + 1u] = v.y; +ss << " scratch[s + 1u] = v.y;\n"; +// 247 | scratch[s + 2u] = v.z; +ss << " scratch[s + 2u] = v.z;\n"; +// 248 | scratch[s + 3u] = v.w; +ss << " scratch[s + 3u] = v.w;\n"; +// 249 | } +ss << " }\n"; +// 250 | for (var i: u32 = local_idx; i < kTileK * kTileN / 4u; i = i + kWorkgroupSize) { +ss << " for (var i: u32 = local_idx; i < kTileK * kTileN / 4u; i = i + kWorkgroupSize) {\n"; +// 251 | let r = i / (kTileN / 4u); +ss << " let r = i / (kTileN / 4u);\n"; +// 252 | let c = (i % (kTileN / 4u)) * 4u; +ss << " let c = (i % (kTileN / 4u)) * 4u;\n"; +// 253 | var v = vec4(); +ss << " var v = vec4();\n"; +// 254 | if (k_start + r < uniforms.K && n_tile_start + c < uniforms.N) { +ss << " if (k_start + r < uniforms.K && n_tile_start + c < uniforms.N) {\n"; +// 255 | v = input_b[(b_batch_offset + (k_start + r) * uniforms.N + n_tile_start + c) / 4u]; +ss << " v = input_b[(b_batch_offset + (k_start + r) * uniforms.N + n_tile_start + c) / 4u];\n"; +// 256 | } +ss << " }\n"; +// 257 | let s = kTileM * kStrideA + r * kStrideB + c; +ss << " let s = kTileM * kStrideA + r * kStrideB + c;\n"; +// 258 | scratch[s] = v.x; +ss << " scratch[s] = v.x;\n"; +// 259 | scratch[s + 1u] = v.y; +ss << " scratch[s + 1u] = v.y;\n"; +// 260 | scratch[s + 2u] = v.z; +ss << " scratch[s + 2u] = v.z;\n"; +// 261 | scratch[s + 3u] = v.w; +ss << " scratch[s + 3u] = v.w;\n"; +// 262 | } +ss << " }\n"; +// 263 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 264 | for (var kk: u32 = 0u; kk < kTileK; kk = kk + kSgMatK) { +ss << " for (var kk: u32 = 0u; kk < kTileK; kk = kk + kSgMatK) {\n"; +// 265 | let b_base = kTileM * kStrideA + kk * kStrideB + n_base; +ss << " let b_base = kTileM * kStrideA + kk * kStrideB + n_base;\n"; +// 266 | var sg_mat_b0: subgroup_matrix_right = +ss << " var sg_mat_b0: subgroup_matrix_right =\n"; +// 267 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 268 | &scratch, b_base + 0 * kSgMatN, kStrideB); +ss << " &scratch, b_base + 0 * kSgMatN, kStrideB);\n"; +// 269 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 270 | var sg_mat_b1: subgroup_matrix_right = +ss << " var sg_mat_b1: subgroup_matrix_right =\n"; +// 271 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 272 | &scratch, b_base + 1 * kSgMatN, kStrideB); +ss << " &scratch, b_base + 1 * kSgMatN, kStrideB);\n"; +// 273 | #endif +} +// 274 | var sg_mat_a0: subgroup_matrix_left = +ss << " var sg_mat_a0: subgroup_matrix_left =\n"; +// 275 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 276 | &scratch, (m_base + 0 * kSgMatM) * kStrideA + kk, kStrideA); +ss << " &scratch, (m_base + 0 * kSgMatM) * kStrideA + kk, kStrideA);\n"; +// 277 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 278 | var sg_mat_a1: subgroup_matrix_left = +ss << " var sg_mat_a1: subgroup_matrix_left =\n"; +// 279 | subgroupMatrixLoad, row_major>( +ss << " subgroupMatrixLoad, row_major>(\n"; +// 280 | &scratch, (m_base + 1 * kSgMatM) * kStrideA + kk, kStrideA); +ss << " &scratch, (m_base + 1 * kSgMatM) * kStrideA + kk, kStrideA);\n"; +// 281 | #endif +} +// 282 | #else +} else { +// 283 | for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { ss << " for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) {\n"; -// 200 | // Load the B right tiles for this K block (KxN_b row-major, stride N_b). +// 284 | // Load the B right tiles for this K block (KxN_b row-major, stride N_b). ss << "\n"; -// 201 | let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n; +// 285 | let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n; ss << " let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n;\n"; -// 202 | var sg_mat_b0: subgroup_matrix_right = +// 286 | var sg_mat_b0: subgroup_matrix_right = ss << " var sg_mat_b0: subgroup_matrix_right =\n"; -// 203 | subgroupMatrixLoad, row_major>( +// 287 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 204 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +// 288 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 0 * kSgMatN, uniforms.N_b);\n"; -// 205 | #if sg_mat_count_n >= 2 +// 289 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 206 | var sg_mat_b1: subgroup_matrix_right = +// 290 | var sg_mat_b1: subgroup_matrix_right = ss << " var sg_mat_b1: subgroup_matrix_right =\n"; -// 207 | subgroupMatrixLoad, row_major>( +// 291 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 208 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); +// 292 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 1 * kSgMatN, uniforms.N_b);\n"; -// 209 | #endif +// 293 | #endif } -// 210 | #if sg_mat_count_n >= 3 +// 294 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 211 | var sg_mat_b2: subgroup_matrix_right = +// 295 | var sg_mat_b2: subgroup_matrix_right = ss << " var sg_mat_b2: subgroup_matrix_right =\n"; -// 212 | subgroupMatrixLoad, row_major>( +// 296 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 213 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +// 297 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 2 * kSgMatN, uniforms.N_b);\n"; -// 214 | #endif +// 298 | #endif } -// 215 | #if sg_mat_count_n >= 4 +// 299 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 216 | var sg_mat_b3: subgroup_matrix_right = +// 300 | var sg_mat_b3: subgroup_matrix_right = ss << " var sg_mat_b3: subgroup_matrix_right =\n"; -// 217 | subgroupMatrixLoad, row_major>( +// 301 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 218 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +// 302 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); ss << " &input_b, b_base + 3 * kSgMatN, uniforms.N_b);\n"; -// 219 | #endif +// 303 | #endif } -// 220 | // -// 221 | // Load the A left tiles (one per M block), row-major with stride K. -// 222 | // a_batch_offset folds this z-slice's flat offset into the column term so -// 223 | // every A load below reads from the correct batch slice. -// 224 | let a_col = a_batch_offset + kb * kSgMatK; +// 304 | // +// 305 | // Load the A left tiles (one per M block), row-major with stride K. +// 306 | // a_batch_offset folds this z-slice's flat offset into the column term so +// 307 | // every A load below reads from the correct batch slice. +// 308 | let a_col = a_batch_offset + kb * kSgMatK; ss << " let a_col = a_batch_offset + kb * kSgMatK;\n"; -// 225 | var sg_mat_a0: subgroup_matrix_left = +// 309 | var sg_mat_a0: subgroup_matrix_left = ss << " var sg_mat_a0: subgroup_matrix_left =\n"; -// 226 | subgroupMatrixLoad, row_major>( +// 310 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 227 | &input_a, (m_base + 0 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 311 | &input_a, (m_base + 0 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 0 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 228 | #if sg_mat_count_m >= 2 +// 312 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 229 | var sg_mat_a1: subgroup_matrix_left = +// 313 | var sg_mat_a1: subgroup_matrix_left = ss << " var sg_mat_a1: subgroup_matrix_left =\n"; -// 230 | subgroupMatrixLoad, row_major>( +// 314 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 231 | &input_a, (m_base + 1 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 315 | &input_a, (m_base + 1 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 1 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 232 | #endif +// 316 | #endif } -// 233 | #if sg_mat_count_m >= 3 +// 317 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 234 | var sg_mat_a2: subgroup_matrix_left = +// 318 | var sg_mat_a2: subgroup_matrix_left = ss << " var sg_mat_a2: subgroup_matrix_left =\n"; -// 235 | subgroupMatrixLoad, row_major>( +// 319 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 236 | &input_a, (m_base + 2 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 320 | &input_a, (m_base + 2 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 2 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 237 | #endif +// 321 | #endif } -// 238 | #if sg_mat_count_m >= 4 +// 322 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 239 | var sg_mat_a3: subgroup_matrix_left = +// 323 | var sg_mat_a3: subgroup_matrix_left = ss << " var sg_mat_a3: subgroup_matrix_left =\n"; -// 240 | subgroupMatrixLoad, row_major>( +// 324 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 241 | &input_a, (m_base + 3 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 325 | &input_a, (m_base + 3 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 3 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 242 | #endif +// 326 | #endif } -// 243 | #if sg_mat_count_m >= 5 +// 327 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 244 | var sg_mat_a4: subgroup_matrix_left = +// 328 | var sg_mat_a4: subgroup_matrix_left = ss << " var sg_mat_a4: subgroup_matrix_left =\n"; -// 245 | subgroupMatrixLoad, row_major>( +// 329 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 246 | &input_a, (m_base + 4 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 330 | &input_a, (m_base + 4 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 4 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 247 | #endif +// 331 | #endif } -// 248 | #if sg_mat_count_m >= 6 +// 332 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 249 | var sg_mat_a5: subgroup_matrix_left = +// 333 | var sg_mat_a5: subgroup_matrix_left = ss << " var sg_mat_a5: subgroup_matrix_left =\n"; -// 250 | subgroupMatrixLoad, row_major>( +// 334 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 251 | &input_a, (m_base + 5 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 335 | &input_a, (m_base + 5 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 5 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 252 | #endif +// 336 | #endif } -// 253 | #if sg_mat_count_m >= 7 +// 337 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 254 | var sg_mat_a6: subgroup_matrix_left = +// 338 | var sg_mat_a6: subgroup_matrix_left = ss << " var sg_mat_a6: subgroup_matrix_left =\n"; -// 255 | subgroupMatrixLoad, row_major>( +// 339 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 256 | &input_a, (m_base + 6 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 340 | &input_a, (m_base + 6 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 6 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 257 | #endif +// 341 | #endif } -// 258 | #if sg_mat_count_m >= 8 +// 342 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 259 | var sg_mat_a7: subgroup_matrix_left = +// 343 | var sg_mat_a7: subgroup_matrix_left = ss << " var sg_mat_a7: subgroup_matrix_left =\n"; -// 260 | subgroupMatrixLoad, row_major>( +// 344 | subgroupMatrixLoad, row_major>( ss << " subgroupMatrixLoad, row_major>(\n"; -// 261 | &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 345 | &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << " &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K);\n"; -// 262 | #endif +// 346 | #endif } -// 263 | // -// 264 | // Accumulate every (M block, N block) pair. -// 265 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +// 347 | // This block stands in for the staged inner loop so both paths close the same braces. +// 348 | { +ss << " {\n"; +// 349 | #endif +} +// 350 | // +// 351 | // Accumulate every (M block, N block) pair. +// 352 | 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"; -// 266 | #if sg_mat_count_n >= 2 +// 353 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 267 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +// 354 | 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"; -// 268 | #endif +// 355 | #endif } -// 269 | #if sg_mat_count_n >= 3 +// 356 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 270 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +// 357 | 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"; -// 271 | #endif +// 358 | #endif } -// 272 | #if sg_mat_count_n >= 4 +// 359 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 273 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +// 360 | 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"; -// 274 | #endif +// 361 | #endif } -// 275 | #if sg_mat_count_m >= 2 +// 362 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 276 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +// 363 | 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"; -// 277 | #if sg_mat_count_n >= 2 +// 364 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 278 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +// 365 | 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"; -// 279 | #endif +// 366 | #endif } -// 280 | #if sg_mat_count_n >= 3 +// 367 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 281 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +// 368 | 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"; -// 282 | #endif +// 369 | #endif } -// 283 | #if sg_mat_count_n >= 4 +// 370 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 284 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +// 371 | 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"; -// 285 | #endif +// 372 | #endif } -// 286 | #endif +// 373 | #endif } -// 287 | #if sg_mat_count_m >= 3 +// 374 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 288 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +// 375 | 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"; -// 289 | #if sg_mat_count_n >= 2 +// 376 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 290 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +// 377 | 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"; -// 291 | #endif +// 378 | #endif } -// 292 | #if sg_mat_count_n >= 3 +// 379 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 293 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +// 380 | 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"; -// 294 | #endif +// 381 | #endif } -// 295 | #if sg_mat_count_n >= 4 +// 382 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 296 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +// 383 | 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"; -// 297 | #endif +// 384 | #endif } -// 298 | #endif +// 385 | #endif } -// 299 | #if sg_mat_count_m >= 4 +// 386 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 300 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +// 387 | 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"; -// 301 | #if sg_mat_count_n >= 2 +// 388 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 302 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +// 389 | 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"; -// 303 | #endif +// 390 | #endif } -// 304 | #if sg_mat_count_n >= 3 +// 391 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 305 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +// 392 | 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"; -// 306 | #endif +// 393 | #endif } -// 307 | #if sg_mat_count_n >= 4 +// 394 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 308 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +// 395 | 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"; -// 309 | #endif +// 396 | #endif } -// 310 | #endif +// 397 | #endif } -// 311 | #if sg_mat_count_m >= 5 +// 398 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 312 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +// 399 | 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"; -// 313 | #if sg_mat_count_n >= 2 +// 400 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 314 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +// 401 | 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"; -// 315 | #endif +// 402 | #endif } -// 316 | #if sg_mat_count_n >= 3 +// 403 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 317 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +// 404 | 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"; -// 318 | #endif +// 405 | #endif } -// 319 | #if sg_mat_count_n >= 4 +// 406 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 320 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +// 407 | 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"; -// 321 | #endif +// 408 | #endif } -// 322 | #endif +// 409 | #endif } -// 323 | #if sg_mat_count_m >= 6 +// 410 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 324 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +// 411 | 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"; -// 325 | #if sg_mat_count_n >= 2 +// 412 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 326 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +// 413 | 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"; -// 327 | #endif +// 414 | #endif } -// 328 | #if sg_mat_count_n >= 3 +// 415 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 329 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +// 416 | 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"; -// 330 | #endif +// 417 | #endif } -// 331 | #if sg_mat_count_n >= 4 +// 418 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 332 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +// 419 | 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"; -// 333 | #endif +// 420 | #endif } -// 334 | #endif +// 421 | #endif } -// 335 | #if sg_mat_count_m >= 7 +// 422 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 336 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +// 423 | 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"; -// 337 | #if sg_mat_count_n >= 2 +// 424 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 338 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +// 425 | 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"; -// 339 | #endif +// 426 | #endif } -// 340 | #if sg_mat_count_n >= 3 +// 427 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 341 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +// 428 | 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"; -// 342 | #endif +// 429 | #endif } -// 343 | #if sg_mat_count_n >= 4 +// 430 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 344 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +// 431 | 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"; -// 345 | #endif +// 432 | #endif } -// 346 | #endif +// 433 | #endif } -// 347 | #if sg_mat_count_m >= 8 +// 434 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 348 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +// 435 | 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"; -// 349 | #if sg_mat_count_n >= 2 +// 436 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 350 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +// 437 | 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"; -// 351 | #endif +// 438 | #endif } -// 352 | #if sg_mat_count_n >= 3 +// 439 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 353 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +// 440 | 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"; -// 354 | #endif +// 441 | #endif } -// 355 | #if sg_mat_count_n >= 4 +// 442 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 356 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +// 443 | 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"; -// 357 | #endif +// 444 | #endif } -// 358 | #endif +// 445 | #endif } -// 359 | } +// 446 | } +ss << " }\n"; +// 447 | } ss << " }\n"; -// 360 | // +// 448 | // ss << "\n"; -// 361 | // Store this subgroup's partial results into its split-K slot of scratch, -// 362 | // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) -// 363 | // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. -// 364 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +// 449 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 450 | // Every subgroup must finish reading the staged slices before scratch holds the output. +// 451 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 452 | #endif +} +// 453 | // Store this subgroup's partial results into its split-K slot of scratch, +// 454 | // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) +// 455 | // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. +// 456 | 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"; -// 365 | #if sg_mat_count_n >= 2 +// 457 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 366 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +// 458 | 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"; -// 367 | #endif +// 459 | #endif } -// 368 | #if sg_mat_count_n >= 3 +// 460 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 369 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +// 461 | 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"; -// 370 | #endif +// 462 | #endif } -// 371 | #if sg_mat_count_n >= 4 +// 463 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 372 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +// 464 | 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"; -// 373 | #endif +// 465 | #endif } -// 374 | #if sg_mat_count_m >= 2 +// 466 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 375 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +// 467 | 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"; -// 376 | #if sg_mat_count_n >= 2 +// 468 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 377 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +// 469 | 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"; -// 378 | #endif +// 470 | #endif } -// 379 | #if sg_mat_count_n >= 3 +// 471 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 380 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +// 472 | 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"; -// 381 | #endif +// 473 | #endif } -// 382 | #if sg_mat_count_n >= 4 +// 474 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 383 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +// 475 | 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"; -// 384 | #endif +// 476 | #endif } -// 385 | #endif +// 477 | #endif } -// 386 | #if sg_mat_count_m >= 3 +// 478 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 387 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +// 479 | 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"; -// 388 | #if sg_mat_count_n >= 2 +// 480 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 389 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +// 481 | 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"; -// 390 | #endif +// 482 | #endif } -// 391 | #if sg_mat_count_n >= 3 +// 483 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 392 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +// 484 | 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"; -// 393 | #endif +// 485 | #endif } -// 394 | #if sg_mat_count_n >= 4 +// 486 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 395 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +// 487 | 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"; -// 396 | #endif +// 488 | #endif } -// 397 | #endif +// 489 | #endif } -// 398 | #if sg_mat_count_m >= 4 +// 490 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 399 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +// 491 | 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"; -// 400 | #if sg_mat_count_n >= 2 +// 492 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 401 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +// 493 | 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"; -// 402 | #endif +// 494 | #endif } -// 403 | #if sg_mat_count_n >= 3 +// 495 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 404 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +// 496 | 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"; -// 405 | #endif +// 497 | #endif } -// 406 | #if sg_mat_count_n >= 4 +// 498 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 407 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +// 499 | 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"; -// 408 | #endif +// 500 | #endif } -// 409 | #endif +// 501 | #endif } -// 410 | #if sg_mat_count_m >= 5 +// 502 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 411 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +// 503 | 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"; -// 412 | #if sg_mat_count_n >= 2 +// 504 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 413 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +// 505 | 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"; -// 414 | #endif +// 506 | #endif } -// 415 | #if sg_mat_count_n >= 3 +// 507 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 416 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +// 508 | 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"; -// 417 | #endif +// 509 | #endif } -// 418 | #if sg_mat_count_n >= 4 +// 510 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 419 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +// 511 | 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"; -// 420 | #endif +// 512 | #endif } -// 421 | #endif +// 513 | #endif } -// 422 | #if sg_mat_count_m >= 6 +// 514 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 423 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +// 515 | 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"; -// 424 | #if sg_mat_count_n >= 2 +// 516 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 425 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +// 517 | 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"; -// 426 | #endif +// 518 | #endif } -// 427 | #if sg_mat_count_n >= 3 +// 519 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 428 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +// 520 | 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"; -// 429 | #endif +// 521 | #endif } -// 430 | #if sg_mat_count_n >= 4 +// 522 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 431 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +// 523 | 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"; -// 432 | #endif +// 524 | #endif } -// 433 | #endif +// 525 | #endif } -// 434 | #if sg_mat_count_m >= 7 +// 526 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 435 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +// 527 | 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"; -// 436 | #if sg_mat_count_n >= 2 +// 528 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 437 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +// 529 | 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"; -// 438 | #endif +// 530 | #endif } -// 439 | #if sg_mat_count_n >= 3 +// 531 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 440 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +// 532 | 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"; -// 441 | #endif +// 533 | #endif } -// 442 | #if sg_mat_count_n >= 4 +// 534 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 443 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +// 535 | 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"; -// 444 | #endif +// 536 | #endif } -// 445 | #endif +// 537 | #endif } -// 446 | #if sg_mat_count_m >= 8 +// 538 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 447 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +// 539 | 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"; -// 448 | #if sg_mat_count_n >= 2 +// 540 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 449 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +// 541 | 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"; -// 450 | #endif +// 542 | #endif } -// 451 | #if sg_mat_count_n >= 3 +// 543 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +// 544 | 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"; -// 453 | #endif +// 545 | #endif } -// 454 | #if sg_mat_count_n >= 4 +// 546 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 455 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +// 547 | 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"; -// 456 | #endif +// 548 | #endif } -// 457 | #endif +// 549 | #endif } -// 458 | // -// 459 | #if split_k >= 2 +// 550 | // +// 551 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 552 | // The barrier makes every cell visible before all threads write out the tile. +// 553 | workgroupBarrier(); +ss << " workgroupBarrier();\n"; +// 554 | for (var i: u32 = local_idx; i < kTileM * kTileN; i = i + kWorkgroupSize) { +ss << " for (var i: u32 = local_idx; i < kTileM * kTileN; i = i + kWorkgroupSize) {\n"; +// 555 | let global_m = m_tile_start + i / kTileN; +ss << " let global_m = m_tile_start + i / kTileN;\n"; +// 556 | let global_n = n_tile_start + i % kTileN; +ss << " let global_n = n_tile_start + i % kTileN;\n"; +// 557 | if (global_m < uniforms.M && global_n < uniforms.N) { +ss << " if (global_m < uniforms.M && global_n < uniforms.N) {\n"; +// 558 | write_output(out_batch_offset + global_m * uniforms.N + global_n, global_n, +ss << " write_output(out_batch_offset + global_m * uniforms.N + global_n, global_n,\n"; +// 559 | output_value_t(scratch[i])); +ss << " output_value_t(scratch[i]));\n"; +// 560 | } +ss << " }\n"; +// 561 | } +ss << " }\n"; +// 562 | #else +} else { +// 563 | #if split_k >= 2 if (__param_split_k >= 2) { -// 460 | // Make all subgroups' partial tiles visible before the reduction. -// 461 | workgroupBarrier(); +// 564 | // Make all subgroups' partial tiles visible before the reduction. +// 565 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 462 | // +// 566 | // ss << "\n"; -// 463 | // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] -// 464 | // tiles into slot 0 of scratch. Each lane owns a strided set of element -// 465 | // positions and sums that position across all kSplitK slots. Writing slot 0 -// 466 | // is safe: a position is only overwritten by the same lane that already read -// 467 | // it as the slot-0 term of its own sum. -// 468 | if (sg_index == 0u) { +// 567 | // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] +// 568 | // tiles into slot 0 of scratch. Each lane owns a strided set of element +// 569 | // positions and sums that position across all kSplitK slots. Writing slot 0 +// 570 | // is safe: a position is only overwritten by the same lane that already read +// 571 | // it as the slot-0 term of its own sum. +// 572 | if (sg_index == 0u) { ss << " if (sg_index == 0u) {\n"; -// 469 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +// 573 | 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"; -// 470 | var acc: f16 = f16(0); +// 574 | var acc: f16 = f16(0); ss << " var acc: f16 = f16(0);\n"; -// 471 | for (var s: u32 = 0; s < kSplitK; s++) { +// 575 | for (var s: u32 = 0; s < kSplitK; s++) { ss << " for (var s: u32 = 0; s < kSplitK; s++) {\n"; -// 472 | acc += scratch[s * kTileM * kTileN + idx]; +// 576 | acc += scratch[s * kTileM * kTileN + idx]; ss << " acc += scratch[s * kTileM * kTileN + idx];\n"; -// 473 | } +// 577 | } ss << " }\n"; -// 474 | scratch[idx] = acc; +// 578 | scratch[idx] = acc; ss << " scratch[idx] = acc;\n"; -// 475 | } +// 579 | } ss << " }\n"; -// 476 | } +// 580 | } ss << " }\n"; -// 477 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +// 581 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << "\n"; -// 478 | workgroupBarrier(); +// 582 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 479 | #endif -} -// 480 | // -// 481 | // Write-out pass: every subgroup writes out whole M-rows from the summed tile -// 482 | // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows -// 483 | // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a -// 484 | // row's N elements (strided by the subgroup size). Rows and columns before the -// 485 | // unshifted tile origin belong to the previous tile and are skipped, so a -// 486 | // shifted tile never writes the same output element twice; M and N are -// 487 | // bounds-checked so non-multiple dimensions only write valid data. -// 488 | let n_skip = n_tile_start - global_base_n; +// 583 | #endif +} +// 584 | // +// 585 | // Write-out pass: every subgroup writes out whole M-rows from the summed tile +// 586 | // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows +// 587 | // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a +// 588 | // row's N elements (strided by the subgroup size). Rows and columns before the +// 589 | // unshifted tile origin belong to the previous tile and are skipped, so a +// 590 | // shifted tile never writes the same output element twice; M and N are +// 591 | // bounds-checked so non-multiple dimensions only write valid data. +// 592 | let n_skip = n_tile_start - global_base_n; ss << " let n_skip = n_tile_start - global_base_n;\n"; -// 489 | let n_count = min(kTileN, uniforms.N - global_base_n); +// 593 | let n_count = min(kTileN, uniforms.N - global_base_n); ss << " let n_count = min(kTileN, uniforms.N - global_base_n);\n"; -// 490 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +// 594 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { ss << " for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) {\n"; -// 491 | let global_m = m_base + r; +// 595 | let global_m = m_base + r; ss << " let global_m = m_base + r;\n"; -// 492 | if (global_m >= m_tile_start && global_m < uniforms.M) { +// 596 | if (global_m >= m_tile_start && global_m < uniforms.M) { ss << " if (global_m >= m_tile_start && global_m < uniforms.M) {\n"; -// 493 | let scratch_base = r * kTileN; +// 597 | let scratch_base = r * kTileN; ss << " let scratch_base = r * kTileN;\n"; -// 494 | let out_base = out_batch_offset + global_m * uniforms.N + global_base_n; +// 598 | let out_base = out_batch_offset + global_m * uniforms.N + global_base_n; ss << " let out_base = out_batch_offset + global_m * uniforms.N + global_base_n;\n"; -// 495 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +// 599 | 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"; -// 496 | write_output(out_base + i, global_base_n + i, +// 600 | write_output(out_base + i, global_base_n + i, ss << " write_output(out_base + i, global_base_n + i,\n"; -// 497 | output_value_t(scratch[scratch_base + i])); +// 601 | output_value_t(scratch[scratch_base + i])); ss << " output_value_t(scratch[scratch_base + i]));\n"; -// 498 | } +// 602 | } ss << " }\n"; -// 499 | } +// 603 | } ss << " }\n"; -// 500 | } +// 604 | } ss << " }\n"; -// 501 | } // MAIN +// 605 | #endif +} +// 606 | } // MAIN MainFunctionEnd(); ss << "\n"; -// 502 | +// 607 | 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 3e837a94da226..2fae15c6fba99 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 @@ -50,12 +50,15 @@ Status ApplyTemplate<"math/subgroup_matrix_gemm_8x16x16.wgsl.template">(ShaderHe template <> struct TemplateParameter<"math/subgroup_matrix_matmul.wgsl.template"> { using type = struct { + int param_sg_grid_m; + int param_sg_grid_n; 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_tile_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 b85499e73881b..187fba7d1acc1 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 @@ -36,7 +36,7 @@ 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.h" // cae16dc433cd1984c18122a20171d9212d83c0e5cb91b51d92e15a215f10cb00 #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 diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h index d44023edb406f..f762e9e766686 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h @@ -10,12 +10,15 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul.wgsl.template">(ShaderHelper& [[maybe_unused]] auto& ss = shader_helper.AdditionalImplementation(); // Extract parameters + auto& __param_sg_grid_m = params.param_sg_grid_m; + auto& __param_sg_grid_n = params.param_sg_grid_n; 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_tile_k = params.param_tile_k; // 1 | // Copyright (c) Microsoft Corporation. All rights reserved. // 2 | // Licensed under the MIT License. @@ -34,1136 +37,1368 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul.wgsl.template">(ShaderHelper& // 15 | // N_b. N_b is the padded row stride of the (possibly column-padded) B buffer and // 16 | // equals N unless the host padded B to an aligned stride to satisfy the load's // 17 | // row-start alignment (see the B padding in subgroup_matrix_matmul.cc). -// 18 | // -// 19 | // Workgroup: split_k subgroups x 32 lanes (split_k in {1,2,4,8}). -// 20 | // Tile: kTileM x kTileN, chosen adaptively by the host config provider. Each -// 21 | // subgroup computes the whole tile as kSgMatCountM x kSgMatCountN subgroup -// 22 | // matrices. K is iterated in sg_mat_k blocks. When split_k > 1 the K blocks are -// 23 | // distributed round-robin across the subgroups, each accumulates its own partial -// 24 | // tile in shared memory, and the partials are summed at write-out. -// 25 | // -// 26 | // Preconditions enforced by the host: K % sg_mat_k == 0; B's row stride N_b is -// 27 | // padded to the row-start alignment the load requires and the tile fits the -// 28 | // matrix (kTileM <= M, kTileN <= N - the host falls back otherwise). M and N may -// 29 | // otherwise be any size: a trailing partial tile is shifted back to end exactly -// 30 | // at M / N_b so no operand load runs past the end of A or B, and the rows/columns -// 31 | // the shifted tile re-covers are skipped at write-out. The shift is required for -// 32 | // correctness, not just safety - a subgroupMatrixLoad whose footprint leaves the -// 33 | // buffer does not degrade to a per-element zero fill: Tint's robustness transform -// 34 | // resets the whole load to offset 0 with the minimum stride, so it silently -// 35 | // returns a different, valid tile, which would corrupt the in-range rows of that -// 36 | // tile too. Output width/stride uses N (not N_b), so a padded B column in -// 37 | // [N, N_b) may be read but is never written. -// 38 | // -// 39 | // Batching: the host dispatches one num_n_tile x num_m_tile tile grid per batch -// 40 | // slice; the batch slice is recovered from the flattened workgroup_idx (the -// 41 | // dispatch may be normalized across x/y/z). A and the output are laid out as -// 42 | // [batch, M, K] / [batch, M, N]; B is either a shared [K, N] weight (batch 1) -// 43 | // or batched as [batch, K, N]. Per-slice flat offsets are derived from M, N, K. -// 44 | // Each slice is an independent tile grid. -// 45 | // -// 46 | #param sg_mat_k -// 47 | #param sg_mat_m -// 48 | #param sg_mat_n -// 49 | #param sg_mat_count_m -// 50 | #param sg_mat_count_n -// 51 | #param split_k -// 52 | // -// 53 | #use .setByOffset -// 54 | // -// 55 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +// 18 | // When tile_k > 0 the workgroup instead stages tile_k deep slices of A and B in +// 19 | // shared memory and the subgroups load their operands from there. +// 20 | // +// 21 | // Workgroup: sg_grid_m x sg_grid_n x split_k subgroups x 32 lanes (split_k in {1,2,4,8}). +// 22 | // Tile: kTileM x kTileN, chosen adaptively by the host config provider. Each +// 23 | // subgroup computes its grid cell as kSgMatCountM x kSgMatCountN subgroup +// 24 | // matrices. K is iterated in sg_mat_k blocks. When split_k > 1 the K blocks are +// 25 | // distributed round-robin across the subgroups, each accumulates its own partial +// 26 | // tile in shared memory, and the partials are summed at write-out. +// 27 | // +// 28 | // Preconditions enforced by the host: K % sg_mat_k == 0; B's row stride N_b is +// 29 | // padded to the row-start alignment the load requires and the tile fits the +// 30 | // matrix (kTileM <= M, kTileN <= N - the host falls back otherwise). M and N may +// 31 | // otherwise be any size: a trailing partial tile is shifted back to end exactly +// 32 | // at M / N_b so no operand load runs past the end of A or B, and the rows/columns +// 33 | // the shifted tile re-covers are skipped at write-out. The shift is required for +// 34 | // correctness, not just safety - a subgroupMatrixLoad whose footprint leaves the +// 35 | // buffer does not degrade to a per-element zero fill: Tint's robustness transform +// 36 | // resets the whole load to offset 0 with the minimum stride, so it silently +// 37 | // returns a different, valid tile, which would corrupt the in-range rows of that +// 38 | // tile too. Output width/stride uses N (not N_b), so a padded B column in +// 39 | // [N, N_b) may be read but is never written. When tile_k > 0 none of this applies +// 40 | // because the staging loads are bounds checked and fill zeros outside M N and K so +// 41 | // the host only requires N % 4 == 0 and split_k == 1 and B is never padded. +// 42 | // +// 43 | // Batching: the host dispatches one num_n_tile x num_m_tile tile grid per batch +// 44 | // slice; the batch slice is recovered from the flattened workgroup_idx (the +// 45 | // dispatch may be normalized across x/y/z). A and the output are laid out as +// 46 | // [batch, M, K] / [batch, M, N]; B is either a shared [K, N] weight (batch 1) +// 47 | // or batched as [batch, K, N]. Per-slice flat offsets are derived from M, N, K. +// 48 | // Each slice is an independent tile grid. +// 49 | // +// 50 | #param sg_mat_k +// 51 | #param sg_mat_m +// 52 | #param sg_mat_n +// 53 | #param sg_mat_count_m +// 54 | #param sg_mat_count_n +// 55 | #param split_k +// 56 | #param sg_grid_m +// 57 | #param sg_grid_n +// 58 | #param tile_k +// 59 | // +// 60 | #use .setByOffset +// 61 | // +// 62 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) ss << __str_32; ss << __param_sg_mat_m; ss << __str_17; -// 56 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +// 63 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) ss << __str_33; ss << __param_sg_mat_n; ss << __str_17; -// 57 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +// 64 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) ss << __str_34; ss << __param_sg_mat_k; ss << __str_17; -// 58 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M +// 65 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per subgroup along M ss << __str_35; ss << __param_sg_mat_count_m; ss << __str_17; -// 59 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N +// 66 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per subgroup along N ss << __str_36; ss << __param_sg_mat_count_n; ss << __str_17; -// 60 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) -ss << __str_37; -// 61 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) -ss << __str_38; -// 62 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K +// 67 | const kSgGridM: u32 = u32(sg_grid_m); // Subgroups splitting the tile along M +ss << __str_218; +ss << __param_sg_grid_m; +ss << __str_17; +// 68 | const kSgGridN: u32 = u32(sg_grid_n); // Subgroups splitting the tile along N +ss << __str_219; +ss << __param_sg_grid_n; +ss << __str_17; +// 69 | const kTileM: u32 = kSgMatM * kSgMatCountM * kSgGridM; // Tile size along M (output rows per workgroup) +ss << __str_220; +// 70 | const kTileN: u32 = kSgMatN * kSgMatCountN * kSgGridN; // Tile size along N (output columns per workgroup) +ss << __str_221; +// 71 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 72 | const kTileK: u32 = u32(tile_k); // K depth of A and B staged in shared memory per step +ss << __str_222; +ss << __param_tile_k; +ss << __str_17; +// 73 | const kWorkgroupSize: u32 = kSgGridM * kSgGridN * kSubgroupSize; // Threads per workgroup since split_k is 1 here +ss << __str_223; +// 74 | const kStrideA: u32 = kTileK + 8u; // Staged A row stride padded by 16 bytes to spread banks +ss << __str_224; +// 75 | const kStrideB: u32 = kTileN + 8u; // Staged B row stride padded by 16 bytes to spread banks +ss << __str_225; +// 76 | #endif +} +// 77 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K ss << __str_39; ss << __param_split_k; ss << __str_17; -// 63 | const kSubgroupSize: u32 = 32; // Lanes per subgroup +// 78 | const kSubgroupSize: u32 = 32; // Lanes per subgroup ss << __str_40; -// 64 | // +// 79 | // ss << __str_5; -// 65 | // Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each -// 66 | // split-K subgroup accumulates its slice of K into its own slot; the slots are -// 67 | // summed during write-back. Storing here also lets us write to global memory -// 68 | // with M/N bounds checks for dimensions that are not multiples of the tile size. -// 69 | var scratch: array; +// 80 | // Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each +// 81 | // split-K subgroup accumulates its slice of K into its own slot; the slots are +// 82 | // summed during write-back. Storing here also lets us write to global memory +// 83 | // with M/N bounds checks for dimensions that are not multiples of the tile size. +// 84 | // When tile_k > 0 it holds the staged kTileM x kTileK slice of A and kTileK x kTileN +// 85 | // slice of B during the K loop and the single output tile after it. +// 86 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 87 | var scratch: array; +ss << __str_226; +// 88 | #else +} else { +// 89 | var scratch: array; ss << __str_41; -// 70 | // -ss << __str_5; -// 71 | $MAIN { +// 90 | #endif +} +// 91 | // +// 92 | $MAIN { MainFunctionStart(); ss << __str_5; -// 72 | // Recover the batch slice and the logical (n_tile, m_tile) from the flattened +// 93 | // Recover the batch slice and the logical (n_tile, m_tile) from the flattened ss << __str_5; -// 73 | // workgroup index. The host dispatches num_n_tile x num_m_tile tiles per batch -// 74 | // slice on (x, y, z=batch), but that grid may be normalized (reshaped across -// 75 | // x/y/z to stay within per-dimension limits), so workgroup_id.x/y/z cannot be -// 76 | // used directly. workgroup_idx is the true linear index, invariant under -// 77 | // normalization; batch slices are contiguous blocks of tiles_per_slice tiles. -// 78 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; -ss << __str_218; -// 79 | let tiles_per_slice = uniforms.num_n_tile * num_m_tile; -ss << __str_219; -// 80 | let batch_id = workgroup_idx / tiles_per_slice; // which batch slice (0 for a shared 2D weight) -ss << __str_220; -// 81 | let slice_idx = workgroup_idx % tiles_per_slice; // tile index within the slice -ss << __str_221; -// 82 | let n_tile = slice_idx % uniforms.num_n_tile; // which kTileN-wide column tile -ss << __str_222; -// 83 | let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile -ss << __str_223; -// 84 | let n_tile_start = n_tile * kTileN; +// 94 | // workgroup index. The host dispatches num_n_tile x num_m_tile tiles per batch +// 95 | // slice on (x, y, z=batch), but that grid may be normalized (reshaped across +// 96 | // x/y/z to stay within per-dimension limits), so workgroup_id.x/y/z cannot be +// 97 | // used directly. workgroup_idx is the true linear index, invariant under +// 98 | // normalization; batch slices are contiguous blocks of tiles_per_slice tiles. +// 99 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; +ss << __str_227; +// 100 | let tiles_per_slice = uniforms.num_n_tile * num_m_tile; +ss << __str_228; +// 101 | let batch_id = workgroup_idx / tiles_per_slice; // which batch slice (0 for a shared 2D weight) +ss << __str_229; +// 102 | let slice_idx = workgroup_idx % tiles_per_slice; // tile index within the slice +ss << __str_230; +// 103 | let n_tile = slice_idx % uniforms.num_n_tile; // which kTileN-wide column tile +ss << __str_231; +// 104 | let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile +ss << __str_232; +// 105 | let n_tile_start = n_tile * kTileN; ss << __str_45; -// 85 | let m_tile_start = m_tile * kTileM; +// 106 | let m_tile_start = m_tile * kTileM; ss << __str_46; -// 86 | // Shift a trailing partial tile back so its operand loads stay inside A and B. -ss << __str_5; -// 87 | // The N shift uses B's padded stride N_b so the shifted origin stays aligned (the -// 88 | // load needs aligned row starts). The extra columns are clipped at write-out. -// 89 | // The max() only guards the u32 subtraction; the host guarantees -// 90 | // the tile fits. -// 91 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); -ss << __str_224; -// 92 | let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM); +// 107 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 108 | let a_batch_offset = batch_id * uniforms.M * uniforms.K; +ss << __str_233; +// 109 | let b_batch_offset = batch_id * uniforms.K * uniforms.N; +ss << __str_234; +// 110 | let out_batch_offset = batch_id * uniforms.M * uniforms.N; +ss << __str_235; +// 111 | let sg_index = local_idx / kSubgroupSize; +ss << __str_50; +// 112 | let m_base = (sg_index % kSgGridM) * kSgMatM * kSgMatCountM; // first tile row of this subgroup's cell +ss << __str_236; +// 113 | let n_base = (sg_index / kSgGridM) * kSgMatN * kSgMatCountN; // first tile column of this subgroup's cell +ss << __str_237; +// 114 | let sg_scratch_base = m_base * kTileN + n_base; // this subgroup's cell in the output tile +ss << __str_238; +// 115 | #else +} else { +// 116 | // Shift a trailing partial tile back so its operand loads stay inside A and B. +// 117 | // The N shift uses B's padded stride N_b so the shifted origin stays aligned (the +// 118 | // load needs aligned row starts). The extra columns are clipped at write-out. +// 119 | // The max() only guards the u32 subtraction; the host guarantees +// 120 | // the tile fits. +// 121 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); +ss << __str_239; +// 122 | let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM); ss << __str_48; -// 93 | // Flat-element offsets into A/B/output for this batch slice, derived from +// 123 | // Flat-element offsets into A/B/output for this batch slice, derived from ss << __str_5; -// 94 | // M/N/K. For a shared 2D weight batch_id is 0, so B collapses to its base. -// 95 | let a_batch_offset = batch_id * uniforms.M * uniforms.K; -ss << __str_225; -// 96 | let b_batch_offset = batch_id * uniforms.K * uniforms.N_b; -ss << __str_226; -// 97 | let out_batch_offset = batch_id * uniforms.M * uniforms.N; -ss << __str_227; -// 98 | let k_blocks = uniforms.K / kSgMatK; +// 124 | // M/N/K. For a shared 2D weight batch_id is 0, so B collapses to its base. +// 125 | let a_batch_offset = batch_id * uniforms.M * uniforms.K; +ss << __str_233; +// 126 | let b_batch_offset = batch_id * uniforms.K * uniforms.N_b; +ss << __str_240; +// 127 | let out_batch_offset = batch_id * uniforms.M * uniforms.N; +ss << __str_235; +// 128 | let k_blocks = uniforms.K / kSgMatK; ss << __str_49; -// 99 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) +// 129 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) ss << __str_50; -// 100 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) +// 130 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) ss << __str_51; -// 101 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's scratch slot +// 131 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's scratch slot ss << __str_52; -// 102 | // -ss << __str_5; -// 103 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). -// 104 | var sg_mat_c0_0: subgroup_matrix_result; +// 132 | #endif +} +// 133 | // +// 134 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). +// 135 | 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; -// 105 | #if sg_mat_count_n >= 2 +// 136 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 106 | var sg_mat_c0_1: subgroup_matrix_result; +// 137 | 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; -// 107 | #endif +// 138 | #endif } -// 108 | #if sg_mat_count_n >= 3 +// 139 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 109 | var sg_mat_c0_2: subgroup_matrix_result; +// 140 | 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; -// 110 | #endif +// 141 | #endif } -// 111 | #if sg_mat_count_n >= 4 +// 142 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 112 | var sg_mat_c0_3: subgroup_matrix_result; +// 143 | 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; -// 113 | #endif +// 144 | #endif } -// 114 | #if sg_mat_count_m >= 2 +// 145 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 115 | var sg_mat_c1_0: subgroup_matrix_result; +// 146 | 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; -// 116 | #if sg_mat_count_n >= 2 +// 147 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 117 | var sg_mat_c1_1: subgroup_matrix_result; +// 148 | 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; -// 118 | #endif +// 149 | #endif } -// 119 | #if sg_mat_count_n >= 3 +// 150 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 120 | var sg_mat_c1_2: subgroup_matrix_result; +// 151 | 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; -// 121 | #endif +// 152 | #endif } -// 122 | #if sg_mat_count_n >= 4 +// 153 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 123 | var sg_mat_c1_3: subgroup_matrix_result; +// 154 | 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; -// 124 | #endif +// 155 | #endif } -// 125 | #endif +// 156 | #endif } -// 126 | #if sg_mat_count_m >= 3 +// 157 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 127 | var sg_mat_c2_0: subgroup_matrix_result; +// 158 | 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; -// 128 | #if sg_mat_count_n >= 2 +// 159 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 129 | var sg_mat_c2_1: subgroup_matrix_result; +// 160 | 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; -// 130 | #endif +// 161 | #endif } -// 131 | #if sg_mat_count_n >= 3 +// 162 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 132 | var sg_mat_c2_2: subgroup_matrix_result; +// 163 | 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; -// 133 | #endif +// 164 | #endif } -// 134 | #if sg_mat_count_n >= 4 +// 165 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 135 | var sg_mat_c2_3: subgroup_matrix_result; +// 166 | 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; -// 136 | #endif +// 167 | #endif } -// 137 | #endif +// 168 | #endif } -// 138 | #if sg_mat_count_m >= 4 +// 169 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 139 | var sg_mat_c3_0: subgroup_matrix_result; +// 170 | 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; -// 140 | #if sg_mat_count_n >= 2 +// 171 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 141 | var sg_mat_c3_1: subgroup_matrix_result; +// 172 | 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; -// 142 | #endif +// 173 | #endif } -// 143 | #if sg_mat_count_n >= 3 +// 174 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 144 | var sg_mat_c3_2: subgroup_matrix_result; +// 175 | 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; -// 145 | #endif +// 176 | #endif } -// 146 | #if sg_mat_count_n >= 4 +// 177 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 147 | var sg_mat_c3_3: subgroup_matrix_result; +// 178 | 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; -// 148 | #endif +// 179 | #endif } -// 149 | #endif +// 180 | #endif } -// 150 | #if sg_mat_count_m >= 5 +// 181 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 151 | var sg_mat_c4_0: subgroup_matrix_result; +// 182 | 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; -// 152 | #if sg_mat_count_n >= 2 +// 183 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 153 | var sg_mat_c4_1: subgroup_matrix_result; +// 184 | 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; -// 154 | #endif +// 185 | #endif } -// 155 | #if sg_mat_count_n >= 3 +// 186 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 156 | var sg_mat_c4_2: subgroup_matrix_result; +// 187 | 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; -// 157 | #endif +// 188 | #endif } -// 158 | #if sg_mat_count_n >= 4 +// 189 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 159 | var sg_mat_c4_3: subgroup_matrix_result; +// 190 | 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; -// 160 | #endif +// 191 | #endif } -// 161 | #endif +// 192 | #endif } -// 162 | #if sg_mat_count_m >= 6 +// 193 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 163 | var sg_mat_c5_0: subgroup_matrix_result; +// 194 | 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; -// 164 | #if sg_mat_count_n >= 2 +// 195 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 165 | var sg_mat_c5_1: subgroup_matrix_result; +// 196 | 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; -// 166 | #endif +// 197 | #endif } -// 167 | #if sg_mat_count_n >= 3 +// 198 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 168 | var sg_mat_c5_2: subgroup_matrix_result; +// 199 | 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; -// 169 | #endif +// 200 | #endif } -// 170 | #if sg_mat_count_n >= 4 +// 201 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 171 | var sg_mat_c5_3: subgroup_matrix_result; +// 202 | 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; -// 172 | #endif +// 203 | #endif } -// 173 | #endif +// 204 | #endif } -// 174 | #if sg_mat_count_m >= 7 +// 205 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 175 | var sg_mat_c6_0: subgroup_matrix_result; +// 206 | 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; -// 176 | #if sg_mat_count_n >= 2 +// 207 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 177 | var sg_mat_c6_1: subgroup_matrix_result; +// 208 | 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; -// 178 | #endif +// 209 | #endif } -// 179 | #if sg_mat_count_n >= 3 +// 210 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 180 | var sg_mat_c6_2: subgroup_matrix_result; +// 211 | 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; -// 181 | #endif +// 212 | #endif } -// 182 | #if sg_mat_count_n >= 4 +// 213 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 183 | var sg_mat_c6_3: subgroup_matrix_result; +// 214 | 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; -// 184 | #endif +// 215 | #endif } -// 185 | #endif +// 216 | #endif } -// 186 | #if sg_mat_count_m >= 8 +// 217 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 187 | var sg_mat_c7_0: subgroup_matrix_result; +// 218 | 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; -// 188 | #if sg_mat_count_n >= 2 +// 219 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 189 | var sg_mat_c7_1: subgroup_matrix_result; +// 220 | 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; -// 190 | #endif +// 221 | #endif } -// 191 | #if sg_mat_count_n >= 3 +// 222 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 192 | var sg_mat_c7_2: subgroup_matrix_result; +// 223 | 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; -// 193 | #endif +// 224 | #endif } -// 194 | #if sg_mat_count_n >= 4 +// 225 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 195 | var sg_mat_c7_3: subgroup_matrix_result; +// 226 | 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; -// 196 | #endif +// 227 | #endif } -// 197 | #endif +// 228 | #endif } -// 198 | // -// 199 | for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { +// 229 | // +// 230 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 231 | for (var k0: u32 = 0u; k0 < uniforms.K; k0 = k0 + kTileK) { +ss << __str_241; +// 232 | // k_start equals k0. Deriving it through a division made this loop up to 40% +ss << __str_5; +// 233 | // faster on an RTX 5080. +// 234 | let k_start = (k0 / kSgMatK) * kSgMatK; +ss << __str_242; +// 235 | // The barrier keeps the previous slice alive until every subgroup has read it. +ss << __str_5; +// 236 | workgroupBarrier(); +ss << __str_243; +// 237 | for (var i: u32 = local_idx; i < kTileM * kTileK / 4u; i = i + kWorkgroupSize) { +ss << __str_244; +// 238 | let r = i / (kTileK / 4u); +ss << __str_245; +// 239 | let c = (i % (kTileK / 4u)) * 4u; +ss << __str_246; +// 240 | var v = vec4(); +ss << __str_247; +// 241 | if (m_tile_start + r < uniforms.M && k_start + c < uniforms.K) { +ss << __str_248; +// 242 | v = input_a[(a_batch_offset + (m_tile_start + r) * uniforms.K + k_start + c) / 4u]; +ss << __str_249; +// 243 | } +ss << __str_202; +// 244 | let s = r * kStrideA + c; +ss << __str_250; +// 245 | scratch[s] = v.x; +ss << __str_251; +// 246 | scratch[s + 1u] = v.y; +ss << __str_252; +// 247 | scratch[s + 2u] = v.z; +ss << __str_253; +// 248 | scratch[s + 3u] = v.w; +ss << __str_254; +// 249 | } +ss << __str_204; +// 250 | for (var i: u32 = local_idx; i < kTileK * kTileN / 4u; i = i + kWorkgroupSize) { +ss << __str_255; +// 251 | let r = i / (kTileN / 4u); +ss << __str_256; +// 252 | let c = (i % (kTileN / 4u)) * 4u; +ss << __str_257; +// 253 | var v = vec4(); +ss << __str_247; +// 254 | if (k_start + r < uniforms.K && n_tile_start + c < uniforms.N) { +ss << __str_258; +// 255 | v = input_b[(b_batch_offset + (k_start + r) * uniforms.N + n_tile_start + c) / 4u]; +ss << __str_259; +// 256 | } +ss << __str_202; +// 257 | let s = kTileM * kStrideA + r * kStrideB + c; +ss << __str_260; +// 258 | scratch[s] = v.x; +ss << __str_251; +// 259 | scratch[s + 1u] = v.y; +ss << __str_252; +// 260 | scratch[s + 2u] = v.z; +ss << __str_253; +// 261 | scratch[s + 3u] = v.w; +ss << __str_254; +// 262 | } +ss << __str_204; +// 263 | workgroupBarrier(); +ss << __str_243; +// 264 | for (var kk: u32 = 0u; kk < kTileK; kk = kk + kSgMatK) { +ss << __str_261; +// 265 | let b_base = kTileM * kStrideA + kk * kStrideB + n_base; +ss << __str_262; +// 266 | var sg_mat_b0: subgroup_matrix_right = +ss << __str_263; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 267 | subgroupMatrixLoad, row_major>( +ss << __str_264; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_101; +// 268 | &scratch, b_base + 0 * kSgMatN, kStrideB); +ss << __str_265; +// 269 | #if sg_mat_count_n >= 2 +if (__param_sg_mat_count_n >= 2) { +// 270 | var sg_mat_b1: subgroup_matrix_right = +ss << __str_266; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_90; +// 271 | subgroupMatrixLoad, row_major>( +ss << __str_264; +ss << __param_sg_mat_n; +ss << __str_54; +ss << __param_sg_mat_k; +ss << __str_101; +// 272 | &scratch, b_base + 1 * kSgMatN, kStrideB); +ss << __str_267; +// 273 | #endif +} +// 274 | var sg_mat_a0: subgroup_matrix_left = +ss << __str_268; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 275 | subgroupMatrixLoad, row_major>( +ss << __str_269; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 276 | &scratch, (m_base + 0 * kSgMatM) * kStrideA + kk, kStrideA); +ss << __str_270; +// 277 | #if sg_mat_count_m >= 2 +if (__param_sg_mat_count_m >= 2) { +// 278 | var sg_mat_a1: subgroup_matrix_left = +ss << __str_271; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_90; +// 279 | subgroupMatrixLoad, row_major>( +ss << __str_269; +ss << __param_sg_mat_k; +ss << __str_54; +ss << __param_sg_mat_m; +ss << __str_101; +// 280 | &scratch, (m_base + 1 * kSgMatM) * kStrideA + kk, kStrideA); +ss << __str_272; +// 281 | #endif +} +// 282 | #else +} else { +// 283 | for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { ss << __str_87; -// 200 | // Load the B right tiles for this K block (KxN_b row-major, stride N_b). +// 284 | // Load the B right tiles for this K block (KxN_b row-major, stride N_b). ss << __str_5; -// 201 | let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n; -ss << __str_228; -// 202 | var sg_mat_b0: subgroup_matrix_right = +// 285 | let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n; +ss << __str_273; +// 286 | var sg_mat_b0: subgroup_matrix_right = ss << __str_89; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_90; -// 203 | subgroupMatrixLoad, row_major>( +// 287 | subgroupMatrixLoad, row_major>( ss << __str_91; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_101; -// 204 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); -ss << __str_229; -// 205 | #if sg_mat_count_n >= 2 +// 288 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); +ss << __str_274; +// 289 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 206 | var sg_mat_b1: subgroup_matrix_right = +// 290 | var sg_mat_b1: subgroup_matrix_right = ss << __str_94; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_90; -// 207 | subgroupMatrixLoad, row_major>( +// 291 | subgroupMatrixLoad, row_major>( ss << __str_91; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_101; -// 208 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); -ss << __str_230; -// 209 | #endif +// 292 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); +ss << __str_275; +// 293 | #endif } -// 210 | #if sg_mat_count_n >= 3 +// 294 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 211 | var sg_mat_b2: subgroup_matrix_right = +// 295 | var sg_mat_b2: subgroup_matrix_right = ss << __str_96; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_90; -// 212 | subgroupMatrixLoad, row_major>( +// 296 | subgroupMatrixLoad, row_major>( ss << __str_91; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_101; -// 213 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); -ss << __str_231; -// 214 | #endif +// 297 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); +ss << __str_276; +// 298 | #endif } -// 215 | #if sg_mat_count_n >= 4 +// 299 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 216 | var sg_mat_b3: subgroup_matrix_right = +// 300 | var sg_mat_b3: subgroup_matrix_right = ss << __str_98; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_90; -// 217 | subgroupMatrixLoad, row_major>( +// 301 | subgroupMatrixLoad, row_major>( ss << __str_91; ss << __param_sg_mat_n; ss << __str_54; ss << __param_sg_mat_k; ss << __str_101; -// 218 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); -ss << __str_232; -// 219 | #endif +// 302 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); +ss << __str_277; +// 303 | #endif } -// 220 | // -// 221 | // Load the A left tiles (one per M block), row-major with stride K. -// 222 | // a_batch_offset folds this z-slice's flat offset into the column term so -// 223 | // every A load below reads from the correct batch slice. -// 224 | let a_col = a_batch_offset + kb * kSgMatK; -ss << __str_233; -// 225 | var sg_mat_a0: subgroup_matrix_left = +// 304 | // +// 305 | // Load the A left tiles (one per M block), row-major with stride K. +// 306 | // a_batch_offset folds this z-slice's flat offset into the column term so +// 307 | // every A load below reads from the correct batch slice. +// 308 | let a_col = a_batch_offset + kb * kSgMatK; +ss << __str_278; +// 309 | var sg_mat_a0: subgroup_matrix_left = ss << __str_107; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 226 | subgroupMatrixLoad, row_major>( +// 310 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 227 | &input_a, (m_base + 0 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 311 | &input_a, (m_base + 0 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_125; -// 228 | #if sg_mat_count_m >= 2 +// 312 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 229 | var sg_mat_a1: subgroup_matrix_left = +// 313 | var sg_mat_a1: subgroup_matrix_left = ss << __str_110; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 230 | subgroupMatrixLoad, row_major>( +// 314 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 231 | &input_a, (m_base + 1 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 315 | &input_a, (m_base + 1 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_126; -// 232 | #endif +// 316 | #endif } -// 233 | #if sg_mat_count_m >= 3 +// 317 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 234 | var sg_mat_a2: subgroup_matrix_left = +// 318 | var sg_mat_a2: subgroup_matrix_left = ss << __str_112; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 235 | subgroupMatrixLoad, row_major>( +// 319 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 236 | &input_a, (m_base + 2 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 320 | &input_a, (m_base + 2 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_127; -// 237 | #endif +// 321 | #endif } -// 238 | #if sg_mat_count_m >= 4 +// 322 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 239 | var sg_mat_a3: subgroup_matrix_left = +// 323 | var sg_mat_a3: subgroup_matrix_left = ss << __str_114; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 240 | subgroupMatrixLoad, row_major>( +// 324 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 241 | &input_a, (m_base + 3 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 325 | &input_a, (m_base + 3 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_128; -// 242 | #endif +// 326 | #endif } -// 243 | #if sg_mat_count_m >= 5 +// 327 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 244 | var sg_mat_a4: subgroup_matrix_left = +// 328 | var sg_mat_a4: subgroup_matrix_left = ss << __str_116; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 245 | subgroupMatrixLoad, row_major>( +// 329 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 246 | &input_a, (m_base + 4 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 330 | &input_a, (m_base + 4 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_129; -// 247 | #endif +// 331 | #endif } -// 248 | #if sg_mat_count_m >= 6 +// 332 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 249 | var sg_mat_a5: subgroup_matrix_left = +// 333 | var sg_mat_a5: subgroup_matrix_left = ss << __str_118; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 250 | subgroupMatrixLoad, row_major>( +// 334 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 251 | &input_a, (m_base + 5 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 335 | &input_a, (m_base + 5 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_130; -// 252 | #endif +// 336 | #endif } -// 253 | #if sg_mat_count_m >= 7 +// 337 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 254 | var sg_mat_a6: subgroup_matrix_left = +// 338 | var sg_mat_a6: subgroup_matrix_left = ss << __str_120; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 255 | subgroupMatrixLoad, row_major>( +// 339 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 256 | &input_a, (m_base + 6 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 340 | &input_a, (m_base + 6 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_131; -// 257 | #endif +// 341 | #endif } -// 258 | #if sg_mat_count_m >= 8 +// 342 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 259 | var sg_mat_a7: subgroup_matrix_left = +// 343 | var sg_mat_a7: subgroup_matrix_left = ss << __str_122; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_90; -// 260 | subgroupMatrixLoad, row_major>( +// 344 | subgroupMatrixLoad, row_major>( ss << __str_108; ss << __param_sg_mat_k; ss << __str_54; ss << __param_sg_mat_m; ss << __str_101; -// 261 | &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K); +// 345 | &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K); ss << __str_132; -// 262 | #endif +// 346 | #endif } -// 263 | // -// 264 | // Accumulate every (M block, N block) pair. -// 265 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); +// 347 | // This block stands in for the staged inner loop so both paths close the same braces. +// 348 | { +ss << __str_279; +// 349 | #endif +} +// 350 | // +// 351 | // Accumulate every (M block, N block) pair. +// 352 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); ss << __str_133; -// 266 | #if sg_mat_count_n >= 2 +// 353 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 267 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); +// 354 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); ss << __str_134; -// 268 | #endif +// 355 | #endif } -// 269 | #if sg_mat_count_n >= 3 +// 356 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 270 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); +// 357 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); ss << __str_135; -// 271 | #endif +// 358 | #endif } -// 272 | #if sg_mat_count_n >= 4 +// 359 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 273 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); +// 360 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); ss << __str_136; -// 274 | #endif +// 361 | #endif } -// 275 | #if sg_mat_count_m >= 2 +// 362 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 276 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); +// 363 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); ss << __str_137; -// 277 | #if sg_mat_count_n >= 2 +// 364 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 278 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); +// 365 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); ss << __str_138; -// 279 | #endif +// 366 | #endif } -// 280 | #if sg_mat_count_n >= 3 +// 367 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 281 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); +// 368 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); ss << __str_139; -// 282 | #endif +// 369 | #endif } -// 283 | #if sg_mat_count_n >= 4 +// 370 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 284 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); +// 371 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); ss << __str_140; -// 285 | #endif +// 372 | #endif } -// 286 | #endif +// 373 | #endif } -// 287 | #if sg_mat_count_m >= 3 +// 374 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 288 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); +// 375 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); ss << __str_141; -// 289 | #if sg_mat_count_n >= 2 +// 376 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 290 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); +// 377 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); ss << __str_142; -// 291 | #endif +// 378 | #endif } -// 292 | #if sg_mat_count_n >= 3 +// 379 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 293 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); +// 380 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); ss << __str_143; -// 294 | #endif +// 381 | #endif } -// 295 | #if sg_mat_count_n >= 4 +// 382 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 296 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); +// 383 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); ss << __str_144; -// 297 | #endif +// 384 | #endif } -// 298 | #endif +// 385 | #endif } -// 299 | #if sg_mat_count_m >= 4 +// 386 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 300 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); +// 387 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); ss << __str_145; -// 301 | #if sg_mat_count_n >= 2 +// 388 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 302 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); +// 389 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); ss << __str_146; -// 303 | #endif +// 390 | #endif } -// 304 | #if sg_mat_count_n >= 3 +// 391 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 305 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); +// 392 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); ss << __str_147; -// 306 | #endif +// 393 | #endif } -// 307 | #if sg_mat_count_n >= 4 +// 394 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 308 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); +// 395 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); ss << __str_148; -// 309 | #endif +// 396 | #endif } -// 310 | #endif +// 397 | #endif } -// 311 | #if sg_mat_count_m >= 5 +// 398 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 312 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); +// 399 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); ss << __str_149; -// 313 | #if sg_mat_count_n >= 2 +// 400 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 314 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); +// 401 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); ss << __str_150; -// 315 | #endif +// 402 | #endif } -// 316 | #if sg_mat_count_n >= 3 +// 403 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 317 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); +// 404 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); ss << __str_151; -// 318 | #endif +// 405 | #endif } -// 319 | #if sg_mat_count_n >= 4 +// 406 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 320 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); +// 407 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); ss << __str_152; -// 321 | #endif +// 408 | #endif } -// 322 | #endif +// 409 | #endif } -// 323 | #if sg_mat_count_m >= 6 +// 410 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 324 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); +// 411 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); ss << __str_153; -// 325 | #if sg_mat_count_n >= 2 +// 412 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 326 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); +// 413 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); ss << __str_154; -// 327 | #endif +// 414 | #endif } -// 328 | #if sg_mat_count_n >= 3 +// 415 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 329 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); +// 416 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); ss << __str_155; -// 330 | #endif +// 417 | #endif } -// 331 | #if sg_mat_count_n >= 4 +// 418 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 332 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); +// 419 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); ss << __str_156; -// 333 | #endif +// 420 | #endif } -// 334 | #endif +// 421 | #endif } -// 335 | #if sg_mat_count_m >= 7 +// 422 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 336 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); +// 423 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); ss << __str_157; -// 337 | #if sg_mat_count_n >= 2 +// 424 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 338 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); +// 425 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); ss << __str_158; -// 339 | #endif +// 426 | #endif } -// 340 | #if sg_mat_count_n >= 3 +// 427 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 341 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); +// 428 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); ss << __str_159; -// 342 | #endif +// 429 | #endif } -// 343 | #if sg_mat_count_n >= 4 +// 430 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 344 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); +// 431 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); ss << __str_160; -// 345 | #endif +// 432 | #endif } -// 346 | #endif +// 433 | #endif } -// 347 | #if sg_mat_count_m >= 8 +// 434 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 348 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); +// 435 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); ss << __str_161; -// 349 | #if sg_mat_count_n >= 2 +// 436 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 350 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); +// 437 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); ss << __str_162; -// 351 | #endif +// 438 | #endif } -// 352 | #if sg_mat_count_n >= 3 +// 439 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 353 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); +// 440 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); ss << __str_163; -// 354 | #endif +// 441 | #endif } -// 355 | #if sg_mat_count_n >= 4 +// 442 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 356 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); +// 443 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); ss << __str_164; -// 357 | #endif +// 444 | #endif } -// 358 | #endif +// 445 | #endif } -// 359 | } +// 446 | } +ss << __str_204; +// 447 | } ss << __str_20; -// 360 | // +// 448 | // ss << __str_5; -// 361 | // Store this subgroup's partial results into its split-K slot of scratch, -// 362 | // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) -// 363 | // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. -// 364 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); +// 449 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 450 | // Every subgroup must finish reading the staged slices before scratch holds the output. +// 451 | workgroupBarrier(); +ss << __str_26; +// 452 | #endif +} +// 453 | // Store this subgroup's partial results into its split-K slot of scratch, +// 454 | // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) +// 455 | // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. +// 456 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); ss << __str_165; -// 365 | #if sg_mat_count_n >= 2 +// 457 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 366 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); +// 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); ss << __str_166; -// 367 | #endif +// 459 | #endif } -// 368 | #if sg_mat_count_n >= 3 +// 460 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 369 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); +// 461 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); ss << __str_167; -// 370 | #endif +// 462 | #endif } -// 371 | #if sg_mat_count_n >= 4 +// 463 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 372 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); +// 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); ss << __str_168; -// 373 | #endif +// 465 | #endif } -// 374 | #if sg_mat_count_m >= 2 +// 466 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { -// 375 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); +// 467 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); ss << __str_169; -// 376 | #if sg_mat_count_n >= 2 +// 468 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 377 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); +// 469 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); ss << __str_170; -// 378 | #endif +// 470 | #endif } -// 379 | #if sg_mat_count_n >= 3 +// 471 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 380 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); +// 472 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); ss << __str_171; -// 381 | #endif +// 473 | #endif } -// 382 | #if sg_mat_count_n >= 4 +// 474 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 383 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); +// 475 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); ss << __str_172; -// 384 | #endif +// 476 | #endif } -// 385 | #endif +// 477 | #endif } -// 386 | #if sg_mat_count_m >= 3 +// 478 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { -// 387 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); +// 479 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); ss << __str_173; -// 388 | #if sg_mat_count_n >= 2 +// 480 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 389 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); +// 481 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); ss << __str_174; -// 390 | #endif +// 482 | #endif } -// 391 | #if sg_mat_count_n >= 3 +// 483 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 392 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); +// 484 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); ss << __str_175; -// 393 | #endif +// 485 | #endif } -// 394 | #if sg_mat_count_n >= 4 +// 486 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 395 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); +// 487 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); ss << __str_176; -// 396 | #endif +// 488 | #endif } -// 397 | #endif +// 489 | #endif } -// 398 | #if sg_mat_count_m >= 4 +// 490 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { -// 399 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); +// 491 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); ss << __str_177; -// 400 | #if sg_mat_count_n >= 2 +// 492 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 401 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); +// 493 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); ss << __str_178; -// 402 | #endif +// 494 | #endif } -// 403 | #if sg_mat_count_n >= 3 +// 495 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 404 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); +// 496 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); ss << __str_179; -// 405 | #endif +// 497 | #endif } -// 406 | #if sg_mat_count_n >= 4 +// 498 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 407 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); +// 499 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); ss << __str_180; -// 408 | #endif +// 500 | #endif } -// 409 | #endif +// 501 | #endif } -// 410 | #if sg_mat_count_m >= 5 +// 502 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { -// 411 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); +// 503 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); ss << __str_181; -// 412 | #if sg_mat_count_n >= 2 +// 504 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 413 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); +// 505 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); ss << __str_182; -// 414 | #endif +// 506 | #endif } -// 415 | #if sg_mat_count_n >= 3 +// 507 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 416 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); +// 508 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); ss << __str_183; -// 417 | #endif +// 509 | #endif } -// 418 | #if sg_mat_count_n >= 4 +// 510 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 419 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); +// 511 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); ss << __str_184; -// 420 | #endif +// 512 | #endif } -// 421 | #endif +// 513 | #endif } -// 422 | #if sg_mat_count_m >= 6 +// 514 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { -// 423 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); +// 515 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); ss << __str_185; -// 424 | #if sg_mat_count_n >= 2 +// 516 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 425 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); +// 517 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); ss << __str_186; -// 426 | #endif +// 518 | #endif } -// 427 | #if sg_mat_count_n >= 3 +// 519 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 428 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); +// 520 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); ss << __str_187; -// 429 | #endif +// 521 | #endif } -// 430 | #if sg_mat_count_n >= 4 +// 522 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 431 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); +// 523 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); ss << __str_188; -// 432 | #endif +// 524 | #endif } -// 433 | #endif +// 525 | #endif } -// 434 | #if sg_mat_count_m >= 7 +// 526 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { -// 435 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); +// 527 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); ss << __str_189; -// 436 | #if sg_mat_count_n >= 2 +// 528 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 437 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); +// 529 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); ss << __str_190; -// 438 | #endif +// 530 | #endif } -// 439 | #if sg_mat_count_n >= 3 +// 531 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 440 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); +// 532 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); ss << __str_191; -// 441 | #endif +// 533 | #endif } -// 442 | #if sg_mat_count_n >= 4 +// 534 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 443 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); +// 535 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); ss << __str_192; -// 444 | #endif +// 536 | #endif } -// 445 | #endif +// 537 | #endif } -// 446 | #if sg_mat_count_m >= 8 +// 538 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { -// 447 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); +// 539 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); ss << __str_193; -// 448 | #if sg_mat_count_n >= 2 +// 540 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { -// 449 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); +// 541 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); ss << __str_194; -// 450 | #endif +// 542 | #endif } -// 451 | #if sg_mat_count_n >= 3 +// 543 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { -// 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); +// 544 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); ss << __str_195; -// 453 | #endif +// 545 | #endif } -// 454 | #if sg_mat_count_n >= 4 +// 546 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { -// 455 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); +// 547 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); ss << __str_196; -// 456 | #endif +// 548 | #endif } -// 457 | #endif +// 549 | #endif } -// 458 | // -// 459 | #if split_k >= 2 +// 550 | // +// 551 | #if tile_k > 0 +if (__param_tile_k > 0) { +// 552 | // The barrier makes every cell visible before all threads write out the tile. +// 553 | workgroupBarrier(); +ss << __str_26; +// 554 | for (var i: u32 = local_idx; i < kTileM * kTileN; i = i + kWorkgroupSize) { +ss << __str_280; +// 555 | let global_m = m_tile_start + i / kTileN; +ss << __str_281; +// 556 | let global_n = n_tile_start + i % kTileN; +ss << __str_282; +// 557 | if (global_m < uniforms.M && global_n < uniforms.N) { +ss << __str_283; +// 558 | write_output(out_batch_offset + global_m * uniforms.N + global_n, global_n, +ss << __str_284; +// 559 | output_value_t(scratch[i])); +ss << __str_285; +// 560 | } +ss << __str_204; +// 561 | } +ss << __str_20; +// 562 | #else +} else { +// 563 | #if split_k >= 2 if (__param_split_k >= 2) { -// 460 | // Make all subgroups' partial tiles visible before the reduction. -// 461 | workgroupBarrier(); +// 564 | // Make all subgroups' partial tiles visible before the reduction. +// 565 | workgroupBarrier(); ss << __str_26; -// 462 | // +// 566 | // ss << __str_5; -// 463 | // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] -// 464 | // tiles into slot 0 of scratch. Each lane owns a strided set of element -// 465 | // positions and sums that position across all kSplitK slots. Writing slot 0 -// 466 | // is safe: a position is only overwritten by the same lane that already read -// 467 | // it as the slot-0 term of its own sum. -// 468 | if (sg_index == 0u) { +// 567 | // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] +// 568 | // tiles into slot 0 of scratch. Each lane owns a strided set of element +// 569 | // positions and sums that position across all kSplitK slots. Writing slot 0 +// 570 | // is safe: a position is only overwritten by the same lane that already read +// 571 | // it as the slot-0 term of its own sum. +// 572 | if (sg_index == 0u) { ss << __str_197; -// 469 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { +// 573 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { ss << __str_198; -// 470 | var acc: f16 = f16(0); +// 574 | var acc: f16 = f16(0); ss << __str_199; -// 471 | for (var s: u32 = 0; s < kSplitK; s++) { +// 575 | for (var s: u32 = 0; s < kSplitK; s++) { ss << __str_200; -// 472 | acc += scratch[s * kTileM * kTileN + idx]; +// 576 | acc += scratch[s * kTileM * kTileN + idx]; ss << __str_201; -// 473 | } +// 577 | } ss << __str_202; -// 474 | scratch[idx] = acc; +// 578 | scratch[idx] = acc; ss << __str_203; -// 475 | } +// 579 | } ss << __str_204; -// 476 | } +// 580 | } ss << __str_20; -// 477 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. +// 581 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << __str_5; -// 478 | workgroupBarrier(); +// 582 | workgroupBarrier(); ss << __str_26; -// 479 | #endif -} -// 480 | // -// 481 | // Write-out pass: every subgroup writes out whole M-rows from the summed tile -// 482 | // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows -// 483 | // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a -// 484 | // row's N elements (strided by the subgroup size). Rows and columns before the -// 485 | // unshifted tile origin belong to the previous tile and are skipped, so a -// 486 | // shifted tile never writes the same output element twice; M and N are -// 487 | // bounds-checked so non-multiple dimensions only write valid data. -// 488 | let n_skip = n_tile_start - global_base_n; +// 583 | #endif +} +// 584 | // +// 585 | // Write-out pass: every subgroup writes out whole M-rows from the summed tile +// 586 | // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows +// 587 | // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a +// 588 | // row's N elements (strided by the subgroup size). Rows and columns before the +// 589 | // unshifted tile origin belong to the previous tile and are skipped, so a +// 590 | // shifted tile never writes the same output element twice; M and N are +// 591 | // bounds-checked so non-multiple dimensions only write valid data. +// 592 | let n_skip = n_tile_start - global_base_n; ss << __str_205; -// 489 | let n_count = min(kTileN, uniforms.N - global_base_n); +// 593 | let n_count = min(kTileN, uniforms.N - global_base_n); ss << __str_206; -// 490 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { +// 594 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { ss << __str_207; -// 491 | let global_m = m_base + r; +// 595 | let global_m = m_base + r; ss << __str_208; -// 492 | if (global_m >= m_tile_start && global_m < uniforms.M) { +// 596 | if (global_m >= m_tile_start && global_m < uniforms.M) { ss << __str_209; -// 493 | let scratch_base = r * kTileN; +// 597 | let scratch_base = r * kTileN; ss << __str_210; -// 494 | let out_base = out_batch_offset + global_m * uniforms.N + global_base_n; -ss << __str_234; -// 495 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { +// 598 | let out_base = out_batch_offset + global_m * uniforms.N + global_base_n; +ss << __str_286; +// 599 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { ss << __str_212; -// 496 | write_output(out_base + i, global_base_n + i, -ss << __str_235; -// 497 | output_value_t(scratch[scratch_base + i])); -ss << __str_236; -// 498 | } +// 600 | write_output(out_base + i, global_base_n + i, +ss << __str_287; +// 601 | output_value_t(scratch[scratch_base + i])); +ss << __str_288; +// 602 | } ss << __str_202; -// 499 | } +// 603 | } ss << __str_204; -// 500 | } +// 604 | } ss << __str_20; -// 501 | } // MAIN +// 605 | #endif +} +// 606 | } // MAIN MainFunctionEnd(); ss << __str_5; -// 502 | +// 607 | return Status::OK(); 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..89e5e2c43676f 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 @@ -28,26 +28,26 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul_pad_b.wgsl.template">(ShaderHe MainFunctionStart(); ss << __str_5; // 13 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); -ss << __str_241; -ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_237); +ss << __str_293; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_289); ss << __str_29; // 14 | let r = global_idx / uniforms.N_b; // padded-output row -ss << __str_242; +ss << __str_294; // 15 | let c = global_idx % uniforms.N_b; // padded-output column -ss << __str_243; +ss << __str_295; // 16 | var v = output_value_t(0); -ss << __str_244; +ss << __str_296; // 17 | if (c < uniforms.N) { // real column -> copy; else zero pad -ss << __str_245; +ss << __str_297; // 18 | v = output_value_t(input_b.getByOffset(r * uniforms.N + c)); -ss << __str_246; -ss << __var_input_b->GetByOffset(__str_238); +ss << __str_298; +ss << __var_input_b->GetByOffset(__str_290); ss << __str_17; // 19 | } ss << __str_11; // 20 | output.setByOffset(global_idx, v); -ss << __str_241; -ss << __var_output->SetByOffset(__str_239, __str_240); +ss << __str_293; +ss << __var_output->SetByOffset(__str_291, __str_292); ss << __str_29; // 21 | } // MAIN MainFunctionEnd(); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h index f86290a3a1457..f01ddcbcd7d81 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h @@ -35,15 +35,15 @@ Status ApplyTemplate<"nn/conv_transpose3d.wgsl.template">(ShaderHelper& shader_h // 12 | #if components == 4 if (__param_components == 4) { // 13 | alias Accumulator = vec4; -ss << __str_262; +ss << __str_314; // 14 | #elif components == 2 } else if (__param_components == 2) { // 15 | alias Accumulator = vec2; -ss << __str_263; +ss << __str_315; // 16 | #else } else { // 17 | alias Accumulator = f32; -ss << __str_264; +ss << __str_316; // 18 | #endif } // 19 | @@ -51,198 +51,198 @@ ss << __str_264; MainFunctionStart(); ss << __str_5; // 21 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); -ss << __str_241; -ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_237); +ss << __str_293; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_289); ss << __str_29; // 22 | let output_indices = output.offsetToIndices(global_idx); -ss << __str_265; -ss << __var_output->OffsetToIndices(__str_239); +ss << __str_317; +ss << __var_output->OffsetToIndices(__str_291); ss << __str_29; // 23 | let batch = output.indicesGet(output_indices, 0); -ss << __str_266; -ss << __var_output->IndicesGet(__str_247, __str_248); +ss << __str_318; +ss << __var_output->IndicesGet(__str_299, __str_300); ss << __str_29; // 24 | #if is_channels_last if (__param_is_channels_last) { // 25 | let output_channel = output.indicesGet(output_indices, 4); -ss << __str_267; -ss << __var_output->IndicesGet(__str_247, __str_249); +ss << __str_319; +ss << __var_output->IndicesGet(__str_299, __str_301); ss << __str_29; // 26 | let spatial_dim = 1u; -ss << __str_268; +ss << __str_320; // 27 | #else } else { // 28 | let output_channel = output.indicesGet(output_indices, 1); -ss << __str_267; -ss << __var_output->IndicesGet(__str_247, __str_250); +ss << __str_319; +ss << __var_output->IndicesGet(__str_299, __str_302); ss << __str_29; // 29 | let spatial_dim = 2u; -ss << __str_269; +ss << __str_321; // 30 | #endif } // 31 | #if is_prepacked if (__param_is_prepacked) { // 32 | let output_channels_per_group = getElementAt(uniforms.w_shape, 3, w.rank); -ss << __str_270; -ss << GetElementAt(__str_251, __str_252, __var_w->Rank()); +ss << __str_322; +ss << GetElementAt(__str_303, __str_304, __var_w->Rank()); ss << __str_29; // 33 | let kernel_shape = vec3(getElementAt(uniforms.w_shape, 0, w.rank), -ss << __str_271; -ss << GetElementAt(__str_251, __str_248, __var_w->Rank()); -ss << __str_272; +ss << __str_323; +ss << GetElementAt(__str_303, __str_300, __var_w->Rank()); +ss << __str_324; // 34 | getElementAt(uniforms.w_shape, 1, w.rank), -ss << __str_273; -ss << GetElementAt(__str_251, __str_250, __var_w->Rank()); -ss << __str_272; +ss << __str_325; +ss << GetElementAt(__str_303, __str_302, __var_w->Rank()); +ss << __str_324; // 35 | getElementAt(uniforms.w_shape, 2, w.rank)); -ss << __str_273; -ss << GetElementAt(__str_251, __str_253, __var_w->Rank()); +ss << __str_325; +ss << GetElementAt(__str_303, __str_305, __var_w->Rank()); ss << __str_17; // 36 | #else } else { // 37 | let output_channels_per_group = getElementAt(uniforms.w_shape, 1, w.rank); -ss << __str_270; -ss << GetElementAt(__str_251, __str_250, __var_w->Rank()); +ss << __str_322; +ss << GetElementAt(__str_303, __str_302, __var_w->Rank()); ss << __str_29; // 38 | let kernel_shape = vec3(getElementAt(uniforms.w_shape, 2, w.rank), -ss << __str_271; -ss << GetElementAt(__str_251, __str_253, __var_w->Rank()); -ss << __str_272; +ss << __str_323; +ss << GetElementAt(__str_303, __str_305, __var_w->Rank()); +ss << __str_324; // 39 | getElementAt(uniforms.w_shape, 3, w.rank), -ss << __str_273; -ss << GetElementAt(__str_251, __str_252, __var_w->Rank()); -ss << __str_272; +ss << __str_325; +ss << GetElementAt(__str_303, __str_304, __var_w->Rank()); +ss << __str_324; // 40 | getElementAt(uniforms.w_shape, 4, w.rank)); -ss << __str_273; -ss << GetElementAt(__str_251, __str_249, __var_w->Rank()); +ss << __str_325; +ss << GetElementAt(__str_303, __str_301, __var_w->Rank()); ss << __str_17; // 41 | #endif } // 42 | let group = output_channel / output_channels_per_group; -ss << __str_274; +ss << __str_326; // 43 | let w_channel = output_channel % output_channels_per_group; -ss << __str_275; +ss << __str_327; // 44 | let corner = vec3( -ss << __str_276; +ss << __str_328; // 45 | i32(output.indicesGet(output_indices, spatial_dim)), -ss << __str_277; -ss << __var_output->IndicesGet(__str_247, __str_254); -ss << __str_278; +ss << __str_329; +ss << __var_output->IndicesGet(__str_299, __str_306); +ss << __str_330; // 46 | i32(output.indicesGet(output_indices, spatial_dim + 1u)), -ss << __str_277; -ss << __var_output->IndicesGet(__str_247, __str_255); -ss << __str_278; +ss << __str_329; +ss << __var_output->IndicesGet(__str_299, __str_307); +ss << __str_330; // 47 | i32(output.indicesGet(output_indices, spatial_dim + 2u))) + vec3(uniforms.pads); -ss << __str_277; -ss << __var_output->IndicesGet(__str_247, __str_256); -ss << __str_279; +ss << __str_329; +ss << __var_output->IndicesGet(__str_299, __str_308); +ss << __str_331; // 48 | var value = Accumulator(0); -ss << __str_280; +ss << __str_332; // 49 | ss << __str_5; // 50 | // Gather input positions using signed integer coordinates so padding and stride // 51 | // alignment remain exact even for float16 tensors with large spatial dimensions. // 52 | for (var kd = 0u; kd < kernel_shape.x; kd++) { -ss << __str_281; +ss << __str_333; // 53 | let pos_d = corner.x - i32(kd) * i32(uniforms.dilations.x); -ss << __str_282; +ss << __str_334; // 54 | if (pos_d < 0 || pos_d % i32(uniforms.strides.x) != 0) { -ss << __str_283; +ss << __str_335; // 55 | continue; -ss << __str_284; +ss << __str_336; // 56 | } ss << __str_20; // 57 | let id = u32(pos_d) / uniforms.strides.x; -ss << __str_285; +ss << __str_337; // 58 | if (id >= getElementAt(uniforms.x_shape, spatial_dim, x.rank)) { -ss << __str_286; -ss << GetElementAt(__str_257, __str_254, __var_x->Rank()); -ss << __str_287; +ss << __str_338; +ss << GetElementAt(__str_309, __str_306, __var_x->Rank()); +ss << __str_339; // 59 | continue; -ss << __str_284; +ss << __str_336; // 60 | } ss << __str_20; // 61 | for (var kh = 0u; kh < kernel_shape.y; kh++) { -ss << __str_288; +ss << __str_340; // 62 | let pos_h = corner.y - i32(kh) * i32(uniforms.dilations.y); -ss << __str_289; +ss << __str_341; // 63 | if (pos_h < 0 || pos_h % i32(uniforms.strides.y) != 0) { -ss << __str_290; +ss << __str_342; // 64 | continue; -ss << __str_291; +ss << __str_343; // 65 | } -ss << __str_292; +ss << __str_344; // 66 | let ih = u32(pos_h) / uniforms.strides.y; -ss << __str_293; +ss << __str_345; // 67 | if (ih >= getElementAt(uniforms.x_shape, spatial_dim + 1u, x.rank)) { -ss << __str_294; -ss << GetElementAt(__str_257, __str_255, __var_x->Rank()); -ss << __str_287; +ss << __str_346; +ss << GetElementAt(__str_309, __str_307, __var_x->Rank()); +ss << __str_339; // 68 | continue; -ss << __str_291; +ss << __str_343; // 69 | } -ss << __str_292; +ss << __str_344; // 70 | for (var kw = 0u; kw < kernel_shape.z; kw++) { -ss << __str_295; +ss << __str_347; // 71 | let pos_w = corner.z - i32(kw) * i32(uniforms.dilations.z); -ss << __str_296; +ss << __str_348; // 72 | if (pos_w < 0 || pos_w % i32(uniforms.strides.z) != 0) { -ss << __str_297; +ss << __str_349; // 73 | continue; -ss << __str_298; +ss << __str_350; // 74 | } ss << __str_204; // 75 | let iw = u32(pos_w) / uniforms.strides.z; -ss << __str_299; +ss << __str_351; // 76 | if (iw >= getElementAt(uniforms.x_shape, spatial_dim + 2u, x.rank)) { -ss << __str_300; -ss << GetElementAt(__str_257, __str_256, __var_x->Rank()); -ss << __str_287; +ss << __str_352; +ss << GetElementAt(__str_309, __str_308, __var_x->Rank()); +ss << __str_339; // 77 | continue; -ss << __str_298; +ss << __str_350; // 78 | } ss << __str_204; // 79 | for (var c = 0u; c < uniforms.input_channels_per_group; c += components) { -ss << __str_301; +ss << __str_353; ss << __param_components; -ss << __str_287; +ss << __str_339; // 80 | let input_channel = group * uniforms.input_channels_per_group + c; -ss << __str_302; +ss << __str_354; // 81 | #if is_channels_last if (__param_is_channels_last) { // 82 | let x_indices = x_indices_t(batch, id, ih, iw, input_channel / components); -ss << __str_303; +ss << __str_355; ss << __param_components; ss << __str_17; // 83 | let x_value = Accumulator(x.getByIndices(x_indices)); -ss << __str_304; -ss << __var_x->GetByIndices(__str_258); +ss << __str_356; +ss << __var_x->GetByIndices(__str_310); ss << __str_17; // 84 | #else } else { // 85 | #if components > 1 if (__param_components > 1) { // 86 | var x_value = Accumulator(0); -ss << __str_305; +ss << __str_357; // 87 | for (var lane = 0u; lane < components; lane++) { -ss << __str_306; +ss << __str_358; ss << __param_components; -ss << __str_307; +ss << __str_359; // 88 | let x_indices = x_indices_t(batch, input_channel + lane, id, ih, iw); -ss << __str_308; +ss << __str_360; // 89 | x_value[lane] = f32(x.getByIndices(x_indices)); -ss << __str_309; -ss << __var_x->GetByIndices(__str_258); +ss << __str_361; +ss << __var_x->GetByIndices(__str_310); ss << __str_17; // 90 | } -ss << __str_310; +ss << __str_362; // 91 | #else } else { // 92 | let x_indices = x_indices_t(batch, input_channel, id, ih, iw); -ss << __str_311; +ss << __str_363; // 93 | let x_value = f32(x.getByIndices(x_indices)); -ss << __str_312; -ss << __var_x->GetByIndices(__str_258); +ss << __str_364; +ss << __var_x->GetByIndices(__str_310); ss << __str_17; // 94 | #endif } @@ -251,23 +251,23 @@ ss << __str_17; // 96 | #if is_prepacked if (__param_is_prepacked) { // 97 | let w_indices = w_indices_t(kd, kh, kw, w_channel, input_channel / components); -ss << __str_313; +ss << __str_365; ss << __param_components; ss << __str_17; // 98 | #else } else { // 99 | let w_indices = w_indices_t(input_channel, w_channel, kd, kh, kw); -ss << __str_314; +ss << __str_366; // 100 | #endif } // 101 | value += x_value * Accumulator(w.getByIndices(w_indices)); -ss << __str_315; -ss << __var_w->GetByIndices(__str_259); +ss << __str_367; +ss << __var_w->GetByIndices(__str_311); ss << __str_17; // 102 | } ss << __str_204; // 103 | } -ss << __str_292; +ss << __str_344; // 104 | } ss << __str_20; // 105 | } @@ -275,24 +275,24 @@ ss << __str_11; // 106 | #if components > 1 if (__param_components > 1) { // 107 | var result = dot(value, Accumulator(1)); -ss << __str_316; +ss << __str_368; // 108 | #else } else { // 109 | var result = value; -ss << __str_317; +ss << __str_369; // 110 | #endif } // 111 | #if has_bias if (__param_has_bias) { // 112 | result += f32(bias.getByOffset(output_channel)); -ss << __str_318; -ss << __var_bias->GetByOffset(__str_260); +ss << __str_370; +ss << __var_bias->GetByOffset(__str_312); ss << __str_17; // 113 | #endif } // 114 | output.setByOffset(global_idx, output_value_t(result)); -ss << __str_241; -ss << __var_output->SetByOffset(__str_239, __str_261); +ss << __str_293; +ss << __var_output->SetByOffset(__str_291, __str_313); ss << __str_29; // 115 | } MainFunctionEnd(); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h index 84976d4a91fa1..1141f90026bf9 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h @@ -43,222 +43,222 @@ Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_help // 18 | // im2col access for src: [N, H_i, W_i, C_i / vec_size] // 19 | // Conceptual Matrix Shape: N * (H_o * W_o) x (K_h * K_w * C_i / vec_size) // 20 | fn load_src(batch : u32, m : u32, k_packed_idx : u32) -> src_value_t { -ss << __str_323; +ss << __str_375; // 21 | if (batch >= uniforms.batch || m >= uniforms.im2col_m || k_packed_idx * vec_size >= uniforms.im2col_k) { -ss << __str_324; +ss << __str_376; ss << __param_vec_size; -ss << __str_325; +ss << __str_377; // 22 | return src_value_t(); -ss << __str_326; +ss << __str_378; // 23 | } ss << __str_11; // 24 | ss << __str_5; // 25 | let channel_i_vec = uniforms.channel_i / vec_size; -ss << __str_327; +ss << __str_379; ss << __param_vec_size; ss << __str_29; // 26 | ss << __str_5; // 27 | // 1. Decompose M index (H_o * W_o) into (h_idx, w_idx) // 28 | let h_idx = m / uniforms.output_w; // Output H index (H_o) -ss << __str_328; +ss << __str_380; // 29 | let w_idx = m % uniforms.output_w; // Output W index (W_o) -ss << __str_329; +ss << __str_381; // 30 | ss << __str_5; // 31 | // 2. Decompose K index into (k_h, k_w, c_i_vec_idx) // 32 | let c_i_vec_idx = k_packed_idx % channel_i_vec; -ss << __str_330; +ss << __str_382; // 33 | let k_h_w_idx = k_packed_idx / channel_i_vec; -ss << __str_331; +ss << __str_383; // 34 | let k_h = k_h_w_idx / uniforms.kernel_w; // Kernel Row -ss << __str_332; +ss << __str_384; // 35 | let k_w = k_h_w_idx % uniforms.kernel_w; // Kernel Column -ss << __str_333; +ss << __str_385; // 36 | ss << __str_5; // 37 | // 3. Calculate the coordinate in the padded input tensor // 38 | let src_h_coord_padded = h_idx * uniforms.strides.x + k_h * uniforms.dilations.x; -ss << __str_334; +ss << __str_386; // 39 | let src_w_coord_padded = w_idx * uniforms.strides.y + k_w * uniforms.dilations.y; -ss << __str_335; +ss << __str_387; // 40 | ss << __str_5; // 41 | // 4. Calculate the coordinate in the original input tensor // 42 | let src_h_coord : i32 = i32(src_h_coord_padded) - i32(uniforms.pads.x); -ss << __str_336; +ss << __str_388; // 43 | let src_w_coord : i32 = i32(src_w_coord_padded) - i32(uniforms.pads.y); -ss << __str_337; +ss << __str_389; // 44 | ss << __str_5; // 45 | // 5. Check for padding/out-of-bounds // 46 | if (src_h_coord < 0 || src_h_coord >= i32(uniforms.src_h) || -ss << __str_338; +ss << __str_390; // 47 | src_w_coord < 0 || src_w_coord >= i32(uniforms.src_w)) { -ss << __str_339; +ss << __str_391; // 48 | return src_value_t(); -ss << __str_326; +ss << __str_378; // 49 | } ss << __str_11; // 50 | ss << __str_5; // 51 | // 6. Calculate final NHWC index // 52 | let src_idx = batch * uniforms.src_h * uniforms.src_w * channel_i_vec + -ss << __str_340; +ss << __str_392; // 53 | u32(src_h_coord) * uniforms.src_w * channel_i_vec + -ss << __str_341; +ss << __str_393; // 54 | u32(src_w_coord) * channel_i_vec + -ss << __str_342; +ss << __str_394; // 55 | c_i_vec_idx; -ss << __str_343; +ss << __str_395; // 56 | return src.getByOffset(src_idx); -ss << __str_344; -ss << __var_src->GetByOffset(__str_319); +ss << __str_396; +ss << __var_src->GetByOffset(__str_371); ss << __str_29; // 57 | } -ss << __str_345; +ss << __str_397; // 58 | ss << __str_5; // 59 | // weight shape: [Co, K_h, K_w, C_i / vec_size] (CoHWCi) // 60 | fn load_weight(n : u32, k_packed_idx : u32) -> weight_value_t { -ss << __str_346; +ss << __str_398; // 61 | if (n < uniforms.im2col_n && k_packed_idx < uniforms.im2col_k / vec_size) { -ss << __str_347; +ss << __str_399; ss << __param_vec_size; -ss << __str_287; +ss << __str_339; // 62 | let weight_idx = n * uniforms.im2col_k / vec_size + -ss << __str_348; +ss << __str_400; ss << __param_vec_size; -ss << __str_349; +ss << __str_401; // 63 | k_packed_idx; -ss << __str_350; +ss << __str_402; // 64 | return weight.getByOffset(weight_idx); -ss << __str_351; -ss << __var_weight->GetByOffset(__str_320); +ss << __str_403; +ss << __var_weight->GetByOffset(__str_372); ss << __str_29; // 65 | } ss << __str_11; // 66 | return weight_value_t(); -ss << __str_352; +ss << __str_404; // 67 | } -ss << __str_345; +ss << __str_397; // 68 | ss << __str_5; // 69 | fn load_bias(n : u32) -> output_element_t { -ss << __str_353; +ss << __str_405; // 70 | #if has_bias if (__param_has_bias) { // 71 | if (n < uniforms.im2col_n) { -ss << __str_354; +ss << __str_406; // 72 | return output_element_t(bias[n]); -ss << __str_355; +ss << __str_407; // 73 | } ss << __str_11; // 74 | #endif } // 75 | return output_element_t(); -ss << __str_356; +ss << __str_408; // 76 | } -ss << __str_345; +ss << __str_397; // 77 | ss << __str_5; // 78 | // output shape: [N, H_o, W_o, C_o] (NHWC) // 79 | fn write_output(batch : u32, m : u32, n : u32, value : output_element_t) { -ss << __str_357; +ss << __str_409; // 80 | if (batch < uniforms.batch && m < uniforms.im2col_m && n < uniforms.im2col_n) { -ss << __str_358; +ss << __str_410; // 81 | let output_idx = batch * uniforms.im2col_m * uniforms.im2col_n + -ss << __str_359; +ss << __str_411; // 82 | m * uniforms.im2col_n + -ss << __str_360; +ss << __str_412; // 83 | n; -ss << __str_361; +ss << __str_413; // 84 | output.setByOffset(output_idx, value); ss << __str_28; -ss << __var_output->SetByOffset(__str_321, __str_322); +ss << __var_output->SetByOffset(__str_373, __str_374); ss << __str_29; // 85 | } ss << __str_11; // 86 | } -ss << __str_345; +ss << __str_397; // 87 | ss << __str_5; // 88 | const TILE_M_SIZE : u32 = tile_m; -ss << __str_362; +ss << __str_414; ss << __param_tile_m; ss << __str_29; // 89 | const TILE_N_SIZE : u32 = tile_n; -ss << __str_363; +ss << __str_415; ss << __param_tile_n; ss << __str_29; // 90 | // In dimension K, the tile consists of 16 scalars, requiring `16 / vec_size` vector loads. ss << __str_5; // 91 | const TILE_K_VEC_SIZE : u32 = 16 / vec_size; -ss << __str_364; +ss << __str_416; ss << __param_vec_size; ss << __str_29; // 92 | // In dimensions M and N, since a workgroup has 64 threads, it advances by `64 / TILE_K_VEC_SIZE`. ss << __str_5; // 93 | const ADVANCE_DIM = 64 / TILE_K_VEC_SIZE; -ss << __str_365; +ss << __str_417; // 94 | ss << __str_5; // 95 | var src_tile : array, TILE_K_VEC_SIZE>; -ss << __str_366; +ss << __str_418; // 96 | var weight_tile : array, TILE_K_VEC_SIZE>; -ss << __str_367; +ss << __str_419; // 97 | ss << __str_5; // 98 | $MAIN { MainFunctionStart(); ss << __str_5; // 99 | let batch = workgroup_idx / (uniforms.M_tiles * uniforms.N_tiles); -ss << __str_368; +ss << __str_420; // 100 | let m_global_base = ((workgroup_idx / uniforms.N_tiles) % uniforms.M_tiles) * TILE_M_SIZE; -ss << __str_369; +ss << __str_421; // 101 | let n_global_base = (workgroup_idx % uniforms.N_tiles) * TILE_N_SIZE; -ss << __str_370; +ss << __str_422; // 102 | ss << __str_5; // 103 | var results : array; -ss << __str_371; +ss << __str_423; // 104 | for (var k_idx = 0u; k_idx < uniforms.K_tiles; k_idx++) { -ss << __str_372; +ss << __str_424; // 105 | for (var src_m = 0u; src_m < TILE_M_SIZE; src_m += ADVANCE_DIM) { -ss << __str_373; +ss << __str_425; // 106 | // Loads a 64 vec of src into the workgroup memory. ss << __str_5; // 107 | let load_src_m = src_m + local_idx / TILE_K_VEC_SIZE; -ss << __str_374; +ss << __str_426; // 108 | let load_src_k = local_idx % TILE_K_VEC_SIZE; -ss << __str_375; +ss << __str_427; // 109 | ss << __str_5; // 110 | src_tile[load_src_k][load_src_m] = load_src(batch, -ss << __str_376; +ss << __str_428; // 111 | m_global_base + load_src_m, -ss << __str_377; +ss << __str_429; // 112 | k_idx * TILE_K_VEC_SIZE + load_src_k); -ss << __str_378; +ss << __str_430; // 113 | } ss << __str_20; // 114 | ss << __str_5; // 115 | for (var weight_n = 0u; weight_n < TILE_N_SIZE; weight_n += ADVANCE_DIM) { -ss << __str_379; +ss << __str_431; // 116 | // Loads a 64 vec of weight into the workgroup memory. ss << __str_5; // 117 | let load_weight_n = weight_n + local_idx / TILE_K_VEC_SIZE; -ss << __str_380; +ss << __str_432; // 118 | let load_weight_k = local_idx % TILE_K_VEC_SIZE; -ss << __str_381; +ss << __str_433; // 119 | ss << __str_5; // 120 | weight_tile[load_weight_k][load_weight_n] = load_weight(n_global_base + load_weight_n, -ss << __str_382; +ss << __str_434; // 121 | k_idx * TILE_K_VEC_SIZE + load_weight_k); -ss << __str_383; +ss << __str_435; // 122 | } ss << __str_20; // 123 | workgroupBarrier(); @@ -266,44 +266,44 @@ ss << __str_26; // 124 | ss << __str_5; // 125 | for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) { -ss << __str_384; +ss << __str_436; // 126 | let weight_data = weight_tile[inner_k_idx][local_idx]; -ss << __str_385; +ss << __str_437; // 127 | #if use_subgroup if (__param_use_subgroup) { // 128 | // The subgroup size is fixed at 32; wrap the index so the 16x64 tiling does not read out of bounds. // 129 | let src_data = src_tile[inner_k_idx][sg_id % TILE_M_SIZE]; -ss << __str_386; +ss << __str_438; // 130 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { -ss << __str_387; +ss << __str_439; // 131 | #if vec_size == 1 if (__param_vec_size == 1) { // 132 | results[m_idx] += output_element_t(weight_data * subgroupShuffle(src_data, m_idx)); -ss << __str_388; +ss << __str_440; // 133 | #else } else { // 134 | results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx))); -ss << __str_389; +ss << __str_441; // 135 | #endif } // 136 | } -ss << __str_292; +ss << __str_344; // 137 | #else } else { // 138 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { -ss << __str_387; +ss << __str_439; // 139 | #if vec_size == 1 if (__param_vec_size == 1) { // 140 | results[m_idx] += output_element_t(weight_data * src_tile[inner_k_idx][m_idx]); -ss << __str_390; +ss << __str_442; // 141 | #else } else { // 142 | results[m_idx] += output_element_t(dot(weight_data, src_tile[inner_k_idx][m_idx])); -ss << __str_391; +ss << __str_443; // 143 | #endif } // 144 | } -ss << __str_292; +ss << __str_344; // 145 | #endif } // 146 | } @@ -315,57 +315,57 @@ ss << __str_11; // 149 | ss << __str_5; // 150 | let m_base = m_global_base; -ss << __str_392; +ss << __str_444; // 151 | let n_base = n_global_base + local_idx; -ss << __str_393; +ss << __str_445; // 152 | ss << __str_5; // 153 | let bias = load_bias(n_base); -ss << __str_394; +ss << __str_446; // 154 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { -ss << __str_395; +ss << __str_447; // 155 | var output_data = results[m_idx] + bias; -ss << __str_396; +ss << __str_448; // 156 | #if activation_kind == 1 if (__param_activation_kind == 1) { // 157 | output_data = max(output_data, output_element_t(0)); -ss << __str_397; +ss << __str_449; // 158 | #elif activation_kind == 2 } else if (__param_activation_kind == 2) { // 159 | output_data = output_element_t(1) / (output_element_t(1) + exp(-output_data)); -ss << __str_398; +ss << __str_450; // 160 | #elif activation_kind == 3 } else if (__param_activation_kind == 3) { // 161 | output_data = clamp(output_data, output_element_t(uniforms.activation_param_0), output_element_t(uniforms.activation_param_1)); -ss << __str_399; +ss << __str_451; // 162 | #elif activation_kind == 4 } else if (__param_activation_kind == 4) { // 163 | output_data = clamp(output_element_t(uniforms.activation_param_0) * output_data + output_element_t(uniforms.activation_param_1), output_element_t(0), output_element_t(1)); -ss << __str_400; +ss << __str_452; // 164 | #elif activation_kind == 5 } else if (__param_activation_kind == 5) { // 165 | output_data = select(output_element_t(uniforms.activation_param_0) * output_data, output_data, output_data >= output_element_t(0)); -ss << __str_401; +ss << __str_453; // 166 | #elif activation_kind == 6 } else if (__param_activation_kind == 6) { // 167 | output_data = tanh(output_data); -ss << __str_402; +ss << __str_454; // 168 | #elif activation_kind == 7 } else if (__param_activation_kind == 7) { // 169 | #if quick_gelu_unit_alpha if (__param_quick_gelu_unit_alpha) { // 170 | output_data = output_data * (output_element_t(1) / (output_element_t(1) + exp(-output_data))); -ss << __str_403; +ss << __str_455; // 171 | #else } else { // 172 | output_data = output_data * (output_element_t(1) / (output_element_t(1) + exp(-(output_element_t(uniforms.activation_param_0) * output_data)))); -ss << __str_404; +ss << __str_456; // 173 | #endif } // 174 | #endif } // 175 | write_output(batch, m_base + m_idx, n_base, output_data); -ss << __str_405; +ss << __str_457; // 176 | } ss << __str_11; // 177 | } // MAIN 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..c7b9d3a1bb76c 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,62 +19,62 @@ 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_459; // 7 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_408; +ss << __str_460; // 8 | let offset = co * uniforms.I * uniforms.H * uniforms.W + -ss << __str_409; +ss << __str_461; // 9 | ci * uniforms.H * uniforms.W + -ss << __str_410; +ss << __str_462; // 10 | h_w; -ss << __str_411; +ss << __str_463; // 11 | return src.getByOffset(offset); -ss << __str_351; -ss << __var_src->GetByOffset(__str_406); +ss << __str_403; +ss << __var_src->GetByOffset(__str_458); ss << __str_29; // 12 | } ss << __str_11; // 13 | return src_element_t(); -ss << __str_412; +ss << __str_464; // 14 | } -ss << __str_345; +ss << __str_397; // 15 | ss << __str_5; // 16 | fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) { -ss << __str_413; +ss << __str_465; // 17 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_408; +ss << __str_460; // 18 | let offset = co * uniforms.H * uniforms.W * uniforms.I + -ss << __str_414; +ss << __str_466; // 19 | h_w * uniforms.I + -ss << __str_415; +ss << __str_467; // 20 | ci; -ss << __str_416; +ss << __str_468; // 21 | output.setByOffset(offset, value); ss << __str_28; -ss << __var_output->SetByOffset(__str_406, __str_322); +ss << __var_output->SetByOffset(__str_458, __str_374); ss << __str_29; // 22 | } ss << __str_11; // 23 | } -ss << __str_345; +ss << __str_397; // 24 | ss << __str_5; // 25 | var data_cache : array, 4>; -ss << __str_417; +ss << __str_469; // 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_470; // 29 | let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64; -ss << __str_419; +ss << __str_471; // 30 | ss << __str_5; // 31 | if (group_co >= uniforms.O || group_ci >= uniforms.I) { -ss << __str_420; +ss << __str_472; // 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_473; // 36 | // load ss << __str_5; // 37 | for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) { -ss << __str_422; +ss << __str_474; // 38 | let load_ci_idx = ci_idx + local_idx / 4; -ss << __str_423; +ss << __str_475; // 39 | let load_h_w_idx = local_idx % 4; -ss << __str_424; +ss << __str_476; // 40 | ss << __str_5; // 41 | data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co, -ss << __str_425; +ss << __str_477; // 42 | group_ci + load_ci_idx, -ss << __str_426; +ss << __str_478; // 43 | h_w_idx * 4 + load_h_w_idx); -ss << __str_427; +ss << __str_479; // 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_480; // 49 | let output_data = data_cache[local_h_w_idx][local_idx]; -ss << __str_429; +ss << __str_481; // 50 | write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data); -ss << __str_430; +ss << __str_482; // 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..7643e0a5e3068 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 @@ -34,112 +34,112 @@ Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, Te MainFunctionStart(); ss << __str_5; // 15 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); -ss << __str_241; -ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_237); +ss << __str_293; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_289); ss << __str_29; // 16 | ss << __str_5; // 17 | let constant_value = -ss << __str_436; +ss << __str_488; // 18 | #if is_float16 if (__param_is_float16) { // 19 | bitcast>(uniforms.constant_value)[0]; -ss << __str_437; +ss << __str_489; // 20 | #else } else { // 21 | bitcast(uniforms.constant_value); -ss << __str_438; +ss << __str_490; // 22 | #endif } // 23 | // 24 | #if dim_value_zero if (__param_dim_value_zero) { // 25 | output[global_idx] = constant_value; -ss << __str_439; +ss << __str_491; // 26 | #else } else { // 27 | let output_indices = output.offsetToIndices(global_idx); -ss << __str_265; -ss << __var_output->OffsetToIndices(__str_239); +ss << __str_317; +ss << __var_output->OffsetToIndices(__str_291); ss << __str_29; // 28 | var input_index = u32(0); -ss << __str_440; +ss << __str_492; // 29 | var use_pad_value = false; -ss << __str_441; +ss << __str_493; // 30 | var in_coord = i32(0); -ss << __str_442; +ss << __str_494; // 31 | ss << __str_5; // 32 | for (var dim = 0; dim < output.rank && !use_pad_value; dim++) { -ss << __str_443; +ss << __str_495; ss << __var_output->Rank(); -ss << __str_444; +ss << __str_496; // 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_497; +ss << GetElementAt(__str_299, __str_483, __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_498; +ss << GetElementAt(__str_484, __str_483, __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_499; +ss << GetElementAt(__str_485, __str_483, __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_500; // 38 | use_pad_value = true; -ss << __str_449; +ss << __str_501; // 39 | #elif pad_mode == PAD_MODE_EDGE } else if (__param_pad_mode == 2) { // 40 | if (output_index < lower_pads) { -ss << __str_450; +ss << __str_502; // 41 | in_coord = 0; -ss << __str_451; +ss << __str_503; // 42 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_452; +ss << __str_504; // 43 | in_coord = data_shape - 1; -ss << __str_453; +ss << __str_505; // 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_500; // 46 | in_coord = output_index - lower_pads; -ss << __str_454; +ss << __str_506; // 47 | if (in_coord < 0) { -ss << __str_455; +ss << __str_507; // 48 | in_coord = -in_coord; -ss << __str_456; +ss << __str_508; // 49 | } -ss << __str_292; +ss << __str_344; // 50 | let _2n_1 = 2 * (data_shape - 1); -ss << __str_457; +ss << __str_509; // 51 | in_coord = in_coord % _2n_1; -ss << __str_458; +ss << __str_510; // 52 | if (in_coord >= data_shape) { -ss << __str_459; +ss << __str_511; // 53 | in_coord = _2n_1 - in_coord; -ss << __str_460; +ss << __str_512; // 54 | } -ss << __str_292; +ss << __str_344; // 55 | #else // PAD_MODE_WRAP } else { // 56 | if (output_index < lower_pads) { -ss << __str_450; +ss << __str_502; // 57 | in_coord = data_shape + output_index - lower_pads; -ss << __str_461; +ss << __str_513; // 58 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_452; +ss << __str_504; // 59 | in_coord = output_index - data_shape - lower_pads; -ss << __str_462; +ss << __str_514; // 60 | #endif // pad_mode } // 61 | } else { -ss << __str_463; +ss << __str_515; // 62 | in_coord = output_index - lower_pads; -ss << __str_464; +ss << __str_516; // 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_517; // 67 | #endif } // 68 | // 69 | input_index += select(u32(in_coord) -ss << __str_466; +ss << __str_518; // 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_519; +ss << GetElementAt(__str_486, __str_483, __var_output->Rank() - 1); ss << __str_5; // 72 | #endif } // 73 | , u32(in_coord), dim == output.rank - 1); -ss << __str_468; +ss << __str_520; ss << __var_output->Rank(); -ss << __str_469; +ss << __str_521; // 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 << __str_293; +ss << __var_output->SetByOffset(__str_291, __str_487); 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..2fae15c6fba99 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 @@ -50,12 +50,15 @@ Status ApplyTemplate<"math/subgroup_matrix_gemm_8x16x16.wgsl.template">(ShaderHe template <> struct TemplateParameter<"math/subgroup_matrix_matmul.wgsl.template"> { using type = struct { + int param_sg_grid_m; + int param_sg_grid_n; 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_tile_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 3bb8396be1a50..0ee59ca627e64 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h @@ -31,18 +31,18 @@ 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" // 8680e1f26bdab33a1f955af885d6c2d143b2ce4384e4c737afbede50d0920ddd // 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/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/math/subgroup_matrix_matmul.h" // 7b145d4725c7154c524392f47106937d880bb0b1ccc4233955c16abf2f1707f0 +#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // 25691a4b28f76e8e3ee12762f3caa70f751ca5e6e6b3a8b3ca679e43afdc4973 +#include "wgsl_template_gen/generated/nn/conv_transpose3d.h" // b9c6705bba231fb713c02c7f05773a86b2747de987f8177b1ae557e4c0da3e60 +#include "wgsl_template_gen/generated/nn/im2col_matmul.h" // f34ed5de5a943093faf863fceb88874399d8228270896c9f058bcb31fab30458 +#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // 94f5309725fa042aa8905bb314e9fe026741ec70213905388ea5874dcbf28801 +#include "wgsl_template_gen/generated/tensor/pad.h" // 2e84305ac4b05367377b043927bb7b2ce95468558e5ae82ffc1bd1f906c69cb8 #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..c258723b1fca0 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 @@ -224,255 +224,307 @@ constexpr const char* __str_214 = " let global_n = global_base_n constexpr const char* __str_215 = " let c_offset = global_m * uniforms.c_stride_m + global_n * uniforms.c_stride_n;\n"; constexpr const char* __str_216 = " val += output_element_t(uniforms.beta) * input_c[c_offset];\n"; constexpr const char* __str_217 = " "; -constexpr const char* __str_218 = " let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM;\n"; -constexpr const char* __str_219 = " let tiles_per_slice = uniforms.num_n_tile * num_m_tile;\n"; -constexpr const char* __str_220 = " let batch_id = workgroup_idx / tiles_per_slice;\n"; -constexpr const char* __str_221 = " let slice_idx = workgroup_idx % tiles_per_slice;\n"; -constexpr const char* __str_222 = " let n_tile = slice_idx % uniforms.num_n_tile;\n"; -constexpr const char* __str_223 = " let m_tile = slice_idx / uniforms.num_n_tile;\n"; -constexpr const char* __str_224 = " let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN);\n"; -constexpr const char* __str_225 = " let a_batch_offset = batch_id * uniforms.M * uniforms.K;\n"; -constexpr const char* __str_226 = " let b_batch_offset = batch_id * uniforms.K * uniforms.N_b;\n"; -constexpr const char* __str_227 = " let out_batch_offset = batch_id * uniforms.M * uniforms.N;\n"; -constexpr const char* __str_228 = " let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n;\n"; -constexpr const char* __str_229 = " &input_b, b_base + 0 * kSgMatN, uniforms.N_b);\n"; -constexpr const char* __str_230 = " &input_b, b_base + 1 * kSgMatN, uniforms.N_b);\n"; -constexpr const char* __str_231 = " &input_b, b_base + 2 * kSgMatN, uniforms.N_b);\n"; -constexpr const char* __str_232 = " &input_b, b_base + 3 * kSgMatN, uniforms.N_b);\n"; -constexpr const char* __str_233 = " let a_col = a_batch_offset + kb * kSgMatK;\n"; -constexpr const char* __str_234 = " let out_base = out_batch_offset + global_m * uniforms.N + global_base_n;\n"; -constexpr const char* __str_235 = " write_output(out_base + i, global_base_n + i,\n"; -constexpr const char* __str_236 = " output_value_t(scratch[scratch_base + i]));\n"; -constexpr const char* __str_237 = "uniforms.output_size"; -constexpr const char* __str_238 = "r * uniforms.N + c"; -constexpr const char* __str_239 = "global_idx"; -constexpr const char* __str_240 = "v"; -constexpr const char* __str_241 = " "; -constexpr const char* __str_242 = " let r = global_idx / uniforms.N_b;\n"; -constexpr const char* __str_243 = " let c = global_idx % uniforms.N_b;\n"; -constexpr const char* __str_244 = " var v = output_value_t(0);\n"; -constexpr const char* __str_245 = " if (c < uniforms.N) {\n"; -constexpr const char* __str_246 = " v = output_value_t("; -constexpr const char* __str_247 = "output_indices"; -constexpr const char* __str_248 = "0"; -constexpr const char* __str_249 = "4"; -constexpr const char* __str_250 = "1"; -constexpr const char* __str_251 = "uniforms.w_shape"; -constexpr const char* __str_252 = "3"; -constexpr const char* __str_253 = "2"; -constexpr const char* __str_254 = "spatial_dim"; -constexpr const char* __str_255 = "spatial_dim + 1u"; -constexpr const char* __str_256 = "spatial_dim + 2u"; -constexpr const char* __str_257 = "uniforms.x_shape"; -constexpr const char* __str_258 = "x_indices"; -constexpr const char* __str_259 = "w_indices"; -constexpr const char* __str_260 = "output_channel"; -constexpr const char* __str_261 = "output_value_t(result)"; -constexpr const char* __str_262 = "alias Accumulator = vec4;\n"; -constexpr const char* __str_263 = "alias Accumulator = vec2;\n"; -constexpr const char* __str_264 = "alias Accumulator = f32;\n"; -constexpr const char* __str_265 = " let output_indices = "; -constexpr const char* __str_266 = " let batch = "; -constexpr const char* __str_267 = " let output_channel = "; -constexpr const char* __str_268 = " let spatial_dim = 1u;\n"; -constexpr const char* __str_269 = " let spatial_dim = 2u;\n"; -constexpr const char* __str_270 = " let output_channels_per_group = "; -constexpr const char* __str_271 = " let kernel_shape = vec3("; -constexpr const char* __str_272 = ",\n"; -constexpr const char* __str_273 = " "; -constexpr const char* __str_274 = " let group = output_channel / output_channels_per_group;\n"; -constexpr const char* __str_275 = " let w_channel = output_channel % output_channels_per_group;\n"; -constexpr const char* __str_276 = " let corner = vec3(\n"; -constexpr const char* __str_277 = " i32("; -constexpr const char* __str_278 = "),\n"; -constexpr const char* __str_279 = ")) + vec3(uniforms.pads);\n"; -constexpr const char* __str_280 = " var value = Accumulator(0);\n"; -constexpr const char* __str_281 = " for (var kd = 0u; kd < kernel_shape.x; kd++) {\n"; -constexpr const char* __str_282 = " let pos_d = corner.x - i32(kd) * i32(uniforms.dilations.x);\n"; -constexpr const char* __str_283 = " if (pos_d < 0 || pos_d % i32(uniforms.strides.x) != 0) {\n"; -constexpr const char* __str_284 = " continue;\n"; -constexpr const char* __str_285 = " let id = u32(pos_d) / uniforms.strides.x;\n"; -constexpr const char* __str_286 = " if (id >= "; -constexpr const char* __str_287 = ") {\n"; -constexpr const char* __str_288 = " for (var kh = 0u; kh < kernel_shape.y; kh++) {\n"; -constexpr const char* __str_289 = " let pos_h = corner.y - i32(kh) * i32(uniforms.dilations.y);\n"; -constexpr const char* __str_290 = " if (pos_h < 0 || pos_h % i32(uniforms.strides.y) != 0) {\n"; -constexpr const char* __str_291 = " continue;\n"; -constexpr const char* __str_292 = " }\n"; -constexpr const char* __str_293 = " let ih = u32(pos_h) / uniforms.strides.y;\n"; -constexpr const char* __str_294 = " if (ih >= "; -constexpr const char* __str_295 = " for (var kw = 0u; kw < kernel_shape.z; kw++) {\n"; -constexpr const char* __str_296 = " let pos_w = corner.z - i32(kw) * i32(uniforms.dilations.z);\n"; -constexpr const char* __str_297 = " if (pos_w < 0 || pos_w % i32(uniforms.strides.z) != 0) {\n"; -constexpr const char* __str_298 = " continue;\n"; -constexpr const char* __str_299 = " let iw = u32(pos_w) / uniforms.strides.z;\n"; -constexpr const char* __str_300 = " if (iw >= "; -constexpr const char* __str_301 = " for (var c = 0u; c < uniforms.input_channels_per_group; c += "; -constexpr const char* __str_302 = " let input_channel = group * uniforms.input_channels_per_group + c;\n"; -constexpr const char* __str_303 = " let x_indices = x_indices_t(batch, id, ih, iw, input_channel / "; -constexpr const char* __str_304 = " let x_value = Accumulator("; -constexpr const char* __str_305 = " var x_value = Accumulator(0);\n"; -constexpr const char* __str_306 = " for (var lane = 0u; lane < "; -constexpr const char* __str_307 = "; lane++) {\n"; -constexpr const char* __str_308 = " let x_indices = x_indices_t(batch, input_channel + lane, id, ih, iw);\n"; -constexpr const char* __str_309 = " x_value[lane] = f32("; -constexpr const char* __str_310 = " }\n"; -constexpr const char* __str_311 = " let x_indices = x_indices_t(batch, input_channel, id, ih, iw);\n"; -constexpr const char* __str_312 = " let x_value = f32("; -constexpr const char* __str_313 = " let w_indices = w_indices_t(kd, kh, kw, w_channel, input_channel / "; -constexpr const char* __str_314 = " let w_indices = w_indices_t(input_channel, w_channel, kd, kh, kw);\n"; -constexpr const char* __str_315 = " value += x_value * Accumulator("; -constexpr const char* __str_316 = " var result = dot(value, Accumulator(1));\n"; -constexpr const char* __str_317 = " var result = value;\n"; -constexpr const char* __str_318 = " result += f32("; -constexpr const char* __str_319 = "src_idx"; -constexpr const char* __str_320 = "weight_idx"; -constexpr const char* __str_321 = "output_idx"; -constexpr const char* __str_322 = "value"; -constexpr const char* __str_323 = "fn load_src(batch : u32, m : u32, k_packed_idx : u32) -> src_value_t {\n"; -constexpr const char* __str_324 = " if (batch >= uniforms.batch || m >= uniforms.im2col_m || k_packed_idx * "; -constexpr const char* __str_325 = " >= uniforms.im2col_k) {\n"; -constexpr const char* __str_326 = " return src_value_t();\n"; -constexpr const char* __str_327 = " let channel_i_vec = uniforms.channel_i / "; -constexpr const char* __str_328 = " let h_idx = m / uniforms.output_w;\n"; -constexpr const char* __str_329 = " let w_idx = m % uniforms.output_w;\n"; -constexpr const char* __str_330 = " let c_i_vec_idx = k_packed_idx % channel_i_vec;\n"; -constexpr const char* __str_331 = " let k_h_w_idx = k_packed_idx / channel_i_vec;\n"; -constexpr const char* __str_332 = " let k_h = k_h_w_idx / uniforms.kernel_w;\n"; -constexpr const char* __str_333 = " let k_w = k_h_w_idx % uniforms.kernel_w;\n"; -constexpr const char* __str_334 = " let src_h_coord_padded = h_idx * uniforms.strides.x + k_h * uniforms.dilations.x;\n"; -constexpr const char* __str_335 = " let src_w_coord_padded = w_idx * uniforms.strides.y + k_w * uniforms.dilations.y;\n"; -constexpr const char* __str_336 = " let src_h_coord : i32 = i32(src_h_coord_padded) - i32(uniforms.pads.x);\n"; -constexpr const char* __str_337 = " let src_w_coord : i32 = i32(src_w_coord_padded) - i32(uniforms.pads.y);\n"; -constexpr const char* __str_338 = " if (src_h_coord < 0 || src_h_coord >= i32(uniforms.src_h) ||\n"; -constexpr const char* __str_339 = " src_w_coord < 0 || src_w_coord >= i32(uniforms.src_w)) {\n"; -constexpr const char* __str_340 = " let src_idx = batch * uniforms.src_h * uniforms.src_w * channel_i_vec +\n"; -constexpr const char* __str_341 = " u32(src_h_coord) * uniforms.src_w * channel_i_vec +\n"; -constexpr const char* __str_342 = " u32(src_w_coord) * channel_i_vec +\n"; -constexpr const char* __str_343 = " c_i_vec_idx;\n"; -constexpr const char* __str_344 = " return "; -constexpr const char* __str_345 = "}\n"; -constexpr const char* __str_346 = "fn load_weight(n : u32, k_packed_idx : u32) -> weight_value_t {\n"; -constexpr const char* __str_347 = " if (n < uniforms.im2col_n && k_packed_idx < uniforms.im2col_k / "; -constexpr const char* __str_348 = " let weight_idx = n * uniforms.im2col_k / "; -constexpr const char* __str_349 = " +\n"; -constexpr const char* __str_350 = " k_packed_idx;\n"; -constexpr const char* __str_351 = " return "; -constexpr const char* __str_352 = " return weight_value_t();\n"; -constexpr const char* __str_353 = "fn load_bias(n : u32) -> output_element_t {\n"; -constexpr const char* __str_354 = " if (n < uniforms.im2col_n) {\n"; -constexpr const char* __str_355 = " return output_element_t(bias[n]);\n"; -constexpr const char* __str_356 = " return output_element_t();\n"; -constexpr const char* __str_357 = "fn write_output(batch : u32, m : u32, n : u32, value : output_element_t) {\n"; -constexpr const char* __str_358 = " if (batch < uniforms.batch && m < uniforms.im2col_m && n < uniforms.im2col_n) {\n"; -constexpr const char* __str_359 = " let output_idx = batch * uniforms.im2col_m * uniforms.im2col_n +\n"; -constexpr const char* __str_360 = " m * uniforms.im2col_n +\n"; -constexpr const char* __str_361 = " n;\n"; -constexpr const char* __str_362 = "const TILE_M_SIZE : u32 = "; -constexpr const char* __str_363 = "const TILE_N_SIZE : u32 = "; -constexpr const char* __str_364 = "const TILE_K_VEC_SIZE : u32 = 16 / "; -constexpr const char* __str_365 = "const ADVANCE_DIM = 64 / TILE_K_VEC_SIZE;\n"; -constexpr const char* __str_366 = "var src_tile : array, TILE_K_VEC_SIZE>;\n"; -constexpr const char* __str_367 = "var weight_tile : array, TILE_K_VEC_SIZE>;\n"; -constexpr const char* __str_368 = " let batch = workgroup_idx / (uniforms.M_tiles * uniforms.N_tiles);\n"; -constexpr const char* __str_369 = " let m_global_base = ((workgroup_idx / uniforms.N_tiles) % uniforms.M_tiles) * TILE_M_SIZE;\n"; -constexpr const char* __str_370 = " let n_global_base = (workgroup_idx % uniforms.N_tiles) * TILE_N_SIZE;\n"; -constexpr const char* __str_371 = " var results : array;\n"; -constexpr const char* __str_372 = " for (var k_idx = 0u; k_idx < uniforms.K_tiles; k_idx++) {\n"; -constexpr const char* __str_373 = " for (var src_m = 0u; src_m < TILE_M_SIZE; src_m += ADVANCE_DIM) {\n"; -constexpr const char* __str_374 = " let load_src_m = src_m + local_idx / TILE_K_VEC_SIZE;\n"; -constexpr const char* __str_375 = " let load_src_k = local_idx % TILE_K_VEC_SIZE;\n"; -constexpr const char* __str_376 = " src_tile[load_src_k][load_src_m] = load_src(batch,\n"; -constexpr const char* __str_377 = " m_global_base + load_src_m,\n"; -constexpr const char* __str_378 = " k_idx * TILE_K_VEC_SIZE + load_src_k);\n"; -constexpr const char* __str_379 = " for (var weight_n = 0u; weight_n < TILE_N_SIZE; weight_n += ADVANCE_DIM) {\n"; -constexpr const char* __str_380 = " let load_weight_n = weight_n + local_idx / TILE_K_VEC_SIZE;\n"; -constexpr const char* __str_381 = " let load_weight_k = local_idx % TILE_K_VEC_SIZE;\n"; -constexpr const char* __str_382 = " weight_tile[load_weight_k][load_weight_n] = load_weight(n_global_base + load_weight_n,\n"; -constexpr const char* __str_383 = " k_idx * TILE_K_VEC_SIZE + load_weight_k);\n"; -constexpr const char* __str_384 = " for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) {\n"; -constexpr const char* __str_385 = " let weight_data = weight_tile[inner_k_idx][local_idx];\n"; -constexpr const char* __str_386 = " let src_data = src_tile[inner_k_idx][sg_id % TILE_M_SIZE];\n"; -constexpr const char* __str_387 = " for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {\n"; -constexpr const char* __str_388 = " results[m_idx] += output_element_t(weight_data * subgroupShuffle(src_data, m_idx));\n"; -constexpr const char* __str_389 = " results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx)));\n"; -constexpr const char* __str_390 = " results[m_idx] += output_element_t(weight_data * src_tile[inner_k_idx][m_idx]);\n"; -constexpr const char* __str_391 = " results[m_idx] += output_element_t(dot(weight_data, src_tile[inner_k_idx][m_idx]));\n"; -constexpr const char* __str_392 = " let m_base = m_global_base;\n"; -constexpr const char* __str_393 = " let n_base = n_global_base + local_idx;\n"; -constexpr const char* __str_394 = " let bias = load_bias(n_base);\n"; -constexpr const char* __str_395 = " for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) {\n"; -constexpr const char* __str_396 = " var output_data = results[m_idx] + bias;\n"; -constexpr const char* __str_397 = " output_data = max(output_data, output_element_t(0));\n"; -constexpr const char* __str_398 = " output_data = output_element_t(1) / (output_element_t(1) + exp(-output_data));\n"; -constexpr const char* __str_399 = " output_data = clamp(output_data, output_element_t(uniforms.activation_param_0), output_element_t(uniforms.activation_param_1));\n"; -constexpr const char* __str_400 = " output_data = clamp(output_element_t(uniforms.activation_param_0) * output_data + output_element_t(uniforms.activation_param_1), output_element_t(0), output_element_t(1));\n"; -constexpr const char* __str_401 = " output_data = select(output_element_t(uniforms.activation_param_0) * output_data, output_data, output_data >= output_element_t(0));\n"; -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_218 = "const kSgGridM: u32 = u32("; +constexpr const char* __str_219 = "const kSgGridN: u32 = u32("; +constexpr const char* __str_220 = "const kTileM: u32 = kSgMatM * kSgMatCountM * kSgGridM;\n"; +constexpr const char* __str_221 = "const kTileN: u32 = kSgMatN * kSgMatCountN * kSgGridN;\n"; +constexpr const char* __str_222 = "const kTileK: u32 = u32("; +constexpr const char* __str_223 = "const kWorkgroupSize: u32 = kSgGridM * kSgGridN * kSubgroupSize;\n"; +constexpr const char* __str_224 = "const kStrideA: u32 = kTileK + 8u;\n"; +constexpr const char* __str_225 = "const kStrideB: u32 = kTileN + 8u;\n"; +constexpr const char* __str_226 = "var scratch: array;\n"; +constexpr const char* __str_227 = " let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM;\n"; +constexpr const char* __str_228 = " let tiles_per_slice = uniforms.num_n_tile * num_m_tile;\n"; +constexpr const char* __str_229 = " let batch_id = workgroup_idx / tiles_per_slice;\n"; +constexpr const char* __str_230 = " let slice_idx = workgroup_idx % tiles_per_slice;\n"; +constexpr const char* __str_231 = " let n_tile = slice_idx % uniforms.num_n_tile;\n"; +constexpr const char* __str_232 = " let m_tile = slice_idx / uniforms.num_n_tile;\n"; +constexpr const char* __str_233 = " let a_batch_offset = batch_id * uniforms.M * uniforms.K;\n"; +constexpr const char* __str_234 = " let b_batch_offset = batch_id * uniforms.K * uniforms.N;\n"; +constexpr const char* __str_235 = " let out_batch_offset = batch_id * uniforms.M * uniforms.N;\n"; +constexpr const char* __str_236 = " let m_base = (sg_index % kSgGridM) * kSgMatM * kSgMatCountM;\n"; +constexpr const char* __str_237 = " let n_base = (sg_index / kSgGridM) * kSgMatN * kSgMatCountN;\n"; +constexpr const char* __str_238 = " let sg_scratch_base = m_base * kTileN + n_base;\n"; +constexpr const char* __str_239 = " let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN);\n"; +constexpr const char* __str_240 = " let b_batch_offset = batch_id * uniforms.K * uniforms.N_b;\n"; +constexpr const char* __str_241 = " for (var k0: u32 = 0u; k0 < uniforms.K; k0 = k0 + kTileK) {\n"; +constexpr const char* __str_242 = " let k_start = (k0 / kSgMatK) * kSgMatK;\n"; +constexpr const char* __str_243 = " workgroupBarrier();\n"; +constexpr const char* __str_244 = " for (var i: u32 = local_idx; i < kTileM * kTileK / 4u; i = i + kWorkgroupSize) {\n"; +constexpr const char* __str_245 = " let r = i / (kTileK / 4u);\n"; +constexpr const char* __str_246 = " let c = (i % (kTileK / 4u)) * 4u;\n"; +constexpr const char* __str_247 = " var v = vec4();\n"; +constexpr const char* __str_248 = " if (m_tile_start + r < uniforms.M && k_start + c < uniforms.K) {\n"; +constexpr const char* __str_249 = " v = input_a[(a_batch_offset + (m_tile_start + r) * uniforms.K + k_start + c) / 4u];\n"; +constexpr const char* __str_250 = " let s = r * kStrideA + c;\n"; +constexpr const char* __str_251 = " scratch[s] = v.x;\n"; +constexpr const char* __str_252 = " scratch[s + 1u] = v.y;\n"; +constexpr const char* __str_253 = " scratch[s + 2u] = v.z;\n"; +constexpr const char* __str_254 = " scratch[s + 3u] = v.w;\n"; +constexpr const char* __str_255 = " for (var i: u32 = local_idx; i < kTileK * kTileN / 4u; i = i + kWorkgroupSize) {\n"; +constexpr const char* __str_256 = " let r = i / (kTileN / 4u);\n"; +constexpr const char* __str_257 = " let c = (i % (kTileN / 4u)) * 4u;\n"; +constexpr const char* __str_258 = " if (k_start + r < uniforms.K && n_tile_start + c < uniforms.N) {\n"; +constexpr const char* __str_259 = " v = input_b[(b_batch_offset + (k_start + r) * uniforms.N + n_tile_start + c) / 4u];\n"; +constexpr const char* __str_260 = " let s = kTileM * kStrideA + r * kStrideB + c;\n"; +constexpr const char* __str_261 = " for (var kk: u32 = 0u; kk < kTileK; kk = kk + kSgMatK) {\n"; +constexpr const char* __str_262 = " let b_base = kTileM * kStrideA + kk * kStrideB + n_base;\n"; +constexpr const char* __str_263 = " var sg_mat_b0: subgroup_matrix_right