From 80d8931f1d835fe29d89cced25f05be6690f81b6 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 14 Jan 2025 00:20:42 +0800 Subject: [PATCH] [webgpu] Use subgroup for matmulnbits (#23224) ### Description This PR applies subgroup to implement matmulnbits when tile_m > 1 for intel devices. With this PR, prefill for 500 tokens prompt for phi3 becomes 3.5s from 8.5s on intel Meteor Lake. --- .../webgpu/quantization/matmul_nbits.cc | 402 ++++++++++++------ .../webgpu/quantization/matmul_nbits.h | 14 +- .../core/providers/webgpu/compute_context.h | 3 + .../core/providers/webgpu/shader_helper.cc | 5 + 4 files changed, 295 insertions(+), 129 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 8abcd78bff..28e4ccec09 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -65,25 +65,9 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { const uint32_t tile_size = WorkgroupSizeX() * components_b_ * 8; // each uint32 has 8 data. const uint32_t a_length_per_tile = tile_size / a.NumComponents(); const uint32_t blocks_per_tile = tile_size / block_size_; - if (tile_m_ == 1) { - shader.AdditionalImplementation() << "fn mm_readA(batch : u32, row : u32, col : u32) -> input_a_value_t {\n" - " if (col < uniforms.input_a_shape[2]) {\n" - << " return " << a.GetByIndices("input_a_indices_t(batch, row, col)") << ";\n" - << " } else {\n" - " return input_a_value_t(0);\n" - " }\n" - "}\n" - << "var sub_a: array;\n" - << "var inter_results: array, " << WorkgroupSizeY() << ">;\n"; - std::string offset = "workgroup_idx * " + std::to_string(WorkgroupSizeY()); - shader.MainFunctionBody() << " let output_indices = " << y.OffsetToIndices(offset) << ";\n" - << " let col = output_indices[2];\n" - " let row = output_indices[1];\n" - " let batch = output_indices[0];\n"; - } else { - ORT_ENFORCE(tile_m_ < WorkgroupSizeY(), "tile_m must be less than or equal to WorkgroupSizeY."); - ORT_ENFORCE(WorkgroupSizeX() == WorkgroupSizeY(), "WorkgroupSizeX must be equal to WorkgroupSizeY."); - + if (tile_m_ > 1 && use_subgroup_) { + ORT_ENFORCE(a.NumComponents() == 4, "input a's components must be equal to 4."); + ORT_ENFORCE(components_b_ == 4, "input b's components must be equal to 4."); shader.AdditionalImplementation() << "fn mm_readA(batch : u32, row : u32, col : u32) -> input_a_value_t {\n" " if (row < uniforms.input_a_shape[1] && col < uniforms.input_a_shape[2]) {\n" << " return " << a.GetByIndices("input_a_indices_t(batch, row, col)") << ";\n" @@ -91,125 +75,297 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { " return input_a_value_t(0);\n" " }\n" "}\n" - << "var sub_a: array," << tile_m_ << ">;\n" + << "var sub_b: array, " << WorkgroupSizeY() << ">;\n" + << "var sub_scale: array, " << WorkgroupSizeY() << ">;\n" << "var inter_results: array, " << WorkgroupSizeY() << ">," << tile_m_ << ">;\n"; shader.MainFunctionBody() << " let col = workgroup_id.x * " << WorkgroupSizeY() << ";\n" << " let row = workgroup_id.y * " << tile_m_ << ";\n" << " let batch = workgroup_id.z;\n"; - } - shader.MainFunctionBody() << " let n_blocks_per_col = uniforms.input_b_shape[1];\n" - << " let num_tiles = (n_blocks_per_col - 1) / " << blocks_per_tile << " + 1;\n" - // Loop over shared dimension. - << " for (var tile: u32 = 0; tile < num_tiles; tile += 1) {\n" - << " let a_col_start = tile * " << a_length_per_tile << ";\n" - << " // load one tile A data into shared memory.\n" - << " for (var a_offset = local_idx; a_offset < " << a_length_per_tile << "; a_offset += " << workgroup_size << ") {\n" - << " let a_col = a_col_start + a_offset;\n"; - if (tile_m_ == 1) { - shader.MainFunctionBody() << " sub_a[a_offset] = mm_readA(batch, row, a_col);\n"; - } else { + shader.MainFunctionBody() << " let n_blocks_per_col = uniforms.input_b_shape[1];\n" + << " let num_tiles = (n_blocks_per_col - 1) / " << blocks_per_tile << " + 1;\n" + // Loop over shared dimension. + << " for (var tile: u32 = 0; tile < num_tiles; tile += 1) {\n" + << " // load one tile B/scale data into shared memory.\n" + // Each thread processes one block. + " let b_col = col + local_id.y;\n" + << " let block = tile * " << blocks_per_tile << " + local_id.x;\n" + << " if (b_col < uniforms.input_b_shape[0] && block < n_blocks_per_col) {\n" + << " sub_b[local_id.y][local_id.x] = " << b.GetByIndices("input_b_indices_t(b_col, block, 0)") << ";\n" + << " sub_scale[local_id.y][local_id.x] = " << scales.GetByOffset("b_col * n_blocks_per_col + block") << ";\n" + << " } else {\n" + " sub_b[local_id.y][local_id.x] = input_b_value_t(0);\n" + " sub_scale[local_id.y][local_id.x] = output_value_t(0);\n" + " }\n" + " workgroupBarrier();\n" + << " var in_y = (local_idx % 32) / 4;\n" + " var in_x = (local_idx / 32) * 4 + local_idx % 4;\n" + << " var word_offset = (local_idx % 4) * " << block_size_ / a.NumComponents() << ";\n" + << " if (sg_size == 8u) {\n" + " in_y = local_idx % 8;\n" + " in_x = local_idx / 8;\n" + << " word_offset = 0u;\n" + " } else if (sg_size == 16u) {\n" + " in_y = (local_idx % 16) / 2;\n" + " in_x = (local_idx / 16) * 2 + local_idx % 2;\n" + << " word_offset = (local_idx % 2) * " << block_size_ / a.NumComponents() << ";\n" + << " } else if (sg_size == 32u) {\n" + " in_y = (local_idx % 32) / 4;\n" + " in_x = (local_idx / 32) * 4 + local_idx % 4;\n" + << " word_offset = (local_idx % 4) * " << block_size_ / a.NumComponents() << ";\n" + << " } else if (sg_size == 64u) {\n" + " in_y = local_idx / 8;\n" + " in_x = local_idx % 8;\n" + << " word_offset = (local_idx % 8) * " << block_size_ / a.NumComponents() << ";\n" + << " }\n"; + if (has_zero_points_) { + const auto& zero_points = shader.AddInput("zero_points", ShaderUsage::UseUniform); + shader.MainFunctionBody() << " let zero_point_bytes_per_col = (n_blocks_per_col + 1) / 2;\n" + " let zero_point_byte_count = b_col * zero_point_bytes_per_col + (block >> 0x1u);\n" + " let zero_point_word_index = zero_point_byte_count >> 0x2u;\n" + " let zero_point_byte_offset = zero_point_byte_count & 0x3u;\n" + " let zero_point_nibble_offset: u32 = block & 0x1u;\n" + " let zero_point_bits_offset = (zero_point_byte_offset << 3) + (zero_point_nibble_offset << 2);\n" + << " let zero_point_word = " << zero_points.GetByOffset("zero_point_word_index") << " >> zero_point_bits_offset;\n" + << " let zero_point = output_element_t((zero_point_word) & 0xFu);\n"; + } else { + // The default zero point is 8 for unsigned 4-bit quantization. + shader.MainFunctionBody() << " let zero_point = output_element_t(8.0);\n"; + } + shader.MainFunctionBody() << " let scale = sub_scale[in_y][in_x];\n" + " let b_data = sub_b[in_y][in_x];\n"; + shader.MainFunctionBody() << " let a_col_start = tile * " << a_length_per_tile << ";\n"; for (uint32_t i = 0; i < tile_m_; i++) { - shader.MainFunctionBody() << " sub_a[" << i << "][a_offset] = mm_readA(batch, row + " << i << ", a_col);\n"; + shader.MainFunctionBody() << " let a_data" << i << " = mm_readA(batch, row + " << i << ", a_col_start + local_idx);\n"; } - } - shader.MainFunctionBody() << " }\n" - " workgroupBarrier();\n" - // Each thread processes one block. - " let b_row = col + local_id.y;\n" - << " let block = tile * " << blocks_per_tile << " + local_id.x;\n"; - if (has_zero_points_) { - const auto& zero_points = shader.AddInput("zero_points", ShaderUsage::UseUniform); - shader.MainFunctionBody() << " let zero_point_bytes_per_col = (n_blocks_per_col + 1) / 2;\n" - " let zero_point_byte_count = b_row * zero_point_bytes_per_col + (block >> 0x1u);\n" - " let zero_point_word_index = zero_point_byte_count >> 0x2u;\n" - " let zero_point_byte_offset = zero_point_byte_count & 0x3u;\n" - " let zero_point_nibble_offset: u32 = block & 0x1u;\n" - " let zero_point_bits_offset = (zero_point_byte_offset << 3) + (zero_point_nibble_offset << 2);\n" - << " let zero_point_word = " << zero_points.GetByOffset("zero_point_word_index") << " >> zero_point_bits_offset;\n" - << " let zero_point = output_element_t((zero_point_word) & 0xFu);\n"; - } else { - // The default zero point is 8 for unsigned 4-bit quantization. - shader.MainFunctionBody() << " let zero_point = output_element_t(8.0);\n"; - } - shader.MainFunctionBody() << " var scale = output_element_t(0);\n" - " var b_data = input_b_value_t(0);\n" - << " if (block < n_blocks_per_col) {\n" - << " scale = " << scales.GetByOffset("b_row * n_blocks_per_col + block") << ";\n" - << " b_data = " << b.GetByIndices("input_b_indices_t(b_row, block, 0)") << ";\n" - << " }\n" - << " var word_offset = local_id.x * " << block_size_ / a.NumComponents() << ";\n" - << " for (var i: u32 = 0; i < " << components_b_ << "; i++) {\n"; - shader.MainFunctionBody() << " let b_value = b_data"; - if (components_b_ > 1) { - shader.MainFunctionBody() << "[i]"; - } - shader.MainFunctionBody() << ";\n" - " let b_value_lower = unpack4xU8(b_value & 0x0F0F0F0Fu);\n" - " let b_value_upper = unpack4xU8((b_value >> 4) & 0x0F0F0F0Fu);\n" - " let b_quantized_values = mat2x4(output_element_t(b_value_lower[0]), output_element_t(b_value_upper[0]), output_element_t(b_value_lower[1]), output_element_t(b_value_upper[1]), output_element_t(b_value_lower[2]), output_element_t(b_value_upper[2]), output_element_t(b_value_lower[3]), output_element_t(b_value_upper[3]));\n" - " let b_dequantized_values = (b_quantized_values - mat2x4("; - for (int i = 0; i < 8; i++) { - shader.MainFunctionBody() << "zero_point"; - if (i < 7) { - shader.MainFunctionBody() << ", "; - } - } - shader.MainFunctionBody() << ")) * scale;\n"; - if (tile_m_ == 1) { - switch (a.NumComponents()) { - case 1: - shader.MainFunctionBody() << " inter_results[local_id.y][local_id.x] += dot(vec4(sub_a[word_offset], sub_a[word_offset + 1], sub_a[word_offset + 2], sub_a[word_offset + 3]), b_dequantized_values[0]) + dot(vec4(sub_a[word_offset + 4], sub_a[word_offset + 5], sub_a[word_offset + 6], sub_a[word_offset + 7]), b_dequantized_values[1]);\n"; - break; - case 2: - shader.MainFunctionBody() << " inter_results[local_id.y][local_id.x] += dot(vec4(sub_a[word_offset], sub_a[word_offset + 1]), b_dequantized_values[0]) + dot(vec4(sub_a[word_offset + 2], sub_a[word_offset + 3]), b_dequantized_values[1]);\n"; - break; - case 4: - shader.MainFunctionBody() << " inter_results[local_id.y][local_id.x] += dot(sub_a[word_offset], b_dequantized_values[0]) + dot(sub_a[word_offset + 1], b_dequantized_values[1]);\n"; - break; - default: - break; - } - } else { + + shader.MainFunctionBody() << " if (sg_size == 8u) {\n"; + shader.MainFunctionBody() << " for (var i: u32 = 0; i < 4; i++) {\n"; + shader.MainFunctionBody() << " let b_value = b_data[i];\n" + " let b_value_lower = unpack4xU8(b_value & 0x0F0F0F0Fu);\n" + " let b_value_upper = unpack4xU8((b_value >> 4) & 0x0F0F0F0Fu);\n" + " let b_quantized_values = mat2x4(output_element_t(b_value_lower[0]), output_element_t(b_value_upper[0]), output_element_t(b_value_lower[1]), output_element_t(b_value_upper[1]), output_element_t(b_value_lower[2]), output_element_t(b_value_upper[2]), output_element_t(b_value_lower[3]), output_element_t(b_value_upper[3]));\n" + " let b_dequantized_values = (b_quantized_values - mat2x4(zero_point, zero_point, zero_point, zero_point, zero_point, zero_point, zero_point, zero_point)) * scale;\n"; for (uint32_t i = 0; i < tile_m_; i++) { + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a0 = subgroupShuffle(a_data" << i << ", i * 2);\n"; + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a1 = subgroupShuffle(a_data" << i << ", i * 2 + 1);\n"; + shader.MainFunctionBody() << " inter_results[" << i << "][in_y][in_x] += dot(a0, b_dequantized_values[0]) + dot(a1, b_dequantized_values[1]);\n"; + } + shader.MainFunctionBody() << " }\n"; + shader.MainFunctionBody() << " } else if (sg_size == 16u) {\n"; + shader.MainFunctionBody() << " for (var i: u32 = 0; i < 4; i++) {\n"; + shader.MainFunctionBody() << " let b_value = b_data[i];\n" + " let b_value_lower = unpack4xU8(b_value & 0x0F0F0F0Fu);\n" + " let b_value_upper = unpack4xU8((b_value >> 4) & 0x0F0F0F0Fu);\n" + " let b_quantized_values = mat2x4(output_element_t(b_value_lower[0]), output_element_t(b_value_upper[0]), output_element_t(b_value_lower[1]), output_element_t(b_value_upper[1]), output_element_t(b_value_lower[2]), output_element_t(b_value_upper[2]), output_element_t(b_value_lower[3]), output_element_t(b_value_upper[3]));\n" + " let b_dequantized_values = (b_quantized_values - mat2x4(zero_point, zero_point, zero_point, zero_point, zero_point, zero_point, zero_point, zero_point)) * scale;\n"; + for (uint32_t i = 0; i < tile_m_; i++) { + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a0 = subgroupShuffle(a_data" << i << ", i * 2);\n"; + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a00 = subgroupShuffle(a_data" << i << ", i * 2 + 8);\n"; + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a1 = subgroupShuffle(a_data" << i << ", i * 2 + 1);\n"; + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a11 = subgroupShuffle(a_data" << i << ", i * 2 + 9);\n"; + shader.MainFunctionBody() << " inter_results[" << i << "][in_y][in_x] += dot(select(a00, a0, local_idx % 2 == 0), b_dequantized_values[0]) + dot(select(a11, a1, local_idx % 2 == 0), b_dequantized_values[1]);\n"; + } + shader.MainFunctionBody() << " word_offset += " << 8 / a.NumComponents() << ";\n" + << " }\n"; + shader.MainFunctionBody() << " } else {\n"; + shader.MainFunctionBody() << " for (var i: u32 = 0; i < 4; i++) {\n"; + shader.MainFunctionBody() << " let b_value = b_data[i];\n" + " let b_value_lower = unpack4xU8(b_value & 0x0F0F0F0Fu);\n" + " let b_value_upper = unpack4xU8((b_value >> 4) & 0x0F0F0F0Fu);\n" + " let b_quantized_values = mat2x4(output_element_t(b_value_lower[0]), output_element_t(b_value_upper[0]), output_element_t(b_value_lower[1]), output_element_t(b_value_upper[1]), output_element_t(b_value_lower[2]), output_element_t(b_value_upper[2]), output_element_t(b_value_lower[3]), output_element_t(b_value_upper[3]));\n" + " let b_dequantized_values = (b_quantized_values - mat2x4(zero_point, zero_point, zero_point, zero_point, zero_point, zero_point, zero_point, zero_point)) * scale;\n"; + for (uint32_t i = 0; i < tile_m_; i++) { + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a0 = subgroupShuffle(a_data" << i << ", word_offset);\n"; + if (i == 0) { + shader.MainFunctionBody() << " var "; + } + shader.MainFunctionBody() << " a1 = subgroupShuffle(a_data" << i << ", word_offset + 1);\n"; + shader.MainFunctionBody() << " inter_results[" << i << "][in_y][in_x] += dot(a0, b_dequantized_values[0]) + dot(a1, b_dequantized_values[1]);\n"; + } + shader.MainFunctionBody() << " word_offset += " << 8 / a.NumComponents() << ";\n"; + shader.MainFunctionBody() << " }\n"; + shader.MainFunctionBody() << " }\n"; + + shader.MainFunctionBody() << " }\n"; + shader.MainFunctionBody() << " if (local_idx < " << WorkgroupSizeY() * tile_m_ << ") {\n" + << " let inner_row = local_idx / " << WorkgroupSizeY() << ";\n" + << " let inner_col = local_idx % " << WorkgroupSizeY() << ";\n" + << " var output_value = output_value_t(0);\n" + << " for (var b = 0u; b < " << WorkgroupSizeX() << "; b++) {\n" + << " output_value += inter_results[inner_row][inner_col][b];\n" + " }\n" + " if (row + inner_row < uniforms.output_shape[1] && col + inner_col < uniforms.output_shape[2]) {\n" + << " " << y.SetByIndices("output_indices_t(batch, row + inner_row, col + inner_col)", "output_value") << ";\n" + << " }\n" + " }\n"; + } else { + if (tile_m_ == 1) { + shader.AdditionalImplementation() << "fn mm_readA(batch : u32, row : u32, col : u32) -> input_a_value_t {\n" + " if (col < uniforms.input_a_shape[2]) {\n" + << " return " << a.GetByIndices("input_a_indices_t(batch, row, col)") << ";\n" + << " } else {\n" + " return input_a_value_t(0);\n" + " }\n" + "}\n" + << "var sub_a: array;\n" + << "var inter_results: array, " << WorkgroupSizeY() << ">;\n"; + std::string offset = "workgroup_idx * " + std::to_string(WorkgroupSizeY()); + shader.MainFunctionBody() << " let output_indices = " << y.OffsetToIndices(offset) << ";\n" + << " let col = output_indices[2];\n" + " let row = output_indices[1];\n" + " let batch = output_indices[0];\n"; + } else { + ORT_ENFORCE(tile_m_ < WorkgroupSizeY(), "tile_m must be less than or equal to WorkgroupSizeY."); + ORT_ENFORCE(WorkgroupSizeX() == WorkgroupSizeY(), "WorkgroupSizeX must be equal to WorkgroupSizeY."); + + shader.AdditionalImplementation() << "fn mm_readA(batch : u32, row : u32, col : u32) -> input_a_value_t {\n" + " if (row < uniforms.input_a_shape[1] && col < uniforms.input_a_shape[2]) {\n" + << " return " << a.GetByIndices("input_a_indices_t(batch, row, col)") << ";\n" + << " } else {\n" + " return input_a_value_t(0);\n" + " }\n" + "}\n" + << "var sub_a: array," << tile_m_ << ">;\n" + << "var inter_results: array, " << WorkgroupSizeY() << ">," << tile_m_ << ">;\n"; + shader.MainFunctionBody() << " let col = workgroup_id.x * " << WorkgroupSizeY() << ";\n" + << " let row = workgroup_id.y * " << tile_m_ << ";\n" + << " let batch = workgroup_id.z;\n"; + } + shader.MainFunctionBody() << " let n_blocks_per_col = uniforms.input_b_shape[1];\n" + << " let num_tiles = (n_blocks_per_col - 1) / " << blocks_per_tile << " + 1;\n" + // Loop over shared dimension. + << " for (var tile: u32 = 0; tile < num_tiles; tile += 1) {\n" + << " let a_col_start = tile * " << a_length_per_tile << ";\n" + << " // load one tile A data into shared memory.\n" + << " for (var a_offset = local_idx; a_offset < " << a_length_per_tile << "; a_offset += " << workgroup_size << ") {\n" + << " let a_col = a_col_start + a_offset;\n"; + if (tile_m_ == 1) { + shader.MainFunctionBody() << " sub_a[a_offset] = mm_readA(batch, row, a_col);\n"; + } else { + for (uint32_t i = 0; i < tile_m_; i++) { + shader.MainFunctionBody() << " sub_a[" << i << "][a_offset] = mm_readA(batch, row + " << i << ", a_col);\n"; + } + } + shader.MainFunctionBody() << " }\n" + " workgroupBarrier();\n" + // Each thread processes one block. + " let b_row = col + local_id.y;\n" + << " let block = tile * " << blocks_per_tile << " + local_id.x;\n"; + if (has_zero_points_) { + const auto& zero_points = shader.AddInput("zero_points", ShaderUsage::UseUniform); + shader.MainFunctionBody() << " let zero_point_bytes_per_col = (n_blocks_per_col + 1) / 2;\n" + " let zero_point_byte_count = b_row * zero_point_bytes_per_col + (block >> 0x1u);\n" + " let zero_point_word_index = zero_point_byte_count >> 0x2u;\n" + " let zero_point_byte_offset = zero_point_byte_count & 0x3u;\n" + " let zero_point_nibble_offset: u32 = block & 0x1u;\n" + " let zero_point_bits_offset = (zero_point_byte_offset << 3) + (zero_point_nibble_offset << 2);\n" + << " let zero_point_word = " << zero_points.GetByOffset("zero_point_word_index") << " >> zero_point_bits_offset;\n" + << " let zero_point = output_element_t((zero_point_word) & 0xFu);\n"; + } else { + // The default zero point is 8 for unsigned 4-bit quantization. + shader.MainFunctionBody() << " let zero_point = output_element_t(8.0);\n"; + } + shader.MainFunctionBody() << " var scale = output_element_t(0);\n" + " var b_data = input_b_value_t(0);\n" + << " if (block < n_blocks_per_col) {\n" + << " scale = " << scales.GetByOffset("b_row * n_blocks_per_col + block") << ";\n" + << " b_data = " << b.GetByIndices("input_b_indices_t(b_row, block, 0)") << ";\n" + << " }\n" + << " var word_offset = local_id.x * " << block_size_ / a.NumComponents() << ";\n" + << " for (var i: u32 = 0; i < " << components_b_ << "; i++) {\n"; + shader.MainFunctionBody() << " let b_value = b_data"; + if (components_b_ > 1) { + shader.MainFunctionBody() << "[i]"; + } + shader.MainFunctionBody() << ";\n" + " let b_value_lower = unpack4xU8(b_value & 0x0F0F0F0Fu);\n" + " let b_value_upper = unpack4xU8((b_value >> 4) & 0x0F0F0F0Fu);\n" + " let b_quantized_values = mat2x4(output_element_t(b_value_lower[0]), output_element_t(b_value_upper[0]), output_element_t(b_value_lower[1]), output_element_t(b_value_upper[1]), output_element_t(b_value_lower[2]), output_element_t(b_value_upper[2]), output_element_t(b_value_lower[3]), output_element_t(b_value_upper[3]));\n" + " let b_dequantized_values = (b_quantized_values - mat2x4("; + for (int i = 0; i < 8; i++) { + shader.MainFunctionBody() << "zero_point"; + if (i < 7) { + shader.MainFunctionBody() << ", "; + } + } + shader.MainFunctionBody() << ")) * scale;\n"; + if (tile_m_ == 1) { switch (a.NumComponents()) { case 1: - shader.MainFunctionBody() << " inter_results[" << i << "][local_id.y][local_id.x] += dot(vec4(sub_a[" << i << "][word_offset], sub_a[" << i << "][word_offset + 1], sub_a[" << i << "][word_offset + 2], sub_a[" << i << "][word_offset + 3]), b_dequantized_values[0]) + dot(vec4(sub_a[" << i << "][word_offset + 4], sub_a[" << i << "][word_offset + 5], sub_a[" << i << "][word_offset + 6], sub_a[" << i << "][word_offset + 7]), b_dequantized_values[1]);\n"; + shader.MainFunctionBody() << " inter_results[local_id.y][local_id.x] += dot(vec4(sub_a[word_offset], sub_a[word_offset + 1], sub_a[word_offset + 2], sub_a[word_offset + 3]), b_dequantized_values[0]) + dot(vec4(sub_a[word_offset + 4], sub_a[word_offset + 5], sub_a[word_offset + 6], sub_a[word_offset + 7]), b_dequantized_values[1]);\n"; break; case 2: - shader.MainFunctionBody() << " inter_results[" << i << "][local_id.y][local_id.x] += dot(vec4(sub_a[" << i << "][word_offset], sub_a[" << i << "][word_offset + 1]), b_dequantized_values[0]) + dot(vec4(sub_a[" << i << "][word_offset + 2], sub_a[" << i << "][word_offset + 3]), b_dequantized_values[1]);\n"; + shader.MainFunctionBody() << " inter_results[local_id.y][local_id.x] += dot(vec4(sub_a[word_offset], sub_a[word_offset + 1]), b_dequantized_values[0]) + dot(vec4(sub_a[word_offset + 2], sub_a[word_offset + 3]), b_dequantized_values[1]);\n"; break; case 4: - shader.MainFunctionBody() << " inter_results[" << i << "][local_id.y][local_id.x] += dot(sub_a[" << i << "][word_offset], b_dequantized_values[0]) + dot(sub_a[" << i << "][word_offset + 1], b_dequantized_values[1]);\n"; + shader.MainFunctionBody() << " inter_results[local_id.y][local_id.x] += dot(sub_a[word_offset], b_dequantized_values[0]) + dot(sub_a[word_offset + 1], b_dequantized_values[1]);\n"; break; default: break; } + } else { + for (uint32_t i = 0; i < tile_m_; i++) { + switch (a.NumComponents()) { + case 1: + shader.MainFunctionBody() << " inter_results[" << i << "][local_id.y][local_id.x] += dot(vec4(sub_a[" << i << "][word_offset], sub_a[" << i << "][word_offset + 1], sub_a[" << i << "][word_offset + 2], sub_a[" << i << "][word_offset + 3]), b_dequantized_values[0]) + dot(vec4(sub_a[" << i << "][word_offset + 4], sub_a[" << i << "][word_offset + 5], sub_a[" << i << "][word_offset + 6], sub_a[" << i << "][word_offset + 7]), b_dequantized_values[1]);\n"; + break; + case 2: + shader.MainFunctionBody() << " inter_results[" << i << "][local_id.y][local_id.x] += dot(vec4(sub_a[" << i << "][word_offset], sub_a[" << i << "][word_offset + 1]), b_dequantized_values[0]) + dot(vec4(sub_a[" << i << "][word_offset + 2], sub_a[" << i << "][word_offset + 3]), b_dequantized_values[1]);\n"; + break; + case 4: + shader.MainFunctionBody() << " inter_results[" << i << "][local_id.y][local_id.x] += dot(sub_a[" << i << "][word_offset], b_dequantized_values[0]) + dot(sub_a[" << i << "][word_offset + 1], b_dequantized_values[1]);\n"; + break; + default: + break; + } + } } - } - shader.MainFunctionBody() << " word_offset += " << 8 / a.NumComponents() << ";\n" - << " }\n" - " workgroupBarrier();\n" - " }\n"; - if (tile_m_ == 1) { - shader.MainFunctionBody() << " if (local_idx < " << WorkgroupSizeY() << ") {\n" - << " var output_value = output_value_t(0);\n" - << " for (var b = 0u; b < " << WorkgroupSizeX() << "; b++) {\n" - << " output_value += inter_results[local_idx][b];\n" - " }\n" - " if (col + local_idx < uniforms.output_shape[2]) {\n" - << " " << y.SetByIndices("output_indices_t(batch, row, col + local_idx)", "output_value") << ";\n" - << " }\n" - " }\n"; - } else { - shader.MainFunctionBody() << " if (local_id.y < " << tile_m_ << ") {\n" - << " var output_value = output_value_t(0);\n" - << " for (var b = 0u; b < " << WorkgroupSizeX() << "; b++) {\n" - << " output_value += inter_results[local_id.y][local_id.x][b];\n" - " }\n" - " if (row + local_id.y < uniforms.output_shape[1] && col + local_id.x < uniforms.output_shape[2]) {\n" - << " " << y.SetByIndices("output_indices_t(batch, row + local_id.y, col + local_id.x)", "output_value") << ";\n" + shader.MainFunctionBody() << " word_offset += " << 8 / a.NumComponents() << ";\n" << " }\n" + " workgroupBarrier();\n" " }\n"; + if (tile_m_ == 1) { + shader.MainFunctionBody() << " if (local_idx < " << WorkgroupSizeY() << ") {\n" + << " var output_value = output_value_t(0);\n" + << " for (var b = 0u; b < " << WorkgroupSizeX() << "; b++) {\n" + << " output_value += inter_results[local_idx][b];\n" + " }\n" + " if (col + local_idx < uniforms.output_shape[2]) {\n" + << " " << y.SetByIndices("output_indices_t(batch, row, col + local_idx)", "output_value") << ";\n" + << " }\n" + " }\n"; + } else { + shader.MainFunctionBody() << " if (local_id.y < " << tile_m_ << ") {\n" + << " var output_value = output_value_t(0);\n" + << " for (var b = 0u; b < " << WorkgroupSizeX() << "; b++) {\n" + << " output_value += inter_results[local_id.y][local_id.x][b];\n" + " }\n" + " if (row + local_id.y < uniforms.output_shape[1] && col + local_id.x < uniforms.output_shape[2]) {\n" + << " " << y.SetByIndices("output_indices_t(batch, row + local_id.y, col + local_id.x)", "output_value") << ";\n" + << " }\n" + " }\n"; + } } } else { const std::string quantized_data_type = QuantizedDataType(a.NumComponents()); @@ -413,7 +569,8 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context // TODO: Support output_number > 1. Some cases are failed when output_number > 1. constexpr uint32_t output_number = 1; const uint32_t tile_m = M > kMinMForTileOptimization ? 4 : 1; - MatMulNBitsProgram program{output_number, block_size, tile_m, gsl::narrow(components_b), has_zero_points}; + const bool use_subgroup = context.Device().HasFeature(wgpu::FeatureName::Subgroups) && context.AdapterInfo().vendor == std::string_view{"intel"} && components_a == 4 && block_size == 32; + MatMulNBitsProgram program{output_number, block_size, tile_m, gsl::narrow(components_b), has_zero_points, use_subgroup}; if (M > kMinMForTileOptimization && block_size == 32) { components = 1; constexpr uint32_t workgroup_size = 64; @@ -423,12 +580,11 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context program.SetDispatchGroupSize((N + workgroup_y - 1) / workgroup_y, (M + tile_m - 1) / tile_m, batch_count); - program.CacheHint("T_M" + std::to_string(tile_m)); + program.CacheHint("T_M" + std::to_string(tile_m) + "Subgroup" + std::to_string(use_subgroup)); } else if (block_size == 32) { components = 1; - constexpr uint32_t workgroup_size = 128; - const uint32_t workgroup_y = N % 8 == 0 ? 8 : N % 4 == 0 ? 4 - : 1; + constexpr uint32_t workgroup_size = 64; + const uint32_t workgroup_y = N % 8 == 0 ? 8 : 1; const uint32_t workgroup_x = workgroup_size / workgroup_y; program.SetWorkgroupSize(workgroup_x, workgroup_y, 1); program.SetDispatchGroupSize(data_size / components / workgroup_y); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index 57615d3dda..ca3c485566 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -14,12 +14,13 @@ using namespace onnxruntime::webgpu; class MatMulNBitsProgram final : public Program { public: - MatMulNBitsProgram(uint32_t output_number, uint32_t block_size, uint32_t tile_m, int components_b, bool has_zero_points) : Program{"MatMulNBits"}, - output_number_{output_number}, - block_size_{block_size}, - tile_m_{tile_m}, - components_b_{components_b}, - has_zero_points_{has_zero_points} { + MatMulNBitsProgram(uint32_t output_number, uint32_t block_size, uint32_t tile_m, int components_b, bool has_zero_points, bool use_subgroup) : Program{"MatMulNBits"}, + output_number_{output_number}, + block_size_{block_size}, + tile_m_{tile_m}, + components_b_{components_b}, + has_zero_points_{has_zero_points}, + use_subgroup_(use_subgroup) { } Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -31,6 +32,7 @@ class MatMulNBitsProgram final : public Program { uint32_t tile_m_; int components_b_; bool has_zero_points_; + bool use_subgroup_; }; class MatMulNBits final : public WebGpuKernel { diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index b7ea8a58e2..2fbd639b27 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/onnxruntime/core/providers/webgpu/compute_context.h @@ -41,6 +41,9 @@ class ComputeContext { inline const wgpu::Limits& DeviceLimits() const { return webgpu_context_.DeviceLimits(); } + inline const wgpu::Device& Device() const { + return webgpu_context_.Device(); + } // // Get the kernel context. diff --git a/onnxruntime/core/providers/webgpu/shader_helper.cc b/onnxruntime/core/providers/webgpu/shader_helper.cc index 5685494556..49c9a84f69 100644 --- a/onnxruntime/core/providers/webgpu/shader_helper.cc +++ b/onnxruntime/core/providers/webgpu/shader_helper.cc @@ -63,6 +63,11 @@ Status ShaderHelper::Init() { " @builtin(workgroup_id) workgroup_id : vec3,\n" " @builtin(local_invocation_index) local_idx : u32,\n" " @builtin(local_invocation_id) local_id : vec3"; + if (device_.HasFeature(wgpu::FeatureName::Subgroups)) { + body_ss_ << ",\n" + " @builtin(subgroup_invocation_id) sg_id : u32,\n" + " @builtin(subgroup_size) sg_size : u32"; + } if (!is_1d_dispatch) { body_ss_ << ",\n" " @builtin(num_workgroups) num_workgroups : vec3";