diff --git a/openequivariance/benchmark/benchmark_configs.py b/openequivariance/benchmark/benchmark_configs.py deleted file mode 100644 index 9fec8f63..00000000 --- a/openequivariance/benchmark/benchmark_configs.py +++ /dev/null @@ -1,141 +0,0 @@ -from openequivariance.benchmark.tpp_creation_utils import ( - FullyConnectedTPProblem as FCTPP, -) -from openequivariance.benchmark.tpp_creation_utils import ChannelwiseTPP as CTPP - -# source: https://github.com/e3nn/e3nn/blob/main/examples/tetris.py -# running tetris will output the layers. I've only extracted the fully connected layers here. -e3nn_torch_tetris = [ - # 0th Layer - FCTPP("1x0e", "1x0e", "150x0e + 50x1o + 50x2e"), # sc - FCTPP("1x0e", "1x0e", "1x0e"), # lin1 - FCTPP("1x0e + 1x1o + 1x2e", "1x0e", "150x0e + 50x1o + 50x2e"), # lin2 - FCTPP("1x0e + 1x1o + 1x2e", "1x0e", "1x0e"), # alpha - # 1st Layer - FCTPP( - "50x0e + 50x1o + 50x2e", "1x0e", "250x0e + 50x1o + 50x1e + 50x2o + 50x2e" - ), # sc - FCTPP("50x0e + 50x1o + 50x2e", "1x0e", "50x0e + 50x1o + 50x2e"), # lin1 - # FCTPP("50x0e + 50x1o + 50x2e", "1x0e + 1x1o + 1x2e", "150x0e + 200x1o + 100x1e + 100x2o + 200x2e"), #tp - FCTPP( - "150x0e + 200x1o + 100x1e + 100x2o + 200x2e", - "1x0e", - "250x0e + 50x1o + 50x1e + 50x2o + 50x2e", - ), # lin2 - FCTPP("150x0e + 200x1o + 100x1e + 100x2o + 200x2e", "1x0e", "1x0e"), # alpha - # 2nd Layer - FCTPP( - "50x0e + 50x1o + 50x1e + 50x2o + 50x2e", - "1x0e", - "50x0o + 250x0e + 50x1o + 50x1e + 50x2o + 50x2e", - ), # sc - FCTPP( - "50x0e + 50x1o + 50x1e + 50x2o + 50x2e", - "1x0e", - "50x0e + 50x1o + 50x1e + 50x2o + 50x2e", - ), # lin1 - FCTPP( - "100x0o + 150x0e + 300x1o + 250x1e + 250x2o + 300x2e", - "1x0e", - "50x0o + 250x0e + 50x1o + 50x1e + 50x2o + 50x2e", - ), # lin2 - FCTPP( - "100x0o + 150x0e + 300x1o + 250x1e + 250x2o + 300x2e", "1x0e", "1x0e" - ), # alpha - # 3rd Layer - FCTPP("50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e", "1x0e", "1x0o + 6x0e"), # sc - FCTPP( - "50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e", - "1x0e", - "50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e", - ), # lin1 - FCTPP("150x0o + 150x0e", "1x0e", "1x0o + 6x0e"), # lin2 - FCTPP("150x0o + 150x0e", "1x0e", "1x0e"), # alpha -] - -# jax version can be found here, but doesn't directly translate -# https://github.com/e3nn/e3nn-jax/blob/main/examples/tetris_point.py - -# source: https://github.com/e3nn/e3nn/blob/f95297952303347a8a3cfe971efe449c710c43b2/examples/tetris_polynomial.py#L66-L68 -e3nn_torch_tetris_polynomial = [ - FCTPP( - "1x0e + 1x1o + 1x2e + 1x3o", - "1x0e + 1x1o + 1x2e + 1x3o", - "64x0e + 24x1e + 24x1o + 16x2e + 16x2o", - label="tetris-poly-1", - ), # tp1 - FCTPP( - "64x0e + 24x1e + 24x1o + 16x2e + 16x2o", - "1x0e + 1x1o + 1x2e", - "0o + 6x0e", - label="tetris-poly-2", - ), # tp2 -] - -# https://github.com/gcorso/DiffDock/blob/b4704d94de74d8cb2acbe7ec84ad234c09e78009/models/tensor_layers.py#L299 -# specific irreps come from Vivek's communication with DiffDock team -diffdock_configs = [ - FCTPP( - "10x1o + 10x1e + 48x0e + 48x0o", - "1x0e + 1x1o", - "10x1o + 10x1e + 48x0e + 48x0o", - shared_weights=False, - label="DiffDock-L=1", - ), - FCTPP( - "10x1o + 10x1e + 48x0e + 48x0o", - "1x0e + 1x1o + 1x2e", - "10x1o + 10x1e + 48x0e + 48x0o", - shared_weights=False, - label="DiffDock-L=2", - ), -] - -mace_conv = [ - ( - "128x0e+128x1o+128x2e", - "1x0e+1x1o+1x2e+1x3o", - "128x0e+128x1o+128x2e+128x3o", - "mace-large", - ), - ("128x0e+128x1o", "1x0e+1x1o+1x2e+1x3o", "128x0e+128x1o+128x2e", "mace-medium"), -] - -nequip_conv = [ - ( - "32x0o + 32x0e + 32x1o + 32x1e + 32x2o + 32x2e", - "0e + 1o + 2e", - "32x0o + 32x0e + 32x1o + 32x1e + 32x2o + 32x2e", - "nequip-lips", - ), - ( - "64x0o + 64x0e + 64x1o + 64x1e", - "0e + 1o", - "64x0o + 64x0e + 64x1o + 64x1e", - "nequip-revmd17-aspirin", - ), - ( - "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e", - "0e + 1o + 2e", - "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e", - "nequip-revmd17-toluene", - ), - ( - "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e", - "0e + 1o + 2e + 3o", - "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e", - "nequip-revmd17-benzene", - ), - ( - "32x0o + 32x0e + 32x1o + 32x1e", - "0e + 1o", - "32x0o + 32x0e + 32x1o + 32x1e", - "nequip-water", - ), -] - -mace_problems = [CTPP(*config) for config in mace_conv] - -mace_nequip_problems = [] -for config in mace_conv + nequip_conv: - mace_nequip_problems.append(CTPP(*config)) diff --git a/openequivariance/benchmark/benchmark_routines/paper_benchmark_uvw.py b/openequivariance/benchmark/benchmark_routines/paper_benchmark_uvw.py deleted file mode 100644 index 18d88ca3..00000000 --- a/openequivariance/benchmark/benchmark_routines/paper_benchmark_uvw.py +++ /dev/null @@ -1,75 +0,0 @@ -import itertools -import logging -import copy -import pathlib -from typing import List -import numpy as np - -from torch._functorch import config - -from openequivariance.benchmark.logging_utils import getLogger -from openequivariance.implementations.E3NNTensorProduct import ( - E3NNTensorProductCompiledCUDAGraphs, -) -from openequivariance.implementations.CUETensorProduct import CUETensorProduct -from openequivariance.implementations.TensorProduct import TensorProduct -from openequivariance.benchmark.TestBenchmarkSuite import ( - TestBenchmarkSuite, - TestDefinition, -) -from openequivariance.benchmark.benchmark_configs import ( - e3nn_torch_tetris_polynomial, - diffdock_configs, -) - -logger = getLogger() - - -@config.patch("donated_buffer", False) -def run_paper_uvw_benchmark(params) -> pathlib.Path: - problems = list(itertools.chain(e3nn_torch_tetris_polynomial, diffdock_configs)) - - float64_problems = copy.deepcopy(problems) - for problem in float64_problems: - problem.irrep_dtype = np.float64 - problem.weight_dtype = np.float64 - - problems += float64_problems - - implementations: List[TensorProduct] = [ - E3NNTensorProductCompiledCUDAGraphs, - CUETensorProduct, - TensorProduct, - ] - - tests = [ - TestDefinition( - implementation, problem, direction, correctness=False, benchmark=True - ) - for problem, direction, implementation in itertools.product( - problems, params.directions, implementations - ) - ] - - bench_suite = TestBenchmarkSuite( - num_warmup=100, - num_iter=100, - bench_batch_size=params.batch_size, - prng_seed=11111, - torch_op=True, - test_name="uvw", - ) - - logger.setLevel(logging.INFO) - data_folder = bench_suite.run(tests, output_folder=params.output_folder) - - if params.plot: - import openequivariance.benchmark.plotting as plotting - - plotting.plot_uvw(data_folder) - - return data_folder - - -if __name__ == "__main__": - run_paper_uvw_benchmark() diff --git a/openequivariance/benchmark/problems.py b/openequivariance/benchmark/problems.py new file mode 100644 index 00000000..efa79529 --- /dev/null +++ b/openequivariance/benchmark/problems.py @@ -0,0 +1,151 @@ +from openequivariance.benchmark.tpp_creation_utils import ( + FullyConnectedTPProblem as FCTPP, +) +from openequivariance.benchmark.tpp_creation_utils import ChannelwiseTPP as CTPP + +# source: https://github.com/e3nn/e3nn/blob/main/examples/tetris.py +# running tetris will output the layers. I've only extracted the fully connected layers here. +_e3nn_torch_tetris = [ + # 0th Layer + FCTPP("1x0e", "1x0e", "150x0e + 50x1o + 50x2e"), # sc + FCTPP("1x0e", "1x0e", "1x0e"), # lin1 + FCTPP("1x0e + 1x1o + 1x2e", "1x0e", "150x0e + 50x1o + 50x2e"), # lin2 + FCTPP("1x0e + 1x1o + 1x2e", "1x0e", "1x0e"), # alpha + # 1st Layer + FCTPP( + "50x0e + 50x1o + 50x2e", "1x0e", "250x0e + 50x1o + 50x1e + 50x2o + 50x2e" + ), # sc + FCTPP("50x0e + 50x1o + 50x2e", "1x0e", "50x0e + 50x1o + 50x2e"), # lin1 + # FCTPP("50x0e + 50x1o + 50x2e", "1x0e + 1x1o + 1x2e", "150x0e + 200x1o + 100x1e + 100x2o + 200x2e"), #tp + FCTPP( + "150x0e + 200x1o + 100x1e + 100x2o + 200x2e", + "1x0e", + "250x0e + 50x1o + 50x1e + 50x2o + 50x2e", + ), # lin2 + FCTPP("150x0e + 200x1o + 100x1e + 100x2o + 200x2e", "1x0e", "1x0e"), # alpha + # 2nd Layer + FCTPP( + "50x0e + 50x1o + 50x1e + 50x2o + 50x2e", + "1x0e", + "50x0o + 250x0e + 50x1o + 50x1e + 50x2o + 50x2e", + ), # sc + FCTPP( + "50x0e + 50x1o + 50x1e + 50x2o + 50x2e", + "1x0e", + "50x0e + 50x1o + 50x1e + 50x2o + 50x2e", + ), # lin1 + FCTPP( + "100x0o + 150x0e + 300x1o + 250x1e + 250x2o + 300x2e", + "1x0e", + "50x0o + 250x0e + 50x1o + 50x1e + 50x2o + 50x2e", + ), # lin2 + FCTPP( + "100x0o + 150x0e + 300x1o + 250x1e + 250x2o + 300x2e", "1x0e", "1x0e" + ), # alpha + # 3rd Layer + FCTPP("50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e", "1x0e", "1x0o + 6x0e"), # sc + FCTPP( + "50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e", + "1x0e", + "50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e", + ), # lin1 + FCTPP("150x0o + 150x0e", "1x0e", "1x0o + 6x0e"), # lin2 + FCTPP("150x0o + 150x0e", "1x0e", "1x0e"), # alpha +] + + +def e3nn_torch_tetris_poly_problems(): + # source: https://github.com/e3nn/e3nn/blob/f95297952303347a8a3cfe971efe449c710c43b2/examples/tetris_polynomial.py#L66-L68 + return [ + FCTPP( + "1x0e + 1x1o + 1x2e + 1x3o", + "1x0e + 1x1o + 1x2e + 1x3o", + "64x0e + 24x1e + 24x1o + 16x2e + 16x2o", + label="tetris-poly-1", + ), # tp1 + FCTPP( + "64x0e + 24x1e + 24x1o + 16x2e + 16x2o", + "1x0e + 1x1o + 1x2e", + "0o + 6x0e", + label="tetris-poly-2", + ), # tp2 + ] + + +# https://github.com/gcorso/DiffDock/blob/b4704d94de74d8cb2acbe7ec84ad234c09e78009/models/tensor_layers.py#L299 +# specific irreps come from Vivek's communication with DiffDock team +def diffdock_problems(): + return [ + FCTPP( + "10x1o + 10x1e + 48x0e + 48x0o", + "1x0e + 1x1o", + "10x1o + 10x1e + 48x0e + 48x0o", + shared_weights=False, + label="DiffDock-L=1", + ), + FCTPP( + "10x1o + 10x1e + 48x0e + 48x0o", + "1x0e + 1x1o + 1x2e", + "10x1o + 10x1e + 48x0e + 48x0o", + shared_weights=False, + label="DiffDock-L=2", + ), + ] + + +def mace_problems(): + return [ + CTPP(*config) + for config in [ + ( + "128x0e+128x1o+128x2e", + "1x0e+1x1o+1x2e+1x3o", + "128x0e+128x1o+128x2e+128x3o", + "mace-large", + ), + ( + "128x0e+128x1o", + "1x0e+1x1o+1x2e+1x3o", + "128x0e+128x1o+128x2e", + "mace-medium", + ), + ] + ] + + +def nequip_problems(): + return [ + CTPP(*config) + for config in [ + ( + "32x0o + 32x0e + 32x1o + 32x1e + 32x2o + 32x2e", + "0e + 1o + 2e", + "32x0o + 32x0e + 32x1o + 32x1e + 32x2o + 32x2e", + "nequip-lips", + ), + ( + "64x0o + 64x0e + 64x1o + 64x1e", + "0e + 1o", + "64x0o + 64x0e + 64x1o + 64x1e", + "nequip-revmd17-aspirin", + ), + ( + "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e", + "0e + 1o + 2e", + "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e", + "nequip-revmd17-toluene", + ), + ( + "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e", + "0e + 1o + 2e + 3o", + "64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e", + "nequip-revmd17-benzene", + ), + ( + "32x0o + 32x0e + 32x1o + 32x1e", + "0e + 1o", + "32x0o + 32x0e + 32x1o + 32x1e", + "nequip-water", + ), + ] + ] diff --git a/tests/batch_test.py b/tests/batch_test.py index fc02a863..9f23b75f 100644 --- a/tests/batch_test.py +++ b/tests/batch_test.py @@ -9,7 +9,7 @@ correctness_backward, correctness_double_backward, ) -from itertools import chain, product +from itertools import product class TPCorrectness: @@ -79,14 +79,18 @@ def test_tp_double_bwd(self, tp_and_problem): class TestProductionModels(TPCorrectness): - from openequivariance.benchmark.benchmark_configs import ( - e3nn_torch_tetris_polynomial, - diffdock_configs, - mace_nequip_problems, + from openequivariance.benchmark.problems import ( + e3nn_torch_tetris_poly_problems, + diffdock_problems, + mace_problems, + nequip_problems, ) - production_model_tpps = list( - chain(mace_nequip_problems, e3nn_torch_tetris_polynomial, diffdock_configs) + production_model_tpps = ( + mace_problems() + + nequip_problems() + + e3nn_torch_tetris_poly_problems() + + diffdock_problems() ) @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class") @@ -219,12 +223,12 @@ def problem(self, request, dtype): class TestSharedWeights(TPCorrectness): - from openequivariance.benchmark.benchmark_configs import ( + from openequivariance.benchmark.problems import ( mace_problems, - diffdock_configs, + diffdock_problems, ) - problems = [mace_problems[0], diffdock_configs[0]] + problems = [mace_problems()[0], diffdock_problems()[0]] def thresh(self, direction): return { diff --git a/tests/benchmark.py b/tests/benchmark.py index fe765792..01126909 100644 --- a/tests/benchmark.py +++ b/tests/benchmark.py @@ -14,6 +14,7 @@ from openequivariance.extlib import DeviceProp from openequivariance.implementations.E3NNTensorProduct import ( E3NNTensorProduct, + E3NNTensorProductCompiledCUDAGraphs, E3NNTensorProductCompiledMaxAutotuneCUDAGraphs, ) from openequivariance.implementations.TensorProduct import TensorProduct @@ -28,9 +29,6 @@ FullyConnectedTPProblem, SingleInstruction, ) -from openequivariance.benchmark.benchmark_routines.paper_benchmark_uvw import ( - run_paper_uvw_benchmark, -) from openequivariance.implementations.convolution.TensorProductConv import ( TensorProductConvAtomic, @@ -42,6 +40,15 @@ from openequivariance.implementations.convolution.CUEConv import CUEConv, CUEConvFused from openequivariance.benchmark.ConvBenchmarkSuite import ConvBenchmarkSuite, load_graph +from openequivariance.benchmark.problems import ( + e3nn_torch_tetris_poly_problems, + diffdock_problems, + mace_problems, + nequip_problems, +) + +from torch._functorch import config + logger = getLogger() CTPP = ChannelwiseTPP @@ -74,13 +81,11 @@ def benchmark_uvu(params): - from openequivariance.benchmark.benchmark_configs import mace_nequip_problems - - float64_problems = copy.deepcopy(mace_nequip_problems) + float64_problems = mace_problems() + nequip_problems() for problem in float64_problems: problem.irrep_dtype = np.float64 problem.weight_dtype = np.float64 - problems = mace_nequip_problems + float64_problems + problems = mace_problems() + nequip_problems() + float64_problems implementations = [implementation_map[impl] for impl in params.implementations] directions = params.directions @@ -157,6 +162,54 @@ def benchmark_uvu(params): plot({"data_folder": data_folder}) +@config.patch("donated_buffer", False) +def benchmark_uvw(params) -> pathlib.Path: + problems = list( + itertools.chain(e3nn_torch_tetris_poly_problems(), diffdock_problems()) + ) + + float64_problems = copy.deepcopy(problems) + for problem in float64_problems: + problem.irrep_dtype = np.float64 + problem.weight_dtype = np.float64 + + problems += float64_problems + + implementations = [ + E3NNTensorProductCompiledCUDAGraphs, + CUETensorProduct, + TensorProduct, + ] + + tests = [ + TestDefinition( + implementation, problem, direction, correctness=False, benchmark=True + ) + for problem, direction, implementation in itertools.product( + problems, params.directions, implementations + ) + ] + + bench_suite = TestBenchmarkSuite( + num_warmup=100, + num_iter=100, + bench_batch_size=params.batch_size, + prng_seed=11111, + torch_op=True, + test_name="uvw", + ) + + logger.setLevel(logging.INFO) + data_folder = bench_suite.run(tests, output_folder=params.output_folder) + + if params.plot: + import openequivariance.benchmark.plotting as plotting + + plotting.plot_uvw(data_folder) + + return data_folder + + def benchmark_roofline(params): implementations = [TensorProduct, CUETensorProduct] directions = ["forward", "backward"] @@ -272,13 +325,8 @@ def benchmark_convolution(params): def benchmark_double_backward(params): - from openequivariance.benchmark.benchmark_configs import ( - mace_nequip_problems, - diffdock_configs, - ) - implementations = [E3NNTensorProduct, CUETensorProduct, TensorProduct] - problems = diffdock_configs + mace_nequip_problems + problems = diffdock_problems() + mace_problems() + nequip_problems() float64_problems = copy.deepcopy(problems) for problem in float64_problems: @@ -306,12 +354,10 @@ def benchmark_double_backward(params): def benchmark_kahan_accuracy(params): - from openequivariance.benchmark.benchmark_configs import mace_problems - filenames = ["carbon_lattice_radius6.0.pickle"] graphs = download_graphs(params, filenames) implementations = [TensorProductConvAtomic, TensorProductConvKahan] - problems = [mace_problems[0]] + problems = [mace_problems()[0]] bench = ConvBenchmarkSuite( problems, test_name="kahan_convolution_accuracy", correctness_threshold=1e-4 @@ -468,7 +514,7 @@ def plot(params): choices=["forward", "backward"], ) parser_uvw.add_argument("--plot", action="store_true", help="Plot the results.") - parser_uvw.set_defaults(func=run_paper_uvw_benchmark) + parser_uvw.set_defaults(func=benchmark_uvw) parser_double_bwd = subparsers.add_parser( "double_backward", help="Run the higher derivative kernel benchmark" diff --git a/tests/conv_test.py b/tests/conv_test.py index ae5f8ca2..8c4c0140 100644 --- a/tests/conv_test.py +++ b/tests/conv_test.py @@ -6,7 +6,7 @@ import numpy as np import openequivariance as oeq from openequivariance.benchmark.ConvBenchmarkSuite import load_graph -from itertools import chain, product +from itertools import product class ConvCorrectness: @@ -106,12 +106,12 @@ def test_tp_double_bwd(self, conv_object, graph): class TestProductionModels(ConvCorrectness): - from openequivariance.benchmark.benchmark_configs import ( + from openequivariance.benchmark.problems import ( mace_problems, - diffdock_configs, + diffdock_problems, ) - production_model_tpps = list(chain(mace_problems, diffdock_configs)) + production_model_tpps = mace_problems() + diffdock_problems() @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class") def problem(self, request, dtype): @@ -207,12 +207,12 @@ def problem(self, request, dtype): class TestAtomicSharedWeights(ConvCorrectness): - from openequivariance.benchmark.benchmark_configs import ( + from openequivariance.benchmark.problems import ( mace_problems, - diffdock_configs, + diffdock_problems, ) - problems = [mace_problems[0], diffdock_configs[0]] + problems = [mace_problems()[0], diffdock_problems()[0]] def thresh(self, direction): return { diff --git a/tests/multidevice_test.py b/tests/multidevice_test.py index 7c017e77..7b7b48c7 100644 --- a/tests/multidevice_test.py +++ b/tests/multidevice_test.py @@ -35,9 +35,9 @@ def test_multidevice(): import openequivariance as oeq # Use MACE-large to test >64KB shared memory allocation - from openequivariance.benchmark.benchmark_configs import mace_problems + from openequivariance.benchmark.problems import mace_problems - problem = mace_problems[0] + problem = mace_problems()[0] local_rank = int(os.environ["LOCAL_RANK"]) device = f"cuda:{local_rank}"