Function edgetam_perceiver_forward

sam3.cpp:5606–5858

}

static bool edgetam_perceiver_forward(
        const sam3_model& model,
        const std::vector<float>& mem_features,  // [64 * H * W]
        const std::vector<float>& mem_pos,        // [64 * H * W]
        int H, int W,
        std::vector<float>& out_latents,          // output: [512 * 64]
        std::vector<float>& out_pos) {            // output: [512 * 64]
    auto t_start = std::chrono::high_resolution_clock::now();

    const auto& perc = model.perceiver;
    const int D        = 64;
    const int N_1d     = 256;     // number of 1D latent tokens
    const int N_2d     = 256;     // number of 2D latent tokens
    const int HW       = H * W;
    const int n_layers = (int)perc.layers.size();
    const int nw       = (int)sqrtf((float)N_2d);  // 16 windows per spatial dim
    const int ws       = H / nw;   // window size (4 for H=64, 2 for H=32)
    const int ws2      = ws * ws;   // tokens per window

    fprintf(stderr, "%s: H=%d W=%d D=%d layers=%d 1d_tokens=%d 2d_tokens=%d ws=%d\n",
            __func__, H, W, D, n_layers, N_1d, N_2d, ws);

    if (ws < 1 || H % ws != 0 || W % ws != 0) {
        fprintf(stderr, "%s: H=%d or W=%d not divisible by window_size=%d (nw=%d)\n",
                __func__, H, W, ws, nw);
        return false;
    }
    if (nw * nw != N_2d) {
        fprintf(stderr, "%s: expected %d 2D windows but got %d\n",
                __func__, N_2d, nw * nw);
        return false;
    }

    // ── Read learnable latent tokens from model weights ─────────────────
    // latents_1d: ggml shape [64, 256] → ne[0]=64(D), ne[1]=256(N)
    std::vector<float> latents_1d_data(D * N_1d);
    sam3_read_f32(perc.latents_1d, latents_1d_data.data(), D * N_1d);

    std::vector<float> latents_2d_data(D * N_2d);
    sam3_read_f32(perc.latents_2d, latents_2d_data.data(), D * N_2d);

    // ── Prepare 2D windowed features on CPU ─────────────────────────────
    // mem_features layout: [D, H*W] (ggml: ne[0]=D, ne[1]=H*W, stored as
    //   element (d, pos) at flat index d + pos*D where pos = w + h*W).
    // Actually, the features from the memory encoder have layout [D, W, H] in
    // ggml 4D, flattened to [D, H*W]. We need to interpret as a 2D spatial grid.
    // The input is already in [D, H*W] layout with spatial stride such that
    // pos = w + h * W.

    // Window partition: [D, H*W] → [D, ws2, N_2d] i.e. [64, 16, 256]
    std::vector<float> feat_windowed(D * ws2 * N_2d);
    edgetam_window_partition_cpu(mem_features.data(), D, H, W, ws, feat_windowed.data());

    // Reshape latents_2d for windowed processing: [D, N_2d] → [D, 1, N_2d]
    // Each window has exactly 1 latent token.
    // The latents_2d weight is [D, 256] = 256 latent tokens.
    // For 2D: batch=256 windows, each with 1 latent and 16 feature tokens.
    // Rearrange latents_2d from [D, 256] to [D, 1, 256] (batch dim = 256)