| 2577 | } |
| 2578 | |
| 2579 | static webgpu_encoded_op ggml_webgpu_soft_max(webgpu_context & ctx, |
| 2580 | ggml_tensor * src0, |
| 2581 | ggml_tensor * src1, |
| 2582 | ggml_tensor * src2, |
| 2583 | ggml_tensor * dst) { |
| 2584 | ggml_webgpu_shader_lib_context shader_lib_ctx = {}; |
| 2585 | shader_lib_ctx.src0 = src0; |
| 2586 | shader_lib_ctx.src1 = src1; |
| 2587 | shader_lib_ctx.src2 = src2; |
| 2588 | shader_lib_ctx.dst = dst; |
| 2589 | shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; |
| 2590 | |
| 2591 | webgpu_pipeline pipeline = ctx->shader_lib->get_soft_max_pipeline(shader_lib_ctx); |
| 2592 | auto * decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get()); |
| 2593 | |
| 2594 | const bool inplace = decisions->inplace; |
| 2595 | const int has_mask = (src1 != nullptr); |
| 2596 | const int has_sink = (src2 != nullptr); |
| 2597 | float max_bias = ggml_get_op_params_f32(dst, 1); |
| 2598 | float n_head_log2 = float(1u << (uint32_t) floor(log2(src0->ne[2]))); |
| 2599 | float m0 = powf(2.0f, -(max_bias) / n_head_log2); |
| 2600 | float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); |
| 2601 | |
| 2602 | std::vector<uint32_t> params = { |
| 2603 | (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)), |
| 2604 | has_mask ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)) : 0, |
| 2605 | has_sink ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)) : 0, |
| 2606 | (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), |
| 2607 | (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)), |
| 2608 | (uint32_t) (src0->nb[2] / ggml_type_size(src0->type)), |
| 2609 | (uint32_t) (src0->nb[3] / ggml_type_size(src0->type)), |
| 2610 | has_mask ? (uint32_t) (src1->nb[1] / ggml_type_size(src1->type)) : 0, |
| 2611 | has_mask ? (uint32_t) (src1->nb[2] / ggml_type_size(src1->type)) : 0, |
| 2612 | has_mask ? (uint32_t) (src1->nb[3] / ggml_type_size(src1->type)) : 0, |
| 2613 | (uint32_t) (dst->nb[1] / ggml_type_size(dst->type)), |
| 2614 | (uint32_t) (dst->nb[2] / ggml_type_size(dst->type)), |
| 2615 | (uint32_t) (dst->nb[3] / ggml_type_size(dst->type)), |
| 2616 | (uint32_t) ggml_nelements(dst), |
| 2617 | (uint32_t) src0->ne[0], |
| 2618 | (uint32_t) src0->ne[1], |
| 2619 | (uint32_t) src0->ne[2], |
| 2620 | has_mask ? (uint32_t) src1->ne[2] : 0, |
| 2621 | has_mask ? (uint32_t) src1->ne[3] : 0, |
| 2622 | ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 0)), // scale |
| 2623 | ggml_webgpu_u32_from_f32(max_bias), |
| 2624 | ggml_webgpu_u32_from_f32(n_head_log2), |
| 2625 | ggml_webgpu_u32_from_f32(m0), |
| 2626 | ggml_webgpu_u32_from_f32(m1) |
| 2627 | }; |
| 2628 | |
| 2629 | std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_bind_group_entry( |
| 2630 | 0, ggml_webgpu_tensor_buf(src0), ggml_webgpu_tensor_align_offset(ctx, src0), |
| 2631 | ggml_webgpu_tensor_binding_size(ctx, src0)) }; |
| 2632 | uint32_t binding_num = 1; |
| 2633 | if (has_mask) { |
| 2634 | entries.push_back(ggml_webgpu_make_bind_group_entry(binding_num, ggml_webgpu_tensor_buf(src1), |
| 2635 | ggml_webgpu_tensor_align_offset(ctx, src1), |
| 2636 | ggml_webgpu_tensor_binding_size(ctx, src1))); |
no test coverage detected