# cunuSHT: GPU Accelerated Spherical Harmonic
Transforms on Arbitrary Pixelizations

Sebastian Belkner, Adriaan Duivenvoorden, Julien Carron, Nathanaël

Schaeffer, Martin Reinecke

► To cite this version:

Sebastian Belkner, Adriaan Duivenvoorden, Julien Carron, Nathanaël Schaeffer, Martin Reinecke. cunuSHT: GPU Accelerated Spherical Harmonic Transforms on Arbitrary Pixelizations. RAS Techniques and Instruments, 2024, 3 (1), pp.711-721. 10.1093/rasti/rzae045. hal-04780694

HAL Id: hal-04780694

https://hal.science/hal-04780694v1

Submitted on 13 Nov 2024

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
cunuSHT: GPU Accelerated Spherical Harmonic Transforms on Arbitrary Pixelizations

Sebastian Belkner<sup>D,1,2</sup> Adriaan J. Duivenvoorden<sup>D,1</sup> Julien Carron<sup>D,2</sup> Nathanael Schaeffer<sup>D,3</sup> and Martin Reinecke<sup>4</sup>

<sup>1</sup>Center for Computational Astrophysics, Flatiron Institute, 162 5th Avenue, New York, NY 10010, USA

<sup>2</sup>Université de Genève, Département de Physique Théorique et CAP, 24 Quai Ansermet, CH-1211 Genève 4, Switzerland

<sup>3</sup>ISTerre, Université de Grenoble 1, CNRS, F-38041 Grenoble, France

<sup>4</sup>Max-Planck Institut für Astrophysik, Karl-Schwarzschild-Str. 1, 85748 Garching, Germany

(Dated: 21st June 2024)

We present `cunuSHT O`, a general-purpose Python package that wraps a highly efficient CUDA implementation of the nonuniform spin-0 spherical harmonic transform. The method is applicable to arbitrary pixelization schemes, including schemes constructed from equally-spaced iso-latitude rings as well as completely nonuniform ones. The algorithm has an asymptotic scaling of $O(\ell_{\max}^3)$ for maximum multipole $\ell_{\max}$ and achieves machine precision accuracy. While `cunuSHT` is developed for applications in cosmology in mind, it is applicable to various other interpolation problems on the sphere. We outperform the fastest available CPU algorithm by a factor of up to 5 for problems with a nonuniform pixelization and $\ell_{\max} > 4 \cdot 10^3$ when comparing a single modern GPU to a modern 32-core CPU. This performance is achieved by utilizing the double Fourier sphere method in combination with the nonuniform fast Fourier transform and by avoiding transfers between the host and device. For scenarios without GPU availability, `cunuSHT` wraps existing CPU libraries. `cunuSHT` is publicly available and includes tests, documentation, and demonstrations.

**keywords:** Nonuniform Spherical Harmonic Transform, Nonuniform Fast Fourier transform, Cosmic Microwave Background Weak Lensing, CUDA

## I. INTRODUCTION

Spherical harmonic transforms (SHTs) are a key ingredient in signal processing for data sets on the 2-sphere. They are extensively used in active research fields such as cosmology (both in studying the cosmic microwave background (CMB) [1–6] and large-scale structure of the Universe [7, 8]), gravitational waves [9], meteorology [10], solar physics [11], or solving partial differential equations on the sphere [12].

Modern data sets routinely require the evaluation of the SHT up to $\ell_{\max} = O(10^4)$ for maps with $N \sim \ell_{\max}^2$ pixels. Here $\ell_{\max}$ denotes the largest multipole $\ell$ considered in the SHT. A direct evaluation of the SHT scales as $O(\ell_{\max}^4)$, which is intractable for large problems. Reducing the computational complexity is thus crucial. One standard optimization is achieved by pixelizing the sphere into rings of constant latitude with equi-angular spaced samples, reducing the problem to $O(\ell_{\max}^3)$. We will refer to this setup as the ring spherical harmonic

transform (rSHT).<sup>1</sup> See [21, 22] for modern implementations of the rSHT. For cases where the transform has to be evaluated on irregularly sampled pixels, ring sampling is not possible. We will refer to this more general setup as the nonuniform spherical harmonic transform (nuSHT). Notable applications of the nuSHT are CMB weak lensing [23], ray tracing [24, 25], or fields where the pixelization changes over time. It should be noted that due to their computational complexity, the rSHT and nuSHT often become the bottleneck in iterative algorithms that repeatedly apply the transforms [26, 27].

In the field of cosmology, nuSHTs are routinely solved via an rSHT and subsequent interpolation to the nonuniform points using bicubic splines [28] or a Taylor series expansion [29]. These methods scale as $O(\ell_{\max}^3)$ but only reach relatively low accuracy. A different $O(\ell_{\max}^3)$ nuSHT method, proposed by [30], makes use of the double Fourier sphere (DFS) method [31] in combination with the nonuniform fast Fourier transform (nuFFT) to achieve an accurate nuSHT algorithm. Several other implementations

<footnote n="1"><sup>1</sup> There exist algorithms for the rSHT that asymptotically scale as $O(\ell_{\max}^2 \log^2(\ell_{\max}))$ [13–17] or $O(\ell_{\max}^2 \log^2(\ell_{\max}) / \log \log(\ell_{\max}))$ [18]. While there have been improvements over the last few years [19, 20] such methods require high memory usage and significant pre-computations. A general purpose implementation that is competitive with the $O(\ell_{\max}^3)$ counterpart has yet to be published.</footnote>
ations of this setup exist [16, 32]. Recently, the lenspyx and DUCC libraries [21, 33] implemented a highly efficient and machine precision accurate implementation of the nuSHT based on the DFS method.

In recent years, there has been a dramatic increase in the availability and usage of graphics processing units (GPUs) dedicated to scientific computing. This development drives the need for GPU based codes and provides an opportunity for increased performance of existing methods. GPUs are optimized for single-instruction multiple-data (SIMD) applications and provide multi-threading well beyond what is achievable with CPUs. Highly parallelizable algorithms can thus greatly benefit from the GPU architecture. In principle, the evaluation of the SHT allows for a large amount of parallel computation, making it a natural target for a GPU implementation. The work presented in [34, 35] was one of the first that explored the use of GPUs for SHTs in the context of cosmology.

