diff --git a/backend/benchmarking/rotate.py b/backend/benchmarking/rotate.py new file mode 100644 index 00000000..e724c136 --- /dev/null +++ b/backend/benchmarking/rotate.py @@ -0,0 +1,57 @@ +from elasticapp._rotations import rotate, rotate_scalar +from elasticapp._PyArrays import Tensor, Matrix +from elastica._rotations import _rotate +import numpy +import time + +# warm up jit for fair comparison +random_1 = numpy.random.random((3, 3, 1)) +random_2 = numpy.random.random((3, 1)) +out1 = _rotate(random_1, 1, random_2) + + +def benchmark_batchsize(funcs: list, batches: list[int], num_iterations: int = 1000): + ret: dict = {} + for batch_size in batches: + random_a = numpy.random.random((3, 3, batch_size)) + random_b = numpy.random.random((3, batch_size)) + + ret[batch_size] = {} + for func_name, func, func_arg_wrap in funcs: + tot = 0.0 + for _ in range(num_iterations): + args = func_arg_wrap(random_a, random_b) + start = time.perf_counter() + func(*args) + tot += time.perf_counter() - start + + ret[batch_size][func_name] = tot / num_iterations + + return ret + + +def _pyelastica_arg_wrap(x, y): + return x, 1.0, y + + +def _elasticapp_arg_wrap(x, y): + return Tensor(x), Matrix(y) + + +results = benchmark_batchsize( + [ + ("pyelastica", _rotate, _pyelastica_arg_wrap), + ("elasticapp_simd", rotate, _elasticapp_arg_wrap), + ("elasticapp_scalar", rotate_scalar, _elasticapp_arg_wrap), + ], + [2**i for i in range(14)], +) +for size, data in results.items(): + pyelastica = data["pyelastica"] + elasticapp_simd = data["elasticapp_simd"] + elasticapp_scalar = data["elasticapp_scalar"] + print(f"{size = }") + print(f"{pyelastica = }") + print(f"{elasticapp_simd = }, ratio: {elasticapp_simd / pyelastica}") + print(f"{elasticapp_scalar = }, ratio: {elasticapp_scalar / pyelastica}") + print() diff --git a/backend/src/_rotations.cpp b/backend/src/_rotations.cpp index d12f6bb2..451d5fed 100644 --- a/backend/src/_rotations.cpp +++ b/backend/src/_rotations.cpp @@ -1,11 +1,14 @@ #include #include +#include "Systems/States/Expressions/backends/blaze/SO3PrimitiveAddAssign/Scalar.hpp" +#include "Systems/States/Expressions/backends/blaze/SO3PrimitiveAddAssign/SIMD.hpp" + #include "Systems/CosseratRods/Traits/Operations/BlazeBackend/InvRotateDivide/Scalar.hpp" #include "Systems/CosseratRods/Traits/Operations/BlazeBackend/InvRotateDivide/SIMD.hpp" -#include "Systems/CosseratRods/Traits/Operations/BlazeBackend/InvRotateDivide/Blaze.hpp" namespace detail = elastica::cosserat_rod::detail; +namespace states = elastica::states; using ElasticaVector = ::blaze::DynamicVector>; @@ -13,6 +16,12 @@ using ElasticaMatrix = ::blaze::DynamicMatrix>; using ElasticaTensor = ::blaze::DynamicTensor; +template +void rotate(ElasticaTensor &director_collection, ElasticaMatrix &axis_collection) +{ + states::detail::SO3AddAssign::apply(director_collection, axis_collection); +} + template ElasticaMatrix inv_rotate_with_span(ElasticaTensor &director_collection, ElasticaVector &span_vector) { @@ -40,9 +49,19 @@ PYBIND11_MODULE(_rotations, m) .. autosummary:: :toctree: _generate + rotate + rotate_scalar inv_rotate + inv_rotate_scalar )pbdoc"; + m.def("rotate", &rotate, R"pbdoc( + Perform the rotate operation (SIMD Variant). + )pbdoc"); + m.def("rotate_scalar", &rotate, R"pbdoc( + Perform the rotate operation (Scalar Variant). + )pbdoc"); + m.def("inv_rotate", &inv_rotate, R"pbdoc( Perform the inverse-rotate operation (SIMD Variant). )pbdoc"); diff --git a/backend/tests/test_rotations.py b/backend/tests/test_rotations.py index e9646f03..07882d94 100644 --- a/backend/tests/test_rotations.py +++ b/backend/tests/test_rotations.py @@ -6,8 +6,49 @@ import numpy as np import pytest from numpy.testing import assert_allclose -from elasticapp._PyArrays import Tensor -from elasticapp._rotations import inv_rotate, inv_rotate_scalar +from elasticapp._PyArrays import Tensor, Matrix +from elasticapp._rotations import rotate, rotate_scalar, inv_rotate, inv_rotate_scalar + + +@pytest.mark.parametrize("rotate_func", [rotate, rotate_scalar]) +def test_rotate_correctness(rotate_func): + blocksize = 16 + + def get_aligned_director_collection(theta_collection): + sins = np.sin(theta_collection) + coss = np.cos(theta_collection) + # Get basic director out, then modify it as you like + dir = np.tile(np.eye(3).reshape(3, 3, 1), blocksize) + dir[0, 0, ...] = coss + # Flip signs on [0,1] and [1,0] to go from our row-wise + # representation to the more commonly used + # columnwise representation, for similar reasons metioned + # before + dir[0, 1, ...] = sins + dir[1, 0, ...] = -sins + dir[1, 1, ...] = coss + + return dir + + base_angle = np.deg2rad(np.linspace(0.0, 90.0, blocksize)) + rotated_by = np.deg2rad(15.0) + 0.0 * base_angle + rotated_about = np.array([0.0, 0.0, 1.0]).reshape(-1, 1) + + director_collection = Tensor(get_aligned_director_collection(base_angle)) + axis_collection = np.tile(rotated_about, blocksize) + axis_collection *= rotated_by + + rotate_func(director_collection, Matrix(axis_collection)) + test_rotated_director_collection = np.asarray(director_collection) + correct_rotation = rotated_by + 1.0 * base_angle + correct_rotated_director_collection = get_aligned_director_collection( + correct_rotation + ) + + assert test_rotated_director_collection.shape == (3, 3, blocksize) + assert_allclose( + test_rotated_director_collection, correct_rotated_director_collection + ) @pytest.mark.parametrize("inv_rotate_func", [inv_rotate, inv_rotate_scalar])