Commit 713319e28 for llama.cpp
commit 713319e28492d319b81e49646151f6bbdb9d4551
Author: Leon <118652405+LeON-Nie-code@users.noreply.github.com>
Date: Tue Nov 3 13:54:00 2026 +0800
webgpu: add ARANGE support (#31455)
diff --git a/docs/ops.md b/docs/ops.md
index 2885441c5..3e14b06cb 100644
--- a/docs/ops.md
+++ b/docs/ops.md
@@ -19,7 +19,7 @@ Legend:
| ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| ADD1 | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
| ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
diff --git a/docs/ops/WebGPU.csv b/docs/ops/WebGPU.csv
index c19396c03..d402092c4 100644
--- a/docs/ops/WebGPU.csv
+++ b/docs/ops/WebGPU.csv
@@ -12326,8 +12326,16 @@
"WebGPU: WebGPU","PAD_REFLECT_1D","type=f32,ne_a=[512,34,2,1],pad_0=10,pad_1=9","support","0","no","WebGPU"
"WebGPU: WebGPU","PAD_REFLECT_1D","type=f32,ne_a=[3000,384,4,1],pad_0=10,pad_1=9","support","0","no","WebGPU"
"WebGPU: WebGPU","ROLL","shift0=3,shift1=-2,shift3=1,shift4=-1","support","0","no","WebGPU"
-"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=10.000000,step=1.000000","support","0","no","WebGPU"
-"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=1048576.000000,step=1.000000","support","0","no","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=10.000000,step=1.000000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=3.000000,stop=3.500000,step=1.000000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-10.000000,stop=10.000000,step=0.300000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=62.750000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=63.000000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=63.250000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=254.750000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=255.000000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=255.250000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=1048576.000000,step=1.000000","support","1","yes","WebGPU"
"WebGPU: WebGPU","TIMESTEP_EMBEDDING","type=f32,ne_a=[2,1,1,1],dim=320,max_period=10000","support","0","no","WebGPU"
"WebGPU: WebGPU","LEAKY_RELU","type=f32,ne_a=[10,5,4,3],negative_slope=0.100000","support","0","no","WebGPU"
"WebGPU: WebGPU","CUMSUM","type=f32,ne=[10,5,4,3]","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 1e63a9fac..d1cd0818e 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
@@ -1233,6 +1233,7 @@ class ggml_webgpu_shader_lib {
std::unordered_map<int, webgpu_pipeline> argsort_pipelines; // key is order
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<ggml_webgpu_row_norm_pipeline_key, webgpu_pipeline, ggml_webgpu_row_norm_pipeline_key_hash>
row_norm_pipelines; // op/inplace
@@ -1943,6 +1944,24 @@ class ggml_webgpu_shader_lib {
return gated_delta_net_pipelines[key];
}
+ webgpu_pipeline get_arange_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ auto it = arange_pipelines.find(1);
+ if (it != arange_pipelines.end()) {
+ return it->second;
+ }
+
+ std::vector<std::string> defines;
+ defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+
+ auto processed = preprocessor.preprocess(wgsl_arange, defines);
+ auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
+ decisions->wg_size = context.max_wg_size;
+ webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, "arange");
+ pipeline.context = decisions;
+ arange_pipelines[1] = pipeline;
+ return arange_pipelines[1];
+ }
+
webgpu_pipeline get_pad_pipeline(const ggml_webgpu_shader_lib_context & context) {
ggml_webgpu_pad_pipeline_key key = {};
key.circular = ggml_get_op_params_i32(context.dst, 8) != 0;
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index 44ff0b10c..bb2deb9da 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -856,6 +856,32 @@ static webgpu_encoded_op ggml_webgpu_set(webgpu_context & ctx,
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}
+static webgpu_encoded_op ggml_webgpu_arange(webgpu_context & ctx, ggml_tensor * dst) {
+ ggml_webgpu_shader_lib_context shader_lib_ctx = {};
+ shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
+
+ webgpu_pipeline pipeline = ctx->shader_lib->get_arange_pipeline(shader_lib_ctx);
+ auto * decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
+
+ const uint32_t ne = (uint32_t) ggml_nelements(dst);
+ std::vector<uint32_t> params = {
+ ne,
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
+ ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 0)),
+ ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 2)),
+ };
+
+ std::vector<wgpu::BindGroupEntry> entries = {
+ ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, dst),
+ };
+
+ uint32_t wg_x;
+ uint32_t wg_y;
+ uint32_t total_wg = CEIL_DIV(ne, decisions->wg_size);
+ compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
+ return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
+}
+
static webgpu_encoded_op ggml_webgpu_pad(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) {
ggml_webgpu_shader_lib_context shader_lib_ctx = {};
shader_lib_ctx.src0 = src;
@@ -3480,6 +3506,8 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
return ggml_webgpu_gated_delta_net(ctx, src0, src1, src2, node->src[3], node->src[4], node->src[5],
node, nullptr);
}
+ case GGML_OP_ARANGE:
+ return ggml_webgpu_arange(ctx, node);
case GGML_OP_PAD:
return ggml_webgpu_pad(ctx, src0, node);
case GGML_OP_ARGMAX:
@@ -4784,6 +4812,9 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
case GGML_OP_COS:
supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && (src0->type == op->type);
break;
+ case GGML_OP_ARANGE:
+ supports_op = op->type == GGML_TYPE_F32 && ggml_is_contiguous(op);
+ break;
case GGML_OP_PAD:
supports_op = op->type == GGML_TYPE_F32 && src0->type == GGML_TYPE_F32;
break;
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/arange.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/arange.wgsl
new file mode 100644
index 000000000..707a66536
--- /dev/null
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/arange.wgsl
@@ -0,0 +1,22 @@
+@group(0) @binding(0)
+var<storage, read_write> dst: array<f32>;
+
+struct Params {
+ ne: u32,
+ offset_dst: u32,
+ start: f32,
+ step: f32,
+};
+
+@group(0) @binding(1)
+var<uniform> params: Params;
+
+@compute @workgroup_size(WG_SIZE)
+fn main(@builtin(num_workgroups) num_wg: vec3<u32>, @builtin(global_invocation_id) gid: vec3<u32>) {
+ let i = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
+ if (i >= params.ne) {
+ return;
+ }
+
+ dst[params.offset_dst + i] = params.start + params.step * f32(i);
+}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 64fdb8259..68b25c876 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12206,6 +12206,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_roll());
test_cases.emplace_back(new test_roll(3, -2, 1, -1, true));
test_cases.emplace_back(new test_arange());
+ test_cases.emplace_back(new test_arange(GGML_TYPE_F32, 3.0f, 3.5f, 1.0f));
+ test_cases.emplace_back(new test_arange(GGML_TYPE_F32, -10.0f, 10.0f, 0.3f));
+ for (float n : {255.0f, 256.0f, 257.0f, 1023.0f, 1024.0f, 1025.0f}) {
+ test_cases.emplace_back(new test_arange(GGML_TYPE_F32, -1.0f, -1.0f + n * 0.25f, 0.25f));
+ }
test_cases.emplace_back(new test_arange(GGML_TYPE_F32, 0.0f, 1048576.0f, 1.0f));
test_cases.emplace_back(new test_timestep_embedding());
test_cases.emplace_back(new test_leaky_relu());
@@ -12759,6 +12764,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
std::vector<std::unique_ptr<test_case>> test_cases;
+ for (float n : {128.0f, 4096.0f, 32768.0f, 1048576.0f}) {
+ test_cases.emplace_back(new test_arange(GGML_TYPE_F32, 0.0f, n, 1.0f));
+ }
+
// delta-net conv input: conv state followed by the transposed qkv of the ubatch
test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {3, 8192, 1, 1}, 512, 0, 16));