Function sam2_window_partition

sam3.cpp:4194–4232

Window partition: [B, H, W, C] -> [B*nW, ws, ws, C] with padding if needed. In ggml layout: input ne = {C, W, H, B}, output ne = {C, ws, ws, B*nW}. Pad H, W to be divisible by ws. Returns (padded_H, padded_W) in pad_hw. The algorithm (4D ggml only): 1. Reshape [C, Wp, Hp, B] -> [C*ws, nW_w, Hp, B] -- group W into windows 2. Permute (0,2,1,3) -> [C*ws, Hp, nW_w, B] -- move H before nW

//
// Window ordering: n = nw_h + nw_w*nW_h  (H-major, reversed from typical).
static struct ggml_tensor* sam2_window_partition(struct ggml_context* ctx,
                                                  struct ggml_tensor* x,
                                                  int ws, int pad_hw[2]) {
    // ggml layout: ne = {C, W, H, B}
    const int64_t C = x->ne[0];
    const int64_t W = x->ne[1];
    const int64_t H = x->ne[2];
    const int64_t B = x->ne[3];

    // Pad to ws-divisible
    int64_t Hp = ((H + ws - 1) / ws) * ws;
    int64_t Wp = ((W + ws - 1) / ws) * ws;
    pad_hw[0] = (int)Hp;
    pad_hw[1] = (int)Wp;

    if (Hp != H || Wp != W) {
        x = ggml_pad(ctx, x, 0, (int)(Wp - W), (int)(Hp - H), 0);
    }

    int64_t nW_w = Wp / ws;
    int64_t nW_h = Hp / ws;
    int64_t nW = nW_w * nW_h;

    // Step 1: Reshape [C, Wp, Hp, B] -> [C*ws, nW_w, Hp, B]
    auto* r1 = ggml_reshape_4d(ctx, x, C * ws, nW_w, Hp, B);
    // Step 2: Permute to [C*ws, Hp, nW_w, B]
    auto* p1 = ggml_permute(ctx, r1, 0, 2, 1, 3);
    // Step 3: Cont
    auto* c1 = ggml_cont(ctx, p1);
    // Step 4: Reshape [C*ws, Hp, nW_w, B] -> [C*ws, ws, nW_h, nW_w*B]
    //   Splits dim 1 (Hp = ws*nW_h) into (ws=local_H, nW_h=window_H_index)
    auto* r2 = ggml_reshape_4d(ctx, c1, C * ws, ws, nW_h, nW_w * B);
    // Step 5 (direct reshape, no permute):
    //   [C*ws, ws, nW_h, nW_w*B] -> [C, ws, ws, nW*B]
    //   Splits dim 0 (C*ws) into (C, ws=local_W) and merges dims 2,3.
    auto* out = ggml_reshape_4d(ctx, r2, C, ws, ws, nW * B);

    return out;
}

// Window unpartition: reverse of sam2_window_partition.
// Input ne = {C, ws, ws, nW*B}, output ne = {C, orig_W, orig_H, B}.