Repository: ggerganov/llama.cpp
File: ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp (C++)

Code before:
    defines.push_back(std::string("SG_MAT_K=") + std::to_string(context.sg_mat_k));

    // Add chosen Q/KV tile sizes
    uint32_t q_tile  = context.sg_mat_m;
    uint32_t kv_tile = std::min(ggml_webgpu_flash_attn_max_kv_tile(context),
                                context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES);
    if (context.key.kv_direct) {
        GGML_ASSERT(kv_tile <= GGML_WEBGPU_KV_SEQ_PAD);
        // Avoids having to use bounds-checks and decreasing performance for direct KV loads
        while (GGML_WEBGPU_KV_SEQ_PAD % kv_tile != 0) {
            kv_tile -= context.sg_mat_n;
        }
    }

    defines.push_back(std::string("Q_TILE=") + std::to_string(q_tile));
    defines.push_back(std::string("KV_TILE=") + std::to_string(kv_tile));

    // workgroup size
    uint32_t wg_size = std::max(context.max_subgroup_size, GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE);

    defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));

    ggml_webgpu_processed_shader result;
    result.wgsl                                         = preprocessor.preprocess(shader_src, defines);
    result.variant                                      = variant;
    auto decisions = std::make_shared<ggml_webgpu_flash_attn_shader_decisions>();
    decisions->q_tile                                   = q_tile;
    decisions->kv_tile                                  = kv_tile;
    decisions->wg_size                                  = wg_size;
    result.decisions                                    = decisions;
    return result;
}

/** Generic **/

struct ggml_webgpu_generic_shader_lib_context {
    int      vec4;
    uint32_t max_wg_size;
};

struct ggml_webgpu_generic_shader_decisions {
    uint32_t wg_size;
};

inline ggml_webgpu_processed_shader ggml_webgpu_preprocess_generic_shader(
    pre_wgsl::Preprocessor &                       preprocessor,
    const char *                                   shader_src,
    const ggml_webgpu_generic_shader_lib_context & context,
    const std::string &                            base_variant) {
    std::vector<std::string> defines;
    std::string              variant = base_variant;

Review comment (suggestion):
I changed this into a `shared_ptr` because this was leaking on the heap since we never deallocated it.

Code after:
    defines.push_back(std::string("SG_MAT_N=") + std::to_string(context.sg_mat_n));
    defines.push_back(std::string("SG_MAT_K=") + std::to_string(context.sg_mat_k));

    // Add chosen Q/KV tile sizes
    uint32_t q_tile  = context.sg_mat_m;
    uint32_t kv_tile = std::min(ggml_webgpu_flash_attn_max_kv_tile(context),
                                context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES);
    if (context.key.kv_direct) {
        GGML_ASSERT(kv_tile <= GGML_WEBGPU_KV_SEQ_PAD);
        // Avoids having to use bounds-checks and decreasing performance for direct KV loads
        while (GGML_WEBGPU_KV_SEQ_PAD % kv_tile != 0) {
            kv_tile -= context.sg_mat_n;
        }
    }

    defines.push_back(std::string("Q_TILE=") + std::to_string(q_tile));
    defines.push_back(std::string("KV_TILE=") + std::to_string(kv_tile));

    // workgroup size
    uint32_t wg_size = std::max(context.max_subgroup_size, GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE);

    defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));

    ggml_webgpu_processed_shader result;
    result.wgsl        = preprocessor.preprocess(shader_src, defines);
    result.variant     = variant;
    auto decisions     = std::make_shared<ggml_webgpu_flash_attn_shader_decisions>();
    decisions->q_tile  = q_tile;
    decisions->kv_tile = kv_tile;
    decisions->wg_size = wg_size;
    result.decisions   = decisions;
    return result;
}

/** Generic **/

struct ggml_webgpu_generic_shader_lib_context {
    int      vec4;
    uint32_t max_wg_size;
};

struct ggml_webgpu_generic_shader_decisions {
    uint32_t wg_size;
};

inline ggml_webgpu_processed_shader ggml_webgpu_preprocess_generic_shader(
    pre_wgsl::Preprocessor &                       preprocessor,
    const char *                                   shader_src,
    const ggml_webgpu_generic_shader_lib_context & context,
    const std::string &                            base_variant) {
    std::vector<std::string> defines;