Robust and efficient rSHT GPU implementations have been developed in recent years. Notable examples are SHTns [22], supporting the spin-0 and 1 rSHTs, and S2HAT<fnref n="2" /> [35], which supports spin-n transforms. S2FFT [36] is a recent JAX implementation of the spin-n rSHT that provides differentiable transforms. A first implementation of the nuSHT in a cosmological context on GPUs was presented in [37].

We present cunuSHT, a CUDA accelerated nuSHT algorithm on the GPU. This is, to our knowledge, the first publicly available nuSHT GPU algorithm that reaches machine precision accuracy and achieves significant speed-up compared to the fastest CPU algorithms. We achieve this by carefully combining existing robust and efficient GPU implementations of the rSHT and nuFFT algorithm. cunuSHT does not require memory allocation or calculation on the host, which allows it to be incorporated in GPU-based software.

The remainder of the paper is organized as follows. In Section II we introduce notation and definitions, and present in qualitative terms our implementation. Section III discusses the implementation on the GPU. Section IV shows benchmarks and results. We conclude in Section V. A series of appendices collects further details.

## II. NONUNIFORM SPHERICAL HARMONIC TRANSFORM

We introduce our notation and conventions in II A, and define the spherical harmonic transform operations that we implement in this paper. We describe the double Fourier sphere method in II B.

### A. Definition and properties

The (spin-0) spherical harmonic functions $Y_{\ell}^{m}(\theta, \phi)$, with quantum numbers $\ell$ and $m$, with $-\ell \le m \le \ell$, are given by,

$$Y_{\ell}^{m}(\theta, \phi) = P_{\ell}^{m}(\theta)e^{im\phi}, \quad (1)$$

where $P_{\ell}^{m}(\theta)$ are the associated Legendre polynomials. A “general” (or “nonuniform”) spherical harmonic transform (nuSHT) is a linear transformation between a set of spherical harmonic coefficients and field values defined at arbitrary locations on the sphere. We distinguish two types of transforms, with nomenclature inspired by nonuniform Fourier transform literature<fnref n="3" /> [38]:

* Type 1 (also “adjoint nuSHT”, the adjoint operation to type 2 below): given as input a set of $N$ values $f_i$, and $N$ locations $(\theta_i, \phi_i)$, desired are the coefficients $c_{\ell m}$ defined by,

$$c_{\ell m} = \sum_{i=1}^{N} f_i Y_{\ell}^{\dagger m}(\theta_i, \phi_i), \quad (2)$$

for $\ell$ up to some band-limit $\ell_{max}$. We want the result to match a target accuracy $\epsilon$ requested by the user.

* Type 2: given as input a set of harmonic coefficients $c_{\ell m}$ up to some band-limit $\ell_{max}$, and a set of $N$ locations $(\theta_i, \phi_i)$, desired are the field values,

$$f_i = \sum_{\ell=0}^{\ell_{max}} \sum_{m=-\ell}^{\ell} c_{\ell m} Y_{\ell}^{m}(\theta_i, \phi_i), \quad (3)$$

again respecting a target accuracy $\epsilon$ as requested by the user.

<footnote n="2">² https://apc.u-paris.fr/APC_CS/Recherche/Adamis/MIDAS09/software/s2hat/s2hat.html</footnote>

<footnote n="3">³ The DUCC package uses the names `adjoint_synthesis_general` and `synthesis_general` for type 1 and type 2.</footnote>
In matrix notation, type 2 may be written as,

$$f = Yc, \quad (4)$$

where the vector **c** collects the harmonic coefficients, the vector **f** the output field values, and the entries of the matrix **Y** are the spherical harmonics. Type 1 is represented by the adjoint matrix $Y^\dagger = [Y^t]^*$. It is worth noting that type 1 is not the inverse to type 2, except in special cases. We also use the qualifiers type 1 and type 2 for the analogous nonuniform (or uniform) Fourier transforms, where the spherical harmonics and coefficients are replaced by their plane wave counterparts.

In typical applications, the total number of points $N$ is comparable to the squared band-limit, $l_{max}^2$. In this case the naive computational complexity of these operations is $O(l_{max}^4)$.

## B. Double Fourier sphere method

We adopt the approach proposed in [33] and implement type 1 and type 2 transforms using the double Fourier sphere (DFS) method. In this approach, the matrix **Y** of the type 2 nuSHT is decomposed into 4 matrices,

$$Y = NFDS. \quad (5)$$

The final operation **N** is a nonuniform Fourier transform of type 2 to the given locations, and the role of **FDS** is to produce the needed Fourier coefficients.

The matrix **S** is an iso-latitude rSHT, that transforms the input harmonic coefficients onto an equi-spaced grid in both $\phi$ and $\theta$, covering the entire sphere. **D** is a “doubling” operation, that extends the range of $\theta$ from $[0, \pi]$ to $[0, 2\pi)$, see Fig. 1. The doubling is performed by extending the meridians across the south pole back up to the north pole. The essential point is that the resulting map, seen as a map on the doubly-periodic torus, has a standard Fourier series with exactly the same Fourier band-limit$^4$ as the spherical harmonic band-limit of the input array $c_{lm}$. Finally, **F** is simply the standard 2D Fourier transform that produces the Fourier coefficients input to **N** from the doubled map.

<figure><img src="image_2.png" /><figcaption>Figure 1. Illustration of the doubling step in the DFS method. The upper half shows a spherical map calculated on a rectangular grid and is mirrored along the $\theta = \pi$ axis. The mirrored image is split in half across the $\phi$-direction in the center, and swapped. The result is a function on the torus, with a 2D ordinary Fourier series having the exact same band-limit as the spherical harmonic series of the original map defined on the sphere. This allows the use of efficient nonuniform Fast Fourier Transform for accurate interpolation.</figcaption></figure>

The adjoint operator $Y^\dagger$, or the type 1 nuSHT is, by definition,

$$Y^\dagger = S^\dagger D^\dagger F^\dagger N^\dagger. \quad (6)$$

