Function edgetam_window_partition_cpu

sam3.cpp:5546–5575

CPU-side window partition: rearrange [D, H*W] features into [256, 16, D] (256 windows of 4x4=16 tokens each, assuming H=W=64 and window_size=4). Input layout: [D, H*W] with inner dim D (ggml convention: ne[0]=D, ne[1]=H*W). Stored row-major: element (d, pos) at flat[d + pos * D], pos = w + h * W. Output layout: [D, 16, 256] — 256 windows, each with 16 tokens of dim D.

// Stored row-major: element (d, pos) at flat[d + pos * D], pos = w + h * W.
// Output layout: [D, 16, 256] — 256 windows, each with 16 tokens of dim D.
static void edgetam_window_partition_cpu(
        const float* src, int D, int H, int W,
        int window_size, float* dst) {
    const int nw_h = H / window_size;  // 16
    const int nw_w = W / window_size;  // 16
    const int ws2  = window_size * window_size;  // 16

    // For each window (wh, ww), for each pixel (ph, pw) within the window:
    //   src_pos = (wh * ws + ph) * W + (ww * ws + pw)
    //   dst: window_idx = wh * nw_w + ww,  token_idx = ph * ws + pw
    //   dst[d + token_idx * D + window_idx * D * ws2]
    for (int wh = 0; wh < nw_h; ++wh) {
        for (int ww = 0; ww < nw_w; ++ww) {
            int win_idx = wh * nw_w + ww;
            for (int ph = 0; ph < window_size; ++ph) {
                for (int pw = 0; pw < window_size; ++pw) {
                    int tok_idx  = ph * window_size + pw;
                    int src_h    = wh * window_size + ph;
                    int src_w    = ww * window_size + pw;
                    int src_pos  = src_w + src_h * W;  // H*W spatial index

                    // Copy D elements
                    const float* s = src + src_pos * D;
                    float* d       = dst + win_idx * (D * ws2) + tok_idx * D;
                    memcpy(d, s, D * sizeof(float));
                }
            }
        }
    }
}

// Inverse of window partition: [D, 16, 256] → [D, H, W] spatial layout.
// Rearranges 256 windows of 16 tokens each back into [D, H*W] feature map.