mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
[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.
This commit is contained in:
Jiajia Qin
GitHub
parent
73f5b0c597
commit
80d8931f1d
4 changed files with 295 additions and 129 deletions
|
|
@ -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<workgroup> sub_a: array<input_a_value_t, " << a_length_per_tile << ">;\n"
|
||||
<< "var<workgroup> inter_results: array<array<output_value_t, " << WorkgroupSizeX() << ">, " << 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<workgroup> sub_a: array<array<input_a_value_t, " << a_length_per_tile << ">," << tile_m_ << ">;\n"
|
||||
<< "var<workgroup> sub_b: array<array<input_b_value_t, " << WorkgroupSizeX() << ">, " << WorkgroupSizeY() << ">;\n"
|
||||
<< "var<workgroup> sub_scale: array<array<output_value_t, " << WorkgroupSizeX() << ">, " << WorkgroupSizeY() << ">;\n"
|
||||
<< "var<workgroup> inter_results: array<array<array<output_value_t, " << WorkgroupSizeX() << ">, " << 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>(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<output_element_t>(";
|
||||
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<output_element_t>(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<output_element_t>(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<output_element_t>(sub_a[word_offset], sub_a[word_offset + 1]), b_dequantized_values[0]) + dot(vec4<output_element_t>(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>(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<output_element_t>(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>(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<output_element_t>(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>(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<output_element_t>(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<workgroup> sub_a: array<input_a_value_t, " << a_length_per_tile << ">;\n"
|
||||
<< "var<workgroup> inter_results: array<array<output_value_t, " << WorkgroupSizeX() << ">, " << 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<workgroup> sub_a: array<array<input_a_value_t, " << a_length_per_tile << ">," << tile_m_ << ">;\n"
|
||||
<< "var<workgroup> inter_results: array<array<array<output_value_t, " << WorkgroupSizeX() << ">, " << 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>(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<output_element_t>(";
|
||||
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<output_element_t>(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<output_element_t>(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<output_element_t>(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<output_element_t>(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<output_element_t>(sub_a[" << i << "][word_offset], sub_a[" << i << "][word_offset + 1]), b_dequantized_values[0]) + dot(vec4<output_element_t>(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<output_element_t>(sub_a[word_offset], sub_a[word_offset + 1]), b_dequantized_values[0]) + dot(vec4<output_element_t>(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<output_element_t>(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<output_element_t>(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<output_element_t>(sub_a[" << i << "][word_offset], sub_a[" << i << "][word_offset + 1]), b_dequantized_values[0]) + dot(vec4<output_element_t>(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<int>(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<int>(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);
|
||||
|
|
|
|||
|
|
@ -14,12 +14,13 @@ using namespace onnxruntime::webgpu;
|
|||
|
||||
class MatMulNBitsProgram final : public Program<MatMulNBitsProgram> {
|
||||
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<MatMulNBitsProgram> {
|
|||
uint32_t tile_m_;
|
||||
int components_b_;
|
||||
bool has_zero_points_;
|
||||
bool use_subgroup_;
|
||||
};
|
||||
|
||||
class MatMulNBits final : public WebGpuKernel {
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -63,6 +63,11 @@ Status ShaderHelper::Init() {
|
|||
" @builtin(workgroup_id) workgroup_id : vec3<u32>,\n"
|
||||
" @builtin(local_invocation_index) local_idx : u32,\n"
|
||||
" @builtin(local_invocation_id) local_id : vec3<u32>";
|
||||
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<u32>";
|
||||
|
|
|
|||
Loading…
Reference in a new issue