VAE stage 3: vectorize direct-conv im2col (VAE decode ~4.7s -> ~3.9s, pixel-identical)

- profile the direct conv into im2col / gemm / permute (GGML_CONV_PROFILE=1);
im2col and gemm were each ~half, permute small
- rewrite im2col: hoist sy/sx and the padding check out of the inner loop, and
for the interior (dilation 1, whole knl_w run in bounds, F32 src) copy the
contiguous run directly; hoist the F32/F16 store branch out too
-> im2col 1716 ms -> 875 ms, output diff mean 0 / max 0
- add GGML_CONV_PATCH_MB diagnostic knob (capping the im2col buffer to L3 made
no difference, left unlimited)
- add stages/stage3_vae_conv notes; README + build.sh updated

README.md

@@ -71,6 +71,13 @@ OpenVINO 的 VAE 解碼約 0.55 s，sd.cpp 約 10 s。第一階段嘗試了 conv
 llamafile 版本把 VAE 解碼從 10.7 s 降到 6.6 s（-38%），加上 conv-direct 後為 5.0 s（-53%）。
 詳細數據與重現步驟見 `stages/stage1_vae/README.md`。
 
+## VAE 階段 3：conv2d im2col 向量化（已完成）
+
+用 `GGML_CONV_PROFILE=1` 拆開 direct conv 的三個階段後發現 im2col 與 GEMM 各佔一半。改寫 im2col：
+把 `sy/sx` 與邊界檢查提出迴圈、對齊的內部區塊直接整段複製（`[ic][ky][kx]` 佈局 + 通道最後的來源），
+im2col 由 1716 ms 降到 875 ms，VAE 解碼 ~4.7 → ~3.9 s，輸出逐像素相同。
+剩下的 GEMM（~2.2 s，F32 llamafile sgemm）才是下一關，需要量化/bf16 AMX 卷積。詳見 `stages/stage3_vae_conv/README.md`。
+
 ## Pipeline 階段 P2：refiner 重新量化 + flash-attn tile（已完成）
 
 參考 GGUF 的 refiner 區塊是 Q8_0，是 P1「refiner w13 慢 1.7x」未解之謎的根因。`tools/requant_refiner.py`
