Commit a691b2345 for llama.cpp
commit a691b2345844ce1b430201fc5fd404ce062e05f2
Author: Leon <118652405+LeON-Nie-code@users.noreply.github.com>
Date: Sun Nov 8 17:50:09 2026 +0800
webgpu : add COUNT_EQUAL support (#31578)
* webgpu : add COUNT_EQUAL support
Assisted-by: Codex
* webgpu : reduce small COUNT_EQUAL inputs in one dispatch
Assisted-by: Codex
* webgpu : simplify COUNT_EQUAL and reuse existing tests
* webgpu : move COUNT_EQUAL pipeline cache declaration
* Update ggml/src/ggml-webgpu/ggml-webgpu.cpp
Co-authored-by: Masashi Yoshimura <yoshimura.masashi.frbs@gmail.com>
* tests : simplify COUNT_EQUAL coverage
---------
Co-authored-by: Masashi Yoshimura <yoshimura.masashi.frbs@gmail.com>
diff --git a/docs/ops.md b/docs/ops.md
index 4eec04e65..463d48250 100644
--- a/docs/ops.md
+++ b/docs/ops.md
@@ -34,7 +34,7 @@ Legend:
| CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| COS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
| CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
| CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
diff --git a/docs/ops/WebGPU.csv b/docs/ops/WebGPU.csv
index 3867ec891..98720366e 100644
--- a/docs/ops/WebGPU.csv
+++ b/docs/ops/WebGPU.csv
@@ -5465,8 +5465,10 @@
"WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[3,2,3,1],ne_kernel=[2,2,1,3],stride=1","support","0","no","WebGPU"
"WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[10,10,9,1],ne_kernel=[3,3,1,9],stride=2","support","0","no","WebGPU"
"WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[129,63,35,1],ne_kernel=[3,3,48,35],stride=1","support","0","no","WebGPU"
-"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1]","support","0","no","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1],same_src=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1],same_src=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1],same_src=1","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1],same_src=1","support","1","yes","WebGPU"
"WebGPU: WebGPU","ARGMAX","type=f32,ne=[32,1,1,1]","support","1","yes","WebGPU"
"WebGPU: WebGPU","ARGMAX","type=f32,ne=[32,513,1,1]","support","1","yes","WebGPU"
"WebGPU: WebGPU","ARGMAX","type=f32,ne=[100,10,1,1]","support","1","yes","WebGPU"
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
index 30653fccf..92ec03fbe 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
@@ -27,6 +27,7 @@
#define GGML_WEBGPU_KV_SEQ_PAD 256u
#define GGML_WEBGPU_ARGSORT_MERGE_MAX_WG_SIZE 512u
+#define GGML_WEBGPU_COUNT_EQUAL_WG_SIZE 256u
// Matrix multiplication parameters
@@ -1262,6 +1263,7 @@ class ggml_webgpu_shader_lib {
std::unordered_map<int, webgpu_pipeline> argsort_merge_pipelines; // key is order
std::unordered_map<int, webgpu_pipeline> cumsum_pipelines; // key is fixed, no variants yet
std::unordered_map<int, webgpu_pipeline> arange_pipelines; // key is fixed, no variants yet
+ std::unordered_map<int, webgpu_pipeline> count_equal_pipelines; // key is src_overlap
std::unordered_map<ggml_webgpu_row_norm_pipeline_key, webgpu_pipeline, ggml_webgpu_row_norm_pipeline_key_hash>
row_norm_pipelines; // op/inplace
@@ -1346,6 +1348,21 @@ class ggml_webgpu_shader_lib {
public:
ggml_webgpu_shader_lib(wgpu::Device device) { this->device = device; }
+ webgpu_pipeline get_count_equal_pipeline(bool src_overlap) {
+ const int key = src_overlap;
+ auto it = count_equal_pipelines.find(key);
+ if (it != count_equal_pipelines.end()) {
+ return it->second;
+ }
+ std::vector<std::string> defines = { "WG_SIZE=" + std::to_string(GGML_WEBGPU_COUNT_EQUAL_WG_SIZE) };
+ if (src_overlap) {
+ defines.push_back("SRC_OVERLAP");
+ }
+ auto processed = preprocessor.preprocess(wgsl_count_equal, defines);
+ count_equal_pipelines[key] = ggml_webgpu_create_pipeline(device, processed, "count_equal");
+ return count_equal_pipelines[key];
+ }
+
webgpu_pipeline get_sum_rows_pipeline(const ggml_webgpu_shader_lib_context & context) {
const bool mean = context.dst->op == GGML_OP_MEAN;
auto it = sum_rows_pipelines.find(mean);
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index 0d78a1b70..31a713bc1 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -3141,6 +3141,52 @@ static webgpu_encoded_op ggml_webgpu_argmax(webgpu_context & ctx, ggml_tensor *
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x);
}
+static webgpu_encoded_op ggml_webgpu_count_equal(webgpu_context & ctx,
+ ggml_tensor * src0,
+ ggml_tensor * src1,
+ ggml_tensor * dst) {
+ const uint32_t ne = ggml_nelements(src0);
+ const uint32_t offset_dst = ggml_webgpu_tensor_misalignment(ctx, dst);
+ std::vector<webgpu_dispatch_desc> dispatches;
+ // Reset both words of the I64 result before accumulating partial counts.
+ dispatches.push_back({
+ ctx->global_ctx->memset_pipeline,
+ { offset_dst, sizeof(int64_t), 0 },
+ { ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, dst) },
+ { 1, 1 }
+ });
+
+ if (ne == 0) {
+ return ggml_backend_webgpu_build_multi(ctx, dispatches);
+ }
+
+ const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1);
+ webgpu_pipeline pipeline = ctx->shader_lib->get_count_equal_pipeline(src_overlap);
+
+ uint32_t offset_src0 = ggml_webgpu_tensor_misalignment(ctx, src0) / sizeof(int32_t);
+ uint32_t offset_src1 = ggml_webgpu_tensor_misalignment(ctx, src1) / sizeof(int32_t);
+ std::vector<wgpu::BindGroupEntry> entries;
+ if (src_overlap) {
+ const ggml_webgpu_merged_binding_range range = ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 });
+ offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, range);
+ offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, range);
+ entries.push_back(ggml_webgpu_make_bind_group_entry(0, ggml_webgpu_tensor_buf(src0), range.offset, range.size));
+ } else {
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0));
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1));
+ }
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, entries.size(), dst));
+ std::vector<uint32_t> params = { offset_src0, offset_src1, offset_dst / (uint32_t) sizeof(uint32_t), ne };
+ uint32_t wg_x;
+ uint32_t wg_y;
+ const uint32_t total_wg = CEIL_DIV(ne, GGML_WEBGPU_COUNT_EQUAL_WG_SIZE);
+ compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
+ dispatches.push_back({
+ pipeline, std::move(params), std::move(entries), { wg_x, wg_y }
+ });
+ return ggml_backend_webgpu_build_multi(ctx, dispatches);
+}
+
static webgpu_encoded_op ggml_webgpu_argsort(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) {
bool is_top_k = dst->op == GGML_OP_TOP_K;
@@ -3580,6 +3626,8 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
return ggml_webgpu_arange(ctx, node);
case GGML_OP_PAD:
return ggml_webgpu_pad(ctx, src0, node);
+ case GGML_OP_COUNT_EQUAL:
+ return ggml_webgpu_count_equal(ctx, src0, src1, node);
case GGML_OP_ARGMAX:
return ggml_webgpu_argmax(ctx, src0, node);
case GGML_OP_ARGSORT:
@@ -5131,6 +5179,10 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
case GGML_OP_SUM:
supports_op = op->type == GGML_TYPE_F32 && src0->type == op->type && ggml_is_contiguously_allocated(src0);
break;
+ case GGML_OP_COUNT_EQUAL:
+ supports_op = src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && ggml_is_contiguous(src0) &&
+ ggml_is_contiguous(src1);
+ break;
case GGML_OP_SUM_ROWS:
case GGML_OP_MEAN:
supports_op = op->type == GGML_TYPE_F32 && src0->type == op->type && ggml_is_contiguous_rows(src0);
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/count_equal.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/count_equal.wgsl
new file mode 100644
index 000000000..1ab6a635b
--- /dev/null
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/count_equal.wgsl
@@ -0,0 +1,57 @@
+struct Params {
+ offset_src0: u32,
+ offset_src1: u32,
+ offset_dst: u32,
+ ne: u32,
+};
+
+#ifdef SRC_OVERLAP
+@group(0) @binding(0)
+var<storage, read_write> merged_src: array<i32>;
+
+@group(0) @binding(1)
+var<storage, read_write> dst: array<atomic<u32>>;
+
+@group(0) @binding(2)
+var<uniform> params: Params;
+#else
+@group(0) @binding(0)
+var<storage, read_write> src0: array<i32>;
+
+@group(0) @binding(1)
+var<storage, read_write> src1: array<i32>;
+
+@group(0) @binding(2)
+var<storage, read_write> dst: array<atomic<u32>>;
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+#endif
+
+var<workgroup> shared_count: array<u32, WG_SIZE>;
+
+@compute @workgroup_size(WG_SIZE)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+ @builtin(num_workgroups) num_wg: vec3<u32>,
+ @builtin(local_invocation_id) lid: vec3<u32>) {
+ var count = 0u;
+ let i = (wid.x + wid.y * num_wg.x) * WG_SIZE + lid.x;
+ if (i < params.ne) {
+#ifdef SRC_OVERLAP
+ count += select(0u, 1u, merged_src[params.offset_src0 + i] == merged_src[params.offset_src1 + i]);
+#else
+ count += select(0u, 1u, src0[params.offset_src0 + i] == src1[params.offset_src1 + i]);
+#endif
+ }
+ shared_count[lid.x] = count;
+ workgroupBarrier();
+ for (var offset: u32 = WG_SIZE >> 1; offset > 0; offset >>= 1) {
+ if (lid.x < offset) {
+ shared_count[lid.x] += shared_count[lid.x + offset];
+ }
+ workgroupBarrier();
+ }
+ if (lid.x == 0) {
+ atomicAdd(&dst[params.offset_dst], shared_count[0]);
+ }
+}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index f3f35329a..852b62701 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -3106,14 +3106,15 @@ struct test_argmax : public test_case {
struct test_count_equal : public test_case {
const ggml_type type;
const std::array<int64_t, 4> ne;
+ const bool same_src;
std::string vars() override {
- return VARS_TO_STR2(type, ne);
+ return VARS_TO_STR3(type, ne, same_src);
}
test_count_equal(ggml_type type = GGML_TYPE_F32,
- std::array<int64_t, 4> ne = {4, 500, 1, 1})
- : type(type), ne(ne) {}
+ std::array<int64_t, 4> ne = {4, 500, 1, 1}, bool same_src = false)
+ : type(type), ne(ne), same_src(same_src) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data());
@@ -3128,7 +3129,7 @@ struct test_count_equal : public test_case {
ggml_tensor * b_argmax = ggml_argmax(ctx, b);
ggml_set_name(b_argmax, "b_argmax");
- ggml_tensor * out = ggml_count_equal(ctx, a_argmax, b_argmax);
+ ggml_tensor * out = ggml_count_equal(ctx, a_argmax, same_src ? a_argmax : b_argmax);
ggml_set_name(out, "out");
return out;
@@ -10963,8 +10964,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 2}, {3, 3, 1, 9}, 2, kernel_type)); // for multiple batches
}
- test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, 500, 1, 1}));
- test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, 5000, 1, 1}));
+ for (bool same_src : {false, true}) {
+ test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, 500, 1, 1}, same_src));
+ test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, 5000, 1, 1}, same_src));
+ }
test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {32, 1, 1, 1}));
test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {32, 513, 1, 1}));
@@ -13233,6 +13236,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {64, 64, 20, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f));
test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {77, 64, 20, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f));
+ for (int64_t n : {256, 1024}) {
+ test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, n, 1, 1}));
+ }
+
test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {32, 10, 1, 1}));
test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {1024, 10, 1, 1}));
test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {32000, 512, 1, 1}));