From 19dca23bf9ca2de6a13bf8d466d2952a5bc9f060 Mon Sep 17 00:00:00 2001 From: Mit Kotak <53411468+mitkotak@users.noreply.github.com> Date: Thu, 13 Jun 2024 13:13:35 -0400 Subject: [PATCH] `o3.experimental.FullTensorProductv2` for `torch.compile(...., fullgraph=True)` (#436) --- .github/CHANGELOG.md | 1 + e3nn/o3/__init__.py | 4 +- e3nn/o3/_irreps.py | 17 +++++ e3nn/o3/experimental/__init__.py | 1 + e3nn/o3/experimental/_full_tp.py | 94 ++++++++++++++++++++++++++ e3nn/util/datatypes.py | 13 ++++ tests/o3/experimental/benchmark_pt2.py | 55 +++++++++++++++ tests/o3/experimental/test_fulltp.py | 15 ++++ 8 files changed, 199 insertions(+), 1 deletion(-) create mode 100644 e3nn/o3/experimental/__init__.py create mode 100644 e3nn/o3/experimental/_full_tp.py create mode 100644 e3nn/util/datatypes.py create mode 100644 tests/o3/experimental/benchmark_pt2.py create mode 100644 tests/o3/experimental/test_fulltp.py diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md index c486d067..f143d51b 100644 --- a/.github/CHANGELOG.md +++ b/.github/CHANGELOG.md @@ -6,6 +6,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- `o3.experimental.FullTensorProductv2` for compatibility with `torch.compile(..., fulgraph=True)` - enable `pip` caching in CI - refactor to use `pyproject.toml` for packaging - refactor `gh` community files diff --git a/e3nn/o3/__init__.py b/e3nn/o3/__init__.py index 2fdbf183..4c94c505 100644 --- a/e3nn/o3/__init__.py +++ b/e3nn/o3/__init__.py @@ -38,6 +38,8 @@ FullTensorProduct, TensorSquare, ) +from .experimental import FullTensorProductv2 + from ._spherical_harmonics import SphericalHarmonics, spherical_harmonics from ._angular_spherical_harmonics import ( SphericalHarmonicsAlphaBeta, @@ -101,7 +103,7 @@ "FullyConnectedTensorProduct", "ElementwiseTensorProduct", "FullTensorProduct", - "TensorSquare", + "FullTensorProductv2" "TensorSquare", "SphericalHarmonics", "spherical_harmonics", "SphericalHarmonicsAlphaBeta", diff --git a/e3nn/o3/_irreps.py b/e3nn/o3/_irreps.py index 02d05bf4..411a341a 100644 --- a/e3nn/o3/_irreps.py +++ b/e3nn/o3/_irreps.py @@ -605,6 +605,23 @@ def sort(self): irreps = Irreps([(mul, ir) for ir, _, mul in out]) return Ret(irreps, p, inv) + def regroup(self) -> "Irreps": + r"""Regroup the same irreps together. + + Equivalent to :meth:`sort` followed by :meth:`simplify`. + + Returns + ------- + irreps: `e3nn.o3.Irreps` + + Examples + -------- + + >>> Irreps("1e + 0e + 1e + 0x2e").regroup() + 1x0e+2x1e + """ + return self.sort().irreps.simplify() + @property def dim(self) -> int: return sum(mul * ir.dim for mul, ir in self) diff --git a/e3nn/o3/experimental/__init__.py b/e3nn/o3/experimental/__init__.py new file mode 100644 index 00000000..7fcd718f --- /dev/null +++ b/e3nn/o3/experimental/__init__.py @@ -0,0 +1 @@ +from ._full_tp import FullTensorProduct as FullTensorProductv2 diff --git a/e3nn/o3/experimental/_full_tp.py b/e3nn/o3/experimental/_full_tp.py new file mode 100644 index 00000000..44134ecd --- /dev/null +++ b/e3nn/o3/experimental/_full_tp.py @@ -0,0 +1,94 @@ +from e3nn.util.datatypes import Path, Chunk +from e3nn import o3 + +import torch +from torch import nn +import numpy as np + + +def _prepare_inputs(input1, input2): + dtype = torch.promote_types(input1.dtype, input2.dtype) + + input1 = input1.to(dtype=dtype) + input2 = input2.to(dtype=dtype) + + leading_shape = torch.broadcast_shapes(input1.shape[:-1], input2.shape[:-1]) + input1 = input1.broadcast_to(leading_shape + (-1,)) + input2 = input2.broadcast_to(leading_shape + (-1,)) + return input1, input2, leading_shape + + +class FullTensorProduct(nn.Module): + def __init__( + self, + irreps_in1: o3.Irreps, + irreps_in2: o3.Irreps, + *, + filter_ir_out: o3.Irreps = None, + irrep_normalization: str = "component", + regroup_output: bool = True, + ): + """Tensor Product adapted from https://github.com/e3nn/e3nn-jax/blob/cf37f3e95264b34587b3a202ea4c3eb82597307e/e3nn_jax/_src/tensor_products.py#L40-L135""" + super(FullTensorProduct, self).__init__() + + if regroup_output: + irreps_in1 = o3.Irreps(irreps_in1).regroup() + irreps_in2 = o3.Irreps(irreps_in2).regroup() + + paths = {} + irreps_out = [] + for (mul_1, ir_1), slice_1 in zip(irreps_in1, irreps_in1.slices()): + for (mul_2, ir_2), slice_2 in zip(irreps_in2, irreps_in2.slices()): + for ir_out in ir_1 * ir_2: + if filter_ir_out is not None and ir_out not in filter_ir_out: + continue + cg = o3.wigner_3j(ir_1.l, ir_2.l, ir_out.l) + if irrep_normalization == "component": + cg *= np.sqrt(ir_out.dim) + elif irrep_normalization == "norm": + cg *= np.sqrt(ir_1.dim * ir_2.dim) + else: + raise ValueError(f"irrep_normalization={irrep_normalization} not supported") + self.register_buffer(f"cg_{ir_1.l}_{ir_2.l}_{ir_out.l}", cg) + paths[(ir_1.l, ir_2.l, ir_out.l)] = Path( + Chunk(mul_1, ir_1.dim, slice_1), Chunk(mul_2, ir_2.dim, slice_2), Chunk(mul_1 * mul_2, ir_out.dim) + ) + irreps_out.append((mul_1 * mul_2, ir_out)) + self.paths = paths + irreps_out = o3.Irreps(irreps_out) + self.irreps_out, _, self.inv = irreps_out.sort() + self.irreps_in1 = irreps_in1 + self.irreps_in2 = irreps_in2 + + def forward( + self, + input1: torch.Tensor, + input2: torch.Tensor, + ) -> torch.Tensor: + input1, input2, leading_shape = _prepare_inputs(input1, input2) + chunks = [] + for (l1, l2, l3), ( + (mul_1, input_dim1, slice_1), + (mul_2, input_dim2, slice_2), + (output_mul, output_dim, _), + ) in self.paths.items(): + x1 = input1[..., slice_1].reshape( + leading_shape + + ( + mul_1, + input_dim1, + ) + ) + x2 = input2[..., slice_2].reshape( + leading_shape + + ( + mul_2, + input_dim2, + ) + ) + cg = getattr(self, f"cg_{l1}_{l2}_{l3}") + chunk = torch.einsum("...ui, ...vj, ijk -> ...uvk", x1, x2, cg) + chunk = torch.reshape(chunk, chunk.shape[:-3] + (output_mul * output_dim,)) + chunks.append(chunk) + + return torch.cat([chunks[i] for i in self.inv], dim=-1) diff --git a/e3nn/util/datatypes.py b/e3nn/util/datatypes.py new file mode 100644 index 00000000..4a8df48f --- /dev/null +++ b/e3nn/util/datatypes.py @@ -0,0 +1,13 @@ +from typing import NamedTuple, Optional + + +class Chunk(NamedTuple): + mul: int + dim: int + slice: Optional[slice] = None + + +class Path(NamedTuple): + input_1_slice: Chunk + input_2_slice: Chunk + output_slice: Chunk diff --git a/tests/o3/experimental/benchmark_pt2.py b/tests/o3/experimental/benchmark_pt2.py new file mode 100644 index 00000000..92fd4ad6 --- /dev/null +++ b/tests/o3/experimental/benchmark_pt2.py @@ -0,0 +1,55 @@ +import torch +from torch._inductor.utils import print_performance + +# Borrowed from https://github.com/pytorch-labs/gpt-fast/blob/db7b273ab86b75358bd3b014f1f022a19aba4797/generate.py#L16-L18 +torch.set_float32_matmul_precision("high") +import torch._dynamo.config +import torch._inductor.config + +torch._inductor.config.coordinate_descent_tuning = True +torch._inductor.config.triton.unique_kernel_names = True + +device = "cuda" +compile_mode = "max-autotune" # Bringing out all of the tricks that Torch 2.0 has but "reduce-overhead" should work as well + +from e3nn import o3, util +import numpy as np +from torch import nn +import time + +LMAX = 8 +CHANNEL = 128 +BATCH = 100 + + +def main(): + for lmax in range(1, LMAX + 1): + irreps = o3.Irreps.spherical_harmonics(lmax) + irreps_x = (CHANNEL * irreps).regroup() + x = irreps_x.randn(BATCH, -1).to(device=device) + irreps_y = irreps + y = irreps_y.randn(BATCH, -1).to(device=device) + print(f"{irreps_x} \otimes {irreps_y}") + + tp = o3.FullTensorProduct(irreps_x, irreps_y) # Doesnt work with fullgraph=True + + tp_jit_compile = util.jit.compile(tp).to(device=device) + + tp_compile = torch.compile(tp, mode=compile_mode).to(device=device) + print( + f"TP JIT lmax {lmax} channel {CHANNEL} batch {BATCH}: {print_performance(lambda: tp_jit_compile(x, y), times=100, repeat=10)*1000:.3f}ms" + ) + + print( + f"TP Torch 2.0 lmax {lmax} channel {CHANNEL} batch {BATCH}: {print_performance(lambda: tp_compile(x, y), times=100, repeat=10)*1000:.3f}ms" + ) + + tp_experimental = o3.experimental.FullTensorProductv2(irreps_x, irreps_y) + tp_experimental_compile = torch.compile(tp_experimental, mode=compile_mode, fullgraph=True).to(device=device) + print( + f"TP Experimental Torch 2.0 lmax {lmax} channel {CHANNEL} batch {BATCH}: {print_performance(lambda: tp_experimental_compile(x, y), times=100, repeat=10)*1000:.3f}ms" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/o3/experimental/test_fulltp.py b/tests/o3/experimental/test_fulltp.py new file mode 100644 index 00000000..a973ea72 --- /dev/null +++ b/tests/o3/experimental/test_fulltp.py @@ -0,0 +1,15 @@ +import torch +from e3nn import o3 +import pytest + + +@pytest.mark.parametrize("irreps_in1", ["0e", "0e + 1e"]) +@pytest.mark.parametrize("irreps_in2", ["2x0e", "2x0e + 3x1e"]) +def test_fulltp(irreps_in1, irreps_in2): + x = o3.Irreps(irreps_in1).randn(10, -1) + y = o3.Irreps(irreps_in2).randn(10, -1) + + tp_pt2 = torch.compile(o3.experimental.FullTensorProductv2(irreps_in1, irreps_in2), fullgraph=True) + tp = o3.FullTensorProduct(irreps_in1, irreps_in2) + + torch.testing.assert_close(tp_pt2(x, y), tp(x, y))