Skip to content

Optimize UNet inference: enable LLAMAFILE SGEMM + fuse SpatialTransformer 1x1 convolutions - #3

Draft
cmdr2 with Copilot wants to merge 3 commits into
unet-perffrom
copilot/optimize-unet-inference
Draft

cmdr2 with Copilot wants to merge 3 commits into
unet-perffrom
copilot/optimize-unet-inference

Conversation

Copilot AI commented Mar 31, 2026

Copy link
Copy Markdown

UNet inference optimized from ~2463ms to ~1332ms per step (~1.85x speedup) on CPU with FP16 weights.

Enable LLAMAFILE optimized SGEMM by default

GGML_LLAMAFILE_DEFAULT was OFF, leaving llamafile's tinyBLAS (cache-blocked AVX2+F16C matrix multiply) unused. Set it ON before add_subdirectory(ggml). This alone accounts for ~1.7x of the speedup since MUL_MAT dominates (~172 GFLOPS across 282 ops in the graph).

Fuse SpatialTransformer 1x1 Conv2d with permutes

For SD1's !use_linear path, proj_in/proj_out are 1x1 Conv2d ops. The original flow:

Conv2d_1x1 → [im2col + mul_mat + permute + cont] → permute + cont → reshape → [transformer] → permute + cont → reshape → Conv2d_1x1 → [im2col + mul_mat + permute + cont]

Now both paths (linear and conv-1x1) share a unified flow — permute once, reshape, apply weight as a linear op via ggml_mul_mat on the reshaped 2D weight, then permute back once:

permute + cont → reshape → mul_mat(w_2d, x) → [transformer] → mul_mat(w_2d, x) → permute + cont → reshape

This eliminates 2× im2col + 2× permute+cont per SpatialTransformer (16 in SD1.5 UNet = 32 fewer data copies).

Minor: add get_params() accessor to GGMLBlock

Needed to access Conv2d weight tensors directly for the reshaped mul_mat path.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants