From 09103666fa1c41833ecd5735100bb94bc5924db6 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Oct 2026 10:18:50 +0800 Subject: [PATCH 1/6] Refactor WebGPU MatMulNBits input preparation for reuse --- .../webgpu/quantization/dp4a_matmul_nbits.cc | 38 ++-- .../webgpu/quantization/dp4a_matmul_nbits.h | 10 +- .../webgpu/quantization/matmul_nbits.cc | 4 +- .../webgpu/quantization/matmul_nbits.h | 5 +- .../subgroup_matrix_matmul_nbits.cc | 149 +++++++-------- .../subgroup_matrix_matmul_nbits.h | 9 +- .../matmul_nbits_input_preparation_test.cc | 170 ++++++++++++++++++ 7 files changed, 293 insertions(+), 92 deletions(-) create mode 100644 onnxruntime/test/providers/webgpu/matmul_nbits_input_preparation_test.cc diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index d9c5f5839a3bf..37660060a87b4 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -106,25 +106,33 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index, - const Tensor* weight_index_indirect) { + const Tensor* weight_index_indirect, DP4AQuantizedInput* quantized_input) { constexpr uint32_t kVec4Components = 4; constexpr uint32_t kVec2Components = 2; constexpr uint32_t kU32Components = 4; - constexpr uint32_t kBlockSizeA = 128; - DP4AMatMulQuantizeProgram quantize_program; - quantize_program.SetWorkgroupSize(64); - uint32_t tile_size = 64 * kVec4Components; - quantize_program.SetDispatchGroupSize((batch_count * M * K + tile_size - 1) / tile_size, 1, 1); - TensorShape a_quant_shape{batch_count, M, K / kU32Components}; - Tensor a_quant = context.CreateGPUTensor(DataTypeImpl::GetType(), a_quant_shape); - TensorShapeVector a_scales_dims({batch_count, 1, M, K / kBlockSizeA}); - Tensor a_scale = context.CreateGPUTensor(a->DataType(), a_scales_dims); - quantize_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}}) - .AddOutputs({{&a_quant, ProgramTensorMetadataDependency::Rank, a_quant.Shape(), 1}, - {&a_scale, ProgramTensorMetadataDependency::Rank, 1}}) - .AddUniformVariable({batch_count * M * K / kU32Components}); - ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program)); + DP4AQuantizedInput local_input; + auto& input = quantized_input ? *quantized_input : local_input; + if (!input.values) { + constexpr uint32_t kBlockSizeA = 128; + DP4AMatMulQuantizeProgram quantize_program; + quantize_program.SetWorkgroupSize(64); + uint32_t tile_size = 64 * kVec4Components; + quantize_program.SetDispatchGroupSize((batch_count * M * K + tile_size - 1) / tile_size, 1, 1); + TensorShape a_quant_shape{batch_count, M, K / kU32Components}; + Tensor a_quant = context.CreateGPUTensor(DataTypeImpl::GetType(), a_quant_shape); + TensorShapeVector a_scales_dims({batch_count, 1, M, K / kBlockSizeA}); + Tensor a_scale = context.CreateGPUTensor(a->DataType(), a_scales_dims); + quantize_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}}) + .AddOutputs({{&a_quant, ProgramTensorMetadataDependency::Rank, a_quant.Shape(), 1}, + {&a_scale, ProgramTensorMetadataDependency::Rank, 1}}) + .AddUniformVariable({batch_count * M * K / kU32Components}); + ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program)); + input.values.emplace(std::move(a_quant)); + input.scales.emplace(std::move(a_scale)); + } + const auto& a_quant = *input.values; + const auto& a_scale = *input.scales; const bool has_zero_points = zero_points != nullptr; const bool has_bias = bias != nullptr; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 66dc5819e37a8..2876802986866 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include "core/providers/webgpu/program.h" #include "core/providers/webgpu/webgpu_kernel.h" @@ -14,6 +15,12 @@ namespace webgpu { using namespace onnxruntime::webgpu; +// Invocation-local workspace, reusable only by projections of the same input tensor. +struct DP4AQuantizedInput { + std::optional values; + std::optional scales; +}; + class DP4AMatMulQuantizeProgram final : public Program { public: DP4AMatMulQuantizeProgram() : Program{"DP4AMatMulQuantize"} {} @@ -117,7 +124,8 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index, - const Tensor* weight_index_indirect = nullptr); + const Tensor* weight_index_indirect = nullptr, + DP4AQuantizedInput* quantized_input = nullptr); // The optional M / has_weight_idx_indirect / y arguments fold the original // dispatch-precondition (DP4A is preferred when M is large enough, or diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 404f06f858869..664e8e873ec92 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -293,7 +293,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, Tensor* y, const uint32_t weight_index, const Tensor* weight_index_indirect, - uint32_t override_M) { + uint32_t override_M, DP4AQuantizedInput* quantized_input) { TensorShape b_shape({N_op, K_op}); MatMulComputeHelper helper; ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true)); @@ -353,7 +353,7 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte values). if (CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, M, has_weight_idx_indirect, y)) { - return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect); + return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, quantized_input); } // This program is optimized for FP32/FP16 Block32 prefill. diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index f8c2f91b78396..e2205f37920dd 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -12,6 +12,8 @@ namespace webgpu { using namespace onnxruntime::webgpu; +struct DP4AQuantizedInput; + class MatMulNBitsWideTileProgram final : public Program { public: MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32) @@ -107,7 +109,8 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, int64_t bits_op, onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index = 0, const Tensor* weight_index_indirect = nullptr, - uint32_t override_M = 0); + uint32_t override_M = 0, + DP4AQuantizedInput* quantized_input = nullptr); } // namespace webgpu } // namespace contrib diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc index 8f6b9f6763108..2396dbefd7fa8 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc @@ -175,56 +175,40 @@ class SubgroupMatrixMatMulNBitsTailCopyProgram final : public Program 0 || has_weight_idx_indirect; +namespace { + +struct MatMulNBitsTiling { + uint32_t m; + uint32_t n; + uint32_t k; + uint32_t workgroup_size; +}; - // Determine tile sizes first (needed for prepack padding). - uint32_t tile_m = 32; - uint32_t tile_n = 64; - uint32_t tile_size_k = 32; - uint32_t workgroup_size = 128; +MatMulNBitsTiling SelectMatMulNBitsTiling(const SubgroupMatrixConfig& config, + uint32_t M, uint32_t N, uint32_t K, bool has_bias) { if (config.Is(8, 16, 16)) { - // Default: 64x64 tile size, 256 thread. - tile_m = 64; - tile_n = 64; - workgroup_size = 256; - // Bias output uses a scratch tile, so keep it at 64x64 for workgroup memory limits. if (!has_bias) { if (M >= 2048 && N % 256 == 0) { - // Large M (M>=2048): 128x256 tile size, 512 threads. - tile_m = 128; - tile_n = 256; - workgroup_size = 512; - } else if (M >= 256 && N % 128 == 0) { - // Mid-sized M (256<=M<2048): 128x128 tile size, 512 threads. - tile_m = 128; - tile_n = 128; - workgroup_size = 512; + return {128, 256, 32, 512}; + } + if (M >= 256 && N % 128 == 0) { + return {128, 128, 32, 512}; } } - } else if (config.Is(16, 16, 16)) { + return {64, 64, 32, 256}; + } + if (config.Is(16, 16, 16)) { if (config.subgroupSize == 64) { + uint32_t tile_m; + uint32_t tile_n; + uint32_t tile_k; // A wave64 subgroup computes a 64x64 micro-tile. Small M distributes waves along N; // tall tiles share each dequantized B tile across row waves. if (M <= 64) { tile_m = 64; // Avoid a fully padded second column wave when N is exactly 64. tile_n = N == 64 ? 64 : 128; - tile_size_k = 64; + tile_k = 64; } else { // Use four waves when that does not add M padding over a two-wave tile and M is large // enough to amortize the larger workgroup. @@ -235,53 +219,74 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te // Larger K blocks reduce barriers. Four-wave tiles retain enough occupancy for BK=128; // larger two-wave grids use BK=64, while small two-wave grids favor BK=32. - tile_size_k = M < 256 ? 32 : (use_four_waves || M <= 384 ? 128 : 64); + tile_k = M < 256 ? 32 : (use_four_waves || M <= 384 ? 128 : 64); } // The K loop and cooperative-matrix loads do not mask a partial BK tile. K is guaranteed // to be divisible by 32, so step down until BK divides it exactly. - if (K % tile_size_k != 0) { - tile_size_k = tile_size_k == 128 && K % 64 == 0 ? 64 : 32; + if (K % tile_k != 0) { + tile_k = tile_k == 128 && K % 64 == 0 ? 64 : 32; } - workgroup_size = (tile_m / 64) * (tile_n / 64) * config.subgroupSize; - } else { - // Wave32: four subgroups in a 2x2 grid. - tile_m = 128; - tile_n = 128; - workgroup_size = 4 * config.subgroupSize; + return {tile_m, tile_n, tile_k, (tile_m / 64) * (tile_n / 64) * config.subgroupSize}; } + // Wave32: four subgroups in a 2x2 grid. + return {128, 128, 32, 4 * config.subgroupSize}; + } + return {32, 64, 32, 128}; +} + +} // namespace + +Status PrepackMatMulNBitsInput(const Tensor* a, uint32_t M, uint32_t K, uint32_t padded_M, + const SubgroupMatrixConfig& config, onnxruntime::webgpu::ComputeContext& context, + Tensor& packed) { + PrepackProgram program{config.M, config.K}; + program.SetWorkgroupSize(config.subgroupSize); + if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + program.SetSubgroupSize(config.subgroupSize); } + ORT_ENFORCE(K % config.K == 0 && padded_M % config.M == 0); + program.SetDispatchGroupSize(padded_M / config.M, K / config.K, 1); + packed = context.CreateGPUTensor(a->DataType(), TensorShape{padded_M, K}); + program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddOutputs({{&packed, ProgramTensorMetadataDependency::Rank, packed.Shape(), 1}}) + .AddUniformVariables({{M}, {K}}) + .CacheHint(config.M, config.K); + return context.RunProgram(program); +} + +Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, + const Tensor* zero_points, const Tensor* bias, + uint32_t M, + uint32_t N, + uint32_t K, + uint32_t nbits, + uint32_t zero_blocks_per_col, + const SubgroupMatrixConfig& config, + onnxruntime::webgpu::ComputeContext& context, + Tensor* y, + const uint32_t weight_index, + const Tensor* weight_index_indirect, + const Tensor* prepacked_a) { + const bool has_zero_points = zero_points != nullptr; + const bool has_bias = bias != nullptr; + const bool has_weight_idx_indirect = weight_index_indirect != nullptr; + const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; + + auto [tile_m, tile_n, tile_size_k, workgroup_size] = SelectMatMulNBitsTiling(config, M, N, K, has_bias); ORT_ENFORCE(workgroup_size <= context.DeviceLimits().maxComputeWorkgroupSizeX && workgroup_size <= context.DeviceLimits().maxComputeInvocationsPerWorkgroup, "Subgroup matrix MatMulNBits workgroup size exceeds device limits: ", workgroup_size); // If applicable, layout optimization of input matrix A(MxK) can be used for SubgroupMatrixLoad. Tensor a_prepack; - if (config.needsPrepack) { - const auto m = config.M; - const auto k = config.K; - - // Optimize the layout of input matrix A(MxK) for SubgroupMatrixLoad. - PrepackProgram prepack_program{m, k}; - prepack_program.SetWorkgroupSize(config.subgroupSize); - if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { - prepack_program.SetSubgroupSize(config.subgroupSize); - } - - // Pad M to workgroup tile size so all subgroups read valid prepacked data. - const uint32_t padded_M = CeilDiv(M, tile_m) * tile_m; - const auto dispatch_group_size_x = padded_M / m; - ORT_ENFORCE(K % k == 0, "K must be a multiple of ", k); - const auto dispatch_group_size_y = K / k; - // Each workgroup will process one subgroup matrix of size m x k. - prepack_program.SetDispatchGroupSize(dispatch_group_size_x, dispatch_group_size_y, 1); - - TensorShape a_prepack_shape{padded_M, K}; - a_prepack = context.CreateGPUTensor(a->DataType(), a_prepack_shape); - prepack_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}}) - .AddOutputs({{&a_prepack, ProgramTensorMetadataDependency::Rank, a_prepack.Shape(), 1}}) - .AddUniformVariables({{M}, {K}}) - .CacheHint(m, k); - ORT_RETURN_IF_ERROR(context.RunProgram(prepack_program)); + if (prepacked_a != nullptr) { + ORT_ENFORCE(config.needsPrepack && prepacked_a->DataType() == a->DataType() && + prepacked_a->Shape().NumDimensions() == 2 && prepacked_a->Shape()[1] == K && + prepacked_a->Shape()[0] >= CeilDiv(M, tile_m) * tile_m); + a = prepacked_a; + } else if (config.needsPrepack) { + ORT_RETURN_IF_ERROR(PrepackMatMulNBitsInput(a, M, K, CeilDiv(M, tile_m) * tile_m, + config, context, a_prepack)); a = &a_prepack; } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h index 240c476675e72..a0e037152e4a2 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h @@ -58,6 +58,12 @@ class SubgroupMatrixMatMulNBitsProgram final : public Program +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" +#include "core/framework/ort_value_name_idx_map.h" +#include "core/graph/model.h" +#include "core/providers/webgpu/allocator.h" +#include "core/providers/webgpu/compute_context.h" +#include "core/providers/webgpu/webgpu_provider_factory_creator.h" +#include "test/test_environment.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { +namespace { + +using namespace webgpu; +namespace quantization = contrib::webgpu; + +class PreparationTestKernel final : public OpKernel { + public: + explicit PreparationTestKernel(const OpKernelInfo& info) : OpKernel(info) {} + Status Compute(OpKernelContext*) const override { return Status::OK(); } +}; + +class PreparationTestContext final : public OpKernelContext { + public: + explicit PreparationTestContext(AllocatorPtr allocator) + : OpKernelContext(nullptr, DefaultLoggingManager().DefaultLogger(), nullptr), allocator_(std::move(allocator)) {} + Status GetTempSpaceAllocator(AllocatorPtr* output) const override { + *output = allocator_; + return Status::OK(); + } + + private: + AllocatorPtr allocator_; +}; + +class MatMulNBitsInputPreparationTest : public ::testing::Test { + protected: + void SetUp() override { + ep_ = WebGpuProviderFactoryCreator::Create(options_)->CreateProvider(); + ep_->SetLogger(&DefaultLoggingManager().DefaultLogger()); + allocator_ = std::make_shared( + 0, [this]() -> const BufferManager& { return Ep().BufferManager(); }, + [this]() -> CommandRecordingState& { return Ep().Recording(); }, false); + } + + WebGpuExecutionProvider& Ep() { return static_cast(*ep_); } + + template + Tensor Upload(const TensorShape& shape, std::vector values) { + Tensor tensor(DataTypeImpl::GetType(), shape, allocator_); + Ep().BufferManager().Upload(Ep().Recording(), values.data(), + reinterpret_cast(tensor.MutableDataRaw()), tensor.SizeInBytes()); + return tensor; + } + + template + std::vector Download(Tensor& tensor) { + std::vector values(tensor.Shape().Size()); + Ep().BufferManager().Download(Ep().Recording(), reinterpret_cast(tensor.MutableDataRaw()), + values.data(), tensor.SizeInBytes()); + return values; + } + + template + void Check(bool subgroup_matrix) { + auto& device = WebGpuContextFactory::GetContext(0); + if constexpr (std::is_same_v) { + if (!device.DeviceHasFeature(wgpu::FeatureName::ShaderF16)) { + GTEST_SKIP() << "FP16 is unavailable."; + } + } + Model model("input_preparation", false, DefaultLoggingManager().DefaultLogger()); + auto& node = model.MainGraph().AddNode("unused", "Identity", "", {}, {}); + auto kernel_def = KernelDefBuilder().SetName("Identity").Provider(kWebGpuExecutionProvider).SinceVersion(1).Build(); + const std::unordered_map initializers; + const OrtValueNameIdxMap values; + const DataTransferManager transfers; + const AllocatorMap allocators; + OpKernelInfo info(node, *kernel_def, Ep(), initializers, values, transfers, allocators, options_); + PreparationTestKernel kernel(info); + PreparationTestContext kernel_context(allocator_); + webgpu::ComputeContext context(device, Ep(), kernel, kernel_context); + constexpr uint32_t M = 65, K = 256; + std::optional config; + if (subgroup_matrix) { + if (!quantization::CanApplySubgroupMatrixMatMulNBits(context, 4, 32, 1, 64, K, 4, + std::is_same_v, config, M, false) || + !config->needsPrepack) { + GTEST_SKIP() << "A compatible subgroup-matrix layout is unavailable."; + } + } else if (!quantization::CanApplyDP4AMatrixMatMulNBits(context, 4, 32, 64, K, 4)) { + GTEST_SKIP() << "DP4A is unavailable."; + } + // Each invocation owns a new workspace; two projections reuse it within that invocation. + for (int invocation = 0; invocation < 2; ++invocation) { + std::vector input(M * K); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = T(static_cast(static_cast(i % 17) - 8 + invocation) * 0.125f); + } + auto a = Upload({M, K}, input); + quantization::DP4AQuantizedInput workspace; + Tensor packed; + if (subgroup_matrix) { + // M=65 exercises padding; 256 rows cover the supported workgroup row tiles. + ASSERT_STATUS_OK(quantization::PrepackMatMulNBitsInput(&a, M, K, 256, *config, context, packed)); + } + const void* quantized_values = nullptr; + const void* quantized_scales = nullptr; + for (uint32_t projection = 0; projection < 2; ++projection) { + const uint32_t N = subgroup_matrix ? 64 : (projection == 0 ? 16 : 64); + const uint32_t bits = projection == 0 ? 8 : 4; + auto b = Upload({N, K / 32, 4 * bits}, + std::vector(N * K * bits / 8, bits == 8 ? 131 : 0x99)); + auto scales = Upload({N, K / 32}, std::vector(N * K / 32, T(0.125f))); + Tensor shared(DataTypeImpl::GetType(), TensorShape{M, N}, allocator_); + Tensor fresh(DataTypeImpl::GetType(), TensorShape{M, N}, allocator_); + if (subgroup_matrix) { + ASSERT_STATUS_OK(quantization::ApplySubgroupMatrixMatMulNBits( + &a, &b, &scales, nullptr, nullptr, M, N, K, bits, K / 32, *config, context, &shared, + 0, nullptr, &packed)); + ASSERT_STATUS_OK(quantization::ApplySubgroupMatrixMatMulNBits( + &a, &b, &scales, nullptr, nullptr, M, N, K, bits, K / 32, *config, context, &fresh, 0)); + } else { + ASSERT_STATUS_OK(quantization::ApplyDP4AMatrixMatMulNBits( + &a, &b, &scales, nullptr, nullptr, 1, M, M, N, K, 32, K / 32, 4, bits, + context, &shared, 0, nullptr, &workspace)); + ASSERT_TRUE(workspace.values.has_value()); + ASSERT_TRUE(workspace.scales.has_value()); + if (projection == 0) { + quantized_values = workspace.values->DataRaw(); + quantized_scales = workspace.scales->DataRaw(); + } else { + EXPECT_EQ(workspace.values->DataRaw(), quantized_values); + EXPECT_EQ(workspace.scales->DataRaw(), quantized_scales); + } + ASSERT_STATUS_OK(quantization::ApplyDP4AMatrixMatMulNBits( + &a, &b, &scales, nullptr, nullptr, 1, M, M, N, K, 32, K / 32, 4, bits, context, &fresh, 0)); + } + EXPECT_EQ(Download(shared), Download(fresh)); + } + } + } + + private: + ConfigOptions options_; + std::unique_ptr ep_; + AllocatorPtr allocator_; +}; + +TEST_F(MatMulNBitsInputPreparationTest, SharedDP4AInputFloat32) { Check(false); } +TEST_F(MatMulNBitsInputPreparationTest, SharedDP4AInputFloat16) { Check(false); } +TEST_F(MatMulNBitsInputPreparationTest, SharedSubgroupMatrixInputFloat16) { Check(true); } + +} // namespace +} // namespace test +} // namespace onnxruntime +#endif From abec758b2b49eba362d0c33ca766d60bfaa783a5 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Oct 2026 10:48:51 +0800 Subject: [PATCH 2/6] Optimize narrow WebGPU MatMulNBits DP4A projections Use rectangular tiles and K-parallel reduction for narrow projections, while retaining the existing full-width load schedule. Cover INT2/4/8 correctness and long-K dispatch, including the existing activation-quantization error in test tolerances. --- .../quantization/dp4a_matmul.wgsl.template | 408 +++++++++--------- ...e => dp4a_matmul_k_parallel.wgsl.template} | 39 +- .../webgpu/quantization/dp4a_matmul_nbits.cc | 164 +++---- .../webgpu/quantization/dp4a_matmul_nbits.h | 99 ++--- .../webgpu/quantization/matmul_nbits.cc | 2 +- .../test/contrib_ops/matmul_2bits_test.cc | 33 ++ .../test/contrib_ops/matmul_4bits_test.cc | 24 +- .../test/contrib_ops/matmul_8bits_test.cc | 25 ++ 8 files changed, 430 insertions(+), 364 deletions(-) rename onnxruntime/contrib_ops/webgpu/quantization/{dp4a_matmul_small_m.wgsl.template => dp4a_matmul_k_parallel.wgsl.template} (82%) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index 19108b848c420..fa28ac0ee568c 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -9,51 +9,33 @@ #param is_qualcomm #param has_weight_idx #param has_weight_idx_indirect +#param tile_size_k_vec +#param tile_size_m +#param tile_size_n +#param workgroup_size #use .getByOffset .setByOffset #include "quantization/dp4a_matmul_common.wgsl.template" -// This shader implements co-operative matrix multiply. The key idea here is to -// assume there is a primitive for medium size matrix multiply a subgroup can perform, -// using all its lanes and pooling all its registers to keep the values in registry. -// -// The entire workgroup which has N subgroups first loads a tile into shared memory, -// Then each subgroup loads a subtile from shared memory into registers and uses -// the medium size matrix multiply primitive to perform the math. -// The values for tile/subtile size are chosen to conform to the resource limits -// of an alderlake/tiger lake gpu. A tile is 64x64, workgroup is 256 threads - -// therefore there are 16 subgroups and 16 lanes in each subgroup. -// K the hidden dimension is paged in from RAM at k tile size which is 64. -// All this puts the shared memory requirement slightly above 16KB. -// WebGPU limit is 16KB, output is moved to registers instead of SHM to make -// everything fit in shared memory. -// -// Each subgroup performs a 16 x 64 x 16 multiply which is implemented with -// subgroup shuffle as a placeholder for the day the medium matrix mul primitive -// becomes available in WGSL. The registry requirements is ~2KB per subgroup, on -// Alderlake/Tigerlake subgroup has 8KB of registry space pooling the -// 512B of registry from each lane. -// -// The medium size matmul is implemented using dot4I8Packed, so the inputs for -// this shader require A to be int8 quantized with block size 64. B is regular -// matmulnbits input with block size 32. +// Each 16-thread group computes a 16x16 output subtile. Rectangular workgroup +// tiles avoid wasting half the columns on narrow projections. Wider K loads +// amortize barriers while retaining the original 32-element dot/scale order. +// A is dynamically INT8-quantized per 128 elements; B uses its model block size. -const tile_size = 64; +const subtiles_n = tile_size_n / 16; const subtile_size = 16; -const tile_size_k = 32; const vec_factor = 4; const u32_factor = 4; -const tile_size_k_vec = 2; // Shared memory -var tile_A : array, tile_size>, tile_size_k_vec>; // 64 x 32 -var scale_A : array; // 64 x 1 -var tile_B : array, tile_size>, tile_size_k_vec>; // 64 x 32 -var scale_B : array; // 64 x 1 +var tile_A : array, tile_size_m>, tile_size_k_vec>; +var scale_A : array; +var tile_B : array, tile_size_n>, tile_size_k_vec>; +var scale_B : array; #if has_zero_points - var zeroes : array; + var zeroes : array; #endif fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) @@ -91,14 +73,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const actual_weight_idx : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16+kidx_v + col); #endif - let block_idx = kidx_v/(block_size/16); + let block_idx = (kidx_v + col)/(block_size/16); let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom4BitsTo8Bits(b_value, zero); - if (col == 0) + if (col % 2u == 0u) { // kidx_v - each kidx_v covers 16 values of k let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); } } #endif @@ -126,14 +108,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif tile_B[col][row] = AlignWithZeroPoint(b_value); - if (col == 0) + if (col % 2u == 0u) { // kidx_v - each kidx_v covers 16 values of k - let block_idx = kidx_v/(block_size/16); + let block_idx = (kidx_v + col)/(block_size/16); let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); #if has_zero_points - zeroes[row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); + zeroes[(col / 2u) * tile_size_n + row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); #endif } } @@ -160,15 +142,17 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const b_weight_offset : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif - let block_idx = kidx_v/(block_size/16); + let block_idx = (kidx_v + col)/(block_size/16); #if has_zero_points let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value, zero); #else tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value); #endif - let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + if (col % 2u == 0u) { + let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); + scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + } } #endif @@ -183,31 +167,33 @@ alias acc_element_t = output_element_t; $MAIN { #if n_bits == 2 - LoadDequantizationTable(local_idx); #if has_zero_points - LoadDequantizationTable(local_idx + 256); - LoadDequantizationTable(local_idx + 512); - LoadDequantizationTable(local_idx + 768); + const table_size = 1024u; +#else + const table_size = 256u; #endif + for (var i = local_idx; i < table_size; i += workgroup_size) { + LoadDequantizationTable(i); + } workgroupBarrier(); #endif - // During the load phase we use all 256 threads to load 64 rows of A/B. - // For each row we load tile_size_k_vec (2) vectorized elements, which are 32 elements of K. let batch = workgroup_idx / (uniforms.num_M_tile * uniforms.num_N_tile); if (batch >= uniforms.batch_count) { return; } - let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size; - let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size; - let load_AorB = u32(local_idx/128); - let load_row = u32((local_idx%128)/2); - let load_col = u32(local_idx%2); + let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size_m; + let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size_n; +#if tile_size_m == 64 && tile_size_n == 64 && tile_size_k_vec == 2 + // Keep the original one-load-per-thread schedule for full-width tiles. + let load_AorB = u32(local_idx / 128); + let load_row = u32((local_idx % 128) / 2); + let load_col = u32(local_idx % 2); +#endif - // During the compute phase, we have the 64x64 tile split into - // subtiles of 16x16. We have a grid of 4x4 subtiles. + // Each group of 16 threads computes one 16x16 output subtile. var subtile_id = u32(local_idx / subtile_size); - var subtile_idx = u32(subtile_id / 4); - var subtile_idy = u32(subtile_id % 4); + var subtile_idx = u32(subtile_id / subtiles_n); + var subtile_idy = u32(subtile_id % subtiles_n); var base_A = subtile_idx * 16; var base_B = subtile_idy * 16; // For each subtile we have 16 threads assigned. @@ -221,8 +207,8 @@ $MAIN { // we would need var lane_outputs: array; if (sg_size == 64) { subtile_id = u32(local_idx / sg_size); - subtile_idx = u32(subtile_id / 4); - subtile_idy = u32(subtile_id % 4); + subtile_idx = u32(subtile_id / subtiles_n); + subtile_idy = u32(subtile_id % subtiles_n); base_A = subtile_idx * sg_size; base_B = subtile_idy * 16; a_idx = sg_id; @@ -235,189 +221,203 @@ $MAIN { var lane_output4: vec4; #endif // K's vectorization is 16 items per index. See input_a/input_b. - // tile_size_k_vec - is the k tile size in vectorized space (1/16). That is - // k tile size is 32. In vectorized space that is 32/16 = 2. + // Load 32 or 128 K elements, represented by 2 or 8 vec4 values. for (var kidx_v:u32 = 0; kidx_v < uniforms.K16; kidx_v+=tile_size_k_vec) { // Load Phase: Populate shared memory for the workgroup. - if (load_AorB == 0) - { +#if tile_size_m == 64 && tile_size_n == 64 && tile_size_k_vec == 2 + if (load_AorB == 0) { loadSHMA(batch, a_global_base, kidx_v, load_row, load_col); - } - else - { + } else { loadSHMB(b_global_base, kidx_v, load_row, load_col); } +#else + for (var i = local_idx; i < (tile_size_m + tile_size_n) * tile_size_k_vec; i += workgroup_size) { + if (i < tile_size_m * tile_size_k_vec) { + loadSHMA(batch, a_global_base, kidx_v, i / tile_size_k_vec, i % tile_size_k_vec); + } else { + let j = i - tile_size_m * tile_size_k_vec; + loadSHMB(b_global_base, kidx_v, j / tile_size_k_vec, j % tile_size_k_vec); + } + } +#endif workgroupBarrier(); - // Compute phase: Perform matmul for this subtile 16 x 32 x 16. - // Step 1: Load from shared memory into registers across entire subgroup. - var own_a0: vec4 = tile_A[0][base_A + a_idx]; - var own_a1: vec4 = tile_A[1][base_A + a_idx]; - var own_scale_a: output_element_t = scale_A[base_A + a_idx]; + // Keep the original per-block scaling and accumulation order. +#if tile_size_k_vec == 2 + { + const k_sub = 0u; +#else + for (var k_sub = 0u; k_sub < tile_size_k_vec / 2u; k_sub++) { +#endif + // Step 1: Load from shared memory into registers across entire subgroup. + var own_a0: vec4 = tile_A[k_sub * 2u][base_A + a_idx]; + var own_a1: vec4 = tile_A[k_sub * 2u + 1u][base_A + a_idx]; + var own_scale_a: output_element_t = scale_A[base_A + a_idx]; #if has_zero_points && n_bits == 8 #if is_qualcomm - if (sg_size == 64) - { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - var zero: i32; - if (sg_id < 16) + if (sg_size == 64) { - own_b0 = tile_B[0][base_B + sg_id]; - own_b1 = tile_B[1][base_B + sg_id]; - own_scale_b = scale_B[base_B + sg_id]; - zero = zeroes[base_B + sg_id]; + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + var zero: i32; + if (sg_id < 16) + { + own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; + own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; + zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; + } + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); + } } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) + else { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i], zeroes[k_sub * tile_size_n + base_B + i]); + } } - } - else - { - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i], zeroes[base_B + i]); - } - } #else - if (sg_size == 16) - { - var own_b0: vec4 = tile_B[0][base_B + sg_id]; - var own_b1: vec4 = tile_B[1][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[base_B + sg_id]; - var zero = zeroes[base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); + if (sg_size == 16) + { + var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; + var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; + var zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); - } - else - { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0], zeroes[base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1], zeroes[base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2], zeroes[base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3], zeroes[base_B + 3]); + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0], zeroes[k_sub * tile_size_n + base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1], zeroes[k_sub * tile_size_n + base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2], zeroes[k_sub * tile_size_n + base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3], zeroes[k_sub * tile_size_n + base_B + 3]); - lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4], zeroes[base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5], zeroes[base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6], zeroes[base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7], zeroes[base_B + 7]); + lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4], zeroes[k_sub * tile_size_n + base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5], zeroes[k_sub * tile_size_n + base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6], zeroes[k_sub * tile_size_n + base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7], zeroes[k_sub * tile_size_n + base_B + 7]); - lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8], zeroes[base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9], zeroes[base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10], zeroes[base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11], zeroes[base_B + 11]); + lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8], zeroes[k_sub * tile_size_n + base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9], zeroes[k_sub * tile_size_n + base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10], zeroes[k_sub * tile_size_n + base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11], zeroes[k_sub * tile_size_n + base_B + 11]); - lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12], zeroes[base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13], zeroes[base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14], zeroes[base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15], zeroes[base_B + 15]); - } + lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12], zeroes[k_sub * tile_size_n + base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13], zeroes[k_sub * tile_size_n + base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14], zeroes[k_sub * tile_size_n + base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15], zeroes[k_sub * tile_size_n + base_B + 15]); + } #endif #else #if is_qualcomm - if (sg_size == 64) - { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - if (sg_id < 16) + if (sg_size == 64) { - own_b0 = tile_B[0][base_B + sg_id]; - own_b1 = tile_B[1][base_B + sg_id]; - own_scale_b = scale_B[base_B + sg_id]; + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + if (sg_id < 16) + { + own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; + own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; + } + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); + } } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); - } - } - else - { - for (var i = 0u; i < 16u; i++) + else { - lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i]); + for (var i = 0u; i < 16u; i++) + { + lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i]); + } } - } #else - if (sg_size == 16) - { - var own_b0: vec4 = tile_B[0][base_B + sg_id]; - var own_b1: vec4 = tile_B[1][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); + if (sg_size == 16) + { + var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; + var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); - } - else - { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3]); + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3]); - lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7]); + lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7]); - lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11]); + lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11]); - lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15]); - } + lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15]); + } #endif #endif + } workgroupBarrier(); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template similarity index 82% rename from onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template rename to onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template index a53e62921b427..2bed945685d16 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template @@ -1,6 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +// K-parallel DP4A projection: each workgroup computes one input row and +// tile_size output columns. Its 128 threads process four columns concurrently, +// using 32 K lanes per column and 32 elements per lane in each K iteration. +// The group stages 1024 activations and their scales in workgroup memory, +// reusing them across output columns, including additional four-column groups. +// Each lane accumulates scaled integer dot products; the 32 lane partials are +// reduced through inter_results to produce each output value after processing K. +// This exposes parallelism along K when small M or N yields few output tiles. + #param acc_f32 #param tile_size #param tile_size_k_vec @@ -15,35 +24,13 @@ #use .getByOffset .setByOffset - #include "quantization/dp4a_matmul_common.wgsl.template" -// This algorithm works to compute dot product of k in parallel, by processing k at each step amongst tile_size_k_vec threads, -// and utilizing the remaining threads in the workgroup to process additional rows of b in parallel (such that the values in shared memory for A can be reused). -// For each load of k, the tile_size_k_vec threads also reload B tile_size/num_concurrent_b_rows times to compute partial dot products of other B rows -// in order to complete all tile_size b rows in this workgroup and also reusing the loaded in register values of a. - -// 1. Each workgroup handles tile_size_k_vec * k_vectorization_in_b (32) columns and num_concurrent_b_rows of matrix B at a time, -// iterating over the columns to compute a partial dot product. -// 2. Uses vec4 vectorization where each K represents 32 elements of matrix B - -// 1. Workgroup Responsibility: -// - Processes one row of matrix A -// - Handles tile_size rows of matrix B -// -// 2. Computation Process: -// - Reads [tile_size][tile_size_k_vec] block of B data at a time -// - Each thread within workgroup computes dot products of 32 A*B elements since each K represents 32 elements of matrix B -// - Stores intermediate results in shared memory (inter_results) -// - Iterates through columns accumulating results in inter_results -// - Performs final reduction sum in inter_results for output -// sub_tile_count is the number of concurrent b rows processed by the workgroup. - const double_tile_size_k_vec = 2 * tile_size_k_vec; -// Precision of the cross-tile accumulator, selected by the -// "enableMatmulFp32Accumulation" provider option. SDP8AI itself is an exact integer -// dot product; it is the accumulation over K that can overflow the f16 max (~65504). +// Narrow outputs use FP32 cross-tile accumulation to limit reduction error; +// other outputs follow "enableMatmulFp32Accumulation". SDP8AI itself is an exact +// integer dot product; accumulation over K can overflow the f16 max (~65504). #if acc_f32 alias acc_element_t = f32; #else @@ -64,7 +51,7 @@ fn loadSHMA(batch: u32, a_global: u32, kidx_v: u32, col: u32) } tile_A[col] = a.getByOffset(batch*uniforms.M*uniforms.K16+a_global*uniforms.K16+k_offset); - if (col < scale_a_size_in_tile_a) + if (col < scale_a_size_in_tile_a && kidx_v/8 + col < uniforms.K16/8) { // kidx_v - covers 16 values of k in input_a scale_A[col] = scales_a.getByOffset(batch*uniforms.M*(uniforms.K/128) + a_global*(uniforms.K/128) + kidx_v/8 + col); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 37660060a87b4..5287fb4214d33 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -10,17 +10,17 @@ namespace contrib { namespace webgpu { Status DP4AMatMulQuantizeProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | + ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform); const auto& scales = shader.AddOutput("scales", ShaderUsage::UseUniform); - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_quantize.wgsl.template", - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales, scales)); + return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_quantize.wgsl.template", WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales, scales)); } Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); + const auto& a = shader.AddInput( + "input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); const auto& scales_a = shader.AddInput("scales_a", ShaderUsage::UseUniform); const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform); const auto& scales_b = shader.AddInput("scales_b", ShaderUsage::UseUniform); @@ -34,25 +34,21 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(block_size, block_size_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/dp4a_matmul.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(block_size, block_size_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), + WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_PARAMETER(tile_size_m, tile_size_m_), + WGSL_TEMPLATE_PARAMETER(tile_size_n, tile_size_n_), WGSL_TEMPLATE_PARAMETER(workgroup_size, WorkgroupSizeX()), + WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } // scale_A components = 1, b components = 4, output components = 1 -Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) const { +Status DP4AMatMulNBitsKParallelProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform); const auto& scales_a = shader.AddInput("scales_a", ShaderUsage::UseUniform); const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform); @@ -68,28 +64,24 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); + ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, + "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); const uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec_; ORT_ENFORCE(tile_size_ % sub_tile_count == 0, "tile_size_ must be divisible by sub_tile_count"); - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_small_m.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), - WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), - WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/dp4a_matmul_k_parallel.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), + WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), + WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), + WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, @@ -139,26 +131,41 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index != 0 || has_weight_idx_indirect; const bool single_scale_weights = (block_size == K * N); - if (has_weight_idx_indirect || M < min_M_for_tile_optimization) { - uint32_t tile_size_k_vec = 32; - uint32_t tile_size_n = 4; + const bool small_n = N <= 32; + // Each K-parallel iteration covers 32 lanes * 32 elements = 1024 K values. + // For narrow N, conservatively require two iterations to amortize reduction + // overhead; shorter K benefits more from the tiled kernel's reuse across M. + const bool use_k_parallel = M < min_M_for_tile_optimization || (small_n && K >= 2048); + // Indirect expert weights can differ by row, so they also require the row-wise kernel. + if (use_k_parallel || has_weight_idx_indirect) { + const uint32_t tile_size_k_vec = 32; + // Keep more workgroups available when there are only a few input rows. + const uint32_t tile_size_n = small_n && M >= min_M_for_tile_optimization ? N : 4; const uint32_t b_components = (nbits == 2 ? kVec2Components : kVec4Components); const bool broadcast_a = dispatch_M > M; - const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32}; + // The narrow projection changes the K reduction order. Keep partial sums + // and the final reduction in FP32, then cast once when storing the output. + const bool acc_f32 = small_n || context.EnableMatmulFp32Accumulation(); + DP4AMatMulNBitsKParallelProgram mul_program{ + tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, + has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile); - mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, - {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(b_components * kU32Components)}, - {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) - .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index, dispatch_M}) + mul_program + .AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, + {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(b_components * kU32Components)}, + {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, + weight_index, dispatch_M}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1}) - .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, broadcast_a, acc_f32); + .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, + has_weight_idx_indirect, broadcast_a, acc_f32); if (has_zero_points) { - mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); + mul_program.AddInput( + {zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } if (has_bias) { mul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); @@ -169,19 +176,26 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor return context.RunProgram(mul_program); } - constexpr uint32_t kTileSize = 64; + // The 64-lane subgroup variant spans 64 rows. Reserve 32-row tiles for + // devices whose subgroup range cannot select that variant. + const uint32_t tile_size_m = N < 64 && context.AdapterInfo().subgroupMaxSize < 64 ? 32 : 64; + const uint32_t tile_size_n = N < 64 ? 32 : 64; + const uint32_t tile_size_k_vec = N < 64 && nbits == 8 && block_size == 32 ? 8 : 2; TensorShape reshaped_y_shape{batch_count, M, N / kVec4Components}; - uint32_t num_M_tile = (M + kTileSize - 1) / kTileSize; - uint32_t num_N_tile = (N + kTileSize - 1) / kTileSize; + uint32_t num_M_tile = CeilDiv(M, tile_size_m); + uint32_t num_N_tile = CeilDiv(N, tile_size_n); bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32}; - mul_program.SetWorkgroupSize(256); + DP4AMatMulNBitsProgram mul_program{ + tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, + has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32}; + mul_program.SetWorkgroupSize(tile_size_m * tile_size_n / 16); mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile); - mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, - {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, - {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + mul_program + .AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, + {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, + {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({{static_cast(batch_count)}, {static_cast(M)}, {static_cast(N)}, @@ -193,9 +207,11 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32); + .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, + has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32); if (has_zero_points) { - mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); + mul_program.AddInput( + {zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } if (has_bias) { mul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); @@ -206,23 +222,15 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor return context.RunProgram(mul_program); } -bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, - uint64_t accuracy_level, - uint32_t block_size, - uint32_t N, - uint32_t K, - uint32_t components_k, - uint32_t M, - bool has_weight_idx_indirect, - const Tensor* y) { +bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level, + uint32_t block_size, uint32_t N, uint32_t K, uint32_t components_k, uint32_t M, + bool has_weight_idx_indirect, const Tensor* y) { // macOS - Avoid using dp4a on Metal, as it does not appear to have native dp4a support. // https://github.com/gpuweb/gpuweb/issues/2677#issuecomment-1713292226 // Use 'vendor' to check for metal; 'backend' is always WEBGPU when running under wasm. - bool use_dp4a = context.HasFeature(wgpu::FeatureName::Subgroups) && - context.AdapterInfo().vendor != std::string_view{"apple"}; - if (!(accuracy_level == 4 && block_size % 32 == 0 && - components_k == 4 && K % 128 == 0 && N % 16 == 0 && - use_dp4a)) { + bool use_dp4a = + context.HasFeature(wgpu::FeatureName::Subgroups) && context.AdapterInfo().vendor != std::string_view{"apple"}; + if (!(accuracy_level == 4 && block_size % 32 == 0 && components_k == 4 && K % 128 == 0 && N % 16 == 0 && use_dp4a)) { return false; } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 2876802986866..9aa06c519ba62 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -30,32 +30,34 @@ class DP4AMatMulQuantizeProgram final : public Program { public: - DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, - bool has_zero_points, bool has_bias, - bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, - bool acc_f32) : Program{"DP4AMatMulNBits"}, - block_size_(block_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - is_qualcomm_(is_qualcomm), - acc_f32_(acc_f32) {} + DP4AMatMulNBitsProgram(uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k_vec, uint32_t block_size, + uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, + bool has_weight_idx_indirect, bool is_qualcomm, bool acc_f32) + : Program{"DP4AMatMulNBits"}, + tile_size_m_(tile_size_m), + tile_size_n_(tile_size_n), + tile_size_k_vec_(tile_size_k_vec), + block_size_(block_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + is_qualcomm_(is_qualcomm), + acc_f32_(acc_f32) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K8", ProgramUniformVariableDataType::Uint32}, - {"K16", ProgramUniformVariableDataType::Uint32}, - {"num_M_tile", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, + {"K8", ProgramUniformVariableDataType::Uint32}, {"K16", ProgramUniformVariableDataType::Uint32}, + {"num_M_tile", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, {"weight_idx", ProgramUniformVariableDataType::Uint32}); private: + uint32_t tile_size_m_; + uint32_t tile_size_n_; + uint32_t tile_size_k_vec_; uint32_t block_size_; uint32_t nbits_; bool has_bias_; @@ -66,36 +68,30 @@ class DP4AMatMulNBitsProgram final : public Program { bool acc_f32_; }; -class DP4AMatMulNBitsSmallMProgram final : public Program { +class DP4AMatMulNBitsKParallelProgram final : public Program { public: - DP4AMatMulNBitsSmallMProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, - bool has_zero_points, bool has_bias, - bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, - bool broadcast_a_row = false, - bool acc_f32 = false) : Program{"DP4AMatMulNBitsSmallMProgram"}, - tile_size_k_vec_(tile_size_k_vec), - tile_size_(tile_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - single_scale_weights_(single_scale_weights), - broadcast_a_row_(broadcast_a_row), - acc_f32_(acc_f32) {} + DP4AMatMulNBitsKParallelProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, bool has_zero_points, + bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + bool single_scale_weights, bool broadcast_a_row = false, bool acc_f32 = false) + : Program{"DP4AMatMulNBitsKParallelProgram"}, + tile_size_k_vec_(tile_size_k_vec), + tile_size_(tile_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + single_scale_weights_(single_scale_weights), + broadcast_a_row_(broadcast_a_row), + acc_f32_(acc_f32) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K16", ProgramUniformVariableDataType::Uint32}, - {"K32", ProgramUniformVariableDataType::Uint32}, - {"block_size", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, + {"K16", ProgramUniformVariableDataType::Uint32}, {"K32", ProgramUniformVariableDataType::Uint32}, + {"block_size", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}, - {"dispatch_M", ProgramUniformVariableDataType::Uint32}); + {"weight_idx", ProgramUniformVariableDataType::Uint32}, {"dispatch_M", ProgramUniformVariableDataType::Uint32}); private: uint32_t tile_size_k_vec_; @@ -132,15 +128,10 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor // unconditionally on FP32 outputs and Qualcomm GPUs) into the feasibility check // so callers don't need a separate wrapper. Defaults make the precondition // trivially satisfied for callers that only want the feasibility check. -bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, - uint64_t accuracy_level, - uint32_t block_size, - uint32_t N, - uint32_t K, - uint32_t components_k, +bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level, + uint32_t block_size, uint32_t N, uint32_t K, uint32_t components_k, uint32_t M = std::numeric_limits::max(), - bool has_weight_idx_indirect = false, - const Tensor* y = nullptr); + bool has_weight_idx_indirect = false, const Tensor* y = nullptr); } // namespace webgpu } // namespace contrib diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 664e8e873ec92..70762af0d163e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -66,7 +66,7 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons WGSL_TEMPLATE_VARIABLE(scales, scales)); } -// Apply similar idea with DP4AMatMulNBitsSmallMProgram algorithm. +// Uses the same K-parallel reduction strategy as DP4AMatMulNBitsKParallelProgram. Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias); const auto& b = shader.AddInput("input_b"); diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 34544efa0ac2e..09e9cf80247ec 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -255,6 +255,18 @@ void RunTest2Bits(const TestOptions2Bits& opts) { test.ConfigEps(std::move(execution_providers)); test.RunWithConfig(); } +#ifdef USE_WEBGPU + if constexpr (std::is_same::value) { + // FP16 WebGPU cases must execute too; an unused OpTester traps in Debug + // and silently omits coverage in Release. + execution_providers.push_back(DefaultWebGpuExecutionProvider()); + SessionOptions session_options; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + test.ConfigEps(std::move(execution_providers)); + test.RunWithConfig(); + } +#endif } template @@ -1520,6 +1532,27 @@ TEST(MatMul2Bits, Float32_2b_BlkLen32_Accuracy0) { TestMatMul2BitsTyped(); } +#ifdef USE_WEBGPU +TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int2) { + // Tiled row/column tails and the narrow-output K-parallel dispatch. + for (const auto& [n, k] : {std::pair{16, 256}, {48, 256}, {80, 256}, {32, 2048}}) { + for (bool has_zero_point : {false, true}) { + TestOptions2Bits opts{}; + opts.M = 65; + opts.N = n; + opts.K = k; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = k == 256 ? 0.05f : 0.1f; + opts.output_rel_error = 0.02f; + RunTest2Bits(opts); + RunTest2Bits(opts); + } + } +} +#endif + #if defined(USE_WEBGPU) && !defined(ORT_USE_EP_API_ADAPTERS) namespace { diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 4d7434ef398dd..c2a59a83aa214 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -452,6 +452,28 @@ TEST(MatMulNBits, Float32_4b_Accuracy0) { TestMatMulNBitsTyped(); } +#ifdef USE_WEBGPU +TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int4) { + // Tiled row/column tails and the narrow-output K-parallel dispatch. + for (const auto& [n, k] : {std::pair{16, 256}, {48, 256}, {80, 256}, {32, 2048}}) { + for (bool has_zero_point : {false, true}) { + TestOptions opts{}; + opts.batch_count = 2; + opts.M = 65; + opts.N = n; + opts.K = k; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = k == 256 ? 0.05f : 0.1f; + opts.output_rel_error = 0.02f; + RunTest(opts); + RunTest(opts); + } + } +} +#endif + TEST(MatMulNBits, Float32_4b_Accuracy1) { TestMatMulNBitsTyped(); TestMatMulNBitsTyped(); @@ -1268,7 +1290,7 @@ TEST(MatMulNBits, Float16_LargeK_AccumulatorOverflow) { // accuracy_level 4, M = 8 -> dp4a_matmul.wgsl.template (where the adapter supports it) // CanApplyDP4AMatrixMatMulNBits also needs subgroups, a non-Apple vendor and M >= 4 // (kMinMForTileOptimization) for an fp16 output, so the M = 2 case below only reaches -// dp4a_matmul_small_m on an fp32-output or Qualcomm adapter; elsewhere it lands on the generic +// dp4a_matmul_k_parallel on an fp32-output or Qualcomm adapter; elsewhere it lands on the generic // kernel and is one more shape for it rather than dp4a coverage. // With the option on and an adapter that has subgroup matrices, the subgroup-matrix path declines // itself (its cooperative-matrix result type is f16 and cannot honour the request) and the dispatch diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index 65ea7f5bf2ecb..864a1917079bb 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -3,6 +3,7 @@ #ifndef ORT_MINIMAL_BUILD #include +#include #include #include "gtest/gtest.h" @@ -1057,6 +1058,30 @@ TEST(MatMulNBits, BFloat16_Int8_SmallMFallback) { } #endif +#if defined(USE_WEBGPU) +TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int8) { + // Small-M reduction, the tiled threshold, and row/column tails in both tile widths. + for (int64_t k : {256, 1920, 2048, 4096}) { + for (const auto& [m, n] : {std::pair{3, 16}, {4, 32}, {65, 48}, {65, 80}}) { + for (bool has_zero_point : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m; + opts.N = n; + opts.K = k; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + // Activation-quantization error for these Gaussian inputs grows with sqrt(K). + opts.output_abs_error = 0.05f * std::sqrt(static_cast(k) / 256.0f); + opts.output_rel_error = 0.02f; + RunTest8Bits(opts); + RunTest8Bits(opts); + } + } + } +} +#endif + #if !defined(USE_CUDA) && !defined(USE_WEBGPU) #ifndef ENABLE_TRAINING // Pre-packing (and therefore cross-session sharing of pre-packed weights) is disabled in a full From 92fe6cd775f9d747be4f2289fa3f92d4c59ec7d7 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Oct 2026 11:43:10 +0800 Subject: [PATCH 3/6] Keep tiled DP4A unchanged in kernel extraction Restore the original 64x64 DP4A program, dispatch and WGSL. Remove small-tile test cases while retaining shared input preparation and focused INT2/4/8 K-parallel coverage. --- .../quantization/dp4a_matmul.wgsl.template | 408 +++++++++--------- .../webgpu/quantization/dp4a_matmul_nbits.cc | 114 ++--- .../webgpu/quantization/dp4a_matmul_nbits.h | 86 ++-- .../test/contrib_ops/matmul_2bits_test.cc | 29 +- .../test/contrib_ops/matmul_4bits_test.cc | 31 +- .../test/contrib_ops/matmul_8bits_test.cc | 36 +- 6 files changed, 354 insertions(+), 350 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index fa28ac0ee568c..19108b848c420 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -9,33 +9,51 @@ #param is_qualcomm #param has_weight_idx #param has_weight_idx_indirect -#param tile_size_k_vec -#param tile_size_m -#param tile_size_n -#param workgroup_size #use .getByOffset .setByOffset #include "quantization/dp4a_matmul_common.wgsl.template" -// Each 16-thread group computes a 16x16 output subtile. Rectangular workgroup -// tiles avoid wasting half the columns on narrow projections. Wider K loads -// amortize barriers while retaining the original 32-element dot/scale order. -// A is dynamically INT8-quantized per 128 elements; B uses its model block size. +// This shader implements co-operative matrix multiply. The key idea here is to +// assume there is a primitive for medium size matrix multiply a subgroup can perform, +// using all its lanes and pooling all its registers to keep the values in registry. +// +// The entire workgroup which has N subgroups first loads a tile into shared memory, +// Then each subgroup loads a subtile from shared memory into registers and uses +// the medium size matrix multiply primitive to perform the math. +// The values for tile/subtile size are chosen to conform to the resource limits +// of an alderlake/tiger lake gpu. A tile is 64x64, workgroup is 256 threads - +// therefore there are 16 subgroups and 16 lanes in each subgroup. +// K the hidden dimension is paged in from RAM at k tile size which is 64. +// All this puts the shared memory requirement slightly above 16KB. +// WebGPU limit is 16KB, output is moved to registers instead of SHM to make +// everything fit in shared memory. +// +// Each subgroup performs a 16 x 64 x 16 multiply which is implemented with +// subgroup shuffle as a placeholder for the day the medium matrix mul primitive +// becomes available in WGSL. The registry requirements is ~2KB per subgroup, on +// Alderlake/Tigerlake subgroup has 8KB of registry space pooling the +// 512B of registry from each lane. +// +// The medium size matmul is implemented using dot4I8Packed, so the inputs for +// this shader require A to be int8 quantized with block size 64. B is regular +// matmulnbits input with block size 32. -const subtiles_n = tile_size_n / 16; +const tile_size = 64; const subtile_size = 16; +const tile_size_k = 32; const vec_factor = 4; const u32_factor = 4; +const tile_size_k_vec = 2; // Shared memory -var tile_A : array, tile_size_m>, tile_size_k_vec>; -var scale_A : array; -var tile_B : array, tile_size_n>, tile_size_k_vec>; -var scale_B : array; +var tile_A : array, tile_size>, tile_size_k_vec>; // 64 x 32 +var scale_A : array; // 64 x 1 +var tile_B : array, tile_size>, tile_size_k_vec>; // 64 x 32 +var scale_B : array; // 64 x 1 #if has_zero_points - var zeroes : array; + var zeroes : array; #endif fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) @@ -73,14 +91,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const actual_weight_idx : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16+kidx_v + col); #endif - let block_idx = (kidx_v + col)/(block_size/16); + let block_idx = kidx_v/(block_size/16); let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom4BitsTo8Bits(b_value, zero); - if (col % 2u == 0u) + if (col == 0) { // kidx_v - each kidx_v covers 16 values of k let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); } } #endif @@ -108,14 +126,14 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif tile_B[col][row] = AlignWithZeroPoint(b_value); - if (col % 2u == 0u) + if (col == 0) { // kidx_v - each kidx_v covers 16 values of k - let block_idx = (kidx_v + col)/(block_size/16); + let block_idx = kidx_v/(block_size/16); let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); + scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); #if has_zero_points - zeroes[(col / 2u) * tile_size_n + row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); + zeroes[row] = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); #endif } } @@ -142,17 +160,15 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) const b_weight_offset : u32 = 0; let b_value = b.getByOffset(b_global * uniforms.K16 + kidx_v + col); #endif - let block_idx = (kidx_v + col)/(block_size/16); + let block_idx = kidx_v/(block_size/16); #if has_zero_points let zero = mm_read_zero(actual_weight_idx, b_global, block_idx, uniforms.N, uniforms.zero_blocks_per_col); tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value, zero); #else tile_B[col][row] = DequantizedFrom2BitsTo8Bits(b_value); #endif - if (col % 2u == 0u) { - let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); - scale_B[(col / 2u) * tile_size_n + row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); - } + let b_scale_offset = actual_weight_idx * uniforms.N * (uniforms.K/block_size); + scale_B[row] = scales_b.getByOffset(b_scale_offset + b_global*(uniforms.K/block_size) + block_idx); } #endif @@ -167,33 +183,31 @@ alias acc_element_t = output_element_t; $MAIN { #if n_bits == 2 + LoadDequantizationTable(local_idx); #if has_zero_points - const table_size = 1024u; -#else - const table_size = 256u; + LoadDequantizationTable(local_idx + 256); + LoadDequantizationTable(local_idx + 512); + LoadDequantizationTable(local_idx + 768); #endif - for (var i = local_idx; i < table_size; i += workgroup_size) { - LoadDequantizationTable(i); - } workgroupBarrier(); #endif + // During the load phase we use all 256 threads to load 64 rows of A/B. + // For each row we load tile_size_k_vec (2) vectorized elements, which are 32 elements of K. let batch = workgroup_idx / (uniforms.num_M_tile * uniforms.num_N_tile); if (batch >= uniforms.batch_count) { return; } - let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size_m; - let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size_n; -#if tile_size_m == 64 && tile_size_n == 64 && tile_size_k_vec == 2 - // Keep the original one-load-per-thread schedule for full-width tiles. - let load_AorB = u32(local_idx / 128); - let load_row = u32((local_idx % 128) / 2); - let load_col = u32(local_idx % 2); -#endif + let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size; + let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size; + let load_AorB = u32(local_idx/128); + let load_row = u32((local_idx%128)/2); + let load_col = u32(local_idx%2); - // Each group of 16 threads computes one 16x16 output subtile. + // During the compute phase, we have the 64x64 tile split into + // subtiles of 16x16. We have a grid of 4x4 subtiles. var subtile_id = u32(local_idx / subtile_size); - var subtile_idx = u32(subtile_id / subtiles_n); - var subtile_idy = u32(subtile_id % subtiles_n); + var subtile_idx = u32(subtile_id / 4); + var subtile_idy = u32(subtile_id % 4); var base_A = subtile_idx * 16; var base_B = subtile_idy * 16; // For each subtile we have 16 threads assigned. @@ -207,8 +221,8 @@ $MAIN { // we would need var lane_outputs: array; if (sg_size == 64) { subtile_id = u32(local_idx / sg_size); - subtile_idx = u32(subtile_id / subtiles_n); - subtile_idy = u32(subtile_id % subtiles_n); + subtile_idx = u32(subtile_id / 4); + subtile_idy = u32(subtile_id % 4); base_A = subtile_idx * sg_size; base_B = subtile_idy * 16; a_idx = sg_id; @@ -221,203 +235,189 @@ $MAIN { var lane_output4: vec4; #endif // K's vectorization is 16 items per index. See input_a/input_b. - // Load 32 or 128 K elements, represented by 2 or 8 vec4 values. + // tile_size_k_vec - is the k tile size in vectorized space (1/16). That is + // k tile size is 32. In vectorized space that is 32/16 = 2. for (var kidx_v:u32 = 0; kidx_v < uniforms.K16; kidx_v+=tile_size_k_vec) { // Load Phase: Populate shared memory for the workgroup. -#if tile_size_m == 64 && tile_size_n == 64 && tile_size_k_vec == 2 - if (load_AorB == 0) { + if (load_AorB == 0) + { loadSHMA(batch, a_global_base, kidx_v, load_row, load_col); - } else { - loadSHMB(b_global_base, kidx_v, load_row, load_col); } -#else - for (var i = local_idx; i < (tile_size_m + tile_size_n) * tile_size_k_vec; i += workgroup_size) { - if (i < tile_size_m * tile_size_k_vec) { - loadSHMA(batch, a_global_base, kidx_v, i / tile_size_k_vec, i % tile_size_k_vec); - } else { - let j = i - tile_size_m * tile_size_k_vec; - loadSHMB(b_global_base, kidx_v, j / tile_size_k_vec, j % tile_size_k_vec); - } + else + { + loadSHMB(b_global_base, kidx_v, load_row, load_col); } -#endif workgroupBarrier(); - // Keep the original per-block scaling and accumulation order. -#if tile_size_k_vec == 2 - { - const k_sub = 0u; -#else - for (var k_sub = 0u; k_sub < tile_size_k_vec / 2u; k_sub++) { -#endif - // Step 1: Load from shared memory into registers across entire subgroup. - var own_a0: vec4 = tile_A[k_sub * 2u][base_A + a_idx]; - var own_a1: vec4 = tile_A[k_sub * 2u + 1u][base_A + a_idx]; - var own_scale_a: output_element_t = scale_A[base_A + a_idx]; + // Compute phase: Perform matmul for this subtile 16 x 32 x 16. + // Step 1: Load from shared memory into registers across entire subgroup. + var own_a0: vec4 = tile_A[0][base_A + a_idx]; + var own_a1: vec4 = tile_A[1][base_A + a_idx]; + var own_scale_a: output_element_t = scale_A[base_A + a_idx]; #if has_zero_points && n_bits == 8 #if is_qualcomm - if (sg_size == 64) + if (sg_size == 64) + { + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + var zero: i32; + if (sg_id < 16) { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - var zero: i32; - if (sg_id < 16) - { - own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; - own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; - zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; - } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); - } + own_b0 = tile_B[0][base_B + sg_id]; + own_b1 = tile_B[1][base_B + sg_id]; + own_scale_b = scale_B[base_B + sg_id]; + zero = zeroes[base_B + sg_id]; } - else + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) { - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i], zeroes[k_sub * tile_size_n + base_B + i]); - } + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); } - #else - if (sg_size == 16) + } + else + { + for (var i = 0u; i < 16u; i++) { - var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; - var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; - var zero = zeroes[k_sub * tile_size_n + base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); + lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i], zeroes[base_B + i]); + } + } + #else + if (sg_size == 16) + { + var own_b0: vec4 = tile_B[0][base_B + sg_id]; + var own_b1: vec4 = tile_B[1][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[base_B + sg_id]; + var zero = zeroes[base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a, subgroupShuffle(zero, 0)); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a, subgroupShuffle(zero, 1)); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a, subgroupShuffle(zero, 2)); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a, subgroupShuffle(zero, 3)); - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a, subgroupShuffle(zero, 4)); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a, subgroupShuffle(zero, 5)); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a, subgroupShuffle(zero, 6)); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a, subgroupShuffle(zero, 7)); - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a, subgroupShuffle(zero, 8)); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a, subgroupShuffle(zero, 9)); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a, subgroupShuffle(zero, 10)); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a, subgroupShuffle(zero, 11)); - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); - } - else - { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0], zeroes[k_sub * tile_size_n + base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1], zeroes[k_sub * tile_size_n + base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2], zeroes[k_sub * tile_size_n + base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3], zeroes[k_sub * tile_size_n + base_B + 3]); + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a, subgroupShuffle(zero, 12)); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a, subgroupShuffle(zero, 13)); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a, subgroupShuffle(zero, 14)); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a, subgroupShuffle(zero, 15)); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0], zeroes[base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1], zeroes[base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2], zeroes[base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3], zeroes[base_B + 3]); - lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4], zeroes[k_sub * tile_size_n + base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5], zeroes[k_sub * tile_size_n + base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6], zeroes[k_sub * tile_size_n + base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7], zeroes[k_sub * tile_size_n + base_B + 7]); + lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4], zeroes[base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5], zeroes[base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6], zeroes[base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7], zeroes[base_B + 7]); - lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8], zeroes[k_sub * tile_size_n + base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9], zeroes[k_sub * tile_size_n + base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10], zeroes[k_sub * tile_size_n + base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11], zeroes[k_sub * tile_size_n + base_B + 11]); + lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8], zeroes[base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9], zeroes[base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10], zeroes[base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11], zeroes[base_B + 11]); - lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12], zeroes[k_sub * tile_size_n + base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13], zeroes[k_sub * tile_size_n + base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14], zeroes[k_sub * tile_size_n + base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15], zeroes[k_sub * tile_size_n + base_B + 15]); - } + lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12], zeroes[base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13], zeroes[base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14], zeroes[base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15], zeroes[base_B + 15]); + } #endif #else #if is_qualcomm - if (sg_size == 64) + if (sg_size == 64) + { + var own_b0: vec4; + var own_b1: vec4; + var own_scale_b: output_element_t; + if (sg_id < 16) { - var own_b0: vec4; - var own_b1: vec4; - var own_scale_b: output_element_t; - if (sg_id < 16) - { - own_b0 = tile_B[k_sub * 2u][base_B + sg_id]; - own_b1 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - own_scale_b = scale_B[k_sub * tile_size_n + base_B + sg_id]; - } - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); - } + own_b0 = tile_B[0][base_B + sg_id]; + own_b1 = tile_B[1][base_B + sg_id]; + own_scale_b = scale_B[base_B + sg_id]; } - else + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + for (var i = 0u; i < 16u; i++) { - for (var i = 0u; i < 16u; i++) - { - lane_outputs[i] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + i], own_a1, tile_B[k_sub * 2u + 1u][base_B + i], own_scale_a * scale_B[k_sub * tile_size_n + base_B + i]); - } + lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a); } - #else - if (sg_size == 16) + } + else + { + for (var i = 0u; i < 16u; i++) { - var own_b0: vec4 = tile_B[k_sub * 2u][base_B + sg_id]; - var own_b1: vec4 = tile_B[k_sub * 2u + 1u][base_B + sg_id]; - var own_scale_b: output_element_t = scale_B[k_sub * tile_size_n + base_B + sg_id]; - // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. - lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); - lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); - lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); - lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); + lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i], own_scale_a * scale_B[base_B + i]); + } + } + #else + if (sg_size == 16) + { + var own_b0: vec4 = tile_B[0][base_B + sg_id]; + var own_b1: vec4 = tile_B[1][base_B + sg_id]; + var own_scale_b: output_element_t = scale_B[base_B + sg_id]; + // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul. + lane_output1[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 0), own_a1, subgroupShuffle(own_b1, 0), subgroupShuffle(own_scale_b, 0) * own_scale_a); + lane_output1[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 1), own_a1, subgroupShuffle(own_b1, 1), subgroupShuffle(own_scale_b, 1) * own_scale_a); + lane_output1[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 2), own_a1, subgroupShuffle(own_b1, 2), subgroupShuffle(own_scale_b, 2) * own_scale_a); + lane_output1[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 3), own_a1, subgroupShuffle(own_b1, 3), subgroupShuffle(own_scale_b, 3) * own_scale_a); - lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); - lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); - lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); - lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); + lane_output2[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 4), own_a1, subgroupShuffle(own_b1, 4), subgroupShuffle(own_scale_b, 4) * own_scale_a); + lane_output2[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 5), own_a1, subgroupShuffle(own_b1, 5), subgroupShuffle(own_scale_b, 5) * own_scale_a); + lane_output2[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 6), own_a1, subgroupShuffle(own_b1, 6), subgroupShuffle(own_scale_b, 6) * own_scale_a); + lane_output2[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 7), own_a1, subgroupShuffle(own_b1, 7), subgroupShuffle(own_scale_b, 7) * own_scale_a); - lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); - lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); - lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); - lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); + lane_output3[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 8), own_a1, subgroupShuffle(own_b1, 8), subgroupShuffle(own_scale_b, 8) * own_scale_a); + lane_output3[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 9), own_a1, subgroupShuffle(own_b1, 9), subgroupShuffle(own_scale_b, 9) * own_scale_a); + lane_output3[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 10), own_a1, subgroupShuffle(own_b1, 10), subgroupShuffle(own_scale_b, 10) * own_scale_a); + lane_output3[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 11), own_a1, subgroupShuffle(own_b1, 11), subgroupShuffle(own_scale_b, 11) * own_scale_a); - lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); - lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); - lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); - lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); - } - else - { - // Code for other subgroup sizes, simply doesnt use subgroups at all. - // Relies on reads from single location tile_B[][base_B + col] by all - // being optimized by the hardware. - lane_output1[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 0], own_a1, tile_B[k_sub * 2u + 1u][base_B + 0], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 0]); - lane_output1[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 1], own_a1, tile_B[k_sub * 2u + 1u][base_B + 1], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 1]); - lane_output1[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 2], own_a1, tile_B[k_sub * 2u + 1u][base_B + 2], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 2]); - lane_output1[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 3], own_a1, tile_B[k_sub * 2u + 1u][base_B + 3], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 3]); + lane_output4[0] += SDP8AI(own_a0, subgroupShuffle(own_b0, 12), own_a1, subgroupShuffle(own_b1, 12), subgroupShuffle(own_scale_b, 12) * own_scale_a); + lane_output4[1] += SDP8AI(own_a0, subgroupShuffle(own_b0, 13), own_a1, subgroupShuffle(own_b1, 13), subgroupShuffle(own_scale_b, 13) * own_scale_a); + lane_output4[2] += SDP8AI(own_a0, subgroupShuffle(own_b0, 14), own_a1, subgroupShuffle(own_b1, 14), subgroupShuffle(own_scale_b, 14) * own_scale_a); + lane_output4[3] += SDP8AI(own_a0, subgroupShuffle(own_b0, 15), own_a1, subgroupShuffle(own_b1, 15), subgroupShuffle(own_scale_b, 15) * own_scale_a); + } + else + { + // Code for other subgroup sizes, simply doesnt use subgroups at all. + // Relies on reads from single location tile_B[][base_B + col] by all + // being optimized by the hardware. + lane_output1[0] += SDP8AI(own_a0, tile_B[0][base_B + 0], own_a1, tile_B[1][base_B + 0], own_scale_a * scale_B[base_B + 0]); + lane_output1[1] += SDP8AI(own_a0, tile_B[0][base_B + 1], own_a1, tile_B[1][base_B + 1], own_scale_a * scale_B[base_B + 1]); + lane_output1[2] += SDP8AI(own_a0, tile_B[0][base_B + 2], own_a1, tile_B[1][base_B + 2], own_scale_a * scale_B[base_B + 2]); + lane_output1[3] += SDP8AI(own_a0, tile_B[0][base_B + 3], own_a1, tile_B[1][base_B + 3], own_scale_a * scale_B[base_B + 3]); - lane_output2[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 4], own_a1, tile_B[k_sub * 2u + 1u][base_B + 4], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 4]); - lane_output2[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 5], own_a1, tile_B[k_sub * 2u + 1u][base_B + 5], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 5]); - lane_output2[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 6], own_a1, tile_B[k_sub * 2u + 1u][base_B + 6], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 6]); - lane_output2[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 7], own_a1, tile_B[k_sub * 2u + 1u][base_B + 7], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 7]); + lane_output2[0] += SDP8AI(own_a0, tile_B[0][base_B + 4], own_a1, tile_B[1][base_B + 4], own_scale_a * scale_B[base_B + 4]); + lane_output2[1] += SDP8AI(own_a0, tile_B[0][base_B + 5], own_a1, tile_B[1][base_B + 5], own_scale_a * scale_B[base_B + 5]); + lane_output2[2] += SDP8AI(own_a0, tile_B[0][base_B + 6], own_a1, tile_B[1][base_B + 6], own_scale_a * scale_B[base_B + 6]); + lane_output2[3] += SDP8AI(own_a0, tile_B[0][base_B + 7], own_a1, tile_B[1][base_B + 7], own_scale_a * scale_B[base_B + 7]); - lane_output3[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 8], own_a1, tile_B[k_sub * 2u + 1u][base_B + 8], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 8]); - lane_output3[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 9], own_a1, tile_B[k_sub * 2u + 1u][base_B + 9], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 9]); - lane_output3[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 10], own_a1, tile_B[k_sub * 2u + 1u][base_B + 10], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 10]); - lane_output3[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 11], own_a1, tile_B[k_sub * 2u + 1u][base_B + 11], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 11]); + lane_output3[0] += SDP8AI(own_a0, tile_B[0][base_B + 8], own_a1, tile_B[1][base_B + 8], own_scale_a * scale_B[base_B + 8]); + lane_output3[1] += SDP8AI(own_a0, tile_B[0][base_B + 9], own_a1, tile_B[1][base_B + 9], own_scale_a * scale_B[base_B + 9]); + lane_output3[2] += SDP8AI(own_a0, tile_B[0][base_B + 10], own_a1, tile_B[1][base_B + 10], own_scale_a * scale_B[base_B + 10]); + lane_output3[3] += SDP8AI(own_a0, tile_B[0][base_B + 11], own_a1, tile_B[1][base_B + 11], own_scale_a * scale_B[base_B + 11]); - lane_output4[0] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 12], own_a1, tile_B[k_sub * 2u + 1u][base_B + 12], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 12]); - lane_output4[1] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 13], own_a1, tile_B[k_sub * 2u + 1u][base_B + 13], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 13]); - lane_output4[2] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 14], own_a1, tile_B[k_sub * 2u + 1u][base_B + 14], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 14]); - lane_output4[3] += SDP8AI(own_a0, tile_B[k_sub * 2u][base_B + 15], own_a1, tile_B[k_sub * 2u + 1u][base_B + 15], own_scale_a * scale_B[k_sub * tile_size_n + base_B + 15]); - } + lane_output4[0] += SDP8AI(own_a0, tile_B[0][base_B + 12], own_a1, tile_B[1][base_B + 12], own_scale_a * scale_B[base_B + 12]); + lane_output4[1] += SDP8AI(own_a0, tile_B[0][base_B + 13], own_a1, tile_B[1][base_B + 13], own_scale_a * scale_B[base_B + 13]); + lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14], own_scale_a * scale_B[base_B + 14]); + lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15], own_scale_a * scale_B[base_B + 15]); + } #endif #endif - } workgroupBarrier(); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 5287fb4214d33..89bb45f513b3c 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -19,8 +19,7 @@ Status DP4AMatMulQuantizeProgram::GenerateShaderCode(ShaderHelper& shader) const } Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& a = shader.AddInput( - "input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); + const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); const auto& scales_a = shader.AddInput("scales_a", ShaderUsage::UseUniform); const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform); const auto& scales_b = shader.AddInput("scales_b", ShaderUsage::UseUniform); @@ -34,17 +33,21 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - return WGSL_TEMPLATE_APPLY( - shader, "quantization/dp4a_matmul.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(block_size, block_size_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_PARAMETER(tile_size_m, tile_size_m_), - WGSL_TEMPLATE_PARAMETER(tile_size_n, tile_size_n_), WGSL_TEMPLATE_PARAMETER(workgroup_size, WorkgroupSizeX()), - WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(block_size, block_size_), + WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), + WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), + WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), + WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(b, b), + WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), + WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } // scale_A components = 1, b components = 4, output components = 1 @@ -64,24 +67,28 @@ Status DP4AMatMulNBitsKParallelProgram::GenerateShaderCode(ShaderHelper& shader) } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, - "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); + ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); const uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec_; ORT_ENFORCE(tile_size_ % sub_tile_count == 0, "tile_size_ must be divisible by sub_tile_count"); - return WGSL_TEMPLATE_APPLY( - shader, "quantization/dp4a_matmul_k_parallel.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), - WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_k_parallel.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), + WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), + WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), + WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), + WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), + WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), + WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(b, b), + WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), + WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, @@ -176,26 +183,19 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor return context.RunProgram(mul_program); } - // The 64-lane subgroup variant spans 64 rows. Reserve 32-row tiles for - // devices whose subgroup range cannot select that variant. - const uint32_t tile_size_m = N < 64 && context.AdapterInfo().subgroupMaxSize < 64 ? 32 : 64; - const uint32_t tile_size_n = N < 64 ? 32 : 64; - const uint32_t tile_size_k_vec = N < 64 && nbits == 8 && block_size == 32 ? 8 : 2; + constexpr uint32_t kTileSize = 64; TensorShape reshaped_y_shape{batch_count, M, N / kVec4Components}; - uint32_t num_M_tile = CeilDiv(M, tile_size_m); - uint32_t num_N_tile = CeilDiv(N, tile_size_n); + uint32_t num_M_tile = (M + kTileSize - 1) / kTileSize; + uint32_t num_N_tile = (N + kTileSize - 1) / kTileSize; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsProgram mul_program{ - tile_size_m, tile_size_n, tile_size_k_vec, block_size, nbits, has_zero_points, - has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32}; - mul_program.SetWorkgroupSize(tile_size_m * tile_size_n / 16); + DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32}; + mul_program.SetWorkgroupSize(256); mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile); - mul_program - .AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, - {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, - {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, + {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, + {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({{static_cast(batch_count)}, {static_cast(M)}, {static_cast(N)}, @@ -207,11 +207,9 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), tile_size_m, tile_size_n, tile_size_k_vec, nbits, - has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32); + .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32); if (has_zero_points) { - mul_program.AddInput( - {zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); + mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } if (has_bias) { mul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); @@ -222,15 +220,23 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor return context.RunProgram(mul_program); } -bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level, - uint32_t block_size, uint32_t N, uint32_t K, uint32_t components_k, uint32_t M, - bool has_weight_idx_indirect, const Tensor* y) { +bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, + uint64_t accuracy_level, + uint32_t block_size, + uint32_t N, + uint32_t K, + uint32_t components_k, + uint32_t M, + bool has_weight_idx_indirect, + const Tensor* y) { // macOS - Avoid using dp4a on Metal, as it does not appear to have native dp4a support. // https://github.com/gpuweb/gpuweb/issues/2677#issuecomment-1713292226 // Use 'vendor' to check for metal; 'backend' is always WEBGPU when running under wasm. - bool use_dp4a = - context.HasFeature(wgpu::FeatureName::Subgroups) && context.AdapterInfo().vendor != std::string_view{"apple"}; - if (!(accuracy_level == 4 && block_size % 32 == 0 && components_k == 4 && K % 128 == 0 && N % 16 == 0 && use_dp4a)) { + bool use_dp4a = context.HasFeature(wgpu::FeatureName::Subgroups) && + context.AdapterInfo().vendor != std::string_view{"apple"}; + if (!(accuracy_level == 4 && block_size % 32 == 0 && + components_k == 4 && K % 128 == 0 && N % 16 == 0 && + use_dp4a)) { return false; } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index 9aa06c519ba62..e627fc7d60dcd 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -30,34 +30,32 @@ class DP4AMatMulQuantizeProgram final : public Program { public: - DP4AMatMulNBitsProgram(uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k_vec, uint32_t block_size, - uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, - bool has_weight_idx_indirect, bool is_qualcomm, bool acc_f32) - : Program{"DP4AMatMulNBits"}, - tile_size_m_(tile_size_m), - tile_size_n_(tile_size_n), - tile_size_k_vec_(tile_size_k_vec), - block_size_(block_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - is_qualcomm_(is_qualcomm), - acc_f32_(acc_f32) {} + DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, + bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, + bool acc_f32) : Program{"DP4AMatMulNBits"}, + block_size_(block_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + is_qualcomm_(is_qualcomm), + acc_f32_(acc_f32) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, - {"K8", ProgramUniformVariableDataType::Uint32}, {"K16", ProgramUniformVariableDataType::Uint32}, - {"num_M_tile", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, + {"K8", ProgramUniformVariableDataType::Uint32}, + {"K16", ProgramUniformVariableDataType::Uint32}, + {"num_M_tile", ProgramUniformVariableDataType::Uint32}, + {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, {"weight_idx", ProgramUniformVariableDataType::Uint32}); private: - uint32_t tile_size_m_; - uint32_t tile_size_n_; - uint32_t tile_size_k_vec_; uint32_t block_size_; uint32_t nbits_; bool has_bias_; @@ -70,28 +68,34 @@ class DP4AMatMulNBitsProgram final : public Program { class DP4AMatMulNBitsKParallelProgram final : public Program { public: - DP4AMatMulNBitsKParallelProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, bool has_zero_points, - bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, - bool single_scale_weights, bool broadcast_a_row = false, bool acc_f32 = false) - : Program{"DP4AMatMulNBitsKParallelProgram"}, - tile_size_k_vec_(tile_size_k_vec), - tile_size_(tile_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - single_scale_weights_(single_scale_weights), - broadcast_a_row_(broadcast_a_row), - acc_f32_(acc_f32) {} + DP4AMatMulNBitsKParallelProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, + bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, + bool broadcast_a_row = false, + bool acc_f32 = false) : Program{"DP4AMatMulNBitsKParallelProgram"}, + tile_size_k_vec_(tile_size_k_vec), + tile_size_(tile_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + single_scale_weights_(single_scale_weights), + broadcast_a_row_(broadcast_a_row), + acc_f32_(acc_f32) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, - {"K16", ProgramUniformVariableDataType::Uint32}, {"K32", ProgramUniformVariableDataType::Uint32}, - {"block_size", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, + {"K16", ProgramUniformVariableDataType::Uint32}, + {"K32", ProgramUniformVariableDataType::Uint32}, + {"block_size", ProgramUniformVariableDataType::Uint32}, + {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}, {"dispatch_M", ProgramUniformVariableDataType::Uint32}); + {"weight_idx", ProgramUniformVariableDataType::Uint32}, + {"dispatch_M", ProgramUniformVariableDataType::Uint32}); private: uint32_t tile_size_k_vec_; diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 09e9cf80247ec..b32bf0fb80c0b 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -1533,22 +1533,19 @@ TEST(MatMul2Bits, Float32_2b_BlkLen32_Accuracy0) { } #ifdef USE_WEBGPU -TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int2) { - // Tiled row/column tails and the narrow-output K-parallel dispatch. - for (const auto& [n, k] : {std::pair{16, 256}, {48, 256}, {80, 256}, {32, 2048}}) { - for (bool has_zero_point : {false, true}) { - TestOptions2Bits opts{}; - opts.M = 65; - opts.N = n; - opts.K = k; - opts.accuracy_level = 4; - opts.has_zero_point = has_zero_point; - opts.has_bias = true; - opts.output_abs_error = k == 256 ? 0.05f : 0.1f; - opts.output_rel_error = 0.02f; - RunTest2Bits(opts); - RunTest2Bits(opts); - } +TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int2) { + for (bool has_zero_point : {false, true}) { + TestOptions2Bits opts{}; + opts.M = 65; + opts.N = 32; + opts.K = 2048; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + RunTest2Bits(opts); + RunTest2Bits(opts); } } #endif diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index c2a59a83aa214..1a02ae1447546 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -453,23 +453,20 @@ TEST(MatMulNBits, Float32_4b_Accuracy0) { } #ifdef USE_WEBGPU -TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int4) { - // Tiled row/column tails and the narrow-output K-parallel dispatch. - for (const auto& [n, k] : {std::pair{16, 256}, {48, 256}, {80, 256}, {32, 2048}}) { - for (bool has_zero_point : {false, true}) { - TestOptions opts{}; - opts.batch_count = 2; - opts.M = 65; - opts.N = n; - opts.K = k; - opts.accuracy_level = 4; - opts.has_zero_point = has_zero_point; - opts.has_bias = true; - opts.output_abs_error = k == 256 ? 0.05f : 0.1f; - opts.output_rel_error = 0.02f; - RunTest(opts); - RunTest(opts); - } +TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int4) { + for (bool has_zero_point : {false, true}) { + TestOptions opts{}; + opts.batch_count = 2; + opts.M = 65; + opts.N = 32; + opts.K = 2048; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + RunTest(opts); + RunTest(opts); } } #endif diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index 864a1917079bb..8d64d6aef7bf6 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -1059,24 +1059,24 @@ TEST(MatMulNBits, BFloat16_Int8_SmallMFallback) { #endif #if defined(USE_WEBGPU) -TEST(MatMulNBits, WebGpu_DP4A_NarrowN_Int8) { - // Small-M reduction, the tiled threshold, and row/column tails in both tile widths. - for (int64_t k : {256, 1920, 2048, 4096}) { - for (const auto& [m, n] : {std::pair{3, 16}, {4, 32}, {65, 48}, {65, 80}}) { - for (bool has_zero_point : {false, true}) { - TestOptions8Bits opts{}; - opts.M = m; - opts.N = n; - opts.K = k; - opts.accuracy_level = 4; - opts.has_zero_point = has_zero_point; - opts.has_bias = true; - // Activation-quantization error for these Gaussian inputs grows with sqrt(K). - opts.output_abs_error = 0.05f * std::sqrt(static_cast(k) / 256.0f); - opts.output_rel_error = 0.02f; - RunTest8Bits(opts); - RunTest8Bits(opts); - } +TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int8) { + // Small-M K tails, both sides of the K-parallel threshold, and both output widths. + constexpr int64_t shapes[][3] = { + {3, 16, 1920}, {4, 32, 1920}, {4, 32, 2048}, {65, 16, 4096}, {65, 32, 4096}}; + for (const auto& [m, n, k] : shapes) { + for (bool has_zero_point : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m; + opts.N = n; + opts.K = k; + opts.accuracy_level = 4; + opts.has_zero_point = has_zero_point; + opts.has_bias = true; + // Activation-quantization error for these Gaussian inputs grows with sqrt(K). + opts.output_abs_error = 0.05f * std::sqrt(static_cast(k) / 256.0f); + opts.output_rel_error = 0.02f; + RunTest8Bits(opts); + RunTest8Bits(opts); } } } From afa23fa00ab33b22e8929e3ba2b64386695184bf Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Oct 2026 16:13:47 +0800 Subject: [PATCH 4/6] Ensure WebGPU DP4A regression tests execute on WebGPU --- .../test/contrib_ops/matmul_2bits_test.cc | 31 +++++++++++++++-- .../test/contrib_ops/matmul_4bits_test.cc | 16 +++++++-- .../test/contrib_ops/matmul_8bits_test.cc | 33 ++++++++++++++++--- 3 files changed, 70 insertions(+), 10 deletions(-) diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index b32bf0fb80c0b..6677a39cd8f9c 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -61,6 +61,7 @@ struct TestOptions2Bits { // Run the node on the CUDA EP instead of the default (CPU / WebGPU) set. bool use_cuda{false}; + bool disable_cpu_ep_fallback{false}; std::optional output_abs_error{}; std::optional output_rel_error{}; @@ -79,7 +80,8 @@ struct TestOptions2Bits { } template -void RunTest2Bits(const TestOptions2Bits& opts) { +void RunTest2Bits(const TestOptions2Bits& opts, + std::vector>&& explicit_eps = {}) { SCOPED_TRACE(opts); ASSERT_FALSE(opts.has_typed_zero_point && !opts.has_zero_point); @@ -230,6 +232,19 @@ void RunTest2Bits(const TestOptions2Bits& opts) { test.SetOutputRelErr("Y", *opts.output_rel_error); } + if (opts.disable_cpu_ep_fallback) { + SessionOptions session_options; + session_options.use_per_session_threads = false; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + } + + if (!explicit_eps.empty()) { + test.ConfigEps(std::move(explicit_eps)); + test.RunWithConfig(); + return; + } + std::vector> execution_providers; if (opts.use_cuda) { #ifdef USE_CUDA @@ -1540,12 +1555,22 @@ TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int2) { opts.N = 32; opts.K = 2048; opts.accuracy_level = 4; + opts.disable_cpu_ep_fallback = true; opts.has_zero_point = has_zero_point; opts.has_bias = true; opts.output_abs_error = 0.1f; opts.output_rel_error = 0.02f; - RunTest2Bits(opts); - RunTest2Bits(opts); + for (bool fp16 : {true, false}) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) GTEST_SKIP() << "WebGPU unavailable."; + std::vector> eps; + eps.push_back(std::move(ep)); + if (fp16) { + RunTest2Bits(opts, std::move(eps)); + } else { + RunTest2Bits(opts, std::move(eps)); + } + } } } #endif diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 1a02ae1447546..0d1f4b9b79f93 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -134,7 +134,7 @@ void RunTest(const TestOptions& opts, "unexpected type for T1"); #ifdef USE_CUDA - if (opts.accuracy_level != 0 && !opts.legacy_shape) { + if (explicit_eps.empty() && opts.accuracy_level != 0 && !opts.legacy_shape) { return; // CUDA EP does not handle accuracy level, so only test one level to avoid unnecessary tests. } #endif @@ -461,12 +461,22 @@ TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int4) { opts.N = 32; opts.K = 2048; opts.accuracy_level = 4; + opts.disable_cpu_ep_fallback = true; opts.has_zero_point = has_zero_point; opts.has_bias = true; opts.output_abs_error = 0.1f; opts.output_rel_error = 0.02f; - RunTest(opts); - RunTest(opts); + for (bool fp16 : {true, false}) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) GTEST_SKIP() << "WebGPU unavailable."; + std::vector> eps; + eps.push_back(std::move(ep)); + if (fp16) { + RunTest(opts, std::move(eps)); + } else { + RunTest(opts, std::move(eps)); + } + } } } #endif diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index 8d64d6aef7bf6..cc6c42219a455 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -58,6 +58,7 @@ struct TestOptions8Bits { bool zp_is_typed{false}; // true: zero_points use same type as scales (T); false: uint8 bool has_g_idx{false}; bool has_bias{false}; + bool disable_cpu_ep_fallback{false}; // When set, RunTest8Bits validates cross-session sharing of the pre-packed weights instead of // doing a single run. The model is run in two CPU sessions that use the same container. @@ -78,11 +79,12 @@ struct TestOptions8Bits { } template -void RunTest8Bits(const TestOptions8Bits& opts) { +void RunTest8Bits(const TestOptions8Bits& opts, + std::vector>&& explicit_eps = {}) { SCOPED_TRACE(opts); #ifdef USE_CUDA - if (opts.accuracy_level != 0) { + if (explicit_eps.empty() && opts.accuracy_level != 0) { return; // CUDA EP does not handle accuracy level, so only test one level to avoid unnecessary tests. } #endif @@ -241,6 +243,19 @@ void RunTest8Bits(const TestOptions8Bits& opts) { return; } + if (opts.disable_cpu_ep_fallback) { + SessionOptions session_options; + session_options.use_per_session_threads = false; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + } + + if (!explicit_eps.empty()) { + test.ConfigEps(std::move(explicit_eps)); + test.RunWithConfig(); + return; + } + std::vector> execution_providers; #ifdef USE_CUDA execution_providers.emplace_back(DefaultCudaExecutionProvider()); @@ -1070,13 +1085,23 @@ TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int8) { opts.N = n; opts.K = k; opts.accuracy_level = 4; + opts.disable_cpu_ep_fallback = true; opts.has_zero_point = has_zero_point; opts.has_bias = true; // Activation-quantization error for these Gaussian inputs grows with sqrt(K). opts.output_abs_error = 0.05f * std::sqrt(static_cast(k) / 256.0f); opts.output_rel_error = 0.02f; - RunTest8Bits(opts); - RunTest8Bits(opts); + for (bool fp16 : {true, false}) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) GTEST_SKIP() << "WebGPU unavailable."; + std::vector> eps; + eps.push_back(std::move(ep)); + if (fp16) { + RunTest8Bits(opts, std::move(eps)); + } else { + RunTest8Bits(opts, std::move(eps)); + } + } } } } From 573e54ba71355ba54f1096cd4432783216d4f652 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Oct 2026 16:39:00 +0800 Subject: [PATCH 5/6] Keep narrow WebGPU DP4A regressions FP32 only --- .../test/contrib_ops/matmul_2bits_test.cc | 16 +++++----------- .../test/contrib_ops/matmul_4bits_test.cc | 16 +++++----------- .../test/contrib_ops/matmul_8bits_test.cc | 16 +++++----------- 3 files changed, 15 insertions(+), 33 deletions(-) diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 6677a39cd8f9c..465327410bac9 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -1560,17 +1560,11 @@ TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int2) { opts.has_bias = true; opts.output_abs_error = 0.1f; opts.output_rel_error = 0.02f; - for (bool fp16 : {true, false}) { - auto ep = DefaultWebGpuExecutionProvider(); - if (!ep) GTEST_SKIP() << "WebGPU unavailable."; - std::vector> eps; - eps.push_back(std::move(ep)); - if (fp16) { - RunTest2Bits(opts, std::move(eps)); - } else { - RunTest2Bits(opts, std::move(eps)); - } - } + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) GTEST_SKIP() << "WebGPU unavailable."; + std::vector> eps; + eps.push_back(std::move(ep)); + RunTest2Bits(opts, std::move(eps)); } } #endif diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 0d1f4b9b79f93..f9a2df44cc836 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -466,17 +466,11 @@ TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int4) { opts.has_bias = true; opts.output_abs_error = 0.1f; opts.output_rel_error = 0.02f; - for (bool fp16 : {true, false}) { - auto ep = DefaultWebGpuExecutionProvider(); - if (!ep) GTEST_SKIP() << "WebGPU unavailable."; - std::vector> eps; - eps.push_back(std::move(ep)); - if (fp16) { - RunTest(opts, std::move(eps)); - } else { - RunTest(opts, std::move(eps)); - } - } + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) GTEST_SKIP() << "WebGPU unavailable."; + std::vector> eps; + eps.push_back(std::move(ep)); + RunTest(opts, std::move(eps)); } } #endif diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index cc6c42219a455..a0b5d18e65bc7 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -1091,17 +1091,11 @@ TEST(MatMulNBits, WebGpu_DP4A_KParallel_Int8) { // Activation-quantization error for these Gaussian inputs grows with sqrt(K). opts.output_abs_error = 0.05f * std::sqrt(static_cast(k) / 256.0f); opts.output_rel_error = 0.02f; - for (bool fp16 : {true, false}) { - auto ep = DefaultWebGpuExecutionProvider(); - if (!ep) GTEST_SKIP() << "WebGPU unavailable."; - std::vector> eps; - eps.push_back(std::move(ep)); - if (fp16) { - RunTest8Bits(opts, std::move(eps)); - } else { - RunTest8Bits(opts, std::move(eps)); - } - } + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) GTEST_SKIP() << "WebGPU unavailable."; + std::vector> eps; + eps.push_back(std::move(ep)); + RunTest8Bits(opts, std::move(eps)); } } } From a541ed1e65b149a97c2688cd3e97dc3d46599403 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Oct 2026 17:07:23 +0800 Subject: [PATCH 6/6] Test FP16 overflow in narrow WebGPU DP4A accumulation --- .../webgpu/contrib_ops/matmul_nbits_test.cc | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 onnxruntime/test/providers/webgpu/contrib_ops/matmul_nbits_test.cc diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_nbits_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_nbits_test.cc new file mode 100644 index 0000000000000..7c805305f9a9e --- /dev/null +++ b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_nbits_test.cc @@ -0,0 +1,112 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS) + +#include +#include +#include + +#include "gtest/gtest.h" + +#include "core/common/inlined_containers.h" +#include "core/providers/webgpu/webgpu_context.h" +#include "core/providers/webgpu/webgpu_provider_options.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { +namespace { + +void RunFloat16NarrowDP4AAccumulationTest(int64_t bits) { + ConfigOptions config_options; + ASSERT_STATUS_OK(config_options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, + webgpu::options::kEnableMatmulFp32Accumulation_OFF)); + auto provider = WebGpuExecutionProviderWithOptions(config_options); + if (!provider) { + GTEST_SKIP() << "WebGPU execution provider is unavailable."; + } + const auto& context = webgpu::WebGpuContextFactory::GetContext(provider->GetDeviceId()); + if (!context.DeviceHasFeature(wgpu::FeatureName::ShaderF16)) { + GTEST_SKIP() << "WebGPU adapter does not support shader-f16."; + } + // Match DP4A dispatch: Apple adapters deliberately use other MatMulNBits paths. + if (!context.DeviceHasFeature(wgpu::FeatureName::Subgroups) || + context.AdapterInfo().vendor == std::string_view{"apple"}) { + GTEST_SKIP() << "WebGPU DP4A path is unavailable."; + } + + // N=32 excludes subgroup-matrix dispatch. M=4 and K=4096 select the narrow + // K-parallel DP4A path, including its unconditional FP32 accumulation. + constexpr int64_t M = 4, N = 32, K = 4096, block_size = 32; + constexpr int64_t blocks = K / block_size; + const int64_t values_per_byte = 8 / bits; + const int64_t blob_size = block_size / values_per_byte; + const int32_t zero_point = 1 << (bits - 1); + InlinedVector a(M * K); + InlinedVector packed_weights(N * blocks * blob_size, 0); + InlinedVector scales(N * blocks, MLFloat16(16.0f)); + InlinedVector expected(M * N); + + // +/-127 quantizes exactly with activation scale 1. Each 32-element block + // contributes +/-65024, which fits FP16, but the first two K iterations sum + // to +/-130048. Later blocks cancel that sum; one omitted negative weight + // leaves exactly +/-2032. FP16 partial sums overflow before cancellation. + for (int64_t m = 0; m < M; ++m) { + const float sign = m % 2 == 0 ? 1.0f : -1.0f; + for (int64_t k = 0; k < K; ++k) { + a[m * K + k] = MLFloat16(sign * 127.0f); + } + for (int64_t n = 0; n < N; ++n) { + expected[m * N + n] = MLFloat16(sign * 2032.0f); + } + } + for (int64_t n = 0; n < N; ++n) { + for (int64_t k = 0; k < K; ++k) { + const int32_t weight = k == K - 1 ? 0 : (k < K / 2 ? 1 : -1); + const auto quantized = static_cast(zero_point + weight); + packed_weights[(n * K + k) / values_per_byte] |= + static_cast(quantized << ((k % values_per_byte) * bits)); + } + } + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("bits", bits); + test.AddAttribute("block_size", block_size); + test.AddAttribute("accuracy_level", 4); + test.AddInput("A", {M, K}, a.data(), a.size()); + test.AddInput("B", {N, blocks, blob_size}, packed_weights.data(), packed_weights.size(), true); + test.AddInput("scales", {N, blocks}, scales.data(), scales.size(), true); + test.AddOutput("Y", {M, N}, expected.data(), expected.size()); + test.SetOutputAbsErr("Y", 0.0f); + test.SetOutputRelErr("Y", 0.0f); + + SessionOptions options; + options.use_per_session_threads = false; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(options).ConfigEp(std::move(provider)).RunWithConfig(); +} + +} // namespace + +TEST(WebGpuMatMulNBitsTest, Float16_Int2_NarrowDP4AAccumulation) { + RunFloat16NarrowDP4AAccumulationTest(2); +} + +TEST(WebGpuMatMulNBitsTest, Float16_Int4_NarrowDP4AAccumulation) { + RunFloat16NarrowDP4AAccumulationTest(4); +} + +TEST(WebGpuMatMulNBitsTest, Float16_Int8_NarrowDP4AAccumulation) { + RunFloat16NarrowDP4AAccumulationTest(8); +} + +} // namespace test +} // namespace onnxruntime + +#endif // !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS)