GitLab

* Replaced PyTorch's slow permutation.
* Split kernel into general and specialized versions (for num_channel <= 8192)
* Enabled float4-based vectorized memory access, when possible.
* Added runtime dispatch logic for kernel specialization.

Aligned attention_fwd_cuda.cu with attention_bwd_cuda.cu in terms of naming conventions and kernel parameters.

Extracted shared host/device functions and declarations into a separate module:
* attention_utils.cuh
* attention_utils.cu

* removing duplicate code from distributed convoloution

* replacing from_numpy with as_tensor

* make preprocess_psi_tensor GPU ready.

* refactoring disco backend code

* removed get_psi as member function and instead put it in _disco_convolution

* setting seeds in tests more consistently

* parametrized test classes to ensure that tests are always run on both CPU and GPU (if available)

* cleaning up

Updated pointer offset calculations to use 64-bit integers to prevent overflow with large batch or image sizes.

https://github.com/NVIDIA/torch-harmonics/compare/main...azrael417:torch-harmonics:tkurth/mauro-rebase

* Added a new CSR array, psi_row_index, containing "ho" values sorted in descending order of CSR row length; this is used to process (ho, wo) points corresponding to longer rows before shorter ones, improving overlap and reducing the tail effect.

* Replaced PyTorch's slow permutation ops with custom kernels, significantly improving performance (especially on GB200).
* Split kernel into general and specialized versions for num_channel <= 16384, significantly reducing memory accesses.
* Enabled float4-based vectorized memory access when pointer alignment and channel size allow, improving throughput.
* Added runtime dispatch logic for kernel specialization.

Leverage the same qdotk_max "trick" for the backward kernel. This avoids 1 loop
and saves about 20% of performance.

Also: Made fwd kernel use modified memory layout with standard shape

Also match the gradient output to the input, in terms of memory layout

Detect memory layout (B,C,H,W) (stride for C should be 1, if not, fix it)

This ensures that the backwards kernel is fast

putting qy in shared is a little faster

Changing internal memory layout means we can leave code in standard shape and
only change layout external to kernel

Introduce new CUDA kernels, `s2_attention_bwd_dkvq_kernel_mbT` and
`s2_attention_kernel_mbT`, for more efficient computation of backward gradients
and forward attention respectively. These changes optimize memory access
patterns and employ coalesced operations by leveraging tensor transpositions.

Forward kernel written by Mauro Bisson
Backwards kernel written by Andrea Paris (aparis@ethz.ch) and Max Rietmann

Parallelization strategy computes 1 output per Warp, with threads computing the
dot-product in parallel. Because inputs are transposed to have channel dimension
last, the dot-product memory access pattern is perfectly coalesced, leading to
excellent performance. This is true across both forward and backward kernels.
Co-authored-by: Mauro Bisson <maurob@nvidia.com>
Co-authored-by: Max Rietmann <mrietmann@nvidia.com>
Co-authored-by: Andrea Paris <aparis@ethz.ch>

fixing bug in quadrature weights for full attention. Adding better unit tests for attention. Cleanup in the cuda code.

* setting imaginary parts of DCT and nyquist frequency to zero in IRSHT variants

This reverts commit 82881276.

* setting imaginary parts of DCT and nyquist frequency to zero in IRSHT variants

* small fix

* making einsum result contiguous

* adding zero frequency to distributed sht

* small hotfix for lobatto grid precomputation routine

* adding lobatto grid to tests

* adding caching

* replacing many numpy calls with torch calls

* bumping up version number to 0.7.6

* Improved computation of Morlet filter basis and switched to a Hann window.

* addresses #064 and some cleanup