@@ -96,6 +103,7 @@ flash-attn 的 V 打包提升為每頭一次、tile 改 128x128，`FLASH_ATTN_EX
 - `z_image_turbo-Q4_K-ffn13.gguf`：由 `tools/fuse_ffn.py` 把 FFN 的 w1、w3 串接成 w13，模型會自動偵測並使用融合路徑。
 - `z_image_turbo-Q4_K-ffn13-rq.gguf`：再由 `tools/requant_refiner.py` 把 refiner 的 16 個 attention/FFN 權重從 Q8_0 改成 Q4_K（速度優先、opt-in，輸出與 base 略有差異）。
 - `GGML_CPU_PROFILE=1`：逐算子計時，輸出到 stderr（僅供分析，預設關閉）。
+- `GGML_CONV_PROFILE=1`：VAE direct-conv 的 im2col/gemm/permute 階段計時（僅供分析，預設關閉）。
 
 ## Profile（Q4_K，優化前）
 
@@ -146,9 +154,11 @@ REPEAT=2 ./SmartTest.sh          # 每列跑 2 次取中位數
   - `ggml-0001-cpu-op-profile.patch`：`ggml-cpu.c`（逐算子計時）
   - `ggml-0002-flash-attn-kpack.patch`：`ops.cpp`（flash-attn K/V 打包提升為每頭一次）
   - `ggml-0003-flash-attn-tiles.patch`：`common.h`（flash-attn tile 128x128）
+  - `ggml-0004-conv2d-im2col.patch`：`ops.cpp`（VAE direct-conv im2col 向量化 + `GGML_CONV_PROFILE`/`GGML_CONV_PATCH_MB` 診斷）
 - `tools/fuse_ffn.py`：把 GGUF 的 FFN w1、w3 串接成 w13（Q4_K/Q8_0 的位元組串接是合法的）。
 - `tools/requant_refiner.py` + `tools/q4k.c`：把 refiner 區塊的 Q8_0 權重重新量化成 Q4_K。
 - `stages/stage_p2_refiner/`：refiner 根因分析、requant 與 flash-attn tile 的數據。
+- `stages/stage3_vae_conv/`：VAE direct-conv 的階段拆解與 im2col 向量化數據。
 - `tools/measure_rss.py`：執行指令並回報行程樹的峰值 RSS 與 wall time。
 - `results/`：`results.csv`、`results.md`、`hardware.txt`、`images/`（每次執行的 512x512 輸出）、`logs/`。
 

build.sh

@@ -141,6 +141,7 @@ if [[ ${#SELECTED[@]} -gt 0 ]]; then
     apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0001-cpu-op-profile.patch"
     apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0002-flash-attn-kpack.patch"
     apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0003-flash-attn-tiles.patch"
+    apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0004-conv2d-im2col.patch"
 
     # ------------------------------------------------------------ build
     for v in "${SELECTED[@]}"; do

patches/ggml-0004-conv2d-im2col.patch

@@ -0,0 +1,157 @@
+diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp
+index bc6918ca..cb67c9b9 100644
+--- a/src/ggml-cpu/ops.cpp
++++ b/src/ggml-cpu/ops.cpp
+@@ -12,6 +12,7 @@
+ #include <algorithm>
+ #include <cfloat>
+ #include <cmath>
++#include <cstdlib>
+ 
+ struct ggml_fp8_e4m3_cpu {
+     uint8_t value;
+@@ -7501,7 +7502,12 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
+     const int64_t patch_total     = dst->ne[3] * dst_w * dst_h;
+ 
+     const int64_t space_per_patch   = knl_n * traits->type_size + c_out * sizeof(float);
+-    const int64_t batch_size        = params->wsize / space_per_patch;
++    int64_t conv_budget             = params->wsize;
++    if (const char * env = getenv("GGML_CONV_PATCH_MB")) {
++        const int64_t cap = (int64_t) atoll(env) * 1024 * 1024;
++        if (cap > 0 && cap < conv_budget) conv_budget = cap;
++    }
++    const int64_t batch_size        = conv_budget / space_per_patch;
+     const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size;
+     const int64_t batch_n           = (patch_total + patches_per_batch - 1) / patches_per_batch;
+ 
+@@ -7509,6 +7515,9 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
+ 
+     void * tmp = params->wdata;
+ 
++    const bool conv_prof = getenv("GGML_CONV_PROFILE") != nullptr;
++    int64_t acc_im2col = 0, acc_gemm = 0, acc_permute = 0;
++
+     for (int64_t batch_i = 0; batch_i < batch_n; ++batch_i) {
+ 
+         const int64_t patch_start_batch = batch_i * patches_per_batch;
+@@ -7516,40 +7525,67 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
+                                               patch_total);
+         const int64_t patch_n           = patch_end_batch - patch_start_batch;
+ 
++        int64_t t_phase = conv_prof ? ggml_time_us() : 0;
+         const int64_t patch_per_thread  = (patch_n + params->nth - 1) / params->nth;
+         const int64_t patch_start       = patch_start_batch + params->ith * patch_per_thread;
+         const int64_t patch_end         = std::min(patch_start + patch_per_thread, patch_end_batch);
+ 
+-        //im2col for a patch
++        // im2col for a patch. Layout is [ic][ky][kx] (dst rows contiguous), and the
++        // common interior case reads knl_w contiguous floats per (ic, ky), so that
++        // run is a straight copy and no per-element bounds check is needed.
++        const int64_t khw     = knl_h * knl_w;
++        const bool    nb0_f32 = src->nb[0] == (int64_t) sizeof(float);
+         for (int64_t p = patch_start; p < patch_end; ++p) {
+-            const int64_t  batch_n     =  p / (dst_w * dst_h);
+-            const int64_t  src_x       = (p / dst_w) % dst_h;
+-            const int64_t  src_y       =  p % dst_w;
++            const int64_t  batch_n =  p / (dst_w * dst_h);
++            const int64_t  src_x   = (p / dst_w) % dst_h;
++            const int64_t  src_y   =  p % dst_w;
+ 
+-            const float * src_base = (const float *)((const char *)src_data + batch_n * src->nb[3]);
+-            char *        dst_row  = (char *) tmp + (p % patches_per_batch) * knl_n * traits->type_size;
++            const char * src_base = (const char *)src_data + batch_n * src->nb[3];
++            char *       dst_row  = (char *) tmp + (p % patches_per_batch) * knl_n * traits->type_size;
++
++            const int64_t sy0 = src_x * stride_y - pad_y;
++            const int64_t sx0 = src_y * stride_x - pad_x;
++            const bool    x_fast = nb0_f32 && dilation_x == 1 && sx0 >= 0 && sx0 + knl_w <= src_w;
+ 
+             for (int64_t ic = 0; ic < c_in; ++ic) {
++                const char * src_ic = src_base + ic * src->nb[2];
++                char *       dst_ic = dst_row + ic * khw * traits->type_size;
+                 for (int64_t ky = 0; ky < knl_h; ++ky) {
+-                    for (int64_t kx = 0; kx < knl_w; ++kx) {
+-                        const int64_t sy = src_x * stride_y + ky * dilation_y - pad_y;
+-                        const int64_t sx = src_y * stride_x + kx * dilation_x - pad_x;
+-
+-                        int64_t dst_idx = ic * (knl_h * knl_w) + ky * knl_w + kx;
++                    const int64_t sy   = sy0 + ky * dilation_y;
++                    char *        drow = dst_ic + ky * knl_w * traits->type_size;
+ 
+-                        float src_val;
+-                        if (sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) {
+-                            src_val = 0.0f;
++                    if (sy < 0 || sy >= src_h) {
++                        if (kernel_type == GGML_TYPE_F32) {
++                            memset(drow, 0, knl_w * sizeof(float));
+                         } else {
+-                            const float * src_ptr = (const float *)((const char *)src_base + sx * src->nb[0] + sy * src->nb[1] + ic * src->nb[2]);
+-                            src_val               = *src_ptr;
++                            ggml_fp16_t * dp = (ggml_fp16_t *) drow;
++                            for (int64_t kx = 0; kx < knl_w; ++kx) dp[kx] = 0;
+                         }
++                        continue;
++                    }
++
++                    const char * src_row = src_ic + sy * src->nb[1];
+ 
+-                        char * element_ptr = dst_row + dst_idx * traits->type_size;
++                    if (x_fast) {
++                        const float * sp = (const float *) src_row + sx0;
+                         if (kernel_type == GGML_TYPE_F32) {
+-                            *(float *) element_ptr = src_val;
+-                        } else if (kernel_type == GGML_TYPE_F16) {
+-                            *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val);
++                            memcpy(drow, sp, knl_w * sizeof(float));
++                        } else {
++                            ggml_fp16_t * dp = (ggml_fp16_t *) drow;
++                            for (int64_t kx = 0; kx < knl_w; ++kx) dp[kx] = GGML_CPU_FP32_TO_FP16(sp[kx]);
++                        }
++                    } else {
++                        for (int64_t kx = 0; kx < knl_w; ++kx) {
++                            const int64_t sx = sx0 + kx * dilation_x;
++                            float src_val = 0.0f;
++                            if (sx >= 0 && sx < src_w) {
++                                src_val = *(const float *)(src_row + sx * src->nb[0]);
++                            }
++                            if (kernel_type == GGML_TYPE_F32) {
++                                ((float *) drow)[kx] = src_val;
++                            } else {
++                                ((ggml_fp16_t *) drow)[kx] = GGML_CPU_FP32_TO_FP16(src_val);
++                            }
+                         }
+                     }
+                 }
+@@ -7558,6 +7594,8 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
+ 
+         ggml_barrier(params->threadpool);
+ 
++        if (conv_prof) { const int64_t t = ggml_time_us(); acc_im2col += t - t_phase; t_phase = t; }
++
+         float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n * traits->type_size);
+ 
+         GGML_ASSERT(gemm_output + patch_n * c_out <= (float*)tmp + params->wsize);
+@@ -7567,6 +7605,8 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
+ 
+         ggml_barrier(params->threadpool);
+ 
++        if (conv_prof) { const int64_t t = ggml_time_us(); acc_gemm += t - t_phase; t_phase = t; }
++
+ 
+         //permute back [OC, N, OH, OW] to [N, OC, OH, OW]
+         const int64_t permute_per_thread = (patch_n + params->nth - 1) / params->nth;
+@@ -7585,6 +7625,14 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
+                 *dst_ptr = value;
+             }
+         }
++        if (conv_prof) { const int64_t t = ggml_time_us(); acc_permute += t - t_phase; }
++    }
++
++    if (conv_prof && params->ith == 0) {
++        fprintf(stderr, "CONVPROF out=%lldx%lld in_c=%lld k=%lldx%lld im2col_us=%lld gemm_us=%lld permute_us=%lld\n",
++                (long long) dst_w, (long long) dst_h, (long long) c_in,
++                (long long) knl_w, (long long) knl_h,
++                (long long) acc_im2col, (long long) acc_gemm, (long long) acc_permute);
+     }
+ }
+ 

stages/stage3_vae_conv/README.md

@@ -0,0 +1,64 @@
+# Stage 3 — VAE conv2d im2col (sd.cpp direct conv)
+
+Follows `stage1_vae`. After stage 1 the VAE decode (512px, `llamafile` + `conv-direct`)
+is ~4.3 s. This stage profiles the direct conv internally and removes the scalar
+im2col overhead.
+
+Config: `llamafile` build, `SD_CPU_EXTRA_BUFT=1`, `--conv-direct auto,vae=on`,
+Q4_K-ffn13-rq, 512x512, 8 steps.
+
+## 1. Phase split inside `ggml_compute_forward_conv_2d_impl`
+
+`GGML_CONV_PROFILE=1` (env-gated, added in `patches/ggml-0004-conv2d-im2col.patch`)
+prints per-call `im2col` / `gemm` / `permute` microseconds. Summed over all VAE convs,
+1-step run:
+
+| phase | before | after |
+|---|---|---|
+| im2col (patch build) | 1716 ms | **875 ms** |
+| GEMM (llamafile sgemm) | 2179 ms | 2238 ms |
+| permute (OC-scattered store) | 202 ms | 216 ms |
+| **VAE decode** | **~4.7 s** | **~3.9 s** |
+
+## 2. What changed in im2col
+
+The old inner loop recomputed `sy`/`sx` and the padding bounds check **per element**
+(`ic × ky × kx`), with a `kernel_type` branch inside the innermost iteration. The
+patch layout is `[ic][ky][kx]` (dst rows contiguous) and the source is channel-last,
+so for the common interior case each `(ic, ky)` reads `knl_w` **contiguous** floats.
+The rewrite:
+
+- hoists `sy0 = src_x*stride_y - pad_y`, `sx0 = src_y*stride_x - pad_x` per patch;
+- adds a fast path (dilation 1, whole `knl_w` run in-bounds, F32 source) that copies
+  the run directly, no per-element bounds check;
+- keeps a general fallback for padding / dilation / non-F32 source;
+- hoists the F32/F16 store branch out of the innermost loop.
+
+Output is **pixel-identical** to the previous build (diff mean 0, max 0).
+
+The same patch adds two diagnostic env knobs (no effect when unset):
+`GGML_CONV_PROFILE=1` (phase timing) and `GGML_CONV_PATCH_MB=<n>` (cap the im2col
+batch buffer; capping to L3 size made no difference, so it is left unlimited).
+
+## 3. Remaining VAE cost
+
+The GEMM is now the largest phase (~2.2 s, ~0.7 TFLOPS effective) and runs through the
+F32 llamafile sgemm. im2col is ~0.9 s. OpenVINO still decodes this VAE in ~0.3 s
+because it runs INT8 AMX. Closing that gap needs a quantized/bf16 conv GEMM, not more
+im2col tuning:
+
+- flatten the 3x3 conv weights to a 2-D `[K=KW*KH*IC, N=OC]` layout at load time so
+  they can live in the AMX weight buffer, then reuse the AMX int8/bf16 GEMM;
+- or a Winograd F(2x2,3x3) transform to cut both the multiply count and the 9x im2col
+  expansion.
+
+## Reproduce
+
+```bash
+./build.sh llamafile    # includes patches/ggml-0004-conv2d-im2col.patch
+GGML_CONV_PROFILE=1 SD_CPU_EXTRA_BUFT=1 work/bin/llamafile/sd-cli \
+  --diffusion-model work/models/z_image_turbo-Q4_K-ffn13-rq.gguf \
+  --vae work/models/ae.safetensors --llm work/models/Qwen3-4B-Instruct-2507-Q4_K_M.gguf \
+  -p "a cat" -W 512 -H 512 --steps 1 --seed 42 -t 16 \
+  --conv-direct auto,vae=on -o out.png
+```