# coding=utf-8
# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
from itertools import accumulate
from typing import Optional, Tuple, Union
import torch
from disco_helpers import optimized_kernels_is_available, pack_psi_dense, preprocess_psi
from torch_harmonics.disco._disco_utils import _get_psi
from torch_harmonics.disco.convolution import (
DiscreteContinuousConv,
_kpacked_device_supported_for_tensor,
_precompute_convolution_tensor_s2,
)
from torch_harmonics.disco.kernels_torch.disco_torch import _disco_s2_transpose_contraction_torch
from torch_harmonics.disco.optimized.disco_optimized import (
_build_kernel_split_csr,
_disco_s2_transpose_contraction_optimized,
_kpacked_build_available,
_maybe_kpack_psi,
_split_csr_python_offsets,
)
from torch_harmonics.quadrature import compute_theta_cutoff, effective_theta_cutoff, precompute_latitudes
# a2a forward orchestration: standard (fused=False) and reordered (fused=True).
from .kernels import (
_distributed_disco_fwd_a2a,
_distributed_disco_fwd_a2a_reordered,
)
# distributed stuff — relative imports to avoid the circular dependency
# through torch_harmonics.distributed.__init__.
from .primitives import (
compute_polar_halo_radius,
compute_split_shapes,
distributed_transpose_azimuth,
gather_from_copy_to_polar_region,
)
from .utils import (
azimuth_group_rank,
azimuth_group_size,
polar_group_rank,
polar_group_size,
)
#: Strategies for combining the polar ranks' contributions, see DistributedDiscreteContinuousConvS2.
_POLAR_MODES = frozenset({"halo-exchange", "reduce-scatter"})
def _split_distributed_convolution_tensor_s2(
idx: torch.Tensor,
vals: torch.Tensor,
in_shape: Tuple[int],
out_shape: Tuple[int],
):
"""
Splits a pre-computed convolution tensor along the latitude dimension for distributed processing.
This function takes a convolution tensor that was generated by the serial routine and filters
it to only include entries corresponding to the local latitude slice assigned to this process.
The filtering is done based on the polar group rank and the computed split shapes.
Parameters
----------
idx : torch.Tensor
Indices of the pre-computed convolution tensor
vals : torch.Tensor
Values of the pre-computed convolution tensor
in_shape : Tuple[int]
Shape of the input tensor (nlat_in, nlon_in)
out_shape : Tuple[int]
Shape of the output tensor (nlat_out, nlon_out)
Returns
-------
idx : torch.Tensor
Filtered indices corresponding to the local latitude slice
vals : torch.Tensor
Filtered values corresponding to the local latitude slice
"""
nlat_in, nlon_in = in_shape
nlat_out, nlon_out = out_shape
comm_size_polar = polar_group_size()
comm_rank_polar = polar_group_rank()
split_shapes = compute_split_shapes(nlat_in, num_chunks=comm_size_polar)
offsets = [0] + list(accumulate(split_shapes))
start_idx = offsets[comm_rank_polar]
end_idx = offsets[comm_rank_polar + 1]
# once normalization is done we can throw away the entries which correspond to input latitudes we do not care about
lats = idx[2] // nlon_in
lons = idx[2] % nlon_in
ilats = torch.argwhere((lats < end_idx) & (lats >= start_idx)).squeeze()
vals = vals[ilats]
# for the indices we need to recompute them to refer to local indices of the input tenor
idx = torch.stack([idx[0, ilats], idx[1, ilats], (lats[ilats] - start_idx) * nlon_in + lons[ilats]], dim=0)
# make results contiguous
idx = idx.contiguous()
vals = vals.contiguous()
return idx, vals
def _split_halo_convolution_tensor_s2(
idx: torch.Tensor,
vals: torch.Tensor,
in_shape: Tuple[int],
out_shape: Tuple[int],
r_lat: int,
):
r"""
Restrict a convolution tensor to this rank's **output** latitudes over a halo-padded input.
The counterpart of :func:`_split_distributed_convolution_tensor_s2` for the halo path, and
the mirror of what it does: that one keeps the local *input* latitudes and leaves the output
axis global, so the rank computes a partial sum over every output latitude which a
reduce-scatter then completes. Here the rank keeps the output latitudes it *owns* and reads
from a halo-padded input band, so its result is already complete and no reduction follows.
Rows are re-keyed to ``[0, nlat_out_local)`` and columns to the halo-padded local input
range ``[lat_in_start - r_lat, lat_in_start + nlat_in_local + r_lat)``. Rows at the poles
are padded on the missing side by :func:`~torch_harmonics.distributed.polar_halo_exchange`,
so their column indices stay in range; entries pointing outside the padded band cannot occur
when ``r_lat`` comes from :func:`~torch_harmonics.distributed.compute_polar_halo_radius`,
and are dropped defensively if they somehow do.
Normalization is unaffected: it is applied to the global tensor before this filter runs, and
it sums over the *input* points of each ``(kernel, output latitude)`` group -- all of which
lie inside the halo for the rows this rank keeps.
Parameters
----------
idx : torch.Tensor
Indices of the pre-computed (global, normalized) convolution tensor.
vals : torch.Tensor
Values of the pre-computed convolution tensor.
in_shape : Tuple[int]
Global input shape ``(nlat_in, nlon_in)``.
out_shape : Tuple[int]
Global output shape ``(nlat_out, nlon_out)``.
r_lat : int
Halo radius in latitude rows.
Returns
-------
idx : torch.Tensor
Indices re-keyed to local output rows and halo-padded input columns.
vals : torch.Tensor
The corresponding values.
"""
nlat_in, nlon_in = in_shape
nlat_out, nlon_out = out_shape
comm_size_polar = polar_group_size()
comm_rank_polar = polar_group_rank()
in_shapes = compute_split_shapes(nlat_in, num_chunks=comm_size_polar)
out_shapes = compute_split_shapes(nlat_out, num_chunks=comm_size_polar)
in_start = sum(in_shapes[:comm_rank_polar])
out_start = sum(out_shapes[:comm_rank_polar])
nlat_out_local = out_shapes[comm_rank_polar]
# the halo-padded input band, clipped at the poles where there is no neighbour to borrow
# from -- polar_halo_exchange zero-pads those rows, so the padded tensor still has
# 2 * r_lat extra rows and the offset below is uniform across ranks
halo_start = in_start - r_lat
lats = idx[2] // nlon_in
lons = idx[2] % nlon_in
keep = (idx[1] >= out_start) & (idx[1] < out_start + nlat_out_local)
# defensive: with a correctly derived r_lat every kept entry already lies in the band
keep = keep & (lats >= halo_start) & (lats < halo_start + nlat_in_local_padded(in_shapes, comm_rank_polar, r_lat))
sel = torch.argwhere(keep).squeeze(-1)
vals = vals[sel]
idx = torch.stack([idx[0, sel], idx[1, sel] - out_start, (lats[sel] - halo_start) * nlon_in + lons[sel]], dim=0)
return idx.contiguous(), vals.contiguous()
def nlat_in_local_padded(in_shapes, rank: int, r_lat: int) -> int:
"""Rows of the halo-padded local input band: the local share plus a halo on each side.
Uniform across ranks, including at the poles: ``polar_halo_exchange`` zero-pads the side
with no neighbour rather than shortening the tensor, so every rank sees the same layout and
the psi column offset does not need a per-rank special case.
"""
return in_shapes[rank] + 2 * r_lat
[docs]
class DistributedDiscreteContinuousConvS2(DiscreteContinuousConv):
"""
Distributed version of Discrete-continuous convolutions (DISCO) on the 2-Sphere as described in :cite:`Ocampo2023`.
We assume the data can be split in polar and azimuthal directions.
.. seealso::
:class:`torch_harmonics.DiscreteContinuousConvS2`
Serial counterpart with full mathematical description and parameter
documentation.
The azimuth direction is all-to-all: a channel swap so the sparse psi contraction
runs against the full ``nlon_in`` row, then back to channel-distributed. The polar
direction is chosen by ``polar_mode=``, which decides whether a rank computes only
the output latitudes it owns (borrowing a halo) or all of them (and reduce-scatters
the partial sums). That choice governs how the K-expanded intermediate scales, so it
is usually the one that decides whether a large model fits. The ``fused=`` flag is
orthogonal to it and mirrors the serial conv:
``fused=False`` (default) — standard a2a: einsum after the
transpose-back; the K-expanded intermediate is saved for backward.
``fused=True`` — reordered a2a: the weight einsum runs before the
collectives on the local azimuth channel shard, via the fused
contraction+einsum op that recomputes the K-expanded in backward
instead of saving it. K× lower activation memory and K× less
collective volume, at the cost of one extra contraction in
backward. CUDA + optimized kernels only.
Parameters
----------
in_channels : int
Number of input channels
out_channels : int
Number of output channels
in_shape : Tuple[int]
Shape of the input tensor
out_shape : Tuple[int]
Shape of the output tensor
kernel_shape : Union[int, Tuple[int], Tuple[int, int]]
Shape of the kernel
basis_type : Optional[str]
Type of basis to use
basis_norm_mode : Optional[str]
Normalization mode for the filter basis
groups : Optional[int]
Number of groups
grid_in : Optional[str]
Grid type for the input tensor
grid_out : Optional[str]
Grid type for the output tensor
bias : Optional[bool]
Whether to use bias
theta_cutoff : Optional[float]
Theta cutoff for the filter basis
optimized_kernel : Optional[bool]
Use the optimized CUDA contraction kernel. Required when ``fused=True``.
fused : bool
Mirrors the serial conv. ``False`` (default): standard all-to-all
(the K-expanded intermediate is saved for backward). ``True``:
reordered all-to-all — the weight einsum runs before the collectives
on the local azimuth channel shard and the K-expanded is recomputed
in backward instead of saved, for K× lower activation memory and K×
less collective volume (CUDA + optimized kernels only).
polar_mode : Optional[str]
How the polar ranks obtain their output latitudes.
``"halo-exchange"`` (default): each rank borrows the input rows its own output rows
reach into and computes them outright. The K-expanded intermediate is then
``(B, C, K, nlat_out / polar_group_size, nlon_out)``, i.e. it shrinks as polar ranks
are added. Requires the filter support to reach no further than the neighbouring
rank; when it does not, construction raises and says so, rather than quietly
selecting the other mode and leaving the memory problem in place.
``"reduce-scatter"``: every rank computes a partial sum over *all* output latitudes
and a reduce-scatter completes and splits it. No bound on the angular reach, so this
is the mode for cutoffs wide enough that a halo cannot serve them -- at the cost of an
intermediate that stays at the full ``nlat_out`` however many polar ranks are used.
Returns
-------
torch.Tensor
Output tensor
References
----------
:cite:`Ocampo2023`
"""
def __init__(
self,
in_channels: int,
out_channels: int,
in_shape: Tuple[int],
out_shape: Tuple[int],
kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
basis_type: Optional[str] = "piecewise linear",
basis_norm_mode: Optional[str] = "nodal",
groups: Optional[int] = 1,
grid_in: Optional[str] = "equiangular",
grid_out: Optional[str] = "equiangular",
bias: Optional[bool] = True,
theta_cutoff: Optional[float] = None,
optimized_kernel: Optional[bool] = True,
fused: bool = False,
polar_mode: Optional[str] = "halo-exchange",
):
super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias, optimized_kernel)
if polar_mode not in _POLAR_MODES:
raise ValueError(f"Unknown polar_mode '{polar_mode}', expected one of {sorted(_POLAR_MODES)}")
self.polar_mode = polar_mode
# fused=True uses the reordered a2a (fused contraction+einsum op with
# K-expanded recompute in backward); it is CUDA + optimized-kernel only.
self.fused = bool(fused)
if self.fused and not (torch.cuda.is_available() and optimized_kernels_is_available() and optimized_kernel):
raise NotImplementedError(
"DistributedDiscreteContinuousConvS2(fused=True) requires CUDA and the "
"optimized DISCO kernels (it uses the fused contraction+einsum op with "
"K-expanded recompute in backward). Use fused=False otherwise."
)
self.nlat_in, self.nlon_in = in_shape
self.nlat_out, self.nlon_out = out_shape
# get the comms grid:
self.comm_size_polar = polar_group_size()
self.comm_rank_polar = polar_group_rank()
self.comm_size_azimuth = azimuth_group_size()
self.comm_rank_azimuth = azimuth_group_rank()
# we need those shapes:
self.lat_in_shapes = compute_split_shapes(self.nlat_in, self.comm_size_polar)
self.lon_in_shapes = compute_split_shapes(self.nlon_in, self.comm_size_azimuth)
self.lat_out_shapes = compute_split_shapes(self.nlat_out, self.comm_size_polar)
self.lon_out_shapes = compute_split_shapes(self.nlon_out, self.comm_size_azimuth)
# compute theta cutoff based on the bandlimit of the input field
if theta_cutoff is None:
self.theta_cutoff = compute_theta_cutoff(self.nlat_out, grid=grid_out)
else:
self.theta_cutoff = theta_cutoff
if self.theta_cutoff <= 0.0:
raise ValueError("Error, theta_cutoff has to be positive.")
# Note that the psi matrix is of shape nlat_out x nlat_in * nlon_in.
# Since the contraction in nlon direction is a convolution, we keep
# it local to all nodes and split along nlat. We further split the
# input dim because this reduces the number of atomic reduction
# calls inside the actual kernel.
# The operator is local: an output latitude only reads input latitudes within
# theta_cutoff of it. When that reach fits inside a neighbouring rank's share, each rank
# can borrow a halo and compute its own output rows outright -- no partial sums, so no
# polar reduction, and the K-expanded intermediate shrinks with the polar group instead
# of staying pinned at the full nlat_out. When the cutoff is wide enough that the reach
# spans more than a neighbour, the immediate-neighbour exchange cannot serve it and we
# fall back to computing every output latitude and reduce-scattering, which has no such
# restriction.
self.r_lat = 0
self.use_halo = self.polar_mode == "halo-exchange"
if self.use_halo:
lats_in, _ = precompute_latitudes(self.nlat_in, grid=grid_in)
lats_out, _ = precompute_latitudes(self.nlat_out, grid=grid_out)
try:
self.r_lat = compute_polar_halo_radius(lats_in, lats_out, effective_theta_cutoff(self.theta_cutoff), self.lat_in_shapes, self.lat_out_shapes)
except ValueError as err:
raise ValueError(
f"{err}\n\n"
"This convolution therefore cannot use polar_mode='halo-exchange'. Either reduce the "
"polar group size or theta_cutoff so the support fits within one neighbour, or pass "
"polar_mode='reduce-scatter', which places no bound on the angular reach but computes "
"every output latitude on every polar rank -- so its intermediate does not shrink as "
"polar ranks are added."
) from err
# set local shapes according to distributed mode
self.nlat_in_local = self.lat_in_shapes[self.comm_rank_polar]
self.nlat_out_local = self.lat_out_shapes[self.comm_rank_polar] if self.use_halo else self.nlat_out
self.nlon_in_local = self.lon_in_shapes[self.comm_rank_azimuth]
self.nlon_out_local = self.lon_out_shapes[self.comm_rank_azimuth]
self.kpacked_device_supported = False
# compute global convolution tensor
idx, vals, _ = _precompute_convolution_tensor_s2(
in_shape,
out_shape,
self.filter_basis,
grid_in=grid_in,
grid_out=grid_out,
theta_cutoff=self.theta_cutoff,
transpose_normalization=False,
basis_norm_mode=basis_norm_mode,
merge_quadrature=True,
)
if self.use_halo:
idx, vals = _split_halo_convolution_tensor_s2(idx, vals, in_shape, out_shape, self.r_lat)
else:
idx, vals = _split_distributed_convolution_tensor_s2(idx, vals, in_shape, out_shape)
self._build_local_psi(idx, vals)
def _build_local_psi(self, idx: torch.Tensor, vals: torch.Tensor):
"""Register psi buffers for the a2a path. The a2a swap makes W
local before the kernel reads col_idx, so no wi pre-shift is
applied here."""
ker_idx = idx[0, ...].contiguous()
row_idx = idx[1, ...].contiguous()
col_idx = idx[2, ...].contiguous()
vals = vals.contiguous()
self.psi_kpacked_K_pad = None
if self.optimized_kernel:
roff_idx = preprocess_psi(
self.kernel_size,
self.nlat_out_local,
ker_idx,
row_idx,
col_idx,
vals,
).contiguous()
self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
split_roff_idx, split_nnz_off, split_ker_idx, split_row_idx, split_col_idx, split_vals = _build_kernel_split_csr(
roff_idx, ker_idx, row_idx, col_idx, vals, self.kernel_size, self.nlat_out_local
)
self.psi_split_row_offsets, self.psi_split_nnz_offsets = _split_csr_python_offsets(split_nnz_off)
self.register_buffer("psi_split_roff_idx", split_roff_idx, persistent=False)
self.register_buffer("psi_split_nnz_off", split_nnz_off, persistent=False)
self.register_buffer("psi_split_ker_idx", split_ker_idx, persistent=False)
self.register_buffer("psi_split_row_idx", split_row_idx, persistent=False)
self.register_buffer("psi_split_col_idx", split_col_idx, persistent=False)
self.register_buffer("psi_split_vals", split_vals, persistent=False)
# Optional K-packed dense layout for the WGMMA / tcgen05 path.
# A2A makes W local before the kernel, so wi_shift=0 like the serial path.
#
# Skipped when the build contains no kpacked kernel; see
# _kpacked_build_available for why the check is build-time, not
# device-time.
if _kpacked_build_available():
psi_packed_idx, psi_packed_vals, psi_packed_count = pack_psi_dense(
self.kernel_size, self.nlat_out_local, self.nlon_in, 0, ker_idx, row_idx, col_idx, vals, roff_idx
)
kpack = _maybe_kpack_psi(psi_packed_idx.contiguous(), psi_packed_vals.contiguous(), psi_packed_count.contiguous())
if kpack is not None:
kpacked_idx, kpacked_vals, kpacked_offset, K_pad = kpack
self.register_buffer("psi_kpacked_idx", kpacked_idx, persistent=False)
self.register_buffer("psi_kpacked_vals", kpacked_vals, persistent=False)
self.register_buffer("psi_kpacked_offset", kpacked_offset, persistent=False)
self.psi_kpacked_K_pad = K_pad
self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
self.register_buffer("psi_row_idx", row_idx, persistent=False)
self.register_buffer("psi_col_idx", col_idx, persistent=False)
self.register_buffer("psi_vals", vals, persistent=False)
if not self.optimized_kernel:
# halo mode keys psi's columns onto the halo-padded input band, so the sparse
# tensor's column extent has to be the padded row count. r_lat is 0 under
# reduce-scatter, where the band is just the local share.
self.psi = _get_psi(
self.kernel_size,
self.psi_idx,
self.psi_vals,
self.nlat_in,
self.nlon_in,
self.nlat_out,
self.nlon_out,
self.nlat_in_local + 2 * self.r_lat,
self.nlat_out_local,
)
def extra_repr(self):
return (
f"in_shape={(self.nlat_in, self.nlon_in)}, "
f"out_shape={(self.nlat_out, self.nlon_out)}, "
f"in_chans={self.groupsize * self.groups}, "
f"out_chans={self.weight.shape[0]}, "
f"filter_basis={self.filter_basis}, "
f"kernel_shape={self.kernel_shape}, "
f"theta_cutoff={self.theta_cutoff}, "
f"groups={self.groups}, fused={self.fused}"
)
@property
def psi_idx(self):
return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
def _refresh_kpacked_device_supported(self):
if not hasattr(self, "psi_vals"):
self.kpacked_device_supported = False
return
self.kpacked_device_supported = _kpacked_device_supported_for_tensor(self.psi_vals)
def _apply(self, fn):
result = super()._apply(fn)
self._refresh_kpacked_device_supported()
return result
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.fused:
# reordered a2a: einsum-first on the local channel shard via the
# fused conv op (K-expanded recomputed in backward, not saved).
out = _distributed_disco_fwd_a2a_reordered(
x,
self.weight,
psi_roff_idx=self.psi_roff_idx,
psi_ker_idx=self.psi_ker_idx,
psi_row_idx=self.psi_row_idx,
psi_col_idx=self.psi_col_idx,
psi_vals=self.psi_vals,
psi_split_roff_idx=self.psi_split_roff_idx,
psi_split_nnz_off=self.psi_split_nnz_off,
psi_split_ker_idx=self.psi_split_ker_idx,
psi_split_row_idx=self.psi_split_row_idx,
psi_split_col_idx=self.psi_split_col_idx,
psi_split_vals=self.psi_split_vals,
psi_split_row_offsets=self.psi_split_row_offsets,
psi_split_nnz_offsets=self.psi_split_nnz_offsets,
psi_kpacked_idx=getattr(self, "psi_kpacked_idx", None),
psi_kpacked_vals=getattr(self, "psi_kpacked_vals", None),
psi_kpacked_offset=getattr(self, "psi_kpacked_offset", None),
psi_kpacked_K_pad=self.psi_kpacked_K_pad,
kpacked_device_supported=self.kpacked_device_supported,
kernel_size=self.kernel_size,
nlat_out_local=self.nlat_out_local,
nlon_out=self.nlon_out,
groups=self.groups,
groupsize=self.groupsize,
comm_size_polar=self.comm_size_polar,
comm_size_azimuth=self.comm_size_azimuth,
comm_rank_azimuth=self.comm_rank_azimuth,
lon_in_shapes=self.lon_in_shapes,
use_halo=self.use_halo,
r_lat=self.r_lat,
)
else:
# standard a2a: contraction then einsum after the transpose-back;
# the K-expanded intermediate is saved for backward.
out = _distributed_disco_fwd_a2a(
x,
self.weight,
psi_roff_idx=getattr(self, "psi_roff_idx", None),
psi_ker_idx=self.psi_ker_idx,
psi_row_idx=self.psi_row_idx,
psi_col_idx=self.psi_col_idx,
psi_vals=self.psi_vals,
psi_kpacked_idx=getattr(self, "psi_kpacked_idx", None),
psi_kpacked_vals=getattr(self, "psi_kpacked_vals", None),
psi_kpacked_offset=getattr(self, "psi_kpacked_offset", None),
psi_kpacked_K_pad=self.psi_kpacked_K_pad,
kpacked_device_supported=self.kpacked_device_supported,
psi_torch=getattr(self, "psi", None),
optimized_kernel=self.optimized_kernel,
kernel_size=self.kernel_size,
nlat_out_local=self.nlat_out_local,
nlon_out=self.nlon_out,
groups=self.groups,
groupsize=self.groupsize,
comm_size_polar=self.comm_size_polar,
comm_size_azimuth=self.comm_size_azimuth,
lon_in_shapes=self.lon_in_shapes,
use_halo=self.use_halo,
r_lat=self.r_lat,
)
if self.bias is not None:
out = out + self.bias.reshape(1, self.bias.shape[0], 1, 1)
return out
[docs]
class DistributedDiscreteContinuousConvTransposeS2(DiscreteContinuousConv):
"""
Distributed version of discrete-continuous transpose convolutions (DISCO) on the 2-Sphere as described in :cite:`Ocampo2023`.
.. seealso::
:class:`torch_harmonics.DiscreteContinuousConvTransposeS2`
Serial counterpart with full mathematical description and parameter
documentation.
Parameters
----------
in_channels : int
Number of input channels
out_channels : int
Number of output channels
in_shape : Tuple[int]
Shape of the input tensor
out_shape : Tuple[int]
Shape of the output tensor
kernel_shape : Union[int, Tuple[int], Tuple[int, int]]
Shape of the kernel
basis_type : Optional[str]
Type of basis to use
basis_norm_mode : Optional[str]
Normalization mode for the filter basis
groups : Optional[int]
Number of groups
grid_in : Optional[str]
Grid type for the input tensor
grid_out : Optional[str]
Grid type for the output tensor
bias : Optional[bool]
Whether to use bias
theta_cutoff : Optional[float]
Theta cutoff for the filter basis
Returns
-------
torch.Tensor
Output tensor
References
----------
:cite:`Ocampo2023`
"""
def __init__(
self,
in_channels: int,
out_channels: int,
in_shape: Tuple[int],
out_shape: Tuple[int],
kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
basis_type: Optional[str] = "piecewise linear",
basis_norm_mode: Optional[str] = "nodal",
groups: Optional[int] = 1,
grid_in: Optional[str] = "equiangular",
grid_out: Optional[str] = "equiangular",
bias: Optional[bool] = True,
theta_cutoff: Optional[float] = None,
optimized_kernel: Optional[bool] = True,
):
super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias, optimized_kernel)
self.nlat_in, self.nlon_in = in_shape
self.nlat_out, self.nlon_out = out_shape
# get the comms grid:
self.comm_size_polar = polar_group_size()
self.comm_rank_polar = polar_group_rank()
self.comm_size_azimuth = azimuth_group_size()
self.comm_rank_azimuth = azimuth_group_rank()
# we need those shapes:
self.lat_in_shapes = compute_split_shapes(self.nlat_in, self.comm_size_polar)
self.lon_in_shapes = compute_split_shapes(self.nlon_in, self.comm_size_azimuth)
self.lat_out_shapes = compute_split_shapes(self.nlat_out, self.comm_size_polar)
self.lon_out_shapes = compute_split_shapes(self.nlon_out, self.comm_size_azimuth)
# bandlimit
if theta_cutoff is None:
self.theta_cutoff = compute_theta_cutoff(self.nlat_in, grid=grid_in)
else:
self.theta_cutoff = theta_cutoff
if self.theta_cutoff <= 0.0:
raise ValueError("Error, theta_cutoff has to be positive.")
# Note that the psi matrix is of shape nlat_out x nlat_in * nlon_in. Since the contraction in nlon direction is a convolution,
# we will keep local to all nodes and split the computation up along nlat. We further split the input dim because this reduces the number
# of atomic reduction calls inside the actual kernel
# set local shapes according to distributed mode:
self.nlat_in_local = self.nlat_in
self.nlat_out_local = self.lat_out_shapes[self.comm_rank_polar]
# compute global convolution tensor
# switch in_shape and out_shape since we want transpose conv
# distributed mode here is swapped because of the transpose
idx, vals, _ = _precompute_convolution_tensor_s2(
out_shape,
in_shape,
self.filter_basis,
grid_in=grid_out,
grid_out=grid_in,
theta_cutoff=self.theta_cutoff,
transpose_normalization=True,
basis_norm_mode=basis_norm_mode,
merge_quadrature=True,
)
# split the convolution tensor along latitude, again, we need to swap the meaning
# of in_shape and out_shape
idx, vals = _split_distributed_convolution_tensor_s2(idx, vals, out_shape, in_shape)
# sort the values
ker_idx = idx[0, ...].contiguous()
row_idx = idx[1, ...].contiguous()
col_idx = idx[2, ...].contiguous()
vals = vals.contiguous()
if self.optimized_kernel:
# preprocessed data-structure for GPU kernel
roff_idx = preprocess_psi(self.kernel_size, self.nlat_in_local, ker_idx, row_idx, col_idx, vals).contiguous()
self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
# save all datastructures
self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
self.register_buffer("psi_row_idx", row_idx, persistent=False)
self.register_buffer("psi_col_idx", col_idx, persistent=False)
self.register_buffer("psi_vals", vals, persistent=False)
# store psi as COO
if not self.optimized_kernel:
self.psi_st = _get_psi(
self.kernel_size,
self.psi_idx,
self.psi_vals,
self.nlat_in,
self.nlon_in,
self.nlat_out,
self.nlon_out,
self.nlat_in_local,
self.nlat_out_local,
semi_transposed=True,
)
def extra_repr(self):
return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, theta_cutoff={self.theta_cutoff}, groups={self.groups}"
@property
def psi_idx(self):
return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
def forward(self, x: torch.Tensor) -> torch.Tensor:
# extract shape
B, C, H, W = x.shape
x = x.reshape(B, self.groups, self.groupsize, H, W)
# do weight multiplication
x = torch.einsum("bgcxy,gock->bgokxy", x, self.weight.reshape(self.groups, self.out_per_group, self.weight.shape[1], self.weight.shape[2])).contiguous()
x = x.reshape(B, self.weight.shape[0], x.shape[-3], H, W)
num_chans = x.shape[1]
# transpose such that lon is local, channels are split
if self.comm_size_azimuth > 1:
x = distributed_transpose_azimuth(x, (1, -1), self.lon_in_shapes)
# Fused gather + copy on the polar group. Forward is the same
# all_gather along nlat that gather_from_polar_region performed;
# backward replaces all_reduce + slice with a single reduce_scatter
# (half the backward comm volume on the K-expanded tensor).
# Guarded like the azimuth transposes above: at size 1 both directions
# are identities, and skipping keeps the `torch.compiler.disable()`d
# wrapper out of the graph so this path stays fullgraph-compilable.
if self.comm_size_polar > 1:
x = gather_from_copy_to_polar_region(x, -2, self.lat_in_shapes)
if self.optimized_kernel:
out = _disco_s2_transpose_contraction_optimized(
x, self.psi_roff_idx, self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx, self.psi_vals, self.kernel_size, self.nlat_out_local, self.nlon_out
)
else:
out = _disco_s2_transpose_contraction_torch(x, self.psi_st.to(x.device), self.nlon_out)
# now we can transpose back the result, so that lon is split and channels are local
if self.comm_size_azimuth > 1:
chan_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
out = distributed_transpose_azimuth(out, (-1, 1), chan_shapes)
if self.bias is not None:
out = out + self.bias.reshape(1, self.bias.shape[0], 1, 1)
return out
