private triton

Author: OutisLi

deepmd/pt/model/descriptor/sezm_nn/so2.py

@@ -958,19 +958,27 @@ def __init__(
         self._rotate_to_local_fn = None
         self._rotate_back_fn = None
         if self.use_triton_infer:
-            from .triton import (
+            from .triton.so2_rotation import (
                 rotate_back_block,
                 rotate_back_dense,
                 rotate_to_local_block,
                 rotate_to_local_dense,
             )
 
             if self.mmax == 1:
-                self._rotate_to_local_fn = rotate_to_local_block
-                self._rotate_back_fn = rotate_back_block
+                self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_block(
+                    x, src, wigner, self.lmax
+                )
+                self._rotate_back_fn = lambda x_local, wigner: rotate_back_block(
+                    x_local, wigner, self.lmax
+                )
             else:
-                self._rotate_to_local_fn = rotate_to_local_dense
-                self._rotate_back_fn = rotate_back_dense
+                self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_dense(
+                    x, src, wigner, self.coeff_index_m, self.ebed_dim_full
+                )
+                self._rotate_back_fn = lambda x_local, wigner: rotate_back_dense(
+                    x_local, wigner, self.coeff_index_m, self.ebed_dim_full
+                )
 
         # === Step 1. Precompute coefficient indices for m-major reduced layout ===
         coeff_index_m = build_m_major_index(self.lmax, self.mmax, device=self.device)
@@ -1460,12 +1468,10 @@ def forward(
                 # ``self._rotate_to_local_fn`` was bound in ``__init__`` (the
                 # block kernel for the m-major ``mmax == 1`` layout, dense
                 # otherwise).
-                x_local = self._rotate_to_local_fn(
-                    x, src, D_full, self.coeff_index_m, self.ebed_dim_full
-                )  # (E, D_m, C_wide)
+                x_local = self._rotate_to_local_fn(x, src, D_full)  # (E, D_m, C_wide)
                 if self.node_wise_grid_product is not None:
                     x_dst_local = self._rotate_to_local_fn(
-                        x, dst, D_full, self.coeff_index_m, self.ebed_dim_full
+                        x, dst, D_full
                     )  # (E, D_m, C_wide)
             else:
                 D_m_prime = project_D_to_m(
@@ -1623,12 +1629,7 @@ def apply_bias_correction(
         with nvtx_range("SO2Conv/rotate_back"):
             Dt_full = edge_cache.Dt_full
             if self.use_triton_infer and not self.training:
-                x_message = self._rotate_back_fn(
-                    x_local,
-                    Dt_full,
-                    self.coeff_index_m,
-                    self.ebed_dim_full,
-                )  # (E, D, C_wide)
+                x_message = self._rotate_back_fn(x_local, Dt_full)  # (E, D, C_wide)
             else:
                 Dt_from_m = project_Dt_from_m(
                     Dt_full=Dt_full,

deepmd/pt/model/descriptor/sezm_nn/triton/__init__.py

@@ -2,20 +2,14 @@
 """Hardware-accelerated SeZM/DPA4 operators.
 
 This package hosts ``make_fx``-composable Triton implementations of SeZM hot
-paths.  The SO(2) rotation API exposes a general dense path that honors arbitrary
-coefficient indices and a block path for the canonical m-major ``mmax=1`` layout.
+paths.  Kernel entry points are internal implementation details of the SeZM
+descriptor; the package-level API only exposes availability.
 """
 
 from .so2_rotation import (
-    rotate_back_block,
-    rotate_back_dense,
-    rotate_to_local_block,
-    rotate_to_local_dense,
+    TRITON_ROTATION_AVAILABLE,
 )
 
 __all__ = [
-    "rotate_back_block",
-    "rotate_back_dense",
-    "rotate_to_local_block",
-    "rotate_to_local_dense",
+    "TRITON_ROTATION_AVAILABLE",
 ]

deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py

@@ -58,8 +58,6 @@
     annotations,
 )
 
-import math
-
 import torch
 from torch import (
     Tensor,
@@ -71,12 +69,6 @@
 
 __all__ = [
     "TRITON_ROTATION_AVAILABLE",
-    "rotate_back_block",
-    "rotate_back_dense",
-    "rotate_back_reference",
-    "rotate_to_local_block",
-    "rotate_to_local_dense",
-    "rotate_to_local_reference",
 ]
 
 try:
@@ -1159,30 +1151,8 @@ def _launch_rotate_back_bwd(
 
 
 # ======================================================================
-# Block-diagonal launch wrappers + layout detection (mmax == 1)
+# Block-diagonal launch wrappers (mmax == 1)
 # ======================================================================
-def _block_layout_lmax(coeff_index: Tensor, dim_full: int) -> int:
-    """Return ``lmax`` if ``(coeff_index, dim_full)`` is the m-major ``mmax=1``
-    layout that the block-diagonal kernels assume, else ``-1``.
-
-    This intentionally checks only shape-level invariants.  The block kernels
-    ignore ``coeff_index`` values, so production callers must only use the block
-    entry points when they own the canonical m-major ``mmax=1`` index.
-    """
-    dim_full = int(dim_full)
-    root = math.isqrt(dim_full)
-    if root * root != dim_full:
-        return -1
-    lmax = root - 1
-    try:
-        numel = int(coeff_index.shape[0])
-    except Exception:  # pragma: no cover - exotic shape proxies
-        return -1
-    if lmax < 1 or numel != 3 * lmax + 1:
-        return -1
-    return lmax
-
-
 def _launch_bd_to_local_fwd(
     x: Tensor, src: Tensor, wigner: Tensor, lmax: int
 ) -> Tensor:
@@ -1619,37 +1589,23 @@ def rotate_back_dense(
     return _rotate_back_op(x_local, wigner, coeff_index, int(dim_full))
 
 
-def rotate_to_local_block(
-    x: Tensor, src: Tensor, wigner: Tensor, coeff_index: Tensor, dim_full: int
-) -> Tensor:
+def rotate_to_local_block(x: Tensor, src: Tensor, wigner: Tensor, lmax: int) -> Tensor:
     """Apply the block-diagonal ``global -> local`` rotation.
 
-    Use this only when the caller owns the invariant that ``coeff_index`` is the
-    canonical m-major ``mmax=1`` index produced by
-    :func:`build_m_major_index`.  The kernel ignores the tensor values in
-    ``coeff_index`` and derives the layout from ``lmax``.
+    Use this when the caller owns the invariant that the reduced layout is the
+    canonical m-major ``mmax=1`` layout for ``lmax``.  The block kernel derives
+    the reduced row order from ``lmax`` and does not consume a coefficient-index
+    tensor.
     """
-    lmax = _block_layout_lmax(coeff_index, dim_full)
-    if lmax < 0:
-        raise ValueError(
-            "rotate_to_local_block requires the m-major mmax=1 coefficient layout."
-        )
-    return _block_to_local_op(x, src, wigner, lmax)
+    return _block_to_local_op(x, src, wigner, int(lmax))
 
 
-def rotate_back_block(
-    x_local: Tensor, wigner: Tensor, coeff_index: Tensor, dim_full: int
-) -> Tensor:
+def rotate_back_block(x_local: Tensor, wigner: Tensor, lmax: int) -> Tensor:
     """Apply the block-diagonal ``local -> global`` rotation.
 
-    Use this only when the caller owns the invariant that ``coeff_index`` is the
-    canonical m-major ``mmax=1`` index produced by
-    :func:`build_m_major_index`.  The kernel ignores the tensor values in
-    ``coeff_index`` and derives the layout from ``lmax``.
+    Use this when the caller owns the invariant that ``x_local`` is ordered in
+    the canonical m-major ``mmax=1`` layout for ``lmax``.  The block kernel
+    derives the reduced column order from ``lmax`` and does not consume a
+    coefficient-index tensor.
     """
-    lmax = _block_layout_lmax(coeff_index, dim_full)
-    if lmax < 0:
-        raise ValueError(
-            "rotate_back_block requires the m-major mmax=1 coefficient layout."
-        )
-    return _block_back_op(x_local, wigner, lmax)
+    return _block_back_op(x_local, wigner, int(lmax))

source/tests/pt/model/test_descriptor_sezm_triton.py

@@ -94,28 +94,6 @@ def test_noncanonical_same_length_uses_dense_reference(self):
             rotate_back_reference(x_local, wigner, coeff_index, dim),
         )
 
-    def test_explicit_block_uses_shape_contract_only(self):
-        device = torch.device("cpu")
-        dtype = torch.float32
-        lmax = 3
-        dim = get_so3_dim_of_lmax(lmax)
-        canonical = build_m_major_index(lmax, 1, device=device)
-        coeff_index = torch.roll(canonical, shifts=1)
-        x = torch.randn(4, dim, 3, device=device, dtype=dtype)
-        src = torch.tensor([0, 2, 1, 3, 0], dtype=torch.long, device=device)
-        wigner = torch.randn(src.numel(), dim, dim, device=device, dtype=dtype)
-        x_local = torch.randn(
-            src.numel(), coeff_index.numel(), 3, device=device, dtype=dtype
-        )
-
-        self.assertEqual(
-            rotate_to_local_block(x, src, wigner, coeff_index, dim).shape, x_local.shape
-        )
-        self.assertEqual(
-            rotate_back_block(x_local, wigner, coeff_index, dim).shape,
-            (src.numel(), dim, 3),
-        )
-
     def test_symbolic_trace_noncanonical_same_length_uses_dense_op(self):
         device = torch.device("cpu")
         dtype = torch.float32
@@ -186,7 +164,7 @@ def _assert_to_local_matches_reference(self, x0, src, w0, coeff_index, dim):
 
         xa = x0.clone().requires_grad_(True)
         wa = w0.clone().requires_grad_(True)
-        out = rotate_to_local_block(xa, src, wa, coeff_index, dim)
+        out = rotate_to_local_block(xa, src, wa, lmax)
         xr = x0.clone().requires_grad_(True)
         wr = w0.clone().requires_grad_(True)
         ref = rotate_to_local_reference(xr, src, wr, coeff_index, dim)
@@ -206,7 +184,7 @@ def _assert_back_matches_reference(self, xl0, w0, coeff_index, dim):
 
         xa = xl0.clone().requires_grad_(True)
         wa = w0.clone().requires_grad_(True)
-        out = rotate_back_block(xa, wa, coeff_index, dim)
+        out = rotate_back_block(xa, wa, lmax)
         xr = xl0.clone().requires_grad_(True)
         wr = w0.clone().requires_grad_(True)
         ref = rotate_back_reference(xr, wr, coeff_index, dim)
@@ -248,7 +226,7 @@ def test_symbolic_make_fx_rotate_to_local_forward_backward_matches_eager(self):
         def forward_and_grad(x, wigner):
             x_req = x.detach().requires_grad_(True)
             w_req = wigner.detach().requires_grad_(True)
-            out = rotate_to_local_block(x_req, src, w_req, coeff_index, dim)
+            out = rotate_to_local_block(x_req, src, w_req, lmax)
             grad_x, grad_wigner = torch.autograd.grad(
                 out,
                 (x_req, w_req),
@@ -291,7 +269,7 @@ def test_symbolic_make_fx_rotate_back_forward_backward_matches_eager(self):
         def forward_and_grad(x_local, wigner):
             x_req = x_local.detach().requires_grad_(True)
             w_req = wigner.detach().requires_grad_(True)
-            out = rotate_back_block(x_req, w_req, coeff_index, dim)
+            out = rotate_back_block(x_req, w_req, lmax)
             grad_x, grad_wigner = torch.autograd.grad(
                 out,
                 (x_req, w_req),