Function sam2_compute_pos_embed

sam3.cpp:4118–4175

Compute Hiera positional embedding on CPU. Returns [E, H, W] in ggml column-major layout: element (e, x, y) at flat index e + x*E + y*E*W. PE = bicubic_interp(pos_embed_bkg, (H, W)) + tile(pos_embed_window, (H/W0, W/W0))

// Returns [E, H, W] in ggml column-major layout: element (e, x, y) at flat index e + x*E + y*E*W.
// PE = bicubic_interp(pos_embed_bkg, (H, W)) + tile(pos_embed_window, (H/W0, W/W0))
static std::vector<float> sam2_compute_pos_embed(const sam3_model& model, int H, int W) {
    const auto& hp = model.hparams;
    const int E = hp.hiera_embed_dim;
    const int bkg_h = hp.hiera_pos_embed_bkg_h;
    const int bkg_w = hp.hiera_pos_embed_bkg_w;
    const int ws = hp.hiera_window_spec[0];

    // Read background PE from GPU.
    // Registered as [bkg_w, bkg_h, E, 1] in ggml; raw bytes match PyTorch NCHW [1,E,H,W].
    // ggml flat index for ne=[W,H,E,1]: w + h*W + e*W*H = same as CHW[e,h,w] = e*H*W + h*W + w.
    // So we can use the raw data directly as CHW layout for bicubic interpolation.
    std::vector<float> bkg_chw(E * bkg_h * bkg_w);
    ggml_backend_tensor_get(model.hiera.pos_embed, bkg_chw.data(), 0,
                            bkg_chw.size() * sizeof(float));

    // Bicubic interpolate to [E, H, W]
    std::vector<float> bkg_interp(E * H * W);
    sam2_bicubic_interpolate_cpu(bkg_chw.data(), E, bkg_h, bkg_w,
                                  bkg_interp.data(), H, W);

    // Read window PE from GPU: registered [ws, ws, E, 1], raw bytes = CHW [E, ws, ws]
    std::vector<float> win_chw(E * ws * ws);
    ggml_backend_tensor_get(model.hiera.pos_embed_window, win_chw.data(), 0,
                            win_chw.size() * sizeof(float));

    // Tile window PE to [E, H, W]
    std::vector<float> win_tiled(E * H * W);
    for (int e = 0; e < E; ++e)
        for (int y = 0; y < H; ++y)
            for (int x = 0; x < W; ++x)
                win_tiled[e * H * W + y * W + x] = win_chw[e * ws * ws + (y % ws) * ws + (x % ws)];

    // Dump intermediates if requested
    const char* dump_dir = getenv("SAM2_DUMP_DIR");
    if (dump_dir) {
        char path[512];
        // Dump bkg_interp as CHW
        snprintf(path, sizeof(path), "%s/cpp_pe_bkg_interp.bin", dump_dir);
        FILE* f = fopen(path, "wb");
        if (f) { fwrite(bkg_interp.data(), sizeof(float), bkg_interp.size(), f); fclose(f); }
        snprintf(path, sizeof(path), "%s/cpp_pe_bkg_interp.shape", dump_dir);
        f = fopen(path, "w"); if (f) { fprintf(f, "%d,%d,%d", E, H, W); fclose(f); }
        // Dump win_tiled as CHW
        snprintf(path, sizeof(path), "%s/cpp_pe_win_tiled.bin", dump_dir);
        f = fopen(path, "wb");
        if (f) { fwrite(win_tiled.data(), sizeof(float), win_tiled.size(), f); fclose(f); }
    }

    // Sum and convert to ggml layout [E, W, H, 1]
    std::vector<float> pe(E * H * W);
    for (int e = 0; e < E; ++e)
        for (int y = 0; y < H; ++y)
            for (int x = 0; x < W; ++x)
                pe[e + x * E + y * E * W] = bkg_interp[e * H * W + y * W + x]
                                            + win_tiled[e * H * W + y * W + x];

    return pe;
}