A specialized Winograd Conv2d op #971

bssrdf · 2024-09-29T00:30:38Z

This PR added a new conv2d op using Winograd algorithm.

Currently ggml's conv2d operator uses im2col and GEMM. There have been efforts to speed up this process using other faster algorithms. Winograd is such a method used by many neural network libraries, e.g. Cudnn . For small kernels, e.g. 3x3, Winograd outperforms GEMM based methods. However, efficient implementation of Winograd on GPUs requires significant engineering efforts. This PR 's Winograd implementation specializes in several ways:

It only supports 3x3 kernel
It only supports channel numbers of multiples of 8
It only supports filter output numbers of multiples of 64
It only supports stride=1, padding = 1 and dilation = 1
It only supports CUDA backend

Other features:

Fused, except the kernel transform which requires additional workspace (can be shared with the weights)

It is mainly used for applications such as stable-diffusion.cpp.

The code is based on openCNN project which uses Apache-2.0 license.

Please review and let me know any problems I'll address. Thanks.

… if not satisfied

…results

…FP32 buffer

JohannesGaessler

The code is based on openCNN project which uses Apache-2.0 license.

Did you get permission from the authors to re-license their code as MIT?

src/ggml-cuda.cu

src/ggml-cuda/conv-winograd.cu

JohannesGaessler · 2024-09-29T16:31:52Z

src/ggml-cuda/conv-winograd.cu

