GitLab

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

* reworked the setup system to use pyproject.toml

* changed the way custom CUDA extension build flag is passed

* reworked the CI

* changed to source distribution only in the CI

* 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>

* 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>