Function sam3_populate_pe_cache

sam3.cpp:10291–10359

Read SAM prompt encoder weights from GPU and cache them in state. Also pre-computes the dense PE grid and no-mask tiled embedding. These never change between PVS calls for the same model.

// Also pre-computes the dense PE grid and no-mask tiled embedding.
// These never change between PVS calls for the same model.
static void sam3_populate_pe_cache(sam3_state& state, const sam3_model& model) {
    if (state.pe_cache_valid) return;

    const int D = model.hparams.sam_embed_dim;  // 256
    const int H = sam3_eff_feat_size(state, model.hparams);
    const int num_pos_feats = D / 2;            // 128
    const int pe_nel = 2 * num_pos_feats;       // 256
    const auto& pe = model.sam_pe;

    state.pe_gauss_cache.resize(pe_nel);
    if (pe.pe_gaussian->type == GGML_TYPE_F16) {
        std::vector<ggml_fp16_t> tmp(pe_nel);
        ggml_backend_tensor_get(pe.pe_gaussian, tmp.data(), 0, pe_nel * sizeof(ggml_fp16_t));
        ggml_fp16_to_fp32_row(tmp.data(), state.pe_gauss_cache.data(), pe_nel);
    } else {
        ggml_backend_tensor_get(pe.pe_gaussian, state.pe_gauss_cache.data(), 0, pe_nel * sizeof(float));
    }

    for (int i = 0; i < 4; ++i) {
        if (pe.point_embed[i]->type == GGML_TYPE_F16) {
            std::vector<ggml_fp16_t> tmp(D);
            ggml_backend_tensor_get(pe.point_embed[i], tmp.data(), 0, D * sizeof(ggml_fp16_t));
            ggml_fp16_to_fp32_row(tmp.data(), state.point_emb_cache[i], D);
        } else {
            ggml_backend_tensor_get(pe.point_embed[i], state.point_emb_cache[i], 0, D * sizeof(float));
        }
    }

    if (pe.not_a_point_embed->type == GGML_TYPE_F16) {
        std::vector<ggml_fp16_t> tmp(D);
        ggml_backend_tensor_get(pe.not_a_point_embed, tmp.data(), 0, D * sizeof(ggml_fp16_t));
        ggml_fp16_to_fp32_row(tmp.data(), state.not_a_point_cache, D);
    } else {
        ggml_backend_tensor_get(pe.not_a_point_embed, state.not_a_point_cache, 0, D * sizeof(float));
    }

    if (pe.no_mask_embed->type == GGML_TYPE_F16) {
        std::vector<ggml_fp16_t> tmp(D);
        ggml_backend_tensor_get(pe.no_mask_embed, tmp.data(), 0, D * sizeof(ggml_fp16_t));
        ggml_fp16_to_fp32_row(tmp.data(), state.no_mask_emb_cache, D);
    } else {
        ggml_backend_tensor_get(pe.no_mask_embed, state.no_mask_emb_cache, 0, D * sizeof(float));
    }

    state.dense_pe_cache.resize(D * H * H);
    for (int row = 0; row < H; ++row) {
        for (int col = 0; col < H; ++col) {
            float x_norm = ((float)col + 0.5f) / (float)H;
            float y_norm = ((float)row + 0.5f) / (float)H;
            float pe_vec[256];
            sam3_pe_encode_coord(pe_vec, x_norm, y_norm,
                                 state.pe_gauss_cache.data(), num_pos_feats);
            for (int d = 0; d < D; ++d)
                state.dense_pe_cache[d + col * D + row * D * H] = pe_vec[d];
        }
    }

    state.dense_nomask_cache.resize(D * H * H);