$N^\dagger$ is a nuFFT of type 1 that produces Fourier frequencies from the input locations and field values. $F^\dagger$ produces from these frequencies a 2D map on the torus. $D^\dagger$ (the adjoint doubling matrix) effectively “folds” this doubled map. The resulting map is $2\pi$-periodic in the $\phi$-direction, and the $\theta$-direction goes again from 0 to $\pi$. The map is then transformed to harmonic coefficients with $S^\dagger$, an iso-latitude type 1 rSHT.

## III. IMPLEMENTATION

We discuss the concrete GPU implementation. Readers interested in the CPU equivalent may consult [33].

A GPU is designed to efficiently apply a single instruction on multiple data (SIMD). On the hardware side, it achieves this with Streaming Multiprocessors (SMs) (at the order of 100), that contain a number of simple processors for arithmetic operations (at the order of 100) that execute “warps” of 32 threads in parallel.

On the software side, a GPU accelerated program is executed via a number of threads that are arranged in thread blocks. The GPU is responsible for distributing

<footnote>$^4$ This may be seen for example from the well-known Fourier representation of the Wigner $d$-matrices [39, 40], and using the relation $P_{\ell}^{m}(\theta) = d_{m0}^{\ell}(\theta) \sqrt{(2\ell+1)/4\pi}$.</footnote>
the thread blocks across the SMs. High throughput is achieved by overloading SMs with many threads as to hide data latency and by ensuring that memory is accessed in multiples of the warp size.

We differentiate between the GPU memory that is “close” to the processor units and can be accessed fast by the device, and host memory, that is managed by the host system of the GPU, and which is generally slow to access by the GPU. We show data transfer benchmarks in Appendix A. Our implementation avoids data transfer and usage of host memory altogether; intermediate results are kept in GPU memory. This is realized by copy- arrays in combination with a C++-binding nanobind, [41], handily providing a nanobind-cupy interface.

Our implementation of the individual operators **N**, **D**, **F**, and **S** and their adjoints are realised as follows.

For the (adjoint) synthesis (S†) S, we use the highly efficient software package SHTns [22], and calculate the SHTs onto a Clenshaw-Curtis (CC) grid. The GPU implementation requires the sample size to be divisible by 4.

For iso-latitude rings, in order to achieve best efficiency for the Legendre transform part (the θ part of the transform, which is the critical part), modern top-performing CPU and GPU codes like SHTns use on-the-fly calculation of $P_{\ell}^{m}(\theta)$ using efficient recurrence formulas put forward recently [42]. This allows to keep memory usage low: indeed, only the recurrence coefficients that are independent of θ need to be stored, which requires only $O(\ell_{\max}^2)$ memory, the same order as the data. It leaves two dimensions along which to parallelize: θ and m, and requires a sequential loop over ℓ to compute the $P_{\ell}^{m}(\theta)$ recursively. When $\ell_{\max}$ is larger than 500 to 1000, this leaves enough parallelization opportunities to efficiently use all of the GPU compute units. The computational complexity stays $O(\ell_{\max}^3)$.

The (adjoint) doubling (D†) D is implemented via CUDA, and we write the arrays in a θ-contiguous memory layout, as required by SHTns to keep high efficiency for the Legendre transform. The computational complexity is $O(\ell_{\max}^2)$.

For the type-1 and type-2 nuFFT in 2-dimensions (N†, N), we use cufinuFFT [43] in double precision. The nuFFT method works by utilizing the Fourier transform convolution theorem, and interpolation or convolution onto a slightly larger, up-sampled grid. Highly accurate versions use kernels whose error ϵ decrease exponentially as a function of the up-sampling factor. The

computational complexity (without planning phase) is $O(\ell_{\max}^2 \log(\ell_{\max}) + \ell_{\max}^2 |\log^2(\epsilon)|)$ in 2 dimensions. It is worth mentioning that we use the guru interface to cufinuFFT to initialize the nuFFT plans. The plans allow for repeated and fast transforms, without reinitialization. However, this planning step is the most time consuming operation and is therefore done before calling the functions.

For the FFT operations **F** (type 2) and **F†** (type 1), we use the package cupyx [44] and its cuFFT<fnref n="5" /> integration therein. For the type 1 Fourier synthesis, we use double precision accuracy. For the type 2 Fourier synthesis, we use single precision accuracy for $\epsilon \le 10^{-6}$, which increases speed at the cost of a negligible decrease in effective accuracy on the final result. The computational complexity is $O(\ell_{\max}^2 \log(\ell_{\max}))$.

FFTs become particularly fast if the prime factorization for the sample size gives many small prime numbers, in the following referred to as a good number. If additional constraints are put on the sample size, good numbers may be more difficult to find, see Appendix B for a discussion and concrete definition.

Ignoring the scaling with accuracy, the overall asymptotic computational complexity (for both type 1 and 2) is,

$$O(\ell_{\max}^3) + O(\ell_{\max}^2 \log(\ell_{\max})) + O(\ell_{\max}^2). \quad (7)$$

Here, we assume that the number of uniform and nonuniform points are about the same.

**Y** (and **Y†**) could be further optimized: **S** and **F** both contain a Fourier transform in φ-direction, effectively cancelling out each other. Avoiding this reduces **F** to a 1-dimensional Fourier transform and **S** to a Legendre transformation. This optimization is implemented in the CPU implementation in DUCC. We leave this optimization to a future study for the GPU implementation. We only expect a large speed up for the **F** operator in cunuSHT, which takes about 10 to 20% of the total runtime.

For CMB weak lensing applications, we also provide a pointing routine implemented via CUDA, see Appendix C.

<footnote n="5">⁵ https://developer.nvidia.com/cufft</footnote>
IV. BENCHMARK

We present the scaling and execution time of our imple-
mentation. We take as a use case an application in CMB
weak lensing [23]. CMB weak lensing describes the de-
flection of primordial CMB photons by mass fluctuations
along the line of sight as they travel through the Uni-
verse. These small deflections (their root-mean-square
is of the order of a few arcminutes) are large enough to
be detectable in CMB sky maps [5, 45, 46]. For some
applications, it is necessary to simulate these deflections
accurately and efficiently [27, 47, 48].

Owing to the deflections, the CMB intensity field
$f\tilde{(\theta, \phi)}$ observed at location $(\theta, \phi)$ is the un-deflected field
$f$ at another location,