+  typedef float(*pointFunction_t)(float *, int);
+
+  template<typename T>
+  __global__ void FX(const T *pInputs, float *pOutputs, int filt_k, 


Suggested change

__global__ void FX(const T *pInputs, float *pOutputs, int filt_k,

__global__ void FX(const T * __restrict__ pInputs, float * __restrict__ pOutputs, int filt_k,

On Pascal this can be a 5x speedup.

JohannesGaessler · 2024-09-29T16:40:05Z

src/ggml-cuda/conv-winograd.cu

+
+}
+
+__device__ __forceinline__ void prefetch_filter_tile(const float *pInputs, float *tiles, int filt_k){


The compiler will rearrange these instructions as it sees fit so there will in effect not be any actual prefetching. For that you need to use asnychronous memcpys (Ampere or newer).

I am not sure here, as this is done by openCNN.

JohannesGaessler · 2024-09-29T16:43:32Z

src/ggml-cuda/conv-winograd.cuh

+__constant__ int access_f_s[2][32];
+__constant__ int access_s[2][32];
+__constant__ int tileid[2][32];


What happens in the case of multiple GPUs? Is the constant memory duplicated across GPUs?

I am pretty ignorant about multi-gpu. I guess they will be duplicated. I don't have a setup to test. Plus, this kernel only works for single GPU, I think.

JohannesGaessler · 2024-09-29T16:45:29Z

src/ggml.c

+static void ggml_compute_forward_winograd_stage0(
+        const struct ggml_compute_params * params,
+              struct ggml_tensor * dst) {
+
+    GGML_ASSERT(false && " CPU backend not implemented!");         
+    return;
+}
+
+static void ggml_compute_forward_winograd_stage1(
+        const struct ggml_compute_params * params,
+              struct ggml_tensor * dst) {
+
+    GGML_ASSERT(false && " CPU backend not implemented!");         
+    return;
+}


If at all possible a CPU implementation should always be done since it serves both as a fallback and as a reference implementation to test other backends against.

A CPU backend should be done, but I am not sure the benefit of it compared to the current im2col+gemm version.

JohannesGaessler · 2024-09-29T16:46:41Z

src/ggml.c

+    bool is_node = false;
+
+    if (a->grad) {
+        is_node = true;
+    }


If #966 is merged first this will need to be removed (should be very straightforward).

Look forward to it...

JohannesGaessler · 2024-09-29T17:16:18Z

Have you done any tests regarding performance? This code does not use tensor cores at all so intuitively I would expect it to be slower than im2col + GEMM with tensor cores.

bssrdf · 2024-09-29T21:19:11Z

Have you done any tests regarding performance? This code does not use tensor cores at all so intuitively I would expect it to be slower than im2col + GEMM with tensor cores.

Thank you for your review, @JohannesGaessler. I leaned a lot from your PRs and comments.

First, I have asked openCNN's author for license issue.

As to performance, I only tested it in SD.cpp as it is developed for it. It is not faster (surprised) than im2col+GEMM with tensor cores (my gpu has them so assuming being used) but definitely not slower. It reduces memory used by VAE quite a lot while increasing UNET param buffer. There is room to further improve its performance as I see several places are not working in an optimal way.

I'll add test cases in test-backend-ops to more rigorously measure performance.

I addressed your other comments above.

slaren · 2024-09-29T21:36:15Z

Have you tried NPP? It is a library bundled with the CUDA toolkit that has all kinds of kernels for image processing. I don't think this can be merged unless the license situation is resolved.

JohannesGaessler · 2024-09-29T21:46:59Z

Generally speaking my stance regarding this PR would be as follows: I think it's good to have convolution operations instead of having to rely on IM2COL. At the same time I want to have a codebase that is easy to maintain - a central factor for me is that there needs to be some benefit for adding code that offsets the increase in maintenance effort. Quite honestly I think that the starting point from OpenCNN is not very good; I would be rather hesitant to add it since the use cases are limited and I think none of the devs on this project would have a very good understanding of how the code works.

And as slaren said, the licensing issue must be resolved or this is a total non-starter anyways.

Have you tried NPP? It is a library bundled with the CUDA toolkit that has all kinds of kernels for image processing.

From what I can tell, there is convolution support.

JohannesGaessler · 2024-09-29T22:03:20Z

I'm already tired so maybe I'm just misreading the docs, but I get the impression that NPP convolutions only support 1-4 input channels.

bssrdf · 2024-09-29T23:53:55Z

Thanks to both of you for reviewing. I am not familiar with the license. In case it is not resolvable, I'll ditch this PR.
Now I am putting it in draft mode, hoping to make it work in more general settings, truly serving as an alternative to im2col approach.

JohannesGaessler · 2024-09-30T06:46:38Z

Also one important question that I forgot to ask: are you going to be available long-term to maintain this code?

bssrdf · 2024-09-30T16:26:55Z

Also one important question that I forgot to ask: are you going to be available long-term to maintain this code?

If this PR makes into the main, I intend to maintain it long term and improve its performance.

bssrdf · 2024-10-04T12:29:18Z

I am getting some puzzling benchmark results. @JohannesGaessler or @slaren, could you help?

I added a conv2d and a winograd case in test-backend-ops, using its eval_perf function to gauge the performance.

Winograd is consistently 10x slower than conv2d(IM2COL++GEMM). However, the same winograd kernel in stable-diffusion improves the running time by 7-8% for UNET and almost 50% for VAE. So I am confused.

What's even more interesting is the same input size for IM2COL itself runs slower than conv2d which has a GEMM after IM2COL.

Maybe the way I tested is not proper for conv2d since it is a composite operator unlike others?

slaren · 2024-10-04T13:15:16Z

Maybe the way I tested is not proper for conv2d since it is a composite operator unlike others?

Yes, this is the reason. The perf mode works by running multiple times the last node of the graph only, everything before that is assumed to be setup for the op and only run once.

bssrdf · 2024-10-04T13:56:01Z

Maybe the way I tested is not proper for conv2d since it is a composite operator unlike others?

Yes, this is the reason. The perf mode works by running multiple times the last node of the graph only, everything before that is assumed to be setup for the op and only run once.

Wow, now I see. Turned out that the conv2d case is just testing the cuda_dup op which is the last node. I need to find another way to compare. Thanks.

Green-Sky · 2024-10-12T17:48:53Z

@bssrdf just took your sd.cpp fork for a spin, and I like the numbers.
for flux's vae, it shrinks vram usage roughly in half, and gives a small speed bump.
before : 94+1664 MB ~0.75s
after : 258+704 MB ~0.53s

with sd2 the unet changes seem to make it use almost double the vram for the model?
vae wins here again though (almost the same values as flux's vae).

images where 512x512.

bssrdf · 2024-10-12T19:19:57Z

@bssrdf just took your sd.cpp fork for a spin, and I like the numbers. for flux's vae, it shrinks vram usage roughly in half, and gives a small speed bump. before : 94+1664 MB ~0.75s after : 258+704 MB ~0.53s

with sd2 the unet changes seem to make it use almost double the vram for the model? vae wins here again though (almost the same values as flux's vae).

images where 512x512.

@Green-Sky, thanks for trying out. Right now, winograd is not ready for SD.cpp due to duplicated UNET and VAE parameter buffer. My winograd implementation has filter weight transform in a separate kernel (inherited from openCNN) so needs a buffer for that. I added filter transform in initialization stage so during unet denoising this step can be skipped (just use the buffer). This is one of the benefits of using winograd (a small time saving though). However, the original/untransformed filter weight still occupies VRAM as they are allocated earlier and there is no mechanism in GGML to release them. I have some ideas to fix this so no VRAM duplication happens; but still the parameter buffers for UNET and VAE will increase in size because the transformed weights are C*K*4*4 and the untransformed are C*K*3*3. The compute buffer sizes are reduced by quite a lot, as you found out, because of no IM2COL. In my tests, VAE runs twice as fast. If you, @FSSRepo or @leejet could help with this, that will be great.

In the meantime, I am adding tensor core support to make it even faster. I think it could be faster by at least 50% once the tensor cores are used.

FSSRepo · 2024-10-19T04:40:47Z

Hello, good luck trying to use the tensor cores to get Winograd working; I think it’s possible. I saw an implementation that used matrix multiplication in Tinygrad. I found a Winograd implementation that omitted redundant calculations, but for that same reason, it wasn't possible to use the tensor cores. In the end, I couldn't improve the performance. After all, I decided to better merge im2col and GEMM (which reduces memory consumption by 50%, but is 40% slower in the VAE phase, which is primarily the one that uses the most memory).

bssrdf · 2024-10-19T19:57:59Z

Hello, good luck trying to use the tensor cores to get Winograd working; I think it’s possible. I saw an implementation that used matrix multiplication in Tinygrad. I found a Winograd implementation that omitted redundant calculations, but for that same reason, it wasn't possible to use the tensor cores. In the end, I couldn't improve the performance. After all, I decided to better merge im2col and GEMM (which reduces memory consumption by 50%, but is 40% slower in the VAE phase, which is primarily the one that uses the most memory).

Thanks for the comments. Right now, using tensor cores, my winograd kernel outperforms cudnn's winograd (not using tensor cores) by 15% for certain inputs (large input size and number of channels and filters). Well, dealing with the bank conflicts are real pain and I have to program in PTX 😄. For such large inputs, the king is cudnn's IMPLICIT_PRECOMP_GEMM using tensor cores; it is 50% faster than my winograd. I hope to get more performance boost once the final "Uncoalesced Shared Accesses" issue is resolved.

JohannesGaessler · 2024-10-19T20:06:05Z

Be aware that if you are using the nvcuda::wmma interface the data layout and memory access pattern are not defined. So your code could end up having bad performance on CUDA architectures that you are not testing. For FP16 the memory access pattern for all GPUs ranging from Turing to Ada Lovelace should be the same though. My recommendation would be to write simple primitives that wrap the tensor core PTX instructions instead. I did an implementation of this for int8 in mma.cuh. This would then give you a defined memory layout and enable you to manipulate the data without having to go through shared memory.

Edit: judging by the edit to your previous comment you are already using PTX.

JohannesGaessler · 2024-10-19T20:08:44Z

For avoiding shared memory bank conflicts with tensor cores, consider padding the data. For the int8 tensor cores I found that a padding of 16 bytes between rows/columns avoids shared memory bank conflicts when loading data from shared memory for use with tensor cores.

bssrdf · 2024-10-19T20:13:13Z

For avoiding shared memory bank conflicts with tensor cores, consider padding the data. For the int8 tensor cores I found that a padding of 16 bytes between rows/columns avoids shared memory bank conflicts when loading data from shared memory for use with tensor cores.

Thanks for the tips, @JohannesGaessler. Yes, I used padding in several places and they really helped with resolving conflicts. I also played around with swizzling but in the end couldn't get it working for the smem layout.

JohannesGaessler · 2024-10-19T20:21:56Z

Also take a look at the ldmatrix instruction. While it doesn't reduce shared memory bank conflicts it does help by issuing fewer instructions which helped with one of the throttle reasons (I forgot which one). But I didn't find it to be universally better than just loading the data as 4 bit values from shared memory.

bssrdf · 2024-10-19T20:39:07Z

Also take a look at the ldmatrix instruction. While it doesn't reduce shared memory bank conflicts it does help by issuing fewer instructions which helped with one of the throttle reasons (I forgot which one). But I didn't find it to be universally better than just loading the data as 4 bit values from shared memory.

I see people used ldmatrix in some GEMM kernels and it didn't help much, so I didn't bother. Yeah, good to know it though.

FSSRepo · 2024-10-20T01:13:59Z

@JohannesGaessler I know it's out of context, but I'm compiling the latest version of stable diffusion.cpp and now it's taking more than 25 minutes to compile the CUDA code. Before (four months ago, I was busy trying to figure out what to do with my life), it took at most 5 minutes or less, and binaries now are more bigger.

JohannesGaessler · 2024-10-20T08:05:14Z

That's probably due to the MMQ changes. For development builds you can edit src/CMakeLists.txt and set CMAKE_CUDA_ARCHITECTURES to only the one that you're actually using.

bssrdf added 10 commits September 26, 2024 18:04

added two winograd ops

5b4e448

added more checking for precondition to use winograd; switch to im2co…

68c251b

… if not satisfied

added source code for winograd kernel

2ccc67d

winograd build ok

02a3cb1

test now passed; for some reason, ggml_conv_2d didn't output correct …

3d80466

…results

remove debugging printouts

893ca79

added a FP16 FX kernel to deal with fp16 filter data; no need to use …

6afbf6e

…FP32 buffer

sync from sd.cpp

0491858

restore test-conv2d.cpp test

93c3da7

restore src/CMakeLists.txt

4e8e0d4

JohannesGaessler reviewed Sep 29, 2024

View reviewed changes

bssrdf marked this pull request as draft September 29, 2024 17:40

bssrdf added 2 commits September 29, 2024 17:09

change mask to unsigned int; add __restric__ to various pointers

e0e94c4

fix indentation

c8700ca

bssrdf added 3 commits September 30, 2024 20:41

skip if already computed in a preprocess step

c5d43a2

added winograd conv2d to backend op tests

00ad37e

add conv2d as a test case in test-backend-op

4f93d67

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

A specialized Winograd Conv2d op #971

A specialized Winograd Conv2d op #971

bssrdf commented Sep 29, 2024 •

edited

Loading

JohannesGaessler left a comment

JohannesGaessler Sep 29, 2024

JohannesGaessler Sep 29, 2024

bssrdf Sep 29, 2024

JohannesGaessler Sep 29, 2024

bssrdf Sep 29, 2024

JohannesGaessler Sep 29, 2024

bssrdf Sep 29, 2024

JohannesGaessler Sep 29, 2024

bssrdf Sep 29, 2024

JohannesGaessler commented Sep 29, 2024

bssrdf commented Sep 29, 2024 •

edited

Loading

slaren commented Sep 29, 2024

JohannesGaessler commented Sep 29, 2024

JohannesGaessler commented Sep 29, 2024 •

edited

Loading

bssrdf commented Sep 29, 2024

JohannesGaessler commented Sep 30, 2024

bssrdf commented Sep 30, 2024

bssrdf commented Oct 4, 2024

slaren commented Oct 4, 2024

bssrdf commented Oct 4, 2024

Green-Sky commented Oct 12, 2024

bssrdf commented Oct 12, 2024 •

edited

Loading

FSSRepo commented Oct 19, 2024

bssrdf commented Oct 19, 2024 •

edited

Loading

JohannesGaessler commented Oct 19, 2024 •

edited

Loading

JohannesGaessler commented Oct 19, 2024 •

edited

Loading

bssrdf commented Oct 19, 2024

JohannesGaessler commented Oct 19, 2024

bssrdf commented Oct 19, 2024

FSSRepo commented Oct 20, 2024 •

edited

Loading

JohannesGaessler commented Oct 20, 2024

	__global__ void FX(const T pInputs, float pOutputs, int filt_k,
	__global__ void FX(const T * __restrict__ pInputs, float * __restrict__ pOutputs, int filt_k,


		}

		__device__ __forceinline__ void prefetch_filter_tile(const float pInputs, float tiles, int filt_k){

A specialized Winograd Conv2d op #971

Are you sure you want to change the base?

A specialized Winograd Conv2d op #971

Conversation

bssrdf commented Sep 29, 2024 • edited Loading

JohannesGaessler left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JohannesGaessler commented Sep 29, 2024

bssrdf commented Sep 29, 2024 • edited Loading

slaren commented Sep 29, 2024

JohannesGaessler commented Sep 29, 2024

JohannesGaessler commented Sep 29, 2024 • edited Loading

bssrdf commented Sep 29, 2024

JohannesGaessler commented Sep 30, 2024

bssrdf commented Sep 30, 2024

bssrdf commented Oct 4, 2024

slaren commented Oct 4, 2024

bssrdf commented Oct 4, 2024

Green-Sky commented Oct 12, 2024

bssrdf commented Oct 12, 2024 • edited Loading

FSSRepo commented Oct 19, 2024

bssrdf commented Oct 19, 2024 • edited Loading

JohannesGaessler commented Oct 19, 2024 • edited Loading

JohannesGaessler commented Oct 19, 2024 • edited Loading

bssrdf commented Oct 19, 2024

JohannesGaessler commented Oct 19, 2024

bssrdf commented Oct 19, 2024

FSSRepo commented Oct 20, 2024 • edited Loading

JohannesGaessler commented Oct 20, 2024

bssrdf commented Sep 29, 2024 •

edited

Loading

bssrdf commented Sep 29, 2024 •

edited

Loading

JohannesGaessler commented Sep 29, 2024 •

edited

Loading

bssrdf commented Oct 12, 2024 •

edited

Loading

bssrdf commented Oct 19, 2024 •

edited

Loading

JohannesGaessler commented Oct 19, 2024 •

edited

Loading

JohannesGaessler commented Oct 19, 2024 •

edited

Loading

FSSRepo commented Oct 20, 2024 •

edited

Loading