Support groups in DenseConvDims

[DhairyaLGandhi]

Also propagate it to the kernel calls.

[CarloLucibello]

Our cpu conv kernels don't seem to support groups, do they?

[CarloLucibello]

Our cpu conv kernels don't seem to support groups, do they?

@DhairyaLGandhi can I get any answers here? Which conv implementation supports groups? Can you point to some specific line of code?

[CarloLucibello]

ah you are implementing groups support right now, don't have to be that mysterious :)

[DhairyaLGandhi]

So 115784b is a "working" version which is very memory wasteful, but I have it here right now in case someone was eager to see it working in NNlib. I will be moving to the non-allocating version of this soon. The issue is Julia doesn't like to look into the stride of deep and lazy wrapper types around SubArray, ReshapedArray etc and can cause races, or return garbage. In the previous version, we had segfaults because the DenseConvDims needs to be extra careful about the encoding dims that include the groups.

[DhairyaLGandhi]

For the record, not using tasks is as fast as the original version for the regular convolutions (ie with groups = 1), and maybe the serial version is faster anyway.

[DhairyaLGandhi]

Oh this is satisfying. Starting Julia with -t2 on my tiny MacBook Pro with a dual core. For context this results in about a 20% reduction compared to doing the sequential thing. It might be safer to not do that, but RFC?

julia> @btime y = conv($(rand(Float32, 28, 28, 100, 2)), $(rand(Float32, 3, 3, 20, 15)), groups = 5);
  947.722 μs (159 allocations: 4.81 MiB)

Either way we now have the system in place to pass it to CUDNN. Should I also add a stub to DepthwiseConvDims for consistency? I might do that actually.

[CarloLucibello]

For context this results in about a 20% reduction compared to doing the sequential thing. It might be safer to not do that, but RFC?

I don't see any safety issue in going multithreaded here, computations are independent.

We should check that using threads doesn't have a performance impact on the standard convolution case (G=1).

[DhairyaLGandhi]

Its a question of computational complexity. The existing cases aren't affected much at all.

[CarloLucibello]

I think the existing cases are affected because spawning threads doesn't come for free, but hopefully could be negligible compared to the convolution cost. If that's not the case, we can still dispatch on the G=1 case.

In any case, we should check we don't have performance regressions in standard convolution before merging
(+ the obvious stuff: existing tests are not passing, new tests are needed, docstring should be updated)

[DhairyaLGandhi]

Way ahead of you :)

[darsnack]

For context this results in about a 20% reduction compared to doing the sequential thing.

🎉

[DhairyaLGandhi]

Addressed all the comments. I've also added more tests and docs. @pxl-th could you rebase your NNlibCUDA PR? Any other comments?

[CarloLucibello]

not sure this respecting the fact that channels_in(cdims) == size(x, 3) and channels_out(cdims) == size(y, 3)

[DhairyaLGandhi]

Done, also made channels_out respect groups. I think storing the groups is needed anyway along with the regular C_out (since we ultimately use the G = 1 definition).

[DhairyaLGandhi]

Ah, almost forgot - we need channels_out to mean the same as master - since that data is needed. See ∇conv_data!

[DhairyaLGandhi]

I've addressed the reviews - thanks @staticfloat! I had a couple of debugging steps that I've cleaned up now. Also comes with more tests, and docs!

Re the added check - its there so we can give a clearer message about groups config.

Re threading - it allows using the newer scheduler which reasons about in terms of tasks, so its pretty smart about how to handle the resources. It is composable too https://julialang.org/blog/2019/07/multithreading/

[DhairyaLGandhi]

@pxl-th do you think you can test your NNlibCUDA PR off of this one? The API is mostly the same, but good to get a check that we are passing into CUDNN properly.

[DhairyaLGandhi]

Let's merge this. I'll wait till tomorrow, so if anyone has final thoughts we can address those.

[pxl-th]

@DhairyaLGandhi I've tested it. Everything seems to work, all the tests are passing.

[DhairyaLGandhi]

Nice, let's get these PRs merged then.