GGML direct conv2d support

[netrunnereve]

Since GGML now has direct conv2d support for CPU and Vulkan we might want to try it out here and see if it helps. Compared to im2col this uses less memory and should run faster on GPUs that don't have matrix cores.

As a quick test I naively switched all instances of ggml_conv_2d in the code with ggml_conv_2d_direct and replaced the ggml directory with the one from llama.cpp. Right now it generates images fine on CPU (it's a bit slower than im2col) but it fails with a segfault on Vulkan.

[Green-Sky]

Upstream was just synced with the vulkan code. @leejet just missed the update by a few hours 😆

edit: There is also support for opencl. I kinda miss cuda support...

[Green-Sky]

Just tried the vulkan backend and yes, it sadly fails.

    case GGML_OP_CONV_2D:
        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 &&
            ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) {
            return ctx->device->pipeline_conv2d_f32;
        }
        return nullptr;

I think the inputs are contiguous, but they are not all f32...

(gdb) p *src0
$5 = {type = GGML_TYPE_F16,

(gdb) p *src1
$6 = {type = GGML_TYPE_F32,

[Green-Sky]

I hacked together a patch for ggml. Will post numbers later.

ggml-org/llama.cpp#14872

[daniandtheweb]

I've just tried the changes on my laptop and unfortunately there's no gains for "old" APUs. The ryzen 7 4700u on my laptop using vulkan for a 512*512 image on SD 1.5 goes from 3.78s/it on master to 4.24s/it with the new direct implementation.
I replaced all instances of ggml_conv_2d with their direct counterpart and pasted the ggml folder from the PR.
I hope that at least discrete GPUs can get a speedup from this.

[Green-Sky]

Yes, I don't expect it to be faster in vulkan, since it is a scalar implementation. Especially when you compare it against im_col+mul with cm1/cm2.

[daniandtheweb]

No coopmat at all on the 4700U, it's a small Vega GPU that uses some ram as vram, so it's going to require a lot of optimization to have some gains. Pytorch takes about half of the time to generate so there's still some margin to improve.

[Green-Sky]

Oh I see.

SD1 (CyberRealistic_V9_FP16)

512x768

method	diffusion buffer	sampling speed	vae buffer	vae speed
direct	1219.65mb	1.11s/it	960mb	1.15s
direct w fa	1217.87mb	1.11s/it	960mb	1.16s
im2col+mul	1220.07mb	1.01it/s	2496mb	2.08s
im2col+mul w fa	1218.29mb	1.01it/s	2496mb	2.10s

Quick testing reveals big wins for VAE, I was sure the memory would go down alot, but I actually did not expect the VAE to be faster too.

But yes, sampling is ~10% slower in this test on my RTX 2070 mobile.

edit: compiled with coopmat1

[stduhpf]

VAE is very slow on Vulkan right now, and the compute buffer very easily exceeds the 4GB allocation limit. For that alone I think these changes are worth the slower diffusion.

[Green-Sky]

direct	im2col

Only a slight difference too.

[Green-Sky]

VAE is very slow on Vulkan right now, and the compute buffer very easily exceeds the 4GB allocation limit. For that alone I think these changes are worth the slower diffusion.

We don't need to switch it for the diffuser. We can make the VAE use direct and the diffusion model continue use im2col. At least until it becomes better.

[daniandtheweb]

VAE decoding goes from 20s on master to 6s with these changes on my laptop so I also agree that it's better only implementing it only on the VAE for now.

[Green-Sky]

flux.1 light

(hyper-8step lora pruned and applied)

f16 (v)ae

768x768

method	diffusion buffer	sampling speed	vae buffer	vae speed
direct	1105.07mb	5.47s/it	1440mb	1.74s
direct fa	505.07mb	4.98s/it	1440mb	1.77s
im2col+mul	1105.07mb	5.46s/it	3744mb	3.34s
im2col+mul fa	505.07mb	4.92s/it	3744mb	3.34s

Flux is expectedly not affected by convolution in the diffuser.
VAE behaves the same.

Bonus: vae cpu numbers

8 threads

method	time
direct	74.47s
im2col+mul fa	71.93s

[rmatif]

I've also added the conv2d op to OpenCL: ggml-org/ggml@7ce741c it seems the previous sync missed it, as this was upstreamed to ggml just a couple of hours after the sync

[daniandtheweb]

I've never touched the vae.hpp file before but changing the Conv2d operations in there to Conv2dDirect operations (leaving the diffusion step with the old Conv2d) doesn't perform like changing every operation in ggml-extend.hpp to Conv2dDirect (10s for vae decoding vs 6s when changing everything). Does anyone know if there's any other call do Conv2d in the vae process that I'm missing other than the ones in vae.hpp?

EDIT: Found it in common.hpp

[netrunnereve]

With a q6_K SDXL Lightning on my RX 470 generating a 512x512 image:

Direct:

  |==================================================| 4/4 - 1.09s/it
[INFO ] stable-diffusion.cpp:1806 - sampling completed, taking 4.37s
[INFO ] stable-diffusion.cpp:1814 - generating 1 latent images completed, taking 4.37s
[INFO ] stable-diffusion.cpp:1817 - decoding 1 latents
[INFO ] stable-diffusion.cpp:1827 - latent 1 decoded, taking 1.40s
[INFO ] stable-diffusion.cpp:1831 - decode_first_stage completed, taking 1.40s
[INFO ] stable-diffusion.cpp:2079 - generate_image completed in 6.04s

Indirect:

  |==================================================| 4/4 - 1.01it/s
[INFO ] stable-diffusion.cpp:1806 - sampling completed, taking 4.01s
[INFO ] stable-diffusion.cpp:1814 - generating 1 latent images completed, taking 4.01s
[INFO ] stable-diffusion.cpp:1817 - decoding 1 latents
[INFO ] stable-diffusion.cpp:1827 - latent 1 decoded, taking 6.91s
[INFO ] stable-diffusion.cpp:1831 - decode_first_stage completed, taking 6.91s
[INFO ] stable-diffusion.cpp:2079 - generate_image completed in 11.19s

This is great. The direct version uses around 1GB less memory during decoding and it's way faster there, while at the same time it's a bit slower at sampling. So yeah it's probably a good idea to use the direct mode only for the vae.