$$
\tilde{f}(\theta, \phi) = f(\theta', \phi'), \qquad (8)
$$

where $(\theta', \phi')$ depends on $(\theta, \phi)$ in a smooth way. Details
on how these angles relate to each other are given in
Appendix C. For the type-2 nuSHT, the inputs are the set
of angles $(\theta', \phi')$ and the spherical harmonic coefficients
of the un-deflected field $f$. For type-1, the input are the
same set of angles and a set of values of $\tilde{f}$. The set up of
our use case is shown in Fig. 2. The top left panel shows
the input CMB map, the deflection field is shown on the
top right. Both are shown in orthographic projection.
The bottom maps show a detail view of $10 \times 10$ degree of
the difference between the input and deflected map (left
panel) and a detail view of approximately $0.5 \times 0.5$ degree
(right panel) of the nonuniform point (orange squares)
relative to the uniform point locations (blue squares).
We use the same setup for the adjoint operation, in which
case the deflected map becomes the input, and the result
becomes the adjoint SHT coefficients.

Benchmarks are run on an NVIDIA A-100 GPU with
80 GB of memory and an Intel Xeon Gold 8358 Processor
with 32 cores. We set the number of threads to 32 for
the CPU benchmarks. If not stated otherwise, we choose
the following parameters (that mostly affect the nuFFT):
an up-sampling factor of 1.25 to reduce the memory us-
age at a small price of increased computation time<sup>6</sup>, a

gpu_method utilizing the hybrid scheme, called shared
memory, and the default kernel evaluation method. The
resulting map (or input map in the case of the adjoint)
is calculated on a Gauss-Legendre grid.

With this implementation, we can solve problem sizes
of up to $\ell_{\max} \sim 9000$ on an A-100 with 80 GB, by using
the pre-computed nuFFT plans and keeping all necessary
intermediate results in memory.

Fig. 3 shows the speed up of the GPU algorithm com-
pared to the CPU as a function of $\ell_{\max}$ for different ac-
curacies. The left panel shows the evaluation of Eq. (5),
the right panel shows Eq. (6). The $\pm 1\sigma$ variance from 5
runs is indicated by error bars.

For type 2 nuSHT, we reach a speed up between 1 and
5 times for single precision and 1 to 3 times for double
precision, with the speed up increasing with increasing
$\ell_{\max}$. This increase is expected due to better paralleliz-
ability for GPUs for higher $\ell_{\max}$. The speed up is mostly
independent of the target accuracy, with the smallest ac-
curacies tending to perform better. For type 1 nuSHT, we
find lower speed up factors, which may indicate further
improvements. However, type 1 directions are generally
expected to perform worse on the GPU due to the shear
amount of threads that have to be written concurrently
to the same memory location. Nevertheless, the GPU
code either performs almost as good as the CPU code
(small $\ell_{\max}$), or better by a factor of up to 3 for large
problem sizes. For high $\ell_{\max}$, the speed up depends on
the accuracy, with lower accuracies performing better.
For double precision ($\epsilon = 10^{-10}$), the speed up is dimin-
ished due to the double precision penalty that we pay in
our implementation.

We can get a better understanding of the resulting
speed up factors by looking at the time spent on each
of the operators on the GPU. This breakdown is shown
in the top panels of Fig. 4 for the type 2 (left panel) and
type 1 (right panel) nuSHT, as a function of $\ell_{\max}$ for dif-
ferent accuracies. The respective total execution times
are shown in the bottom panels together with the empir-
ically fitted computational complexity model, Eq. (D1).
At the top panels, for each problem size $\ell_{\max}$, each bar
represents a benchmark with an accuracy of (from left to
right) $10^{-10}$, $10^{-6}$, $10^{-2}$. Each bar represent the mean
over 5 runs. For small problem sizes, doubling dominates
and becomes almost negligible for large $\ell_{\max}$. S only
takes about 20% of the execution time for large $\ell_{\max}$,
even though it has the worst asymptotic computational

<footnote n="6">⁶ This somewhat low up-sampling factor reduces the minimal possible accuracy, in double precision, to $10^{-10}$ due to its dependence on the size of the up-sampled grid, but can easily be changed by the user, if needed.</footnote>
<figure><img src="image_3.png" /><figcaption>Figure 2. Setup and simulated data, here for a problem size of $ℓ_{max} = 3095$, $N \approx 2 \cdot 10^7$. The top left (right) panel shows a typical CMB temperature map (a typical deflection field) in orthographic projection. The bottom left plot shows a 10 × 10 degree detail view of the difference between the undeflected and deflected CMB in Cartesian projection. On the bottom right, we show a 0.5 × 0.5 degree detail view comparison between the uniform grid (blue squares), nonuniform grid (orange squares), and the relation between them as indicated by the black arrows.</figcaption></figure>

complexity. This highlights the quality of the rSHT implementation by SHTns.

The choice of accuracy has an impact on the total execution time, as seen in the bottom panels of Fig. 4. The highest accuracy ($\epsilon = 10^{-10}$, red line) takes at most twice as long compared to the low accuracy ($\epsilon = 10^{-2}$, purple line). Our results suggests that improvements in $F$ might be possible: due to FFT being in principle a bandwidth limited routine, we would expect the execution time of $F$ to be comparable to that of $D$, and only a fraction of that of $S$. For type 1 (right column), the total execution time overall takes longer. The breakdown shows that for the low and intermediate accuracy cases, less time is spent in the $F^\dagger$ call. This is a consequence of our use of single precision arithmetic FFTs for the $\epsilon \le 10^{-6}$ case, which,

as mentioned before, is possible for the type 1 nuSHT.

It is interesting to compare this breakdown to the CPU implementation of DUCC. While we cannot expect both breakdowns to be exactly the same due to the different nature of the hardware, large differences may imply possible improvements. Fig. 5 shows the computation time of the CPU implementation of DUCC for the type 2 nuSHT (left column) and type 1 nuSHT (right column) as a function of $ℓ_{max}$ for different accuracies. Looking at the top left panel, we see that $S$ increases with increasing $ℓ_{max}$, as expected from its $O(ℓ_{max}^3)$ scaling. The bottom panel of the left figure shows that the total execution time of the high-accuracy run is at most twice as long as the low accuracy one.

Many optimizations have gone into DUCC's and SHTns's
<figure><img src="image_4.png" /><figcaption>Figure 3. cunuSHT execution time comparison against DUCC for type 2 (type 1) nuSHT shown in the left (right) panel, comparing one A-100 against one Intel Xeon Gold 8358 Processor with 32 cores. A speed up factor > 1 means that the GPU is faster, and we show the result for different target accuracies. The error bars show the ±1σ variance calculated from 5 runs. The GPU takes over at $ℓ_{max} \sim 3 \cdot 10^2$ for type 2, and between $ℓ_{max} \sim 4 \cdot 10^2$ for type 1. The speed up increases up to 5 (3) for large $ℓ_{max}$ left (right) panel.</figcaption></figure>

SHT routines (S, S†) over the years and it is safe to assume that they are close to optimal. Comparing the breakdown of the CPU and GPU implementation for S (green bars) shows that DUCC spends more than twice as long with this operator. This hints to potential sub-optimalities in the implementations of the GPU operators. This also becomes apparent when we look at the operator N (and N†). The scaling with the problem size is much more pronounced for the GPU code.

Fig. 6 shows the effective accuracy as a function of target accuracy for both CPU and GPU. The effective accuracy $\epsilon_{eff}$ is calculated by solving Eq. (3) in a brute-force manner (giving us the values $f^{true}$) and comparing it against Eq. (5) (giving us the values $\hat{f}^{est}$),

$$ \epsilon_{eff} = \frac{1}{f^{true}} \sqrt{\sum_{i=1}^{N_t} (f_i^{true} - \hat{f}_i^{est})^2}, \quad (9) $$

with

$$ \bar{f}^{true} = \sqrt{\sum_{i=1}^{N_t} (f_i^{true})^2}. \quad (10) $$

We choose $N_t = 10^6$ random points on the sphere, giv-

ing us good sampling across the full sphere, and a sufficiently low variance on $\epsilon_{eff}$. Both algorithms achieve good effective accuracies, with the CPU code being more conservative. For the GPU implementation, we note that only low effective accuracies, $\epsilon_{eff} > 10^{-4}$, can be achieved when nuFFT is executed in single precision. When an accuracy of $\epsilon = 10^{-6}$ is desired, a single precision nuFFT evaluation is thus insufficient. We find that executing N and N† in double precision solves this. There is an associated penalty in efficiency due to the double precision calculations being approximately twice as slow. The nuFFT implementation in DUCC circumvents this by allowing for double precision accuracy on the pointing and intermediate results while using single precision accuracy on the Fourier coefficients, hereby reducing the execution time for single precision accuracies.

## V. CONCLUSION

We presented cunuSHT, a GPU accelerated implementation of the spherical harmonic transform on arbitrary pixelization that is, to our knowledge, the first of its kind to achieve faster execution when compared against CPU based algorithms. cunuSHT achieves machine precision
<figure><img src="image_5.png" /><figcaption>Figure 4. Breakdown and total execution time of the GPU implementation cunuSHT. The top panels show the percentage of time spent with the individual operators for type 2 (left column) and type 1 (right column) nuSHT. For each of the problem sizes we show this for different accuracies ($10^{-10}$ on the left, $10^{-6}$ in the center, $10^{-2}$ on the right). The bottom panels show the total execution time, with $\pm 1\sigma$ variance as shaded area, for different accuracies and the empirically fitted computational complexity model, Eq. (D1).</figcaption></figure>

accuracy by transforming the problem of interpolating
on the sphere into a problem of computing a nonuni-
form fast Fourier transform on the torus. Comparing
our implementation executed on an A-100 to the fast-
est available CPU implementations to date running on
a single Intel Xeon Gold 8358 Processor with 32 cores,

we find that our implementation is up to 5 times faster.
We used highly efficient, publicly available packages that
are well tested and robust: SHTns for rSHTs, cufinuFFT
for nuFFTs. We found that, although it has the highest
asymptotic complexity, the high-quality rSHT implemen-
ted in SHTns is not the bottleneck. Our code does not
<figure><img src="image_6.png" /><figcaption>Figure 5. Same as Fig. 4, but for the CPU implementation DUCC. It is important to note that DUCC implements the individual operators more efficiently by avoiding redundant Fourier transforms and by effectively combining the doubling and Fourier transform operations into one. To reflect this, we have grouped the D and F contributions in the top panel.</figcaption></figure>

require intermediate transfers between host and device, allowing it to be incorporated within larger GPU-based algorithms. Many applications in cosmology that we have in mind typically require spin-1 to 3 transforms. We thus plan to extend this package to spin-$n$ transforms in the near future. There are, in principle, no obstacles to the generalization to spin-$n$ by implementing the corresponding Wigner-$d$ transform.

cunuSHT is a general purpose package distributed via pypi, and also works on standard pixelization schemes such as HEALPix, and can also perform rSHTs. Demonstrations of our package on GitHub present how to integrate it into existing pipelines.
ACKNOWLEDGMENT

The authors thank Alex Barnett for helpful discus-
sions about nuFFT, Lehman Garrison for help on the
cupy-nanobind implementation, and Libin Lu for general
discussion and upgrades to cufinuFFT. This work is sup-
ported by the Research Analyst grant from the Simons
Foundation and the computing resources of the Flatiron
Institute. The Flatiron Institute is supported by the Si-
mons Foundation. SB and JC acknowledges support from
a SNSF Eccellenza Professorial Fellowship (No. 186879).

<figure><img src="image_7.png" /><figcaption>Figure 6. Effective accuracy as a function of target accuracy for the CPU (blue), and GPU (orange) implementation.</figcaption></figure>

[1] E. Hivon, K. M. Górski, C. B. Netterfield, B. P. Crill, S. Prunet, and F. Hansen, MASTER of the Cosmic Microwave Background Anisotropy Power Spectrum: A Fast Method for Statistical Analysis of Large and Complex Cosmic Microwave Background Data Sets, Astrophys. J. 567, 2 (2002), arXiv:astro-ph/0105302 [astro-ph].

[2] N. Aghanim et al. (Planck), Planck 2018 results. VIII. Gravitational lensing, Astron. Astrophys. 641, A8 (2020), arXiv:1807.06210 [astro-ph.CO].

[3] P. A. R. Ade et al. (SPTpol, BICEP, Keck), A demonstration of improved constraints on primordial gravitational waves with delensing, Phys. Rev. D 103, 022004 (2021), arXiv:2011.08163 [astro-ph.CO].

[4] P. A. R. Ade et al. (BICEP2, Keck Array), BICEP2 / Keck Array X: Constraints on Primordial Gravitational Waves using Planck, WMAP, and New BICEP2/Keck Observations through the 2015 Season, Phys. Rev. Lett. 121, 221301 (2018), arXiv:1810.05216 [astro-ph.CO].

[5] F. J. Qu et al. (ACT), The Atacama Cosmology Telescope: A Measurement of the DR6 CMB Lensing Power Spectrum and Its Implications for Structure Growth, Astrophys. J. 962, 112 (2024), arXiv:2304.05202 [astro-ph.CO].

[6] S. Hanany et al. (NASA PICO), PICO: Probe of Inflation and Cosmic Origins, (2019), arXiv:1902.10541 [astro-ph.IM].

[7] C. A. Scharf and O. Lahav, Spherical Harmonic Analysis of the 2-JANSKY IRAS Galaxy Redshift Survey, MNRAS 264, 439 (1993).

[8] C. Hikage, M. Takada, T. Hamana, and D. Spergel, Shear power spectrum reconstruction using the pseudo-spectrum method, MNRAS 412, 65 (2011), arXiv:1004.3542 [astro-ph.CO].

[9] N. Deppe, W. Throwe, L. E. Kidder, N. L. Vu, K. C. Nelli, C. Armaza, M. S. Bonilla, F. Hébert, Y. Kim, P. Kumar, G. Lovelace, A. Macedo, J. Moxon, E. O'Shea, H. P. Pfeiffer, M. A. Scheel, S. A. Teukolsky, N. A. Wittek, et al., SpECTRE v2024.06.18, 10.5281/zenodo.12098412 (2024).

[10] N. P. Wedi, M. Hamrud, and G. Mozdzynski, A fast spherical harmonics transform for global nwp and climate models, Monthly Weather Review 141, 3450 (2013).

[11] A. S. Brun and M. Rempel, Large Scale Flows in the Solar Convection Zone, "Space Science Reviews" 144, 151 (2009).

[12] G. L. Browning, J. J. Hack, and P. N. Swarztrauber, A comparison of three numerical methods for solving differential equations on the sphere, Monthly Weather Review 117, 1058 (1989).

[13] J. Driscoll and D. Healy, Computing fourier transforms and convolutions on the 2-sphere, Advances in Applied Mathematics 15, 202 (1994).

[14] D. Potts, G. Steidl, and M. Tasche, Fast and stable algorithms for discrete spherical fourier transforms, Linear Algebra and its Applications 275-276, 433 (1998), proceedings of the Sixth Conference of the International Linear Algebra Society.

[15] D. Potts, G. Steifdl, and M. Tasche, Fast algorithms for discrete polynomial transforms, Mathematics of Computation 67 (1998).
[16] J. Keiner, S. Kunis, and D. Potts, Using NFFT3 - a software library for various nonequispaced fast Fourier transforms, ACM Trans. Math. Software **36**, Article 19, 1 (2009).

[17] D. S. Seljebotn, Wavemoth-fast spherical harmonic transforms by butterfly matrix compression, The Astrophysical Journal Supplement Series **199**, 5 (2012).

[18] N. Hale and A. Townsend, A fast FFT-based discrete Legendre transform, IMA Journal of Numerical Analysis **36**, 1670 (2015), https://academic.oup.com/imajna/article-pdf/36/4/1670/7920439/drv060.pdf.

[19] R. M. Slevinsky, Fast and backward stable transforms between spherical harmonic expansions and bivariate fourier series, Applied and Computational Harmonic Analysis **47**, 585 (2019).

[20] F. Yin, J. Wu, J. Song, and J. Yang, A high accurate and stable legendre transform based on block partitioning and butterfly algorithm for nwp, Mathematics **7**, 10.3390/math7100966 (2019).

[21] M. Reinecke and D. S. Seljebotn, Libsharp - spherical harmonic transforms revisited, Astronomy & Astrophysics **554**, A112 (2013).

[22] N. Schaeffer, Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, Geochemistry **14**, 751 (2012).

[23] A. Lewis and A. Challinor, Weak gravitational lensing of the cmb, Phys. Rept. **429**, 1 (2006), arXiv:astro-ph/0601594 [astro-ph].

[24] G. Fabbian, M. Calabrese, and C. Carbone, Cmb weak-lensing beyond the born approximation: a numerical approach, Journal of Cosmology and Astroparticle Physics **2018** (02), 050–050.

[25] F. Ferlito, C. T. Davies, V. Springel, M. Reinecke, A. Greco, A. M. Delgado, S. D. M. White, C. Hernández-Aguayo, S. Bose, and L. Hernquist, Ray-tracing vs. born approximation in full-sky weak lensing simulations of the millenniumtng project (2024), arXiv:2406.08540.

[26] B. D. Wandelt, D. L. Larson, and A. Lakshminarayanan, Global, exact cosmic microwave background data analysis using Gibbs sampling, Phys. Rev. D **70**, 083511 (2004), arXiv:astro-ph/0310080 [astro-ph].

[27] J. Carron and A. Lewis, Maximum a posteriori CMB lensing reconstruction, Phys. Rev. D **96**, 063510 (2017), arXiv:1704.08230 [astro-ph.CO].

[28] A. Lewis, Lensed CMB simulation and parameter estimation, Phys. Rev. **D71**, 083008 (2005), arXiv:astro-ph/0502469 [astro-ph].

[29] S. K. Næss and T. Louis, Lensing simulations by taylor expansion — not so inefficient after all, Journal of Cosmology and Astroparticle Physics **2013** (09), 001–001.

[30] S. Basak, S. Prunet, and K. Benabed, Simulating weak lensing on cmb maps, (2008), arXiv:0811.1677 [astro-ph].

[31] P. E. Merilees, The pseudospectral approximation applied to the shallow water equations on a sphere, Atmosphere **11**, 13 (1973), https://doi.org/10.1080/00046973.1973.9648342.

[32] A. Townsend, H. Wilber, and G. B. Wright, Computing with Functions in Spherical and Polar Geometries I. The Sphere, SIAM Journal on Scientific Computing **38**, C403 (2016), arXiv:1510.08094 [math.NA].

[33] M. Reinecke, S. Belkner, and J. Carron, Improved cosmic microwave background (de-)lensing using general spherical harmonic transforms, Astron. Astrophys. **678**, A165 (2023), arXiv:2304.10431 [astro-ph.CO].

[34] I. O. Hupca, J. Falcou, L. Grigori, and R. Stompor, Spherical harmonic transform with gpus, in *Lecture Notes in Computer Science* (Springer Berlin Heidelberg, 2012) p. 355–366.

[35] M. Szydlarski, P. Esterie, J. Falcou, L. Grigori, and R. Stompor, Parallel spherical harmonic transforms on heterogeneous architectures (gpus/multi-core cpus), (2013), arXiv:1106.0159 [cs.DC].

[36] M. A. Price and J. D. McEwen, Differentiable and accelerated spherical harmonic and wigner transforms, Journal of Computational Physics **510**, 113109 (2024).

[37] A. Baleato Lizancos and M. White, Harmonic analysis of discrete tracers of large-scale structure, JCAP **05**, 010, arXiv:2312.12285 [astro-ph.CO].

[38] A. H. Barnett, J. Magland, and L. af Klinteberg, A Parallel Nonuniform Fast Fourier Transform Library Based on an “Exponential of Semicircle” Kernel, SIAM Journal on Scientific Computing **41**, C479 (2019), arXiv:1808.06736 [math.NA].

[39] T. Risbo, Fourier transform summation of legendre series and d-functions, Journal of Geodesy **70**, 383 (1996).

[40] K. M. Huffenberger and B. D. Wandelt, Fast and Exact Spin-s Spherical Harmonic Transforms, Astrophys. J. Suppl. **189**, 255 (2010), arXiv:1007.3514 [astro-ph.IM].

[41] W. Jakob, nanobind: tiny and efficient c++/python bindings (2022), https://github.com/wjakob/nanobind.

[42] K. ISHIOKA, A new recurrence formula for efficient computation of spherical harmonic transform, Journal of the Meteorological Society of Japan. Ser. II **96**, 241 (2018).

[43] Y. hsuan Shih, G. Wright, J. Andén, J. Blaschke, and A. H. Barnett, cufinufft: a load-balanced gpu library for general-purpose nonuniform ffts, 2021 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW), 688 (2021).

[44] R. Okuta, Y. Unno, D. Nishino, S. Hido, and C. Loomis, Cupy: A numpy-compatible library for nvidia gpu calculations, in *Proceedings of Workshop on Machine Learn-
ing Systems (LearningSys) in The Thirty-first Annual
Conference on Neural Information Processing Systems
(NIPS) (2017).

[45] N. Aghanim et al. (Planck), Planck 2018 results. VIII. Gravitational lensing, Astron. Astrophys. **641**, A8 (2020), arXiv:1807.06210 [astro-ph.CO].

[46] Z. Pan et al. (SPT), Measurement of gravitational lensing of the cosmic microwave background using SPT-3G 2018 data, Phys. Rev. D **108**, 122005 (2023), arXiv:2308.11608 [astro-ph.CO].

[47] C. M. Hirata and U. Seljak, Reconstruction of lensing from the cosmic microwave background polarization, Phys. Rev. **D68**, 083002 (2003), arXiv:astro-ph/0306354 [astro-ph].

[48] S. Belkner, J. Carron, L. Legrand, C. Umiltà, C. Pryke, and C. Bischoff (CMB-S4), CMB-S4: Iterative Internal Delensing and r Constraints, Astrophys. J. **964**, 148 (2024), arXiv:2310.06729 [astro-ph.CO].

## Appendix A: Data Transfer

All benchmarks are done without measuring the time to transfer the data to device and back. This can be a large part of the overall computation and should be avoided, as is shown in Fig. A1 for different problem sizes. We note here that cunuSHT provides means to keep everything on the GPU without having to transfer the data. Thus, these transfer times can in principle be avoided.

## Appendix B: Good numbers

For fast Fourier transforms, execution time depends strongly on the largest prime factor of the transform length; the smaller it is, the better. Luckily enough, many good numbers for CPU algorithms exist. Assuming Clenshaw-Curtis quadrature, the constraint on the number of rings $N_r$ to sufficiently sample a map with band limit $\ell_{max}$ is $N_r \ge \ell_{max} + 1$. Therefore, after applying the double Fourier sphere method, we must have at least $2N_r + 2$ samples, and we are free to choose an arbitrarily larger number which happens to be a good FFT size. For the GPU algorithm, however, there is a caveat: due to the additional constraint on $N_r$ to be divisible by 4 for the rSHT operation, the selection of good numbers is reduced. While there are still good numbers for FFT satisfying both constraints, they are sparser. The full list

<figure><img src="image_8.png" /><figcaption>Figure A1. type 2 nuSHT transfer times “Host-to-Device” (H2D) (green) and “Device-to-Host” (D2H) (yellow) as a function of $\ell_{max}$. The shaded areas show the $\pm 1\sigma$ variance in transfer time, calculated from ten runs. Note that these transfer times are not generally part of the routine and can completely be avoided. H2D transfer contains the SHT coefficients, while D2H contains the much larger map. For type 1 nuSHT, the transfer times are equally long, but the H2D and D2H are swapped due to the input and output of the function.</figcaption></figure>

of $N_r$ that are multiples of 4 up to $N_r \le 10000$ for which $\ell_{max} = N_r - 1$ can be factored into primes up to 11 is 4, 8, **12**, 16, 28, 36, **56**, 64, 76, **100**, 136, 148, 176, **232**, 244, **276**, 316, 344, **364**, 376, **496**, **540**, 568, 676, 736, **848**, 876, **892**, **1156**, 1216, 1324, **1332**, **1376**, 1576, 1716, **1816**, 1876, **2080**, 2188, **2476**, **2696**, 2836, 3088, **3268**, 3376, 3676, **4236**, 4376, **4456**, **4852**, 5104, **5776**, 6076, 6616, **6656**, **6876**, 7204, **7624**, 7876, **8020**, 8576, **9076**, 9376. The numbers with the highest prime factor being 11 are highlighted with an underline as efficient calculation of cunuSHT only applies to the CPU code here.

