deepmd.pt_expt.kernels.cutile.sezm.indexing

Coefficient layout of the reduced SO(2) basis as the cuTile kernels see it.

The canonical reduced layout is defined once, in deepmd.dpmodel.descriptor.dpa4_nn.indexing, and is reused here rather than restated. This module adds the two things the tile programming model needs on top of it: the structural non-zeros of the reduced rotation, which fix the coefficient numbering every generated kernel shares, and the power-of-two padded extents, which every tile must have.

Three layouts appear throughout the package and the kernels must agree on them exactly:

full basis

(lmax + 1)^2 spherical-harmonic coefficients ordered by (l, m), the layout of the node features and of the per-edge Wigner-D blocks.

reduced m-major

3 * lmax + 1 rows at mmax == 1: the m = 0 degrees first, then m = -1, then m = +1. The mixing stack operates here, and the layout is block diagonal over |m|.

padded groups

each |m| group widened to a power-of-two degree count. Padded degrees carry exact zeros in both the activation and the weight, so a padded contraction equals the exact one.

Classes

SO2TileLayout

Every extent the generated kernels of one block layout need.

Functions

m_major_index(→ list[int])	Return the reduced rows as indices into the full (lmax + 1)^2 basis.
rotation_pairs(→ list[tuple[int, int]])	Enumerate the structural non-zeros of the reduced rotation.

Module Contents

deepmd.pt_expt.kernels.cutile.sezm.indexing.m_major_index(lmax: int, mmax: int = 1) → list[int]

Return the reduced rows as indices into the full (lmax + 1)^2 basis.

Parameters:

lmaxint

Maximum spherical-harmonic degree.

mmaxint

Maximum retained absolute order.

Returns:

list[int]

Full-basis index of each reduced row, in reduced-row order. Python ints, because the generated kernels embed them as literals.

deepmd.pt_expt.kernels.cutile.sezm.indexing.rotation_pairs(lmax: int, mmax: int = 1) → list[tuple[int, int]]

Enumerate the structural non-zeros of the reduced rotation.

The Wigner-D matrix is block diagonal in the degree, so projecting onto the reduced rows couples reduced row r only to the full-basis rows of its own degree. The returned order defines the coefficient numbering used by every generated kernel that touches the rotation.

Parameters:

lmaxint

Maximum spherical-harmonic degree.

mmaxint

Maximum retained absolute order.

Returns:

list[tuple[int, int]]

(reduced_row, full_row) pairs, sum_r (2 * l_r + 1) of them.

class deepmd.pt_expt.kernels.cutile.sezm.indexing.SO2TileLayout

Every extent the generated kernels of one block layout need.

Attributes:

lmax, focus_dim, n_layers

Configuration of the convolution block.

lmax: int

focus_dim: int

n_layers: int

property key: tuple[int, int]

Shape key for the launch-configuration tables.

property n_m0: int

Real degree count of the m = 0 group.

property n_m1: int

Real degree count of the |m| = 1 group.

property pad_m0: int

m = 0 degree count rounded up to a power of two.

property pad_m1: int

|m| = 1 degree count rounded up to a power of two.

property width_m0: int

Padded m = 0 group width in channels, a tile extent of the stack.

property width_m1: int

Padded |m| = 1 group width in channels, a tile extent of the stack.

property dim: int

Full-basis coefficient count.

property row: int

Reduced-layout width in channels, the stride of the activation.

property kernel_size: int

Compact per-edge degree-mixing kernel size at rank one.

property n_gated: int

Number of gated layers; the final layer is the identity layer.