GitLab

* 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

- doing manual type conversion in custom autograd
- fixing stride issues in backward pass by making some output tensors contiguous

* updating amp to new torch.amp
* using amp autocrats to FP32 for disco convolution kernels
* implemented reduce_scatter routines but disabled those because of memory fluctuations which can cause OOM on big networks

* adding cuda kernels for disco conv

* making psi_idx an attribute

* adding license headers

* adding author files

* reorganizing files

* draft implementation

* added conditional installation to setup.py

* formatting changes

* removing triton kernel in DISCO convolution

* updated github actions

* updated Readme and changelog

* adding another guard for the cuda installation

* renaming the  cuda extension

* simplifying setup.py

* minor bugfix

* Bbonev/cuda disco cleanup (#32)

* cleanup of disco convolutions based on CUDA extension

* fixing unittest

* changing version to experimental 0.7.0a

* initial rewrite of the distributed convolution with CUDA

* fixing streams

* need to fix install options

* fixing streams

* undid setup.py changes

* reset setup.py

* including CUDAStream

* adjusted the precomputation of theta_cutoff. If you rely on this, your models will not be backwards-compatible.

* adjusting theta_cutoff in the unittest

* adding newly refactored kernels for faster compile

* Tkurth/cuda disco distributed fix (#34)

* attempt to make disco distributed

* working distributed convolutions

* fixing distributed conv

* working distributed disco

* removing irrelevant extra argument

* using stream functions from at instead of c10

* using stream functions from at instead of c10, small fix

* Bbonev/disc even filters (#35)

* initial working commit with new convention of counting collocation points across the diameter instead of across the radius

* fixed a bug in the computation of the even kernels

* changing heuristic for computing theta_cutoff

* Fixing unittest

* Readability improvements

* reworked normalization of filter basis functions

* implemented discrete normalization of disco filters

* relaxing tolerances in convolution unit test

* bugfix to correctly support unequal scale factors in latitudes and longitudes

* hotfix to a bug in the imports

* Bbonev/distributed disco refactor (#37)

* cleaned up normalization code in convolution

* formatting changes in distributed convolution

* Fixing default theta_cutoff to be the same in distributed and local case

* fixed distributed convolution to support the same normalization as non-distributed one

* readability improvements

* fixed initial scale of convolution parameter weights and fixed naming of the normalization routine

* Updated Readme.md

* added comment in Dockerfile regarding older architectures

---------
Co-authored-by: Thorsten Kurth <tkurth@nvidia.com>
Co-authored-by: Boris Bonev <bbonev@nvidia.com>

* initial implementation of distributed DISCO layer

* working distributed convolution

* working refactored serial conv transpose with torch kernel

* working distributed conv and transposed conv when using the python kernel

* working distributed convolution with torch kernel

* fixed triton kernel tests

* adding print statement to debug CI

* adjusting tolerances in local convolution unittest

---------
Co-authored-by: Boris Bonev <bbonev@nvidia.com>

* Moved convolutions and exposed them directly

* Added transposition to the unit test

* Minor bugfix in CPU version of DISCO transpose code

* Adding convolution tests to CI

* Added gradient check

* Checking the weight grad as well

* Added test for anisotropic kernels

* Fixing the precomputation of the Psi matrix

* moving tests to a parent folder

* removing tqdm

* deleting deprecated distributed tests

* Moving distributed test

* adapting workfolow

* Added some more comments to make the code more understandable

* more detailed explanation for the derivation of the rotation angles

* added another comment

---------
Co-authored-by: Thorsten Kurth <tkurth@nvidia.com>

* Adding prototype implementation of disco convolutions

* adding s2convolutions.py

* Somewhat functional first implementation in Triton

* Somewhat functional first implementation in Triton

* Made batched execution work on GPU and refactored precomputation of kernel values and supports.

* Adding the Lobatto grid to the grid selection method in quadrature

* fixing bug in triton kernel when iterating over non-zeros

* Updating with working custom triton implementation of DISCO convolution

* Merged both forward and backward DISCO contraction into a single kernel

* Some cleanup and minor bugfix

* bugfixes to the reference implementation

* adding weights to s2conv

* removing unnecessary imports

* suggestion for torch harmonics, math should be checked thoroughly

* Intermediate working reference implementation for the transpose DISCO convolution. Fixed normalization of kernels

* adjusting cutoff frequency in disco convolution

* fixing transpose DISCO contraction

* moving triton to install_requires

---------
Co-authored-by: Thorsten Kurth <tkurth@nvidia.com>