One way to enlarge this list could be for example by using a Fejér grid as intermediate grid, as opposed to the currently used Clenshaw-Curtis grid. This would effectively reduce the number of required samples to $2N_r$ after doubling, resulting in a much wider choice of sample sizes that are also good FFT lengths. Hence restrictions on the number of good numbers are expected to go away in the future.
<figure><img src="image_9.png" /><figcaption>Figure C1. Execution time of the pointing routine on the GPU as a function of the problem size. The shaded area shows the ±1σ uncertainty calculated from 5 runs.</figcaption></figure>

<table><caption>Table D1. Empirical fits of Eq. (7) for type 2 and type 1 nuSHT, for different accuracies for both CPU and GPU.</caption><thead><tr><th></th><th>type 2</th><th>Target acc.</th><th>α</th><th>β</th></tr></thead><tbody><tr><td rowspan="3">CPU</td><td>10<sup>-10</sup></td><td>1.52 · 10<sup>-1</sup></td><td>6.96 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-6</sup></td><td>1.41 · 10<sup>-1</sup></td><td>6.16 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-2</sup></td><td>1.67 · 10<sup>-1</sup></td><td>4.11 · 10<sup>-1</sup></td></tr><tr><td rowspan="3">GPU</td><td>10<sup>-10</sup></td><td>1.04 · 10<sup>-6</sup></td><td>2.72 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-6</sup></td><td>1.06 · 10<sup>-2</sup></td><td>1.60 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-2</sup></td><td>1.39 · 10<sup>-4</sup></td><td>1.19 · 10<sup>-1</sup></td></tr></tbody></table><table><thead><tr><th></th><th>type 1</th><th>Target acc.</th><th>α</th><th>β</th></tr></thead><tbody><tr><td rowspan="3">CPU</td><td>10<sup>-10</sup></td><td>1.24 · 10<sup>-1</sup></td><td>7.34 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-6</sup></td><td>1.06 · 10<sup>-1</sup></td><td>6.67 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-2</sup></td><td>1.68 · 10<sup>-1</sup></td><td>3.72 · 10<sup>-1</sup></td></tr><tr><td rowspan="3">GPU</td><td>10<sup>-10</sup></td><td>1.31 · 10<sup>-8</sup></td><td>4.33 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-6</sup></td><td>1.07 · 10<sup>-6</sup></td><td>2.84 · 10<sup>-1</sup></td></tr><tr><td>10<sup>-2</sup></td><td>2.16 · 10<sup>-5</sup></td><td>1.64 · 10<sup>-1</sup></td></tr></tbody></table>

## Appendix C: CMB weak lensing pointing

Let $\hat{e}_{\theta}, \hat{e}_{\phi}$ and $\hat{n}$ form the right-handed unit basis vectors at the point on the sphere parametrized by $\theta, \phi$ (so, the components of $\hat{n}$ are $\sin\theta \cos\phi, \sin\theta \sin\phi, \cos\theta$). In our benchmark CMB lensing application, the deflected positions $\hat{n}'$ that define the angles $\theta'$ and $\phi'$ at which the CMB field must be evaluated are given by,

$$ \hat{n}' = \cos(\alpha)\hat{n} + \frac{\sin(\alpha)}{\alpha} (\alpha_{\theta}\hat{e}_{\theta} + \alpha_{\phi}\hat{e}_{\phi}), \quad (C1) $$

where $\alpha = \sqrt{\alpha_{\theta}^2 + \alpha_{\phi}^2}$, and $\alpha_{\theta}$ and $\alpha_{\phi}$ form the gradient of the lensing potential $\Phi$,

$$ \alpha_{\theta}(\hat{n}) + i\alpha_{\phi}(\hat{n}) = \left( \frac{\partial}{\partial\theta} + \frac{i}{\sin\theta} \frac{\partial}{\partial\phi} \right) \Phi(\hat{n}). \quad (C2) $$

cunuSHT first obtains $\alpha_{\theta}$ and $\alpha_{\phi}$ from this equation using SHTns. Then, threading across each ring, we solve for $\theta', \phi'$ in Eq. (C1) on the fly. A benchmark as a function of $\ell_{max}$ is shown in Fig. C1.

## Appendix D: Computational complexity

Table D1 shows the results for the empirical fit of the type 2 and type 1 nuSHT computational complexity models,

$$ C(l) = \alpha \frac{l^3}{l_{\text{norm}}^3} + \beta \frac{l^2 \log(l)}{l_{\text{norm}}^2 \log(l_{\text{norm}})}, \quad (\text{D1}) $$

with $l_{norm} = 6067$ the normalization for the unknown prefactors. Compared to Eq. (7) we only account for the terms with the largest complexities. The fits are shown in Fig. 3, Fig. 4, and Fig. 5.

## Appendix E: Code examples

The following code block shows a minimum working example for calculating a type 2 and type 1 nuSHT on the GPU for SHT coefficients alm that are deflected by a deflection field dlm_dlm_scaled, for an accuracy of epsilon and band limit lmax. Both routines will set up the plans, so that the actual nusht2dX() call may be called repeatedly.
1 import cunusht as cu

2

3 lenjob_geominfo = ('gl', {'lmax': lmax})

4 kwargs = {

5 'geominfo_deflection': lenjob_geominfo,

6 'epsilon': epsilon,

7 'nuFFTtype': 2,

8 }

9 t = cu.getTransformer( backend='GPU')(**kwargs)

10 ptg = t.dlm2pointing(dlm

11 lenmap = t.nusht2d2(alm, ptg, lmax, lenmap)

1 import cunusht as cu

2

3 lenjob_geominfo = ('gl', {'lmax': lmax})

4 kwargs = {

5 'geominfo_deflection': lenjob_geominfo,

6 'epsilon': epsilon,

7 'nuFFTtype': 1,

8 }

9 t = cu.getTransformer( backend='GPU')(**kwargs)

10 ptg = t.dlm2pointing(dlm_sca

11 alm = t.nusht2d1(alm, ptg, lmax, lenmap)