[CK tests] Extend conv GPU reference

[johannes-graner]

Proposed changes

This PR extends GPU reference implementation support for convolution operations with elementwise fusion and output operations. The changes enable GPU-accelerated reference implementations for tests involving scale, bias, batchnorm, clamp, and bilinear operations across forward, backward data, and backward weight convolutions.

Key improvements:

• Extended naive_conv_fwd_gpu, naive_conv_bwd_data_gpu, and naive_conv_bwd_weight_gpu to support elementwise operations, clamp, scale, and bilinear fusion
• Significantly improved test execution times across 11 test suites

Performance impact:

Test Name	Before (s)	After (s)	Speedup
test_convnd_fwd	99.0	5.4	18.3x
test_convnd_bwd_data	41.0	13.0	3.2x
test_grouped_conv_bwd_data_scale	36.0	29.0	1.2x
test_grouped_convnd_fwd_clamp	352.0	211.0	1.7x
test_grouped_convnd_fwd_scale	108.0	36.0	3.0x
test_grouped_convnd_fwd_bias_clamp	291.0	235.0	1.2x
test_grouped_convnd_fwd_gk_bias_clamp	290.0	228.0	1.3x
test_grouped_convnd_fwd_bilinear	140.0	41.0	3.4x
test_grouped_convnd_fwd_scaleadd_ab	171.0	22.0	7.8x
test_grouped_conv_bwd_data_bilinear	4.9	3.3	1.5x
test_grouped_convnd_bwd_weight_bilinear	6.7	2.5	2.7x

These improvements reduce total execution time for these tests from 1540 seconds to 826 seconds, saving approximately 12 minutes.

Checklist

• I have added tests relevant to the introduced functionality, and the unit tests are passing locally
• I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
• I have added inline documentation which enables the maintainers with understanding the motivation
• I have removed the stale documentation which is no longer relevant after this pull request
• (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
• I have run clang-format on all changed files
• Any dependent changes have been merged

Discussion

The implementation focuses on extending the GPU reference path to match the functionality available in the CPU reference path.

Additional improvement is possible by using GPU for verification and tensor initialization. This PR is already large, so those improvements are deferred.

The batchnorm profiler and tests are not changed since the tests are flaky. In order to keep this PR focused, those changes are also deferred.

[bartekxk]

Can we change these names? out_gn0, wei_g0, wei_gkc0 etc they are not clear for me

[johannes-graner]

Good point, I'll change them.

[bartekxk]

MAybe we can create some common function which will call function with proper number of parameters

[johannes-graner]

That would definitely make it nicer to look at. Thanks for the suggestion!

[bartekxk]

or use some unpack

[bartekxk]

why we need this?

[johannes-graner]

In order to make the naive implementation as simple as possible, the actual conv kernels only operate on packed data. To support all the various layouts, we first have to transform the non-packed tensors into packed tensor, run the kernel, and then transform back to the correct layout.

The loop performs this transformation for all the tensors used in the convolution (for bwd_data, this is the out and weight tensors), which is more than one in the bilinear convolutions.