diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index a7ff36030f..d1cf780835 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -106,6 +106,11 @@ struct ggml_webgpu_generic_shader_decisions { bool inplace = false; }; +struct ggml_webgpu_get_rows_shader_decisions { + uint32_t wg_size = 0; + bool vectorized = false; +}; + struct ggml_webgpu_binary_shader_decisions { uint32_t wg_size = 0; bool inplace = false; @@ -1551,8 +1556,8 @@ class ggml_webgpu_shader_lib { return argsort_merge_pipelines[order]; } - webgpu_pipeline get_get_rows_pipeline(const ggml_webgpu_shader_lib_context & context) { - const bool vectorized = context.src0->type == GGML_TYPE_F32 && context.dst->ne[0] % 4 == 0; + webgpu_pipeline get_get_rows_pipeline(const ggml_webgpu_shader_lib_context & context, bool vec4_aligned) { + const bool vectorized = context.src0->type == GGML_TYPE_F32 && context.dst->ne[0] % 4 == 0 && vec4_aligned; ggml_webgpu_get_rows_pipeline_key key = {}; key.src_type = context.src0->type; key.vectorized = (int) vectorized; @@ -1669,8 +1674,9 @@ class ggml_webgpu_shader_lib { defines.push_back("WG_SIZE=" + std::to_string(context.max_wg_size)); auto processed = preprocessor.preprocess(wgsl_get_rows, defines); - auto decisions = std::make_shared(); + auto decisions = std::make_shared(); decisions->wg_size = context.max_wg_size; + decisions->vectorized = vectorized; webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); pipeline.context = decisions; get_rows_pipelines[key] = pipeline; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 8b060c41a1..9b494d421f 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -1518,15 +1518,24 @@ static webgpu_encoded_op ggml_webgpu_get_rows(webgpu_context & ctx, shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_get_rows_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); + const uint32_t offset_src = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)); + const uint32_t offset_dst = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)); + const uint32_t stride_src1 = (uint32_t) (src->nb[1] / ggml_type_size(src->type)); + const uint32_t stride_src2 = (uint32_t) (src->nb[2] / ggml_type_size(src->type)); + const uint32_t stride_src3 = (uint32_t) (src->nb[3] / ggml_type_size(src->type)); - std::vector params = { (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)), + const bool vec4_aligned = offset_src % 4 == 0 && offset_dst % 4 == 0 && stride_src1 % 4 == 0 && + stride_src2 % 4 == 0 && stride_src3 % 4 == 0; + + webgpu_pipeline pipeline = ctx->shader_lib->get_get_rows_pipeline(shader_lib_ctx, vec4_aligned); + auto * decisions = static_cast(pipeline.context.get()); + + std::vector params = { offset_src, (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, idx) / ggml_type_size(idx->type)), - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), - (uint32_t) (src->nb[1] / ggml_type_size(src->type)), - (uint32_t) (src->nb[2] / ggml_type_size(src->type)), - (uint32_t) (src->nb[3] / ggml_type_size(src->type)), + offset_dst, + stride_src1, + stride_src2, + stride_src3, (uint32_t) (idx->nb[0] / ggml_type_size(idx->type)), (uint32_t) (idx->nb[1] / ggml_type_size(idx->type)), (uint32_t) (idx->nb[2] / ggml_type_size(idx->type)), @@ -1544,7 +1553,7 @@ static webgpu_encoded_op ggml_webgpu_get_rows(webgpu_context & ctx, ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx), ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst) }; - uint32_t blocks_per_row = (uint32_t) (dst->ne[0] / (src->type == GGML_TYPE_F32 && dst->ne[0] % 4 == 0 ? 4 : 1)); + uint32_t blocks_per_row = (uint32_t) (dst->ne[0] / (decisions->vectorized ? 4 : 1)); uint32_t total_rows = (uint32_t) (dst->ne[1] * dst->ne[2] * dst->ne[3]); uint32_t total_threads = float_parallel ? blocks_per_row * total_rows : total_rows; uint32_t wg_x = CEIL_DIV(total_threads, decisions->wg_size); @@ -4333,22 +4342,12 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const src0->type == GGML_TYPE_F32 && (src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32)); break; case GGML_OP_GET_ROWS: - { - const size_t storage_alignment = - ctx->webgpu_global_ctx->capabilities.limits.minStorageBufferOffsetAlignment; - const size_t src_address_unit = - src0->type == GGML_TYPE_F32 && op->ne[0] % 4 == 0 ? 4 * sizeof(float) : ggml_type_size(src0->type); - if (ggml_webgpu_tensor_misalignment(src0, storage_alignment) % src_address_unit != 0) { - break; - } - if (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || - ggml_webgpu_supported_qtype(src0->type)) { - supports_op = (op->type == GGML_TYPE_F32); - } else if (src0->type == GGML_TYPE_I32) { - supports_op = op->type == GGML_TYPE_I32; - } - break; + if (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || ggml_webgpu_supported_qtype(src0->type)) { + supports_op = (op->type == GGML_TYPE_F32); + } else if (src0->type == GGML_TYPE_I32) { + supports_op = op->type == GGML_TYPE_I32; } + break; case GGML_OP_MUL_MAT: { switch (src1->type) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 260ffef66d..30792e4096 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2338,24 +2338,26 @@ struct test_get_rows : public test_case { const int be2; // batch size const bool v; // view src1 const bool vs0; // view src0 + const int offset_cols; // // column offset of the view src0 std::string vars() override { - return VARS_TO_STR8(type, n, m, r, be1, be2, v, vs0); + return VARS_TO_STR9(type, n, m, r, be1, be2, v, vs0, offset_cols); } - test_get_rows(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int be1 = 1, int be2 = 1, bool v = false, bool vs0 = false) - : type(type), n(n), m(m), r(r), be1(be1), be2(be2), v(v), vs0(vs0) {} + test_get_rows(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int be1 = 1, int be2 = 1, bool v = false, bool vs0 = false, int offset_cols = 0) + : type(type), n(n), m(m), r(r), be1(be1), be2(be2), v(v), vs0(vs0), offset_cols(offset_cols) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * in; if (vs0) { const int offset_rows = 3; const int padded_m = m + offset_rows; - ggml_tensor * in_padded = ggml_new_tensor_4d(ctx, type, n, padded_m, be1, be2); + const int padded_n = n + offset_cols; + ggml_tensor * in_padded = ggml_new_tensor_4d(ctx, type, padded_n, padded_m, be1, be2); ggml_set_name(in_padded, "in_padded"); in = ggml_view_4d(ctx, in_padded, n, m, be1, be2, in_padded->nb[1], in_padded->nb[2], in_padded->nb[3], - offset_rows * in_padded->nb[1]); + offset_cols * in_padded->nb[0] + offset_rows * in_padded->nb[1]); ggml_set_name(in, "in_view"); } else { in = ggml_new_tensor_4d(ctx, type, n, m, be1, be2); @@ -9043,6 +9045,7 @@ static std::vector> make_test_cases_eval() { } } } + test_cases.emplace_back(new test_get_rows(GGML_TYPE_F32, 256, 8, 2, 1, 1, false, true, 3)); test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 8, 2, 1, false)); test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 70000, 4, 1, false)); // row count > CUDA grid-y limit (65535)