ggml-webgpu: add mulmat with overlapping src0/src1 (e.g., for minimax-01) (#27321)
This commit is contained in:
@@ -954,10 +954,11 @@ struct ggml_webgpu_mul_mat_vec_pipeline_key {
|
||||
int vectorized;
|
||||
uint32_t num_cols;
|
||||
bool use_mmvq;
|
||||
bool src_overlap;
|
||||
|
||||
bool operator==(const ggml_webgpu_mul_mat_vec_pipeline_key & other) const {
|
||||
return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized &&
|
||||
num_cols == other.num_cols && use_mmvq == other.use_mmvq;
|
||||
num_cols == other.num_cols && use_mmvq == other.use_mmvq && src_overlap == other.src_overlap;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -969,6 +970,7 @@ struct ggml_webgpu_mul_mat_vec_pipeline_key_hash {
|
||||
ggml_webgpu_hash_combine(seed, key.vectorized);
|
||||
ggml_webgpu_hash_combine(seed, key.num_cols);
|
||||
ggml_webgpu_hash_combine(seed, key.use_mmvq);
|
||||
ggml_webgpu_hash_combine(seed, key.src_overlap);
|
||||
return seed;
|
||||
}
|
||||
};
|
||||
@@ -977,6 +979,7 @@ struct ggml_webgpu_mul_mat_vec_shader_decisions {
|
||||
uint32_t wg_size;
|
||||
uint32_t outputs_per_wg;
|
||||
uint32_t vec_size;
|
||||
bool src_overlap = false;
|
||||
};
|
||||
|
||||
struct ggml_webgpu_quantize_q8_pipeline_key {
|
||||
@@ -998,10 +1001,11 @@ struct ggml_webgpu_mul_mat_pipeline_key {
|
||||
ggml_type src1_type;
|
||||
int vectorized;
|
||||
int use_subgroup_matrix;
|
||||
bool src_overlap;
|
||||
|
||||
bool operator==(const ggml_webgpu_mul_mat_pipeline_key & other) const {
|
||||
return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized &&
|
||||
use_subgroup_matrix == other.use_subgroup_matrix;
|
||||
use_subgroup_matrix == other.use_subgroup_matrix && src_overlap == other.src_overlap;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -1012,6 +1016,7 @@ struct ggml_webgpu_mul_mat_pipeline_key_hash {
|
||||
ggml_webgpu_hash_combine(seed, key.src1_type);
|
||||
ggml_webgpu_hash_combine(seed, key.vectorized);
|
||||
ggml_webgpu_hash_combine(seed, key.use_subgroup_matrix);
|
||||
ggml_webgpu_hash_combine(seed, key.src_overlap);
|
||||
return seed;
|
||||
}
|
||||
};
|
||||
@@ -1034,6 +1039,7 @@ struct ggml_webgpu_mul_mat_shader_decisions {
|
||||
uint32_t subgroup_matrix_n;
|
||||
|
||||
uint32_t mul_mat_wg_size;
|
||||
bool src_overlap = false;
|
||||
};
|
||||
|
||||
/** MUL_MAT_ID **/
|
||||
@@ -1950,7 +1956,7 @@ class ggml_webgpu_shader_lib {
|
||||
return quantize_q8_pipelines[key];
|
||||
}
|
||||
|
||||
webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context) {
|
||||
webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
|
||||
ggml_webgpu_mul_mat_vec_pipeline_key key = {};
|
||||
key.src0_type = context.src0->type;
|
||||
key.src1_type = context.src1->type;
|
||||
@@ -1961,6 +1967,7 @@ class ggml_webgpu_shader_lib {
|
||||
key.num_cols = context.dst->ne[1];
|
||||
key.use_mmvq =
|
||||
ggml_webgpu_can_use_mmvq(context.src0, context.src1, context.supports_dot_product, context.vendor);
|
||||
key.src_overlap = src_overlap;
|
||||
|
||||
auto it = mul_mat_vec_pipelines.find(key);
|
||||
if (it != mul_mat_vec_pipelines.end()) {
|
||||
@@ -2068,6 +2075,11 @@ class ggml_webgpu_shader_lib {
|
||||
defines.push_back("Q8_1_T");
|
||||
}
|
||||
|
||||
if (key.src_overlap) {
|
||||
defines.push_back("SRC_OVERLAP");
|
||||
variant += "_src_overlap";
|
||||
}
|
||||
|
||||
defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));
|
||||
defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg));
|
||||
defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION");
|
||||
@@ -2089,7 +2101,7 @@ class ggml_webgpu_shader_lib {
|
||||
return mul_mat_vec_pipelines[key];
|
||||
}
|
||||
|
||||
webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context) {
|
||||
webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
|
||||
ggml_webgpu_mul_mat_pipeline_key key = {};
|
||||
key.src0_type = context.src0->type;
|
||||
key.src1_type = context.src1->type;
|
||||
@@ -2098,6 +2110,7 @@ class ggml_webgpu_shader_lib {
|
||||
1 :
|
||||
0;
|
||||
key.use_subgroup_matrix = context.supports_subgroup_matrix;
|
||||
key.src_overlap = src_overlap;
|
||||
|
||||
auto it = mul_mat_fast_pipelines.find(key);
|
||||
if (it != mul_mat_fast_pipelines.end()) {
|
||||
@@ -2216,6 +2229,11 @@ class ggml_webgpu_shader_lib {
|
||||
variant += "_vectorized";
|
||||
}
|
||||
|
||||
if (key.src_overlap) {
|
||||
defines.push_back("SRC_OVERLAP");
|
||||
variant += "_src_overlap";
|
||||
}
|
||||
|
||||
if (!key.use_subgroup_matrix) {
|
||||
defines.push_back("WORKGROUP_SIZE_M=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_M) + "u");
|
||||
defines.push_back("WORKGROUP_SIZE_N=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_N) + "u");
|
||||
|
||||
Reference in New Issue
Block a user