diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index dcb499fe4e..ca99478c32 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -84,6 +84,8 @@ extend_descrpt_stat, ) +_DEGREE_GAIN_INIT_STD = 0.1 + def np_softmax(x: Array, axis: int = -1) -> Array: xp = array_api_compat.array_namespace(x) @@ -98,6 +100,121 @@ def np_normalize(x: Array, axis: int = -1) -> Array: return x / xp.linalg.vector_norm(x, axis=axis, keepdims=True) +def build_dpa1_moment_basis( + rr: Array, + diff: Array, + switch: Array, + radial_stddev: Array, + valid_mask: Array, + lmax: int, + protection: float, +) -> Array: + """Build the DPA1 Cartesian moment basis through angular degree ``lmax``. + + Parameters + ---------- + rr + Normalized environment matrix with shape ``(..., 4)``. + diff + Neighbor displacement vectors with shape ``(..., 3)``. + switch + Smooth cutoff values with shape ``(..., 1)``. + radial_stddev + Scalar environment standard deviation with shape ``(..., 1)``. + valid_mask + Valid non-excluded neighbor mask with shape ``(...)``. + lmax + Maximum angular degree. Supported values are 1 through 4. + protection + Distance protection added to the radial denominator. + + Returns + ------- + Array + Moment basis with shape ``(..., (lmax + 1) ** 2)``. + """ + if lmax == 1: + return rr + + xp = array_api_compat.array_namespace(rr, diff, switch, radial_stddev) + distance_squared = xp.sum(diff * diff, axis=-1, keepdims=True) + direction_mask = distance_squared > 0.0 + safe_distance = xp.sqrt( + xp.where(direction_mask, distance_squared, xp.ones_like(distance_squared)) + ) + basis_mask = valid_mask[..., None] & direction_mask + denominator = xp.where( + basis_mask, + safe_distance + protection, + xp.ones_like(safe_distance), + ) + direction = diff / denominator * xp.astype(basis_mask, diff.dtype) + radial = switch / denominator / radial_stddev * xp.astype(basis_mask, switch.dtype) + + x, y, z = direction[..., 0], direction[..., 1], direction[..., 2] + q = x * x + y * y + z * z + sqrt_three = math.sqrt(3.0) + degree_two = xp.stack( + [ + sqrt_three * x * y, + sqrt_three * y * z, + 0.5 * (3.0 * z * z - q), + sqrt_three * x * z, + 0.5 * sqrt_three * (x * x - y * y), + ], + axis=-1, + ) + blocks = [rr, radial * degree_two] + if lmax >= 3: + degree_three = xp.stack( + [ + math.sqrt(5.0 / 8.0) * y * (3.0 * x * x - y * y), + math.sqrt(15.0) * x * y * z, + math.sqrt(3.0 / 8.0) * y * (5.0 * z * z - q), + 0.5 * z * (5.0 * z * z - 3.0 * q), + math.sqrt(3.0 / 8.0) * x * (5.0 * z * z - q), + 0.5 * math.sqrt(15.0) * z * (x * x - y * y), + math.sqrt(5.0 / 8.0) * x * (x * x - 3.0 * y * y), + ], + axis=-1, + ) + blocks.append(radial * degree_three) + if lmax >= 4: + z2 = z * z + x2_minus_y2 = x * x - y * y + degree_four = xp.stack( + [ + 0.5 * math.sqrt(35.0) * x * y * x2_minus_y2, + 0.25 * math.sqrt(70.0) * y * z * (3.0 * x * x - y * y), + 0.5 * math.sqrt(5.0) * x * y * (7.0 * z2 - q), + 0.25 * math.sqrt(10.0) * y * z * (7.0 * z2 - 3.0 * q), + 0.125 * (35.0 * z2 * z2 - 30.0 * z2 * q + 3.0 * q * q), + 0.25 * math.sqrt(10.0) * x * z * (7.0 * z2 - 3.0 * q), + 0.25 * math.sqrt(5.0) * x2_minus_y2 * (7.0 * z2 - q), + 0.25 * math.sqrt(70.0) * x * z * (x * x - 3.0 * y * y), + 0.125 * math.sqrt(35.0) * (x**4 - 6.0 * x * x * y * y + y**4), + ], + axis=-1, + ) + blocks.append(radial * degree_four) + return xp.concat(blocks, axis=-1) + + +def build_dpa1_degree_weights( + raw_gain: Array, + lmax: int, + reference: Array, +) -> Array: + """Expand non-negative per-degree Gram weights to packed moment rows.""" + xp = array_api_compat.array_namespace(raw_gain, reference) + device = array_api_compat.device(reference) + blocks = [xp.ones((4,), dtype=reference.dtype, device=device)] + for degree in range(2, lmax + 1): + weight = raw_gain[degree - 2] * raw_gain[degree - 2] + blocks.append(xp.broadcast_to(weight, (2 * degree + 1,))) + return xp.concat(blocks) + + @BaseDescriptor.register("se_atten") @BaseDescriptor.register("dpa1") class DescrptDPA1(NativeOP, BaseDescriptor): @@ -233,6 +350,9 @@ class DescrptDPA1(NativeOP, BaseDescriptor): Whether to use bias in the type embedding layer. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + lmax: int + Maximum angular degree of the Cartesian moment basis. Supported + values are 1 through 4. spin (Only support None to keep consistent with other backend references.) (Not used in this version. Not-none option is not implemented.) @@ -289,6 +409,7 @@ def __init__( type_map: list[str] | None = None, # consistent with argcheck, not used though seed: int | list[int] | None = None, + lmax: int = 1, ) -> None: ## seed, uniform_seed, not included. # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -333,6 +454,7 @@ def __init__( ln_eps=ln_eps, seed=child_seed(seed, 0), trainable=trainable, + lmax=lmax, ) self.use_econf_tebd = use_econf_tebd self.use_tebd_bias = use_tebd_bias @@ -980,7 +1102,7 @@ def serialize(self) -> dict: data = { "@class": "Descriptor", "type": "dpa1", - "@version": 3 if self.compress else 2, + "@version": 4 if obj.lmax != 1 else (3 if self.compress else 2), "rcut": obj.rcut, "rcut_smth": obj.rcut_smth, "sel": obj.sel, @@ -1023,6 +1145,11 @@ def serialize(self) -> dict: "trainable": self.trainable, "spin": None, } + if obj.lmax != 1: + data["lmax"] = obj.lmax + data["@variables"]["degree_gain_raw"] = to_numpy_array( + obj.adam_degree_gain_raw + ) if obj.tebd_input_mode in ["strip"]: data.update({"embeddings_strip": obj.embeddings_strip.serialize()}) if self.compress: @@ -1061,7 +1188,7 @@ def serialize(self) -> dict: def deserialize(cls, data: dict) -> "DescrptDPA1": """Deserialize from dict.""" data = data.copy() - check_version_compatibility(data.pop("@version"), 3, 1) + check_version_compatibility(data.pop("@version"), 4, 1) data.pop("@class") data.pop("type") variables = data.pop("@variables") @@ -1078,10 +1205,16 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": # compat with version 1 if "use_tebd_bias" not in data: data["use_tebd_bias"] = True + data.setdefault("lmax", 1) obj = cls(**data) obj.se_atten["davg"] = variables["davg"] obj.se_atten["dstd"] = variables["dstd"] + if obj.se_atten.lmax > 1: + obj.se_atten.adam_degree_gain_raw = np.asarray( + variables["degree_gain_raw"], + dtype=PRECISION_DICT[obj.se_atten.precision], + ) obj.se_atten.embeddings = NetworkCollection.deserialize(embeddings) if tebd_input_mode in ["strip"]: obj.se_atten.embeddings_strip = NetworkCollection.deserialize( @@ -1222,6 +1355,9 @@ class DescrptBlockSeAtten(NativeOP, DescriptorBlock): Random seed for parameter initialization. trainable : bool, optional If the parameters are trainable. + lmax : int, optional + Maximum angular degree of the Cartesian moment basis. Supported values + are 1 through 4. """ def __init__( @@ -1253,6 +1389,7 @@ def __init__( smooth: bool = True, seed: int | list[int] | None = None, trainable: bool = True, + lmax: int = 1, ) -> None: self.rcut = rcut self.rcut_smth = rcut_smth @@ -1265,6 +1402,9 @@ def __init__( self.neuron = neuron self.filter_neuron = self.neuron self.axis_neuron = axis_neuron + if lmax not in (1, 2, 3, 4): + raise ValueError(f"`lmax` must be between 1 and 4, got {lmax}") + self.lmax = int(lmax) self.tebd_dim = tebd_dim self.tebd_input_mode = tebd_input_mode self.resnet_dt = resnet_dt @@ -1284,6 +1424,16 @@ def __init__( self.normalize = normalize self.temperature = temperature self.smooth = smooth + self.trainable = bool(trainable) + if self.lmax > 1: + gain_rng = np.random.default_rng(child_seed(seed, 3)) + self.adam_degree_gain_raw = gain_rng.normal( + loc=0.0, + scale=_DEGREE_GAIN_INIT_STD, + size=(self.lmax - 1,), + ).astype(PRECISION_DICT[self.precision]) + else: + self.adam_degree_gain_raw = None # order matters, placed after the assignment of self.ntypes self.reinit_exclude(exclude_types) @@ -1638,6 +1788,23 @@ def call( rr = rr * xp.astype(exclude_mask[:, :, None], rr.dtype) # nfnl x nnei x 1 ss = rr[..., 0:1] + moment_basis = rr + if self.lmax > 1: + diff_flat = xp.reshape(diff, (nf * nloc, nnei, 3)) + radial_stddev = xp.take( + self.stddev, + xp.reshape(atype, (-1,)), + axis=0, + )[..., 0:1] + moment_basis = build_dpa1_moment_basis( + rr, + diff_flat, + sw, + radial_stddev, + nlist_mask, + self.lmax, + self.env_protection, + ) geo_gr = None if self.tebd_input_mode in ["concat"]: # nfnl x tebd_dim @@ -1731,7 +1898,7 @@ def call( self.compress_data[0], self.compress_info[0], ss_scalar, - rr, + moment_basis, gg_t, self.filter_neuron[-1], ) @@ -1754,15 +1921,21 @@ def call( gg = self.dpa1_attention( gg, nlist_mask, input_r=input_r, sw=sw ) # shape is [nframes*nloc, self.neei, out_size] - # nfnl x ng x 4 - # gr = xp.einsum("lni,lnj->lij", gg, rr) - gr = xp.sum(gg[:, :, :, None] * rr[:, :, None, :], axis=1) + # nfnl x ng x moment_dim + gr = xp.sum(gg[:, :, :, None] * moment_basis[:, :, None, :], axis=1) g2 = xp.reshape(gg, (nf, nloc, self.nnei, self.filter_neuron[-1])) else: gr = xp.permute_dims(geo_gr, (0, 2, 1)) g2 = None gr /= self.nnei gr1 = gr[:, : self.axis_neuron, :] + if self.lmax > 1: + degree_weights = build_dpa1_degree_weights( + self.adam_degree_gain_raw, + self.lmax, + gr, + ) + gr1 = gr1 * degree_weights[None, None, :] # nfnl x ng x ng1 # grrg = xp.einsum("lid,ljd->lij", gr, gr1) grrg = xp.sum(gr[:, :, None, :] * gr1[:, None, :, :], axis=3) @@ -1774,7 +1947,7 @@ def call( xp.reshape(grrg, (nf, nloc, self.filter_neuron[-1] * self.axis_neuron)), g2, xp.reshape(dmatrix, (nf, nloc, self.nnei, 4))[..., 1:], - xp.reshape(gr[..., 1:], (nf, nloc, self.filter_neuron[-1], 3)), + xp.reshape(gr[..., 1:4], (nf, nloc, self.filter_neuron[-1], 3)), xp.reshape(sw, (nf, nloc, nnei, 1)), ) @@ -1869,6 +2042,18 @@ def call_graph( ) # (E, 4), (E, 1) sw zeroed on padding # radial channel ss = rr[:, 0:1] # (E, 1) + moment_basis = rr + if self.lmax > 1: + radial_stddev = xp.take(self.stddev[:, 0, 0:1], center_type, axis=0) + moment_basis = build_dpa1_moment_basis( + rr, + graph.edge_vec, + sw_e, + radial_stddev, + graph.edge_mask, + self.lmax, + self.env_protection, + ) if self.tebd_input_mode == "concat": # neighbor / center type embeddings; ghost type == owner type so # gathering by the LOCAL owner (src) reproduces the dense neighbor tebd. @@ -1897,11 +2082,18 @@ def call_graph( ) # zero padding/guard edges BEFORE the segment sum gg = gg * xp.astype(graph.edge_mask[:, None], gg.dtype) - # outer product (replaces the dense gg[:,:,:,None] * rr[:,:,None,:]) - outer = gg[:, :, None] * rr[:, None, :] # (E, ng, 4) + # outer product (replaces the dense neighbor-axis moment reduction) + outer = gg[:, :, None] * moment_basis[:, None, :] # (E, ng, moment_dim) # neighbor-axis reduction -> segment_sum over centers; divide by nnei - gr = segment_sum(outer, dst, n_total) / self.nnei # (N, ng, 4) + gr = segment_sum(outer, dst, n_total) / self.nnei gr1 = gr[:, : self.axis_neuron, :] + if self.lmax > 1: + degree_weights = build_dpa1_degree_weights( + self.adam_degree_gain_raw, + self.lmax, + gr, + ) + gr1 = gr1 * degree_weights[None, None, :] # nf x nloc x (ng x ng1) grrg = xp.sum(gr[:, :, None, :] * gr1[:, None, :, :], axis=3) # (N, ng, ng1) ng = self.neuron[-1] @@ -1912,7 +2104,7 @@ def call_graph( # equivariant single-particle representation, dense-ABI slice gr[..., 1:] # (N, ng, 3); not cast, mirroring the dense block which leaves rot_mat in # the working precision before the descriptor-level @cast_precision. - rot_mat = gr[:, :, 1:] + rot_mat = gr[:, :, 1:4] return grrg, rot_mat def _graph_edge_gg_strip( @@ -2133,7 +2325,7 @@ def serialize(self) -> dict: data = { "@class": "DescriptorBlock", "type": "dpa1", - "@version": 1, + "@version": 2 if obj.lmax != 1 else 1, "rcut": obj.rcut, "rcut_smth": obj.rcut_smth, "sel": obj.sel, @@ -2155,6 +2347,7 @@ def serialize(self) -> dict: "trainable_ln": obj.trainable_ln, "ln_eps": obj.ln_eps, "smooth": obj.smooth, + "trainable": obj.trainable, "type_one_side": obj.type_one_side, # make deterministic "precision": np.dtype(PRECISION_DICT[obj.precision]).name, @@ -2168,6 +2361,11 @@ def serialize(self) -> dict: "dstd": to_numpy_array(obj["dstd"]), }, } + if obj.lmax != 1: + data["lmax"] = obj.lmax + data["@variables"]["degree_gain_raw"] = to_numpy_array( + obj.adam_degree_gain_raw + ) if obj.tebd_input_mode in ["strip"]: data.update({"embeddings_strip": obj.embeddings_strip.serialize()}) return data @@ -2176,7 +2374,7 @@ def serialize(self) -> dict: def deserialize(cls, data: dict) -> "DescrptDPA1": """Deserialize from dict.""" data = data.copy() - check_version_compatibility(data.pop("@version"), 1, 1) + check_version_compatibility(data.pop("@version"), 2, 1) data.pop("@class") data.pop("type") variables = data.pop("@variables") @@ -2188,10 +2386,16 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": embeddings_strip = data.pop("embeddings_strip") else: embeddings_strip = None + data.setdefault("lmax", 1) obj = cls(**data) obj["davg"] = variables["davg"] obj["dstd"] = variables["dstd"] + if obj.lmax > 1: + obj.adam_degree_gain_raw = np.asarray( + variables["degree_gain_raw"], + dtype=PRECISION_DICT[obj.precision], + ) obj.embeddings = NetworkCollection.deserialize(embeddings) if tebd_input_mode in ["strip"]: obj.embeddings_strip = NetworkCollection.deserialize(embeddings_strip) diff --git a/deepmd/dpmodel/descriptor/se_atten_v2.py b/deepmd/dpmodel/descriptor/se_atten_v2.py index 6942a5bac3..ec7e3e4f77 100644 --- a/deepmd/dpmodel/descriptor/se_atten_v2.py +++ b/deepmd/dpmodel/descriptor/se_atten_v2.py @@ -122,6 +122,9 @@ class DescrptSeAttenV2(DescrptDPA1): A list of strings. Give the name to each type of atoms. seed : int, Optional Random seed for initializing the network parameters. + lmax : int + Maximum angular degree of the Cartesian moment basis. Supported + values are 1 through 4. """ def __init__( @@ -158,6 +161,7 @@ def __init__( type_map: list[str] | None = None, # consistent with argcheck, not used though seed: int | list[int] | None = None, + lmax: int = 1, ) -> None: DescrptDPA1.__init__( self, @@ -195,6 +199,7 @@ def __init__( type_map=type_map, # consistent with argcheck, not used though seed=seed, + lmax=lmax, ) self.compress = False @@ -204,7 +209,7 @@ def serialize(self) -> dict: data = { "@class": "Descriptor", "type": "se_atten_v2", - "@version": 3 if self.compress else 2, + "@version": 4 if obj.lmax != 1 else (3 if self.compress else 2), "rcut": obj.rcut, "rcut_smth": obj.rcut_smth, "sel": obj.sel, @@ -246,6 +251,11 @@ def serialize(self) -> dict: "trainable": self.trainable, "spin": None, } + if obj.lmax != 1: + data["lmax"] = obj.lmax + data["@variables"]["degree_gain_raw"] = to_numpy_array( + obj.adam_degree_gain_raw + ) if self.compress: type_embd_data = ( self.type_embd_data @@ -282,7 +292,7 @@ def serialize(self) -> dict: def deserialize(cls, data: dict) -> "DescrptSeAttenV2": """Deserialize from dict.""" data = data.copy() - check_version_compatibility(data.pop("@version"), 3, 1) + check_version_compatibility(data.pop("@version"), 4, 1) data.pop("@class") data.pop("type") variables = data.pop("@variables") @@ -295,10 +305,16 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2": # compat with version 1 if "use_tebd_bias" not in data: data["use_tebd_bias"] = True + data.setdefault("lmax", 1) obj = cls(**data) obj.se_atten["davg"] = variables["davg"] obj.se_atten["dstd"] = variables["dstd"] + if obj.se_atten.lmax > 1: + obj.se_atten.adam_degree_gain_raw = np.asarray( + variables["degree_gain_raw"], + dtype=PRECISION_DICT[obj.se_atten.precision], + ) obj.se_atten.embeddings = NetworkCollection.deserialize(embeddings) obj.se_atten.embeddings_strip = NetworkCollection.deserialize(embeddings_strip) obj.type_embedding = TypeEmbedNet.deserialize(type_embedding) diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py index 8af278caa7..ebfdd9981d 100644 --- a/deepmd/kernels/cuda/dpa1/canonical.py +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -70,6 +70,7 @@ def _forward_fake( type_embedding: torch.Tensor, average: torch.Tensor, inverse_stddev: torch.Tensor, + degree_gain: torch.Tensor, table: torch.Tensor, gate_table: torch.Tensor, type_one_side: int, @@ -86,12 +87,14 @@ def _forward_fake( rcut_smooth: float, protection: float, neighbors: float, + basis_dim: int, ) -> tuple[torch.Tensor, ...]: del ( source, destination_row_ptr, average, inverse_stddev, + degree_gain, gate_table, type_one_side, smooth, @@ -117,7 +120,7 @@ def _forward_fake( width, 3, ), - edge_vec.new_empty(node_count, 4, width), + edge_vec.new_empty(node_count, basis_dim, width), ) @@ -131,6 +134,7 @@ def _backward_fake( atype: torch.Tensor, average: torch.Tensor, inverse_stddev: torch.Tensor, + degree_gain: torch.Tensor, table: torch.Tensor, gate_table: torch.Tensor, type_one_side: int, @@ -155,6 +159,7 @@ def _backward_fake( atype, average, inverse_stddev, + degree_gain, table, gate_table, type_one_side, @@ -218,9 +223,9 @@ def _cpu_forward(*args: Any) -> tuple[torch.Tensor, ...]: edge_mask, destination_order, destination_row_ptr, - *tail[:11], + *tail[:12], True, - *tail[11:], + *tail[12:], ) @@ -245,9 +250,9 @@ def _cpu_backward(*args: Any) -> torch.Tensor: edge_mask, destination_order, destination_row_ptr, - *tail[:8], + *tail[:9], True, - *tail[8:], + *tail[9:], ) @@ -329,6 +334,11 @@ def dpa1_canonical_compress_energy_force( compress_data = desc.compress_data[0].contiguous() gate_table = desc.type_embd_data.contiguous() inverse_stddev = torch.reciprocal(se.stddev[:, 0, :]).contiguous() + degree_gain = ( + se.adam_degree_gain_raw.to(torch.float32).contiguous() + if se.adam_degree_gain_raw is not None + else compress_data.new_empty(0) + ) from torch.fx.experimental.proxy_tensor import ( disable_proxy_modes_tracing, ) @@ -346,6 +356,7 @@ def dpa1_canonical_compress_energy_force( type_embedding, se.mean[:, 0, :].contiguous(), inverse_stddev, + degree_gain, compress_data, gate_table, int(se.type_one_side), @@ -362,6 +373,7 @@ def dpa1_canonical_compress_energy_force( float(se.rcut_smth), float(se.env_protection), float(se.nnei), + (int(se.lmax) + 1) ** 2, ) *hidden, head = fit.nets[0].layers @@ -426,6 +438,7 @@ def dpa1_canonical_compress_energy_force( atype, se.mean[:, 0, :].contiguous(), inverse_stddev, + degree_gain, compress_data, gate_table, int(se.type_one_side), diff --git a/deepmd/kernels/cuda/dpa1/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py index 1bad0501d1..27a8c958c8 100644 --- a/deepmd/kernels/cuda/dpa1/graph_compress.py +++ b/deepmd/kernels/cuda/dpa1/graph_compress.py @@ -26,6 +26,11 @@ import torch +from deepmd.dpmodel.descriptor.dpa1 import ( + build_dpa1_degree_weights, + build_dpa1_moment_basis, +) + __all__ = [ "dpa1_graph_compress", "dpa1_graph_compress_energy_force", @@ -101,6 +106,7 @@ def _forward_fake( type_embedding: torch.Tensor, davg: torch.Tensor, inverse_stddev: torch.Tensor, + degree_gain: torch.Tensor, table: torch.Tensor, gate_table: torch.Tensor, type_one_side: int, @@ -118,6 +124,7 @@ def _forward_fake( rcut_smth: float, protection: float, nnei: float, + basis_dim: int, ) -> tuple[torch.Tensor, ...]: n_node = atype.shape[0] ng = table.shape[1] // 6 @@ -133,7 +140,7 @@ def _forward_fake( dtype=torch.float32, device=dev, ), - torch.empty(n_node, 4, ng, dtype=torch.float32, device=dev), + torch.empty(n_node, basis_dim, ng, dtype=torch.float32, device=dev), ) @@ -161,6 +168,7 @@ def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: type_embedding, davg, inverse_stddev, + degree_gain, table, gate_table, type_one_side, @@ -178,6 +186,7 @@ def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: rcut_smth, protection, nnei, + _basis_dim, ) = inputs (gr,) = output[2:] ctx.save_for_backward( @@ -190,6 +199,7 @@ def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: atype, davg, inverse_stddev, + degree_gain, table, gate_table, ) @@ -228,6 +238,7 @@ def _backward( atype, davg, inverse_stddev, + degree_gain, table, gate_table, ) = ctx.saved_tensors @@ -259,6 +270,7 @@ def _backward( atype, davg, inverse_stddev, + degree_gain, table, gate_table, type_one_side, @@ -275,7 +287,7 @@ def _backward( protection, nnei, ) - return (d_edge_vec,) + (None,) * 25 + return (d_edge_vec,) + (None,) * 27 # ====================================================================== @@ -431,10 +443,11 @@ def _cpu_env_and_gg( rcut: float, rcut_smth: float, protection: float, + basis_dim: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Environment matrix, tabulated geometric net and strip gate (fp32). - Returns ``(rr, outer)`` where ``outer`` is ``(E, 4, ng)`` before the + Returns ``(rr, outer)`` where ``outer`` is ``(E, basis_dim, ng)`` before the neighbor-axis reduction. """ ev = edge_vec.to(torch.float32) @@ -447,6 +460,18 @@ def _cpu_env_and_gg( sw = u**3 * (-6 * u**2 + 15 * u - 10) + 1.0 em = torch.cat([sw / q, ev * (sw / q**2)], dim=-1) rr = (em - davg[center_type]) / dstd[center_type] + moment_basis = rr + if basis_dim > 4: + lmax = {9: 2, 16: 3, 25: 4}[basis_dim] + moment_basis = build_dpa1_moment_basis( + rr, + ev, + sw, + dstd[center_type, 0:1], + edge_mask, + lmax, + protection, + ) ss = rr[:, 0:1] pair_idx = _cpu_pair_idx(edge_index, atype, type_one_side, ntypes) gate = gate_table[pair_idx] @@ -454,19 +479,20 @@ def _cpu_env_and_gg( gate = gate * sw n_edge = edge_vec.shape[0] em_x = ss.reshape(n_edge, 1) - em_rr = rr.reshape(n_edge, 1, 4) + em_rr = moment_basis.reshape(n_edge, 1, basis_dim) two_embed = gate.reshape(n_edge, 1, ng) outer = _cpu_tabulate_se_atten(table, info, em_x, em_rr, two_embed, ng) - return rr, outer + return moment_basis, outer def _cpu_outputs( - rr: torch.Tensor, + moment_basis: torch.Tensor, outer: torch.Tensor, edge_index: torch.Tensor, edge_mask: torch.Tensor, atype: torch.Tensor, type_embedding: torch.Tensor | None, + degree_gain: torch.Tensor, ng: int, axis: int, concat_tebd: int, @@ -474,11 +500,22 @@ def _cpu_outputs( n_node: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: outer = outer * edge_mask[:, None, None].to(outer.dtype) - gr = torch.zeros(n_node, 4, ng, dtype=outer.dtype, device=outer.device) + gr = torch.zeros( + n_node, + moment_basis.shape[-1], + ng, + dtype=outer.dtype, + device=outer.device, + ) gr.index_add_(0, edge_index[1], outer) gr = gr / nnei gr_t = gr.permute(0, 2, 1) - grrg = torch.matmul(gr_t, gr[:, :, :axis]).reshape(n_node, ng * axis) + gr_axis = gr[:, :, :axis] + if moment_basis.shape[-1] > 4: + lmax = {9: 2, 16: 3, 25: 4}[moment_basis.shape[-1]] + degree_weights = build_dpa1_degree_weights(degree_gain, lmax, gr) + gr_axis = gr_axis * degree_weights.view(1, -1, 1) + grrg = torch.matmul(gr_t, gr_axis).reshape(n_node, ng * axis) rot_mat = gr_t[:, :, 1:4].contiguous() if concat_tebd: grrg = torch.cat([grrg, type_embedding[atype]], dim=-1) @@ -495,6 +532,7 @@ def _cpu_forward( type_embedding: torch.Tensor, davg: torch.Tensor, inverse_stddev: torch.Tensor, + degree_gain: torch.Tensor, table: torch.Tensor, gate_table: torch.Tensor, type_one_side: int, @@ -512,6 +550,7 @@ def _cpu_forward( rcut_smth: float, protection: float, nnei: float, + basis_dim: int, ) -> tuple[torch.Tensor, ...]: n_node = atype.shape[0] ng = table.shape[1] // 6 @@ -521,9 +560,9 @@ def _cpu_forward( return ( table.new_empty(0, out_dim), table.new_empty(0, ng, 3), - table.new_empty(0, 4, ng), + table.new_empty(0, basis_dim, ng), ) - rr, outer = _cpu_env_and_gg( + moment_basis, outer = _cpu_env_and_gg( edge_vec, edge_index, edge_mask, @@ -540,14 +579,16 @@ def _cpu_forward( rcut, rcut_smth, protection, + basis_dim, ) grrg, rot_mat, gr = _cpu_outputs( - rr, + moment_basis, outer, edge_index, edge_mask, atype, type_embedding, + degree_gain, ng, axis, concat_tebd, @@ -571,6 +612,7 @@ def _cpu_backward( atype: torch.Tensor, davg: torch.Tensor, inverse_stddev: torch.Tensor, + degree_gain: torch.Tensor, table: torch.Tensor, gate_table: torch.Tensor, type_one_side: int, @@ -589,12 +631,13 @@ def _cpu_backward( ) -> torch.Tensor: n_node = atype.shape[0] ng = table.shape[1] // 6 + basis_dim = gr.shape[1] if n_node == 0: return torch.zeros_like(edge_vec) ntypes = gate_table.shape[0] if type_one_side else round(gate_table.shape[0] ** 0.5) ev = edge_vec.detach().clone().requires_grad_(True) with torch.enable_grad(): - rr, outer = _cpu_env_and_gg( + moment_basis, outer = _cpu_env_and_gg( ev, edge_index, edge_mask, @@ -611,14 +654,16 @@ def _cpu_backward( rcut, rcut_smth, protection, + basis_dim, ) grrg, rot_mat, _gr = _cpu_outputs( - rr, + moment_basis, outer, edge_index, edge_mask, atype, None, + degree_gain, ng, axis, 0, @@ -720,6 +765,11 @@ def dpa1_graph_compress( ) if graph.destination_order is None or graph.destination_row_ptr is None: raise ValueError("dpa1_graph_compress requires destination CSR topology") + degree_gain = ( + se.adam_degree_gain_raw.to(torch.float32).contiguous() + if se.adam_degree_gain_raw is not None + else compress_data.new_empty(0) + ) grrg, rot_mat, _moment = torch.ops.deepmd.dpa1_graph_compress( graph.edge_vec.contiguous(), graph.edge_index.contiguous(), @@ -730,6 +780,7 @@ def dpa1_graph_compress( type_embedding.contiguous(), se.mean[:, 0, :].contiguous(), torch.reciprocal(se.stddev[:, 0, :]).contiguous(), + degree_gain, compress_data, gate_table, int(se.type_one_side), @@ -747,6 +798,7 @@ def dpa1_graph_compress( float(se.rcut_smth), float(se.env_protection), float(se.nnei), + (int(se.lmax) + 1) ** 2, ) if pad: # Drop the padding channels: the descriptor is stored channel-major @@ -862,6 +914,11 @@ def dpa1_graph_compress_energy_force( *hidden, head = fit.nets[0].layers fempty = hidden[0].w.new_empty(0) inverse_stddev = torch.reciprocal(se.stddev[:, 0, :]).contiguous() + degree_gain = ( + se.adam_degree_gain_raw.to(torch.float32).contiguous() + if se.adam_degree_gain_raw is not None + else compress_data.new_empty(0) + ) descriptor, _rotation, moment = torch.ops.deepmd.dpa1_graph_compress( edge_vec, graph.edge_index.contiguous(), @@ -872,6 +929,7 @@ def dpa1_graph_compress_energy_force( type_embedding.contiguous(), se.mean[:, 0, :].contiguous(), inverse_stddev, + degree_gain, compress_data, gate_table, int(se.type_one_side), @@ -889,6 +947,7 @@ def dpa1_graph_compress_energy_force( float(se.rcut_smth), float(se.env_protection), float(se.nnei), + (int(se.lmax) + 1) ** 2, ) weights = [layer.w.contiguous() for layer in hidden] biases = [ @@ -960,6 +1019,7 @@ def dpa1_graph_compress_energy_force( atype, se.mean[:, 0, :].contiguous(), inverse_stddev, + degree_gain, compress_data, gate_table, int(se.type_one_side), diff --git a/deepmd/kernels/cuda/dpa1/graph_descriptor.py b/deepmd/kernels/cuda/dpa1/graph_descriptor.py index 8d1ff2f486..5be3b62f30 100644 --- a/deepmd/kernels/cuda/dpa1/graph_descriptor.py +++ b/deepmd/kernels/cuda/dpa1/graph_descriptor.py @@ -57,6 +57,10 @@ import torch +from deepmd.dpmodel.descriptor.dpa1 import ( + build_dpa1_degree_weights, + build_dpa1_moment_basis, +) from deepmd.kernels.triton.dpa1.activation import ( ACT_CODES, ) @@ -87,6 +91,7 @@ def _forward_fake( type_embedding: torch.Tensor, davg: torch.Tensor, dstd: torch.Tensor, + degree_gain: torch.Tensor, w1: torch.Tensor, b1: torch.Tensor, idt1: torch.Tensor, @@ -109,6 +114,7 @@ def _forward_fake( rcut_smth: float, protection: float, nnei: float, + basis_dim: int, ) -> tuple[torch.Tensor, ...]: n_edge = edge_vec.shape[0] n_node = atype.shape[0] @@ -131,7 +137,7 @@ def _forward_fake( dtype=torch.float32, device=dev, ), - torch.empty(n_node, 4, ng, dtype=torch.float32, device=dev), + torch.empty(n_node, basis_dim, ng, dtype=torch.float32, device=dev), torch.empty(n_edge, dtype=torch.int32, device=dev), torch.empty(n_pairs, w1.shape[1], dtype=torch.float32, device=dev), torch.empty(n2, e_pad, dtype=torch.float32, device=dev), @@ -165,6 +171,7 @@ def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: type_embedding, davg, dstd, + degree_gain, w1, b1, idt1, @@ -187,6 +194,7 @@ def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: rcut_smth, protection, nnei, + _basis_dim, ) = inputs # gr, edge_order, pair_table, pre2_saved, g_saved; the type embedding is # not saved -- the backward re-reads layer 1 through the folded pair table @@ -200,6 +208,7 @@ def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: atype, davg, dstd, + degree_gain, w1, b1, idt1, @@ -244,6 +253,7 @@ def _backward( atype, davg, dstd, + degree_gain, w1, b1, idt1, @@ -279,6 +289,7 @@ def _backward( atype, davg, dstd, + degree_gain, w1, b1, idt1, @@ -300,7 +311,7 @@ def _backward( protection, nnei, ) - return (d_edge_vec,) + (None,) * 28 + return (d_edge_vec,) + (None,) * 30 # ====================================================================== @@ -386,6 +397,7 @@ def _cpu_embedding( rcut: float, rcut_smth: float, protection: float, + basis_dim: int, ) -> tuple[torch.Tensor, ...]: """Environment matrix, embedding MLP and (strip) type-pair gate (fp32). @@ -416,16 +428,29 @@ def _cpu_embedding( if smooth: gate = gate * sw gg = g * (1.0 + gate) - return rr, pre2, pre3, g, gg + moment_basis = rr + if basis_dim > 4: + lmax = {9: 2, 16: 3, 25: 4}[basis_dim] + moment_basis = build_dpa1_moment_basis( + rr, + ev, + sw, + dstd[center_type, 0:1], + edge_mask, + lmax, + protection, + ) + return moment_basis, pre2, pre3, g, gg def _cpu_outputs( - rr: torch.Tensor, + moment_basis: torch.Tensor, gg: torch.Tensor, edge_index: torch.Tensor, edge_mask: torch.Tensor, atype: torch.Tensor, type_embedding: torch.Tensor | None, + degree_gain: torch.Tensor, ng: int, axis: int, concat_tebd: int, @@ -433,12 +458,23 @@ def _cpu_outputs( n_node: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ggm = gg * edge_mask[:, None].to(gg.dtype) - outer = rr[:, :, None] * ggm[:, None, :] # (E, 4, ng) - gr = torch.zeros(n_node, 4, ng, dtype=gg.dtype, device=gg.device) + outer = moment_basis[:, :, None] * ggm[:, None, :] + gr = torch.zeros( + n_node, + moment_basis.shape[-1], + ng, + dtype=gg.dtype, + device=gg.device, + ) gr.index_add_(0, edge_index[1], outer) gr = gr / nnei - gr_t = gr.permute(0, 2, 1) # (N, ng, 4) - grrg = torch.matmul(gr_t, gr[:, :, :axis]).reshape(n_node, ng * axis) + gr_t = gr.permute(0, 2, 1) + gr_axis = gr[:, :, :axis] + if moment_basis.shape[-1] > 4: + lmax = {9: 2, 16: 3, 25: 4}[moment_basis.shape[-1]] + degree_weights = build_dpa1_degree_weights(degree_gain, lmax, gr) + gr_axis = gr_axis * degree_weights.view(1, -1, 1) + grrg = torch.matmul(gr_t, gr_axis).reshape(n_node, ng * axis) rot_mat = gr_t[:, :, 1:4].contiguous() if concat_tebd: grrg = torch.cat([grrg, type_embedding[atype]], dim=-1) @@ -453,6 +489,7 @@ def _cpu_forward( type_embedding: torch.Tensor, davg: torch.Tensor, dstd: torch.Tensor, + degree_gain: torch.Tensor, w1: torch.Tensor, b1: torch.Tensor, idt1: torch.Tensor, @@ -475,6 +512,7 @@ def _cpu_forward( rcut_smth: float, protection: float, nnei: float, + basis_dim: int, ) -> tuple[torch.Tensor, ...]: n_edge = edge_vec.shape[0] n_node = atype.shape[0] @@ -482,7 +520,7 @@ def _cpu_forward( ntypes = type_embedding.shape[0] strip = gate_table.shape[0] > 0 pair_table = _cpu_pair_table(type_embedding, w1, b1, type_one_side, strip) - rr, pre2, pre3, g, gg = _cpu_embedding( + moment_basis, pre2, pre3, g, gg = _cpu_embedding( edge_vec, edge_index, edge_mask, @@ -508,14 +546,16 @@ def _cpu_forward( rcut, rcut_smth, protection, + basis_dim, ) grrg, rot_mat, gr = _cpu_outputs( - rr, + moment_basis, gg, edge_index, edge_mask, atype, type_embedding, + degree_gain, ng, axis, concat_tebd, @@ -548,6 +588,7 @@ def _cpu_backward( atype: torch.Tensor, davg: torch.Tensor, dstd: torch.Tensor, + degree_gain: torch.Tensor, w1: torch.Tensor, b1: torch.Tensor, idt1: torch.Tensor, @@ -571,12 +612,13 @@ def _cpu_backward( ) -> torch.Tensor: n_node = atype.shape[0] ng = w3.shape[1] + basis_dim = gr.shape[1] strip = gate_table.shape[0] > 0 n_pairs = gate_table.shape[0] if strip else pair_table.shape[0] ntypes = n_pairs if type_one_side else round(n_pairs**0.5) ev = edge_vec.detach().clone().requires_grad_(True) with torch.enable_grad(): - rr, _pre2, _pre3, _g, gg = _cpu_embedding( + moment_basis, _pre2, _pre3, _g, gg = _cpu_embedding( ev, edge_index, edge_mask, @@ -602,14 +644,16 @@ def _cpu_backward( rcut, rcut_smth, protection, + basis_dim, ) grrg, rot_mat, _gr = _cpu_outputs( - rr, + moment_basis, gg, edge_index, edge_mask, atype, None, + degree_gain, ng, axis, 0, @@ -724,6 +768,11 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: else: gate_table = empty.reshape(0, 0) smooth = 0 + degree_gain = ( + se.adam_degree_gain_raw.to(torch.float32).contiguous() + if se.adam_degree_gain_raw is not None + else empty + ) w1, w2, w3 = (layer.w.contiguous() for layer in layers) grrg, rot_mat, *_aux = torch.ops.deepmd.dpa1_graph_descriptor( graph.edge_vec.contiguous(), @@ -734,6 +783,7 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: # mean / stddev are slot-independent; slot 0 is the canonical (T, 4). se.mean[:, 0, :].contiguous(), se.stddev[:, 0, :].contiguous(), + degree_gain, w1, optional(layers[0].b), optional(layers[0].idt), @@ -756,5 +806,6 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: float(se.rcut_smth), float(se.env_protection), float(se.nnei), + (int(se.lmax) + 1) ** 2, ) return grrg, rot_mat diff --git a/deepmd/kernels/cuda/dpa1/graph_energy_force.py b/deepmd/kernels/cuda/dpa1/graph_energy_force.py index 2cc262290e..a257a96c92 100644 --- a/deepmd/kernels/cuda/dpa1/graph_energy_force.py +++ b/deepmd/kernels/cuda/dpa1/graph_energy_force.py @@ -69,6 +69,7 @@ def _fake( type_embedding: torch.Tensor, davg: torch.Tensor, dstd: torch.Tensor, + degree_gain: torch.Tensor, w1: torch.Tensor, b1: torch.Tensor, idt1: torch.Tensor, @@ -90,6 +91,7 @@ def _fake( rcut_smth: float, protection: float, nnei: float, + basis_dim: int, fit_ws: list[torch.Tensor], fit_bs: list[torch.Tensor], fit_idts: list[torch.Tensor], @@ -130,6 +132,7 @@ def _cpu( type_embedding: torch.Tensor, davg: torch.Tensor, dstd: torch.Tensor, + degree_gain: torch.Tensor, w1: torch.Tensor, b1: torch.Tensor, idt1: torch.Tensor, @@ -151,6 +154,7 @@ def _cpu( rcut_smth: float, protection: float, nnei: float, + basis_dim: int, fit_ws: list[torch.Tensor], fit_bs: list[torch.Tensor], fit_idts: list[torch.Tensor], @@ -174,6 +178,7 @@ def _cpu( type_embedding, davg, dstd, + degree_gain, w1, b1, idt1, @@ -196,6 +201,7 @@ def _cpu( rcut_smth, protection, nnei, + basis_dim, ) ) atom_e_raw, fit_saved = torch.ops.deepmd.graph_fitting( @@ -237,6 +243,7 @@ def _cpu( atype, davg, dstd, + degree_gain, w1, b1, idt1, @@ -389,6 +396,11 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: type_embedding.contiguous(), se.mean[:, 0, :].contiguous(), se.stddev[:, 0, :].contiguous(), + ( + se.adam_degree_gain_raw.to(torch.float32).contiguous() + if se.adam_degree_gain_raw is not None + else empty + ), w1, optional(layers[0].b), optional(layers[0].idt), @@ -410,6 +422,7 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: float(se.rcut_smth), float(se.env_protection), float(se.nnei), + (int(se.lmax) + 1) ** 2, [layer.w.contiguous() for layer in hidden], [layer.b.contiguous() if layer.b is not None else fempty for layer in hidden], [ diff --git a/deepmd/kernels/triton/dpa1/se_conv.py b/deepmd/kernels/triton/dpa1/se_conv.py index 96390baf86..052ce496ff 100644 --- a/deepmd/kernels/triton/dpa1/se_conv.py +++ b/deepmd/kernels/triton/dpa1/se_conv.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later # pyright: reportMissingImports=false -# ruff: noqa: ANN001, ANN202 +# ruff: noqa: ANN001, ANN202, RUF005 """Fused environment convolution for the DPA1 (``se_atten``) descriptor. For the attention-free (``attn_layer == 0``), strip-embedding ``se_atten`` path @@ -8,17 +8,17 @@ ``h2 = act(z2) * idt + resnet(h1)`` (last embedding layer + resnet) ``gg = h2 * (1 + tt[idx] * sw)`` (type-pair gate + smooth cutoff) - ``xyz[k, c] = sum_j rr[j, k] * gg[j, c]`` (moment accumulation) + ``moment[k, c] = sum_j basis[j, k] * gg[j, c]`` (moment accumulation) where ``act`` is the last layer's activation (``tanh`` or ``silu``), ``z2`` is the last embedding pre-activation ``h1 @ W2 + b2``, ``h1`` is the penultimate activation feeding the resnet, ``idt`` is the last layer's per- channel timestep (all ones when ``resnet_dt`` is off), ``tt`` is the type-pair embedding table with per-edge row index ``idx``, ``sw`` is the smooth radial -cutoff, and ``rr`` is the ``(s, s*x/r, s*y/r, s*z/r)`` environment matrix. The +cutoff, and ``basis`` is the 4/9/16/25-row angular moment basis. The ``1 / nnei`` normalization of the moment is applied by the descriptor after this operator (matching the eager and tabulated paths), so the operator returns the -unnormalized moment ``rr^T @ gg``. +unnormalized moment ``basis^T @ gg``. The last-layer resnet takes one of three shapes, selected by ``resnet_mult`` (``= ng // H1`` when the layer adds a residual, else ``0``): @@ -32,15 +32,15 @@ covers both; the no-residual case skips the ``h1`` read. The eager path materializes three ``(E, ng)`` tensors (``h2``, the gathered type -feature, and ``gg``) and then runs a batched ``rr^T @ gg`` matmul whose ``M = 4`` -contraction uses cuBLAS poorly. This operator fuses the whole tail into one +feature, and ``gg``) and then runs a batched ``basis^T @ gg`` matmul whose small +geometric contraction uses cuBLAS poorly. This operator fuses the whole tail into one node-parallel kernel: each program owns one node, streams its neighbors, forms ``gg`` in registers (never materializing any ``(E, ng)`` tensor, and gathering -the type feature inline) and accumulates the four moment rows. The two embedding +the type feature inline) and accumulates all moment rows. The two embedding GEMMs (``h0 @ W1`` and ``h1 @ W2``) stay on cuBLAS -- in the fp32 regime this descriptor runs in, Triton ``tl.dot`` has no tensor-core path and cannot beat cuBLAS, so only the memory-bound, non-GEMM tail is fused. The trailing -``ng x ng x 4`` Gram contraction that forms the final descriptor is likewise +small-axis Gram contraction that forms the final descriptor is likewise left on cuBLAS. Design notes and pitfalls @@ -68,7 +68,7 @@ :func:`.tile_configs.resolve_conv_config`, whose default is deliberately small. - **Inference-only autograd.** The registered backward returns gradients for - the coordinate-bearing inputs (``z2``, ``h1``, ``sw``, ``rr``) that carry the + the coordinate-bearing inputs (``z2``, ``h1``, ``sw``, ``basis``) that carry the force; the timestep, type table and index do not depend on coordinates. Training keeps the dense reference path (the gate is inference-only), so their gradients are never required here. @@ -128,7 +128,7 @@ def _se_conv_reference( tt: Tensor, idx: Tensor, sw: Tensor, - rr: Tensor, + basis: Tensor, resnet_mult: int, act: int, gated: int, @@ -151,24 +151,24 @@ def _se_conv_reference( gg = h2 * (1.0 + gg_t * sw.unsqueeze(-1)) else: gg = h2 - return torch.matmul(rr.transpose(1, 2), gg) + return torch.matmul(basis.transpose(1, 2), gg) def _se_conv_reference_backward( - grad_xyz: Tensor, + grad_moment: Tensor, z2: Tensor, h1: Tensor, idt: Tensor, tt: Tensor, idx: Tensor, sw: Tensor, - rr: Tensor, + basis: Tensor, resnet_mult: int, act: int, gated: int, ) -> tuple[Tensor, Tensor, Tensor, Tensor]: """Closed-form gradient of :func:`_se_conv_reference` w.r.t. the coordinate- - bearing inputs ``(z2, h1, sw, rr)``. + bearing inputs ``(z2, h1, sw, basis)``. A closed form (rather than a nested ``torch.autograd.grad``) is used so the fallback composes under ``make_fx`` / ``torch.export``: the tracer runs this @@ -195,9 +195,9 @@ def _se_conv_reference_backward( else: fac = torch.ones_like(h2) gg = h2 * fac - # d(xyz = rr^T @ gg): grad_gg = rr @ grad_xyz, grad_rr = gg @ grad_xyz^T. - grad_gg = torch.matmul(rr, grad_xyz) - grad_rr = torch.matmul(gg, grad_xyz.transpose(1, 2)) + # d(moment = basis^T @ gg): contract each factor with the upstream tensor. + grad_gg = torch.matmul(basis, grad_moment) + grad_basis = torch.matmul(gg, grad_moment.transpose(1, 2)) grad_h2 = grad_gg * fac grad_z2 = grad_h2 * idt * act_grad # fac = 1 + gg_t * sw, so d gg / d sw = h2 * gg_t (zero without the gate). @@ -212,7 +212,7 @@ def _se_conv_reference_backward( grad_h1 = grad_h2 else: grad_h1 = torch.zeros_like(h1) - return grad_z2, grad_h1, grad_sw, grad_rr + return grad_z2, grad_h1, grad_sw, grad_basis # ====================================================================== @@ -235,18 +235,19 @@ def _se_conv_fwd_kernel( tt_ptr, # (P, NG) type-pair embedding table idx_ptr, # (N * NNEI,) per-edge type-pair row index sw_ptr, # (N, NNEI) smooth radial cutoff - rr_ptr, # (N, NNEI, 4) environment matrix - out_ptr, # (N, 4, NG) accumulated moments + basis_ptr, # (N, NNEI, BASIS_DIM) angular moment basis + out_ptr, # (N, BASIS_DIM, NG) accumulated moments NNEI, H1: tl.constexpr, NG: tl.constexpr, NGP: tl.constexpr, + BASIS_DIM: tl.constexpr, RESNET_MULT: tl.constexpr, ACT: tl.constexpr, GATED: tl.constexpr, BN: tl.constexpr, ): - """Accumulate the four unnormalized moment rows over a node's neighbors. + """Accumulate the unnormalized moment rows over a node's neighbors. ``gg`` is formed in registers per neighbor block and never written to global memory; the type feature is gathered inline through ``idx``. The @@ -272,21 +273,25 @@ def _se_conv_fwd_kernel( # Accumulate in the input precision so the kernel serves both fp32 # (the eager DPA1 path) and fp64 (the float64 pt_expt / export path). acc_ty = z2_ptr.dtype.element_ty - acc0 = tl.zeros((NGP,), dtype=acc_ty) - acc1 = tl.zeros((NGP,), dtype=acc_ty) - acc2 = tl.zeros((NGP,), dtype=acc_ty) - acc3 = tl.zeros((NGP,), dtype=acc_ty) + accumulators = () + for _ in tl.static_range(BASIS_DIM): + accumulators = accumulators + (tl.zeros((NGP,), dtype=acc_ty),) for n0 in range(0, NNEI, BN): offs = n0 + tl.arange(0, BN) nmask = offs < NNEI m = (nmask[:, None] & cm[None, :]) if NGP != NG else nmask[:, None] e = node * NNEI + offs z2 = tl.load(z2_ptr + e[:, None] * NG + rc[None, :], mask=m, other=0.0) - base = e * 4 - s = tl.load(rr_ptr + base + 0, mask=nmask, other=0.0) - rx = tl.load(rr_ptr + base + 1, mask=nmask, other=0.0) - ry = tl.load(rr_ptr + base + 2, mask=nmask, other=0.0) - rz = tl.load(rr_ptr + base + 3, mask=nmask, other=0.0) + basis_base = e * BASIS_DIM + basis_rows = () + for row in tl.static_range(BASIS_DIM): + basis_rows = basis_rows + ( + tl.load( + basis_ptr + basis_base + row, + mask=nmask, + other=0.0, + ), + ) h2 = activation(z2, ACT) * idt[None, :] if RESNET_MULT > 0: # concat[h1, h1] (doubling) and identity share one addressing @@ -304,35 +309,40 @@ def _se_conv_fwd_kernel( gg = h2 * (1.0 + ggt * sw[:, None]) else: gg = h2 - acc0 += tl.sum(s[:, None] * gg, axis=0) - acc1 += tl.sum(rx[:, None] * gg, axis=0) - acc2 += tl.sum(ry[:, None] * gg, axis=0) - acc3 += tl.sum(rz[:, None] * gg, axis=0) - ob = node * 4 * NG - tl.store(out_ptr + ob + 0 * NG + rc, acc0, mask=cm) - tl.store(out_ptr + ob + 1 * NG + rc, acc1, mask=cm) - tl.store(out_ptr + ob + 2 * NG + rc, acc2, mask=cm) - tl.store(out_ptr + ob + 3 * NG + rc, acc3, mask=cm) + updated_accumulators = () + for row in tl.static_range(BASIS_DIM): + updated_accumulators = updated_accumulators + ( + accumulators[row] + tl.sum(basis_rows[row][:, None] * gg, axis=0), + ) + accumulators = updated_accumulators + output_base = node * BASIS_DIM * NG + for row in tl.static_range(BASIS_DIM): + tl.store( + out_ptr + output_base + row * NG + rc, + accumulators[row], + mask=cm, + ) @triton.jit def _se_conv_bwd_kernel( - gout_ptr, # (N, 4, NG) upstream gradient of the moments + gout_ptr, # (N, BASIS_DIM, NG) upstream gradient of the moments z2_ptr, h1_ptr, idt_ptr, tt_ptr, idx_ptr, sw_ptr, - rr_ptr, + basis_ptr, dz2_ptr, # (N, NNEI, NG) dh1_ptr, # (N, NNEI, H1); written only when RESNET_MULT > 0 dsw_ptr, # (N, NNEI) - drr_ptr, # (N, NNEI, 4) + dbasis_ptr, # (N, NNEI, BASIS_DIM) NNEI, H1: tl.constexpr, NG: tl.constexpr, NGP: tl.constexpr, H1P: tl.constexpr, + BASIS_DIM: tl.constexpr, RESNET_MULT: tl.constexpr, ACT: tl.constexpr, GATED: tl.constexpr, @@ -355,24 +365,38 @@ def _se_conv_bwd_kernel( rc = tl.arange(0, NGP) cm = rc < NG idt = tl.load(idt_ptr + rc, mask=cm, other=0.0) - gb = node * (4 * NG) - g0 = tl.load(gout_ptr + gb + 0 * NG + rc, mask=cm, other=0.0) - g1 = tl.load(gout_ptr + gb + 1 * NG + rc, mask=cm, other=0.0) - g2 = tl.load(gout_ptr + gb + 2 * NG + rc, mask=cm, other=0.0) - g3 = tl.load(gout_ptr + gb + 3 * NG + rc, mask=cm, other=0.0) + gb = node * BASIS_DIM * NG + moment_gradients = () + for row in tl.static_range(BASIS_DIM): + moment_gradients = moment_gradients + ( + tl.load( + gout_ptr + gb + row * NG + rc, + mask=cm, + other=0.0, + ), + ) # Doubling with padded H1 needs the moment gradient at both residual # halves; hoist the per-node loads at columns ``rj`` and ``rj + H1``. if RESNET_MULT == 2 and H1P != H1: rj = tl.arange(0, H1P) hm = rj < H1 - g0lo = tl.load(gout_ptr + gb + 0 * NG + rj, mask=hm, other=0.0) - g1lo = tl.load(gout_ptr + gb + 1 * NG + rj, mask=hm, other=0.0) - g2lo = tl.load(gout_ptr + gb + 2 * NG + rj, mask=hm, other=0.0) - g3lo = tl.load(gout_ptr + gb + 3 * NG + rj, mask=hm, other=0.0) - g0hi = tl.load(gout_ptr + gb + 0 * NG + rj + H1, mask=hm, other=0.0) - g1hi = tl.load(gout_ptr + gb + 1 * NG + rj + H1, mask=hm, other=0.0) - g2hi = tl.load(gout_ptr + gb + 2 * NG + rj + H1, mask=hm, other=0.0) - g3hi = tl.load(gout_ptr + gb + 3 * NG + rj + H1, mask=hm, other=0.0) + moment_gradients_lo = () + moment_gradients_hi = () + for row in tl.static_range(BASIS_DIM): + moment_gradients_lo = moment_gradients_lo + ( + tl.load( + gout_ptr + gb + row * NG + rj, + mask=hm, + other=0.0, + ), + ) + moment_gradients_hi = moment_gradients_hi + ( + tl.load( + gout_ptr + gb + row * NG + rj + H1, + mask=hm, + other=0.0, + ), + ) for n0 in range(0, NNEI, BN): offs = n0 + tl.arange(0, BN) nmask = offs < NNEI @@ -380,11 +404,16 @@ def _se_conv_bwd_kernel( e = node * NNEI + offs ec = e[:, None] * NG + rc[None, :] z2 = tl.load(z2_ptr + ec, mask=m, other=0.0) - base = e * 4 - s = tl.load(rr_ptr + base + 0, mask=nmask, other=0.0) - rx = tl.load(rr_ptr + base + 1, mask=nmask, other=0.0) - ry = tl.load(rr_ptr + base + 2, mask=nmask, other=0.0) - rz = tl.load(rr_ptr + base + 3, mask=nmask, other=0.0) + basis_base = e * BASIS_DIM + basis_rows = () + for row in tl.static_range(BASIS_DIM): + basis_rows = basis_rows + ( + tl.load( + basis_ptr + basis_base + row, + mask=nmask, + other=0.0, + ), + ) a, ad = activation_grad(z2, ACT) h2 = a * idt[None, :] if RESNET_MULT > 0: @@ -402,12 +431,9 @@ def _se_conv_bwd_kernel( else: fac = 1.0 gg = h2 * fac - dgg = ( - s[:, None] * g0[None, :] - + rx[:, None] * g1[None, :] - + ry[:, None] * g2[None, :] - + rz[:, None] * g3[None, :] - ) + dgg = tl.zeros((BN, NGP), dtype=z2_ptr.dtype.element_ty) + for row in tl.static_range(BASIS_DIM): + dgg += basis_rows[row][:, None] * moment_gradients[row][None, :] grad_h2 = dgg * fac tl.store( dz2_ptr + ec, @@ -427,18 +453,15 @@ def _se_conv_bwd_kernel( elif RESNET_MULT == 2: # Padded doubling: recompute the two residual halves directly. hmask = nmask[:, None] & (rj < H1)[None, :] - dgg_lo = ( - s[:, None] * g0lo[None, :] - + rx[:, None] * g1lo[None, :] - + ry[:, None] * g2lo[None, :] - + rz[:, None] * g3lo[None, :] - ) - dgg_hi = ( - s[:, None] * g0hi[None, :] - + rx[:, None] * g1hi[None, :] - + ry[:, None] * g2hi[None, :] - + rz[:, None] * g3hi[None, :] - ) + dgg_lo = tl.zeros((BN, H1P), dtype=z2_ptr.dtype.element_ty) + dgg_hi = tl.zeros((BN, H1P), dtype=z2_ptr.dtype.element_ty) + for row in tl.static_range(BASIS_DIM): + dgg_lo += ( + basis_rows[row][:, None] * moment_gradients_lo[row][None, :] + ) + dgg_hi += ( + basis_rows[row][:, None] * moment_gradients_hi[row][None, :] + ) if GATED: ggt_lo = tl.load( tt_ptr + idx[:, None] * NG + rj[None, :], mask=hmask, other=0.0 @@ -457,10 +480,12 @@ def _se_conv_bwd_kernel( if GATED: # sw enters only through the gate; concat has no sw gradient. tl.store(dsw_ptr + e, tl.sum(dgg * h2 * ggt, axis=1), mask=nmask) - tl.store(drr_ptr + base + 0, tl.sum(gg * g0[None, :], axis=1), mask=nmask) - tl.store(drr_ptr + base + 1, tl.sum(gg * g1[None, :], axis=1), mask=nmask) - tl.store(drr_ptr + base + 2, tl.sum(gg * g2[None, :], axis=1), mask=nmask) - tl.store(drr_ptr + base + 3, tl.sum(gg * g3[None, :], axis=1), mask=nmask) + for row in tl.static_range(BASIS_DIM): + tl.store( + dbasis_ptr + basis_base + row, + tl.sum(gg * moment_gradients[row][None, :], axis=1), + mask=nmask, + ) # ====================================================================== @@ -481,7 +506,7 @@ def _se_conv_fwd_impl( tt: Tensor, idx: Tensor, sw: Tensor, - rr: Tensor, + basis: Tensor, resnet_mult: int, act: int, gated: int, @@ -489,9 +514,12 @@ def _se_conv_fwd_impl( num_warps: int, ) -> Tensor: if not _use_triton(z2): - return _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated) + return _se_conv_reference( + z2, h1, idt, tt, idx, sw, basis, resnet_mult, act, gated + ) nfnl, nnei, ng = z2.shape - out = torch.empty((nfnl, 4, ng), dtype=z2.dtype, device=z2.device) + basis_dim = basis.shape[-1] + out = torch.empty((nfnl, basis_dim, ng), dtype=z2.dtype, device=z2.device) wrap_triton(_se_conv_fwd_kernel)[(nfnl,)]( z2, h1, @@ -499,12 +527,13 @@ def _se_conv_fwd_impl( tt, idx, sw, - rr, + basis, out, nnei, H1=h1.shape[-1], NG=ng, NGP=triton.next_power_of_2(ng), + BASIS_DIM=basis_dim, RESNET_MULT=resnet_mult, ACT=act, GATED=gated, @@ -515,14 +544,14 @@ def _se_conv_fwd_impl( def _se_conv_bwd_impl( - grad_xyz: Tensor, + grad_moment: Tensor, z2: Tensor, h1: Tensor, idt: Tensor, tt: Tensor, idx: Tensor, sw: Tensor, - rr: Tensor, + basis: Tensor, resnet_mult: int, act: int, gated: int, @@ -531,7 +560,17 @@ def _se_conv_bwd_impl( ) -> tuple[Tensor, Tensor, Tensor, Tensor]: if not _use_triton(z2): return _se_conv_reference_backward( - grad_xyz, z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated + grad_moment, + z2, + h1, + idt, + tt, + idx, + sw, + basis, + resnet_mult, + act, + gated, ) nfnl, nnei, ng = z2.shape dz2 = torch.empty_like(z2) @@ -541,32 +580,33 @@ def _se_conv_bwd_impl( # Concat (gated == 0) uses no sw gate; the kernel skips the dsw store, so # pre-zero it here. dsw = torch.empty_like(sw) if gated else torch.zeros_like(sw) - drr = torch.empty_like(rr) + dbasis = torch.empty_like(basis) wrap_triton(_se_conv_bwd_kernel)[(nfnl,)]( - grad_xyz.contiguous(), + grad_moment.contiguous(), z2, h1, idt, tt, idx, sw, - rr, + basis, dz2, dh1, dsw, - drr, + dbasis, nnei, H1=h1.shape[-1], NG=ng, NGP=triton.next_power_of_2(ng), H1P=triton.next_power_of_2(h1.shape[-1]), + BASIS_DIM=basis.shape[-1], RESNET_MULT=resnet_mult, ACT=act, GATED=gated, BN=block_n, num_warps=num_warps, ) - return dz2, dh1, dsw, drr + return dz2, dh1, dsw, dbasis _se_conv_op = triton_op("dpa1_triton::se_conv", mutates_args=())(_se_conv_fwd_impl) @@ -576,25 +616,39 @@ def _se_conv_bwd_impl( @_se_conv_op.register_fake -def _(z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps): - return z2.new_empty((z2.shape[0], 4, z2.shape[2])) +def _(z2, h1, idt, tt, idx, sw, basis, resnet_mult, act, gated, block_n, num_warps): + return z2.new_empty((z2.shape[0], basis.shape[2], z2.shape[2])) @_se_conv_bwd_op.register_fake def _( - grad_xyz, z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps + grad_moment, + z2, + h1, + idt, + tt, + idx, + sw, + basis, + resnet_mult, + act, + gated, + block_n, + num_warps, ): return ( torch.empty_like(z2), torch.empty_like(h1), torch.empty_like(sw), - torch.empty_like(rr), + torch.empty_like(basis), ) def _se_conv_setup_context(ctx, inputs, output): - z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps = inputs - ctx.save_for_backward(z2, h1, idt, tt, idx, sw, rr) + z2, h1, idt, tt, idx, sw, basis, resnet_mult, act, gated, block_n, num_warps = ( + inputs + ) + ctx.save_for_backward(z2, h1, idt, tt, idx, sw, basis) ctx.resnet_mult = resnet_mult ctx.act = act ctx.gated = gated @@ -602,17 +656,17 @@ def _se_conv_setup_context(ctx, inputs, output): ctx.num_warps = num_warps -def _se_conv_backward(ctx, grad_xyz): - z2, h1, idt, tt, idx, sw, rr = ctx.saved_tensors - grad_z2, grad_h1, grad_sw, grad_rr = _se_conv_bwd_op( - grad_xyz.contiguous(), +def _se_conv_backward(ctx, grad_moment): + z2, h1, idt, tt, idx, sw, basis = ctx.saved_tensors + grad_z2, grad_h1, grad_sw, grad_basis = _se_conv_bwd_op( + grad_moment.contiguous(), z2, h1, idt, tt, idx, sw, - rr, + basis, ctx.resnet_mult, ctx.act, ctx.gated, @@ -627,7 +681,7 @@ def _se_conv_backward(ctx, grad_xyz): None, None, grad_sw, - grad_rr, + grad_basis, None, None, None, @@ -673,7 +727,7 @@ def se_conv( tt: Tensor, idx: Tensor, sw: Tensor, - rr: Tensor, + basis: Tensor, resnet_mult: int, act: int, gated: int, @@ -697,8 +751,9 @@ def se_conv( Per-edge row index into ``tt`` with shape (N * nnei,), dtype int64. sw : Tensor Smooth radial cutoff with shape (N, nnei). - rr : Tensor - Environment matrix ``(s, s*x/r, s*y/r, s*z/r)`` with shape (N, nnei, 4). + basis : Tensor + Angular moment basis with shape (N, nnei, basis_dim), where + ``basis_dim`` is ``(lmax + 1) ** 2`` for ``lmax`` from 1 through 4. resnet_mult : int Residual structure of the last layer: ``2`` (width doubling), ``1`` (identity), or ``0`` (no residual). @@ -714,8 +769,8 @@ def se_conv( Returns ------- Tensor - Unnormalized moments ``xyz`` with shape (N, 4, ng), equal to - ``rr^T @ gg``. The ``1 / nnei`` normalization is applied by the + Unnormalized moments with shape (N, basis_dim, ng), equal to + ``basis^T @ gg``. The ``1 / nnei`` normalization is applied by the descriptor after this operator. Notes @@ -726,10 +781,24 @@ def se_conv( gradient through the registered backward. """ block_n, num_warps = resolve_conv_config( - int(z2.shape[-1]), int(h1.shape[-1]), triton_infer_level() + int(z2.shape[-1]), + int(h1.shape[-1]), + int(basis.shape[-1]), + triton_infer_level(), ) return _se_conv_op( - z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps + z2, + h1, + idt, + tt, + idx, + sw, + basis, + resnet_mult, + act, + gated, + block_n, + num_warps, ) @@ -739,7 +808,7 @@ def se_atten_conv( tt: Tensor | None, idx: Tensor | None, sw: Tensor | None, - rr: Tensor, + basis: Tensor, gated: int, ) -> Tensor: """Fuse the embedding net's final layer, type gate and moment accumulation. @@ -771,15 +840,15 @@ def se_atten_conv( concat. sw : Tensor or None Smooth radial cutoff with shape (N, nnei); ``None`` for concat. - rr : Tensor - Environment matrix with shape (N, nnei, 4). + basis : Tensor + Angular moment basis with shape (N, nnei, basis_dim). gated : int ``1`` (strip) applies the type-pair gate; ``0`` (concat) skips it. Returns ------- Tensor - Unnormalized moments ``xyz`` with shape (N, 4, ng). + Unnormalized moments with shape (N, basis_dim, ng). """ *head, last = embedding_net.layers h = ss @@ -798,7 +867,16 @@ def se_atten_conv( if not gated: tt, idx, sw = concat_gate_placeholders(z2, ng) return se_conv( - z2.contiguous(), h.contiguous(), idt, tt, idx, sw, rr, resnet_mult, act, gated + z2.contiguous(), + h.contiguous(), + idt, + tt, + idx, + sw, + basis, + resnet_mult, + act, + gated, ) @@ -808,7 +886,7 @@ def se_atten_conv( def _autotune_conv(model: torch.nn.Module, level: int, device: torch.device) -> None: """Sweep the fused-convolution launch table for a model about to be frozen. - Collects the ``(ng, H1)`` shape key of every eligible ``se_atten`` + Collects the ``(ng, H1, basis_dim)`` shape key of every eligible ``se_atten`` descriptor in ``model`` and sweeps the keys the built-in / freeze-time tables do not yet cover on the target GPU, registering the winners so the ``resolve_conv_config`` lookups made while tracing bake tuned launches into @@ -819,22 +897,28 @@ def _autotune_conv(model: torch.nn.Module, level: int, device: torch.device) -> ) device_name = torch.cuda.get_device_name(device) - keys: set[tuple[int, int]] = set() + keys: set[tuple[int, int, int]] = set() for module in model.modules(): eligible = getattr(module, "_fused_eligible", None) if not (callable(eligible) and eligible("triton")): continue weight = module.se_atten.embeddings[0].layers[-1].w - keys.add((int(weight.shape[1]), int(weight.shape[0]))) - tuned: dict[tuple[int, int], tuple[int, int]] = {} - for ng, h1 in sorted(keys): + keys.add( + ( + int(weight.shape[1]), + int(weight.shape[0]), + (int(getattr(module.se_atten, "lmax", 1)) + 1) ** 2, + ) + ) + tuned: dict[tuple[int, int, int], tuple[int, int]] = {} + for ng, h1, basis_dim in sorted(keys): # The sweep needs a residual last layer (ng in {H1, 2*H1}); other shapes # keep the default. Skip keys the tables already cover. - if has_conv_config(ng, h1) or ng not in (h1, 2 * h1): + if has_conv_config(ng, h1, basis_dim) or ng not in (h1, 2 * h1): continue - config = sweep(ng, h1, device=device) - register_conv_config(device_name, ng, h1, config) - tuned[(ng, h1)] = config + config = sweep(ng, h1, basis_dim=basis_dim, device=device) + register_conv_config(device_name, ng, h1, basis_dim, config) + tuned[(ng, h1, basis_dim)] = config if tuned: log.info("DPA1 se_conv: tuned launch configs %s on %s.", tuned, device_name) else: diff --git a/deepmd/kernels/triton/dpa1/sweep_tile_configs.py b/deepmd/kernels/triton/dpa1/sweep_tile_configs.py index 51e81bb2c9..03c0b3d2da 100644 --- a/deepmd/kernels/triton/dpa1/sweep_tile_configs.py +++ b/deepmd/kernels/triton/dpa1/sweep_tile_configs.py @@ -2,8 +2,8 @@ # ruff: noqa: T201 r"""Sweep the launch configuration of a DPA1 fused environment convolution. -Both fused kernels have a two-parameter launch configuration resolved per -``(ng, H1)`` by :mod:`.tile_configs`: ``se_conv`` (node-parallel) is keyed by +Both fused kernels have a two-parameter launch configuration resolved by +:mod:`.tile_configs`: ``se_conv`` (node-parallel) is keyed by the per-neighbor block width ``BLOCK_N`` and a warp count; ``edge_conv`` (edge-parallel) by the per-block edge count ``BLOCK_E`` and a warp count. This module measures the candidate configurations for one channel width on synthetic @@ -20,16 +20,17 @@ :: python -m deepmd.kernels.triton.dpa1.sweep_tile_configs \\ - --kind {conv,edge} --ng NG --h1 H1 [--device cuda:0] + --kind {conv,edge} --ng NG --h1 H1 [--basis-dim {4,9,16,25}] [--device cuda:0] -The printed ``(ng, h1): (BLOCK, num_warps)`` line is appended, under the device -name from ``torch.cuda.get_device_name``, to the relevant built-in table in +The printed key is ``(ng, h1, basis_dim)`` for ``se_conv`` and ``(ng, h1)`` for +``edge_conv``. Append it under the device name from +``torch.cuda.get_device_name`` to the relevant built-in table in :mod:`.tile_configs` (``_CONV_BUILTIN`` / ``_EDGE_BUILTIN``). Regeneration note ----------------- Any change to a kernel body invalidates its existing table entries; rerun the -sweep for every covered ``(ng, H1)`` and refresh the table. +sweep for every covered key and refresh the table. """ from __future__ import ( @@ -81,7 +82,12 @@ def _make_inputs( - nodes: int, nnei: int, ng: int, h1: int, device: torch.device + nodes: int, + nnei: int, + ng: int, + h1: int, + basis_dim: int, + device: torch.device, ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: p = 4096 z2 = torch.randn(nodes, nnei, ng, dtype=torch.float32, device=device) @@ -90,8 +96,8 @@ def _make_inputs( tt = torch.randn(p, ng, dtype=torch.float32, device=device) * 0.3 idx = torch.randint(0, p, (nodes * nnei,), dtype=torch.int64, device=device) sw = torch.rand(nodes, nnei, dtype=torch.float32, device=device) - rr = torch.randn(nodes, nnei, 4, dtype=torch.float32, device=device) - return z2, h1t, idt, tt, idx, sw, rr + basis = torch.randn(nodes, nnei, basis_dim, dtype=torch.float32, device=device) + return z2, h1t, idt, tt, idx, sw, basis def _bench(fn: Callable[[], object], iters: int = 40, warmup: int = 15) -> float: @@ -110,6 +116,7 @@ def _bench(fn: Callable[[], object], iters: int = 40, warmup: int = 15) -> float def sweep( ng: int, h1: int, + basis_dim: int = 4, nnei: int = 181, nodes: int = 4096, device: torch.device | None = None, @@ -122,6 +129,8 @@ def sweep( Embedding channel width. h1 : int Penultimate embedding width; ``ng == 2 * h1`` is required. + basis_dim : int + Angular moment width. Supported values are 4, 9, 16, and 25. nnei : int Neighbor count used to size the synthetic input. nodes : int @@ -138,34 +147,61 @@ def sweep( raise ValueError( "se_conv sweep requires a residual last layer (ng in {h1, 2*h1})" ) + if basis_dim not in (4, 9, 16, 25): + raise ValueError("se_conv sweep requires basis_dim in {4, 9, 16, 25}") resnet_mult = ng // h1 device = device or torch.device("cuda") torch.backends.cuda.matmul.allow_tf32 = False - # The launch configuration is memory/register bound and keyed by (ng, H1) - # only; the activation adds a few cheap elementwise ops and does not shift + # The launch configuration is memory/register bound and keyed by + # (ng, H1, basis_dim); the activation adds cheap elementwise ops and does not shift # the optimum, so the sweep times the ``tanh`` path (act = 0). act = 0 - # The launch configuration is keyed by (ng, H1) and is independent of the + # The launch configuration is independent of the # tebd-input mode; the strip gate (gated = 1) is the register-heaviest case, # so its optimum is a safe upper bound for concat (gated = 0). gated = 1 - z2, h1t, idt, tt, idx, sw, rr = _make_inputs(nodes, nnei, ng, h1, device) - ref = _se_conv_reference(z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated) + z2, h1t, idt, tt, idx, sw, basis = _make_inputs( + nodes, nnei, ng, h1, basis_dim, device + ) + ref = _se_conv_reference(z2, h1t, idt, tt, idx, sw, basis, resnet_mult, act, gated) gout = torch.randn_like(ref) def fwd_bwd(bn: int, nw: int) -> None: _se_conv_fwd_impl( - z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated, bn, nw + z2, h1t, idt, tt, idx, sw, basis, resnet_mult, act, gated, bn, nw ) _se_conv_bwd_impl( - gout, z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated, bn, nw + gout, + z2, + h1t, + idt, + tt, + idx, + sw, + basis, + resnet_mult, + act, + gated, + bn, + nw, ) results: list[tuple[float, int, int]] = [] for bn, nw in itertools.product(_BLOCK_N_CANDIDATES, _WARP_CANDIDATES): try: out = _se_conv_fwd_impl( - z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated, bn, nw + z2, + h1t, + idt, + tt, + idx, + sw, + basis, + resnet_mult, + act, + gated, + bn, + nw, ) rel = (out - ref).abs().max().item() / ref.abs().max().item() if rel > _REL_TOL: @@ -287,6 +323,7 @@ def main() -> None: parser.add_argument("--kind", choices=("conv", "edge"), default="conv") parser.add_argument("--ng", type=int, required=True) parser.add_argument("--h1", type=int, required=True) + parser.add_argument("--basis-dim", type=int, choices=(4, 9, 16, 25), default=4) parser.add_argument("--nnei", type=int, default=181) parser.add_argument("--nodes", type=int, default=4096) parser.add_argument("--device", type=str, default="cuda:0") @@ -296,9 +333,21 @@ def main() -> None: if args.kind == "edge": block, nw = sweep_edge(args.ng, args.h1, device=device) else: - block, nw = sweep(args.ng, args.h1, args.nnei, args.nodes, device) + block, nw = sweep( + args.ng, + args.h1, + basis_dim=args.basis_dim, + nnei=args.nnei, + nodes=args.nodes, + device=device, + ) print(f'\n"{torch.cuda.get_device_name()}": {{') - print(f" ({args.ng}, {args.h1}): ({block}, {nw}),") + key = ( + f"({args.ng}, {args.h1}, {args.basis_dim})" + if args.kind == "conv" + else f"({args.ng}, {args.h1})" + ) + print(f" {key}: ({block}, {nw}),") print("}") diff --git a/deepmd/kernels/triton/dpa1/tile_configs.py b/deepmd/kernels/triton/dpa1/tile_configs.py index 0218446231..3e5b0d7dbe 100644 --- a/deepmd/kernels/triton/dpa1/tile_configs.py +++ b/deepmd/kernels/triton/dpa1/tile_configs.py @@ -2,21 +2,22 @@ """Launch-configuration resolution for the DPA1 fused environment convolutions. Two memory-bound kernels are configured here, each reducing to a block width -and a warp count keyed by the channel width ``ng`` and the resnet width ``H1``: +and a warp count: - ``se_conv`` (strip / dense, node-parallel): one program owns a node and streams its neighbors, so the launch is ``(BLOCK_N, num_warps)`` -- neighbors per block. ``BLOCK_N`` bounds the live ``(BLOCK_N, channels)`` register footprint of the backward pass; oversized blocks spill and collapse - throughput, so the universal default is kept small. + throughput, so the universal default is kept small. Its key is + ``(ng, H1, basis_dim)`` because the 4/9/16/25-row moments have different + register pressure. - ``edge_conv`` (concat / graph, edge-parallel): one program owns a block of edges and scatters them into their center nodes, so the launch is ``(BLOCK_E, num_warps)`` -- edges per block. ``BLOCK_E`` bounds the live ``(BLOCK_E, channels)`` register footprint of both passes. -The optimum depends on ``(ng, H1)`` (the register footprint) but is insensitive -to the neighbor / edge count, which only sets the loop trip count or the grid -size. Table keys are therefore ``(ng, H1)``. +The optimum is insensitive to the neighbor / edge count, which only sets the +loop trip count or grid size. Level policy (see :func:`deepmd.kernels.utils.triton_infer_level`): @@ -46,17 +47,23 @@ # edge_conv (edge-parallel) universal default. EDGE_DEFAULT_CONFIG: Config = (8, 4) -# Per-GPU built-in tables keyed by (ng, H1). Values are the fastest spill-free -# forward+backward configuration produced by the fp32 sweep in -# :mod:`.sweep_tile_configs` for that channel width. -_CONV_BUILTIN: dict[str, dict[tuple[int, int], Config]] = { +# Per-GPU built-in tables keyed by (ng, H1, basis_dim). Values are the fastest +# spill-free forward+backward configuration produced by the fp32 sweep in +# :mod:`.sweep_tile_configs` for that shape. +_CONV_BUILTIN: dict[str, dict[tuple[int, int, int], Config]] = { "NVIDIA H20": { - (32, 16): (32, 2), - (64, 32): (16, 2), - (128, 64): (16, 2), - (256, 128): (16, 4), - (100, 50): (16, 2), - (200, 100): (16, 4), + (32, 16, 4): (32, 2), + (64, 32, 4): (16, 2), + (128, 64, 4): (16, 2), + (256, 128, 4): (16, 4), + (100, 50, 4): (16, 2), + (200, 100, 4): (16, 4), + (32, 16, 9): (64, 2), + (64, 32, 9): (32, 2), + (128, 64, 9): (16, 2), + (256, 128, 9): (16, 4), + (100, 50, 9): (16, 2), + (200, 100, 9): (16, 4), }, } _EDGE_BUILTIN: dict[str, dict[tuple[int, int], Config]] = { @@ -77,7 +84,7 @@ # shape keys the built-in tables do not cover. Process-local: the freeze traces # on the target GPU, so these are baked into the exported ``.pt2``; they never # persist across processes. Same schema as the built-in tables. -_CONV_RUNTIME: dict[str, dict[tuple[int, int], Config]] = {} +_CONV_RUNTIME: dict[str, dict[tuple[int, int, int], Config]] = {} _EDGE_RUNTIME: dict[str, dict[tuple[int, int], Config]] = {} @@ -123,27 +130,45 @@ def _resolve( # --- se_conv (node-parallel) ------------------------------------------------ -def register_conv_config(device_name: str, ng: int, h1: int, config: Config) -> None: - """Register a freshly swept ``se_conv`` launch for ``(ng, h1)``. +def register_conv_config( + device_name: str, + ng: int, + h1: int, + basis_dim: int, + config: Config, +) -> None: + """Register a freshly swept ``se_conv`` launch for ``(ng, h1, basis_dim)``. Used by the freeze-time autotuner so a subsequent :func:`resolve_conv_config` (made while tracing) bakes the tuned launch into the exported artifact. """ - _register(_CONV_RUNTIME, device_name, ng, h1, config) + key = (int(ng), int(h1), int(basis_dim)) + _CONV_RUNTIME.setdefault(device_name, {})[key] = config -def has_conv_config(ng: int, h1: int) -> bool: - """Whether a tuned ``se_conv`` entry (built-in or freeze-time) covers ``(ng, h1)``.""" - return _covered(_CONV_BUILTIN, _CONV_RUNTIME, ng, h1) +def has_conv_config(ng: int, h1: int, basis_dim: int) -> bool: + """Whether a tuned ``se_conv`` entry covers ``(ng, h1, basis_dim)``.""" + if not torch.cuda.is_available(): + return False + name = torch.cuda.get_device_name() + key = (int(ng), int(h1), int(basis_dim)) + return key in _CONV_RUNTIME.get(name, {}) or key in _CONV_BUILTIN.get(name, {}) -def resolve_conv_config(ng: int, h1: int, level: int) -> Config: +def resolve_conv_config(ng: int, h1: int, basis_dim: int, level: int) -> Config: """Resolve the ``(BLOCK_N, num_warps)`` for a fused ``se_conv`` launch. Level 1 forces the universal default; level ``>= 2`` consults the freeze-time and per-GPU tables with fallback. ``BLOCK_N`` is a power of two. """ - return _resolve(_CONV_BUILTIN, _CONV_RUNTIME, DEFAULT_CONFIG, ng, h1, level) + if level < 2 or not torch.cuda.is_available(): + return DEFAULT_CONFIG + name = torch.cuda.get_device_name() + key = (int(ng), int(h1), int(basis_dim)) + runtime_dev = _CONV_RUNTIME.get(name, {}) + if key in runtime_dev: + return runtime_dev[key] + return _CONV_BUILTIN.get(name, {}).get(key, DEFAULT_CONFIG) # --- edge_conv (edge-parallel) ---------------------------------------------- diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index 7b0ff403c5..d21d012335 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -168,8 +168,10 @@ def share_params( # must share, even if not do stat self.mean = base_class.mean self.stddev = base_class.stddev - # self.load_state_dict(base_class.state_dict()) # this does not work, because it only inits the model - # the following will successfully link all the params except buffers + # Direct parameters are not part of the child-module registry. + for item in self._parameters: + self._parameters[item] = base_class._parameters[item] + # Child modules carry their own parameters and buffers. for item in self._modules: self._modules[item] = base_class._modules[item] else: diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 9751db621c..e18176e6a2 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -70,10 +70,17 @@ class DescrptDPA1(BaseDescriptor, torch.nn.Module): This descriptor, :math:`\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}`, is given by .. math:: - \mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<, - - where :math:`\hat{\mathcal{G}}^i` represents the embedding matrix:math:`\mathcal{G}^i` - after additional self-attention mechanism and :math:`\mathcal{R}^i` is defined by the full case in the se_e2_a descriptor. + \mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T + \mathcal{B}_{L}^i (\mathcal{B}_{L}^i)^T \hat{\mathcal{G}}^i_<, + + where :math:`\hat{\mathcal{G}}^i` represents the embedding matrix + :math:`\mathcal{G}^i` + after additional self-attention mechanism. For :math:`L=1`, + :math:`\mathcal{B}_{L}^i` is the full environment matrix + :math:`\mathcal{R}^i` of the se_e2_a descriptor. For :math:`L=2,3,4`, + norm-normalized real spherical harmonics are appended to each environment + row. Their per-degree inner products are :math:`P_l(\cos\theta)`, adding + higher angular correlations without enumerating neighbor pairs. Note that we obtain :math:`\mathcal{G}^i` using the type embedding method by default in this descriptor. To perform the self-attention mechanism, the queries :math:`\mathcal{Q}^{i,l} \in \mathbb{R}^{N_c\times d_k}`, @@ -134,6 +141,9 @@ class DescrptDPA1(BaseDescriptor, torch.nn.Module): Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}` axis_neuron: int Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix) + lmax: int + Maximum angular degree of the aggregated moment basis. Supported + values are 1 through 4. tebd_dim: int Dimension of the type embedding tebd_input_mode: str @@ -254,6 +264,7 @@ def __init__( # not implemented spin: Any | None = None, type: str | None = None, + lmax: int = 1, ) -> None: super().__init__() # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -280,6 +291,7 @@ def __init__( ntypes, neuron=neuron, axis_neuron=axis_neuron, + lmax=lmax, tebd_dim=tebd_dim, tebd_input_mode=tebd_input_mode, set_davg_zero=set_davg_zero, @@ -500,7 +512,7 @@ def serialize(self) -> dict: data = { "@class": "Descriptor", "type": "dpa1", - "@version": 2, + "@version": 4 if obj.lmax != 1 else 2, "rcut": obj.rcut, "rcut_smth": obj.rcut_smth, "sel": obj.sel, @@ -544,6 +556,12 @@ def serialize(self) -> dict: "trainable": self.trainable, "spin": None, } + if obj.adam_degree_gain_raw is not None: + data["@variables"]["degree_gain_raw"] = ( + obj.adam_degree_gain_raw.detach().cpu().numpy() + ) + if obj.lmax != 1: + data["lmax"] = obj.lmax if obj.tebd_input_mode in ["strip"]: data.update({"embeddings_strip": obj.filter_layers_strip.serialize()}) return data @@ -551,7 +569,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "DescrptDPA1": data = data.copy() - check_version_compatibility(data.pop("@version"), 3, 1) + check_version_compatibility(data.pop("@version"), 4, 1) data.pop("@class") data.pop("type") variables = data.pop("@variables") @@ -568,6 +586,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": # compat with version 1 if "use_tebd_bias" not in data: data["use_tebd_bias"] = True + data.setdefault("lmax", 1) obj = cls(**data) def t_cvt(xx: Any) -> torch.Tensor: @@ -578,6 +597,10 @@ def t_cvt(xx: Any) -> torch.Tensor: ) obj.se_atten["davg"] = t_cvt(variables["davg"]) obj.se_atten["dstd"] = t_cvt(variables["dstd"]) + if obj.se_atten.adam_degree_gain_raw is not None: + obj.se_atten.adam_degree_gain_raw.data.copy_( + t_cvt(variables["degree_gain_raw"]) + ) obj.se_atten.filter_layers = NetworkCollection.deserialize(embeddings) if tebd_input_mode in ["strip"]: obj.se_atten.filter_layers_strip = NetworkCollection.deserialize( diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 76c1db65e4..fe20fffb83 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -52,6 +52,9 @@ from deepmd.pt.utils.exclude_mask import ( PairExcludeMask, ) +from deepmd.pt.utils.utils import ( + get_generator, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -82,6 +85,168 @@ def tabulate_fusion_se_atten( torch.ops.deepmd.tabulate_fusion_se_atten = tabulate_fusion_se_atten +_DEGREE_GAIN_INIT_STD = 0.1 + + +def _safe_direction( + diff: torch.Tensor, + protection: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Scale displacements by the protected distance denominator. + + Parameters + ---------- + diff + Neighbor displacement vectors with shape ``(..., 3)``. + protection + Distance protection added to the denominator. + + Returns + ------- + torch.Tensor + Protected displacement coordinates with shape ``(..., 3)``. + torch.Tensor + Unprotected distances with shape ``(..., 1)``. + torch.Tensor + Nonzero-distance mask with shape ``(..., 1)``. + """ + distance_squared = torch.sum(diff * diff, dim=-1, keepdim=True) + direction_mask = distance_squared > 0.0 + safe_distance = torch.sqrt( + torch.where( + direction_mask, + distance_squared, + torch.ones_like(distance_squared), + ) + ) + denominator = torch.where( + direction_mask, + safe_distance + protection, + torch.ones_like(safe_distance), + ) + return diff / denominator, safe_distance, direction_mask + + +def _compute_angular_radial( + distance: torch.Tensor, + direction_mask: torch.Tensor, + switch: torch.Tensor, + stddev: torch.Tensor, + valid_mask: torch.Tensor, + protection: float, +) -> torch.Tensor: + """Compute the zero-mean radial factor for non-scalar moments. + + Non-scalar SO(3) components cannot carry an additive statistical mean. The + factor therefore applies the scalar environment standard deviation without + subtracting its mean, which keeps every angular component zero at the + cutoff. + """ + basis_mask = valid_mask.unsqueeze(-1) & direction_mask + safe_distance = torch.where( + basis_mask, + distance + protection, + torch.ones_like(distance), + ) + return switch / safe_distance / stddev * basis_mask.to(dtype=switch.dtype) + + +def _build_moment_basis( + rr: torch.Tensor, + direction: torch.Tensor, + radial: torch.Tensor, + lmax: int, +) -> torch.Tensor: + """Build the Cartesian moment basis through angular degree ``lmax``. + + The first four rows preserve the existing scalar and vector environment + matrix exactly. Higher-degree rows are norm-normalized real spherical + harmonics in the ``m=-l,...,l`` order: + + ``Y_l(u) @ Y_l(v) = P_l(u @ v)``. + + Parameters + ---------- + rr + Normalized environment matrix with shape ``(ncenter, nnei, 4)``. + direction + Protected displacement coordinates with shape + ``(ncenter, nnei, 3)``. + radial + Zero-mean normalized radial amplitude with shape + ``(ncenter, nnei, 1)``. Invalid and excluded neighbors are zero. + lmax + Maximum angular degree. Supported values are ``1`` through ``4``. + + Returns + ------- + torch.Tensor + Moment basis with shape ``(ncenter, nnei, (lmax + 1) ** 2)``. + """ + if lmax == 1: + return rr + + x, y, z = direction.unbind(dim=-1) + q = x * x + y * y + z * z + sqrt_three = 3.0**0.5 + degree_two = torch.stack( + ( + sqrt_three * x * y, + sqrt_three * y * z, + 0.5 * (3.0 * z * z - q), + sqrt_three * x * z, + 0.5 * sqrt_three * (x * x - y * y), + ), + dim=-1, + ) + blocks = [rr, radial * degree_two] + if lmax >= 3: + degree_three = torch.stack( + ( + (5.0 / 8.0) ** 0.5 * y * (3.0 * x * x - y * y), + 15.0**0.5 * x * y * z, + (3.0 / 8.0) ** 0.5 * y * (5.0 * z * z - q), + 0.5 * z * (5.0 * z * z - 3.0 * q), + (3.0 / 8.0) ** 0.5 * x * (5.0 * z * z - q), + 0.5 * 15.0**0.5 * z * (x * x - y * y), + (5.0 / 8.0) ** 0.5 * x * (x * x - 3.0 * y * y), + ), + dim=-1, + ) + blocks.append(radial * degree_three) + if lmax >= 4: + z2 = z * z + x2_minus_y2 = x * x - y * y + degree_four = torch.stack( + ( + 0.5 * 35.0**0.5 * x * y * x2_minus_y2, + 0.25 * 70.0**0.5 * y * z * (3.0 * x * x - y * y), + 0.5 * 5.0**0.5 * x * y * (7.0 * z2 - q), + 0.25 * 10.0**0.5 * y * z * (7.0 * z2 - 3.0 * q), + 0.125 * (35.0 * z2 * z2 - 30.0 * z2 * q + 3.0 * q * q), + 0.25 * 10.0**0.5 * x * z * (7.0 * z2 - 3.0 * q), + 0.25 * 5.0**0.5 * x2_minus_y2 * (7.0 * z2 - q), + 0.25 * 70.0**0.5 * x * z * (x * x - 3.0 * y * y), + 0.125 * 35.0**0.5 * (x**4 - 6.0 * x * x * y * y + y**4), + ), + dim=-1, + ) + blocks.append(radial * degree_four) + return torch.cat(blocks, dim=-1) + + +def _build_degree_weights( + raw_gain: torch.Tensor, + lmax: int, + reference: torch.Tensor, +) -> torch.Tensor: + """Expand non-negative per-degree Gram weights to packed moment rows.""" + blocks = [torch.ones(4, dtype=reference.dtype, device=reference.device)] + for degree in range(2, lmax + 1): + blocks.append(raw_gain[degree - 2].square().expand(2 * degree + 1)) + return torch.cat(blocks) + + @DescriptorBlock.register("se_atten") class DescrptBlockSeAtten(DescriptorBlock): def __init__( @@ -114,6 +279,7 @@ def __init__( seed: int | list[int] | None = None, type: str | None = None, trainable: bool = True, + lmax: int = 1, ) -> None: r"""Construct an embedding net of type `se_atten`. @@ -132,6 +298,9 @@ def __init__( Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}` axis_neuron : int Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix) + lmax : int + Maximum angular degree of the aggregated moment basis. Supported + values are 1 through 4. tebd_dim : int Dimension of the type embedding tebd_input_mode : str @@ -189,6 +358,9 @@ def __init__( self.neuron = neuron self.filter_neuron = self.neuron self.axis_neuron = axis_neuron + if lmax not in (1, 2, 3, 4): + raise ValueError(f"`lmax` must be between 1 and 4, got {lmax}") + self.lmax = int(lmax) self.tebd_dim = tebd_dim self.tebd_input_mode = tebd_input_mode self.set_davg_zero = set_davg_zero @@ -208,6 +380,23 @@ def __init__( self.env_protection = env_protection self.trainable_ln = trainable_ln self.seed = seed + if self.lmax > 1: + self.adam_degree_gain_raw = nn.Parameter( + torch.empty( + self.lmax - 1, + dtype=self.prec, + device=env.DEVICE, + ), + requires_grad=trainable, + ) + nn.init.normal_( + self.adam_degree_gain_raw, + mean=0.0, + std=_DEGREE_GAIN_INIT_STD, + generator=get_generator(child_seed(seed, 3)), + ) + else: + self.register_parameter("adam_degree_gain_raw", None) # to keep consistent with default value in this backends if ln_eps is None: ln_eps = 1e-5 @@ -611,6 +800,30 @@ def forward( rr = dmatrix rr = rr * exclude_mask[:, :, None] ss = rr[:, :, :1] + diff_flat = diff.view(nfnl, nnei, 3) + nlist_mask_flat = nlist_mask.view(nfnl, nnei) + moment_radial = rr[..., :1] + direction = diff_flat + if self.lmax > 1: + direction, distance, direction_mask = _safe_direction( + diff_flat, + self.env_protection, + ) + radial_stddev = self.stddev[atype][..., :1].view(nfnl, nnei, 1) + moment_radial = _compute_angular_radial( + distance, + direction_mask, + sw.view(nfnl, nnei, 1), + radial_stddev, + nlist_mask_flat, + self.env_protection, + ) + moment_basis = _build_moment_basis( + rr, + direction, + moment_radial, + self.lmax, + ) # Whether the pair representation ``g2`` is produced this pass. The # compressed and fused strip paths form only the moment tensor and # leave ``g2`` empty; the flag drives the return below. @@ -650,19 +863,19 @@ def forward( # embedding layer (its timestep and residual) and the moment # reduction collapse into one node-parallel Triton kernel while # the two embedding GEMMs stay on cuBLAS. - xyz_scatter = se_atten_conv( + moment = se_atten_conv( self.filter_layers.networks[0], ss, None, None, None, - rr, + moment_basis, gated=0, ) # ``gg`` (the pair representation g2) is not formed on this path; # it is returned as ``None``, so a zero-element placeholder keeps # ``gg`` a tensor without the full (nf, nloc, nnei, ng) allocation. - gg = xyz_scatter.new_empty(0) + gg = moment.new_empty(0) g2_is_none = True else: # nfnl x nnei x ng @@ -673,8 +886,8 @@ def forward( gg = self.dpa1_attention( gg, nlist_mask, input_r=input_r, sw=sw ) # shape is [nframes*nloc, self.neei, out_size] - # nfnl x 4 x ng - xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) + # nfnl x moment_dim x ng + moment = torch.matmul(moment_basis.permute(0, 2, 1), gg) elif self.tebd_input_mode in ["strip"]: assert self.filter_layers_strip is not None assert type_embedding is not None @@ -733,19 +946,19 @@ def forward( # the operator is a ``triton_op`` composable only under eager # and ``make_fx`` / ``torch.compile``. sw_eff = sw if self.smooth else torch.ones_like(sw) - xyz_scatter = se_atten_conv( + moment = se_atten_conv( self.filter_layers.networks[0], ss, tt_full, tebd_idx, sw_eff.reshape(nfnl, self.nnei), - rr, + moment_basis, gated=1, ) # ``gg`` (the pair representation g2) is not formed on this path; # it is returned as ``None``, so a zero-element placeholder keeps # ``gg`` a tensor without the full (nf, nloc, nnei, ng) allocation. - gg = xyz_scatter.new_empty(0) + gg = moment.new_empty(0) g2_is_none = True else: # (nf x nl) x nnei x ng @@ -755,11 +968,11 @@ def forward( if self.geo_compress: ss = ss.reshape(-1, 1) gg_t = gg_t.reshape(-1, gg_t.size(-1)) - xyz_scatter = torch.ops.deepmd.tabulate_fusion_se_atten( + moment = torch.ops.deepmd.tabulate_fusion_se_atten( self.compress_data[0].contiguous(), self.compress_info[0].cpu().contiguous(), ss.contiguous(), - rr.contiguous(), + moment_basis.contiguous(), gg_t.contiguous(), self.filter_neuron[-1], self.is_sorted, @@ -784,17 +997,25 @@ def forward( gg = self.dpa1_attention( gg, nlist_mask, input_r=input_r, sw=sw ) # shape is [nframes*nloc, self.neei, out_size] - # nfnl x 4 x ng - xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) + # nfnl x moment_dim x ng + moment = torch.matmul(moment_basis.permute(0, 2, 1), gg) else: raise NotImplementedError - xyz_scatter = xyz_scatter / self.nnei - xyz_scatter_1 = xyz_scatter.permute(0, 2, 1) - rot_mat = xyz_scatter_1[:, :, 1:4] - xyz_scatter_2 = xyz_scatter[:, :, 0 : self.axis_neuron] + moment = moment / self.nnei + moment_t = moment.permute(0, 2, 1) + rot_mat = moment_t[:, :, 1:4] + moment_axis = moment[:, :, 0 : self.axis_neuron] + if self.lmax > 1: + assert self.adam_degree_gain_raw is not None + degree_weights = _build_degree_weights( + self.adam_degree_gain_raw, + self.lmax, + moment, + ) + moment_axis = moment_axis * degree_weights.view(1, -1, 1) result = torch.matmul( - xyz_scatter_1, xyz_scatter_2 + moment_t, moment_axis ) # shape is [nframes*nloc, self.filter_neuron[-1], self.axis_neuron] return ( diff --git a/deepmd/pt/model/descriptor/se_atten_v2.py b/deepmd/pt/model/descriptor/se_atten_v2.py index e27938b166..02ccf36f1a 100644 --- a/deepmd/pt/model/descriptor/se_atten_v2.py +++ b/deepmd/pt/model/descriptor/se_atten_v2.py @@ -70,6 +70,7 @@ def __init__( # not implemented spin: Any | None = None, type: str | None = None, + lmax: int = 1, ) -> None: r"""Construct smooth version of embedding net of type `se_atten_v2`. @@ -88,6 +89,9 @@ def __init__( Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}` axis_neuron : int Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix) + lmax : int + Maximum angular degree of the aggregated moment basis. Supported + values are 1 through 4. tebd_dim : int Dimension of the type embedding set_davg_zero : bool @@ -160,6 +164,7 @@ def __init__( ntypes, neuron=neuron, axis_neuron=axis_neuron, + lmax=lmax, tebd_dim=tebd_dim, tebd_input_mode="strip", set_davg_zero=set_davg_zero, @@ -196,7 +201,7 @@ def serialize(self) -> dict: data = { "@class": "Descriptor", "type": "se_atten_v2", - "@version": 2, + "@version": 4 if obj.lmax != 1 else 2, "rcut": obj.rcut, "rcut_smth": obj.rcut_smth, "sel": obj.sel, @@ -237,12 +242,18 @@ def serialize(self) -> dict: "trainable": self.trainable, "spin": None, } + if obj.adam_degree_gain_raw is not None: + data["@variables"]["degree_gain_raw"] = ( + obj.adam_degree_gain_raw.detach().cpu().numpy() + ) + if obj.lmax != 1: + data["lmax"] = obj.lmax return data @classmethod def deserialize(cls, data: dict) -> "DescrptSeAttenV2": data = data.copy() - check_version_compatibility(data.pop("@version"), 3, 1) + check_version_compatibility(data.pop("@version"), 4, 1) data.pop("@class") data.pop("type") variables = data.pop("@variables") @@ -255,6 +266,7 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2": # compat with version 1 if "use_tebd_bias" not in data: data["use_tebd_bias"] = True + data.setdefault("lmax", 1) obj = cls(**data) def t_cvt(xx: Any) -> torch.Tensor: @@ -265,6 +277,10 @@ def t_cvt(xx: Any) -> torch.Tensor: ) obj.se_atten["davg"] = t_cvt(variables["davg"]) obj.se_atten["dstd"] = t_cvt(variables["dstd"]) + if obj.se_atten.adam_degree_gain_raw is not None: + obj.se_atten.adam_degree_gain_raw.data.copy_( + t_cvt(variables["degree_gain_raw"]) + ) obj.se_atten.filter_layers = NetworkCollection.deserialize(embeddings) obj.se_atten.filter_layers_strip = NetworkCollection.deserialize( embeddings_strip diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index b6a7f9f84f..80e6f3a088 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -10,6 +10,10 @@ cast_precision, ) from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP +from deepmd.dpmodel.descriptor.dpa1 import ( + build_dpa1_degree_weights, + build_dpa1_moment_basis, +) from deepmd.dpmodel.utils.env_mat_stat import ( merge_env_stat, ) @@ -95,12 +99,10 @@ def _env_mat( ) -> tuple: """Environment-matrix prologue shared by every fused path (strip / concat). - Returns ``(nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, - type_embedding)``: ``rr`` the ``(nfnl, nnei, 4)`` environment matrix - (excluded edges zeroed), ``ss`` its radial channel ``(nfnl, nnei, 1)``, - ``sw`` the ``(nfnl, nnei, 1)`` smooth cutoff (zeroed on excluded/padding - edges), and ``nlist_masked`` the neighbor indices with excluded/padding - entries mapped to ``0`` (for downstream gathers). + Returns ``(nf, nloc, nnei, ng, nfnl, rr, moment_basis, ss, sw, + nlist_masked, type_embedding)``. ``rr`` is the four-component environment + matrix, while ``moment_basis`` has four or nine components according to + ``lmax``. Excluded and padding edges are zeroed in both tensors. """ se = desc.se_atten nf, nloc, nnei = nlist.shape @@ -108,7 +110,7 @@ def _env_mat( # Fused env-matrix operator, captured opaquely under the pt_expt trace and # resolving to the Triton kernel at CUDA runtime; identical outputs to the # array-API ``EnvMat.call`` below. - rr, _diff, sw = _env_mat_triton( + rr, diff, sw = _env_mat_triton( coord_ext, nlist, atype_ext[:, :nloc], @@ -121,7 +123,7 @@ def _env_mat( use_exp_switch=se.env_mat.use_exp_switch, ) else: - rr, _diff, sw = se.env_mat.call( + rr, diff, sw = se.env_mat.call( coord_ext, atype_ext, nlist, se.mean[...], se.stddev[...] ) nf, nloc, nnei, _ = rr.shape @@ -144,9 +146,34 @@ def _env_mat( ) rr = rr.view(nfnl, nnei, 4) * exclude_mask[:, :, None].to(rr.dtype) ss = rr[:, :, :1] + moment_basis = rr + if se.lmax > 1: + diff = diff.view(nfnl, nnei, 3) + radial_stddev = se.stddev[:, :, :1][atype_ext[:, :nloc]].view(nfnl, nnei, 1) + moment_basis = build_dpa1_moment_basis( + rr, + diff, + sw, + radial_stddev, + nlist_mask, + se.lmax, + se.env_protection, + ) type_embedding = desc.type_embedding.call() - return nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, type_embedding + return ( + nf, + nloc, + nnei, + ng, + nfnl, + rr, + moment_basis, + ss, + sw, + nlist_masked, + type_embedding, + ) def _strip_pair_index( @@ -224,15 +251,22 @@ def _grrg_from_moment( ) -> Any: """Strip-mode epilogue: symmetry-invariant contraction of the moment. - Consumes the unnormalized moment ``xyz_scatter`` (nfnl, 4, ng), applies the - ``1 / nnei`` normalization, forms the ``G^T G`` descriptor and the rotation - matrix, and appends the center type embedding when ``concat_output_tebd``. + Consumes the unnormalized moment, applies the ``1 / nnei`` normalization, + forms the ``G^T G`` descriptor and the rotation matrix, and appends the + center type embedding when ``concat_output_tebd``. """ se = desc.se_atten xyz_scatter = xyz_scatter / se.nnei xyz_scatter_1 = xyz_scatter.permute(0, 2, 1) rot_mat = xyz_scatter_1[:, :, 1:4] xyz_scatter_2 = xyz_scatter[:, :, 0 : se.axis_neuron] + if se.lmax > 1: + degree_weights = build_dpa1_degree_weights( + se.adam_degree_gain_raw, + se.lmax, + xyz_scatter, + ) + xyz_scatter_2 = xyz_scatter_2 * degree_weights.view(1, -1, 1) result = torch.matmul(xyz_scatter_1, xyz_scatter_2) result = result.view(nf, nloc, ng * se.axis_neuron) rot_mat = rot_mat.view(nf, nloc, ng, 3) @@ -275,6 +309,7 @@ class DescrptDPA1(DescrptDPA1DP): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + self._promote_degree_gain() # Persisted graph-routing knob (first-class training configuration): # ``disable_graph_lower()`` used to flip only the plain dpmodel bool, # which a Trainer checkpoint restart silently reset (the fresh model @@ -289,6 +324,24 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: torch.zeros((), dtype=torch.bool, device="cpu"), ) + def _promote_degree_gain(self) -> None: + """Promote the dpmodel raw degree gains from a buffer to a Parameter.""" + block = self.se_atten + raw = block._buffers.get("adam_degree_gain_raw") + if raw is None: + return + del block._buffers["adam_degree_gain_raw"] + block.adam_degree_gain_raw = torch.nn.Parameter( + raw, + requires_grad=bool(block.trainable), + ) + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA1": + obj = super().deserialize(data) + obj._promote_degree_gain() + return obj + def disable_graph_lower(self) -> None: """Persisted variant of the dpmodel escape hatch (see base class). @@ -560,6 +613,7 @@ def call_graph( fused and triton_infer_level() >= 1 and not self.geo_compress + and self.se_atten.lmax == 1 and not self.se_atten.exclude_types and self._fused_eligible("triton") ): @@ -599,6 +653,11 @@ def _fused_eligible(self, backend: str) -> bool: precomputed outside the kernel). ``triton`` serves any layer stack (the head layers run on cuBLAS), so only the last activation must be inlined. + + Production CUDA binaries currently instantiate only ``lmax=1``. + The lmax 2/3/4 kernel sources are retained behind the + ``DEEPMD_ENABLE_DPA1_HIGH_LMAX`` CMake option and use the portable + reference path unless a dedicated experimental build enables them. """ se = self.se_atten if se.attn_layer != 0: @@ -611,6 +670,7 @@ def _fused_eligible(self, backend: str) -> bool: # matters, served up to 256 by the padded bucket dispatch. return ( se.tebd_input_mode == "strip" + and se.lmax == 1 and not se.exclude_types and se.mean.dtype == torch.float32 and self.compress_data[0].dtype == torch.float32 @@ -619,6 +679,8 @@ def _fused_eligible(self, backend: str) -> bool: and 0 < int(se.axis_neuron) <= min(16, int(se.neuron[-1])) and cuda_compress_available() ) + if se.lmax != 1: + return False widths = [int(layer.w.shape[1]) for layer in layers] first = layers[0] first_has_residual = bool(first.resnet) and first.w.shape[1] in ( @@ -677,9 +739,19 @@ def _call_triton( Composes under ``make_fx`` / ``torch.export`` so the operator is baked into the pt_expt ``.pt2``. """ - nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, type_embedding = _env_mat( - self, coord_ext, atype_ext, nlist - ) + ( + nf, + nloc, + nnei, + ng, + nfnl, + rr, + moment_basis, + ss, + sw, + nlist_masked, + type_embedding, + ) = _env_mat(self, coord_ext, atype_ext, nlist) se = self.se_atten strip = se.tebd_input_mode == "strip" # Embedding-net input: the radial channel (strip) or the radial-plus- @@ -727,7 +799,7 @@ def _call_triton( # unused by the kernel (``gated == 0``). gated = 0 tt_full, tebd_idx, sw_eff = concat_gate_placeholders(z2, ng) - # Unnormalized moment (nfnl, 4, ng); _grrg_from_moment applies 1 / nnei. + # Unnormalized moment; _grrg_from_moment applies 1 / nnei. xyz_scatter = se_conv( z2.contiguous(), h.contiguous(), @@ -735,7 +807,7 @@ def _call_triton( tt_full, tebd_idx, sw_eff, - rr, + moment_basis, resnet_mult, act, gated, @@ -760,9 +832,19 @@ def _call_compressed( nlist: torch.Tensor, ) -> Any: """Compressed forward for DPA1 descriptor (strip only).""" - nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, type_embedding = _env_mat( - self, coord_ext, atype_ext, nlist - ) + ( + nf, + nloc, + nnei, + ng, + nfnl, + rr, + moment_basis, + ss, + sw, + nlist_masked, + type_embedding, + ) = _env_mat(self, coord_ext, atype_ext, nlist) tebd_idx = _strip_pair_index( self, atype_ext, nlist_masked, type_embedding, nf, nloc, nnei ) @@ -777,7 +859,7 @@ def _call_compressed( self.compress_data[0].contiguous(), self.compress_info[0].cpu().contiguous(), ss.reshape(-1, 1).contiguous(), - rr.contiguous(), + moment_basis.contiguous(), gg_t.reshape(-1, gg_t.size(-1)).contiguous(), self.se_atten.neuron[-1], is_sorted, @@ -795,7 +877,7 @@ def _call_compressed( rr.view(-1, self.se_atten.nnei, 4)[:, :, 1:4], dim=-1 ) gg = self.se_atten.dpa1_attention(gg, nlist_mask, input_r=input_r, sw=sw) - xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) + xyz_scatter = torch.matmul(moment_basis.permute(0, 2, 1), gg) return _grrg_from_moment( self, @@ -822,7 +904,8 @@ def _call_graph_compress_reference( Evaluates the tabulated geometric embedding with the original fused table operator ``deepmd::tabulate_fusion_se_atten``, treating each edge as a one-neighbor block (``nloc = E``, ``nnei = 1``) so the operator - returns the per-edge moment outer product ``(E, 4, ng)``; a + returns the per-edge moment outer product with four or nine basis + rows; a ``segment_sum`` over edge centers then forms the per-node moment. This matches the dense :meth:`DescrptDPA1._call_compressed` (same table, same gate) to the fp32 summation-order floor, and composes under autograd so @@ -863,6 +946,17 @@ def _call_graph_compress_reference( sw = u**3 * (-6 * u**2 + 15 * u - 10) + 1.0 em = torch.cat([sw / q, ev * (sw / q**2)], dim=-1) rr = (em - se.mean[:, 0, :][center_type]) / se.stddev[:, 0, :][center_type] + moment_basis = rr + if se.lmax > 1: + moment_basis = build_dpa1_moment_basis( + rr, + ev, + sw, + se.stddev[:, 0, 0:1][center_type], + graph.edge_mask, + se.lmax, + se.env_protection, + ) # === Step 2. Strip type-pair gate from the precomputed table === ntypes = type_embedding.shape[0] @@ -879,7 +973,7 @@ def _call_graph_compress_reference( self.compress_data[0].contiguous(), self.compress_info[0].cpu().contiguous(), rr[:, 0:1].contiguous(), - rr.reshape(-1, 1, 4).contiguous(), + moment_basis.reshape(-1, 1, moment_basis.shape[-1]).contiguous(), gate.contiguous(), ng, is_sorted, @@ -887,12 +981,25 @@ def _call_graph_compress_reference( # === Step 4. Moment reduction and G^T G contraction === outer = outer * graph.edge_mask[:, None, None].to(outer.dtype) - gr = torch.zeros(n_total, 4, ng, dtype=outer.dtype, device=outer.device) + gr = torch.zeros( + n_total, + moment_basis.shape[-1], + ng, + dtype=outer.dtype, + device=outer.device, + ) gr.index_add_(0, dst, outer) gr = gr / se.nnei - gr_perm = gr.permute(0, 2, 1) # (N, ng, 4) + gr_perm = gr.permute(0, 2, 1) rot_mat = gr_perm[:, :, 1:4] - gr_sub = gr[:, :, : se.axis_neuron] # (N, 4, axis) + gr_sub = gr[:, :, : se.axis_neuron] + if se.lmax > 1: + degree_weights = build_dpa1_degree_weights( + se.adam_degree_gain_raw, + se.lmax, + gr, + ) + gr_sub = gr_sub * degree_weights.view(1, -1, 1) grrg = torch.matmul(gr_perm, gr_sub).reshape(n_total, ng * se.axis_neuron) grrg = grrg.to(graph.edge_vec.dtype) if self.concat_output_tebd: diff --git a/deepmd/pt_expt/descriptor/se_atten_v2.py b/deepmd/pt_expt/descriptor/se_atten_v2.py index c50bf8292d..34beb7322b 100644 --- a/deepmd/pt_expt/descriptor/se_atten_v2.py +++ b/deepmd/pt_expt/descriptor/se_atten_v2.py @@ -22,6 +22,24 @@ class DescrptSeAttenV2(DescrptSeAttenV2DP): _update_sel_cls = UpdateSel + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + DescrptDPA1._promote_degree_gain(self) + + @classmethod + def deserialize(cls, data: dict) -> "DescrptSeAttenV2": + obj = super().deserialize(data) + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + DescrptDPA1._promote_degree_gain(obj) + return obj + def share_params(self, *args: Any, **kwargs: Any) -> None: from deepmd.pt_expt.descriptor.dpa1 import ( DescrptDPA1, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 92285f951f..c6c7fcfd78 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1300,6 +1300,7 @@ def descrpt_se_atten_common_args() -> list[Argument]: doc_rcut_smth = "Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`" doc_neuron = "Number of neurons in each hidden layer of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built." doc_axis_neuron = "Size of the submatrix of `G` (the embedding matrix) used to build the descriptor." + doc_lmax = "Maximum angular degree of the aggregated moment basis. The PyTorch backend supports integer values from 1 through 4." doc_activation_function = f'The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' doc_type_one_side = r"If 'False', type embeddings of both neighbor and central atoms are considered. If 'True', only type embeddings of neighbor atoms are considered. Default is 'False'." @@ -1330,6 +1331,12 @@ def descrpt_se_atten_common_args() -> list[Argument]: alias=["n_axis_neuron"], doc=doc_axis_neuron, ), + Argument( + "lmax", + int, + optional=True, + doc=supported_backends("pt") + doc_lmax, + ), Argument( "activation_function", str, diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 05881436a1..0d0e3c5605 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -20,6 +20,16 @@ Attention-based descriptor $\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}$, whic where $\hat{\mathcal{G}}^i$ represents the embedding matrix $\mathcal{G}^i$ after additional self-attention mechanism and $\mathcal{R}^i$ is defined by the full case in the [`se_e2_a`](./train-se-e2-a.md). Note that we obtain $\mathcal{G}^i$ using the type embedding method by default in this descriptor. By default, we concat $s(r_{ij})$ and the type embeddings of central and neighboring atoms $\mathcal{A}^i$ and $\mathcal{A}^j$ as input of the embedding network $\mathcal{N}_{e,2}$: +The PyTorch and DP model implementations optionally extend $\mathcal{R}^i$ with +norm-normalized real spherical harmonics through degree four. Setting +{ref}`lmax ` to `2`, `3`, or `4` +changes the moment dimension from four to nine, sixteen, or twenty-five while +preserving the neighbor-linear reduction. Rows of degree two and above use the +zero-mean radial factor $s(r_{ij})/\sigma_s$ and therefore do not introduce an +explicit neighbor-pair or angle axis. Each added degree has a trainable, +non-negative Gram weight initialized near zero, so the descriptor begins close +to its lower-degree form and can enable useful angular orders during training. + ```math (\mathcal{G}^i)_j = \mathcal{N}_{e,2}(\{s(r_{ij}), \mathcal{A}^i, \mathcal{A}^j\}) \quad \mathrm{or}\quad(\mathcal{G}^i)_j = \mathcal{N}_{e,2}(\{s(r_{ij}), \mathcal{A}^j\}) ``` @@ -102,6 +112,7 @@ An example of the DPA-1 descriptor is provided as follows - **{ref}`sel `** gives the maximum possible number of neighbors in the cut-off radius. It is an int. Note that this number highly affects the efficiency of training, which we usually use less than 200. (We use 120 for training 56 elements in [OC2M dataset](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md)) - The {ref}`neuron ` specifies the size of the embedding net. From left to right the members denote the sizes of each hidden layer from the input end to the output end, respectively. If the outer layer is twice the size of the inner layer, then the inner layer is copied and concatenated, then a [ResNet architecture](https://arxiv.org/abs/1512.03385) is built between them. - The {ref}`axis_neuron ` specifies the size of the submatrix of the embedding matrix, the axis matrix as explained in the [DeepPot-SE paper](https://arxiv.org/abs/1805.09003) +- {ref}`lmax ` selects the maximum angular degree of the moment basis from `1` through `4`. Higher values add quadrupole, octupole, and hexadecapole components with trainable degree weights. These options are currently supported by the PyTorch and DP model implementations. - If the option {ref}`resnet_dt ` is set to `true`, then a timestep is used in the ResNet. - {ref}`seed ` gives the random seed that is used to generate random numbers when initializing the model parameters. - {ref}`attn ` sets the length of a hidden vector during scale-dot attention computation. @@ -194,6 +205,12 @@ Model compression is supported only when the descriptor attention depth {ref}`at Model compression is supported for any {ref}`attn_layer ` value when {ref}`tebd_input_mode ` is `"strip"`. When `attn_layer` is 0, both the type embedding and geometric parts are compressed. When `attn_layer` is not 0, only the type embedding is compressed while the geometric part keeps the neural network implementation (a warning is emitted during compression). +In the pt_expt CUDA graph lower, automatic fused routing is limited to `lmax` +1 for both compressed and uncompressed descriptors. The `lmax` 2, 3, and 4 +CUDA specializations are retained for experimental builds enabled with +`DEEPMD_ENABLE_DPA1_HIGH_LMAX=ON`; the production eligibility gate continues +to route those descriptors through the portable reference path. + ## Training example Here we upload the AlMgCu example shown in the paper, you can download it here: diff --git a/source/lib/include/tabulate.h b/source/lib/include/tabulate.h index ab57ddf6e9..4ec0e5c80b 100644 --- a/source/lib/include/tabulate.h +++ b/source/lib/include/tabulate.h @@ -1,8 +1,30 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once +#include +#include + +#include "errors.h" + namespace deepmd { +inline bool is_supported_se_a_basis_dimension( + const std::int64_t ndescrpt) noexcept { + return ndescrpt == 4 || ndescrpt == 9 || ndescrpt == 16 || ndescrpt == 25; +} + +namespace detail { + +inline void check_se_a_basis_dimension(const std::int64_t ndescrpt) { + if (!is_supported_se_a_basis_dimension(ndescrpt)) { + throw deepmd_exception( + "The environment basis dimension must be 4, 9, 16, or 25, got " + + std::to_string(ndescrpt)); + } +} + +} // namespace detail + template void tabulate_fusion_se_a_cpu(FPTYPE* out, const FPTYPE* table, @@ -13,7 +35,8 @@ void tabulate_fusion_se_a_cpu(FPTYPE* out, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted = true); + const bool is_sorted = true, + const int ndescrpt = 4); template void tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, @@ -28,7 +51,8 @@ void tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted = true); + const bool is_sorted = true, + const int ndescrpt = 4); template void tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, @@ -43,7 +67,8 @@ void tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted = true); + const bool is_sorted = true, + const int ndescrpt = 4); template void tabulate_fusion_se_t_cpu(FPTYPE* out, @@ -157,7 +182,8 @@ void tabulate_fusion_se_a_gpu(FPTYPE* out, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted = true); + const bool is_sorted = true, + const int ndescrpt = 4); template void tabulate_fusion_se_a_grad_gpu(FPTYPE* dy_dem_x, @@ -172,7 +198,8 @@ void tabulate_fusion_se_a_grad_gpu(FPTYPE* dy_dem_x, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted = true); + const bool is_sorted = true, + const int ndescrpt = 4); template void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy, @@ -187,7 +214,8 @@ void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted = true); + const bool is_sorted = true, + const int ndescrpt = 4); template void tabulate_fusion_se_t_gpu(FPTYPE* out, diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index 5534ecb8ea..a2ce1c7a6f 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -1,5 +1,10 @@ #include +#if GOOGLE_CUDA +#include +#include +#endif + #include "device.h" #include "tabulate.h" @@ -324,9 +329,9 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( FPTYPE var[6]; for (int ii = 0; ii < nnei; ii++) { FPTYPE xx = em_x[block_idx * nnei + ii]; - if (xx == ago && em[block_idx * nnei * 4 + ii * 4 + 1] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 2] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 3] == 0. && is_sorted) { + const int em_base = block_idx * nnei * MTILE + ii * MTILE; + if (xx == ago && em[em_base + 1] == 0. && em[em_base + 2] == 0. && + em[em_base + 3] == 0. && is_sorted) { unloop = true; breakpoint = ii; } @@ -353,8 +358,7 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( #else #error "should not touch here" #endif - += (nnei - breakpoint) * - em[block_idx * nnei * MTILE + ii * MTILE + kk] * res; + += (nnei - breakpoint) * em[em_base + kk] * res; } if (unloop) { break; @@ -374,7 +378,7 @@ __global__ void tabulate_fusion_se_a_fifth_order_polynomial( } } -template +template __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( FPTYPE* dy_dem_x, FPTYPE* dy_dem, @@ -399,15 +403,18 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( int warp_idx = GpuShuffleSync(0xffffffff, threadIdx.x / WARP_SIZE, 0); int lane_idx = threadIdx.x % WARP_SIZE; __shared__ int breakpoint; - FPTYPE* iteratorA = (FPTYPE*)&_data[0]; // dy - for (int ii = 0; ii < MTILE; ii++) { - for (int jj = thread_idx; jj < last_layer_size; jj += blockDim.x) { - iteratorA[ii * last_layer_size + jj] = - dy[block_idx * MTILE * last_layer_size + ii * last_layer_size + jj]; + const FPTYPE* iteratorA = dy + block_idx * MTILE * last_layer_size; + if (CACHE_DY) { + FPTYPE* shared_dy = (FPTYPE*)&_data[0]; + for (int ii = 0; ii < MTILE; ii++) { + for (int jj = thread_idx; jj < last_layer_size; jj += blockDim.x) { + shared_dy[ii * last_layer_size + jj] = + iteratorA[ii * last_layer_size + jj]; + } } + __syncthreads(); + iteratorA = shared_dy; } - __syncthreads(); - // Sorted padding must be folded at the first sentinel for the whole atom, // exactly as in the sequential CPU implementation. A warp-local search can // select several later sentinels because neighbor indices are striped over @@ -418,10 +425,9 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( if (is_sorted) { const FPTYPE ago = em_x[block_idx * nnei + nnei - 1]; for (int ii = 0; ii < nnei; ++ii) { - if (ago == em_x[block_idx * nnei + ii] && - em[block_idx * nnei * 4 + ii * 4 + 1] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 2] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 3] == 0.) { + const int em_base = block_idx * nnei * MTILE + ii * MTILE; + if (ago == em_x[block_idx * nnei + ii] && em[em_base + 1] == 0. && + em[em_base + 2] == 0. && em[em_base + 3] == 0.) { breakpoint = ii; break; } @@ -436,16 +442,17 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( for (int tile = 0; tile < nnei && tile <= breakpoint; tile += KTILE) { const int ii = tile + warp_idx; const bool active = ii < nnei && ii <= breakpoint; + const int em_base = block_idx * nnei * MTILE + ii * MTILE; FPTYPE Csub = (FPTYPE)0.; FPTYPE sum[MTILE] = {(FPTYPE)0.}; if (active) { const int repeat_count = ii == breakpoint ? nnei - breakpoint : 1; FPTYPE xx = em_x[block_idx * nnei + ii]; int table_idx = 0; - FPTYPE reg_em[MTILE] = {em[block_idx * nnei * MTILE + ii * MTILE + 0], - em[block_idx * nnei * MTILE + ii * MTILE + 1], - em[block_idx * nnei * MTILE + ii * MTILE + 2], - em[block_idx * nnei * MTILE + ii * MTILE + 3]}; + FPTYPE reg_em[MTILE]; + for (int kk = 0; kk < MTILE; ++kk) { + reg_em[kk] = em[em_base + kk]; + } FPTYPE extrapolate_delta = (FPTYPE)0.; locate_xx_se_a(xx, table_idx, lower, upper, max, stride0, stride1, extrapolate_delta); @@ -466,10 +473,10 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( for (int kk = 0; kk < MTILE; kk++) { sum[kk] += repeat_count * iteratorA[kk * last_layer_size + jj] * res; } - res = reg_em[0] * iteratorA[0 * last_layer_size + jj]; - res += reg_em[1] * iteratorA[1 * last_layer_size + jj]; - res += reg_em[2] * iteratorA[2 * last_layer_size + jj]; - res += reg_em[3] * iteratorA[3 * last_layer_size + jj]; + res = (FPTYPE)0.; + for (int kk = 0; kk < MTILE; ++kk) { + res += reg_em[kk] * iteratorA[kk * last_layer_size + jj]; + } Csub += repeat_count * res_grad * (enable_se_atten ? res * t + res : res); if (enable_se_atten) { @@ -490,7 +497,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( warp_reduce(Csub); if (lane_idx == 0) { for (int kk = 0; kk < MTILE; kk++) { - dy_dem[block_idx * nnei * MTILE + ii * MTILE + kk] = sum[kk]; + dy_dem[em_base + kk] = sum[kk]; } dy_dem_x[block_idx * nnei + ii] = Csub; } @@ -523,20 +530,26 @@ __global__ void tabulate_fusion_se_a_grad_grad_fifth_order_polynomial( FPTYPE ago = GpuShuffleSync(0xffffffff, em_x[block_idx * nnei + nnei - 1], 0); bool unloop = false; int breakpoint = nnei - 1; +#if GOOGLE_CUDA + FPTYPE sum[MTILE] = {(FPTYPE)0.}; +#elif TENSORFLOW_USE_ROCM FPTYPE* iteratorC = (FPTYPE*)&_data[0]; for (int kk = 0; kk < MTILE; kk++) { iteratorC[kk * last_layer_size + thread_idx] = (FPTYPE)0.; } __syncthreads(); +#else +#error "should not touch here" +#endif int mark_table_idx = -1; FPTYPE var[6]; for (int ii = 0; ii < nnei; ii++) { FPTYPE xx = em_x[block_idx * nnei + ii]; FPTYPE dz_xx = dz_dy_dem_x[block_idx * nnei + ii]; - if (xx == ago && em[block_idx * nnei * 4 + ii * 4 + 1] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 2] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 3] == 0. && is_sorted) { + const int em_base = block_idx * nnei * MTILE + ii * MTILE; + if (xx == ago && em[em_base + 1] == 0. && em[em_base + 2] == 0. && + em[em_base + 3] == 0. && is_sorted) { unloop = true; breakpoint = ii; } @@ -587,7 +600,13 @@ __global__ void tabulate_fusion_se_a_grad_grad_fifth_order_polynomial( */ for (int kk = 0; kk < MTILE; kk++) { int em_index = block_idx * nnei * MTILE + ii * MTILE + kk; +#if GOOGLE_CUDA + sum[kk] += +#elif TENSORFLOW_USE_ROCM iteratorC[kk * last_layer_size + thread_idx] += +#else +#error "should not touch here" +#endif (nnei - breakpoint) * (em[em_index] * (res_grad * dz_xx + two_grad) + dz_dy_dem[em_index] * res); } @@ -598,7 +617,14 @@ __global__ void tabulate_fusion_se_a_grad_grad_fifth_order_polynomial( } for (int ii = 0; ii < MTILE; ii++) { dz_dy[block_idx * MTILE * last_layer_size + ii * last_layer_size + - thread_idx] = iteratorC[ii * last_layer_size + thread_idx]; + thread_idx] = +#if GOOGLE_CUDA + sum[ii]; +#elif TENSORFLOW_USE_ROCM + iteratorC[ii * last_layer_size + thread_idx]; +#else +#error "should not touch here" +#endif } } @@ -1056,6 +1082,139 @@ __global__ void tabulate_fusion_se_r_grad_grad_fifth_order_polynomial( } } +template +void launch_tabulate_fusion_se_a(FPTYPE* out, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted) { + tabulate_fusion_se_a_fifth_order_polynomial +#if GOOGLE_CUDA + <<>> +#elif TENSORFLOW_USE_ROCM + <<>> +#else +#error "should not touch here" +#endif + (out, table, em_x, em, two_embed, table_info[0], table_info[1], + table_info[2], table_info[3], table_info[4], nnei, last_layer_size, + is_sorted); +} + +#if GOOGLE_CUDA +namespace { + +struct CudaSharedMemoryLimits { + size_t standard; + size_t opt_in; +}; + +CudaSharedMemoryLimits get_cuda_shared_memory_limits(const int device) { + static std::mutex cache_mutex; + static std::unordered_map cache; + std::lock_guard lock(cache_mutex); + const auto cached = cache.find(device); + if (cached != cache.end()) { + return cached->second; + } + + cudaDeviceProp properties; + DPErrcheck(cudaGetDeviceProperties(&properties, device)); + const CudaSharedMemoryLimits limits{ + properties.sharedMemPerBlock, + properties.sharedMemPerBlockOptin, + }; + cache.emplace(device, limits); + return limits; +} + +} // namespace +#endif + +template +void launch_tabulate_fusion_se_a_grad(FPTYPE* dy_dem_x, + FPTYPE* dy_dem, + FPTYPE* dy_dtwo, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const FPTYPE* dy, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted) { +#if GOOGLE_CUDA + const size_t shared_memory = sizeof(FPTYPE) * MTILE * last_layer_size; + int device = 0; + DPErrcheck(cudaGetDevice(&device)); + const CudaSharedMemoryLimits limits = get_cuda_shared_memory_limits(device); + const size_t shared_memory_limit = + limits.standard > limits.opt_in ? limits.standard : limits.opt_in; + if (shared_memory <= shared_memory_limit) { + auto kernel = + tabulate_fusion_se_a_grad_fifth_order_polynomial; + if (shared_memory > limits.standard) { + DPErrcheck(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(shared_memory))); + } + kernel<<>>( + dy_dem_x, dy_dem, dy_dtwo, table, em_x, em, two_embed, dy, + table_info[0], table_info[1], table_info[2], table_info[3], + table_info[4], nnei, last_layer_size, is_sorted); + } else { + tabulate_fusion_se_a_grad_fifth_order_polynomial + <<>>(dy_dem_x, dy_dem, dy_dtwo, table, em_x, em, + two_embed, dy, table_info[0], table_info[1], + table_info[2], table_info[3], table_info[4], + nnei, last_layer_size, is_sorted); + } +#elif TENSORFLOW_USE_ROCM + tabulate_fusion_se_a_grad_fifth_order_polynomial + <<>>( + dy_dem_x, dy_dem, dy_dtwo, table, em_x, em, two_embed, dy, + table_info[0], table_info[1], table_info[2], table_info[3], + table_info[4], nnei, last_layer_size, is_sorted); +#else +#error "should not touch here" +#endif +} + +template +void launch_tabulate_fusion_se_a_grad_grad(FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const FPTYPE* dz_dy_dem_x, + const FPTYPE* dz_dy_dem, + const FPTYPE* dz_dy_dtwo, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted) { + tabulate_fusion_se_a_grad_grad_fifth_order_polynomial +#if GOOGLE_CUDA + <<>>( +#elif TENSORFLOW_USE_ROCM + <<>>( +#else +#error "should not touch here" +#endif + dz_dy, table, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, dz_dy_dtwo, + table_info[0], table_info[1], table_info[2], table_info[3], + table_info[4], nnei, last_layer_size, is_sorted); +} + namespace deepmd { template void tabulate_fusion_se_a_gpu(FPTYPE* out, @@ -1067,7 +1226,9 @@ void tabulate_fusion_se_a_gpu(FPTYPE* out, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted) { + const bool is_sorted, + const int ndescrpt) { + detail::check_se_a_basis_dimension(ndescrpt); if (nloc <= 0) { return; } @@ -1076,21 +1237,28 @@ void tabulate_fusion_se_a_gpu(FPTYPE* out, if (nnei <= 0) { // The descriptor does not carry the empty neighbor dimension, so its // mathematically empty reduction must be materialized explicitly. - DPErrcheck(gpuMemset(out, 0, sizeof(FPTYPE) * nloc * MM * last_layer_size)); + DPErrcheck( + gpuMemset(out, 0, sizeof(FPTYPE) * nloc * ndescrpt * last_layer_size)); DPErrcheck(gpuDeviceSynchronize()); return; } - tabulate_fusion_se_a_fifth_order_polynomial -#if GOOGLE_CUDA - <<>> -#elif TENSORFLOW_USE_ROCM - <<>> -#else -#error "should not touch here" -#endif - (out, table, em_x, em, two_embed, table_info[0], table_info[1], - table_info[2], table_info[3], table_info[4], nnei, last_layer_size, - is_sorted); + if (ndescrpt == 4) { + launch_tabulate_fusion_se_a(out, table, table_info, em_x, em, + two_embed, nloc, nnei, + last_layer_size, is_sorted); + } else if (ndescrpt == 9) { + launch_tabulate_fusion_se_a(out, table, table_info, em_x, em, + two_embed, nloc, nnei, + last_layer_size, is_sorted); + } else if (ndescrpt == 16) { + launch_tabulate_fusion_se_a(out, table, table_info, em_x, em, + two_embed, nloc, nnei, + last_layer_size, is_sorted); + } else { + launch_tabulate_fusion_se_a(out, table, table_info, em_x, em, + two_embed, nloc, nnei, + last_layer_size, is_sorted); + } DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); } @@ -1108,14 +1276,16 @@ void tabulate_fusion_se_a_grad_gpu(FPTYPE* dy_dem_x, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted) { + const bool is_sorted, + const int ndescrpt) { + detail::check_se_a_basis_dimension(ndescrpt); if (nloc <= 0 || nnei <= 0) { return; } DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); DPErrcheck(gpuMemset(dy_dem_x, 0, sizeof(FPTYPE) * nloc * nnei)); - DPErrcheck(gpuMemset(dy_dem, 0, sizeof(FPTYPE) * nloc * nnei * 4)); + DPErrcheck(gpuMemset(dy_dem, 0, sizeof(FPTYPE) * nloc * nnei * ndescrpt)); if (two_embed != nullptr && is_sorted) { // The sorted-padding fast path writes only the first sentinel. Explicitly // clear the unused tail because framework output buffers are uninitialized. @@ -1123,11 +1293,23 @@ void tabulate_fusion_se_a_grad_gpu(FPTYPE* dy_dem_x, gpuMemset(dy_dtwo, 0, sizeof(FPTYPE) * nloc * nnei * last_layer_size)); } - tabulate_fusion_se_a_grad_fifth_order_polynomial - <<>>( - dy_dem_x, dy_dem, dy_dtwo, table, em_x, em, two_embed, dy, - table_info[0], table_info[1], table_info[2], table_info[3], - table_info[4], nnei, last_layer_size, is_sorted); + if (ndescrpt == 4) { + launch_tabulate_fusion_se_a_grad( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + } else if (ndescrpt == 9) { + launch_tabulate_fusion_se_a_grad( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + } else if (ndescrpt == 16) { + launch_tabulate_fusion_se_a_grad( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + } else { + launch_tabulate_fusion_se_a_grad( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + } DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); } @@ -1145,7 +1327,9 @@ void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted) { + const bool is_sorted, + const int ndescrpt) { + detail::check_se_a_basis_dimension(ndescrpt); if (nloc <= 0) { return; } @@ -1154,17 +1338,30 @@ void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy, if (nnei <= 0) { // Unlike the neighbor-shaped inputs, dz_dy remains non-empty and must be // initialized to the zero second derivative of an empty reduction. - DPErrcheck( - gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * MM * last_layer_size)); + DPErrcheck(gpuMemset(dz_dy, 0, + sizeof(FPTYPE) * nloc * ndescrpt * last_layer_size)); DPErrcheck(gpuDeviceSynchronize()); return; } - DPErrcheck(gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size)); - tabulate_fusion_se_a_grad_grad_fifth_order_polynomial - <<>>( - dz_dy, table, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, dz_dy_dtwo, - table_info[0], table_info[1], table_info[2], table_info[3], - table_info[4], nnei, last_layer_size, is_sorted); + DPErrcheck( + gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * ndescrpt * last_layer_size)); + if (ndescrpt == 4) { + launch_tabulate_fusion_se_a_grad_grad( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + } else if (ndescrpt == 9) { + launch_tabulate_fusion_se_a_grad_grad( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + } else if (ndescrpt == 16) { + launch_tabulate_fusion_se_a_grad_grad( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + } else { + launch_tabulate_fusion_se_a_grad_grad( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + } DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); } @@ -1420,7 +1617,8 @@ template void tabulate_fusion_se_a_gpu(float* out, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void tabulate_fusion_se_a_gpu(double* out, const double* table, const double* table_info, @@ -1430,7 +1628,8 @@ template void tabulate_fusion_se_a_gpu(double* out, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void tabulate_fusion_se_a_grad_gpu(float* dy_dem_x, float* dy_dem, float* dy_dtwo, @@ -1443,7 +1642,8 @@ template void tabulate_fusion_se_a_grad_gpu(float* dy_dem_x, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void tabulate_fusion_se_a_grad_gpu(double* dy_dem_x, double* dy_dem, double* dy_dtwo, @@ -1456,7 +1656,8 @@ template void tabulate_fusion_se_a_grad_gpu(double* dy_dem_x, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void tabulate_fusion_se_a_grad_grad_gpu( float* dz_dy, const float* table, @@ -1470,7 +1671,8 @@ template void tabulate_fusion_se_a_grad_grad_gpu( const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void tabulate_fusion_se_a_grad_grad_gpu( double* dz_dy, const double* table, @@ -1484,7 +1686,8 @@ template void tabulate_fusion_se_a_grad_grad_gpu( const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void tabulate_fusion_se_t_gpu(float* out, const float* table, diff --git a/source/lib/src/tabulate.cc b/source/lib/src/tabulate.cc index 54ee9e113f..ff011868dd 100644 --- a/source/lib/src/tabulate.cc +++ b/source/lib/src/tabulate.cc @@ -3,10 +3,11 @@ #include -#include #include #include +#include #include + /* This inline function was designed to get the table info and bias value for current input xx! lower: indicate the lower boundary of the first table; @@ -109,9 +110,13 @@ inline void locate_xx_se_t(const FPTYPE& lower, } } -template -inline FPTYPE dot(FPTYPE a[4], FPTYPE b[4]) { - return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] + a[3] * b[3]; +template +inline FPTYPE dot(const FPTYPE (&a)[size], const FPTYPE (&b)[size]) { + FPTYPE result = (FPTYPE)0.; + for (int ii = 0; ii < size; ++ii) { + result += a[ii] * b[ii]; + } + return result; } template @@ -152,19 +157,21 @@ inline FPTYPE extrapolated_polynomial5(const FPTYPE& a0, return polynomial5(a0, a1, a2, a3, a4, a5, xx) + grad * extrapolate_delta; } -template -void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, - const FPTYPE* table, - const FPTYPE* table_info, - const FPTYPE* em_x, - const FPTYPE* em, - const FPTYPE* two_embed, - const int nloc, - const int nnei, - const int last_layer_size, - const bool is_sorted) { +namespace { + +template +void tabulate_fusion_se_a_cpu_impl(FPTYPE* out, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted) { bool enable_se_atten = two_embed != nullptr; - memset(out, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size); + memset(out, 0, sizeof(FPTYPE) * nloc * NDESCRPT * last_layer_size); // An empty neighbor axis is a valid empty reduction. Return after // initializing the non-empty descriptor output instead of inspecting the // nonexistent last neighbor below. @@ -177,17 +184,17 @@ void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, const FPTYPE stride0 = table_info[3]; const FPTYPE stride1 = table_info[4]; // for every atom, execute a small manual gemm ~ -// FPTYPE * res = new FPTYPE[4 * last_layer_size]; +// FPTYPE * res = new FPTYPE[NDESCRPT * last_layer_size]; #pragma omp parallel for for (int ii = 0; ii < nloc; ii++) { - FPTYPE ll[4] = {0}; + FPTYPE ll[NDESCRPT] = {0}; FPTYPE ago = em_x[ii * nnei + nnei - 1]; bool unloop = false; for (int jj = 0; jj < nnei; jj++) { - ll[0] = em[ii * nnei * 4 + jj * 4 + 0]; - ll[1] = em[ii * nnei * 4 + jj * 4 + 1]; - ll[2] = em[ii * nnei * 4 + jj * 4 + 2]; - ll[3] = em[ii * nnei * 4 + jj * 4 + 3]; + const int em_base = ii * nnei * NDESCRPT + jj * NDESCRPT; + for (int mm = 0; mm < NDESCRPT; ++mm) { + ll[mm] = em[em_base + mm]; + } FPTYPE xx = em_x[ii * nnei + jj]; if (ago == xx && ll[1] == 0. && ll[2] == 0. && ll[3] == 0. && is_sorted) { unloop = true; @@ -211,24 +218,10 @@ void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, var = var * t + var; } - if (unloop) { - out[ii * last_layer_size * 4 + 0 * last_layer_size + kk] += - (nnei - jj) * var * ll[0]; - out[ii * last_layer_size * 4 + 1 * last_layer_size + kk] += - (nnei - jj) * var * ll[1]; - out[ii * last_layer_size * 4 + 2 * last_layer_size + kk] += - (nnei - jj) * var * ll[2]; - out[ii * last_layer_size * 4 + 3 * last_layer_size + kk] += - (nnei - jj) * var * ll[3]; - } else { - out[ii * last_layer_size * 4 + 0 * last_layer_size + kk] += - var * ll[0]; - out[ii * last_layer_size * 4 + 1 * last_layer_size + kk] += - var * ll[1]; - out[ii * last_layer_size * 4 + 2 * last_layer_size + kk] += - var * ll[2]; - out[ii * last_layer_size * 4 + 3 * last_layer_size + kk] += - var * ll[3]; + const FPTYPE scale = unloop ? (nnei - jj) * var : var; + for (int mm = 0; mm < NDESCRPT; ++mm) { + out[ii * last_layer_size * NDESCRPT + mm * last_layer_size + kk] += + scale * ll[mm]; } } if (unloop) { @@ -238,20 +231,20 @@ void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, } } -template -void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, - FPTYPE* dy_dem, - FPTYPE* dy_dtwo, - const FPTYPE* table, - const FPTYPE* table_info, - const FPTYPE* em_x, - const FPTYPE* em, - const FPTYPE* two_embed, - const FPTYPE* dy, - const int nloc, - const int nnei, - const int last_layer_size, - const bool is_sorted) { +template +void tabulate_fusion_se_a_grad_cpu_impl(FPTYPE* dy_dem_x, + FPTYPE* dy_dem, + FPTYPE* dy_dtwo, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const FPTYPE* dy, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted) { // Every gradient output has a zero-sized neighbor axis in this case. Avoid // both zero-length memory operations on potentially null tensor pointers and // the last-neighbor lookup in the atom loop. @@ -260,7 +253,7 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, } bool enable_se_atten = two_embed != nullptr; memset(dy_dem_x, 0, sizeof(FPTYPE) * nloc * nnei); - memset(dy_dem, 0, sizeof(FPTYPE) * nloc * nnei * 4); + memset(dy_dem, 0, sizeof(FPTYPE) * nloc * nnei * NDESCRPT); if (enable_se_atten) { memset(dy_dtwo, 0, sizeof(FPTYPE) * nloc * nnei * last_layer_size); } @@ -270,19 +263,19 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, FPTYPE const stride0 = table_info[3]; FPTYPE const stride1 = table_info[4]; // for every atom, execute a small gemm~ -// FPTYPE * res = new FPTYPE[4 * last_layer_size]; +// FPTYPE * res = new FPTYPE[NDESCRPT * last_layer_size]; #pragma omp parallel for for (int ii = 0; ii < nloc; ii++) { - FPTYPE ll[4]; - FPTYPE rr[4]; + FPTYPE ll[NDESCRPT] = {0}; + FPTYPE rr[NDESCRPT] = {0}; FPTYPE ago = em_x[ii * nnei + nnei - 1]; bool unloop = false; for (int jj = 0; jj < nnei; jj++) { // construct the dy/dx - ll[0] = em[ii * nnei * 4 + jj * 4 + 0]; - ll[1] = em[ii * nnei * 4 + jj * 4 + 1]; - ll[2] = em[ii * nnei * 4 + jj * 4 + 2]; - ll[3] = em[ii * nnei * 4 + jj * 4 + 3]; + const int em_base = ii * nnei * NDESCRPT + jj * NDESCRPT; + for (int mm = 0; mm < NDESCRPT; ++mm) { + ll[mm] = em[em_base + mm]; + } FPTYPE xx = em_x[ii * nnei + jj]; if (ago == xx && ll[1] == 0. && ll[2] == 0. && ll[3] == 0. && is_sorted) { unloop = true; @@ -293,10 +286,10 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, extrapolate_delta); FPTYPE grad = (FPTYPE)0.0; for (int kk = 0; kk < last_layer_size; kk++) { - rr[0] = dy[ii * last_layer_size * 4 + 0 * last_layer_size + kk]; - rr[1] = dy[ii * last_layer_size * 4 + 1 * last_layer_size + kk]; - rr[2] = dy[ii * last_layer_size * 4 + 2 * last_layer_size + kk]; - rr[3] = dy[ii * last_layer_size * 4 + 3 * last_layer_size + kk]; + for (int mm = 0; mm < NDESCRPT; ++mm) { + rr[mm] = + dy[ii * last_layer_size * NDESCRPT + mm * last_layer_size + kk]; + } FPTYPE a0 = table[table_idx * last_layer_size * 6 + 6 * kk + 0]; FPTYPE a1 = table[table_idx * last_layer_size * 6 + 6 * kk + 1]; FPTYPE a2 = table[table_idx * last_layer_size * 6 + 6 * kk + 2]; @@ -317,10 +310,9 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, FPTYPE dotllrr = dot(ll, rr); if (unloop) { grad += g * dotllrr * (nnei - jj); - dy_dem[ii * nnei * 4 + jj * 4 + 0] += res * rr[0] * (nnei - jj); - dy_dem[ii * nnei * 4 + jj * 4 + 1] += res * rr[1] * (nnei - jj); - dy_dem[ii * nnei * 4 + jj * 4 + 2] += res * rr[2] * (nnei - jj); - dy_dem[ii * nnei * 4 + jj * 4 + 3] += res * rr[3] * (nnei - jj); + for (int mm = 0; mm < NDESCRPT; ++mm) { + dy_dem[em_base + mm] += res * rr[mm] * (nnei - jj); + } if (enable_se_atten) { // Forward folds the complete padding tail using only this first // sentinel's two-embedding value. Its gradient therefore owns the @@ -331,10 +323,9 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, } } else { grad += g * dotllrr; - dy_dem[ii * nnei * 4 + jj * 4 + 0] += res * rr[0]; - dy_dem[ii * nnei * 4 + jj * 4 + 1] += res * rr[1]; - dy_dem[ii * nnei * 4 + jj * 4 + 2] += res * rr[2]; - dy_dem[ii * nnei * 4 + jj * 4 + 3] += res * rr[3]; + for (int mm = 0; mm < NDESCRPT; ++mm) { + dy_dem[em_base + mm] += res * rr[mm]; + } if (enable_se_atten) { dy_dtwo[ii * nnei * last_layer_size + jj * last_layer_size + kk] += resold * dotllrr; @@ -349,22 +340,22 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, } } -template -void deepmd::tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, - const FPTYPE* table, - const FPTYPE* table_info, - const FPTYPE* em_x, - const FPTYPE* em, - const FPTYPE* two_embed, - const FPTYPE* dz_dy_dem_x, - const FPTYPE* dz_dy_dem, - const FPTYPE* dz_dy_dtwo, - const int nloc, - const int nnei, - const int last_layer_size, - const bool is_sorted) { +template +void tabulate_fusion_se_a_grad_grad_cpu_impl(FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const FPTYPE* dz_dy_dem_x, + const FPTYPE* dz_dy_dem, + const FPTYPE* dz_dy_dtwo, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted) { bool enable_se_atten = two_embed != nullptr; - memset(dz_dy, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size); + memset(dz_dy, 0, sizeof(FPTYPE) * nloc * NDESCRPT * last_layer_size); // The second-order output retains the descriptor shape, so initialize the // empty reduction to zero before returning. if (nnei <= 0) { @@ -376,22 +367,19 @@ void deepmd::tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, const FPTYPE stride0 = table_info[3]; const FPTYPE stride1 = table_info[4]; // for every atom, execute a small manual gemm ~ -// FPTYPE * res = new FPTYPE[4 * last_layer_size]; +// FPTYPE * res = new FPTYPE[NDESCRPT * last_layer_size]; #pragma omp parallel for for (int ii = 0; ii < nloc; ii++) { - FPTYPE ll[4]; - FPTYPE hh[4]; + FPTYPE ll[NDESCRPT] = {0}; + FPTYPE hh[NDESCRPT] = {0}; FPTYPE ago = em_x[ii * nnei + nnei - 1]; bool unloop = false; for (int jj = 0; jj < nnei; jj++) { - ll[0] = em[ii * nnei * 4 + jj * 4 + 0]; - ll[1] = em[ii * nnei * 4 + jj * 4 + 1]; - ll[2] = em[ii * nnei * 4 + jj * 4 + 2]; - ll[3] = em[ii * nnei * 4 + jj * 4 + 3]; - hh[0] = dz_dy_dem[ii * nnei * 4 + jj * 4 + 0]; - hh[1] = dz_dy_dem[ii * nnei * 4 + jj * 4 + 1]; - hh[2] = dz_dy_dem[ii * nnei * 4 + jj * 4 + 2]; - hh[3] = dz_dy_dem[ii * nnei * 4 + jj * 4 + 3]; + const int em_base = ii * nnei * NDESCRPT + jj * NDESCRPT; + for (int mm = 0; mm < NDESCRPT; ++mm) { + ll[mm] = em[em_base + mm]; + hh[mm] = dz_dy_dem[em_base + mm]; + } FPTYPE xx = em_x[ii * nnei + jj]; FPTYPE dz_xx = dz_dy_dem_x[ii * nnei + jj]; if (ago == xx && ll[1] == 0. && ll[2] == 0. && ll[3] == 0. && is_sorted) { @@ -445,28 +433,10 @@ void deepmd::tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, * If `enable_se_atten` is true, `var` will be `var * t + var`, and * `var'` will be `(var_grad * t + var_grad) * dz_xx`. */ - if (unloop) { - dz_dy[ii * last_layer_size * 4 + 0 * last_layer_size + kk] += - (nnei - jj) * - (var * hh[0] + (dz_xx * var_grad + two_grad) * ll[0]); - dz_dy[ii * last_layer_size * 4 + 1 * last_layer_size + kk] += - (nnei - jj) * - (var * hh[1] + (dz_xx * var_grad + two_grad) * ll[1]); - dz_dy[ii * last_layer_size * 4 + 2 * last_layer_size + kk] += - (nnei - jj) * - (var * hh[2] + (dz_xx * var_grad + two_grad) * ll[2]); - dz_dy[ii * last_layer_size * 4 + 3 * last_layer_size + kk] += - (nnei - jj) * - (var * hh[3] + (dz_xx * var_grad + two_grad) * ll[3]); - } else { - dz_dy[ii * last_layer_size * 4 + 0 * last_layer_size + kk] += - var * hh[0] + (dz_xx * var_grad + two_grad) * ll[0]; - dz_dy[ii * last_layer_size * 4 + 1 * last_layer_size + kk] += - var * hh[1] + (dz_xx * var_grad + two_grad) * ll[1]; - dz_dy[ii * last_layer_size * 4 + 2 * last_layer_size + kk] += - var * hh[2] + (dz_xx * var_grad + two_grad) * ll[2]; - dz_dy[ii * last_layer_size * 4 + 3 * last_layer_size + kk] += - var * hh[3] + (dz_xx * var_grad + two_grad) * ll[3]; + const FPTYPE scale = unloop ? (FPTYPE)(nnei - jj) : (FPTYPE)1.; + for (int mm = 0; mm < NDESCRPT; ++mm) { + dz_dy[ii * last_layer_size * NDESCRPT + mm * last_layer_size + kk] += + scale * (var * hh[mm] + (dz_xx * var_grad + two_grad) * ll[mm]); } } if (unloop) { @@ -476,6 +446,125 @@ void deepmd::tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, } } +} // namespace + +template +void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted, + const int ndescrpt) { + deepmd::detail::check_se_a_basis_dimension(ndescrpt); + switch (ndescrpt) { + case 4: + tabulate_fusion_se_a_cpu_impl(out, table, table_info, em_x, em, + two_embed, nloc, nnei, + last_layer_size, is_sorted); + return; + case 9: + tabulate_fusion_se_a_cpu_impl(out, table, table_info, em_x, em, + two_embed, nloc, nnei, + last_layer_size, is_sorted); + return; + case 16: + tabulate_fusion_se_a_cpu_impl(out, table, table_info, em_x, + em, two_embed, nloc, nnei, + last_layer_size, is_sorted); + return; + case 25: + tabulate_fusion_se_a_cpu_impl(out, table, table_info, em_x, + em, two_embed, nloc, nnei, + last_layer_size, is_sorted); + return; + } +} + +template +void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, + FPTYPE* dy_dem, + FPTYPE* dy_dtwo, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const FPTYPE* dy, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted, + const int ndescrpt) { + deepmd::detail::check_se_a_basis_dimension(ndescrpt); + switch (ndescrpt) { + case 4: + tabulate_fusion_se_a_grad_cpu_impl( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + return; + case 9: + tabulate_fusion_se_a_grad_cpu_impl( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + return; + case 16: + tabulate_fusion_se_a_grad_cpu_impl( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + return; + case 25: + tabulate_fusion_se_a_grad_cpu_impl( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); + return; + } +} + +template +void deepmd::tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* two_embed, + const FPTYPE* dz_dy_dem_x, + const FPTYPE* dz_dy_dem, + const FPTYPE* dz_dy_dtwo, + const int nloc, + const int nnei, + const int last_layer_size, + const bool is_sorted, + const int ndescrpt) { + deepmd::detail::check_se_a_basis_dimension(ndescrpt); + switch (ndescrpt) { + case 4: + tabulate_fusion_se_a_grad_grad_cpu_impl( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + return; + case 9: + tabulate_fusion_se_a_grad_grad_cpu_impl( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + return; + case 16: + tabulate_fusion_se_a_grad_grad_cpu_impl( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + return; + case 25: + tabulate_fusion_se_a_grad_grad_cpu_impl( + dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + return; + } +} + template void deepmd::tabulate_fusion_se_t_cpu(FPTYPE* out, const FPTYPE* table, @@ -908,7 +997,8 @@ template void deepmd::tabulate_fusion_se_a_cpu(float* out, const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void deepmd::tabulate_fusion_se_a_cpu( double* out, const double* table, @@ -919,7 +1009,8 @@ template void deepmd::tabulate_fusion_se_a_cpu( const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void deepmd::tabulate_fusion_se_a_grad_cpu( float* dy_dem_x, float* dy_dem, @@ -933,7 +1024,8 @@ template void deepmd::tabulate_fusion_se_a_grad_cpu( const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void deepmd::tabulate_fusion_se_a_grad_cpu( double* dy_dem_x, double* dy_dem, @@ -947,7 +1039,8 @@ template void deepmd::tabulate_fusion_se_a_grad_cpu( const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void deepmd::tabulate_fusion_se_a_grad_grad_cpu( float* dz_dy, const float* table, @@ -961,7 +1054,8 @@ template void deepmd::tabulate_fusion_se_a_grad_grad_cpu( const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void deepmd::tabulate_fusion_se_a_grad_grad_cpu( double* dz_dy, const double* table, @@ -975,7 +1069,8 @@ template void deepmd::tabulate_fusion_se_a_grad_grad_cpu( const int nloc, const int nnei, const int last_layer_size, - const bool is_sorted); + const bool is_sorted, + const int ndescrpt); template void deepmd::tabulate_fusion_se_t_cpu( float* out, diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index a6560e4fec..5e1a46ec91 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -12,6 +12,9 @@ option( # only within its own subtree, which does not cover this target. if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) find_package(CUDAToolkit REQUIRED) + option(DEEPMD_ENABLE_DPA1_HIGH_LMAX + "Instantiate experimental DPA1 CUDA kernels for lmax greater than one" + OFF) if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) # CUDA 12.9 CCCL fails to compile CUB/Thrust with -arch=all. if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.9" AND CUDAToolkit_VERSION @@ -22,11 +25,16 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) endif() endif() enable_language(CUDA) + set(DPA1_GRAPH_COMPRESS_KERNEL_SRC + dpa1_graph_compress_c8.cu dpa1_graph_compress_c16.cu + dpa1_graph_compress_c32.cu dpa1_graph_compress_c64.cu + dpa1_graph_compress_c128.cu dpa1_graph_compress_c256.cu) list( APPEND OP_SRC dpa1_graph_descriptor.cu dpa1_graph_compress.cu + ${DPA1_GRAPH_COMPRESS_KERNEL_SRC} graph_fitting.cu edge_force_virial.cu dpa1_graph_energy_force.cu) @@ -56,7 +64,14 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) # libtorch headers require C++17; the CUDA sources must match. set_target_properties(deepmd_op_pt PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON) - set_source_files_properties(dpa1_graph_compress.cu + if(DEEPMD_ENABLE_DPA1_HIGH_LMAX) + target_compile_definitions(deepmd_op_pt + PRIVATE DEEPMD_ENABLE_DPA1_HIGH_LMAX=1) + endif() + # The compressed DPA1 kernels are instantiated one translation unit per + # channel width so their topology and angular-degree variants compile in + # parallel. + set_source_files_properties(${DPA1_GRAPH_COMPRESS_KERNEL_SRC} PROPERTIES COMPILE_OPTIONS "--use_fast_math") endif() if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI}) diff --git a/source/op/pt/dpa1_graph_common.cuh b/source/op/pt/dpa1_graph_common.cuh index c8fc3f8f01..880be334c5 100644 --- a/source/op/pt/dpa1_graph_common.cuh +++ b/source/op/pt/dpa1_graph_common.cuh @@ -119,48 +119,75 @@ DEV_INLINE float4 weight4(const float* p) { return __ldg(reinterpret_cast(p)); } -// Generic N-row float4 fragment load / store (N a multiple of 4). These back -// the edge-fragment loads of the backward kernels, whose per-thread edge count -// (EPT) is a template parameter; the forwards keep the fixed load8 / store8. +// Generic fragment load / store. Multiples of four use vector transactions; +// the two-edge low-shared-memory backward fallback uses scalar transactions. template DEV_INLINE void loadN(const float* p, float (&a)[N]) { + if constexpr (N % 4 == 0) { #pragma unroll - for (int i = 0; i < N; i += 4) { - const float4 v = *reinterpret_cast(p + i); - a[i] = v.x; - a[i + 1] = v.y; - a[i + 2] = v.z; - a[i + 3] = v.w; + for (int i = 0; i < N; i += 4) { + const float4 v = *reinterpret_cast(p + i); + a[i] = v.x; + a[i + 1] = v.y; + a[i + 2] = v.z; + a[i + 3] = v.w; + } + } else { +#pragma unroll + for (int i = 0; i < N; ++i) { + a[i] = p[i]; + } } } template DEV_INLINE void storeN(float* p, const float (&a)[N]) { + if constexpr (N % 4 == 0) { +#pragma unroll + for (int i = 0; i < N; i += 4) { + *reinterpret_cast(p + i) = + make_float4(a[i], a[i + 1], a[i + 2], a[i + 3]); + } + } else { #pragma unroll - for (int i = 0; i < N; i += 4) { - *reinterpret_cast(p + i) = - make_float4(a[i], a[i + 1], a[i + 2], a[i + 3]); + for (int i = 0; i < N; ++i) { + p[i] = a[i]; + } } } template DEV_INLINE void loadN_streaming(const float* p, float (&a)[N]) { + if constexpr (N % 4 == 0) { +#pragma unroll + for (int i = 0; i < N; i += 4) { + const float4 v = __ldcs(reinterpret_cast(p + i)); + a[i] = v.x; + a[i + 1] = v.y; + a[i + 2] = v.z; + a[i + 3] = v.w; + } + } else { #pragma unroll - for (int i = 0; i < N; i += 4) { - const float4 v = __ldcs(reinterpret_cast(p + i)); - a[i] = v.x; - a[i + 1] = v.y; - a[i + 2] = v.z; - a[i + 3] = v.w; + for (int i = 0; i < N; ++i) { + a[i] = __ldcs(p + i); + } } } template DEV_INLINE void storeN_streaming(float* p, const float (&a)[N]) { + if constexpr (N % 4 == 0) { #pragma unroll - for (int i = 0; i < N; i += 4) { - __stcs(reinterpret_cast(p + i), - make_float4(a[i], a[i + 1], a[i + 2], a[i + 3])); + for (int i = 0; i < N; i += 4) { + __stcs(reinterpret_cast(p + i), + make_float4(a[i], a[i + 1], a[i + 2], a[i + 3])); + } + } else { +#pragma unroll + for (int i = 0; i < N; ++i) { + __stcs(p + i, a[i]); + } } } @@ -216,28 +243,32 @@ __global__ void edge_order_scatter_kernel(long n_edge, // Per-edge staging: environment-matrix row, type-pair index, switch value // and center node. // ====================================================================== +template struct EdgeStage { - float r0, r1, r2, r3; // normalized environment-matrix row - float sw; // raw smooth-switch value (type-pair gate factor) + float basis[BASIS_DIM]; // normalized angular moment basis + float sw; // raw smooth-switch value (type-pair gate factor) int pair_idx; int dst; bool valid; }; -DEV_INLINE EdgeStage stage_edge(long e, - long n_edge, - int ntypes, - int one_side, - float rcut, - float rcut_smth, - float protection, - const float* __restrict__ edge_vec, - const long* __restrict__ edge_index, - const bool* __restrict__ edge_mask, - const long* __restrict__ atype, - const float* __restrict__ davg, - const float* __restrict__ inv_dstd) { - EdgeStage s; +template +DEV_INLINE EdgeStage stage_edge(long e, + long n_edge, + int ntypes, + int one_side, + float rcut, + float rcut_smth, + float protection, + const float* __restrict__ edge_vec, + const long* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const long* __restrict__ atype, + const float* __restrict__ davg, + const float* __restrict__ inv_dstd) { + static_assert(BASIS_DIM == 4 || BASIS_DIM == 9, + "Uncompressed DPA1 CUDA kernels support only lmax <= 2."); + EdgeStage s; const long src = edge_index[e]; const long dst = edge_index[n_edge + e]; s.dst = (int)dst; @@ -257,15 +288,107 @@ DEV_INLINE EdgeStage stage_edge(long e, const float t0 = sw * rq, iq2 = sw * rq * rq; const float* av = davg + (long)ct * 4; const float* isd = inv_dstd + (long)ct * 4; - s.r0 = (t0 - av[0]) * isd[0]; - s.r1 = (x * iq2 - av[1]) * isd[1]; - s.r2 = (y * iq2 - av[2]) * isd[2]; - s.r3 = (z * iq2 - av[3]) * isd[3]; + s.basis[0] = (t0 - av[0]) * isd[0]; + s.basis[1] = (x * iq2 - av[1]) * isd[1]; + s.basis[2] = (y * iq2 - av[2]) * isd[2]; + s.basis[3] = (z * iq2 - av[3]) * isd[3]; + if constexpr (BASIS_DIM > 4) { + const float vx = x * rq; + const float vy = y * rq; + const float vz = z * rq; + const float v2 = vx * vx + vy * vy + vz * vz; + const float radial = s.valid && len > 0.f ? t0 * isd[0] : 0.f; + constexpr float sqrt_three = 1.7320508075688772935f; + s.basis[4] = radial * sqrt_three * vx * vy; + s.basis[5] = radial * sqrt_three * vy * vz; + s.basis[6] = radial * 0.5f * (3.f * vz * vz - v2); + s.basis[7] = radial * sqrt_three * vx * vz; + s.basis[8] = radial * 0.5f * sqrt_three * (vx * vx - vy * vy); + } s.sw = sw; s.pair_idx = one_side ? nt : ct * ntypes + nt; return s; } +template +DEV_INLINE void moment_basis_edge_gradient(const float (&d_basis)[BASIS_DIM], + float d_radial, + float d_switch, + float x, + float y, + float z, + float len, + float protection, + float sw, + float dsw, + float inv_nnei, + const float* __restrict__ inv_dstd, + float* __restrict__ output) { + static_assert(BASIS_DIM == 4 || BASIS_DIM == 9, + "Uncompressed DPA1 CUDA kernels support only lmax <= 2."); + const float inv_len = len > 0.f ? 1.f / len : 0.f; + const float inv_q = 1.f / (len + protection); + const float g0 = (d_basis[0] * inv_nnei + d_radial) * inv_dstd[0]; + const float gx = d_basis[1] * inv_nnei * inv_dstd[1]; + const float gy = d_basis[2] * inv_nnei * inv_dstd[2]; + const float gz = d_basis[3] * inv_nnei * inv_dstd[3]; + const float directional = gx * x + gy * y + gz * z; + const float coefficient = + (g0 * inv_q * (dsw - sw * inv_q) + + directional * inv_q * inv_q * (dsw - 2.f * sw * inv_q) + + d_switch * dsw) * + inv_len; + const float vector_scale = sw * inv_q * inv_q; + float grad_x = coefficient * x + vector_scale * gx; + float grad_y = coefficient * y + vector_scale * gy; + float grad_z = coefficient * z + vector_scale * gz; + + if constexpr (BASIS_DIM == 9) { + const float nx = x * inv_len; + const float ny = y * inv_len; + const float nz = z * inv_len; + const float vx = x * inv_q; + const float vy = y * inv_q; + const float vz = z * inv_q; + const float v2 = vx * vx + vy * vy + vz * vz; + constexpr float sqrt_three = 1.7320508075688772935f; + const float y2[5] = { + sqrt_three * vx * vy, + sqrt_three * vy * vz, + 0.5f * (3.f * vz * vz - v2), + sqrt_three * vx * vz, + 0.5f * sqrt_three * (vx * vx - vy * vy), + }; + float radial_partial = 0.f; +#pragma unroll + for (int row = 0; row < 5; ++row) { + radial_partial = + fmaf(d_basis[4 + row] * inv_nnei, y2[row], radial_partial); + } + const float d4 = d_basis[4] * inv_nnei; + const float d5 = d_basis[5] * inv_nnei; + const float d6 = d_basis[6] * inv_nnei; + const float d7 = d_basis[7] * inv_nnei; + const float d8 = d_basis[8] * inv_nnei; + const float grad_vx = sqrt_three * (d4 * vy + d7 * vz + d8 * vx) - d6 * vx; + const float grad_vy = sqrt_three * (d4 * vx + d5 * vz - d8 * vy) - d6 * vy; + const float grad_vz = sqrt_three * (d5 * vy + d7 * vx) + 2.f * d6 * vz; + const float protected_dot = grad_vx * vx + grad_vy * vy + grad_vz * vz; + const float amplitude = sw * inv_q * inv_dstd[0]; + const float amplitude_grad = inv_dstd[0] * inv_q * (dsw - sw * inv_q); + grad_x += radial_partial * amplitude_grad * nx + + amplitude * inv_q * (grad_vx - protected_dot * nx); + grad_y += radial_partial * amplitude_grad * ny + + amplitude * inv_q * (grad_vy - protected_dot * ny); + grad_z += radial_partial * amplitude_grad * nz + + amplitude * inv_q * (grad_vz - protected_dot * nz); + } + + output[0] = grad_x; + output[1] = grad_y; + output[2] = grad_z; +} + // Parallel CSR run scan over one tile: row r is a run head iff // dst[r] != dst[r - 1]; per-warp head ballots give every row its run index // as a popcount prefix. All threads must call; barriers inside. @@ -278,7 +401,7 @@ DEV_INLINE void scan_runs(int tid, int* run_begin, int* run_of, int* n_runs) { - constexpr int TILE = NW * 32; // edges per tile + constexpr int TILE = NW * 32; const int r = tid; bool head = false; if (r < TILE) { @@ -329,11 +452,11 @@ DEV_INLINE void scan_runs(int tid, // width (edges per tile) is a template parameter so a forward can keep the // 128-edge tile while its backward uses a narrower tile (fewer edges per // thread, so the per-thread register footprint drops below the spill wall). -template +template struct EdgeTablesT { - float rr[TILE][4]; // env-mat row premultiplied by mask / nnei - float radial[TILE]; // rr0 (unmasked MLP / table input) - float sw[TILE]; // raw switch value (strip gate factor) + float basis[TILE][BASIS_DIM]; // moment basis premultiplied by mask / nnei + float radial[TILE]; // rr0 (unmasked MLP / table input) + float sw[TILE]; // raw switch value (strip gate factor) int pair_idx[TILE]; int dst[TILE]; int run_node[TILE]; @@ -343,7 +466,7 @@ struct EdgeTablesT { int n_runs; }; -template +template DEV_INLINE void stage_tile(int tid, long tile_base, int rows, @@ -361,19 +484,19 @@ DEV_INLINE void stage_tile(int tid, const float* __restrict__ davg, const float* __restrict__ inv_dstd, const int* __restrict__ order, - EdgeTablesT& T) { + EdgeTablesT& T) { if (tid < TILE) { if (tid < rows) { const int e = order[tile_base + tid]; - const auto s = - stage_edge(e, n_edge, ntypes, one_side, rcut, rcut_smth, protection, - edge_vec, edge_index, edge_mask, atype, davg, inv_dstd); + const auto s = stage_edge( + e, n_edge, ntypes, one_side, rcut, rcut_smth, protection, edge_vec, + edge_index, edge_mask, atype, davg, inv_dstd); const float mm = (s.valid ? 1.f : 0.f) * inv_nnei; - T.radial[tid] = s.r0; - T.rr[tid][0] = s.r0 * mm; - T.rr[tid][1] = s.r1 * mm; - T.rr[tid][2] = s.r2 * mm; - T.rr[tid][3] = s.r3 * mm; + T.radial[tid] = s.basis[0]; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + T.basis[tid][k] = s.basis[k] * mm; + } T.sw[tid] = s.sw; T.pair_idx[tid] = s.pair_idx; T.dst[tid] = s.dst; @@ -381,10 +504,10 @@ DEV_INLINE void stage_tile(int tid, // Tail rows of a partial tile: finite MLP input, zero moment weight, // sentinel dst excluded from every run. T.radial[tid] = 0.f; - T.rr[tid][0] = 0.f; - T.rr[tid][1] = 0.f; - T.rr[tid][2] = 0.f; - T.rr[tid][3] = 0.f; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + T.basis[tid][k] = 0.f; + } T.sw[tid] = 0.f; T.pair_idx[tid] = 0; T.dst[tid] = -1; @@ -415,35 +538,50 @@ DEV_INLINE float gate_factor(const float* __restrict__ gate_table, // rot_mat[n, i, :] = gr[n, 1:4, i] // plus the appended center type embedding when concat_tebd is set. // ====================================================================== +DEV_INLINE float moment_row_weight(int row, + const float* __restrict__ degree_gain_raw) { + if (row < 4) { + return 1.0f; + } + const float gain = __ldg(degree_gain_raw); + return gain * gain; +} + +template __global__ void gram_kernel(int n_node, int ng, int axis, int tebd_dim, int concat_tebd, const float* __restrict__ gr, + const float* __restrict__ degree_gain_raw, const float* __restrict__ type_embedding, const long* __restrict__ atype, float* __restrict__ grrg, float* __restrict__ rot_mat) { const int n = blockIdx.x; - extern __shared__ float s_gr[]; // (4, ng) - for (int t = threadIdx.x; t < 4 * ng; t += blockDim.x) { - s_gr[t] = gr[(long)n * 4 * ng + t]; + extern __shared__ float s_gr[]; // (BASIS_DIM, ng) + for (int t = threadIdx.x; t < BASIS_DIM * ng; t += blockDim.x) { + s_gr[t] = gr[(long)n * BASIS_DIM * ng + t]; } __syncthreads(); const int out_dim = ng * axis + (concat_tebd ? tebd_dim : 0); float* out = grrg + (long)n * out_dim; for (int i = threadIdx.x; i < ng; i += blockDim.x) { - const float g0 = s_gr[0 * ng + i], g1 = s_gr[1 * ng + i]; - const float g2 = s_gr[2 * ng + i], g3 = s_gr[3 * ng + i]; for (int j = 0; j < axis; ++j) { - out[i * axis + j] = g0 * s_gr[0 * ng + j] + g1 * s_gr[1 * ng + j] + - g2 * s_gr[2 * ng + j] + g3 * s_gr[3 * ng + j]; + float value = 0.f; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + const float weight = + BASIS_DIM == 4 ? 1.0f : moment_row_weight(k, degree_gain_raw); + value = fmaf(s_gr[k * ng + i] * weight, s_gr[k * ng + j], value); + } + out[i * axis + j] = value; } if (rot_mat != nullptr) { - rot_mat[((long)n * ng + i) * 3 + 0] = g1; - rot_mat[((long)n * ng + i) * 3 + 1] = g2; - rot_mat[((long)n * ng + i) * 3 + 2] = g3; + rot_mat[((long)n * ng + i) * 3 + 0] = s_gr[1 * ng + i]; + rot_mat[((long)n * ng + i) * 3 + 1] = s_gr[2 * ng + i]; + rot_mat[((long)n * ng + i) * 3 + 2] = s_gr[3 * ng + i]; } } if (concat_tebd) { @@ -460,6 +598,7 @@ __global__ void gram_kernel(int n_node, // + (k >= 1) R[c, k - 1]. // The concat tebd tail of d(grrg) is a constant feature and carries no // gradient; grrg_stride skips it. +template __global__ void gram_backward_kernel(int n_node, int ng, int axis, @@ -467,45 +606,52 @@ __global__ void gram_backward_kernel(int n_node, const float* __restrict__ d_grrg, const float* __restrict__ d_rot, const float* __restrict__ gr, + const float* __restrict__ degree_gain_raw, float* __restrict__ dgr) { const int n = blockIdx.x; - extern __shared__ float sm[]; // [0, 4*ng): gr; [4*ng, ...): d_grrg row + extern __shared__ float sm[]; float* s_gr = sm; - float* s_dg = sm + 4 * ng; - for (int t = threadIdx.x; t < 4 * ng; t += blockDim.x) { - s_gr[t] = gr[(long)n * 4 * ng + t]; + float* s_dg = sm + BASIS_DIM * ng; + for (int t = threadIdx.x; t < BASIS_DIM * ng; t += blockDim.x) { + s_gr[t] = gr[(long)n * BASIS_DIM * ng + t]; } for (int t = threadIdx.x; t < ng * axis; t += blockDim.x) { s_dg[t] = d_grrg[(long)n * grrg_stride + t]; } __syncthreads(); for (int c = threadIdx.x; c < ng; c += blockDim.x) { - float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f; + float acc[BASIS_DIM] = {}; for (int j = 0; j < axis; ++j) { const float d = s_dg[c * axis + j]; - acc0 = fmaf(d, s_gr[0 * ng + j], acc0); - acc1 = fmaf(d, s_gr[1 * ng + j], acc1); - acc2 = fmaf(d, s_gr[2 * ng + j], acc2); - acc3 = fmaf(d, s_gr[3 * ng + j], acc3); +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + acc[k] = fmaf(d, s_gr[k * ng + j], acc[k]); + } } if (c < axis) { for (int i = 0; i < ng; ++i) { const float d = s_dg[i * axis + c]; - acc0 = fmaf(d, s_gr[0 * ng + i], acc0); - acc1 = fmaf(d, s_gr[1 * ng + i], acc1); - acc2 = fmaf(d, s_gr[2 * ng + i], acc2); - acc3 = fmaf(d, s_gr[3 * ng + i], acc3); +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + acc[k] = fmaf(d, s_gr[k * ng + i], acc[k]); + } + } + } + if constexpr (BASIS_DIM > 4) { +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + acc[k] *= moment_row_weight(k, degree_gain_raw); } } if (d_rot) { - acc1 += d_rot[((long)n * ng + c) * 3 + 0]; - acc2 += d_rot[((long)n * ng + c) * 3 + 1]; - acc3 += d_rot[((long)n * ng + c) * 3 + 2]; + acc[1] += d_rot[((long)n * ng + c) * 3 + 0]; + acc[2] += d_rot[((long)n * ng + c) * 3 + 1]; + acc[3] += d_rot[((long)n * ng + c) * 3 + 2]; + } +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dgr[((long)n * BASIS_DIM + k) * ng + c] = acc[k]; } - dgr[((long)n * 4 + 0) * ng + c] = acc0; - dgr[((long)n * 4 + 1) * ng + c] = acc1; - dgr[((long)n * 4 + 2) * ng + c] = acc2; - dgr[((long)n * 4 + 3) * ng + c] = acc3; } } diff --git a/source/op/pt/dpa1_graph_compress.cu b/source/op/pt/dpa1_graph_compress.cu index c06009b2a4..d41c833092 100644 --- a/source/op/pt/dpa1_graph_compress.cu +++ b/source/op/pt/dpa1_graph_compress.cu @@ -1,23 +1,10 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // -// Geometrically compressed DPA1 graph descriptor over destination CSR edges. +// Torch bindings and runtime dispatch of the compressed DPA1 CUDA descriptor. // -// A warp owns one center node. Widths from 16 through 64 use two 16-lane -// sub-warps on alternating edges; each lane evaluates one or more spline -// channels. Wider tables retain one edge per warp to bound register pressure. -// The node moment and its Gram contraction remain in the same kernel. -// -// The backward recomputes the inexpensive spline value/derivative, contracts -// the descriptor gradient into the four environment channels, and writes each -// edge gradient exactly once. It is inference-oriented (one backward); the -// registered Python autograd bridge continues to expose the edge-vector -// gradient for the level-1 graph path. -// -// Every specialization has balanced (two CTA/SM launch bound) and occupancy -// (four CTA/SM launch bound) resource variants. The first uncaptured call times -// 128- and 256-thread launches on a bounded node sample and caches the selected -// variant per device and workload class. Device-family defaults remain valid -// when timing is disabled or CUDA Graph capture is active. +// The CUDA kernels are instantiated in one translation unit per channel width. +// This translation unit only validates tensors, builds the plain launch +// argument bundle, dispatches by width, and registers the operators. #include #include @@ -27,1308 +14,117 @@ #include #include -#include "dpa1_graph_compress_tuning.h" - -namespace { - -using deepmd::dpa1_compress_tuning::device_properties; -using deepmd::dpa1_compress_tuning::KernelDirection; -using deepmd::dpa1_compress_tuning::LaunchConfig; -using deepmd::dpa1_compress_tuning::ResourcePolicy; -using deepmd::dpa1_compress_tuning::select_launch_config; -using deepmd::dpa1_compress_tuning::TuningKey; -using deepmd::dpa1_compress_tuning::type_count_class; -using deepmd::dpa1_compress_tuning::workload_degree_class; -using deepmd::dpa1_compress_tuning::workload_size_class; - -constexpr int kThreads = 256; -constexpr int kWarpSize = 32; - -#define COMPRESS_CHECK_LAUNCH(name) \ - do { \ - const cudaError_t error = cudaGetLastError(); \ - TORCH_CHECK(error == cudaSuccess, name, ": ", cudaGetErrorString(error)); \ - } while (0) - -__device__ __forceinline__ float switch_value(float radius, - float lower, - float upper) { - const float coordinate = - __fdividef(fminf(fmaxf(radius, lower), upper) - lower, upper - lower); - const float square = coordinate * coordinate; - return square * coordinate * (-6.0f * square + 15.0f * coordinate - 10.0f) + - 1.0f; -} - -__device__ __forceinline__ float switch_derivative(float radius, - float lower, - float upper) { - if (radius <= lower || radius >= upper) { - return 0.0f; - } - const float coordinate = __fdividef(radius - lower, upper - lower); - const float square = coordinate * coordinate; - return __fdividef( - -30.0f * square * square + 60.0f * square * coordinate - 30.0f * square, - upper - lower); -} - -struct TableLocation { - int index; - float coordinate; - float extrapolation; -}; - -__device__ __forceinline__ int high_tail_index( - float lower, float upper, float table_max, float stride0, float stride1) { - const float boundary = nextafterf(table_max, lower); - const int first_stride = static_cast(__fdividef(upper - lower, stride0)); - return first_stride + static_cast(__fdividef(boundary - upper, stride1)); -} - -__device__ __forceinline__ TableLocation locate_table(float radial, - float lower, - float upper, - float table_max, - float stride0, - float stride1) { - TableLocation location; - location.coordinate = radial; - location.extrapolation = 0.0f; - if (radial < lower) { - location.index = 0; - location.coordinate = 0.0f; - location.extrapolation = radial - lower; - } else if (radial < upper) { - location.index = static_cast(__fdividef(radial - lower, stride0)); - location.coordinate -= location.index * stride0 + lower; - } else if (radial < table_max) { - const int first_stride = - static_cast(__fdividef(upper - lower, stride0)); - location.index = - first_stride + static_cast(__fdividef(radial - upper, stride1)); - location.coordinate -= (location.index - first_stride) * stride1 + upper; - } else { - const int first_stride = - static_cast(__fdividef(upper - lower, stride0)); - location.index = high_tail_index(lower, upper, table_max, stride0, stride1); - location.coordinate = - table_max - ((location.index - first_stride) * stride1 + upper); - location.extrapolation = radial - table_max; - } - return location; -} - -__device__ __forceinline__ void load_coefficients(const float* table, - const TableLocation& location, - int channel, - int width, - float2& c01, - float2& c23, - float2& c45) { - const long offset = static_cast(location.index) * width * 6 + - static_cast(channel) * 6; - c01 = __ldg(reinterpret_cast(table + offset)); - c23 = __ldg(reinterpret_cast(table + offset + 2)); - c45 = __ldg(reinterpret_cast(table + offset + 4)); -} - -__device__ __forceinline__ float evaluate_table_forward( - const float* table, const TableLocation& location, int channel, int width) { - float2 c01, c23, c45; - load_coefficients(table, location, channel, width, c01, c23, c45); - const float value = - c01.x + (c01.y + (c23.x + (c23.y + (c45.x + c45.y * location.coordinate) * - location.coordinate) * - location.coordinate) * - location.coordinate) * - location.coordinate; - if (location.extrapolation == 0.0f) { - return value; - } - const float derivative = - c01.y + - (2.0f * c23.x + - (3.0f * c23.y + (4.0f * c45.x + 5.0f * c45.y * location.coordinate) * - location.coordinate) * - location.coordinate) * - location.coordinate; - return value + derivative * location.extrapolation; -} - -__device__ __forceinline__ float2 evaluate_table_backward( - const float* table, const TableLocation& location, int channel, int width) { - float2 c01, c23, c45; - load_coefficients(table, location, channel, width, c01, c23, c45); - float value = c45.y; - float derivative = 0.0f; - derivative = fmaf(derivative, location.coordinate, value); - value = fmaf(value, location.coordinate, c45.x); - derivative = fmaf(derivative, location.coordinate, value); - value = fmaf(value, location.coordinate, c23.y); - derivative = fmaf(derivative, location.coordinate, value); - value = fmaf(value, location.coordinate, c23.x); - derivative = fmaf(derivative, location.coordinate, value); - value = fmaf(value, location.coordinate, c01.y); - derivative = fmaf(derivative, location.coordinate, value); - value = fmaf(value, location.coordinate, c01.x); - return make_float2(value + derivative * location.extrapolation, derivative); -} - -struct EdgeEnvironment { - float radial; - float r0; - float r1; - float r2; - float r3; - float switch_factor; - float x; - float y; - float z; - float radius; - int pair_index; -}; - -template -__device__ __forceinline__ EdgeEnvironment -load_environment(long edge, - int center_type, - int ntypes, - bool one_side, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - const float* edge_vec, - const index_t* edge_index, - const long* atype, - const float* average, - const float* inverse_stddev) { - EdgeEnvironment environment; - const long source = static_cast(edge_index[edge]); - const int neighbor_type = static_cast(atype[source]); - environment.x = edge_vec[edge * 3 + 0]; - environment.y = edge_vec[edge * 3 + 1]; - environment.z = edge_vec[edge * 3 + 2]; - const float square_length = environment.x * environment.x + - environment.y * environment.y + - environment.z * environment.z; - environment.radius = - square_length > 0.0f ? square_length * rsqrtf(square_length) : 0.0f; - const float denominator = environment.radius + protection; - environment.switch_factor = - switch_value(environment.radius, rcut_smooth, rcut); - const float inverse_radius = __fdividef(1.0f, denominator); - const float radial_scale = - environment.switch_factor * inverse_radius * inverse_radius; - const float* center_average = average + static_cast(center_type) * 4; - const float* center_inverse_stddev = - inverse_stddev + static_cast(center_type) * 4; - environment.radial = - (environment.switch_factor * inverse_radius - center_average[0]) * - center_inverse_stddev[0]; - environment.r0 = environment.radial * inverse_neighbors; - environment.r1 = (environment.x * radial_scale - center_average[1]) * - center_inverse_stddev[1] * inverse_neighbors; - environment.r2 = (environment.y * radial_scale - center_average[2]) * - center_inverse_stddev[2] * inverse_neighbors; - environment.r3 = (environment.z * radial_scale - center_average[3]) * - center_inverse_stddev[3] * inverse_neighbors; - environment.pair_index = - one_side ? neighbor_type : center_type * ntypes + neighbor_type; - return environment; -} - -__device__ __forceinline__ EdgeEnvironment -broadcast_environment(EdgeEnvironment value, int source_lane, unsigned mask) { - value.radial = __shfl_sync(mask, value.radial, source_lane); - value.r0 = __shfl_sync(mask, value.r0, source_lane); - value.r1 = __shfl_sync(mask, value.r1, source_lane); - value.r2 = __shfl_sync(mask, value.r2, source_lane); - value.r3 = __shfl_sync(mask, value.r3, source_lane); - value.switch_factor = __shfl_sync(mask, value.switch_factor, source_lane); - value.x = __shfl_sync(mask, value.x, source_lane); - value.y = __shfl_sync(mask, value.y, source_lane); - value.z = __shfl_sync(mask, value.z, source_lane); - value.radius = __shfl_sync(mask, value.radius, source_lane); - value.pair_index = __shfl_sync(mask, value.pair_index, source_lane); - return value; -} - -__device__ __forceinline__ TableLocation broadcast_location(TableLocation value, - int source_lane, - unsigned mask) { - value.index = __shfl_sync(mask, value.index, source_lane); - value.coordinate = __shfl_sync(mask, value.coordinate, source_lane); - value.extrapolation = __shfl_sync(mask, value.extrapolation, source_lane); - return value; -} - -__device__ __forceinline__ void store_edge_gradient( - long edge, - const EdgeEnvironment& environment, - float partial0, - float partial1, - float partial2, - float partial3, - float partial_radial, - float partial_switch, - float inverse_neighbors, - float inverse_stddev0, - float inverse_stddev1, - float inverse_stddev2, - float inverse_stddev3, - float rcut, - float rcut_smooth, - float protection, - float* edge_gradient) { - const float inverse_denominator = - __fdividef(1.0f, environment.radius + protection); - const float inverse_length = - environment.radius > 0.0f ? __fdividef(1.0f, environment.radius) : 0.0f; - const float switch_gradient = - switch_derivative(environment.radius, rcut_smooth, rcut); - const float gradient_radial = - (partial0 * inverse_neighbors + partial_radial) * inverse_stddev0; - const float gradient_x = partial1 * inverse_neighbors * inverse_stddev1; - const float gradient_y = partial2 * inverse_neighbors * inverse_stddev2; - const float gradient_z = partial3 * inverse_neighbors * inverse_stddev3; - const float directional = gradient_x * environment.x + - gradient_y * environment.y + - gradient_z * environment.z; - const float coefficient = - (gradient_radial * inverse_denominator * - (switch_gradient - environment.switch_factor * inverse_denominator) + - directional * inverse_denominator * inverse_denominator * - (switch_gradient - - 2.0f * environment.switch_factor * inverse_denominator) + - partial_switch * switch_gradient) * - inverse_length; - const float vector_scale = - environment.switch_factor * inverse_denominator * inverse_denominator; - edge_gradient[edge * 3 + 0] = - coefficient * environment.x + vector_scale * gradient_x; - edge_gradient[edge * 3 + 1] = - coefficient * environment.y + vector_scale * gradient_y; - edge_gradient[edge * 3 + 2] = - coefficient * environment.z + vector_scale * gradient_z; -} - -template -struct ChannelPolicy { - static constexpr bool use_half_warp = Width >= 16 && Width <= 64; - static constexpr int accumulation_groups = - use_half_warp ? Width / 16 : (Width + 31) / 32; - static constexpr int gradient_groups = (Width + 31) / 32; -}; - -template -__device__ __forceinline__ long edge_at_csr_position( - long position, const index_t* destination_order) { - if constexpr (Canonical) { - return position; - } - return static_cast(destination_order[position]); -} - -template -__device__ __forceinline__ bool edge_is_active(long edge, - const bool* edge_mask) { - if constexpr (Masked) { - return edge_mask[edge]; - } - return true; -} - -template -__global__ -__launch_bounds__(kThreads, MinimumBlocks) void compressed_forward_kernel( - long node_count, - int ntypes, - bool one_side, - bool smooth, - int axis, - bool concatenate_type_embedding, - bool write_rotation, - int type_embedding_dim, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const float* __restrict__ edge_vec, - const index_t* __restrict__ edge_index, - const bool* __restrict__ edge_mask, - const index_t* __restrict__ destination_order, - const long* __restrict__ destination_row_ptr, - const long* __restrict__ atype, - const float* __restrict__ type_embedding, - const float* __restrict__ average, - const float* __restrict__ inverse_stddev, - const float* __restrict__ table, - const float* __restrict__ gate_table, - float* __restrict__ descriptor, - float* __restrict__ rotation, - float* __restrict__ moment) { - constexpr unsigned kWarpMask = 0xffffffffu; - constexpr int kGroups = ChannelPolicy::accumulation_groups; - const int lane = threadIdx.x & 31; - const int warp = threadIdx.x >> 5; - const int warps_per_block = blockDim.x / kWarpSize; - const long node = static_cast(blockIdx.x) * warps_per_block + warp; - if (node >= node_count) { - return; - } - - float accumulator0[kGroups] = {}; - float accumulator1[kGroups] = {}; - float accumulator2[kGroups] = {}; - float accumulator3[kGroups] = {}; - int center_type = lane == 0 ? static_cast(atype[node]) : 0; - center_type = __shfl_sync(kWarpMask, center_type, 0); - const long begin = destination_row_ptr[node]; - const long end = destination_row_ptr[node + 1]; - - if constexpr (ChannelPolicy::use_half_warp) { - const int half = lane >> 4; - const int half_lane = lane & 15; - const int leader = half * 16; - const unsigned half_mask = half == 0 ? 0x0000ffffu : 0xffff0000u; - for (long position = begin + half; position < end; position += 2) { - const long edge = - edge_at_csr_position(position, destination_order); - if (!edge_is_active(edge, edge_mask)) { - continue; - } - EdgeEnvironment environment{}; - TableLocation location{}; - if (half_lane == 0) { - environment = load_environment(edge, center_type, ntypes, one_side, - rcut, rcut_smooth, protection, - inverse_neighbors, edge_vec, edge_index, - atype, average, inverse_stddev); - location = locate_table(environment.radial, lower, upper, table_max, - stride0, stride1); - } - environment = broadcast_environment(environment, leader, half_mask); - location = broadcast_location(location, leader, half_mask); -#pragma unroll - for (int group = 0; group < kGroups; ++group) { - const int channel = group * 16 + half_lane; - const float table_value = - evaluate_table_forward(table, location, channel, Width); - const float gate = - __ldg(gate_table + - static_cast(environment.pair_index) * Width + channel); - const float effective_gate = - smooth ? gate * environment.switch_factor : gate; - const float embedding = table_value * (1.0f + effective_gate); - accumulator0[group] = - fmaf(environment.r0, embedding, accumulator0[group]); - accumulator1[group] = - fmaf(environment.r1, embedding, accumulator1[group]); - accumulator2[group] = - fmaf(environment.r2, embedding, accumulator2[group]); - accumulator3[group] = - fmaf(environment.r3, embedding, accumulator3[group]); - } - } -#pragma unroll - for (int group = 0; group < kGroups; ++group) { - accumulator0[group] += - __shfl_xor_sync(kWarpMask, accumulator0[group], 16); - accumulator1[group] += - __shfl_xor_sync(kWarpMask, accumulator1[group], 16); - accumulator2[group] += - __shfl_xor_sync(kWarpMask, accumulator2[group], 16); - accumulator3[group] += - __shfl_xor_sync(kWarpMask, accumulator3[group], 16); - } - } else { - for (long position = begin; position < end; ++position) { - const long edge = - edge_at_csr_position(position, destination_order); - if (!edge_is_active(edge, edge_mask)) { - continue; - } - EdgeEnvironment environment{}; - TableLocation location{}; - if (lane == 0) { - environment = load_environment(edge, center_type, ntypes, one_side, - rcut, rcut_smooth, protection, - inverse_neighbors, edge_vec, edge_index, - atype, average, inverse_stddev); - location = locate_table(environment.radial, lower, upper, table_max, - stride0, stride1); - } - environment = broadcast_environment(environment, 0, kWarpMask); - location = broadcast_location(location, 0, kWarpMask); -#pragma unroll - for (int group = 0; group < kGroups; ++group) { - const int channel = group * 32 + lane; - if (channel < Width) { - const float table_value = - evaluate_table_forward(table, location, channel, Width); - const float gate = __ldg( - gate_table + static_cast(environment.pair_index) * Width + - channel); - const float effective_gate = - smooth ? gate * environment.switch_factor : gate; - const float embedding = table_value * (1.0f + effective_gate); - accumulator0[group] = - fmaf(environment.r0, embedding, accumulator0[group]); - accumulator1[group] = - fmaf(environment.r1, embedding, accumulator1[group]); - accumulator2[group] = - fmaf(environment.r2, embedding, accumulator2[group]); - accumulator3[group] = - fmaf(environment.r3, embedding, accumulator3[group]); - } - } - } - } - - const long moment_base = node * 4 * Width; - if constexpr (ChannelPolicy::use_half_warp) { - if (lane < 16) { -#pragma unroll - for (int group = 0; group < kGroups; ++group) { - const int channel = group * 16 + lane; - moment[moment_base + 0 * Width + channel] = accumulator0[group]; - moment[moment_base + 1 * Width + channel] = accumulator1[group]; - moment[moment_base + 2 * Width + channel] = accumulator2[group]; - moment[moment_base + 3 * Width + channel] = accumulator3[group]; - } - } - } else { -#pragma unroll - for (int group = 0; group < kGroups; ++group) { - const int channel = group * 32 + lane; - if (channel < Width) { - moment[moment_base + 0 * Width + channel] = accumulator0[group]; - moment[moment_base + 1 * Width + channel] = accumulator1[group]; - moment[moment_base + 2 * Width + channel] = accumulator2[group]; - moment[moment_base + 3 * Width + channel] = accumulator3[group]; - } - } - } - - const int output_dim = - Width * axis + (concatenate_type_embedding ? type_embedding_dim : 0); - float* output = descriptor + node * output_dim; - if constexpr (ChannelPolicy::use_half_warp) { -#pragma unroll - for (int group = 0; group < kGroups; ++group) { - const int channel = group * 16 + (lane & 15); - for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { - const float axis0 = - __shfl_sync(kWarpMask, accumulator0[0], axis_channel); - const float axis1 = - __shfl_sync(kWarpMask, accumulator1[0], axis_channel); - const float axis2 = - __shfl_sync(kWarpMask, accumulator2[0], axis_channel); - const float axis3 = - __shfl_sync(kWarpMask, accumulator3[0], axis_channel); - if (lane < 16) { - output[channel * axis + axis_channel] = - accumulator0[group] * axis0 + accumulator1[group] * axis1 + - accumulator2[group] * axis2 + accumulator3[group] * axis3; - } - } - if (write_rotation && lane < 16) { - float* rotation_row = rotation + (node * Width + channel) * 3; - rotation_row[0] = accumulator1[group]; - rotation_row[1] = accumulator2[group]; - rotation_row[2] = accumulator3[group]; - } - } - } else { -#pragma unroll - for (int group = 0; group < kGroups; ++group) { - const int channel = group * 32 + lane; - for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { - const float axis0 = - __shfl_sync(kWarpMask, accumulator0[0], axis_channel); - const float axis1 = - __shfl_sync(kWarpMask, accumulator1[0], axis_channel); - const float axis2 = - __shfl_sync(kWarpMask, accumulator2[0], axis_channel); - const float axis3 = - __shfl_sync(kWarpMask, accumulator3[0], axis_channel); - if (channel < Width) { - output[channel * axis + axis_channel] = - accumulator0[group] * axis0 + accumulator1[group] * axis1 + - accumulator2[group] * axis2 + accumulator3[group] * axis3; - } - } - if (write_rotation && channel < Width) { - float* rotation_row = rotation + (node * Width + channel) * 3; - rotation_row[0] = accumulator1[group]; - rotation_row[1] = accumulator2[group]; - rotation_row[2] = accumulator3[group]; - } - } - } - if (concatenate_type_embedding) { - for (int channel = lane; channel < type_embedding_dim; channel += 32) { - output[Width * axis + channel] = - type_embedding[static_cast(center_type) * type_embedding_dim + - channel]; - } - } -} - -template -__global__ -__launch_bounds__(kThreads, MinimumBlocks) void compressed_backward_kernel( - long node_count, - long edge_count, - int ntypes, - bool one_side, - bool smooth, - int axis, - int descriptor_stride, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const float* __restrict__ descriptor_gradient, - const float* __restrict__ rotation_gradient, - const float* __restrict__ moment, - const float* __restrict__ edge_vec, - const index_t* __restrict__ edge_index, - const bool* __restrict__ edge_mask, - const index_t* __restrict__ destination_order, - const long* __restrict__ destination_row_ptr, - const long* __restrict__ atype, - const float* __restrict__ average, - const float* __restrict__ inverse_stddev, - const float* __restrict__ table, - const float* __restrict__ gate_table, - float* __restrict__ edge_gradient) { - constexpr unsigned kWarpMask = 0xffffffffu; - constexpr int kGradientGroups = ChannelPolicy::gradient_groups; - constexpr int kEdgeGroups = ChannelPolicy::accumulation_groups; - const int lane = threadIdx.x & 31; - const int warp = threadIdx.x >> 5; - const int warps_per_block = blockDim.x / kWarpSize; - const long node = static_cast(blockIdx.x) * warps_per_block + warp; - if (node >= node_count) { - return; - } - - float gradient0[kGradientGroups] = {}; - float gradient1[kGradientGroups] = {}; - float gradient2[kGradientGroups] = {}; - float gradient3[kGradientGroups] = {}; - float own_moment0[kGradientGroups] = {}; - float own_moment1[kGradientGroups] = {}; - float own_moment2[kGradientGroups] = {}; - float own_moment3[kGradientGroups] = {}; - const long moment_base = node * 4 * Width; - const float* node_descriptor_gradient = - descriptor_gradient + node * descriptor_stride; - -#pragma unroll - for (int group = 0; group < kGradientGroups; ++group) { - const int channel = group * 32 + lane; - if (channel < Width) { - own_moment0[group] = __ldg(moment + moment_base + 0 * Width + channel); - own_moment1[group] = __ldg(moment + moment_base + 1 * Width + channel); - own_moment2[group] = __ldg(moment + moment_base + 2 * Width + channel); - own_moment3[group] = __ldg(moment + moment_base + 3 * Width + channel); - } - for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { - const float axis0 = __shfl_sync(kWarpMask, own_moment0[0], axis_channel); - const float axis1 = __shfl_sync(kWarpMask, own_moment1[0], axis_channel); - const float axis2 = __shfl_sync(kWarpMask, own_moment2[0], axis_channel); - const float axis3 = __shfl_sync(kWarpMask, own_moment3[0], axis_channel); - if (channel < Width) { - const float value = - __ldg(node_descriptor_gradient + channel * axis + axis_channel); - gradient0[group] = fmaf(value, axis0, gradient0[group]); - gradient1[group] = fmaf(value, axis1, gradient1[group]); - gradient2[group] = fmaf(value, axis2, gradient2[group]); - gradient3[group] = fmaf(value, axis3, gradient3[group]); - } - } - if (channel < axis) { - for (int input = 0; input < Width; ++input) { - const float value = - __ldg(node_descriptor_gradient + input * axis + channel); - gradient0[group] = - fmaf(value, __ldg(moment + moment_base + 0 * Width + input), - gradient0[group]); - gradient1[group] = - fmaf(value, __ldg(moment + moment_base + 1 * Width + input), - gradient1[group]); - gradient2[group] = - fmaf(value, __ldg(moment + moment_base + 2 * Width + input), - gradient2[group]); - gradient3[group] = - fmaf(value, __ldg(moment + moment_base + 3 * Width + input), - gradient3[group]); - } - } - if (rotation_gradient != nullptr && channel < Width) { - const long rotation_offset = (node * Width + channel) * 3; - gradient1[group] += __ldg(rotation_gradient + rotation_offset + 0); - gradient2[group] += __ldg(rotation_gradient + rotation_offset + 1); - gradient3[group] += __ldg(rotation_gradient + rotation_offset + 2); - } - } - - int center_type = lane == 0 ? static_cast(atype[node]) : 0; - center_type = __shfl_sync(kWarpMask, center_type, 0); - const float inverse_stddev0 = - __ldg(inverse_stddev + static_cast(center_type) * 4 + 0); - const float inverse_stddev1 = - __ldg(inverse_stddev + static_cast(center_type) * 4 + 1); - const float inverse_stddev2 = - __ldg(inverse_stddev + static_cast(center_type) * 4 + 2); - const float inverse_stddev3 = - __ldg(inverse_stddev + static_cast(center_type) * 4 + 3); - const long begin = destination_row_ptr[node]; - const long end = destination_row_ptr[node + 1]; +#include "dpa1_graph_compress_launch.h" - if constexpr (ChannelPolicy::use_half_warp) { - const int half = lane >> 4; - const int half_lane = lane & 15; - const int leader = half * 16; - const unsigned half_mask = half == 0 ? 0x0000ffffu : 0xffff0000u; - float edge_gradient0[kEdgeGroups] = {}; - float edge_gradient1[kEdgeGroups] = {}; - float edge_gradient2[kEdgeGroups] = {}; - float edge_gradient3[kEdgeGroups] = {}; -#pragma unroll - for (int group = 0; group < kEdgeGroups; ++group) { - const int channel = group * 16 + half_lane; - const int owner_group = channel / 32; - const int owner_lane = channel & 31; - edge_gradient0[group] = - __shfl_sync(kWarpMask, gradient0[owner_group], owner_lane); - edge_gradient1[group] = - __shfl_sync(kWarpMask, gradient1[owner_group], owner_lane); - edge_gradient2[group] = - __shfl_sync(kWarpMask, gradient2[owner_group], owner_lane); - edge_gradient3[group] = - __shfl_sync(kWarpMask, gradient3[owner_group], owner_lane); - } +#ifndef DEEPMD_ENABLE_DPA1_HIGH_LMAX +#define DEEPMD_ENABLE_DPA1_HIGH_LMAX 0 +#endif - for (long position = begin + half; position < end; position += 2) { - const long edge = - edge_at_csr_position(position, destination_order); - if (!edge_is_active(edge, edge_mask)) { - if (half_lane == 0) { - edge_gradient[edge * 3 + 0] = 0.0f; - edge_gradient[edge * 3 + 1] = 0.0f; - edge_gradient[edge * 3 + 2] = 0.0f; - } - continue; - } - EdgeEnvironment environment{}; - TableLocation location{}; - if (half_lane == 0) { - environment = load_environment(edge, center_type, ntypes, one_side, - rcut, rcut_smooth, protection, - inverse_neighbors, edge_vec, edge_index, - atype, average, inverse_stddev); - location = locate_table(environment.radial, lower, upper, table_max, - stride0, stride1); - } - environment = broadcast_environment(environment, leader, half_mask); - location = broadcast_location(location, leader, half_mask); - float partial0 = 0.0f; - float partial1 = 0.0f; - float partial2 = 0.0f; - float partial3 = 0.0f; - float partial_radial = 0.0f; - float partial_switch = 0.0f; -#pragma unroll - for (int group = 0; group < kEdgeGroups; ++group) { - const int channel = group * 16 + half_lane; - const float d0 = edge_gradient0[group]; - const float d1 = edge_gradient1[group]; - const float d2 = edge_gradient2[group]; - const float d3 = edge_gradient3[group]; - const float descriptor_product = - environment.r0 * d0 + environment.r1 * d1 + environment.r2 * d2 + - environment.r3 * d3; - const float2 table_value = - evaluate_table_backward(table, location, channel, Width); - const float gate = - __ldg(gate_table + - static_cast(environment.pair_index) * Width + channel); - const float effective_gate = - smooth ? gate * environment.switch_factor : gate; - const float embedding = table_value.x * (1.0f + effective_gate); - if (smooth) { - partial_switch = - fmaf(descriptor_product * table_value.x, gate, partial_switch); - } - partial_radial = fmaf(descriptor_product * (1.0f + effective_gate), - table_value.y, partial_radial); - partial0 = fmaf(embedding, d0, partial0); - partial1 = fmaf(embedding, d1, partial1); - partial2 = fmaf(embedding, d2, partial2); - partial3 = fmaf(embedding, d3, partial3); - } -#pragma unroll - for (int offset = 8; offset > 0; offset >>= 1) { - partial0 += __shfl_down_sync(half_mask, partial0, offset, 16); - partial1 += __shfl_down_sync(half_mask, partial1, offset, 16); - partial2 += __shfl_down_sync(half_mask, partial2, offset, 16); - partial3 += __shfl_down_sync(half_mask, partial3, offset, 16); - partial_radial += - __shfl_down_sync(half_mask, partial_radial, offset, 16); - partial_switch += - __shfl_down_sync(half_mask, partial_switch, offset, 16); - } - if (half_lane == 0) { - store_edge_gradient(edge, environment, partial0, partial1, partial2, - partial3, partial_radial, partial_switch, - inverse_neighbors, inverse_stddev0, inverse_stddev1, - inverse_stddev2, inverse_stddev3, rcut, rcut_smooth, - protection, edge_gradient); - } - } - } else { - for (long position = begin; position < end; ++position) { - const long edge = - edge_at_csr_position(position, destination_order); - if (!edge_is_active(edge, edge_mask)) { - if (lane == 0) { - edge_gradient[edge * 3 + 0] = 0.0f; - edge_gradient[edge * 3 + 1] = 0.0f; - edge_gradient[edge * 3 + 2] = 0.0f; - } - continue; - } - EdgeEnvironment environment{}; - TableLocation location{}; - if (lane == 0) { - environment = load_environment(edge, center_type, ntypes, one_side, - rcut, rcut_smooth, protection, - inverse_neighbors, edge_vec, edge_index, - atype, average, inverse_stddev); - location = locate_table(environment.radial, lower, upper, table_max, - stride0, stride1); - } - environment = broadcast_environment(environment, 0, kWarpMask); - location = broadcast_location(location, 0, kWarpMask); - float partial0 = 0.0f; - float partial1 = 0.0f; - float partial2 = 0.0f; - float partial3 = 0.0f; - float partial_radial = 0.0f; - float partial_switch = 0.0f; -#pragma unroll - for (int group = 0; group < kGradientGroups; ++group) { - const int channel = group * 32 + lane; - if (channel < Width) { - const float d0 = gradient0[group]; - const float d1 = gradient1[group]; - const float d2 = gradient2[group]; - const float d3 = gradient3[group]; - const float descriptor_product = - environment.r0 * d0 + environment.r1 * d1 + environment.r2 * d2 + - environment.r3 * d3; - const float2 table_value = - evaluate_table_backward(table, location, channel, Width); - const float gate = __ldg( - gate_table + static_cast(environment.pair_index) * Width + - channel); - const float effective_gate = - smooth ? gate * environment.switch_factor : gate; - const float embedding = table_value.x * (1.0f + effective_gate); - if (smooth) { - partial_switch = - fmaf(descriptor_product * table_value.x, gate, partial_switch); - } - partial_radial = fmaf(descriptor_product * (1.0f + effective_gate), - table_value.y, partial_radial); - partial0 = fmaf(embedding, d0, partial0); - partial1 = fmaf(embedding, d1, partial1); - partial2 = fmaf(embedding, d2, partial2); - partial3 = fmaf(embedding, d3, partial3); - } - } -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - partial0 += __shfl_down_sync(kWarpMask, partial0, offset); - partial1 += __shfl_down_sync(kWarpMask, partial1, offset); - partial2 += __shfl_down_sync(kWarpMask, partial2, offset); - partial3 += __shfl_down_sync(kWarpMask, partial3, offset); - partial_radial += __shfl_down_sync(kWarpMask, partial_radial, offset); - partial_switch += __shfl_down_sync(kWarpMask, partial_switch, offset); - } - if (lane == 0) { - store_edge_gradient(edge, environment, partial0, partial1, partial2, - partial3, partial_radial, partial_switch, - inverse_neighbors, inverse_stddev0, inverse_stddev1, - inverse_stddev2, inverse_stddev3, rcut, rcut_smooth, - protection, edge_gradient); - } - } - } -} +namespace { -template -__global__ void zero_padding_kernel( - long node_count, - long edge_count, - const index_t* __restrict__ destination_order, - const long* __restrict__ destination_row_ptr, - float* __restrict__ edge_gradient) { - const long valid_edge_count = destination_row_ptr[node_count]; - for (long position = valid_edge_count + blockIdx.x * blockDim.x + threadIdx.x; - position < edge_count; - position += static_cast(blockDim.x) * gridDim.x) { - const long edge = - edge_at_csr_position(position, destination_order); - edge_gradient[edge * 3 + 0] = 0.0f; - edge_gradient[edge * 3 + 1] = 0.0f; - edge_gradient[edge * 3 + 2] = 0.0f; +using deepmd_dpa1_compress::Arguments; +using deepmd_dpa1_compress::IndexKind; + +cudaError_t dispatch_width(int width, + bool backward, + const Arguments& arguments, + cudaStream_t stream) { +#define DPA1_COMPRESS_DISPATCH(width_value) \ + if (width == width_value) { \ + return backward ? deepmd_dpa1_compress::launch_backward_c##width_value( \ + arguments, stream) \ + : deepmd_dpa1_compress::launch_forward_c##width_value( \ + arguments, stream); \ } + DPA1_COMPRESS_FOR_EACH_CHANNEL(DPA1_COMPRESS_DISPATCH) +#undef DPA1_COMPRESS_DISPATCH + TORCH_CHECK(false, "dpa1_graph_compress: unsupported width ", width); + return cudaErrorInvalidValue; } -template -void launch_forward_variant(long node_count, - int threads, - int ntypes, - bool one_side, - bool smooth, - int axis, - bool concatenate_type_embedding, - bool write_rotation, - int type_embedding_dim, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_index, - const torch::Tensor& edge_mask, - const torch::Tensor& destination_order, - const torch::Tensor& destination_row_ptr, - const torch::Tensor& atype, - const torch::Tensor& type_embedding, - const torch::Tensor& average, - const torch::Tensor& inverse_stddev, - const torch::Tensor& table, - const torch::Tensor& gate_table, - torch::Tensor& descriptor, - torch::Tensor& rotation, - torch::Tensor& moment, - cudaStream_t stream) { - const int warps_per_block = threads / kWarpSize; - const int blocks = - static_cast((node_count + warps_per_block - 1) / warps_per_block); - compressed_forward_kernel - <<>>( - node_count, ntypes, one_side, smooth, axis, - concatenate_type_embedding, write_rotation, type_embedding_dim, rcut, - rcut_smooth, protection, inverse_neighbors, lower, upper, table_max, - stride0, stride1, edge_vec.data_ptr(), - edge_index.data_ptr(), - edge_mask.numel() ? edge_mask.data_ptr() : nullptr, - destination_order.numel() ? destination_order.data_ptr() - : nullptr, - destination_row_ptr.data_ptr(), atype.data_ptr(), - type_embedding.data_ptr(), average.data_ptr(), - inverse_stddev.data_ptr(), table.data_ptr(), - gate_table.data_ptr(), descriptor.data_ptr(), - write_rotation ? rotation.data_ptr() : nullptr, - moment.data_ptr()); -} - -template -void launch_forward(long node_count, - int ntypes, - bool one_side, - bool smooth, - int axis, - bool concatenate_type_embedding, - bool write_rotation, - int type_embedding_dim, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_index, - const torch::Tensor& edge_mask, - const torch::Tensor& destination_order, - const torch::Tensor& destination_row_ptr, - const torch::Tensor& atype, - const torch::Tensor& type_embedding, - const torch::Tensor& average, - const torch::Tensor& inverse_stddev, - const torch::Tensor& table, - const torch::Tensor& gate_table, - torch::Tensor& descriptor, - torch::Tensor& rotation, - torch::Tensor& moment, - cudaStream_t stream) { - const int device = edge_vec.get_device(); - const cudaDeviceProp& properties = device_properties(device); - const TuningKey key = { - device, - static_cast(KernelDirection::kForward), - Width, - axis, - Canonical ? 1 : 0, - static_cast(sizeof(index_t)), - (one_side ? 1 : 0) | (smooth ? 2 : 0) | - (concatenate_type_embedding ? 4 : 0) | (write_rotation ? 8 : 0) | - (Masked ? 16 : 0), - concatenate_type_embedding ? type_embedding_dim : 0, - type_count_class(ntypes), - workload_size_class(node_count, properties.multiProcessorCount), - workload_degree_class(node_count, edge_vec.size(0)), - }; - const auto launch = [&](const LaunchConfig& config, long count) { - if (config.resource == ResourcePolicy::kOccupancy) { - launch_forward_variant( - count, config.threads, ntypes, one_side, smooth, axis, - concatenate_type_embedding, write_rotation, type_embedding_dim, rcut, - rcut_smooth, protection, inverse_neighbors, lower, upper, table_max, - stride0, stride1, edge_vec, edge_index, edge_mask, destination_order, - destination_row_ptr, atype, type_embedding, average, inverse_stddev, - table, gate_table, descriptor, rotation, moment, stream); - } else { - launch_forward_variant( - count, config.threads, ntypes, one_side, smooth, axis, - concatenate_type_embedding, write_rotation, type_embedding_dim, rcut, - rcut_smooth, protection, inverse_neighbors, lower, upper, table_max, - stride0, stride1, edge_vec, edge_index, edge_mask, destination_order, - destination_row_ptr, atype, type_embedding, average, inverse_stddev, - table, gate_table, descriptor, rotation, moment, stream); - } - }; - const LaunchConfig config = - select_launch_config(key, properties, node_count, stream, launch); - launch(config, node_count); - COMPRESS_CHECK_LAUNCH("dpa1_graph_compress forward"); -} - -template -void launch_backward_variant(long node_count, - int threads, - long edge_count, - int ntypes, - bool one_side, - bool smooth, - int axis, - int descriptor_stride, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const torch::Tensor& descriptor_gradient, - const float* rotation_gradient, - const torch::Tensor& moment, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_index, - const torch::Tensor& edge_mask, - const torch::Tensor& destination_order, - const torch::Tensor& destination_row_ptr, - const torch::Tensor& atype, - const torch::Tensor& average, - const torch::Tensor& inverse_stddev, - const torch::Tensor& table, - const torch::Tensor& gate_table, - torch::Tensor& edge_gradient, - cudaStream_t stream) { - const int warps_per_block = threads / kWarpSize; - const int blocks = - static_cast((node_count + warps_per_block - 1) / warps_per_block); - compressed_backward_kernel - <<>>( - node_count, edge_count, ntypes, one_side, smooth, axis, - descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, - lower, upper, table_max, stride0, stride1, - descriptor_gradient.data_ptr(), rotation_gradient, - moment.data_ptr(), edge_vec.data_ptr(), - edge_index.data_ptr(), - edge_mask.numel() ? edge_mask.data_ptr() : nullptr, - destination_order.numel() ? destination_order.data_ptr() - : nullptr, - destination_row_ptr.data_ptr(), atype.data_ptr(), - average.data_ptr(), inverse_stddev.data_ptr(), - table.data_ptr(), gate_table.data_ptr(), - edge_gradient.data_ptr()); +void check_launch(const char* operation, const cudaError_t error) { + TORCH_CHECK(error == cudaSuccess, operation, ": ", cudaGetErrorString(error)); } -template -void launch_backward(long node_count, - long edge_count, - int ntypes, - bool one_side, - bool smooth, - int axis, - int descriptor_stride, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const torch::Tensor& descriptor_gradient, - const float* rotation_gradient, - const torch::Tensor& moment, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_index, - const torch::Tensor& edge_mask, - const torch::Tensor& destination_order, - const torch::Tensor& destination_row_ptr, - const torch::Tensor& atype, - const torch::Tensor& average, - const torch::Tensor& inverse_stddev, - const torch::Tensor& table, - const torch::Tensor& gate_table, - torch::Tensor& edge_gradient, - cudaStream_t stream) { +Arguments make_common_arguments(long node_count, + int basis_dim, + int type_count, + bool one_side, + bool smooth, + int axis, + bool canonical, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& degree_gain, + const torch::Tensor& table, + const torch::Tensor& gate_table) { const int device = edge_vec.get_device(); - const cudaDeviceProp& properties = device_properties(device); - const TuningKey key = { - device, - static_cast(KernelDirection::kBackward), - Width, - axis, - Canonical ? 1 : 0, - static_cast(sizeof(index_t)), - (one_side ? 1 : 0) | (smooth ? 2 : 0) | - (rotation_gradient != nullptr ? 8 : 0) | (Masked ? 16 : 0), - descriptor_stride, - type_count_class(ntypes), - workload_size_class(node_count, properties.multiProcessorCount), - workload_degree_class(node_count, edge_count), - }; - const auto launch = [&](const LaunchConfig& config, long count) { - if (config.resource == ResourcePolicy::kOccupancy) { - launch_backward_variant( - count, config.threads, edge_count, ntypes, one_side, smooth, axis, - descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, - lower, upper, table_max, stride0, stride1, descriptor_gradient, - rotation_gradient, moment, edge_vec, edge_index, edge_mask, - destination_order, destination_row_ptr, atype, average, - inverse_stddev, table, gate_table, edge_gradient, stream); - } else { - launch_backward_variant( - count, config.threads, edge_count, ntypes, one_side, smooth, axis, - descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, - lower, upper, table_max, stride0, stride1, descriptor_gradient, - rotation_gradient, moment, edge_vec, edge_index, edge_mask, - destination_order, destination_row_ptr, atype, average, - inverse_stddev, table, gate_table, edge_gradient, stream); - } - }; - const LaunchConfig config = - select_launch_config(key, properties, node_count, stream, launch); - launch(config, node_count); - COMPRESS_CHECK_LAUNCH("dpa1_graph_compress backward"); - zero_padding_kernel<<<1, kThreads, 0, stream>>>( - node_count, edge_count, - destination_order.numel() ? destination_order.data_ptr() - : nullptr, - destination_row_ptr.data_ptr(), edge_gradient.data_ptr()); - COMPRESS_CHECK_LAUNCH("dpa1_graph_compress padding"); -} - -template -void dispatch_forward(int width, - long node_count, - int ntypes, - bool one_side, - bool smooth, - int axis, - bool canonical, - bool masked, - bool concatenate_type_embedding, - bool write_rotation, - int type_embedding_dim, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_index, - const torch::Tensor& edge_mask, - const torch::Tensor& destination_order, - const torch::Tensor& destination_row_ptr, - const torch::Tensor& atype, - const torch::Tensor& type_embedding, - const torch::Tensor& average, - const torch::Tensor& inverse_stddev, - const torch::Tensor& table, - const torch::Tensor& gate_table, - torch::Tensor& descriptor, - torch::Tensor& rotation, - torch::Tensor& moment, - cudaStream_t stream) { -#define DISPATCH_WIDTH(value) \ - if (width == value) { \ - if (canonical && !masked) { \ - launch_forward( \ - node_count, ntypes, one_side, smooth, axis, \ - concatenate_type_embedding, write_rotation, type_embedding_dim, \ - rcut, rcut_smooth, protection, inverse_neighbors, lower, upper, \ - table_max, stride0, stride1, edge_vec, edge_index, edge_mask, \ - destination_order, destination_row_ptr, atype, type_embedding, \ - average, inverse_stddev, table, gate_table, descriptor, rotation, \ - moment, stream); \ - } else if (canonical) { \ - launch_forward( \ - node_count, ntypes, one_side, smooth, axis, \ - concatenate_type_embedding, write_rotation, type_embedding_dim, \ - rcut, rcut_smooth, protection, inverse_neighbors, lower, upper, \ - table_max, stride0, stride1, edge_vec, edge_index, edge_mask, \ - destination_order, destination_row_ptr, atype, type_embedding, \ - average, inverse_stddev, table, gate_table, descriptor, rotation, \ - moment, stream); \ - } else { \ - launch_forward( \ - node_count, ntypes, one_side, smooth, axis, \ - concatenate_type_embedding, write_rotation, type_embedding_dim, \ - rcut, rcut_smooth, protection, inverse_neighbors, lower, upper, \ - table_max, stride0, stride1, edge_vec, edge_index, edge_mask, \ - destination_order, destination_row_ptr, atype, type_embedding, \ - average, inverse_stddev, table, gate_table, descriptor, rotation, \ - moment, stream); \ - } \ - return; \ - } - DISPATCH_WIDTH(8) - DISPATCH_WIDTH(16) - DISPATCH_WIDTH(32) - DISPATCH_WIDTH(64) - DISPATCH_WIDTH(128) - DISPATCH_WIDTH(256) -#undef DISPATCH_WIDTH - TORCH_CHECK(false, "dpa1_graph_compress: unsupported width ", width); -} - -template -void dispatch_backward(int width, - long node_count, - long edge_count, - int ntypes, - bool one_side, - bool smooth, - int axis, - bool canonical, - bool masked, - int descriptor_stride, - float rcut, - float rcut_smooth, - float protection, - float inverse_neighbors, - float lower, - float upper, - float table_max, - float stride0, - float stride1, - const torch::Tensor& descriptor_gradient, - const float* rotation_gradient, - const torch::Tensor& moment, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_index, - const torch::Tensor& edge_mask, - const torch::Tensor& destination_order, - const torch::Tensor& destination_row_ptr, - const torch::Tensor& atype, - const torch::Tensor& average, - const torch::Tensor& inverse_stddev, - const torch::Tensor& table, - const torch::Tensor& gate_table, - torch::Tensor& edge_gradient, - cudaStream_t stream) { -#define DISPATCH_WIDTH(value) \ - if (width == value) { \ - if (canonical && !masked) { \ - launch_backward( \ - node_count, edge_count, ntypes, one_side, smooth, axis, \ - descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, \ - lower, upper, table_max, stride0, stride1, descriptor_gradient, \ - rotation_gradient, moment, edge_vec, edge_index, edge_mask, \ - destination_order, destination_row_ptr, atype, average, \ - inverse_stddev, table, gate_table, edge_gradient, stream); \ - } else if (canonical) { \ - launch_backward( \ - node_count, edge_count, ntypes, one_side, smooth, axis, \ - descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, \ - lower, upper, table_max, stride0, stride1, descriptor_gradient, \ - rotation_gradient, moment, edge_vec, edge_index, edge_mask, \ - destination_order, destination_row_ptr, atype, average, \ - inverse_stddev, table, gate_table, edge_gradient, stream); \ - } else { \ - launch_backward( \ - node_count, edge_count, ntypes, one_side, smooth, axis, \ - descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, \ - lower, upper, table_max, stride0, stride1, descriptor_gradient, \ - rotation_gradient, moment, edge_vec, edge_index, edge_mask, \ - destination_order, destination_row_ptr, atype, average, \ - inverse_stddev, table, gate_table, edge_gradient, stream); \ - } \ - return; \ - } - DISPATCH_WIDTH(8) - DISPATCH_WIDTH(16) - DISPATCH_WIDTH(32) - DISPATCH_WIDTH(64) - DISPATCH_WIDTH(128) - DISPATCH_WIDTH(256) -#undef DISPATCH_WIDTH - TORCH_CHECK(false, "dpa1_graph_compress_backward: unsupported width ", width); + const cudaDeviceProp* properties = at::cuda::getDeviceProperties(device); + TORCH_CHECK(properties != nullptr, + "dpa1_graph_compress: cannot query CUDA device properties"); + + Arguments arguments; + arguments.node_count = node_count; + arguments.edge_count = edge_vec.size(0); + arguments.device = device; + arguments.device_major = properties->major; + arguments.multiprocessor_count = properties->multiProcessorCount; + arguments.basis_dim = basis_dim; + arguments.type_count = type_count; + arguments.axis = axis; + arguments.one_side = one_side; + arguments.smooth = smooth; + arguments.canonical = canonical; + arguments.masked = edge_mask.numel() != 0; + arguments.index_kind = edge_index.scalar_type() == torch::kInt32 + ? IndexKind::kInt32 + : IndexKind::kInt64; + arguments.rcut = rcut; + arguments.rcut_smooth = rcut_smooth; + arguments.protection = protection; + arguments.inverse_neighbors = inverse_neighbors; + arguments.lower = lower; + arguments.upper = upper; + arguments.table_max = table_max; + arguments.stride0 = stride0; + arguments.stride1 = stride1; + arguments.edge_vec = edge_vec.data_ptr(); + arguments.edge_index = + arguments.index_kind == IndexKind::kInt32 + ? static_cast(edge_index.data_ptr()) + : static_cast(edge_index.data_ptr()); + arguments.edge_mask = arguments.masked ? edge_mask.data_ptr() : nullptr; + arguments.destination_order = + destination_order.numel() == 0 + ? nullptr + : (arguments.index_kind == IndexKind::kInt32 + ? static_cast(destination_order.data_ptr()) + : static_cast( + destination_order.data_ptr())); + arguments.destination_row_ptr = destination_row_ptr.data_ptr(); + arguments.atype = atype.data_ptr(); + arguments.average = average.data_ptr(); + arguments.inverse_stddev = inverse_stddev.data_ptr(); + arguments.degree_gain = + degree_gain.numel() == 0 ? nullptr : degree_gain.data_ptr(); + arguments.table = table.data_ptr(); + arguments.gate_table = gate_table.data_ptr(); + return arguments; } void validate_inputs(const torch::Tensor& edge_vec, @@ -1389,6 +185,7 @@ std::tuple dpa1_graph_compress( torch::Tensor type_embedding, torch::Tensor average, torch::Tensor inverse_stddev, + torch::Tensor degree_gain, torch::Tensor table, torch::Tensor gate_table, int64_t type_one_side, @@ -1405,7 +202,8 @@ std::tuple dpa1_graph_compress( double rcut, double rcut_smooth, double protection, - double neighbors) { + double neighbors, + int64_t basis_dim) { const long node_count = atype.size(0); const int width = static_cast(table.size(1) / 6); validate_inputs(edge_vec, edge_index, edge_mask, destination_order, @@ -1415,6 +213,25 @@ std::tuple dpa1_graph_compress( type_embedding.scalar_type() == torch::kFloat32, "dpa1_graph_compress: type_embedding must be contiguous fp32 " "on CUDA"); +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + TORCH_CHECK( + basis_dim == 4 || basis_dim == 9 || basis_dim == 16 || basis_dim == 25, + "dpa1_graph_compress: basis_dim must be 4, 9, 16, or 25"); + TORCH_CHECK(basis_dim <= 9 || (width >= 16 && width <= 128), + "dpa1_graph_compress: basis_dim 16 and 25 require width 16, 32, " + "64, or 128"); +#else + TORCH_CHECK( + basis_dim == 4, + "dpa1_graph_compress: this build instantiates only lmax=1; rebuild " + "with DEEPMD_ENABLE_DPA1_HIGH_LMAX=ON for lmax=2/3/4"); +#endif + const int degree_gain_size = + basis_dim == 4 ? 0 : (basis_dim == 9 ? 1 : (basis_dim == 16 ? 2 : 3)); + TORCH_CHECK(degree_gain.is_contiguous() && + degree_gain.scalar_type() == torch::kFloat32 && + degree_gain.numel() == degree_gain_size, + "dpa1_graph_compress: degree_gain has an invalid shape or dtype"); const int ntypes = static_cast(type_embedding.size(0)); const int type_embedding_dim = static_cast(type_embedding.size(1)); const int output_dim = width * static_cast(axis) + @@ -1423,33 +240,30 @@ std::tuple dpa1_graph_compress( auto descriptor = torch::empty({node_count, output_dim}, options); auto rotation = torch::empty({write_rotation ? node_count : 0, width, 3}, options); - auto moment = torch::empty({node_count, 4, width}, options); + auto moment = torch::empty({node_count, basis_dim, width}, options); if (node_count == 0) { return {descriptor, rotation, moment}; } const auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); const auto stream = at::cuda::getCurrentCUDAStream(); - - auto launch = [&](auto index_tag) { - using index_t = decltype(index_tag); - dispatch_forward( - width, node_count, ntypes, type_one_side != 0, smooth != 0, - static_cast(axis), canonical, edge_mask.numel() != 0, - concatenate_type_embedding != 0, write_rotation != 0, - type_embedding_dim, static_cast(rcut), - static_cast(rcut_smooth), static_cast(protection), - static_cast(1.0 / neighbors), static_cast(lower), - static_cast(upper), static_cast(table_max), - static_cast(stride0), static_cast(stride1), - edge_vec_float, edge_index, edge_mask, destination_order, - destination_row_ptr, atype, type_embedding, average, inverse_stddev, - table, gate_table, descriptor, rotation, moment, stream); - }; - if (edge_index.scalar_type() == torch::kInt32) { - launch(int{}); - } else { - launch(long{}); - } + Arguments arguments = make_common_arguments( + node_count, static_cast(basis_dim), ntypes, type_one_side != 0, + smooth != 0, static_cast(axis), canonical, static_cast(rcut), + static_cast(rcut_smooth), static_cast(protection), + static_cast(1.0 / neighbors), static_cast(lower), + static_cast(upper), static_cast(table_max), + static_cast(stride0), static_cast(stride1), edge_vec_float, + edge_index, edge_mask, destination_order, destination_row_ptr, atype, + average, inverse_stddev, degree_gain, table, gate_table); + arguments.concatenate_type_embedding = concatenate_type_embedding != 0; + arguments.write_rotation = write_rotation != 0; + arguments.type_embedding_dim = type_embedding_dim; + arguments.type_embedding = type_embedding.data_ptr(); + arguments.descriptor = descriptor.data_ptr(); + arguments.rotation = write_rotation ? rotation.data_ptr() : nullptr; + arguments.moment_out = moment.data_ptr(); + check_launch("dpa1_graph_compress forward", + dispatch_width(width, false, arguments, stream)); return {descriptor, rotation, moment}; } @@ -1465,6 +279,7 @@ torch::Tensor dpa1_graph_compress_backward( torch::Tensor atype, torch::Tensor average, torch::Tensor inverse_stddev, + torch::Tensor degree_gain, torch::Tensor table, torch::Tensor gate_table, int64_t type_one_side, @@ -1481,11 +296,26 @@ torch::Tensor dpa1_graph_compress_backward( double protection, double neighbors) { const long node_count = atype.size(0); - const long edge_count = edge_vec.size(0); const int width = static_cast(table.size(1) / 6); + const int basis_dim = static_cast(moment.size(1)); validate_inputs(edge_vec, edge_index, edge_mask, destination_order, destination_row_ptr, atype, average, inverse_stddev, table, gate_table, width, static_cast(axis)); +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + TORCH_CHECK( + basis_dim == 4 || basis_dim == 9 || basis_dim == 16 || basis_dim == 25, + "dpa1_graph_compress_backward: basis dimension must be 4, 9, " + "16, or 25"); + TORCH_CHECK( + basis_dim <= 9 || (width >= 16 && width <= 128), + "dpa1_graph_compress_backward: basis dimensions 16 and 25 require " + "width 16, 32, 64, or 128"); +#else + TORCH_CHECK( + basis_dim == 4, + "dpa1_graph_compress_backward: this build instantiates only lmax=1; " + "rebuild with DEEPMD_ENABLE_DPA1_HIGH_LMAX=ON for lmax=2/3/4"); +#endif if (node_count == 0) { return torch::zeros_like(edge_vec); } @@ -1506,27 +336,23 @@ torch::Tensor dpa1_graph_compress_backward( auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); auto edge_gradient = torch::empty_like(edge_vec_float); const auto stream = at::cuda::getCurrentCUDAStream(); - - auto launch = [&](auto index_tag) { - using index_t = decltype(index_tag); - dispatch_backward( - width, node_count, edge_count, ntypes, type_one_side != 0, smooth != 0, - static_cast(axis), canonical, edge_mask.numel() != 0, - static_cast(descriptor_gradient_float.size(1)), - static_cast(rcut), static_cast(rcut_smooth), - static_cast(protection), static_cast(1.0 / neighbors), - static_cast(lower), static_cast(upper), - static_cast(table_max), static_cast(stride0), - static_cast(stride1), descriptor_gradient_float, - rotation_gradient_ptr, moment, edge_vec_float, edge_index, edge_mask, - destination_order, destination_row_ptr, atype, average, inverse_stddev, - table, gate_table, edge_gradient, stream); - }; - if (edge_index.scalar_type() == torch::kInt32) { - launch(int{}); - } else { - launch(long{}); - } + Arguments arguments = make_common_arguments( + node_count, basis_dim, ntypes, type_one_side != 0, smooth != 0, + static_cast(axis), canonical, static_cast(rcut), + static_cast(rcut_smooth), static_cast(protection), + static_cast(1.0 / neighbors), static_cast(lower), + static_cast(upper), static_cast(table_max), + static_cast(stride0), static_cast(stride1), edge_vec_float, + edge_index, edge_mask, destination_order, destination_row_ptr, atype, + average, inverse_stddev, degree_gain, table, gate_table); + arguments.descriptor_stride = + static_cast(descriptor_gradient_float.size(1)); + arguments.descriptor_gradient = descriptor_gradient_float.data_ptr(); + arguments.rotation_gradient = rotation_gradient_ptr; + arguments.moment = moment.data_ptr(); + arguments.edge_gradient = edge_gradient.data_ptr(); + check_launch("dpa1_graph_compress backward", + dispatch_width(width, true, arguments, stream)); return edge_gradient.to(edge_vec.scalar_type()); } @@ -1538,6 +364,7 @@ std::tuple dpa1_canonical_compress( torch::Tensor type_embedding, torch::Tensor average, torch::Tensor inverse_stddev, + torch::Tensor degree_gain, torch::Tensor table, torch::Tensor gate_table, int64_t type_one_side, @@ -1553,7 +380,8 @@ std::tuple dpa1_canonical_compress( double rcut, double rcut_smooth, double protection, - double neighbors) { + double neighbors, + int64_t basis_dim) { TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), "dpa1_canonical_compress: source and edge_vec storage must " "share the edge axis"); @@ -1564,10 +392,10 @@ std::tuple dpa1_canonical_compress( auto destination_order = torch::empty({0}, source.options()); return dpa1_graph_compress( edge_vec, source, edge_mask, destination_order, destination_row_ptr, - atype, type_embedding, average, inverse_stddev, table, gate_table, - type_one_side, concatenate_type_embedding, write_rotation, smooth, axis, - true, lower, upper, table_max, stride0, stride1, rcut, rcut_smooth, - protection, neighbors); + atype, type_embedding, average, inverse_stddev, degree_gain, table, + gate_table, type_one_side, concatenate_type_embedding, write_rotation, + smooth, axis, true, lower, upper, table_max, stride0, stride1, rcut, + rcut_smooth, protection, neighbors, basis_dim); } torch::Tensor dpa1_canonical_compress_backward( @@ -1580,6 +408,7 @@ torch::Tensor dpa1_canonical_compress_backward( torch::Tensor atype, torch::Tensor average, torch::Tensor inverse_stddev, + torch::Tensor degree_gain, torch::Tensor table, torch::Tensor gate_table, int64_t type_one_side, @@ -1605,9 +434,9 @@ torch::Tensor dpa1_canonical_compress_backward( return dpa1_graph_compress_backward( descriptor_gradient, rotation_gradient, moment, edge_vec, source, edge_mask, destination_order, destination_row_ptr, atype, average, - inverse_stddev, table, gate_table, type_one_side, smooth, axis, true, - lower, upper, table_max, stride0, stride1, rcut, rcut_smooth, protection, - neighbors); + inverse_stddev, degree_gain, table, gate_table, type_one_side, smooth, + axis, true, lower, upper, table_max, stride0, stride1, rcut, rcut_smooth, + protection, neighbors); } TORCH_LIBRARY_FRAGMENT(deepmd, library) { @@ -1616,11 +445,12 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor edge_mask, Tensor destination_order, " "Tensor destination_row_ptr, Tensor atype, " "Tensor type_embedding, Tensor average, Tensor inverse_stddev, " - "Tensor table, Tensor gate_table, int type_one_side, " + "Tensor degree_gain, Tensor table, Tensor gate_table, int type_one_side, " "int concatenate_type_embedding, int write_rotation, int smooth, " "int axis, bool canonical, float lower, float upper, float table_max, " "float stride0, float stride1, " - "float rcut, float rcut_smooth, float protection, float neighbors) " + "float rcut, float rcut_smooth, float protection, float neighbors, " + "int basis_dim) " "-> (Tensor descriptor, Tensor rotation, Tensor moment)"); library.impl("dpa1_graph_compress", torch::kCUDA, &dpa1_graph_compress); library.def( @@ -1628,7 +458,7 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor? rotation_gradient, Tensor moment, Tensor edge_vec, " "Tensor edge_index, Tensor edge_mask, Tensor destination_order, " "Tensor destination_row_ptr, Tensor atype, Tensor average, " - "Tensor inverse_stddev, Tensor table, " + "Tensor inverse_stddev, Tensor degree_gain, Tensor table, " "Tensor gate_table, int type_one_side, int smooth, int axis, " "bool canonical, float lower, float upper, float table_max, float " "stride0, " @@ -1639,11 +469,12 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { library.def( "dpa1_canonical_compress(Tensor edge_vec, Tensor source, " "Tensor destination_row_ptr, Tensor atype, Tensor type_embedding, " - "Tensor average, Tensor inverse_stddev, Tensor table, " + "Tensor average, Tensor inverse_stddev, Tensor degree_gain, Tensor " + "table, " "Tensor gate_table, int type_one_side, int concatenate_type_embedding, " "int write_rotation, int smooth, int axis, float lower, float upper, " "float table_max, float stride0, float stride1, float rcut, " - "float rcut_smooth, float protection, float neighbors) -> " + "float rcut_smooth, float protection, float neighbors, int basis_dim) -> " "(Tensor descriptor, Tensor rotation, Tensor moment)"); library.impl("dpa1_canonical_compress", torch::kCUDA, &dpa1_canonical_compress); @@ -1651,7 +482,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "dpa1_canonical_compress_backward(Tensor descriptor_gradient, " "Tensor? rotation_gradient, Tensor moment, Tensor edge_vec, " "Tensor source, Tensor destination_row_ptr, Tensor atype, " - "Tensor average, Tensor inverse_stddev, Tensor table, " + "Tensor average, Tensor inverse_stddev, Tensor degree_gain, Tensor " + "table, " "Tensor gate_table, int type_one_side, int smooth, int axis, " "float lower, float upper, float table_max, float stride0, " "float stride1, float rcut, float rcut_smooth, float protection, " diff --git a/source/op/pt/dpa1_graph_compress_c128.cu b/source/op/pt/dpa1_graph_compress_c128.cu new file mode 100644 index 0000000000..b90a4e30ed --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_c128.cu @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compressed DPA1 CUDA specializations for 128 embedding channels. + +#include "dpa1_graph_compress_kernel.cuh" + +namespace deepmd_dpa1_compress { + +DPA1_COMPRESS_DEFINE_CHANNEL(128) + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_c16.cu b/source/op/pt/dpa1_graph_compress_c16.cu new file mode 100644 index 0000000000..4bf38d4843 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_c16.cu @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compressed DPA1 CUDA specializations for 16 embedding channels. + +#include "dpa1_graph_compress_kernel.cuh" + +namespace deepmd_dpa1_compress { + +DPA1_COMPRESS_DEFINE_CHANNEL(16) + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_c256.cu b/source/op/pt/dpa1_graph_compress_c256.cu new file mode 100644 index 0000000000..12801db989 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_c256.cu @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compressed DPA1 CUDA specializations for 256 embedding channels. + +#include "dpa1_graph_compress_kernel.cuh" + +namespace deepmd_dpa1_compress { + +DPA1_COMPRESS_DEFINE_CHANNEL(256) + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_c32.cu b/source/op/pt/dpa1_graph_compress_c32.cu new file mode 100644 index 0000000000..30f8bac9b0 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_c32.cu @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compressed DPA1 CUDA specializations for 32 embedding channels. + +#include "dpa1_graph_compress_kernel.cuh" + +namespace deepmd_dpa1_compress { + +DPA1_COMPRESS_DEFINE_CHANNEL(32) + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_c64.cu b/source/op/pt/dpa1_graph_compress_c64.cu new file mode 100644 index 0000000000..65cf47d8c8 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_c64.cu @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compressed DPA1 CUDA specializations for 64 embedding channels. + +#include "dpa1_graph_compress_kernel.cuh" + +namespace deepmd_dpa1_compress { + +DPA1_COMPRESS_DEFINE_CHANNEL(64) + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_c8.cu b/source/op/pt/dpa1_graph_compress_c8.cu new file mode 100644 index 0000000000..0d8a9c49d5 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_c8.cu @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compressed DPA1 CUDA specializations for 8 embedding channels. + +#include "dpa1_graph_compress_kernel.cuh" + +namespace deepmd_dpa1_compress { + +DPA1_COMPRESS_DEFINE_CHANNEL(8) + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_kernel.cuh b/source/op/pt/dpa1_graph_compress_kernel.cuh new file mode 100644 index 0000000000..91362858d8 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_kernel.cuh @@ -0,0 +1,1168 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernels of the geometrically compressed DPA1 graph descriptor. +// +// This header is included by exactly one translation unit per channel width, +// which lets the expensive basis, topology, index, and resource-policy +// specializations compile in parallel. +// +// A warp owns one center node. Widths from 16 through 64 use two 16-lane +// sub-warps on alternating edges; each lane evaluates one or more spline +// channels. Wider tables retain one edge per warp to bound register pressure. +// The node moment and its Gram contraction remain in the same kernel. +// +// The backward recomputes the inexpensive spline value/derivative, contracts +// the descriptor gradient into the four environment channels, and writes each +// edge gradient exactly once. It is inference-oriented (one backward); the +// registered Python autograd bridge continues to expose the edge-vector +// gradient for the level-1 graph path. +// +// Every specialization has balanced (two CTA/SM launch bound) and occupancy +// (four CTA/SM launch bound) resource variants. The first uncaptured call times +// 128- and 256-thread launches on a bounded node sample and caches the selected +// variant per device and workload class. Device-family defaults remain valid +// when timing is disabled or CUDA Graph capture is active. + +#pragma once + +#include + +#include + +#include "dpa1_graph_compress_launch.h" + +#ifndef DEEPMD_ENABLE_DPA1_HIGH_LMAX +#define DEEPMD_ENABLE_DPA1_HIGH_LMAX 0 +#endif + +#include "dpa1_graph_compress_tuning.h" +#include "dpa1_moment_basis.cuh" + +// The lmax=2/3/4 compressed kernels are retained below as an experimental +// implementation, but their template instances are disabled by default. +// Current DPA1 production artifacts use lmax=1; the CMake option can restore +// the higher-degree binary specializations without recovering deleted code. + +namespace deepmd_dpa1_compress { +namespace { + +using deepmd::dpa1_compress_tuning::DeviceProperties; +using deepmd::dpa1_compress_tuning::KernelDirection; +using deepmd::dpa1_compress_tuning::LaunchConfig; +using deepmd::dpa1_compress_tuning::ResourcePolicy; +using deepmd::dpa1_compress_tuning::select_launch_config; +using deepmd::dpa1_compress_tuning::TuningKey; +using deepmd::dpa1_compress_tuning::type_count_class; +using deepmd::dpa1_compress_tuning::workload_degree_class; +using deepmd::dpa1_compress_tuning::workload_size_class; + +constexpr int kThreads = 256; +constexpr int kWarpSize = 32; + +__device__ __forceinline__ float switch_value(float radius, + float lower, + float upper) { + const float coordinate = + __fdividef(fminf(fmaxf(radius, lower), upper) - lower, upper - lower); + const float square = coordinate * coordinate; + return square * coordinate * (-6.0f * square + 15.0f * coordinate - 10.0f) + + 1.0f; +} + +__device__ __forceinline__ float switch_derivative(float radius, + float lower, + float upper) { + if (radius <= lower || radius >= upper) { + return 0.0f; + } + const float coordinate = __fdividef(radius - lower, upper - lower); + const float square = coordinate * coordinate; + return __fdividef( + -30.0f * square * square + 60.0f * square * coordinate - 30.0f * square, + upper - lower); +} + +struct TableLocation { + int index; + float coordinate; + float extrapolation; +}; + +__device__ __forceinline__ int high_tail_index( + float lower, float upper, float table_max, float stride0, float stride1) { + const float boundary = nextafterf(table_max, lower); + const int first_stride = static_cast(__fdividef(upper - lower, stride0)); + return first_stride + static_cast(__fdividef(boundary - upper, stride1)); +} + +__device__ __forceinline__ TableLocation locate_table(float radial, + float lower, + float upper, + float table_max, + float stride0, + float stride1) { + TableLocation location; + location.coordinate = radial; + location.extrapolation = 0.0f; + if (radial < lower) { + location.index = 0; + location.coordinate = 0.0f; + location.extrapolation = radial - lower; + } else if (radial < upper) { + location.index = static_cast(__fdividef(radial - lower, stride0)); + location.coordinate -= location.index * stride0 + lower; + } else if (radial < table_max) { + const int first_stride = + static_cast(__fdividef(upper - lower, stride0)); + location.index = + first_stride + static_cast(__fdividef(radial - upper, stride1)); + location.coordinate -= (location.index - first_stride) * stride1 + upper; + } else { + const int first_stride = + static_cast(__fdividef(upper - lower, stride0)); + location.index = high_tail_index(lower, upper, table_max, stride0, stride1); + location.coordinate = + table_max - ((location.index - first_stride) * stride1 + upper); + location.extrapolation = radial - table_max; + } + return location; +} + +__device__ __forceinline__ void load_coefficients(const float* table, + const TableLocation& location, + int channel, + int width, + float2& c01, + float2& c23, + float2& c45) { + const long offset = static_cast(location.index) * width * 6 + + static_cast(channel) * 6; + c01 = __ldg(reinterpret_cast(table + offset)); + c23 = __ldg(reinterpret_cast(table + offset + 2)); + c45 = __ldg(reinterpret_cast(table + offset + 4)); +} + +__device__ __forceinline__ float evaluate_table_forward( + const float* table, const TableLocation& location, int channel, int width) { + float2 c01, c23, c45; + load_coefficients(table, location, channel, width, c01, c23, c45); + const float value = + c01.x + (c01.y + (c23.x + (c23.y + (c45.x + c45.y * location.coordinate) * + location.coordinate) * + location.coordinate) * + location.coordinate) * + location.coordinate; + if (location.extrapolation == 0.0f) { + return value; + } + const float derivative = + c01.y + + (2.0f * c23.x + + (3.0f * c23.y + (4.0f * c45.x + 5.0f * c45.y * location.coordinate) * + location.coordinate) * + location.coordinate) * + location.coordinate; + return value + derivative * location.extrapolation; +} + +__device__ __forceinline__ float2 evaluate_table_backward( + const float* table, const TableLocation& location, int channel, int width) { + float2 c01, c23, c45; + load_coefficients(table, location, channel, width, c01, c23, c45); + float value = c45.y; + float derivative = 0.0f; + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c45.x); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c23.y); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c23.x); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c01.y); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c01.x); + return make_float2(value + derivative * location.extrapolation, derivative); +} + +template +struct EdgeEnvironment { + float radial; + float basis[BasisDim]; + float switch_factor; + float x; + float y; + float z; + float radius; + int pair_index; +}; + +template +__device__ __forceinline__ EdgeEnvironment load_environment( + long edge, + int center_type, + int ntypes, + bool one_side, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + const float* edge_vec, + const index_t* edge_index, + const long* atype, + const float* average, + const float* inverse_stddev) { + EdgeEnvironment environment; + const long source = static_cast(edge_index[edge]); + const int neighbor_type = static_cast(atype[source]); + environment.x = edge_vec[edge * 3 + 0]; + environment.y = edge_vec[edge * 3 + 1]; + environment.z = edge_vec[edge * 3 + 2]; + const float square_length = environment.x * environment.x + + environment.y * environment.y + + environment.z * environment.z; + environment.radius = + square_length > 0.0f ? square_length * rsqrtf(square_length) : 0.0f; + const float denominator = environment.radius + protection; + environment.switch_factor = + switch_value(environment.radius, rcut_smooth, rcut); + const float inverse_denominator = __fdividef(1.0f, denominator); + const float radial_scale = + environment.switch_factor * inverse_denominator * inverse_denominator; + const float* center_average = average + static_cast(center_type) * 4; + const float* center_inverse_stddev = + inverse_stddev + static_cast(center_type) * 4; + environment.radial = + (environment.switch_factor * inverse_denominator - center_average[0]) * + center_inverse_stddev[0]; + environment.basis[0] = environment.radial * inverse_neighbors; + environment.basis[1] = (environment.x * radial_scale - center_average[1]) * + center_inverse_stddev[1] * inverse_neighbors; + environment.basis[2] = (environment.y * radial_scale - center_average[2]) * + center_inverse_stddev[2] * inverse_neighbors; + environment.basis[3] = (environment.z * radial_scale - center_average[3]) * + center_inverse_stddev[3] * inverse_neighbors; + if constexpr (BasisDim > 4) { + const float vx = environment.x * inverse_denominator; + const float vy = environment.y * inverse_denominator; + const float vz = environment.z * inverse_denominator; + const float radial = environment.radius > 0.0f + ? environment.switch_factor * inverse_denominator * + center_inverse_stddev[0] * inverse_neighbors + : 0.0f; + deepmd::dpa1::fill_angular_basis(environment.basis, vx, vy, vz, + radial); + } + environment.pair_index = + one_side ? neighbor_type : center_type * ntypes + neighbor_type; + return environment; +} + +template +__device__ __forceinline__ EdgeEnvironment broadcast_environment( + EdgeEnvironment value, int source_lane, unsigned mask) { + value.radial = __shfl_sync(mask, value.radial, source_lane); +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + value.basis[k] = __shfl_sync(mask, value.basis[k], source_lane); + } + value.switch_factor = __shfl_sync(mask, value.switch_factor, source_lane); + value.x = __shfl_sync(mask, value.x, source_lane); + value.y = __shfl_sync(mask, value.y, source_lane); + value.z = __shfl_sync(mask, value.z, source_lane); + value.radius = __shfl_sync(mask, value.radius, source_lane); + value.pair_index = __shfl_sync(mask, value.pair_index, source_lane); + return value; +} + +__device__ __forceinline__ TableLocation broadcast_location(TableLocation value, + int source_lane, + unsigned mask) { + value.index = __shfl_sync(mask, value.index, source_lane); + value.coordinate = __shfl_sync(mask, value.coordinate, source_lane); + value.extrapolation = __shfl_sync(mask, value.extrapolation, source_lane); + return value; +} + +template +__device__ __forceinline__ void store_edge_gradient( + long edge, + const EdgeEnvironment& environment, + const float (&partial_basis)[BasisDim], + float partial_radial, + float partial_switch, + float inverse_neighbors, + float inverse_stddev0, + float inverse_stddev1, + float inverse_stddev2, + float inverse_stddev3, + float rcut, + float rcut_smooth, + float protection, + float* edge_gradient) { + const float inverse_denominator = + __fdividef(1.0f, environment.radius + protection); + const float inverse_length = + environment.radius > 0.0f ? __fdividef(1.0f, environment.radius) : 0.0f; + const float switch_gradient = + switch_derivative(environment.radius, rcut_smooth, rcut); + const float gradient_radial = + (partial_basis[0] * inverse_neighbors + partial_radial) * inverse_stddev0; + const float gradient_x = + partial_basis[1] * inverse_neighbors * inverse_stddev1; + const float gradient_y = + partial_basis[2] * inverse_neighbors * inverse_stddev2; + const float gradient_z = + partial_basis[3] * inverse_neighbors * inverse_stddev3; + const float directional = gradient_x * environment.x + + gradient_y * environment.y + + gradient_z * environment.z; + const float coefficient = + (gradient_radial * inverse_denominator * + (switch_gradient - environment.switch_factor * inverse_denominator) + + directional * inverse_denominator * inverse_denominator * + (switch_gradient - + 2.0f * environment.switch_factor * inverse_denominator) + + partial_switch * switch_gradient) * + inverse_length; + const float vector_scale = + environment.switch_factor * inverse_denominator * inverse_denominator; + float output_x = coefficient * environment.x + vector_scale * gradient_x; + float output_y = coefficient * environment.y + vector_scale * gradient_y; + float output_z = coefficient * environment.z + vector_scale * gradient_z; + if constexpr (BasisDim > 4) { + deepmd::dpa1::add_angular_edge_gradient( + partial_basis, inverse_neighbors, environment.x, environment.y, + environment.z, environment.radius, inverse_denominator, inverse_stddev0, + environment.switch_factor, switch_gradient, output_x, output_y, + output_z); + } + edge_gradient[edge * 3 + 0] = output_x; + edge_gradient[edge * 3 + 1] = output_y; + edge_gradient[edge * 3 + 2] = output_z; +} + +template +struct ChannelPolicy { + static constexpr bool use_half_warp = Width >= 16 && Width <= 64; + static constexpr int accumulation_groups = + use_half_warp ? Width / 16 : (Width + 31) / 32; + static constexpr int gradient_groups = (Width + 31) / 32; +}; + +template +__device__ __forceinline__ long edge_at_csr_position( + long position, const index_t* destination_order) { + if constexpr (Canonical) { + return position; + } + return static_cast(destination_order[position]); +} + +template +__device__ __forceinline__ bool edge_is_active(long edge, + const bool* edge_mask) { + if constexpr (Masked) { + return edge_mask[edge]; + } + return true; +} + +template +__global__ +__launch_bounds__(kThreads, MinimumBlocks) void compressed_forward_kernel( + long node_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + bool concatenate_type_embedding, + bool write_rotation, + int type_embedding_dim, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const float* __restrict__ edge_vec, + const index_t* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const index_t* __restrict__ destination_order, + const long* __restrict__ destination_row_ptr, + const long* __restrict__ atype, + const float* __restrict__ type_embedding, + const float* __restrict__ average, + const float* __restrict__ inverse_stddev, + const float* __restrict__ table, + const float* __restrict__ gate_table, + const float* __restrict__ degree_gain_raw, + float* __restrict__ descriptor, + float* __restrict__ rotation, + float* __restrict__ moment) { + constexpr unsigned kWarpMask = 0xffffffffu; + constexpr int kGroups = ChannelPolicy::accumulation_groups; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int warps_per_block = blockDim.x / kWarpSize; + const long node = static_cast(blockIdx.x) * warps_per_block + warp; + if (node >= node_count) { + return; + } + + float accumulator[BasisDim][kGroups] = {}; + int center_type = lane == 0 ? static_cast(atype[node]) : 0; + center_type = __shfl_sync(kWarpMask, center_type, 0); + const long begin = destination_row_ptr[node]; + const long end = destination_row_ptr[node + 1]; + + if constexpr (ChannelPolicy::use_half_warp) { + const int half = lane >> 4; + const int half_lane = lane & 15; + const int leader = half * 16; + const unsigned half_mask = half == 0 ? 0x0000ffffu : 0xffff0000u; + for (long position = begin + half; position < end; position += 2) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (half_lane == 0) { + environment = load_environment( + edge, center_type, ntypes, one_side, rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, atype, average, + inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = + broadcast_environment(environment, leader, half_mask); + location = broadcast_location(location, leader, half_mask); +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 16 + half_lane; + const float table_value = + evaluate_table_forward(table, location, channel, Width); + const float gate = + __ldg(gate_table + + static_cast(environment.pair_index) * Width + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value * (1.0f + effective_gate); +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + accumulator[k][group] = + fmaf(environment.basis[k], embedding, accumulator[k][group]); + } + } + } +#pragma unroll + for (int group = 0; group < kGroups; ++group) { +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + accumulator[k][group] += + __shfl_xor_sync(kWarpMask, accumulator[k][group], 16); + } + } + } else { + for (long position = begin; position < end; ++position) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (lane == 0) { + environment = load_environment( + edge, center_type, ntypes, one_side, rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, atype, average, + inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = broadcast_environment(environment, 0, kWarpMask); + location = broadcast_location(location, 0, kWarpMask); +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 32 + lane; + if (channel < Width) { + const float table_value = + evaluate_table_forward(table, location, channel, Width); + const float gate = __ldg( + gate_table + static_cast(environment.pair_index) * Width + + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value * (1.0f + effective_gate); +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + accumulator[k][group] = + fmaf(environment.basis[k], embedding, accumulator[k][group]); + } + } + } + } + } + + const long moment_base = node * BasisDim * Width; + if constexpr (ChannelPolicy::use_half_warp) { + if (lane < 16) { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 16 + lane; +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + moment[moment_base + k * Width + channel] = accumulator[k][group]; + } + } + } + } else { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 32 + lane; + if (channel < Width) { +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + moment[moment_base + k * Width + channel] = accumulator[k][group]; + } + } + } + } + + const int output_dim = + Width * axis + (concatenate_type_embedding ? type_embedding_dim : 0); + float* output = descriptor + node * output_dim; + if constexpr (ChannelPolicy::use_half_warp) { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 16 + (lane & 15); + for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { + float value = 0.0f; +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + const float axis_value = + __shfl_sync(kWarpMask, accumulator[k][0], axis_channel); + const float weight = + BasisDim == 4 ? 1.0f + : deepmd::dpa1::degree_weight(k, degree_gain_raw); + value = fmaf(accumulator[k][group] * weight, axis_value, value); + } + if (lane < 16) { + output[channel * axis + axis_channel] = value; + } + } + if (write_rotation && lane < 16) { + float* rotation_row = rotation + (node * Width + channel) * 3; + rotation_row[0] = accumulator[1][group]; + rotation_row[1] = accumulator[2][group]; + rotation_row[2] = accumulator[3][group]; + } + } + } else { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 32 + lane; + for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { + float value = 0.0f; +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + const float axis_value = + __shfl_sync(kWarpMask, accumulator[k][0], axis_channel); + const float weight = + BasisDim == 4 ? 1.0f + : deepmd::dpa1::degree_weight(k, degree_gain_raw); + value = fmaf(accumulator[k][group] * weight, axis_value, value); + } + if (channel < Width) { + output[channel * axis + axis_channel] = value; + } + } + if (write_rotation && channel < Width) { + float* rotation_row = rotation + (node * Width + channel) * 3; + rotation_row[0] = accumulator[1][group]; + rotation_row[1] = accumulator[2][group]; + rotation_row[2] = accumulator[3][group]; + } + } + } + if (concatenate_type_embedding) { + for (int channel = lane; channel < type_embedding_dim; channel += 32) { + output[Width * axis + channel] = + type_embedding[static_cast(center_type) * type_embedding_dim + + channel]; + } + } +} + +template +__global__ +__launch_bounds__(kThreads, MinimumBlocks) void compressed_backward_kernel( + long node_count, + long edge_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + int descriptor_stride, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const float* __restrict__ descriptor_gradient, + const float* __restrict__ rotation_gradient, + const float* __restrict__ moment, + const float* __restrict__ edge_vec, + const index_t* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const index_t* __restrict__ destination_order, + const long* __restrict__ destination_row_ptr, + const long* __restrict__ atype, + const float* __restrict__ average, + const float* __restrict__ inverse_stddev, + const float* __restrict__ table, + const float* __restrict__ gate_table, + const float* __restrict__ degree_gain_raw, + float* __restrict__ edge_gradient) { + constexpr unsigned kWarpMask = 0xffffffffu; + constexpr int kGradientGroups = ChannelPolicy::gradient_groups; + constexpr int kEdgeGroups = ChannelPolicy::accumulation_groups; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int warps_per_block = blockDim.x / kWarpSize; + const long node = static_cast(blockIdx.x) * warps_per_block + warp; + if (node >= node_count) { + return; + } + + float gradient[BasisDim][kGradientGroups] = {}; + const long moment_base = node * BasisDim * Width; + const float* node_descriptor_gradient = + descriptor_gradient + node * descriptor_stride; + float axis_moment[BasisDim] = {}; + if (lane < Width) { +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + axis_moment[k] = __ldg(moment + moment_base + k * Width + lane); + } + } + +#pragma unroll + for (int group = 0; group < kGradientGroups; ++group) { + const int channel = group * 32 + lane; + for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { + float axis_value[BasisDim]; +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + axis_value[k] = __shfl_sync(kWarpMask, axis_moment[k], axis_channel); + } + if (channel < Width) { + const float value = + __ldg(node_descriptor_gradient + channel * axis + axis_channel); +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + gradient[k][group] = fmaf(value, axis_value[k], gradient[k][group]); + } + } + } + if (channel < axis) { + for (int input = 0; input < Width; ++input) { + const float value = + __ldg(node_descriptor_gradient + input * axis + channel); +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + gradient[k][group] = + fmaf(value, __ldg(moment + moment_base + k * Width + input), + gradient[k][group]); + } + } + } + if constexpr (BasisDim > 4) { +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + gradient[k][group] *= deepmd::dpa1::degree_weight(k, degree_gain_raw); + } + } + if (rotation_gradient != nullptr && channel < Width) { + const long rotation_offset = (node * Width + channel) * 3; + gradient[1][group] += __ldg(rotation_gradient + rotation_offset + 0); + gradient[2][group] += __ldg(rotation_gradient + rotation_offset + 1); + gradient[3][group] += __ldg(rotation_gradient + rotation_offset + 2); + } + } + + int center_type = lane == 0 ? static_cast(atype[node]) : 0; + center_type = __shfl_sync(kWarpMask, center_type, 0); + const float inverse_stddev0 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 0); + const float inverse_stddev1 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 1); + const float inverse_stddev2 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 2); + const float inverse_stddev3 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 3); + const long begin = destination_row_ptr[node]; + const long end = destination_row_ptr[node + 1]; + + if constexpr (ChannelPolicy::use_half_warp) { + const int half = lane >> 4; + const int half_lane = lane & 15; + const int leader = half * 16; + const unsigned half_mask = half == 0 ? 0x0000ffffu : 0xffff0000u; + float edge_moment_gradient[BasisDim][kEdgeGroups] = {}; +#pragma unroll + for (int group = 0; group < kEdgeGroups; ++group) { + const int channel = group * 16 + half_lane; + const int owner_group = channel / 32; + const int owner_lane = channel & 31; +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + edge_moment_gradient[k][group] = + __shfl_sync(kWarpMask, gradient[k][owner_group], owner_lane); + } + } + + for (long position = begin + half; position < end; position += 2) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + if (half_lane == 0) { + edge_gradient[edge * 3 + 0] = 0.0f; + edge_gradient[edge * 3 + 1] = 0.0f; + edge_gradient[edge * 3 + 2] = 0.0f; + } + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (half_lane == 0) { + environment = load_environment( + edge, center_type, ntypes, one_side, rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, atype, average, + inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = + broadcast_environment(environment, leader, half_mask); + location = broadcast_location(location, leader, half_mask); + float partial_basis[BasisDim] = {}; + float partial_radial = 0.0f; + float partial_switch = 0.0f; +#pragma unroll + for (int group = 0; group < kEdgeGroups; ++group) { + const int channel = group * 16 + half_lane; + float descriptor_product = 0.0f; +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + descriptor_product = + fmaf(environment.basis[k], edge_moment_gradient[k][group], + descriptor_product); + } + const float2 table_value = + evaluate_table_backward(table, location, channel, Width); + const float gate = + __ldg(gate_table + + static_cast(environment.pair_index) * Width + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value.x * (1.0f + effective_gate); + if (smooth) { + partial_switch = + fmaf(descriptor_product * table_value.x, gate, partial_switch); + } + partial_radial = fmaf(descriptor_product * (1.0f + effective_gate), + table_value.y, partial_radial); +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + partial_basis[k] = + fmaf(embedding, edge_moment_gradient[k][group], partial_basis[k]); + } + } +#pragma unroll + for (int offset = 8; offset > 0; offset >>= 1) { +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + partial_basis[k] += + __shfl_down_sync(half_mask, partial_basis[k], offset, 16); + } + partial_radial += + __shfl_down_sync(half_mask, partial_radial, offset, 16); + partial_switch += + __shfl_down_sync(half_mask, partial_switch, offset, 16); + } + if (half_lane == 0) { + store_edge_gradient( + edge, environment, partial_basis, partial_radial, partial_switch, + inverse_neighbors, inverse_stddev0, inverse_stddev1, + inverse_stddev2, inverse_stddev3, rcut, rcut_smooth, protection, + edge_gradient); + } + } + } else { + for (long position = begin; position < end; ++position) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + if (lane == 0) { + edge_gradient[edge * 3 + 0] = 0.0f; + edge_gradient[edge * 3 + 1] = 0.0f; + edge_gradient[edge * 3 + 2] = 0.0f; + } + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (lane == 0) { + environment = load_environment( + edge, center_type, ntypes, one_side, rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, atype, average, + inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = broadcast_environment(environment, 0, kWarpMask); + location = broadcast_location(location, 0, kWarpMask); + float partial_basis[BasisDim] = {}; + float partial_radial = 0.0f; + float partial_switch = 0.0f; +#pragma unroll + for (int group = 0; group < kGradientGroups; ++group) { + const int channel = group * 32 + lane; + if (channel < Width) { + float descriptor_product = 0.0f; +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + descriptor_product = fmaf(environment.basis[k], gradient[k][group], + descriptor_product); + } + const float2 table_value = + evaluate_table_backward(table, location, channel, Width); + const float gate = __ldg( + gate_table + static_cast(environment.pair_index) * Width + + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value.x * (1.0f + effective_gate); + if (smooth) { + partial_switch = + fmaf(descriptor_product * table_value.x, gate, partial_switch); + } + partial_radial = fmaf(descriptor_product * (1.0f + effective_gate), + table_value.y, partial_radial); +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + partial_basis[k] = + fmaf(embedding, gradient[k][group], partial_basis[k]); + } + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int k = 0; k < BasisDim; ++k) { + partial_basis[k] += + __shfl_down_sync(kWarpMask, partial_basis[k], offset); + } + partial_radial += __shfl_down_sync(kWarpMask, partial_radial, offset); + partial_switch += __shfl_down_sync(kWarpMask, partial_switch, offset); + } + if (lane == 0) { + store_edge_gradient( + edge, environment, partial_basis, partial_radial, partial_switch, + inverse_neighbors, inverse_stddev0, inverse_stddev1, + inverse_stddev2, inverse_stddev3, rcut, rcut_smooth, protection, + edge_gradient); + } + } + } +} + +template +__global__ void zero_padding_kernel( + long node_count, + long edge_count, + const index_t* __restrict__ destination_order, + const long* __restrict__ destination_row_ptr, + float* __restrict__ edge_gradient) { + const long valid_edge_count = destination_row_ptr[node_count]; + for (long position = valid_edge_count + blockIdx.x * blockDim.x + threadIdx.x; + position < edge_count; + position += static_cast(blockDim.x) * gridDim.x) { + const long edge = + edge_at_csr_position(position, destination_order); + edge_gradient[edge * 3 + 0] = 0.0f; + edge_gradient[edge * 3 + 1] = 0.0f; + edge_gradient[edge * 3 + 2] = 0.0f; + } +} + +template +void launch_forward_variant(const Arguments& arguments, + long node_count, + int threads, + cudaStream_t stream) { + const int warps_per_block = threads / kWarpSize; + const int blocks = + static_cast((node_count + warps_per_block - 1) / warps_per_block); + compressed_forward_kernel<<>>( + node_count, arguments.type_count, arguments.one_side, arguments.smooth, + arguments.axis, arguments.concatenate_type_embedding, + arguments.write_rotation, arguments.type_embedding_dim, arguments.rcut, + arguments.rcut_smooth, arguments.protection, arguments.inverse_neighbors, + arguments.lower, arguments.upper, arguments.table_max, arguments.stride0, + arguments.stride1, arguments.edge_vec, + static_cast(arguments.edge_index), arguments.edge_mask, + static_cast(arguments.destination_order), + arguments.destination_row_ptr, arguments.atype, arguments.type_embedding, + arguments.average, arguments.inverse_stddev, arguments.table, + arguments.gate_table, arguments.degree_gain, arguments.descriptor, + arguments.rotation, arguments.moment_out); +} + +template +cudaError_t launch_forward(const Arguments& arguments, cudaStream_t stream) { + const DeviceProperties properties = {arguments.device_major, + arguments.multiprocessor_count}; + const TuningKey key = { + arguments.device, + static_cast(KernelDirection::kForward), + Width, + BasisDim, + arguments.axis, + Canonical ? 1 : 0, + static_cast(sizeof(index_t)), + (arguments.one_side ? 1 : 0) | (arguments.smooth ? 2 : 0) | + (arguments.concatenate_type_embedding ? 4 : 0) | + (arguments.write_rotation ? 8 : 0) | (Masked ? 16 : 0), + arguments.concatenate_type_embedding ? arguments.type_embedding_dim : 0, + type_count_class(arguments.type_count), + workload_size_class(arguments.node_count, + properties.multiprocessor_count), + workload_degree_class(arguments.node_count, arguments.edge_count), + }; + const auto launch = [&](const LaunchConfig& config, long count) { + if (config.resource == ResourcePolicy::kOccupancy) { + launch_forward_variant( + arguments, count, config.threads, stream); + } else { + launch_forward_variant( + arguments, count, config.threads, stream); + } + }; + const LaunchConfig config = select_launch_config( + key, properties, arguments.node_count, stream, launch); + launch(config, arguments.node_count); + return cudaGetLastError(); +} + +template +void launch_backward_variant(const Arguments& arguments, + long node_count, + int threads, + cudaStream_t stream) { + const int warps_per_block = threads / kWarpSize; + const int blocks = + static_cast((node_count + warps_per_block - 1) / warps_per_block); + compressed_backward_kernel<<>>( + node_count, arguments.edge_count, arguments.type_count, + arguments.one_side, arguments.smooth, arguments.axis, + arguments.descriptor_stride, arguments.rcut, arguments.rcut_smooth, + arguments.protection, arguments.inverse_neighbors, arguments.lower, + arguments.upper, arguments.table_max, arguments.stride0, + arguments.stride1, arguments.descriptor_gradient, + arguments.rotation_gradient, arguments.moment, arguments.edge_vec, + static_cast(arguments.edge_index), arguments.edge_mask, + static_cast(arguments.destination_order), + arguments.destination_row_ptr, arguments.atype, arguments.average, + arguments.inverse_stddev, arguments.table, arguments.gate_table, + arguments.degree_gain, arguments.edge_gradient); +} + +template +cudaError_t launch_backward(const Arguments& arguments, cudaStream_t stream) { + const DeviceProperties properties = {arguments.device_major, + arguments.multiprocessor_count}; + const TuningKey key = { + arguments.device, + static_cast(KernelDirection::kBackward), + Width, + BasisDim, + arguments.axis, + Canonical ? 1 : 0, + static_cast(sizeof(index_t)), + (arguments.one_side ? 1 : 0) | (arguments.smooth ? 2 : 0) | + (arguments.rotation_gradient != nullptr ? 8 : 0) | (Masked ? 16 : 0), + arguments.descriptor_stride, + type_count_class(arguments.type_count), + workload_size_class(arguments.node_count, + properties.multiprocessor_count), + workload_degree_class(arguments.node_count, arguments.edge_count), + }; + const auto launch = [&](const LaunchConfig& config, long count) { + if (config.resource == ResourcePolicy::kOccupancy) { + launch_backward_variant( + arguments, count, config.threads, stream); + } else { + launch_backward_variant( + arguments, count, config.threads, stream); + } + }; + const LaunchConfig config = select_launch_config( + key, properties, arguments.node_count, stream, launch); + launch(config, arguments.node_count); + const cudaError_t backward_error = cudaGetLastError(); + if (backward_error != cudaSuccess) { + return backward_error; + } + zero_padding_kernel<<<1, kThreads, 0, stream>>>( + arguments.node_count, arguments.edge_count, + static_cast(arguments.destination_order), + arguments.destination_row_ptr, arguments.edge_gradient); + return cudaGetLastError(); +} + +template +cudaError_t dispatch_forward_topology(const Arguments& arguments, + cudaStream_t stream) { + if (arguments.canonical && !arguments.masked) { + return launch_forward(arguments, + stream); + } + if (arguments.canonical) { + return launch_forward(arguments, + stream); + } + return launch_forward(arguments, + stream); +} + +template +cudaError_t dispatch_backward_topology(const Arguments& arguments, + cudaStream_t stream) { + if (arguments.canonical && !arguments.masked) { + return launch_backward(arguments, + stream); + } + if (arguments.canonical) { + return launch_backward(arguments, + stream); + } + return launch_backward(arguments, + stream); +} + +template +cudaError_t dispatch_forward_basis(const Arguments& arguments, + cudaStream_t stream) { + if (arguments.basis_dim == 4) { + return dispatch_forward_topology(arguments, stream); + } +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + if (arguments.basis_dim == 9) { + return dispatch_forward_topology(arguments, stream); + } + if constexpr (Width >= 16 && Width <= 128) { + if (arguments.basis_dim == 16) { + return dispatch_forward_topology(arguments, stream); + } + if (arguments.basis_dim == 25) { + return dispatch_forward_topology(arguments, stream); + } + } +#endif + return cudaErrorInvalidValue; +} + +template +cudaError_t dispatch_backward_basis(const Arguments& arguments, + cudaStream_t stream) { + if (arguments.basis_dim == 4) { + return dispatch_backward_topology(arguments, stream); + } +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + if (arguments.basis_dim == 9) { + return dispatch_backward_topology(arguments, stream); + } + if constexpr (Width >= 16 && Width <= 128) { + if (arguments.basis_dim == 16) { + return dispatch_backward_topology(arguments, stream); + } + if (arguments.basis_dim == 25) { + return dispatch_backward_topology(arguments, stream); + } + } +#endif + return cudaErrorInvalidValue; +} + +template +cudaError_t dispatch_forward_index(const Arguments& arguments, + cudaStream_t stream) { + if (arguments.index_kind == IndexKind::kInt32) { + return dispatch_forward_basis(arguments, stream); + } + return dispatch_forward_basis(arguments, stream); +} + +template +cudaError_t dispatch_backward_index(const Arguments& arguments, + cudaStream_t stream) { + if (arguments.index_kind == IndexKind::kInt32) { + return dispatch_backward_basis(arguments, stream); + } + return dispatch_backward_basis(arguments, stream); +} + +} // namespace + +#define DPA1_COMPRESS_DEFINE_CHANNEL(width) \ + cudaError_t launch_forward_c##width(const Arguments& arguments, \ + cudaStream_t stream) { \ + return dispatch_forward_index(arguments, stream); \ + } \ + cudaError_t launch_backward_c##width(const Arguments& arguments, \ + cudaStream_t stream) { \ + return dispatch_backward_index(arguments, stream); \ + } + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_launch.h b/source/op/pt/dpa1_graph_compress_launch.h new file mode 100644 index 0000000000..40e4e12e92 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_launch.h @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Host-side launch interface for the geometrically compressed DPA1 CUDA +// descriptor. +// +// The Torch-facing translation unit validates tensors and converts them into +// this plain pointer bundle. Each supported channel width owns one CUDA +// translation unit that instantiates the corresponding forward and backward +// kernels. Keeping this interface free of ATen headers avoids reparsing the +// PyTorch C++ API in every specialization. + +#pragma once + +#include + +namespace deepmd_dpa1_compress { + +enum class IndexKind : int { + kInt32 = 0, + kInt64 = 1, +}; + +struct Arguments { + long node_count = 0; + long edge_count = 0; + int device = 0; + int device_major = 0; + int multiprocessor_count = 0; + int basis_dim = 4; + int type_count = 0; + int axis = 0; + int type_embedding_dim = 0; + int descriptor_stride = 0; + bool one_side = false; + bool smooth = false; + bool canonical = false; + bool masked = false; + bool concatenate_type_embedding = false; + bool write_rotation = false; + IndexKind index_kind = IndexKind::kInt64; + + float rcut = 0.0f; + float rcut_smooth = 0.0f; + float protection = 0.0f; + float inverse_neighbors = 0.0f; + float lower = 0.0f; + float upper = 0.0f; + float table_max = 0.0f; + float stride0 = 0.0f; + float stride1 = 0.0f; + + const float* edge_vec = nullptr; + const void* edge_index = nullptr; + const bool* edge_mask = nullptr; + const void* destination_order = nullptr; + const long* destination_row_ptr = nullptr; + const long* atype = nullptr; + const float* type_embedding = nullptr; + const float* average = nullptr; + const float* inverse_stddev = nullptr; + const float* degree_gain = nullptr; + const float* table = nullptr; + const float* gate_table = nullptr; + + const float* descriptor_gradient = nullptr; + const float* rotation_gradient = nullptr; + const float* moment = nullptr; + + float* descriptor = nullptr; + float* rotation = nullptr; + float* moment_out = nullptr; + float* edge_gradient = nullptr; +}; + +#define DPA1_COMPRESS_FOR_EACH_CHANNEL(macro) \ + macro(8) macro(16) macro(32) macro(64) macro(128) macro(256) + +#define DPA1_COMPRESS_DECLARE_CHANNEL(width) \ + cudaError_t launch_forward_c##width(const Arguments& arguments, \ + cudaStream_t stream); \ + cudaError_t launch_backward_c##width(const Arguments& arguments, \ + cudaStream_t stream); + +DPA1_COMPRESS_FOR_EACH_CHANNEL(DPA1_COMPRESS_DECLARE_CHANNEL) + +#undef DPA1_COMPRESS_DECLARE_CHANNEL + +} // namespace deepmd_dpa1_compress diff --git a/source/op/pt/dpa1_graph_compress_tuning.h b/source/op/pt/dpa1_graph_compress_tuning.h index 304ee79093..f2a4f36bf8 100644 --- a/source/op/pt/dpa1_graph_compress_tuning.h +++ b/source/op/pt/dpa1_graph_compress_tuning.h @@ -5,7 +5,6 @@ #pragma once #include -#include #include #include @@ -31,10 +30,16 @@ struct LaunchConfig { int threads; }; +struct DeviceProperties { + int major; + int multiprocessor_count; +}; + struct TuningKey { int device; int direction; int width; + int basis_dim; int axis; int canonical; int index_bytes; @@ -46,8 +51,9 @@ struct TuningKey { bool operator==(const TuningKey& other) const { return device == other.device && direction == other.direction && - width == other.width && axis == other.axis && - canonical == other.canonical && index_bytes == other.index_bytes && + width == other.width && basis_dim == other.basis_dim && + axis == other.axis && canonical == other.canonical && + index_bytes == other.index_bytes && model_flags == other.model_flags && model_stride == other.model_stride && type_class == other.type_class && size_class == other.size_class && @@ -58,10 +64,10 @@ struct TuningKey { struct TuningKeyHash { std::size_t operator()(const TuningKey& key) const { std::size_t value = 0; - const std::array fields = { - key.device, key.direction, key.width, key.axis, - key.canonical, key.index_bytes, key.model_flags, key.model_stride, - key.type_class, key.size_class, key.degree_class}; + const std::array fields = { + key.device, key.direction, key.width, key.basis_dim, + key.axis, key.canonical, key.index_bytes, key.model_flags, + key.model_stride, key.type_class, key.size_class, key.degree_class}; for (const int field : fields) { value ^= std::hash{}(field) + 0x9e3779b9 + (value << 6) + (value >> 2); @@ -81,23 +87,9 @@ tuning_cache() { return cache; } -inline const cudaDeviceProp& device_properties(int device) { - constexpr int kMaximumCachedDevices = 64; - TORCH_CHECK(device >= 0 && device < kMaximumCachedDevices, - "dpa1_graph_compress: unsupported CUDA device index ", device); - static std::array initialized; - static std::array properties; - std::call_once(initialized[device], [device] { - TORCH_CHECK( - cudaGetDeviceProperties(&properties[device], device) == cudaSuccess, - "dpa1_graph_compress: cannot query CUDA device properties"); - }); - return properties[device]; -} - -inline LaunchConfig architecture_fallback(const cudaDeviceProp& properties) { +inline LaunchConfig architecture_fallback(const DeviceProperties& properties) { if (properties.major >= 9) { - if (properties.multiProcessorCount <= 80) { + if (properties.multiprocessor_count <= 80) { return {ResourcePolicy::kOccupancy, 256}; } return {ResourcePolicy::kBalanced, 256}; @@ -145,7 +137,7 @@ inline int type_count_class(int type_count) { template LaunchConfig select_launch_config(const TuningKey& key, - const cudaDeviceProp& properties, + const DeviceProperties& properties, long node_count, cudaStream_t stream, const LaunchFunction& launch) { @@ -172,7 +164,7 @@ LaunchConfig select_launch_config(const TuningKey& key, const long sample_node_count = std::min( node_count, - std::max(4096L, static_cast(properties.multiProcessorCount) * 64)); + std::max(4096L, static_cast(properties.multiprocessor_count) * 64)); constexpr std::array kCandidates = {{ {ResourcePolicy::kBalanced, 128}, {ResourcePolicy::kBalanced, 256}, diff --git a/source/op/pt/dpa1_graph_descriptor.cu b/source/op/pt/dpa1_graph_descriptor.cu index 4bf2d28148..ed0c3d901f 100644 --- a/source/op/pt/dpa1_graph_descriptor.cu +++ b/source/op/pt/dpa1_graph_descriptor.cu @@ -88,6 +88,15 @@ #include "dpa1_graph_common.cuh" +#ifndef DEEPMD_ENABLE_DPA1_HIGH_LMAX +#define DEEPMD_ENABLE_DPA1_HIGH_LMAX 0 +#endif + +// Degree-two and higher implementations remain available for future +// experiments, but production builds intentionally omit their template +// instantiations. DPA1 deployment currently uses lmax=1; the CMake option +// restores the retained kernels when that contract changes. + namespace { // Activation codes follow deepmd.kernels.triton.dpa1.activation.ACT_CODES: @@ -160,9 +169,9 @@ __global__ void pair_table_kernel(int n_pairs, // the gather streams coalesced 16-byte words. Strip mode has no per-pair // layer-1 term (the type embedding enters through the gate instead), so the // table degenerates to its single bias row. -template +template DEV_INLINE void layer1_tile(int tid, - const EdgeTablesT& T, + const EdgeTablesT& T, int strip, const float* __restrict__ pair_table, const float* __restrict__ w1, @@ -200,7 +209,7 @@ DEV_INLINE void layer1_tile(int tid, // tile rows with all four moment components, one atomicAdd per // (run, component, channel) // ====================================================================== -template +template struct FwdSmem { static constexpr int STRIDE = TILE + 4; // The g tile overlays the (dead) h1 tile; NG >= N1, so the union is sized @@ -211,13 +220,13 @@ struct FwdSmem { float g[NG][STRIDE]; // stages 3-4 (walk view) } u; float h2[N2][STRIDE]; - EdgeTablesT T; + EdgeTablesT T; }; -static_assert(sizeof(FwdSmem<32, 64, 128, 128>) > 96 * 1024); -static_assert(sizeof(FwdSmem<32, 64, 128, 64>) <= 96 * 1024); +static_assert(sizeof(FwdSmem<32, 64, 128, 128, 4>) > 96 * 1024); +static_assert(sizeof(FwdSmem<32, 64, 128, 64, 4>) <= 96 * 1024); -template +template __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_forward_kernel( long n_edge, int ntypes, @@ -253,7 +262,7 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_forward_kernel( constexpr int TILE = 16 * EPT; // edges per tile (16 edge-lanes x EPT) constexpr int STRIDE = TILE + 4; extern __shared__ char smem_raw[]; - auto& S = *reinterpret_cast*>(smem_raw); + auto& S = *reinterpret_cast*>(smem_raw); const int tid = threadIdx.x; const int tx = tid % 16, ty = tid / 16; const int erow = tx * EPT; @@ -264,11 +273,12 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_forward_kernel( const long tile_base = tile * TILE; const int rows = (int)min((long)TILE, n_edge - tile_base); - stage_tile(tid, tile_base, rows, n_edge, ntypes, one_side, rcut, - rcut_smth, protection, inv_nnei, edge_vec, edge_index, - edge_mask, atype, davg, inv_dstd, order, S.T); - layer1_tile(tid, S.T, strip, pair_table, w1, idt1, - S.u.h1); + stage_tile(tid, tile_base, rows, n_edge, ntypes, one_side, + rcut, rcut_smth, protection, inv_nnei, edge_vec, + edge_index, edge_mask, atype, davg, inv_dstd, + order, S.T); + layer1_tile(tid, S.T, strip, pair_table, + w1, idt1, S.u.h1); __syncthreads(); // === Step 2. GEMM2 -> h2 tile; spill pre2 (transposed (N2, E)) === @@ -413,28 +423,21 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_forward_kernel( const int span = re - rb; const int beg = rb + span * slice / kSlices; const int end = rb + span * (slice + 1) / kSlices; - float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f; + float accum[BASIS_DIM] = {}; for (int r = beg; r < end; ++r) { - const float4 rv = *reinterpret_cast(&S.T.rr[r][0]); const float gv = gcol[r]; - a0 = fmaf(rv.x, gv, a0); - a1 = fmaf(rv.y, gv, a1); - a2 = fmaf(rv.z, gv, a2); - a3 = fmaf(rv.w, gv, a3); +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + accum[k] = fmaf(S.T.basis[r][k], gv, accum[k]); + } } if (node >= 0) { - const long base = ((long)node * 4) * NG + c; - if (a0 != 0.f) { - atomicAdd(&gr[base + 0 * NG], a0); - } - if (a1 != 0.f) { - atomicAdd(&gr[base + 1 * NG], a1); - } - if (a2 != 0.f) { - atomicAdd(&gr[base + 2 * NG], a2); - } - if (a3 != 0.f) { - atomicAdd(&gr[base + 3 * NG], a3); + const long base = ((long)node * BASIS_DIM) * NG + c; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + if (accum[k] != 0.f) { + atomicAdd(&gr[base + k * NG], accum[k]); + } } } } @@ -473,21 +476,32 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_forward_kernel( // shared memory // 5. analytic environment backward -> d_edge_vec // ====================================================================== -template +template struct BwdSmem { - static constexpr int kResidentRuns = 4; + static constexpr int kResidentRuns = + BASIS_DIM <= 9 ? 4 : (BASIS_DIM == 16 ? 2 : 1); static constexpr int STRIDE = TILE + 4; - float h2[N2][STRIDE]; // act(pre2) tile; raw dh2 after stage 3 - float x_t[N2][STRIDE]; // dpre3 block tile; dpre2 rows [0, N2) - float drr_banks[8][4][TILE]; // per-warp env-gradient banks + float h2[N2][STRIDE]; // act(pre2) tile; raw dh2 after stage 3 + float x_t[N2][STRIDE]; // dpre3 block tile; dpre2 rows [0, N2) + float drr_banks[8][BASIS_DIM][TILE]; float d_radial[TILE]; float d_sw[TILE]; // strip gate: dE/d(sw) through gate * sw - float dgr_rows[kResidentRuns][4 * NG]; + float dgr_rows[kResidentRuns][BASIS_DIM * NG]; int run_slot[TILE]; // resident dgr slot per row (-1: global read) - EdgeTablesT T; + EdgeTablesT T; }; -template +static_assert(sizeof(BwdSmem<32, 64, 128, 64, 9>) > 64 * 1024); +static_assert(sizeof(BwdSmem<32, 64, 128, 32, 9>) <= 48 * 1024); + +template __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( long n_edge, int ntypes, @@ -523,9 +537,10 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( long e_pad) { constexpr int kBlocks = NG / N2; // channel blocks: 1 (identity) or 2 constexpr int TILE = 16 * EPT; // edges per tile (16 edge-lanes x EPT) - constexpr int kResidentRuns = BwdSmem::kResidentRuns; + constexpr int kResidentRuns = + BwdSmem::kResidentRuns; extern __shared__ char smem_raw[]; - auto& S = *reinterpret_cast*>(smem_raw); + auto& S = *reinterpret_cast*>(smem_raw); const int tid = threadIdx.x; const int tx = tid % 16, ty = tid / 16; const int erow = tx * EPT; @@ -536,15 +551,16 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( const long tile_base = tile * TILE; const int rows = (int)min((long)TILE, n_edge - tile_base); - stage_tile(tid, tile_base, rows, n_edge, ntypes, one_side, rcut, - rcut_smth, protection, inv_nnei, edge_vec, edge_index, - edge_mask, atype, davg, inv_dstd, order, S.T); + stage_tile(tid, tile_base, rows, n_edge, ntypes, one_side, + rcut, rcut_smth, protection, inv_nnei, edge_vec, + edge_index, edge_mask, atype, davg, inv_dstd, + order, S.T); if (tid < TILE) { S.run_slot[tid] = S.T.run_of[tid] < kResidentRuns ? S.T.run_of[tid] : -1; S.d_radial[tid] = 0.f; S.d_sw[tid] = 0.f; } - for (int t = tid; t < 8 * 4 * TILE; t += kThreads) { + for (int t = tid; t < 8 * BASIS_DIM * TILE; t += kThreads) { (&S.drr_banks[0][0][0])[t] = 0.f; } { @@ -554,8 +570,8 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( if (d < 0) { continue; } - for (int t = tid; t < 4 * NG; t += kThreads) { - S.dgr_rows[s][t] = __ldg(dgr + d * 4 * NG + t); + for (int t = tid; t < BASIS_DIM * NG; t += kThreads) { + S.dgr_rows[s][t] = __ldg(dgr + d * BASIS_DIM * NG + t); } } } @@ -602,15 +618,13 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( // Per-edge moment weights, hoisted to registers once per tile (the dgv / // drr / residual-fold loops otherwise re-read the float4 per channel). - float rr_frag[EPT][4]; + float rr_frag[EPT][BASIS_DIM]; #pragma unroll - for (int i = 0; i < EPT; ++i) { - const float4 rv = *reinterpret_cast(&S.T.rr[erow + i][0]); - rr_frag[i][0] = rv.x; - rr_frag[i][1] = rv.y; - rr_frag[i][2] = rv.z; - rr_frag[i][3] = rv.w; - } + for (int i = 0; i < EPT; ++i) +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + rr_frag[i][k] = S.T.basis[erow + i][k]; + } #define RR_(i, k) (rr_frag[i][k]) // === Step 2. N2-channel blocks: moment backward + dgrad3 === @@ -637,10 +651,14 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( const int slot0 = S.run_slot[erow]; const bool uniform_run = S.T.run_of[erow] == S.T.run_of[erow + EPT - 1] && slot0 >= 0; - float dr0[EPT], dr1[EPT], dr2[EPT], dr3[EPT], dsw[EPT]; + float dr[BASIS_DIM][EPT], dsw[EPT]; #pragma unroll for (int i = 0; i < EPT; ++i) { - dr0[i] = dr1[i] = dr2[i] = dr3[i] = dsw[i] = 0.f; + dsw[i] = 0.f; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dr[k][i] = 0.f; + } } // Optional prefetch (PIPE): issue the group's four channel fragments // of the g / pre3 spill together so their (L2-missing) global @@ -674,12 +692,18 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( } if (uniform_run) { const float* row = S.dgr_rows[slot0]; - const float d0 = row[0 * NG + c], d1 = row[1 * NG + c]; - const float d2 = row[2 * NG + c], d3 = row[3 * NG + c]; + float dval[BASIS_DIM]; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dval[k] = row[k * NG + c]; + } #pragma unroll for (int i = 0; i < EPT; ++i) { - dgv[i] = RR_(i, 0) * d0 + RR_(i, 1) * d1 + RR_(i, 2) * d2 + - RR_(i, 3) * d3; + dgv[i] = 0.f; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dgv[i] = fmaf(RR_(i, k), dval[k], dgv[i]); + } } #pragma unroll for (int i = 0; i < EPT; ++i) { @@ -720,10 +744,10 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( dgv[i] *= 1.f + geff; } dp[i] = dgv[i] * is * aprime; - dr0[i] = fmaf(gg, d0, dr0[i]); - dr1[i] = fmaf(gg, d1, dr1[i]); - dr2[i] = fmaf(gg, d2, dr2[i]); - dr3[i] = fmaf(gg, d3, dr3[i]); +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dr[k][i] = fmaf(gg, dval[k], dr[k][i]); + } } } else { #pragma unroll @@ -732,11 +756,14 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( const int slot = S.run_slot[e]; const float* row = slot >= 0 ? S.dgr_rows[slot] - : dgr + (long)max(S.T.dst[e], 0) * 4 * NG; - const float d0 = row[0 * NG + c], d1 = row[1 * NG + c]; - const float d2 = row[2 * NG + c], d3 = row[3 * NG + c]; - dgv[i] = RR_(i, 0) * d0 + RR_(i, 1) * d1 + RR_(i, 2) * d2 + - RR_(i, 3) * d3; + : dgr + (long)max(S.T.dst[e], 0) * BASIS_DIM * NG; + float dval[BASIS_DIM]; + dgv[i] = 0.f; +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dval[k] = row[k * NG + c]; + dgv[i] = fmaf(RR_(i, k), dval[k], dgv[i]); + } float gval, aprime; if constexpr (ACT == 0) { gval = gv[i]; @@ -770,10 +797,10 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( dgv[i] *= 1.f + geff; } dp[i] = dgv[i] * is * aprime; - dr0[i] = fmaf(gg, d0, dr0[i]); - dr1[i] = fmaf(gg, d1, dr1[i]); - dr2[i] = fmaf(gg, d2, dr2[i]); - dr3[i] = fmaf(gg, d3, dr3[i]); +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dr[k][i] = fmaf(gg, dval[k], dr[k][i]); + } } } // Layer-3 residual fold: dh2[c mod N2] accumulates dgs over the @@ -793,25 +820,20 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( // fragment), then accumulate race-free into this warp's bank. #pragma unroll for (int i = 0; i < EPT; ++i) { - dr0[i] += __shfl_xor_sync(0xffffffffu, dr0[i], 16); - dr1[i] += __shfl_xor_sync(0xffffffffu, dr1[i], 16); - dr2[i] += __shfl_xor_sync(0xffffffffu, dr2[i], 16); - dr3[i] += __shfl_xor_sync(0xffffffffu, dr3[i], 16); +#pragma unroll + for (int k = 0; k < BASIS_DIM; ++k) { + dr[k][i] += __shfl_xor_sync(0xffffffffu, dr[k][i], 16); + } dsw[i] += __shfl_xor_sync(0xffffffffu, dsw[i], 16); } if ((ty & 1) == 0) { const int w = tid >> 5; - float* b0 = &S.drr_banks[w][0][erow]; - float* b1 = &S.drr_banks[w][1][erow]; - float* b2v = &S.drr_banks[w][2][erow]; - float* b3v = &S.drr_banks[w][3][erow]; #pragma unroll - for (int i = 0; i < EPT; ++i) { - b0[i] += dr0[i]; - b1[i] += dr1[i]; - b2v[i] += dr2[i]; - b3v[i] += dr3[i]; - } + for (int k = 0; k < BASIS_DIM; ++k) +#pragma unroll + for (int i = 0; i < EPT; ++i) { + S.drr_banks[w][k][erow + i] += dr[k][i]; + } if (strip && smooth) { #pragma unroll for (int i = 0; i < EPT; ++i) { @@ -962,32 +984,19 @@ __global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( const float y = edge_vec[(long)e * 3 + 1]; const float z = edge_vec[(long)e * 3 + 2]; const float len = sqrtf(x * x + y * y + z * z); - const float q = len + protection; const float sw = switch_val(len, rcut_smth, rcut); const float dsw = switch_deriv(len, rcut_smth, rcut); const float* isd = inv_dstd + (long)atype[S.T.dst[tid]] * 4; - float dr[4] = {0.f, 0.f, 0.f, 0.f}; + float d_basis[BASIS_DIM] = {}; #pragma unroll for (int w = 0; w < 8; ++w) #pragma unroll - for (int k = 0; k < 4; ++k) { - dr[k] += S.drr_banks[w][k][tid]; + for (int k = 0; k < BASIS_DIM; ++k) { + d_basis[k] += S.drr_banks[w][k][tid]; } - const float g0 = (dr[0] * inv_nnei + S.d_radial[tid]) * isd[0]; - const float gx = dr[1] * inv_nnei * isd[1]; - const float gy = dr[2] * inv_nnei * isd[2]; - const float gz = dr[3] * inv_nnei * isd[3]; - const float inv_len = len > 0.f ? 1.f / len : 0.f; - const float rq = 1.f / q; - const float gdot = gx * x + gy * y + gz * z; - const float coef = - (g0 * rq * (dsw - sw * rq) + - gdot * rq * rq * (dsw - 2.f * sw * rq) + S.d_sw[tid] * dsw) * - inv_len; - const float s2 = sw * rq * rq; - d_edge_vec[(long)e * 3 + 0] = coef * x + s2 * gx; - d_edge_vec[(long)e * 3 + 1] = coef * y + s2 * gy; - d_edge_vec[(long)e * 3 + 2] = coef * z + s2 * gz; + moment_basis_edge_gradient( + d_basis, S.d_radial[tid], S.d_sw[tid], x, y, z, len, protection, sw, + dsw, inv_nnei, isd, d_edge_vec + (long)e * 3); } } __syncthreads(); @@ -1072,14 +1081,15 @@ struct LaunchArgs { cudaStream_t stream; }; -template +template void launch_forward(const LaunchArgs& a, torch::Tensor& gr, torch::Tensor& pre2_saved, torch::Tensor& g_saved) { constexpr int TILE = 16 * EPT; - auto kernel = dpa1_graph_forward_kernel; - const size_t smem = sizeof(FwdSmem); + auto kernel = + dpa1_graph_forward_kernel; + const size_t smem = sizeof(FwdSmem); const cudaError_t attribute_error = cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); TORCH_CHECK( @@ -1101,7 +1111,7 @@ void launch_forward(const LaunchArgs& a, DPA1_CHECK_LAUNCH("dpa1_graph_descriptor forward"); } -template +template void launch_forward_portable(const LaunchArgs& a, torch::Tensor& gr, torch::Tensor& pre2_saved, @@ -1109,10 +1119,11 @@ void launch_forward_portable(const LaunchArgs& a, constexpr int kWideEdgesPerThread = 8; constexpr int kNarrowEdgesPerThread = 4; constexpr int kWideTile = 16 * kWideEdgesPerThread; - constexpr size_t kWideSharedMemory = sizeof(FwdSmem); + constexpr size_t kWideSharedMemory = + sizeof(FwdSmem); constexpr size_t kPortableSharedMemoryFloor = 48 * 1024; if constexpr (kWideSharedMemory <= kPortableSharedMemoryFloor) { - launch_forward( + launch_forward( a, gr, pre2_saved, g_saved); return; } @@ -1120,30 +1131,38 @@ void launch_forward_portable(const LaunchArgs& a, const size_t device_limit = std::max(properties->sharedMemPerBlock, properties->sharedMemPerBlockOptin); if (kWideSharedMemory <= device_limit) { - launch_forward( + launch_forward( a, gr, pre2_saved, g_saved); } else { constexpr int kNarrowTile = 16 * kNarrowEdgesPerThread; constexpr size_t kNarrowSharedMemory = - sizeof(FwdSmem); + sizeof(FwdSmem); TORCH_CHECK(kNarrowSharedMemory <= device_limit, "dpa1_graph_descriptor forward requires ", kNarrowSharedMemory, " bytes of dynamic shared memory, but the device supports ", device_limit); - launch_forward( + launch_forward( a, gr, pre2_saved, g_saved); } } -template +template void launch_backward(const LaunchArgs& a, const torch::Tensor& dgr, const torch::Tensor& pre2_saved, const torch::Tensor& g_saved, torch::Tensor& d_edge_vec) { constexpr int TILE = 16 * EPT; - auto kernel = dpa1_graph_backward_kernel; - const size_t smem = sizeof(BwdSmem); + auto kernel = + dpa1_graph_backward_kernel; + const size_t smem = sizeof(BwdSmem); const cudaError_t attribute_error = cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); TORCH_CHECK( @@ -1166,13 +1185,52 @@ void launch_backward(const LaunchArgs& a, DPA1_CHECK_LAUNCH("dpa1_graph_descriptor backward"); } -// The backward uses four edges per thread once N1 reaches 16 to bound the -// register footprint of widening stacks. The N1 == 8 specialisation keeps -// eight edges per thread to avoid doubling the tile count. NG >= 128 -// prefetches saved g / pre3 rows into the available registers. +// The four-row backward keeps eight edges per thread for N1 == 8 and otherwise +// uses four. Nine-row specializations use four to bound register pressure and +// fall back to two only when the device cannot host the preferred shared tile. +// Four-row NG >= 128 specializations prefetch saved g / pre3 fragments. constexpr int backward_edges_per_thread(int n1) { return n1 >= 16 ? 4 : 8; } constexpr int backward_prefetch(int ng) { return ng >= 128 ? 1 : 0; } +template +void launch_backward_portable(const LaunchArgs& a, + const torch::Tensor& dgr, + const torch::Tensor& pre2_saved, + const torch::Tensor& g_saved, + torch::Tensor& d_edge_vec) { + if constexpr (BASIS_DIM == 4) { + constexpr int kEdgesPerThread = backward_edges_per_thread(N1); + constexpr int kPrefetch = backward_prefetch(NG); + launch_backward(a, dgr, pre2_saved, g_saved, d_edge_vec); + } else { + constexpr int kPreferredEdgesPerThread = BASIS_DIM == 25 ? 2 : 4; + constexpr int kPreferredTile = 16 * kPreferredEdgesPerThread; + constexpr size_t kPreferredSharedMemory = + sizeof(BwdSmem); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + const size_t device_limit = std::max(properties->sharedMemPerBlock, + properties->sharedMemPerBlockOptin); + if (kPreferredSharedMemory <= device_limit) { + launch_backward(a, dgr, pre2_saved, g_saved, d_edge_vec); + return; + } + + constexpr int kFallbackEdgesPerThread = 2; + constexpr int kFallbackTile = 16 * kFallbackEdgesPerThread; + constexpr size_t kFallbackSharedMemory = + sizeof(BwdSmem); + TORCH_CHECK(kFallbackSharedMemory <= device_limit, + "dpa1_graph_descriptor backward requires at least ", + kFallbackSharedMemory, + " bytes of dynamic shared memory, but the device supports ", + device_limit); + launch_backward(a, dgr, pre2_saved, g_saved, d_edge_vec); + } +} + } // namespace // Width / activation instantiation table shared by the two entry points: @@ -1239,6 +1297,7 @@ dpa1_graph_descriptor(torch::Tensor edge_vec, torch::Tensor type_embedding, torch::Tensor davg, torch::Tensor dstd, + torch::Tensor degree_gain, torch::Tensor w1, torch::Tensor b1, torch::Tensor idt1, @@ -1260,7 +1319,8 @@ dpa1_graph_descriptor(torch::Tensor edge_vec, double rcut, double rcut_smth, double protection, - double nnei) { + double nnei, + int64_t basis_dim) { const auto widths = check_widths(w1, w2, w3); const long n_edge = edge_vec.size(0); const long n_node = atype.size(0); @@ -1285,6 +1345,22 @@ dpa1_graph_descriptor(torch::Tensor edge_vec, "dpa1_graph_descriptor: edge_mask must be bool"); TORCH_CHECK(act == 0 || act == 1, "dpa1_graph_descriptor: act must be 0 (tanh) or 1 (silu)"); +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + TORCH_CHECK(basis_dim == 4 || basis_dim == 9, + "dpa1_graph_descriptor: basis_dim must be 4 or 9"); +#else + TORCH_CHECK(basis_dim == 4, + "dpa1_graph_descriptor: this build instantiates only lmax=1; " + "rebuild with DEEPMD_ENABLE_DPA1_HIGH_LMAX=ON for lmax=2"); +#endif + const int degree_gain_size = basis_dim == 4 ? 0 : 1; + TORCH_CHECK( + degree_gain.is_contiguous() && + degree_gain.scalar_type() == torch::kFloat32 && + degree_gain.numel() == degree_gain_size, + "dpa1_graph_descriptor: degree_gain has an invalid shape or dtype"); + const float* degree_gain_ptr = + degree_gain.numel() ? degree_gain.data_ptr() : nullptr; auto stream = at::cuda::getCurrentCUDAStream(); auto [order, pair_table] = build_order_and_pair_table( @@ -1294,7 +1370,7 @@ dpa1_graph_descriptor(torch::Tensor edge_vec, auto f32 = torch::TensorOptions().dtype(torch::kFloat32).device(edge_vec.device()); - auto gr = torch::zeros({n_node, 4, NG}, f32); + auto gr = torch::zeros({n_node, basis_dim, NG}, f32); // Backward operands, spilled transposed so both directions stream // coalesced rows. E_pad rounds up to the tile size so partial-tile float4 // stores stay in bounds; for silu the g slot holds pre3 (its derivative @@ -1335,22 +1411,50 @@ dpa1_graph_descriptor(torch::Tensor edge_vec, order, stream}; -#define DPA1_LAUNCH_FWD(N1, N2, NG, ACT, STRIP) \ - launch_forward_portable(args, gr, pre2_saved, g_saved) +#define DPA1_LAUNCH_FWD4(N1, N2, NG, ACT, STRIP) \ + launch_forward_portable(args, gr, pre2_saved, \ + g_saved) +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX +#define DPA1_LAUNCH_FWD9(N1, N2, NG, ACT, STRIP) \ + launch_forward_portable(args, gr, pre2_saved, \ + g_saved) +#endif if (n_edge > 0) { - DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_FWD); + if (basis_dim == 4) { + DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_FWD4); +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + } else { + DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_FWD9); +#endif + } } -#undef DPA1_LAUNCH_FWD +#undef DPA1_LAUNCH_FWD4 +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX +#undef DPA1_LAUNCH_FWD9 +#endif const int out_dim = NG * (int)axis + (concat_tebd ? tebd_dim : 0); auto grrg = torch::empty({n_node, out_dim}, f32); auto rot_mat = torch::empty({write_rotation ? n_node : 0, NG, 3}, f32); if (n_node > 0) { - gram_kernel<<<(int)n_node, 128, 4 * NG * sizeof(float), stream>>>( - (int)n_node, NG, (int)axis, tebd_dim, (int)concat_tebd, - gr.data_ptr(), type_embedding.data_ptr(), - atype.data_ptr(), grrg.data_ptr(), - write_rotation ? rot_mat.data_ptr() : nullptr); + const size_t smem = basis_dim * NG * sizeof(float); + if (basis_dim == 4) { + gram_kernel<4><<<(int)n_node, 128, smem, stream>>>( + (int)n_node, NG, (int)axis, tebd_dim, (int)concat_tebd, + gr.data_ptr(), degree_gain_ptr, + type_embedding.data_ptr(), atype.data_ptr(), + grrg.data_ptr(), + write_rotation ? rot_mat.data_ptr() : nullptr); +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + } else { + gram_kernel<9><<<(int)n_node, 128, smem, stream>>>( + (int)n_node, NG, (int)axis, tebd_dim, (int)concat_tebd, + gr.data_ptr(), degree_gain_ptr, + type_embedding.data_ptr(), atype.data_ptr(), + grrg.data_ptr(), + write_rotation ? rot_mat.data_ptr() : nullptr); +#endif + } DPA1_CHECK_LAUNCH("dpa1_graph_descriptor gram"); } return {grrg, rot_mat, gr, order, pair_table, pre2_saved, g_saved}; @@ -1373,6 +1477,7 @@ torch::Tensor dpa1_graph_descriptor_backward( torch::Tensor atype, torch::Tensor davg, torch::Tensor dstd, + torch::Tensor degree_gain, torch::Tensor w1, torch::Tensor b1, torch::Tensor idt1, @@ -1397,12 +1502,24 @@ torch::Tensor dpa1_graph_descriptor_backward( const long n_edge = edge_vec.size(0); const long n_node = atype.size(0); const int NG = widths.ng; + const int basis_dim = (int)gr.size(1); const bool strip = gate_table.numel() > 0; +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + TORCH_CHECK(basis_dim == 4 || basis_dim == 9, + "dpa1_graph_descriptor_backward: basis dimension must be 4 or 9"); +#else + TORCH_CHECK( + basis_dim == 4, + "dpa1_graph_descriptor_backward: this build instantiates only lmax=1; " + "rebuild with DEEPMD_ENABLE_DPA1_HIGH_LMAX=ON for lmax=2"); +#endif + const float* degree_gain_ptr = + degree_gain.numel() ? degree_gain.data_ptr() : nullptr; auto stream = at::cuda::getCurrentCUDAStream(); auto f32 = torch::TensorOptions().dtype(torch::kFloat32).device(edge_vec.device()); - auto dgr = torch::empty({n_node, 4, NG}, f32); + auto dgr = torch::empty({n_node, basis_dim, NG}, f32); auto d_grrg_c = d_grrg.to(torch::kFloat32).contiguous(); torch::Tensor d_rot_c; const float* d_rot_ptr = nullptr; @@ -1411,11 +1528,20 @@ torch::Tensor dpa1_graph_descriptor_backward( d_rot_ptr = d_rot_c.data_ptr(); } if (n_node > 0) { - const size_t smem = (4 * NG + NG * (int)axis) * sizeof(float); - gram_backward_kernel<<<(int)n_node, 128, smem, stream>>>( - (int)n_node, NG, (int)axis, (int)d_grrg_c.size(1), - d_grrg_c.data_ptr(), d_rot_ptr, gr.data_ptr(), - dgr.data_ptr()); + const size_t smem = (basis_dim * NG + NG * (int)axis) * sizeof(float); + if (basis_dim == 4) { + gram_backward_kernel<4><<<(int)n_node, 128, smem, stream>>>( + (int)n_node, NG, (int)axis, (int)d_grrg_c.size(1), + d_grrg_c.data_ptr(), d_rot_ptr, gr.data_ptr(), + degree_gain_ptr, dgr.data_ptr()); +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + } else { + gram_backward_kernel<9><<<(int)n_node, 128, smem, stream>>>( + (int)n_node, NG, (int)axis, (int)d_grrg_c.size(1), + d_grrg_c.data_ptr(), d_rot_ptr, gr.data_ptr(), + degree_gain_ptr, dgr.data_ptr()); +#endif + } DPA1_CHECK_LAUNCH("dpa1_graph_descriptor gram backward"); } @@ -1460,28 +1586,46 @@ torch::Tensor dpa1_graph_descriptor_backward( order, stream}; -#define DPA1_LAUNCH_BWD(N1, N2, NG, ACT, STRIP) \ - launch_backward(args, dgr, pre2_saved, g_saved, \ - d_edge_vec) +#define DPA1_LAUNCH_BWD4(N1, N2, NG, ACT, STRIP) \ + launch_backward_portable(args, dgr, pre2_saved, \ + g_saved, d_edge_vec) +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX +#define DPA1_LAUNCH_BWD9(N1, N2, NG, ACT, STRIP) \ + launch_backward_portable(args, dgr, pre2_saved, \ + g_saved, d_edge_vec) +#endif if (n_edge > 0) { - DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_BWD); + if (basis_dim == 4) { + DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_BWD4); +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX + } else { + DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_BWD9); +#endif + } } -#undef DPA1_LAUNCH_BWD +#undef DPA1_LAUNCH_BWD4 +#if DEEPMD_ENABLE_DPA1_HIGH_LMAX +#undef DPA1_LAUNCH_BWD9 +#endif // Return the gradient in the coordinate's precision (a no-op when fp32). return d_edge_vec.to(edge_vec.scalar_type()); } +#undef DPA1_DISPATCH_WIDTH_ACT +#undef DPA1_DISPATCH_ONE + TORCH_LIBRARY_FRAGMENT(deepmd, m) { m.def( "dpa1_graph_descriptor(Tensor edge_vec, Tensor edge_index, " "Tensor edge_mask, Tensor atype, Tensor type_embedding, Tensor davg, " - "Tensor dstd, Tensor w1, Tensor b1, Tensor idt1, Tensor w2, Tensor b2, " + "Tensor dstd, Tensor degree_gain, Tensor w1, Tensor b1, Tensor idt1, " + "Tensor w2, Tensor b2, " "Tensor idt2, Tensor w3, Tensor b3, Tensor idt3, Tensor gate_table, " "int act, int type_one_side, int concat_tebd, int write_rotation, int " "smooth, int axis, int resnet2, int resnet3, float rcut, float " "rcut_smth, " - "float protection, float nnei) -> (Tensor grrg, Tensor rot_mat, " + "float protection, float nnei, int basis_dim) -> (Tensor grrg, Tensor " + "rot_mat, " "Tensor gr, Tensor edge_order, Tensor pair_table, Tensor pre2_saved, " "Tensor g_saved)"); m.impl("dpa1_graph_descriptor", torch::kCUDA, &dpa1_graph_descriptor); @@ -1489,7 +1633,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, m) { "dpa1_graph_descriptor_backward(Tensor d_grrg, Tensor? d_rot_mat, " "Tensor gr, Tensor edge_order, Tensor pair_table, Tensor pre2_saved, " "Tensor g_saved, Tensor edge_vec, Tensor edge_index, Tensor edge_mask, " - "Tensor atype, Tensor davg, Tensor dstd, Tensor w1, Tensor b1, " + "Tensor atype, Tensor davg, Tensor dstd, Tensor degree_gain, Tensor w1, " + "Tensor b1, " "Tensor idt1, Tensor w2, Tensor b2, Tensor idt2, Tensor w3, Tensor b3, " "Tensor idt3, Tensor gate_table, int act, int type_one_side, " "int smooth, int axis, int resnet2, int resnet3, float rcut, " diff --git a/source/op/pt/dpa1_graph_energy_force.cu b/source/op/pt/dpa1_graph_energy_force.cu index d1e9d56d3f..04942e1493 100644 --- a/source/op/pt/dpa1_graph_energy_force.cu +++ b/source/op/pt/dpa1_graph_energy_force.cu @@ -43,6 +43,7 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, torch::Tensor type_embedding, torch::Tensor davg, torch::Tensor dstd, + torch::Tensor degree_gain, torch::Tensor w1, torch::Tensor b1, torch::Tensor idt1, @@ -64,6 +65,7 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, double rcut_smth, double protection, double nnei, + int64_t basis_dim, std::vector fit_ws, std::vector fit_bs, std::vector fit_idts, @@ -88,10 +90,10 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, // === Step 1. Descriptor forward: edge stream -> (N, nd) descriptor. === auto desc = dpa1_graph_descriptor( - edge_vec_f, edge_index, edge_mask, atype, type_embedding, davg, dstd, w1, - b1, idt1, w2, b2, idt2, w3, b3, idt3, gate_table, act, type_one_side, - concat_tebd, /*write_rotation=*/0, smooth, axis, resnet2, resnet3, rcut, - rcut_smth, protection, nnei); + edge_vec_f, edge_index, edge_mask, atype, type_embedding, davg, dstd, + degree_gain, w1, b1, idt1, w2, b2, idt2, w3, b3, idt3, gate_table, act, + type_one_side, concat_tebd, /*write_rotation=*/0, smooth, axis, resnet2, + resnet3, rcut, rcut_smth, protection, nnei, basis_dim); const torch::Tensor& grrg = std::get<0>(desc); const torch::Tensor& gr = std::get<2>(desc); const torch::Tensor& edge_order = std::get<3>(desc); @@ -122,9 +124,9 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, std::get<1>(fit) = torch::Tensor(); auto g_e = dpa1_graph_descriptor_backward( d_grrg, std::nullopt, gr, edge_order, pair_table, pre2_saved, g_saved, - edge_vec_f, edge_index, edge_mask, atype, davg, dstd, w1, b1, idt1, w2, - b2, idt2, w3, b3, idt3, gate_table, act, type_one_side, smooth, axis, - resnet2, resnet3, rcut, rcut_smth, protection, nnei); + edge_vec_f, edge_index, edge_mask, atype, davg, dstd, degree_gain, w1, b1, + idt1, w2, b2, idt2, w3, b3, idt3, gate_table, act, type_one_side, smooth, + axis, resnet2, resnet3, rcut, rcut_smth, protection, nnei); // === Step 5. Scatter dE/d(edge_vec) into force / virial / atom virial. === // g_e and edge_vec_f are already in the compute precision; the per-node force @@ -143,12 +145,14 @@ TORCH_LIBRARY_FRAGMENT(deepmd, m) { "edge_mask, Tensor destination_order, Tensor destination_row_ptr, " "Tensor source_order, Tensor source_row_ptr, Tensor atype, Tensor " "n_node, Tensor ownership, Tensor type_embedding, " - "Tensor davg, Tensor dstd, Tensor w1, Tensor b1, Tensor idt1, Tensor w2, " + "Tensor davg, Tensor dstd, Tensor degree_gain, Tensor w1, Tensor b1, " + "Tensor idt1, Tensor w2, " "Tensor " "b2, Tensor idt2, Tensor w3, Tensor b3, Tensor idt3, Tensor gate_table, " "int act, int type_one_side, int concat_tebd, int smooth, int axis, int " "resnet2, int resnet3, float rcut, float rcut_smth, float protection, " - "float nnei, Tensor[] fit_ws, Tensor[] fit_bs, Tensor[] fit_idts, int[] " + "float nnei, int basis_dim, Tensor[] fit_ws, Tensor[] fit_bs, Tensor[] " + "fit_idts, int[] " "fit_resnets, Tensor w_head, Tensor b_head, Tensor bias_atom_e, int " "fit_act, SymInt node_capacity, bool do_atomic_virial) -> (Tensor, " "Tensor, Tensor, Tensor, Tensor)"); diff --git a/source/op/pt/dpa1_moment_basis.cuh b/source/op/pt/dpa1_moment_basis.cuh new file mode 100644 index 0000000000..28eac4837b --- /dev/null +++ b/source/op/pt/dpa1_moment_basis.cuh @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compact Cartesian real spherical harmonics for DPA1 moment aggregation. +// +// Rows [l^2, (l+1)^2) contain degree l in m=-l,...,l order. Degrees 2-4 +// use norm-normalized regular solid harmonics. On unit vectors, their inner +// product equals P_l(u dot v). + +#pragma once + +#include + +namespace deepmd::dpa1 { + +#define DPA1_MOMENT_INLINE __device__ __forceinline__ + +struct Jet3 { + float value; + float dx; + float dy; + float dz; +}; + +DPA1_MOMENT_INLINE float degree_weight( + int row, const float* __restrict__ degree_gain_raw) { + if (row < 4) { + return 1.0f; + } + const int degree_index = row < 9 ? 0 : (row < 16 ? 1 : 2); + const float gain = __ldg(degree_gain_raw + degree_index); + return gain * gain; +} + +DPA1_MOMENT_INLINE Jet3 operator+(const Jet3& a, const Jet3& b) { + return {a.value + b.value, a.dx + b.dx, a.dy + b.dy, a.dz + b.dz}; +} + +DPA1_MOMENT_INLINE Jet3 operator-(const Jet3& a, const Jet3& b) { + return {a.value - b.value, a.dx - b.dx, a.dy - b.dy, a.dz - b.dz}; +} + +DPA1_MOMENT_INLINE Jet3 operator*(const Jet3& a, const Jet3& b) { + return {a.value * b.value, a.dx * b.value + a.value * b.dx, + a.dy * b.value + a.value * b.dy, a.dz * b.value + a.value * b.dz}; +} + +DPA1_MOMENT_INLINE Jet3 operator*(float scale, const Jet3& value) { + return {scale * value.value, scale * value.dx, scale * value.dy, + scale * value.dz}; +} + +DPA1_MOMENT_INLINE Jet3 operator*(const Jet3& value, float scale) { + return scale * value; +} + +template +DPA1_MOMENT_INLINE void evaluate_angular_basis(const T& x, + const T& y, + const T& z, + T* output) { + static_assert(BasisDim == 4 || BasisDim == 9 || BasisDim == 16 || + BasisDim == 25); + if constexpr (BasisDim == 4) { + return; + } + + constexpr float sqrt3 = 1.7320508075688772935f; + const T x2 = x * x; + const T y2 = y * y; + const T z2 = z * z; + const T q = x2 + y2 + z2; + const T x2_minus_y2 = x2 - y2; + output[0] = sqrt3 * x * y; + output[1] = sqrt3 * y * z; + output[2] = 0.5f * (3.0f * z2 - q); + output[3] = sqrt3 * x * z; + output[4] = (0.5f * sqrt3) * x2_minus_y2; + + if constexpr (BasisDim >= 16) { + constexpr float sqrt5_over8 = 0.79056941504209483299f; + constexpr float sqrt15 = 3.8729833462074168852f; + constexpr float sqrt3_over8 = 0.61237243569579452455f; + output[5] = sqrt5_over8 * y * (3.0f * x2 - y2); + output[6] = sqrt15 * x * y * z; + output[7] = sqrt3_over8 * y * (5.0f * z2 - q); + output[8] = 0.5f * z * (5.0f * z2 - 3.0f * q); + output[9] = sqrt3_over8 * x * (5.0f * z2 - q); + output[10] = (0.5f * sqrt15) * z * x2_minus_y2; + output[11] = sqrt5_over8 * x * (x2 - 3.0f * y2); + } + + if constexpr (BasisDim == 25) { + constexpr float sqrt35 = 5.9160797830996160426f; + constexpr float sqrt70 = 8.3666002653407554798f; + constexpr float sqrt5 = 2.2360679774997896964f; + constexpr float sqrt10 = 3.1622776601683793320f; + output[12] = (0.5f * sqrt35) * x * y * x2_minus_y2; + output[13] = (0.25f * sqrt70) * y * z * (3.0f * x2 - y2); + output[14] = (0.5f * sqrt5) * x * y * (7.0f * z2 - q); + output[15] = (0.25f * sqrt10) * y * z * (7.0f * z2 - 3.0f * q); + output[16] = 0.125f * (35.0f * z2 * z2 - 30.0f * z2 * q + 3.0f * q * q); + output[17] = (0.25f * sqrt10) * x * z * (7.0f * z2 - 3.0f * q); + output[18] = (0.25f * sqrt5) * x2_minus_y2 * (7.0f * z2 - q); + output[19] = (0.25f * sqrt70) * x * z * (x2 - 3.0f * y2); + output[20] = (0.125f * sqrt35) * (x2 * x2 - 6.0f * x2 * y2 + y2 * y2); + } +} + +template +DPA1_MOMENT_INLINE void fill_angular_basis( + float* basis, float x, float y, float z, float radial) { + if constexpr (BasisDim > 4) { + float angular[BasisDim - 4]; + evaluate_angular_basis(x, y, z, angular); +#pragma unroll + for (int row = 0; row < BasisDim - 4; ++row) { + basis[4 + row] = radial * angular[row]; + } + } +} + +template +DPA1_MOMENT_INLINE void angular_basis_vjp(const float (&d_basis)[BasisDim], + float inverse_neighbors, + float x, + float y, + float z, + float& radial_partial, + float& grad_x, + float& grad_y, + float& grad_z) { + if constexpr (BasisDim > 4) { + const Jet3 x_jet{x, 1.0f, 0.0f, 0.0f}; + const Jet3 y_jet{y, 0.0f, 1.0f, 0.0f}; + const Jet3 z_jet{z, 0.0f, 0.0f, 1.0f}; + Jet3 angular[BasisDim - 4]; + evaluate_angular_basis(x_jet, y_jet, z_jet, angular); +#pragma unroll + for (int row = 0; row < BasisDim - 4; ++row) { + const float gradient = d_basis[4 + row] * inverse_neighbors; + radial_partial = fmaf(gradient, angular[row].value, radial_partial); + grad_x = fmaf(gradient, angular[row].dx, grad_x); + grad_y = fmaf(gradient, angular[row].dy, grad_y); + grad_z = fmaf(gradient, angular[row].dz, grad_z); + } + } +} + +template +DPA1_MOMENT_INLINE void add_angular_edge_gradient( + const float (&d_basis)[BasisDim], + float inverse_neighbors, + float x, + float y, + float z, + float radius, + float inverse_denominator, + float inverse_stddev0, + float switch_value, + float switch_gradient, + float& output_x, + float& output_y, + float& output_z) { + if constexpr (BasisDim > 4) { + if (radius <= 0.0f) { + return; + } + const float inverse_radius = 1.0f / radius; + const float nx = x * inverse_radius; + const float ny = y * inverse_radius; + const float nz = z * inverse_radius; + const float vx = x * inverse_denominator; + const float vy = y * inverse_denominator; + const float vz = z * inverse_denominator; + float radial_partial = 0.0f; + float grad_vx = 0.0f; + float grad_vy = 0.0f; + float grad_vz = 0.0f; + angular_basis_vjp(d_basis, inverse_neighbors, vx, vy, vz, + radial_partial, grad_vx, grad_vy, grad_vz); + const float protected_dot = grad_vx * vx + grad_vy * vy + grad_vz * vz; + const float amplitude = + switch_value * inverse_denominator * inverse_stddev0; + const float amplitude_gradient = + inverse_stddev0 * inverse_denominator * + (switch_gradient - switch_value * inverse_denominator); + output_x += + radial_partial * amplitude_gradient * nx + + amplitude * inverse_denominator * (grad_vx - protected_dot * nx); + output_y += + radial_partial * amplitude_gradient * ny + + amplitude * inverse_denominator * (grad_vy - protected_dot * ny); + output_z += + radial_partial * amplitude_gradient * nz + + amplitude * inverse_denominator * (grad_vz - protected_dot * nz); + } +} + +#undef DPA1_MOMENT_INLINE + +} // namespace deepmd::dpa1 diff --git a/source/op/pt/graph_ops.h b/source/op/pt/graph_ops.h index 669513b2b5..418364ef0b 100644 --- a/source/op/pt/graph_ops.h +++ b/source/op/pt/graph_ops.h @@ -33,6 +33,7 @@ dpa1_graph_descriptor(torch::Tensor edge_vec, torch::Tensor type_embedding, torch::Tensor davg, torch::Tensor dstd, + torch::Tensor degree_gain, torch::Tensor w1, torch::Tensor b1, torch::Tensor idt1, @@ -54,7 +55,8 @@ dpa1_graph_descriptor(torch::Tensor edge_vec, double rcut, double rcut_smth, double protection, - double nnei); + double nnei, + int64_t basis_dim); // dE/d(edge_vec) from dE/d(grrg); consumes the saved tensors of the forward. torch::Tensor dpa1_graph_descriptor_backward( @@ -71,6 +73,7 @@ torch::Tensor dpa1_graph_descriptor_backward( torch::Tensor atype, torch::Tensor davg, torch::Tensor dstd, + torch::Tensor degree_gain, torch::Tensor w1, torch::Tensor b1, torch::Tensor idt1, diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index bf5ee3f87b..dadd3aab68 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -48,6 +48,12 @@ void CheckTabulateDevices(const torch::Tensor& table_tensor, } } +void CheckTabulateFusionSeABasisDimension(const int64_t ndescrpt) { + TORCH_CHECK(deepmd::is_supported_se_a_basis_dimension(ndescrpt), + "The environment basis dimension must be 4, 9, 16, or 25, got ", + ndescrpt); +} + template void TabulateFusionSeAForward(const torch::Tensor& table_tensor, const torch::Tensor& table_info_tensor, @@ -92,12 +98,14 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor, const int64_t nloc = em_tensor.size(0); const int64_t nnei = em_tensor.size(1); + const int64_t ndescrpt = em_tensor.size(2); + CheckTabulateFusionSeABasisDimension(ndescrpt); // compute if (device == "GPU") { #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em, two_embed, nloc, nnei, last_layer_size, - is_sorted); + is_sorted, ndescrpt); #else throw std::runtime_error( "The input tensor is on the GPU, but the GPU support for the " @@ -106,7 +114,7 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor, } else if (device == "CPU") { deepmd::tabulate_fusion_se_a_cpu(descriptor, table, table_info, em_x, em, two_embed, nloc, nnei, last_layer_size, - is_sorted); + is_sorted, ndescrpt); } } @@ -155,13 +163,15 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor, const FPTYPE* dy = dy_tensor.view({-1}).data_ptr(); const int64_t nloc = em_tensor.size(0); const int64_t nnei = em_tensor.size(1); + const int64_t ndescrpt = em_tensor.size(2); const int64_t last_layer_size = descriptor_tensor.size(2); + CheckTabulateFusionSeABasisDimension(ndescrpt); // compute if (device == "GPU") { #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_a_grad_gpu( dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, - nloc, nnei, last_layer_size, is_sorted); + nloc, nnei, last_layer_size, is_sorted, ndescrpt); #else throw std::runtime_error( "The input tensor is on the GPU, but the GPU support for the " @@ -170,7 +180,7 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor, } else if (device == "CPU") { deepmd::tabulate_fusion_se_a_grad_cpu( dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, - nloc, nnei, last_layer_size, is_sorted); + nloc, nnei, last_layer_size, is_sorted, ndescrpt); } } @@ -224,13 +234,15 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor, : dz_dy_dtwo_tensor.view({-1}).data_ptr(); const int64_t nloc = em_tensor.size(0); const int64_t nnei = em_tensor.size(1); + const int64_t ndescrpt = em_tensor.size(2); const int64_t last_layer_size = descriptor_tensor.size(2); + CheckTabulateFusionSeABasisDimension(ndescrpt); // compute if (device == "GPU") { #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_a_grad_grad_gpu( dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, - dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted, ndescrpt); #else throw std::runtime_error( "The input tensor is on the GPU, but the GPU support for the " @@ -242,7 +254,7 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor, } else if (device == "CPU") { deepmd::tabulate_fusion_se_a_grad_grad_cpu( dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, - dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); + dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted, ndescrpt); } } @@ -898,8 +910,8 @@ class TabulateFusionSeAOp auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) .device(table_tensor.device()); - torch::Tensor descriptor_tensor = - torch::empty({em_tensor.size(0), 4, last_layer_size}, options); + torch::Tensor descriptor_tensor = torch::empty( + {em_tensor.size(0), em_tensor.size(2), last_layer_size}, options); // compute // Keep the sorted fold enabled: exclusions are uniform within each se_a // type-pair invocation, and compressed forward_lower sorts its nlist first. @@ -1086,8 +1098,8 @@ class TabulateFusionSeAttenOp auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) .device(table_tensor.device()); - torch::Tensor descriptor_tensor = - torch::empty({em_tensor.size(0), 4, last_layer_size}, options); + torch::Tensor descriptor_tensor = torch::empty( + {em_tensor.size(0), em_tensor.size(2), last_layer_size}, options); // compute TabulateFusionSeAForward( table_tensor, table_info_tensor, em_x_tensor, em_tensor, diff --git a/source/tests/common/dpmodel/test_descriptor_dpa1.py b/source/tests/common/dpmodel/test_descriptor_dpa1.py index 67d214203a..1df473c087 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa1.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa1.py @@ -44,6 +44,27 @@ def test_get_numb_attn_layer(self) -> None: em2 = DescrptDPA1(self.rcut, self.rcut_smth, self.sel, ntypes=2, attn_layer=2) self.assertEqual(em2.get_numb_attn_layer(), 2) + def test_lmax_two_serialization(self) -> None: + descriptor = DescrptDPA1( + self.rcut, + self.rcut_smth, + self.sel, + ntypes=2, + attn_layer=0, + lmax=2, + ) + serialized = descriptor.serialize() + self.assertEqual(serialized["@version"], 4) + self.assertEqual(descriptor.se_atten.serialize()["@version"], 2) + restored = DescrptDPA1.deserialize(serialized) + + actual = descriptor.call(self.coord_ext, self.atype_ext, self.nlist) + expected = restored.call(self.coord_ext, self.atype_ext, self.nlist) + + self.assertEqual(restored.se_atten.lmax, 2) + for index in (0, 1, 4): + np.testing.assert_allclose(actual[index], expected[index]) + def test_multiple_frames(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py index 11db9a7785..15e2911ac7 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py @@ -19,7 +19,7 @@ class TestDpa1BlockCallGraph: - def _make(self, sel, type_one_side=False): + def _make(self, sel, type_one_side=False, lmax=1): return DescrptDPA1( rcut=4.0, rcut_smth=0.5, @@ -29,6 +29,7 @@ def _make(self, sel, type_one_side=False): axis_neuron=2, neuron=[6, 12], type_one_side=type_one_side, + lmax=lmax, ) def setup_method(self) -> None: @@ -39,14 +40,15 @@ def setup_method(self) -> None: @pytest.mark.parametrize("type_one_side", [False, True]) # tebd concat branch @pytest.mark.parametrize("sel", [[20], [3]]) # non-binding AND binding - def test_block_graph_equals_dense_any_sel(self, sel, type_one_side) -> None: + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) + def test_block_graph_equals_dense_any_sel(self, sel, type_one_side, lmax) -> None: """Graph block output is bit-exact with the dense block on the same nlist. ``type_one_side`` toggles the concat branch in the block: when True the per-edge feature concatenates only the NEIGHBOR tebd (no center tebd), so both the graph and dense paths must agree for either branch. """ - dd = self._make(sel, type_one_side=type_one_side) + dd = self._make(sel, type_one_side=type_one_side, lmax=lmax) blk = dd.se_atten # build the dense nlist exactly as the descriptor would ( @@ -180,7 +182,13 @@ def setup_method(self) -> None: self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5 self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64) - def _make(self, type_one_side: bool, smooth: bool, attn_layer: int) -> DescrptDPA1: + def _make( + self, + type_one_side: bool, + smooth: bool, + attn_layer: int, + lmax: int = 1, + ) -> DescrptDPA1: return DescrptDPA1( rcut=4.0, rcut_smth=0.5, @@ -192,6 +200,7 @@ def _make(self, type_one_side: bool, smooth: bool, attn_layer: int) -> DescrptDP tebd_input_mode="strip", type_one_side=type_one_side, smooth_type_embedding=smooth, + lmax=lmax, ) def _assert_parity(self, dd: DescrptDPA1, compact: bool) -> None: @@ -240,9 +249,10 @@ def _assert_parity(self, dd: DescrptDPA1, compact: bool) -> None: "type_one_side", [False, True] ) # two-side vs one-side strip table @pytest.mark.parametrize("smooth", [False, True]) # gg_t switch-smoothing branch - def test_strip_attn0_equals_dense(self, type_one_side, smooth) -> None: + @pytest.mark.parametrize("lmax", [1, 2, 3, 4]) + def test_strip_attn0_equals_dense(self, type_one_side, smooth, lmax) -> None: """attn_layer=0: no attention, so strip parity is bit-exact for both smooth values.""" - dd = self._make(type_one_side, smooth, attn_layer=0) + dd = self._make(type_one_side, smooth, attn_layer=0, lmax=lmax) self._assert_parity(dd, compact=True) @pytest.mark.parametrize( diff --git a/source/tests/pt/model/test_compressed_descriptor_se_atten.py b/source/tests/pt/model/test_compressed_descriptor_se_atten.py index c3f4444351..ad58d2e60f 100644 --- a/source/tests/pt/model/test_compressed_descriptor_se_atten.py +++ b/source/tests/pt/model/test_compressed_descriptor_se_atten.py @@ -45,10 +45,42 @@ def eval_pt_descriptor( return result -@parameterized(("float32", "float64"), (True, False)) +def eval_pt_descriptor_with_gradient( + pt_obj: Any, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, +) -> tuple[torch.Tensor, torch.Tensor]: + coord = ( + torch.from_numpy(coords).to(PT_DEVICE).reshape(1, -1, 3).requires_grad_(True) + ) + ext_coords, ext_atype, mapping = extend_coord_with_ghosts_pt( + coord, + torch.from_numpy(atype).to(PT_DEVICE).reshape(1, -1), + torch.from_numpy(box).to(PT_DEVICE).reshape(1, 3, 3), + pt_obj.get_rcut(), + ) + nlist = build_neighbor_list_pt( + ext_coords, + ext_atype, + natoms[0], + pt_obj.get_rcut(), + pt_obj.get_sel(), + distinguish_types=False, + ) + result = pt_obj(ext_coords, ext_atype, nlist, mapping=mapping)[0] + (first_derivative,) = torch.autograd.grad( + result.sum(), + coord, + ) + return result, first_derivative + + +@parameterized(("float32", "float64"), (True, False), (1, 2, 3, 4)) class TestDescriptorSeAtten(unittest.TestCase): def setUp(self) -> None: - (self.dtype, self.type_one_side) = self.param + (self.dtype, self.type_one_side, self.lmax) = self.param if self.dtype == "float32": self.atol = 1e-5 elif self.dtype == "float64": @@ -104,6 +136,7 @@ def setUp(self) -> None: precision=self.dtype, type_one_side=self.type_one_side, tebd_input_mode="strip", + lmax=self.lmax, ) def test_compressed_forward(self) -> None: @@ -132,6 +165,79 @@ def test_compressed_forward(self) -> None: rtol=self.atol, ) + def test_compressed_coordinate_gradient(self) -> None: + if self.lmax == 1: + self.skipTest("The lmax=1 derivative path is covered by operator tests.") + + dense = eval_pt_descriptor_with_gradient( + self.se_atten, + self.natoms, + self.coords, + self.atype, + self.box, + ) + self.se_atten.enable_compression(0.5) + compressed = eval_pt_descriptor_with_gradient( + self.se_atten, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + derivative_atol = 1e-4 if self.dtype == "float32" else 1e-9 + for dense_value, compressed_value in zip(dense, compressed, strict=True): + torch.testing.assert_close( + compressed_value, + dense_value, + atol=derivative_atol, + rtol=derivative_atol, + ) + + def test_compressed_excluded_type_holes(self) -> None: + if (self.dtype, self.type_one_side, self.lmax) != ("float64", True, 2): + self.skipTest("A single representative configuration is sufficient.") + + descriptor = DescrptDPA1( + self.rcut, + self.rcut_smth, + self.sel, + self.ntypes, + self.neuron, + self.axis_neuron, + 4, + attn=8, + attn_layer=0, + seed=self.seed, + precision=self.dtype, + type_one_side=self.type_one_side, + tebd_input_mode="strip", + exclude_types=[(0, 1)], + lmax=self.lmax, + ) + self.assertFalse(descriptor.se_atten.is_sorted) + dense = eval_pt_descriptor( + descriptor, + self.natoms, + self.coords, + self.atype, + self.box, + ) + descriptor.enable_compression(0.5) + compressed = eval_pt_descriptor( + descriptor, + self.natoms, + self.coords, + self.atype, + self.box, + ) + torch.testing.assert_close( + compressed, + dense, + atol=self.atol, + rtol=self.atol, + ) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt/model/test_descriptor_dpa1_triton.py b/source/tests/pt/model/test_descriptor_dpa1_triton.py index 391c4921e9..c9486fe4c5 100644 --- a/source/tests/pt/model/test_descriptor_dpa1_triton.py +++ b/source/tests/pt/model/test_descriptor_dpa1_triton.py @@ -68,7 +68,17 @@ ) -def _rand_conv_inputs(nfnl, nnei, ng, resnet_mult, has_idt, ntype_pair, device, seed=0): +def _rand_conv_inputs( + nfnl, + nnei, + ng, + resnet_mult, + has_idt, + ntype_pair, + device, + seed=0, + basis_dim=4, +): gen = torch.Generator(device=device).manual_seed(seed) h1_dim = ng // resnet_mult if resnet_mult > 0 else ng z2 = torch.randn(nfnl, nnei, ng, device=device, generator=gen) @@ -81,22 +91,31 @@ def _rand_conv_inputs(nfnl, nnei, ng, resnet_mult, has_idt, ntype_pair, device, tt = torch.randn(ntype_pair, ng, device=device, generator=gen) * 0.3 idx = torch.randint(0, ntype_pair, (nfnl * nnei,), device=device, generator=gen) sw = torch.rand(nfnl, nnei, device=device, generator=gen) - rr = torch.randn(nfnl, nnei, 4, device=device, generator=gen) - return z2, h1, idt, tt, idx, sw, rr + basis = torch.randn(nfnl, nnei, basis_dim, device=device, generator=gen) + return z2, h1, idt, tt, idx, sw, basis class TestSeConvConfig(unittest.TestCase): """Launch-configuration resolution (CPU-safe, no kernel launch).""" def test_level1_returns_default(self) -> None: - self.assertEqual(resolve_conv_config(128, 64, level=1), DEFAULT_CONFIG) + self.assertEqual(resolve_conv_config(128, 64, 4, level=1), DEFAULT_CONFIG) def test_level0_returns_default(self) -> None: - self.assertEqual(resolve_conv_config(128, 64, level=0), DEFAULT_CONFIG) + self.assertEqual(resolve_conv_config(128, 64, 4, level=0), DEFAULT_CONFIG) def test_level2_unknown_shape_falls_back(self) -> None: # An unswept channel width can only fall back to the universal default. - self.assertEqual(resolve_conv_config(777, 333, level=2), DEFAULT_CONFIG) + self.assertEqual(resolve_conv_config(777, 333, 9, level=2), DEFAULT_CONFIG) + + def test_level2_distinguishes_basis_width(self) -> None: + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_name() != "NVIDIA H20" + ): + self.skipTest("The built-in launch configurations are H20-specific.") + self.assertEqual(resolve_conv_config(64, 32, 4, level=2), (16, 2)) + self.assertEqual(resolve_conv_config(64, 32, 9, level=2), (32, 2)) @_GPU @@ -115,36 +134,69 @@ def setUp(self) -> None: # path (padded to 128) and the direct doubling fold at the true ``H1``. # ``act`` selects the inlined activation: 0 = tanh, 1 = silu. self.cases = [ - (ng, mult, has_idt, act, gated) + (ng, mult, has_idt, act, gated, basis_dim) for ng in (128, 100) for mult in (2, 1, 0) for has_idt in (False, True) for act in (0, 1) for gated in (1, 0) + for basis_dim in (4, 9, 16, 25) ] def test_forward_matches_reference(self) -> None: - for ng, mult, has_idt, act, gated in self.cases: - with self.subTest(ng=ng, mult=mult, has_idt=has_idt, act=act, gated=gated): - z2, h1, idt, tt, idx, sw, rr = _rand_conv_inputs( - 512, 47, ng, mult, has_idt, 169, self.device + for ng, mult, has_idt, act, gated, basis_dim in self.cases: + with self.subTest( + ng=ng, + mult=mult, + has_idt=has_idt, + act=act, + gated=gated, + basis_dim=basis_dim, + ): + z2, h1, idt, tt, idx, sw, basis = _rand_conv_inputs( + 512, + 47, + ng, + mult, + has_idt, + 169, + self.device, + basis_dim=basis_dim, ) - ref = _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, mult, act, gated) - got = se_conv(z2, h1, idt, tt, idx, sw, rr, mult, act, gated) + ref = _se_conv_reference( + z2, h1, idt, tt, idx, sw, basis, mult, act, gated + ) + got = se_conv(z2, h1, idt, tt, idx, sw, basis, mult, act, gated) rel = (got - ref).abs().max() / ref.abs().max() self.assertLess(rel.item(), 1e-5) def test_backward_matches_reference(self) -> None: - for ng, mult, has_idt, act, gated in self.cases: - with self.subTest(ng=ng, mult=mult, has_idt=has_idt, act=act, gated=gated): - z2, h1, idt, tt, idx, sw, rr = _rand_conv_inputs( - 512, 47, ng, mult, has_idt, 169, self.device + for ng, mult, has_idt, act, gated, basis_dim in self.cases: + with self.subTest( + ng=ng, + mult=mult, + has_idt=has_idt, + act=act, + gated=gated, + basis_dim=basis_dim, + ): + z2, h1, idt, tt, idx, sw, basis = _rand_conv_inputs( + 512, + 47, + ng, + mult, + has_idt, + 169, + self.device, + basis_dim=basis_dim, ) gout = torch.randn_like( - _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, mult, act, gated) + _se_conv_reference( + z2, h1, idt, tt, idx, sw, basis, mult, act, gated + ) ) ref_in = [ - t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, rr) + t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, basis) ] _se_conv_reference( ref_in[0], @@ -159,7 +211,7 @@ def test_backward_matches_reference(self) -> None: gated, ).backward(gout) got_in = [ - t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, rr) + t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, basis) ] se_conv( got_in[0], @@ -174,7 +226,7 @@ def test_backward_matches_reference(self) -> None: gated, ).backward(gout) for name, a, b in zip( - ("z2", "h1", "sw", "rr"), + ("z2", "h1", "sw", "basis"), (t.grad for t in ref_in), (t.grad for t in got_in), strict=True, @@ -189,27 +241,29 @@ def test_backward_matches_reference(self) -> None: self.assertLess(rel.item(), 1e-5, msg=f"grad {name}") def test_make_fx_force_trace(self) -> None: - z2, h1, idt, tt, idx, sw, rr = _rand_conv_inputs( - 512, 47, 128, 2, False, 169, self.device + z2, h1, idt, tt, idx, sw, basis = _rand_conv_inputs( + 512, 47, 128, 2, False, 169, self.device, basis_dim=25 ) gout = torch.randn_like( - _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, 2, 0, 1) + _se_conv_reference(z2, h1, idt, tt, idx, sw, basis, 2, 0, 1) ) - def force_fn(z2, h1, idt, tt, idx, sw, rr): + def force_fn(z2, h1, idt, tt, idx, sw, basis): z2r = z2.detach().requires_grad_(True) - out = se_conv(z2r, h1, idt, tt, idx, sw, rr, 2, 0, 1) - (grad_z2,) = torch.autograd.grad(out, z2r, gout) - return out, grad_z2 + basisr = basis.detach().requires_grad_(True) + out = se_conv(z2r, h1, idt, tt, idx, sw, basisr, 2, 0, 1) + grad_z2, grad_basis = torch.autograd.grad(out, (z2r, basisr), gout) + return out, grad_z2, grad_basis - traced = make_fx(force_fn)(z2, h1, idt, tt, idx, sw, rr) + traced = make_fx(force_fn)(z2, h1, idt, tt, idx, sw, basis) n_fwd = sum(1 for n in traced.graph.nodes if "se_conv.default" in str(n.target)) n_bwd = sum(1 for n in traced.graph.nodes if "se_conv_bwd" in str(n.target)) self.assertEqual(n_fwd, 1) self.assertEqual(n_bwd, 1) - out_ref, _ = force_fn(z2, h1, idt, tt, idx, sw, rr) - out_traced, _ = traced(z2, h1, idt, tt, idx, sw, rr) - torch.testing.assert_close(out_traced, out_ref) + reference = force_fn(z2, h1, idt, tt, idx, sw, basis) + actual = traced(z2, h1, idt, tt, idx, sw, basis) + for actual_value, reference_value in zip(actual, reference, strict=True): + torch.testing.assert_close(actual_value, reference_value) def _rand_edge_inputs( @@ -445,6 +499,7 @@ def _build_dpa1( resnet_dt=False, activation_function="tanh", tebd_input_mode="strip", + lmax=1, ): des = DescrptDPA1( rcut=6.0, @@ -461,6 +516,7 @@ def _build_dpa1( activation_function=activation_function, precision="float32", seed=1, + lmax=lmax, ).to(device) des.eval() return des @@ -516,6 +572,27 @@ def test_parity_doubling(self) -> None: def test_parity_identity(self) -> None: self._assert_parity(_build_dpa1(self.device, [8, 16, 16])) + def test_parity_lmax_two(self) -> None: + saved = os.environ.get("DP_TRITON_INFER") + os.environ["DP_TRITON_INFER"] = "2" + try: + self._assert_parity(_build_dpa1(self.device, [8, 16, 32], lmax=2)) + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + + def test_parity_lmax_two_concat(self) -> None: + self._assert_parity( + _build_dpa1( + self.device, + [8, 16, 32], + tebd_input_mode="concat", + lmax=2, + ) + ) + def test_parity_resnet_dt(self) -> None: self._assert_parity(_build_dpa1(self.device, [8, 16, 32], resnet_dt=True)) diff --git a/source/tests/pt/model/test_dpa1.py b/source/tests/pt/model/test_dpa1.py index bdb306f4de..531580250b 100644 --- a/source/tests/pt/model/test_dpa1.py +++ b/source/tests/pt/model/test_dpa1.py @@ -1,14 +1,24 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +import math import unittest import numpy as np import torch from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DPDescrptDPA1 +from deepmd.dpmodel.descriptor.dpa1 import ( + build_dpa1_moment_basis, +) from deepmd.pt.model.descriptor.dpa1 import ( DescrptDPA1, ) +from deepmd.pt.model.descriptor.se_atten import ( + _build_degree_weights, + _build_moment_basis, + _compute_angular_radial, + _safe_direction, +) from deepmd.pt.utils import ( env, ) @@ -170,3 +180,430 @@ def test_jit( # dd1 = DescrptDPA1.deserialize(dd0.serialize()) model = torch.jit.script(dd0) # model = torch.jit.script(dd1) + + +class TestDPA1AngularMoments(unittest.TestCase): + """Test the dense PyTorch angular moment basis of DPA1.""" + + dtype = torch.float64 + + @staticmethod + def _build_descriptor(lmax: int) -> DescrptDPA1: + return DescrptDPA1( + rcut=3.0, + rcut_smth=2.5, + sel=4, + ntypes=1, + neuron=[4, 8, 8], + axis_neuron=4, + lmax=lmax, + tebd_dim=2, + tebd_input_mode="strip", + set_davg_zero=True, + attn_layer=0, + precision="float64", + concat_output_tebd=False, + seed=11, + ).to(env.DEVICE) + + @classmethod + def _evaluate( + cls, + descriptor: DescrptDPA1, + neighbors: torch.Tensor, + *, + requires_grad: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + coord = torch.cat( + [ + torch.zeros( + (1, 3), + dtype=cls.dtype, + device=env.DEVICE, + ), + neighbors, + ] + ).reshape(1, -1) + coord.requires_grad_(requires_grad) + atype = torch.zeros((1, 5), dtype=torch.long, device=env.DEVICE) + nlist = torch.tensor( + [[[1, 2, 3, 4]]], + dtype=torch.long, + device=env.DEVICE, + ) + result = descriptor(coord, atype, nlist)[0] + return result, coord + + @classmethod + def _square_directions(cls) -> torch.Tensor: + return torch.tensor( + [ + [1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, -1.0, 0.0], + ], + dtype=cls.dtype, + device=env.DEVICE, + ) + + @classmethod + def _tetrahedral_directions(cls) -> torch.Tensor: + return torch.tensor( + [ + [1.0, 1.0, 1.0], + [1.0, -1.0, -1.0], + [-1.0, 1.0, -1.0], + [-1.0, -1.0, 1.0], + ], + dtype=cls.dtype, + device=env.DEVICE, + ) / math.sqrt(3.0) + + def test_higher_degree_addition_theorem(self) -> None: + generator = torch.Generator(device=env.DEVICE).manual_seed(37) + left = torch.randn( + 32, + 3, + dtype=self.dtype, + device=env.DEVICE, + generator=generator, + ) + right = torch.randn( + 32, + 3, + dtype=self.dtype, + device=env.DEVICE, + generator=generator, + ) + left = torch.nn.functional.normalize(left, dim=-1) + right = torch.nn.functional.normalize(right, dim=-1) + rr = torch.zeros(32, 1, 4, dtype=self.dtype, device=env.DEVICE) + radial = torch.ones(32, 1, 1, dtype=self.dtype, device=env.DEVICE) + left_basis = _build_moment_basis(rr, left[:, None, :], radial, 4) + right_basis = _build_moment_basis(rr, right[:, None, :], radial, 4) + cosine = torch.sum(left * right, dim=-1) + expected = { + 2: 0.5 * (3.0 * cosine**2 - 1.0), + 3: 0.5 * (5.0 * cosine**3 - 3.0 * cosine), + 4: 0.125 * (35.0 * cosine**4 - 30.0 * cosine**2 + 3.0), + } + for degree in range(2, 5): + current = torch.sum( + left_basis[:, 0, degree * degree : (degree + 1) ** 2] + * right_basis[:, 0, degree * degree : (degree + 1) ** 2], + dim=-1, + ) + torch.testing.assert_close( + current, expected[degree], atol=1e-12, rtol=1e-12 + ) + + def test_tetrahedral_and_octahedral_degree_signatures(self) -> None: + tetrahedral = self._tetrahedral_directions() + octahedral = torch.tensor( + [ + [1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, -1.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 0.0, -1.0], + ], + dtype=self.dtype, + device=env.DEVICE, + ) + + def moments(direction: torch.Tensor) -> torch.Tensor: + rr = torch.zeros( + direction.shape[0], + 1, + 4, + dtype=self.dtype, + device=env.DEVICE, + ) + radial = torch.ones( + direction.shape[0], + 1, + 1, + dtype=self.dtype, + device=env.DEVICE, + ) + return _build_moment_basis(rr, direction[:, None, :], radial, 4).sum(dim=0)[ + 0 + ] + + tetrahedral_moment = moments(tetrahedral) + octahedral_moment = moments(octahedral) + torch.testing.assert_close( + tetrahedral_moment[4:9], + torch.zeros_like(tetrahedral_moment[4:9]), + atol=1e-12, + rtol=0.0, + ) + self.assertGreater( + torch.linalg.vector_norm(tetrahedral_moment[9:16]).item(), 1.0 + ) + torch.testing.assert_close( + octahedral_moment[1:16], + torch.zeros_like(octahedral_moment[1:16]), + atol=1e-12, + rtol=0.0, + ) + self.assertGreater( + torch.linalg.vector_norm(octahedral_moment[16:25]).item(), 1.0 + ) + + def test_degree_gain_zero_recovers_lmax_one(self) -> None: + neighbors = self._square_directions() + degree_one = self._build_descriptor(lmax=1).eval() + degree_four = self._build_descriptor(lmax=4).eval() + assert degree_four.se_atten.adam_degree_gain_raw is not None + degree_four.se_atten.adam_degree_gain_raw.data.zero_() + result_one, _ = self._evaluate(degree_one, neighbors) + result_four, _ = self._evaluate(degree_four, neighbors) + torch.testing.assert_close(result_four, result_one, atol=1e-12, rtol=1e-12) + + raw_gain = degree_four.se_atten.adam_degree_gain_raw + raw_gain.data.copy_( + torch.tensor([0.1, 0.2, 0.3], dtype=self.dtype, device=env.DEVICE) + ) + degree_weights = _build_degree_weights(raw_gain, 4, result_four) + self.assertTrue(bool(torch.all(degree_weights >= 0.0))) + output, _ = self._evaluate(degree_four, neighbors) + (gradient,) = torch.autograd.grad(output.sum(), raw_gain) + self.assertGreater(torch.linalg.vector_norm(gradient).item(), 0.0) + + def test_level_zero_shares_degree_gain_parameter(self) -> None: + base = self._build_descriptor(lmax=4) + branch = self._build_descriptor(lmax=4) + base_gain = base.se_atten.adam_degree_gain_raw + branch_gain = branch.se_atten.adam_degree_gain_raw + assert base_gain is not None + assert branch_gain is not None + self.assertIsNot(branch_gain, base_gain) + + branch.share_params(base, shared_level=0, resume=True) + + self.assertIs(branch.se_atten.adam_degree_gain_raw, base_gain) + + def test_lmax_two_resolves_quadrupole_collision(self) -> None: + square = self._square_directions() + tetrahedral = self._tetrahedral_directions() + + degree_one = self._build_descriptor(lmax=1).eval() + square_l1, _ = self._evaluate(degree_one, square) + tetrahedral_l1, _ = self._evaluate(degree_one, tetrahedral) + torch.testing.assert_close(square_l1, tetrahedral_l1, atol=1e-12, rtol=1e-12) + + degree_two = self._build_descriptor(lmax=2).eval() + square_l2, _ = self._evaluate(degree_two, square) + tetrahedral_l2, _ = self._evaluate(degree_two, tetrahedral) + self.assertGreater( + torch.linalg.vector_norm(square_l2 - tetrahedral_l2).item(), + 1e-8, + ) + + serialized = degree_two.serialize() + self.assertEqual(serialized["@version"], 4) + restored = DescrptDPA1.deserialize(serialized).to(env.DEVICE).eval() + restored_square, _ = self._evaluate(restored, square) + torch.testing.assert_close(restored_square, square_l2) + scripted_square, _ = self._evaluate(torch.jit.script(degree_two), square) + torch.testing.assert_close(scripted_square, square_l2) + + dp_descriptor = DPDescrptDPA1.deserialize(serialized) + dp_coord = np.concatenate( + (np.zeros((1, 3)), square.detach().cpu().numpy()), + axis=0, + ).reshape(1, -1) + dp_atype = np.zeros((1, 5), dtype=np.int64) + dp_nlist = np.array([[[1, 2, 3, 4]]], dtype=np.int64) + dp_square = dp_descriptor.call(dp_coord, dp_atype, dp_nlist)[0] + np.testing.assert_allclose( + square_l2.detach().cpu().numpy(), + dp_square, + atol=1e-12, + rtol=1e-12, + ) + + def test_env_protection_regularizes_higher_degree_derivatives(self) -> None: + protection = 1e-2 + gradient_norms = [] + for scale in (1e-4, 1e-7): + diff = ( + torch.tensor( + [[[1.0, 0.4, -0.2]]], + dtype=self.dtype, + device=env.DEVICE, + ) + * scale + ).requires_grad_(True) + direction, distance, direction_mask = _safe_direction( + diff, + protection, + ) + radial = _compute_angular_radial( + distance, + direction_mask, + torch.ones((1, 1, 1), dtype=self.dtype, device=env.DEVICE), + torch.ones((1, 1, 1), dtype=self.dtype, device=env.DEVICE), + torch.ones((1, 1), dtype=torch.bool, device=env.DEVICE), + protection, + ) + basis = _build_moment_basis( + torch.zeros((1, 1, 4), dtype=self.dtype, device=env.DEVICE), + direction, + radial, + 2, + ) + expected_basis = build_dpa1_moment_basis( + np.zeros((1, 1, 4)), + diff.detach().cpu().numpy(), + np.ones((1, 1, 1)), + np.ones((1, 1, 1)), + np.ones((1, 1), dtype=bool), + 2, + protection, + ) + np.testing.assert_allclose( + basis.detach().cpu().numpy(), + expected_basis, + atol=1e-12, + rtol=1e-12, + ) + cotangent = torch.tensor( + [0.2, -0.7, 1.1, 0.3, -0.4], + dtype=self.dtype, + device=env.DEVICE, + ) + (gradient,) = torch.autograd.grad( + (basis[..., 4:] * cotangent).sum(), + diff, + ) + self.assertTrue(torch.isfinite(gradient).all()) + gradient_norms.append(torch.linalg.vector_norm(gradient)) + + self.assertLess( + gradient_norms[1].item(), + gradient_norms[0].item() * 1e-2, + ) + + def test_higher_degrees_are_rotation_and_permutation_invariant(self) -> None: + neighbors = torch.tensor( + [ + [1.1, 0.2, -0.1], + [-0.4, 0.9, 0.3], + [0.2, -0.5, 1.2], + [-0.7, -0.3, -0.8], + ], + dtype=self.dtype, + device=env.DEVICE, + ) + rotation = torch.tensor( + [ + [-2.0 / 3.0, 2.0 / 15.0, 11.0 / 15.0], + [2.0 / 3.0, -1.0 / 3.0, 2.0 / 3.0], + [1.0 / 3.0, 14.0 / 15.0, 2.0 / 15.0], + ], + dtype=self.dtype, + device=env.DEVICE, + ) + + for lmax in (2, 3, 4): + with self.subTest(lmax=lmax): + descriptor = self._build_descriptor(lmax=lmax).eval() + descriptor.se_atten.mean[..., 0] = 0.25 + descriptor.se_atten.stddev[..., 0] = 1.75 + reference, _ = self._evaluate(descriptor, neighbors) + rotated, _ = self._evaluate(descriptor, neighbors @ rotation.T) + permuted, _ = self._evaluate( + descriptor, + neighbors[[2, 0, 3, 1]], + ) + torch.testing.assert_close( + rotated, + reference, + atol=1e-10, + rtol=1e-10, + ) + torch.testing.assert_close( + permuted, + reference, + atol=1e-10, + rtol=1e-10, + ) + + def test_higher_degree_coordinate_derivatives_are_finite(self) -> None: + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.0, 0.2, 0.1], [-0.3, 0.9, -0.2]]], + dtype=self.dtype, + device=env.DEVICE, + requires_grad=True, + ) + atype = torch.zeros((1, 3), dtype=torch.long, device=env.DEVICE) + nlist = torch.tensor( + [[[1, 2, -1, -1]]], + dtype=torch.long, + device=env.DEVICE, + ) + for lmax in (2, 3, 4): + with self.subTest(lmax=lmax): + descriptor = self._build_descriptor(lmax=lmax) + assert descriptor.se_atten.adam_degree_gain_raw is not None + descriptor.se_atten.adam_degree_gain_raw.data.copy_( + torch.tensor( + [0.7, -0.5, 0.9][: lmax - 1], + dtype=self.dtype, + device=env.DEVICE, + ) + ) + current_coord = coord.detach().clone().requires_grad_(True) + result = descriptor(current_coord.reshape(1, -1), atype, nlist)[0] + cotangent = torch.linspace( + -0.8, + 1.1, + result.numel(), + dtype=self.dtype, + device=env.DEVICE, + ).reshape_as(result) + (first_derivative,) = torch.autograd.grad( + (result * cotangent).sum(), + current_coord, + create_graph=True, + ) + epsilon = 1e-6 + finite_difference = torch.empty_like(current_coord) + flat_coord = current_coord.detach().reshape(-1) + for index in range(flat_coord.numel()): + positive = flat_coord.clone() + negative = flat_coord.clone() + positive[index] += epsilon + negative[index] -= epsilon + positive_value = descriptor( + positive.reshape(1, -1), + atype, + nlist, + )[0] + negative_value = descriptor( + negative.reshape(1, -1), + atype, + nlist, + )[0] + finite_difference.reshape(-1)[index] = ( + (positive_value * cotangent).sum() + - (negative_value * cotangent).sum() + ) / (2.0 * epsilon) + (second_derivative,) = torch.autograd.grad( + first_derivative.square().sum(), + current_coord, + ) + torch.testing.assert_close( + first_derivative, + finite_difference, + atol=2e-8, + rtol=2e-8, + ) + self.assertTrue(torch.isfinite(first_derivative).all()) + self.assertTrue(torch.isfinite(second_derivative).all()) + self.assertGreater(first_derivative.abs().max().item(), 0.0) + self.assertGreater(second_derivative.abs().max().item(), 0.0) diff --git a/source/tests/pt/model/test_se_atten_v2.py b/source/tests/pt/model/test_se_atten_v2.py index e2d34ead5e..97e9da8b55 100644 --- a/source/tests/pt/model/test_se_atten_v2.py +++ b/source/tests/pt/model/test_se_atten_v2.py @@ -33,6 +33,48 @@ class TestDescrptSeAttenV2(unittest.TestCase, TestCaseSingleFrameWithNlist): def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) + def test_lmax_two_serialization(self) -> None: + descriptor = DescrptSeAttenV2( + self.rcut, + self.rcut_smth, + self.sel_mix, + self.nt, + lmax=2, + attn_layer=0, + precision="float64", + seed=GLOBAL_SEED, + ).to(env.DEVICE) + coord = torch.tensor(self.coord_ext, dtype=torch.float64, device=env.DEVICE) + atype = torch.tensor(self.atype_ext, dtype=torch.long, device=env.DEVICE) + nlist = torch.tensor(self.nlist, dtype=torch.long, device=env.DEVICE) + + result = descriptor(coord, atype, nlist)[0] + restored = DescrptSeAttenV2.deserialize(descriptor.serialize()).to(env.DEVICE) + restored_result = restored(coord, atype, nlist)[0] + + self.assertEqual(restored.se_atten.lmax, 2) + torch.testing.assert_close(restored_result, result) + + def test_lmax_four_degree_gain_serialization(self) -> None: + descriptor = DescrptSeAttenV2( + self.rcut, + self.rcut_smth, + self.sel_mix, + self.nt, + lmax=4, + attn_layer=0, + precision="float64", + seed=GLOBAL_SEED, + ).to(env.DEVICE) + restored = DescrptSeAttenV2.deserialize(descriptor.serialize()).to(env.DEVICE) + self.assertEqual(restored.se_atten.lmax, 4) + assert descriptor.se_atten.adam_degree_gain_raw is not None + assert restored.se_atten.adam_degree_gain_raw is not None + torch.testing.assert_close( + restored.se_atten.adam_degree_gain_raw, + descriptor.se_atten.adam_degree_gain_raw, + ) + def test_consistency( self, ) -> None: diff --git a/source/tests/pt/test_tabulate_fusion_se_atten.py b/source/tests/pt/test_tabulate_fusion_se_atten.py index 925d346880..e8325091cd 100644 --- a/source/tests/pt/test_tabulate_fusion_se_atten.py +++ b/source/tests/pt/test_tabulate_fusion_se_atten.py @@ -1644,6 +1644,145 @@ def test_second_order_backward(self) -> None: (self.em_x_tensor, self.em_tensor, self.two_embed_tensor), ) + def test_extended_component_basis(self) -> None: + for basis_dim in (9, 16, 25): + with self.subTest(basis_dim=basis_dim): + base_basis = self.em_tensor.detach() + base_output = self.expected_descriptor_tensor + extra_basis = [] + extra_output = [] + for row in range(basis_dim - 4): + first = row % 4 + second = (row + 1) % 4 + scale = 0.25 * (row + 1) + extra_basis.append( + scale * base_basis[..., first] - base_basis[..., second] + ) + extra_output.append( + scale * base_output[:, first] - base_output[:, second] + ) + basis = torch.cat( + (base_basis, torch.stack(extra_basis, dim=-1)), + dim=-1, + ).requires_grad_(True) + expected_output = torch.cat( + (base_output, torch.stack(extra_output, dim=1)), + dim=1, + ) + + output = torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor, + basis, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + )[0] + unsorted_output = torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor, + basis, + self.two_embed_tensor, + self.last_layer_size, + False, + )[0] + torch.testing.assert_close( + output, + expected_output, + atol=self.prec, + rtol=self.prec, + ) + torch.testing.assert_close( + unsorted_output, + expected_output, + atol=self.prec, + rtol=self.prec, + ) + + (basis_grad,) = torch.autograd.grad( + output.sum(), + basis, + retain_graph=True, + ) + expected_grad = self.expected_dy_dem[..., :1].expand_as(basis) + torch.testing.assert_close( + basis_grad, + expected_grad, + atol=self.prec, + rtol=self.prec, + ) + assert_second_order_backward_matches_finite_difference( + output, + (self.em_x_tensor, basis, self.two_embed_tensor), + ) + + @unittest.skipIf(not torch.cuda.is_available(), "CUDA is not available") + def test_wide_lmax_four_backward_launch(self) -> None: + """Exercise backward above the default CUDA shared-memory limit.""" + (dtype,) = self.param + width = 256 if dtype == torch.float64 else 512 + generator = torch.Generator(device=env.DEVICE).manual_seed(17) + table = torch.randn( + 20, + width * 6, + dtype=dtype, + device=env.DEVICE, + generator=generator, + ) + table_info = torch.tensor( + [0.0, 1.0, 2.0, 0.1, 0.2], + dtype=dtype, + device="cpu", + ) + em_x = torch.full( + (1, 2), + 0.05, + dtype=dtype, + device=env.DEVICE, + requires_grad=True, + ) + basis = torch.randn( + 1, + 2, + 25, + dtype=dtype, + device=env.DEVICE, + generator=generator, + requires_grad=True, + ) + two_embed = torch.randn( + 2, + width, + dtype=dtype, + device=env.DEVICE, + generator=generator, + requires_grad=True, + ) + output = torch.ops.deepmd.tabulate_fusion_se_atten( + table, + table_info, + em_x, + basis, + two_embed, + width, + False, + )[0] + first_gradients = torch.autograd.grad( + output.sum(), + (em_x, basis, two_embed), + create_graph=True, + ) + second_gradients = torch.autograd.grad( + sum(gradient.sum() for gradient in first_gradients), + (em_x, basis, two_embed), + allow_unused=True, + ) + for gradient in (*first_gradients, *second_gradients): + if gradient is not None: + self.assertTrue(torch.isfinite(gradient).all()) + @parameterized((torch.float64, torch.float32)) @unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index 8fd4e77598..73c68a23ae 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -44,6 +44,35 @@ def test_get_numb_attn_layer(self, attn) -> None: ) assert dd.get_numb_attn_layer() == attn + @pytest.mark.parametrize("trainable", [True, False]) + def test_degree_gain_preserves_trainable_after_deserialization( + self, trainable: bool + ) -> None: + descriptor = DescrptDPA1( + self.rcut, + self.rcut_smth, + self.sel_mix, + self.nt, + lmax=4, + attn_layer=0, + precision="float64", + seed=GLOBAL_SEED, + trainable=trainable, + ) + restored = DescrptDPA1.deserialize(descriptor.serialize()) + key = "se_atten.adam_degree_gain_raw" + assert key in dict(descriptor.named_parameters()) + assert key in dict(restored.named_parameters()) + assert restored.se_atten.trainable is trainable + assert dict(restored.named_parameters())[key].requires_grad is trainable + assert all( + parameter.requires_grad is trainable for parameter in restored.parameters() + ) + torch.testing.assert_close( + dict(restored.named_parameters())[key], + dict(descriptor.named_parameters())[key], + ) + @pytest.mark.parametrize("idt", [False, True]) # resnet_dt @pytest.mark.parametrize("sm", [False, True]) # smooth_type_embedding @pytest.mark.parametrize("to", [False, True]) # type_one_side diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index 97920c65d8..e14dc81610 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -63,6 +63,14 @@ def _cuda_ops_loaded() -> bool: "CUDA and the compiled deepmd op library are required", ) +# Instantiating the lmax 2/3/4 CUDA specializations dominates build time, so +# they sit behind the DEEPMD_ENABLE_DPA1_HIGH_LMAX CMake option, which is off by +# default, and `DescrptDPA1._fused_eligible` correspondingly serves only lmax=1. +# The tests below are verified locally against a build with that option enabled; +# they cannot run against a default build and are therefore not exercised in +# continuous integration. +_HIGH_LMAX = unittest.skip("requires a build with DEEPMD_ENABLE_DPA1_HIGH_LMAX=ON") + class _CudaLevel: """Context manager pinning ``DP_CUDA_INFER`` and restoring it on exit.""" @@ -93,6 +101,7 @@ def _build_dpa1_expt( ntypes=2, axis_neuron=4, tebd_dim=8, + lmax=1, ): from deepmd.pt_expt.descriptor.dpa1 import ( DescrptDPA1, @@ -114,6 +123,7 @@ def _build_dpa1_expt( smooth_type_embedding=smooth, precision="float32", seed=1, + lmax=lmax, ).to(device) des.eval() return des @@ -185,6 +195,27 @@ def _assert_parity(self, des) -> None: def test_parity_two_side_tanh(self) -> None: self._assert_parity(_build_dpa1_expt(self.device, [32, 64, 128])) + @_HIGH_LMAX + def test_parity_lmax_two(self) -> None: + descriptor = _build_dpa1_expt(self.device, [16, 32, 64], lmax=2) + assert descriptor.se_atten.adam_degree_gain_raw is not None + descriptor.se_atten.adam_degree_gain_raw.data.fill_(0.7) + self._assert_parity(descriptor) + + @_HIGH_LMAX + def test_parity_lmax_two_wide(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [32, 64, 128], lmax=2)) + + def test_lmax_four_declines_uncompressed_cuda(self) -> None: + descriptor = _build_dpa1_expt( + self.device, + [16, 16, 16], + act="silu", + tebd_input_mode="strip", + lmax=4, + ) + self.assertFalse(descriptor._fused_eligible("cuda")) + def test_parity_one_side_silu_resnet_dt(self) -> None: self._assert_parity( _build_dpa1_expt( @@ -305,6 +336,17 @@ def test_parity_strip_smooth_two_side(self) -> None: _build_dpa1_expt(self.device, [32, 64, 128], tebd_input_mode="strip") ) + @_HIGH_LMAX + def test_parity_strip_lmax_two(self) -> None: + self._assert_strip_parity( + _build_dpa1_expt( + self.device, + [16, 32, 64], + tebd_input_mode="strip", + lmax=2, + ) + ) + def test_parity_strip_smooth_one_side_silu(self) -> None: self._assert_strip_parity( _build_dpa1_expt( @@ -383,6 +425,7 @@ def _build_compressed_dpa1( ntypes=2, axis_neuron=4, tebd_dim=8, + lmax=1, ): """Strip DPA1 with the geometric embedding tabulated (``geo_compress``).""" des = _build_dpa1_expt( @@ -396,6 +439,7 @@ def _build_compressed_dpa1( ntypes=ntypes, axis_neuron=axis_neuron, tebd_dim=tebd_dim, + lmax=lmax, ) des.enable_compression(min_nbor_dist) des.to(device) @@ -490,6 +534,31 @@ def test_parity_two_side_smooth(self) -> None: # NG = 64: eight warps active in the moment backward. self._assert_parity(_build_compressed_dpa1(self.device, [16, 32, 64])) + @_HIGH_LMAX + def test_parity_lmax_two(self) -> None: + descriptor = _build_compressed_dpa1(self.device, [16, 32, 64], lmax=2) + assert descriptor.se_atten.adam_degree_gain_raw is not None + descriptor.se_atten.adam_degree_gain_raw.data.fill_(0.7) + self._assert_parity(descriptor) + + @_HIGH_LMAX + def test_parity_high_degrees(self) -> None: + for lmax in (3, 4): + with self.subTest(lmax=lmax): + descriptor = _build_compressed_dpa1( + self.device, + [16, 32, 64], + lmax=lmax, + ) + assert descriptor.se_atten.adam_degree_gain_raw is not None + descriptor.se_atten.adam_degree_gain_raw.data.copy_( + torch.tensor( + [0.7, -0.5, 0.9][: lmax - 1], + device=self.device, + ) + ) + self._assert_parity(descriptor) + def test_parity_wide_two_side(self) -> None: # NG = 128: the moment backward covers the table in two channel blocks. self._assert_parity(_build_compressed_dpa1(self.device, [32, 64, 128])) @@ -568,7 +637,7 @@ def run(graph) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: for output32, output64 in zip(outputs32, outputs64, strict=True): torch.testing.assert_close(output32, output64, atol=1e-6, rtol=1e-6) - def test_compact_canonical_descriptor_parity(self) -> None: + def _assert_compact_canonical_descriptor_parity(self, lmax: int) -> None: """Source-only topology matches the generic canonical operator.""" from deepmd.dpmodel.utils.neighbor_graph import ( canonicalize_neighbor_graph, @@ -583,7 +652,11 @@ def test_compact_canonical_descriptor_parity(self) -> None: canonical_graph_from_neighbor_graph, ) - des = _build_compressed_dpa1(self.device, [16, 32, 64]) + des = _build_compressed_dpa1( + self.device, + [16, 32, 64], + lmax=lmax, + ) graph, atype, _ = self._graph_and_dense(des) graph = canonicalize_neighbor_graph( dataclasses.replace(graph, n_local=graph.n_node), @@ -593,6 +666,11 @@ def test_compact_canonical_descriptor_parity(self) -> None: type_embedding = des.type_embedding.call() se = des.se_atten inverse_stddev = torch.reciprocal(se.stddev[:, 0, :]).contiguous() + degree_gain = ( + se.adam_degree_gain_raw.to(torch.float32).contiguous() + if se.adam_degree_gain_raw is not None + else des.compress_data[0].new_empty(0) + ) lower, upper, table_max, stride0, stride1 = ( float(value) for value in des.compress_info[0].tolist()[:5] ) @@ -612,6 +690,7 @@ def test_compact_canonical_descriptor_parity(self) -> None: type_embedding, se.mean[:, 0, :].contiguous(), inverse_stddev, + degree_gain, des.compress_data[0].contiguous(), des.type_embd_data.contiguous(), int(se.type_one_side), @@ -628,6 +707,7 @@ def test_compact_canonical_descriptor_parity(self) -> None: float(se.rcut_smth), float(se.env_protection), float(se.nnei), + (int(se.lmax) + 1) ** 2, ) torch.testing.assert_close(descriptor, generic_descriptor) @@ -660,6 +740,7 @@ def test_compact_canonical_descriptor_parity(self) -> None: atype, se.mean[:, 0, :].contiguous(), inverse_stddev, + degree_gain, des.compress_data[0].contiguous(), des.type_embd_data.contiguous(), int(se.type_one_side), @@ -687,6 +768,13 @@ def test_compact_canonical_descriptor_parity(self) -> None: torch.zeros_like(compact_gradient[physical_edge_count:]), ) + def test_compact_canonical_descriptor_parity(self) -> None: + self._assert_compact_canonical_descriptor_parity(1) + + @_HIGH_LMAX + def test_compact_canonical_descriptor_parity_lmax_four(self) -> None: + self._assert_compact_canonical_descriptor_parity(4) + def test_adaptive_resource_selection_large_graph(self) -> None: """First-use tuning preserves the reference on a non-trivial graph.""" generator = torch.Generator(device=self.device).manual_seed(37) @@ -1141,7 +1229,7 @@ def _graph(self, des): graph = from_dense_quartet(ec, nl, mp, compact=True, canonicalize=True) return graph, self.atype.reshape(-1).to(self.device) - def test_parity_vs_separate_ops(self) -> None: + def _assert_parity_vs_separate_ops(self, lmax: int) -> None: from deepmd.kernels.cuda.dpa1.graph_energy_force import ( dpa1_graph_energy_force, ) @@ -1150,7 +1238,14 @@ def test_parity_vs_separate_ops(self) -> None: ) # A doubling stack exercises the retiled backward inside the fusion. - des = _build_dpa1_expt(self.device, [8, 16, 32], act="silu") + des = _build_dpa1_expt( + self.device, + [8, 16, 32], + act="silu", + lmax=lmax, + ) + if des.se_atten.adam_degree_gain_raw is not None: + des.se_atten.adam_degree_gain_raw.data.fill_(0.7) fit = self._build_fitting(des.get_dim_out()) graph, atype = self._graph(des) tebd = des.type_embedding.call() @@ -1203,6 +1298,37 @@ def test_parity_vs_separate_ops(self) -> None: torch.testing.assert_close(virial, r_virial, atol=1e-4, rtol=1e-4) torch.testing.assert_close(atom_vir, r_atom_vir, atol=1e-4, rtol=1e-4) + def test_parity_vs_separate_ops(self) -> None: + self._assert_parity_vs_separate_ops(1) + + @_HIGH_LMAX + def test_parity_vs_separate_ops_lmax_two(self) -> None: + self._assert_parity_vs_separate_ops(2) + + def test_level_two_graph_export(self) -> None: + """The fused energy-force CPU implementation preserves its operator ABI.""" + from deepmd.pt_expt.model import ( + EnergyModel, + ) + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + ) + + cpu = torch.device("cpu") + des = _build_dpa1_expt(cpu, [8, 16, 32], act="silu") + original_device = self.device + self.device = cpu + fit = self._build_fitting(des.get_dim_out()) + self.device = original_device + model = EnergyModel(des, fit, type_map=["A", "B"]).eval() + with _CudaLevel("2"): + exported, _metadata, _model_json, _output_keys = _trace_and_export( + {"model": model.serialize()}, + do_atomic_virial=True, + lower_kind="graph", + ) + self.assertIsInstance(exported, torch.export.ExportedProgram) + def test_missing_csr_declines_energy_force_fusion(self) -> None: """The caller can fall back when optional CSR views are absent.""" for compressed in (False, True): @@ -1406,7 +1532,7 @@ def _graph(self, des): graph = from_dense_quartet(ec, nl, mp, compact=True, canonicalize=True) return graph, self.atype.reshape(-1).to(self.device) - def test_parity_vs_separate_ops(self) -> None: + def _assert_parity_vs_separate_ops(self, lmax: int) -> None: from deepmd.kernels.cuda.dpa1.graph_compress import ( dpa1_graph_compress, dpa1_graph_compress_energy_force, @@ -1428,7 +1554,15 @@ def test_parity_vs_separate_ops(self) -> None: act="silu", ntypes=4, axis_neuron=16, + lmax=lmax, ) + if des.se_atten.adam_degree_gain_raw is not None: + des.se_atten.adam_degree_gain_raw.data.copy_( + torch.tensor( + [0.7, -0.5, 0.9][: lmax - 1], + device=self.device, + ) + ) self.assertTrue(mega_eligible(des)) fit = self._build_fitting(des.get_dim_out(), ntypes=4) graph, atype = self._graph(des) @@ -1483,6 +1617,13 @@ def test_parity_vs_separate_ops(self) -> None: torch.testing.assert_close(virial, r_virial, atol=1e-4, rtol=1e-4) torch.testing.assert_close(atom_vir, r_atom_vir, atol=1e-4, rtol=1e-4) + def test_parity_vs_separate_ops(self) -> None: + self._assert_parity_vs_separate_ops(1) + + @_HIGH_LMAX + def test_parity_vs_separate_ops_lmax_four(self) -> None: + self._assert_parity_vs_separate_ops(4) + def test_compact_canonical_model_trace(self) -> None: """The eight-tensor deployment forward composes under symbolic make_fx.""" from deepmd.pt_expt.model import ( diff --git a/source/tests/pt_expt/descriptor/test_se_atten_v2.py b/source/tests/pt_expt/descriptor/test_se_atten_v2.py index 8f98e67510..0fe76c2c5f 100644 --- a/source/tests/pt_expt/descriptor/test_se_atten_v2.py +++ b/source/tests/pt_expt/descriptor/test_se_atten_v2.py @@ -36,6 +36,22 @@ def setup_method(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.device = env.DEVICE + def test_frozen_degree_gain_remains_frozen_after_deserialization(self) -> None: + descriptor = DescrptSeAttenV2( + self.rcut, + self.rcut_smth, + self.sel_mix, + self.nt, + lmax=4, + attn_layer=0, + trainable=False, + precision="float64", + seed=GLOBAL_SEED, + ) + restored = DescrptSeAttenV2.deserialize(descriptor.serialize()) + assert restored.se_atten.trainable is False + assert not any(parameter.requires_grad for parameter in restored.parameters()) + @pytest.mark.parametrize("idt", [False, True]) # resnet_dt @pytest.mark.parametrize("to", [False, True]) # type_one_side @pytest.mark.parametrize("prec", ["float64"]) # precision