From 473cffef8375812e31b3a65056c0332c0e75d82d Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 3 Jan 2025 23:42:25 -0800 Subject: [PATCH] Add: Small Matrix (6x6, F) Add the implementation for the new small matrix (and small vector) types and their operators. Specialize for now for the most common type we need in accelerator physics, a 6x6, F-ordered, 1-indexed matrix. --- docs/source/usage/api.rst | 7 + src/Base/CMakeLists.txt | 1 + src/Base/SmallMatrix.H | 307 ++++++++++++++++++++++++++++ src/Base/SmallMatrix.cpp | 42 ++++ src/amrex/extensions/SmallMatrix.py | 142 +++++++++++++ src/amrex/space1d/__init__.py | 2 + src/amrex/space2d/__init__.py | 2 + src/amrex/space3d/__init__.py | 2 + src/pyAMReX.cpp | 15 +- 9 files changed, 515 insertions(+), 5 deletions(-) create mode 100644 src/Base/SmallMatrix.H create mode 100644 src/Base/SmallMatrix.cpp create mode 100644 src/amrex/extensions/SmallMatrix.py diff --git a/docs/source/usage/api.rst b/docs/source/usage/api.rst index 0a12511c..ce537ea0 100644 --- a/docs/source/usage/api.rst +++ b/docs/source/usage/api.rst @@ -174,6 +174,13 @@ Data Containers :members: :undoc-members: +Small Matrices and Vectors +"""""""""""""""""""""""""" + +.. autoclass:: amrex.space3d.SmallMatrix_6x6_F_SI1_double + :members: + :undoc-members: + Utility """"""" diff --git a/src/Base/CMakeLists.txt b/src/Base/CMakeLists.txt index 0db7d7cb..e4d774e3 100644 --- a/src/Base/CMakeLists.txt +++ b/src/Base/CMakeLists.txt @@ -25,6 +25,7 @@ foreach(D IN LISTS AMReX_SPACEDIM) IndexType.cpp IntVect.cpp RealVect.cpp + SmallMatrix.cpp MultiFab.cpp ParallelDescriptor.cpp ParmParse.cpp diff --git a/src/Base/SmallMatrix.H b/src/Base/SmallMatrix.H new file mode 100644 index 00000000..bc1973e0 --- /dev/null +++ b/src/Base/SmallMatrix.H @@ -0,0 +1,307 @@ +/* Copyright 2021-2025 The AMReX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#pragma once + +#include "pyAMReX.H" + +#include + +#include +#include +#include +#include +#include +#include + + +namespace +{ + // helper type traits + template + struct get_value_type { using value_type = T; }; + template + struct get_value_type> { using value_type = T; }; + template + using get_value_type_t = typename get_value_type::value_type; + + // helper to check if Array4 is of constant value type T + template + constexpr bool is_not_const () + { + return std::is_same_v< + std::remove_cv_t< + T + >, + T + > && + std::is_same_v< + std::remove_cv_t< + get_value_type_t + >, + get_value_type_t + >; + } +} + +namespace pyAMReX +{ + using namespace amrex; + + /** CPU: __array_interface__ v3 + * + * https://numpy.org/doc/stable/reference/arrays.interface.html + */ + template< + class T, + int NRows, + int NCols, + Order ORDER = Order::F, + int StartIndex = 0 + > + py::dict + array_interface (SmallMatrix const & m) + { + auto d = py::dict(); + // provide C index order for shape and strides + auto shape = m.ordering == Order::F ? py::make_tuple( + py::ssize_t(NRows), + py::ssize_t(NCols) // fastest varying index + ) : py::make_tuple( + py::ssize_t(NCols), + py::ssize_t(NRows) // fastest varying index + ); + // buffer protocol strides are in bytes + auto const strides = m.ordering == Order::F ? py::make_tuple( + py::ssize_t(sizeof(T) * NCols), + py::ssize_t(sizeof(T)) // fastest varying index + ) : py::make_tuple( + py::ssize_t(sizeof(T) * NRows), + py::ssize_t(sizeof(T)) // fastest varying index + ); + bool const read_only = false; // note: we could decide on is_not_const, + // but many libs, e.g. PyTorch, do not + // support read-only and will raise + // warnings, casting to read-write + d["data"] = py::make_tuple(std::intptr_t(&m.template get<0>()), read_only); + // note: if we want to keep the same global indexing with non-zero + // box small_end as in AMReX, then we can explore playing with + // this offset as well + //d["offset"] = 0; // default + //d["mask"] = py::none(); // default + + d["shape"] = shape; + // we could also set this after checking the strides are C-style contiguous: + //if (is_contiguous(shape, strides)) + // d["strides"] = py::none(); // C-style contiguous + //else + d["strides"] = strides; + + // type description + // for more complicated types, e.g., tuples/structs + //d["descr"] = ...; + // we currently only need this + using T_no_cv = std::remove_cv_t; + d["typestr"] = py::format_descriptor::format(); + + d["version"] = 3; + return d; + } + + template< + class T, + int NRows, + int NCols, + Order ORDER = Order::F, + int StartIndex = 0 + > + void make_SmallMatrix(py::module &m, std::string typestr) + { + using namespace amrex; + + using T_no_cv = std::remove_cv_t; + + // dispatch simpler via: py::format_descriptor::format() naming + // but note the _const suffix that might be needed + auto const sm_name = std::string("SmallMatrix_") + .append(std::to_string(NRows)).append("x").append(std::to_string(NCols)) + .append("_").append(ORDER == Order::F ? "F" : "C") + .append("_SI").append(std::to_string(StartIndex)) + .append("_").append(typestr); + using SM = SmallMatrix; + py::class_< SM > py_sm(m, sm_name.c_str()); + py_sm + .def("__repr__", + [sm_name](SM const &) { + return ""; + } + ) + .def("__str__", + [sm_name](SM const & sm) { + std::stringstream ss; + ss << sm; + return ss.str(); + } + ) + + .def_property_readonly("size", [](SM const &){ return SM::row_size * SM::column_size; }) + .def_property_readonly("row_size", [](SM const &){ return SM::row_size; }) + .def_property_readonly("column_size", [](SM const &){ return SM::column_size; }) + .def_property_readonly("order", [](SM const &){ return SM::ordering == Order::F ? "F" : "C"; }) // NumPy name + .def_property_readonly("starting_index", [](SM const &){ return SM::starting_index; }) + + .def_static("zero", [](){ return SM::Zero(); }) + + .def(py::init([](){ return SM{}; })) // zero-init + .def(py::init()) // copy-init + + /* init from a numpy or other buffer protocol array: copy + */ + .def(py::init([](py::array_t & arr) + { + py::buffer_info buf = arr.request(); + + constexpr bool is_vector = SM::column_size == 1 || SM::row_size == 1; + constexpr int sm_dim = is_vector ? 1 : 2; + if (buf.ndim != sm_dim) + throw std::runtime_error("The SmallMatrix to create is " + std::to_string(sm_dim) + + "D, but the passed array is " + std::to_string(buf.ndim) + "D."); + if (buf.size != SM::column_size * SM::row_size) + throw std::runtime_error("Array size mismatch: Expected " + std::to_string(SM::column_size * SM::row_size) + + " elements, but passed " + std::to_string(buf.size) + " elements."); + + if (buf.format != py::format_descriptor::format()) + throw std::runtime_error("Incompatible format: expected '" + + py::format_descriptor::format() + + "' and received '" + buf.format + "'!"); + + // TODO: check that strides are either exact or None in buf + // TODO: transpose if SM order is not C? + + auto sm = std::make_unique< SM >(); + auto * src = static_cast(buf.ptr); + std::copy(src, src + buf.size, &sm->template get<0>()); + + // todo: we could check and store here if the array buffer we got is read-only + + return sm; + })) + + /* init from __cuda_array_interface__: device-to-host copy + * TODO + */ + + + // CPU: __array_interface__ v3 + // https://numpy.org/doc/stable/reference/arrays.interface.html + .def_property_readonly("__array_interface__", [](SM const & sm) { + return pyAMReX::array_interface(sm); + }) + + // CPU: __array_function__ interface (TODO) + // + // NEP 18 — A dispatch mechanism for NumPy's high level array functions. + // https://numpy.org/neps/nep-0018-array-function-protocol.html + // This enables code using NumPy to be directly operated on Array4 arrays. + // __array_function__ feature requires NumPy 1.16 or later. + + + // Nvidia GPUs: __cuda_array_interface__ v3 + // https://numba.readthedocs.io/en/latest/cuda/cuda_array_interface.html + .def_property_readonly("__cuda_array_interface__", [](SM const & sm) { + auto d = pyAMReX::array_interface(sm); + + // data: + // Because the user of the interface may or may not be in the same context, the most common case is to use cuPointerGetAttribute with CU_POINTER_ATTRIBUTE_DEVICE_POINTER in the CUDA driver API (or the equivalent CUDA Runtime API) to retrieve a device pointer that is usable in the currently active context. + // TODO For zero-size arrays, use 0 here. + + // None or integer + // An optional stream upon which synchronization must take place at the point of consumption, either by synchronizing on the stream or enqueuing operations on the data on the given stream. Integer values in this entry are as follows: + // 0: This is disallowed as it would be ambiguous between None and the default stream, and also between the legacy and per-thread default streams. Any use case where 0 might be given should either use None, 1, or 2 instead for clarity. + // 1: The legacy default stream. + // 2: The per-thread default stream. + // Any other integer: a cudaStream_t represented as a Python integer. + // When None, no synchronization is required. + d["stream"] = py::none(); + + d["version"] = 3; + return d; + }) + + + // TODO: __dlpack__ __dlpack_device__ + // DLPack protocol (CPU, NVIDIA GPU, AMD GPU, Intel GPU, etc.) + // https://dmlc.github.io/dlpack/latest/ + // https://data-apis.org/array-api/latest/design_topics/data_interchange.html + // https://github.com/data-apis/consortium-feedback/issues/1 + // https://github.com/dmlc/dlpack/blob/master/include/dlpack/dlpack.h + // https://docs.cupy.dev/en/stable/user_guide/interoperability.html#dlpack-data-exchange-protocol + + .def("dot", &SM::dot) + .def("prod", &SM::product) // NumPy name + .def("set_val", &SM::setVal) + .def("sum", &SM::sum) + .def_property_readonly("T", &SM::transpose) // NumPy name + + // operators + .def(py::self + py::self) + .def(py::self - py::self) + .def(py::self * amrex::Real()) + .def(amrex::Real() * py::self) + .def(-py::self) + + // getter + .def("__getitem__", [](SM & sm, std::array const & key){ + return sm(key[0], key[1]); + }) + ; + // setter + if constexpr (is_not_const()) + { + py_sm + .def("__setitem__", [](SM & sm, std::array const & key, T const value){ + sm(key[0], key[1]) = value; + }) + ; + } + + // square matrix + if constexpr (NRows == NCols) + { + py_sm + .def_static("identity", [](){ return SM::Identity(); }) + .def("trace", &SM::trace) + .def("transpose_in_place", &SM::transposeInPlace) + ; + } + + // vector + if constexpr (NRows == 1 || NCols == 1) + { + py_sm + .def("__getitem__", [](SM & sm, int key){ + if (key < SM::starting_index || key >= SM::column_size * SM::row_size + SM::starting_index) + throw std::runtime_error("Index out of bounds: " + std::to_string(key)); + return sm(key); + }) + .def("__setitem__", [](SM & sm, int key, T const value){ + if (key < SM::starting_index || key >= SM::column_size * SM::row_size + SM::starting_index) + throw std::runtime_error("Index out of bounds: " + std::to_string(key)); + sm(key) = value; + }) + ; + } else { + using SV = SmallMatrix; + using SRV = SmallMatrix; + + // operators for matrix-matrix & matrix-vector + py_sm + .def(py::self * py::self) + .def(py::self * SV()) + .def(SRV() * py::self) + ; + } + } +} diff --git a/src/Base/SmallMatrix.cpp b/src/Base/SmallMatrix.cpp new file mode 100644 index 00000000..cb8c886e --- /dev/null +++ b/src/Base/SmallMatrix.cpp @@ -0,0 +1,42 @@ +/* Copyright 2021-2025 The AMReX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyAMReX.H" + +#include "SmallMatrix.H" + + +void init_SmallMatrix (py::module &m) +{ + using namespace pyAMReX; + + // 6x6 Matrix as commonly used in accelerator physics + { + constexpr int NRows = 6; + constexpr int NCols = 6; + constexpr Order ORDER = Order::F; + constexpr int StartIndex = 1; + + // Vector + make_SmallMatrix< float, NRows, 1, ORDER, StartIndex >(m, "float"); + make_SmallMatrix< double, NRows, 1, ORDER, StartIndex >(m, "double"); + make_SmallMatrix< long double, NRows, 1, ORDER, StartIndex >(m, "longdouble"); + + // RowVector + make_SmallMatrix< float, 1, NCols, ORDER, StartIndex >(m, "float"); + make_SmallMatrix< double, 1, NCols, ORDER, StartIndex >(m, "double"); + make_SmallMatrix< long double, 1, NCols, ORDER, StartIndex >(m, "longdouble"); + + // Matrix + make_SmallMatrix< float, NRows, NCols, ORDER, StartIndex >(m, "float"); + make_SmallMatrix< double, NRows, NCols, ORDER, StartIndex >(m, "double"); + make_SmallMatrix< long double, NRows, NCols, ORDER, StartIndex >(m, "longdouble"); + /* + make_SmallMatrix< float const, NRows, NCols, ORDER, StartIndex >(m, "float_const"); + make_SmallMatrix< double const, NRows, NCols, ORDER, StartIndex >(m, "double_const"); + make_SmallMatrix< long double const, NRows, NCols, ORDER, StartIndex >(m, "longdouble_const"); + */ + } +} diff --git a/src/amrex/extensions/SmallMatrix.py b/src/amrex/extensions/SmallMatrix.py new file mode 100644 index 00000000..8cedb93f --- /dev/null +++ b/src/amrex/extensions/SmallMatrix.py @@ -0,0 +1,142 @@ +""" +This file is part of pyAMReX + +Copyright 2025 AMReX community +Authors: Axel Huebl +License: BSD-3-Clause-LBNL +""" + + +def smallmatrix_to_numpy(self, copy=False, order="F"): + """ + Provide a NumPy view into an SmallMatrix. + + Note on the order of indices: + By default, this is as in AMReX in Fortran contiguous order, indexing as + x,y,z. This has performance implications for use in external libraries such + as cupy. + The order="C" option will index as z,y,x and perform better with cupy. + https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074 + + Parameters + ---------- + self : amrex.SmallMatrix_* + A SmallMatrix class in pyAMReX + copy : bool, optional + Copy the data if true, otherwise create a view (default). + order : string, optional + F order (default) or C. C is faster with external libraries. + + Returns + ------- + np.array + A NumPy 2-dimensional array. + """ + import numpy as np + + if copy: + data = np.array(self, copy=True) + else: + data = np.array(self, copy=False) + + # TODO: Check self.order == "F" ? + if order == "F": + return data.T + elif order == "C": + return data + else: + raise ValueError("The order argument must be F or C.") + + +def smallmatrix_to_cupy(self, copy=False, order="F"): + """ + Provide a CuPy view into an SmallMatrix. + + Note on the order of indices: + By default, this is as in AMReX in Fortran contiguous order, indexing as + x,y,z. This has performance implications for use in external libraries such + as cupy. + The order="C" option will index as z,y,x and perform better with cupy. + https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074 + + Parameters + ---------- + self : amrex.SmallMatrix_* + A SmallMatrix class in pyAMReX + copy : bool, optional + Copy the data if true, otherwise create a view (default). + order : string, optional + F order (default) or C. C is faster with external libraries. + + Returns + ------- + cupy.array + A cupy 2-dimensional array. + + Raises + ------ + ImportError + Raises an exception if cupy is not installed + """ + import cupy as cp + + # TODO: Check self.order == "F" ? + if order == "F": + return cp.array(self, copy=copy).T + elif order == "C": + return cp.array(self, copy=copy) + else: + raise ValueError("The order argument must be F or C.") + + +def smallmatrix_to_xp(self, copy=False, order="F"): + """ + Provide a NumPy or CuPy view into a SmallMatrix, depending on amr.Config.have_gpu . + + This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code: + https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code + + Note on the order of indices: + By default, this is as in AMReX in Fortran contiguous order, indexing as + x,y,z. This has performance implications for use in external libraries such + as cupy. + The order="C" option will index as z,y,x and perform better with cupy. + https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074 + + Parameters + ---------- + self : amrex.SmallMatrix_* + A SmallMatrix class in pyAMReX + copy : bool, optional + Copy the data if true, otherwise create a view (default). + order : string, optional + F order (default) or C. C is faster with external libraries. + + Returns + ------- + xp.array + A NumPy or CuPy 2-dimensional array. + """ + import inspect + + amr = inspect.getmodule(self) + return ( + self.to_cupy(copy, order) if amr.Config.have_gpu else self.to_numpy(copy, order) + ) + + +def register_SmallMatrix_extension(amr): + """SmallMatrix helper methods""" + import inspect + import sys + + # register member functions for every Array4_* type + for _, SmallMatrix_type in inspect.getmembers( + sys.modules[amr.__name__], + lambda member: inspect.isclass(member) + and member.__module__ == amr.__name__ + and member.__name__.startswith("SmallMatrix_"), + ): + SmallMatrix_type.to_numpy = smallmatrix_to_numpy + SmallMatrix_type.to_cupy = smallmatrix_to_cupy + SmallMatrix_type.to_xp = smallmatrix_to_xp diff --git a/src/amrex/space1d/__init__.py b/src/amrex/space1d/__init__.py index 8af62c81..2fcbde32 100644 --- a/src/amrex/space1d/__init__.py +++ b/src/amrex/space1d/__init__.py @@ -50,11 +50,13 @@ def Print(*args, **kwargs): from ..extensions.MultiFab import register_MultiFab_extension from ..extensions.ParticleContainer import register_ParticleContainer_extension from ..extensions.PODVector import register_PODVector_extension +from ..extensions.SmallMatrix import register_SmallMatrix_extension from ..extensions.StructOfArrays import register_SoA_extension register_Array4_extension(amrex_1d_pybind) register_MultiFab_extension(amrex_1d_pybind) register_PODVector_extension(amrex_1d_pybind) +register_SmallMatrix_extension(amrex_1d_pybind) register_SoA_extension(amrex_1d_pybind) register_AoS_extension(amrex_1d_pybind) register_ParticleContainer_extension(amrex_1d_pybind) diff --git a/src/amrex/space2d/__init__.py b/src/amrex/space2d/__init__.py index 10e089b5..39a0ad59 100644 --- a/src/amrex/space2d/__init__.py +++ b/src/amrex/space2d/__init__.py @@ -50,11 +50,13 @@ def Print(*args, **kwargs): from ..extensions.MultiFab import register_MultiFab_extension from ..extensions.ParticleContainer import register_ParticleContainer_extension from ..extensions.PODVector import register_PODVector_extension +from ..extensions.SmallMatrix import register_SmallMatrix_extension from ..extensions.StructOfArrays import register_SoA_extension register_Array4_extension(amrex_2d_pybind) register_MultiFab_extension(amrex_2d_pybind) register_PODVector_extension(amrex_2d_pybind) +register_SmallMatrix_extension(amrex_2d_pybind) register_SoA_extension(amrex_2d_pybind) register_AoS_extension(amrex_2d_pybind) register_ParticleContainer_extension(amrex_2d_pybind) diff --git a/src/amrex/space3d/__init__.py b/src/amrex/space3d/__init__.py index 5c4254a4..f1cef356 100644 --- a/src/amrex/space3d/__init__.py +++ b/src/amrex/space3d/__init__.py @@ -50,11 +50,13 @@ def Print(*args, **kwargs): from ..extensions.MultiFab import register_MultiFab_extension from ..extensions.ParticleContainer import register_ParticleContainer_extension from ..extensions.PODVector import register_PODVector_extension +from ..extensions.SmallMatrix import register_SmallMatrix_extension from ..extensions.StructOfArrays import register_SoA_extension register_Array4_extension(amrex_3d_pybind) register_MultiFab_extension(amrex_3d_pybind) register_PODVector_extension(amrex_3d_pybind) +register_SmallMatrix_extension(amrex_3d_pybind) register_SoA_extension(amrex_3d_pybind) register_AoS_extension(amrex_3d_pybind) register_ParticleContainer_extension(amrex_3d_pybind) diff --git a/src/pyAMReX.cpp b/src/pyAMReX.cpp index ed52c2be..226a4989 100644 --- a/src/pyAMReX.cpp +++ b/src/pyAMReX.cpp @@ -14,6 +14,7 @@ // forward declarations of exposed classes void init_Algorithm(py::module&); void init_AMReX(py::module&); +void init_AmrMesh(py::module &); void init_Arena(py::module&); void init_Array4(py::module&); void init_BaseFab(py::module&); @@ -27,8 +28,9 @@ void init_FArrayBox(py::module&); void init_Geometry(py::module&); void init_IndexType(py::module &); void init_IntVect(py::module &); -void init_RealVect(py::module &); -void init_AmrMesh(py::module &); +#ifdef AMREX_USE_MPI +void init_MPMD(py::module &); +#endif void init_MultiFab(py::module &); void init_ParallelDescriptor(py::module &); void init_ParmParse(py::module &); @@ -36,13 +38,12 @@ void init_ParticleContainer(py::module &); void init_Periodicity(py::module &); void init_PlotFileUtil(py::module &); void init_PODVector(py::module &); +void init_RealVect(py::module &); +void init_SmallMatrix(py::module &); void init_Utility(py::module &); void init_Vector(py::module &); void init_Version(py::module &); void init_VisMF(py::module &); -#ifdef AMREX_USE_MPI -void init_MPMD(py::module &); -#endif #if AMREX_SPACEDIM == 1 PYBIND11_MODULE(amrex_1d_pybind, m) { @@ -81,6 +82,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) { Periodicity PlotFileUtil PODVector + SmallMatrix StructOfArrays Utility Vector @@ -88,6 +90,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) { )pbdoc"; // note: order from parent to child classes and argument usage + init_AMReX(m); init_Arena(m); init_Dim3(m); @@ -98,6 +101,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) { init_Box(m); init_Periodicity(m); init_Array4(m); + init_SmallMatrix(m); init_BoxArray(m); init_ParmParse(m); init_CoordSys(m); @@ -117,6 +121,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) { #ifdef AMREX_USE_MPI init_MPMD(m); #endif + // Wrappers around standalone functions init_PlotFileUtil(m); init_Utility(m);