diff --git a/symforce/codegen/__init__.py b/symforce/codegen/__init__.py index 526e7f6c7..3c5b22286 100644 --- a/symforce/codegen/__init__.py +++ b/symforce/codegen/__init__.py @@ -7,5 +7,8 @@ Package for executable code generation from symbolic expressions. """ -from .codegen_config import CodegenConfig, CppConfig, PythonConfig from .codegen import Codegen, LinearizationMode, GeneratedPaths +from .codegen_config import CodegenConfig + +from .backends.cpp.cpp_config import CppConfig +from .backends.python.python_config import PythonConfig diff --git a/symforce/codegen/backends/README.md b/symforce/codegen/backends/README.md new file mode 100644 index 000000000..7efa4bbf1 --- /dev/null +++ b/symforce/codegen/backends/README.md @@ -0,0 +1,14 @@ +# Code Generation Backends + +SymForce takes symbolic functions and generates runtime functions for multiple target backends. It is quite straightforward to add new backends. Before you do this, you should be familiar with [SymPy printing](https://docs.sympy.org/latest/modules/printing.html) for getting code from symbolic expressions, and with [Jinja templating](https://realpython.com/primer-on-jinja-templating/) for rendering output files. + +The minimal steps to support a new backend are: + + 1. Choose a name for your backend (for example 'julia') and create a corresponding package in `symforce/codegen/backends`. + 2. Implement a subclass of `sympy.CodePrinter` that emits backend math code while traversing symbolic expressions. Sometimes SymPy already contains the backend and the best pattern is to inherit from it and customize as needed. The best way to do this is by looking at existing backends as examples. + 3. Implement a subclass of `symforce.codegen.codegen_config.CodegenConfig`. This is the spec that users pass to the `Codegen` object to use your backend. Again, see existing examples. Optionally import your config in `symforce/codegen/__init__.py`. + 4. Create a `templates` directory containing jinja templates that are used to generate the actual output files. They specify the high level structure and APIs around the math code. Your codegen config has a `templates_to_render` method that should match your templates. A typical start is just one function template. + 5. Add your backend's extensions to `FileType` in `symforce/codegen/template_util.py`, filling out relevant methods there. + 6. Add tests to `test/symforce_codegen_test.py`. + +This will result in being able to generate functions for your backend that deal with scalars and arrays, but the `sym` geometry and camera classes. To implement those, follow the examples in `geo_package_codegen` and `cam_package_codegen`. diff --git a/symforce/codegen/python_templates/ops/__init__.py.jinja b/symforce/codegen/backends/__init__.py similarity index 100% rename from symforce/codegen/python_templates/ops/__init__.py.jinja rename to symforce/codegen/backends/__init__.py diff --git a/symforce/codegen/backends/cpp/README.md b/symforce/codegen/backends/cpp/README.md new file mode 100644 index 000000000..e69de29bb diff --git a/symforce/codegen/backends/cpp/__init__.py b/symforce/codegen/backends/cpp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/symforce/codegen/printers/cpp_code_printer.py b/symforce/codegen/backends/cpp/cpp_code_printer.py similarity index 98% rename from symforce/codegen/printers/cpp_code_printer.py rename to symforce/codegen/backends/cpp/cpp_code_printer.py index 00c453172..035722a8e 100644 --- a/symforce/codegen/printers/cpp_code_printer.py +++ b/symforce/codegen/backends/cpp/cpp_code_printer.py @@ -3,14 +3,12 @@ # This source code is under the Apache 2.0 license found in the LICENSE file. # ---------------------------------------------------------------------------- +import sympy from sympy.printing.c import get_math_macros from sympy.printing.cxx import CXX11CodePrinter from symforce import typing as T -# Everything in this file is SymPy, not SymEngine (even when SymForce is on the SymEngine backend) -import sympy - class CppCodePrinter(CXX11CodePrinter): """ diff --git a/symforce/codegen/backends/cpp/cpp_config.py b/symforce/codegen/backends/cpp/cpp_config.py new file mode 100644 index 000000000..bf6d5f993 --- /dev/null +++ b/symforce/codegen/backends/cpp/cpp_config.py @@ -0,0 +1,74 @@ +# ---------------------------------------------------------------------------- +# SymForce - Copyright 2022, Skydio, Inc. +# This source code is under the Apache 2.0 license found in the LICENSE file. +# ---------------------------------------------------------------------------- +from dataclasses import dataclass +from pathlib import Path +from sympy.printing.codeprinter import CodePrinter + +from symforce import typing as T +from symforce.codegen.codegen_config import CodegenConfig + +CURRENT_DIR = Path(__file__).parent + + +@dataclass +class CppConfig(CodegenConfig): + """ + Code generation config for the C++ backend. + + Args: + doc_comment_line_prefix: Prefix applied to each line in a docstring + line_length: Maximum allowed line length in docstrings; used for formatting docstrings. + use_eigen_types: Use eigen_lcm types for vectors instead of lists + autoformat: Run a code formatter on the generated code + cse_optimizations: Optimizations argument to pass to sm.cse + support_complex: Generate code that can work with std::complex or with regular float types + force_no_inline: Mark generated functions as `__attribute__((noinline))` + zero_initialization_sparsity_threshold: Threshold between 0 and 1 for the sparsity below + which we'll initialize an output matrix to 0, so we + don't have to generate a line to set each zero + element to 0 individually + explicit_template_instantiation_types: Explicity instantiates templated functions in a `.cc` + file for each given type. This allows the generated function to be compiled in its own + translation unit. Useful for large functions which take a long time to compile. + """ + + doc_comment_line_prefix: str = " * " + line_length: int = 100 + use_eigen_types: bool = True + support_complex: bool = False + force_no_inline: bool = False + zero_initialization_sparsity_threshold: float = 0.5 + explicit_template_instantiation_types: T.Optional[T.Sequence[str]] = None + + @classmethod + def backend_name(cls) -> str: + return "cpp" + + @classmethod + def template_dir(cls) -> Path: + return CURRENT_DIR / "templates" + + def templates_to_render(self, generated_file_name: str) -> T.List[T.Tuple[str, str]]: + # Generate code into a header (since the code is templated) + templates = [("function/FUNCTION.h.jinja", f"{generated_file_name}.h")] + + # Generate a cc file only if we need explicit instantiation. + if self.explicit_template_instantiation_types is not None: + templates.append(("function/FUNCTION.cc.jinja", f"{generated_file_name}.cc")) + + return templates + + def printer(self) -> CodePrinter: + # NOTE(hayk): Is there any benefit to this being lazy? + from symforce.codegen.backends.cpp import cpp_code_printer + + if self.support_complex: + return cpp_code_printer.ComplexCppCodePrinter() + else: + return cpp_code_printer.CppCodePrinter() + + @staticmethod + def format_data_accessor(prefix: str, index: int) -> str: + return f"{prefix}.Data()[{index}]" diff --git a/symforce/codegen/cpp_templates/cam_package/CLASS.cc.jinja b/symforce/codegen/backends/cpp/templates/cam_package/CLASS.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/cam_package/CLASS.cc.jinja rename to symforce/codegen/backends/cpp/templates/cam_package/CLASS.cc.jinja diff --git a/symforce/codegen/cpp_templates/cam_package/CLASS.h.jinja b/symforce/codegen/backends/cpp/templates/cam_package/CLASS.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/cam_package/CLASS.h.jinja rename to symforce/codegen/backends/cpp/templates/cam_package/CLASS.h.jinja diff --git a/symforce/codegen/cpp_templates/cam_package/camera.h.jinja b/symforce/codegen/backends/cpp/templates/cam_package/camera.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/cam_package/camera.h.jinja rename to symforce/codegen/backends/cpp/templates/cam_package/camera.h.jinja diff --git a/symforce/codegen/cpp_templates/cam_package/posed_camera.h.jinja b/symforce/codegen/backends/cpp/templates/cam_package/posed_camera.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/cam_package/posed_camera.h.jinja rename to symforce/codegen/backends/cpp/templates/cam_package/posed_camera.h.jinja diff --git a/symforce/codegen/cpp_templates/epsilon.cc.jinja b/symforce/codegen/backends/cpp/templates/epsilon.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/epsilon.cc.jinja rename to symforce/codegen/backends/cpp/templates/epsilon.cc.jinja diff --git a/symforce/codegen/cpp_templates/epsilon.h.jinja b/symforce/codegen/backends/cpp/templates/epsilon.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/epsilon.h.jinja rename to symforce/codegen/backends/cpp/templates/epsilon.h.jinja diff --git a/symforce/codegen/cpp_templates/function/FUNCTION.cc.jinja b/symforce/codegen/backends/cpp/templates/function/FUNCTION.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/function/FUNCTION.cc.jinja rename to symforce/codegen/backends/cpp/templates/function/FUNCTION.cc.jinja diff --git a/symforce/codegen/cpp_templates/function/FUNCTION.h.jinja b/symforce/codegen/backends/cpp/templates/function/FUNCTION.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/function/FUNCTION.h.jinja rename to symforce/codegen/backends/cpp/templates/function/FUNCTION.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/CLASS.cc.jinja b/symforce/codegen/backends/cpp/templates/geo_package/CLASS.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/CLASS.cc.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/CLASS.cc.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/CLASS.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/CLASS.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/CLASS.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/CLASS.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/custom_methods/pose2.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/custom_methods/pose2.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/custom_methods/pose2.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/custom_methods/pose2.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/custom_methods/pose3.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/custom_methods/pose3.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/custom_methods/pose3.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/custom_methods/pose3.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/custom_methods/rot2.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/custom_methods/rot2.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/custom_methods/rot2.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/custom_methods/rot2.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/custom_methods/rot3.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/custom_methods/rot3.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/custom_methods/rot3.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/custom_methods/rot3.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/lie_group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/lie_group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/lie_group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/lie_group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/matrix/group_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/group_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/matrix/group_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/group_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/matrix/group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/matrix/group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/matrix/lie_group_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/lie_group_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/matrix/lie_group_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/lie_group_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/matrix/lie_group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/lie_group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/matrix/lie_group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/lie_group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/matrix/storage_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/storage_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/matrix/storage_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/storage_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/matrix/storage_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/storage_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/matrix/storage_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/matrix/storage_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/scalar/group_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/group_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/scalar/group_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/group_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/scalar/group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/scalar/group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/scalar/lie_group_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/lie_group_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/scalar/lie_group_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/lie_group_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/scalar/lie_group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/lie_group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/scalar/lie_group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/lie_group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/scalar/storage_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/storage_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/scalar/storage_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/storage_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/scalar/storage_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/storage_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/scalar/storage_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/scalar/storage_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/geo_package/ops/storage_ops.h.jinja b/symforce/codegen/backends/cpp/templates/geo_package/ops/storage_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/geo_package/ops/storage_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/geo_package/ops/storage_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/keys.h.jinja b/symforce/codegen/backends/cpp/templates/keys.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/keys.h.jinja rename to symforce/codegen/backends/cpp/templates/keys.h.jinja diff --git a/symforce/codegen/cpp_templates/ops/CLASS/group_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/ops/CLASS/group_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/ops/CLASS/group_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/ops/CLASS/group_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/ops/CLASS/group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/ops/CLASS/group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/ops/CLASS/group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/ops/CLASS/group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/ops/CLASS/lie_group_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/ops/CLASS/lie_group_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/ops/CLASS/lie_group_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/ops/CLASS/lie_group_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/ops/CLASS/lie_group_ops.h.jinja b/symforce/codegen/backends/cpp/templates/ops/CLASS/lie_group_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/ops/CLASS/lie_group_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/ops/CLASS/lie_group_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/ops/CLASS/storage_ops.cc.jinja b/symforce/codegen/backends/cpp/templates/ops/CLASS/storage_ops.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/ops/CLASS/storage_ops.cc.jinja rename to symforce/codegen/backends/cpp/templates/ops/CLASS/storage_ops.cc.jinja diff --git a/symforce/codegen/cpp_templates/ops/CLASS/storage_ops.h.jinja b/symforce/codegen/backends/cpp/templates/ops/CLASS/storage_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/ops/CLASS/storage_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/ops/CLASS/storage_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/tests/cam_function_codegen_cpp_test.cc.jinja b/symforce/codegen/backends/cpp/templates/tests/cam_function_codegen_cpp_test.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/tests/cam_function_codegen_cpp_test.cc.jinja rename to symforce/codegen/backends/cpp/templates/tests/cam_function_codegen_cpp_test.cc.jinja diff --git a/symforce/codegen/cpp_templates/tests/cam_package_cpp_test.cc.jinja b/symforce/codegen/backends/cpp/templates/tests/cam_package_cpp_test.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/tests/cam_package_cpp_test.cc.jinja rename to symforce/codegen/backends/cpp/templates/tests/cam_package_cpp_test.cc.jinja diff --git a/symforce/codegen/cpp_templates/tests/geo_package_cpp_test.cc.jinja b/symforce/codegen/backends/cpp/templates/tests/geo_package_cpp_test.cc.jinja similarity index 100% rename from symforce/codegen/cpp_templates/tests/geo_package_cpp_test.cc.jinja rename to symforce/codegen/backends/cpp/templates/tests/geo_package_cpp_test.cc.jinja diff --git a/symforce/codegen/cpp_templates/type_ops.h.jinja b/symforce/codegen/backends/cpp/templates/type_ops.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/type_ops.h.jinja rename to symforce/codegen/backends/cpp/templates/type_ops.h.jinja diff --git a/symforce/codegen/cpp_templates/typedefs.h.jinja b/symforce/codegen/backends/cpp/templates/typedefs.h.jinja similarity index 100% rename from symforce/codegen/cpp_templates/typedefs.h.jinja rename to symforce/codegen/backends/cpp/templates/typedefs.h.jinja diff --git a/symforce/codegen/cpp_templates/util/util.jinja b/symforce/codegen/backends/cpp/templates/util/util.jinja similarity index 100% rename from symforce/codegen/cpp_templates/util/util.jinja rename to symforce/codegen/backends/cpp/templates/util/util.jinja diff --git a/symforce/codegen/backends/python/README.md b/symforce/codegen/backends/python/README.md new file mode 100644 index 000000000..e69de29bb diff --git a/symforce/codegen/backends/python/__init__.py b/symforce/codegen/backends/python/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/symforce/codegen/printers/python_code_printer.py b/symforce/codegen/backends/python/python_code_printer.py similarity index 96% rename from symforce/codegen/printers/python_code_printer.py rename to symforce/codegen/backends/python/python_code_printer.py index c6ae5223d..3c2d34822 100644 --- a/symforce/codegen/printers/python_code_printer.py +++ b/symforce/codegen/backends/python/python_code_printer.py @@ -3,10 +3,8 @@ # This source code is under the Apache 2.0 license found in the LICENSE file. # ---------------------------------------------------------------------------- -from sympy.printing.pycode import PythonCodePrinter as _PythonCodePrinter - -# Everything in this file is SymPy, not SymEngine (even when SymForce is on the SymEngine backend) import sympy +from sympy.printing.pycode import PythonCodePrinter as _PythonCodePrinter class PythonCodePrinter(_PythonCodePrinter): diff --git a/symforce/codegen/backends/python/python_config.py b/symforce/codegen/backends/python/python_config.py new file mode 100644 index 000000000..d2a805e06 --- /dev/null +++ b/symforce/codegen/backends/python/python_config.py @@ -0,0 +1,58 @@ +# ---------------------------------------------------------------------------- +# SymForce - Copyright 2022, Skydio, Inc. +# This source code is under the Apache 2.0 license found in the LICENSE file. +# ---------------------------------------------------------------------------- +from dataclasses import dataclass +from pathlib import Path +from sympy.printing.codeprinter import CodePrinter + +from symforce import typing as T +from symforce.codegen.codegen_config import CodegenConfig + + +CURRENT_DIR = Path(__file__).parent + + +@dataclass +class PythonConfig(CodegenConfig): + """ + Code generation config for the Python backend. + + Args: + doc_comment_line_prefix: Prefix applied to each line in a docstring + line_length: Maximum allowed line length in docstrings; used for formatting docstrings. + use_eigen_types: Use eigen_lcm types for vectors instead of lists + autoformat: Run a code formatter on the generated code + cse_optimizations: Optimizations argument to pass to sm.cse + use_numba: Add the `@numba.njit` decorator to generated functions. This will greatly + speed up functions by compiling them to machine code, but has large overhead + on the first call and some overhead on subsequent calls, so it should not be + used for small functions or functions that are only called a handfull of + times. + matrix_is_1D: geo.Matrix symbols get formatted as a 1D array + """ + + doc_comment_line_prefix: str = "" + line_length: int = 100 + use_eigen_types: bool = True + use_numba: bool = False + matrix_is_1d: bool = True + + @classmethod + def backend_name(cls) -> str: + return "python" + + @classmethod + def template_dir(cls) -> Path: + return CURRENT_DIR / "templates" + + def templates_to_render(self, generated_file_name: str) -> T.List[T.Tuple[str, str]]: + return [ + ("function/FUNCTION.py.jinja", f"{generated_file_name}.py"), + ("function/__init__.py.jinja", "__init__.py"), + ] + + def printer(self) -> CodePrinter: + from symforce.codegen.backends.python import python_code_printer + + return python_code_printer.PythonCodePrinter() diff --git a/symforce/codegen/python_templates/cam_package/CLASS.py.jinja b/symforce/codegen/backends/python/templates/cam_package/CLASS.py.jinja similarity index 100% rename from symforce/codegen/python_templates/cam_package/CLASS.py.jinja rename to symforce/codegen/backends/python/templates/cam_package/CLASS.py.jinja diff --git a/symforce/codegen/python_templates/cam_package/ops/CLASS/__init__.py.jinja b/symforce/codegen/backends/python/templates/cam_package/ops/CLASS/__init__.py.jinja similarity index 100% rename from symforce/codegen/python_templates/cam_package/ops/CLASS/__init__.py.jinja rename to symforce/codegen/backends/python/templates/cam_package/ops/CLASS/__init__.py.jinja diff --git a/symforce/codegen/python_templates/cam_package/ops/CLASS/camera_ops.py.jinja b/symforce/codegen/backends/python/templates/cam_package/ops/CLASS/camera_ops.py.jinja similarity index 100% rename from symforce/codegen/python_templates/cam_package/ops/CLASS/camera_ops.py.jinja rename to symforce/codegen/backends/python/templates/cam_package/ops/CLASS/camera_ops.py.jinja diff --git a/symforce/codegen/python_templates/function/FUNCTION.py.jinja b/symforce/codegen/backends/python/templates/function/FUNCTION.py.jinja similarity index 100% rename from symforce/codegen/python_templates/function/FUNCTION.py.jinja rename to symforce/codegen/backends/python/templates/function/FUNCTION.py.jinja diff --git a/symforce/codegen/python_templates/function/__init__.py.jinja b/symforce/codegen/backends/python/templates/function/__init__.py.jinja similarity index 100% rename from symforce/codegen/python_templates/function/__init__.py.jinja rename to symforce/codegen/backends/python/templates/function/__init__.py.jinja diff --git a/symforce/codegen/python_templates/geo_package/CLASS.py.jinja b/symforce/codegen/backends/python/templates/geo_package/CLASS.py.jinja similarity index 100% rename from symforce/codegen/python_templates/geo_package/CLASS.py.jinja rename to symforce/codegen/backends/python/templates/geo_package/CLASS.py.jinja diff --git a/symforce/codegen/python_templates/geo_package/__init__.py.jinja b/symforce/codegen/backends/python/templates/geo_package/__init__.py.jinja similarity index 100% rename from symforce/codegen/python_templates/geo_package/__init__.py.jinja rename to symforce/codegen/backends/python/templates/geo_package/__init__.py.jinja diff --git a/symforce/codegen/python_templates/geo_package/custom_methods/pose2.py.jinja b/symforce/codegen/backends/python/templates/geo_package/custom_methods/pose2.py.jinja similarity index 100% rename from symforce/codegen/python_templates/geo_package/custom_methods/pose2.py.jinja rename to symforce/codegen/backends/python/templates/geo_package/custom_methods/pose2.py.jinja diff --git a/symforce/codegen/python_templates/geo_package/custom_methods/pose3.py.jinja b/symforce/codegen/backends/python/templates/geo_package/custom_methods/pose3.py.jinja similarity index 100% rename from symforce/codegen/python_templates/geo_package/custom_methods/pose3.py.jinja rename to symforce/codegen/backends/python/templates/geo_package/custom_methods/pose3.py.jinja diff --git a/symforce/codegen/python_templates/geo_package/custom_methods/rot2.py.jinja b/symforce/codegen/backends/python/templates/geo_package/custom_methods/rot2.py.jinja similarity index 100% rename from symforce/codegen/python_templates/geo_package/custom_methods/rot2.py.jinja rename to symforce/codegen/backends/python/templates/geo_package/custom_methods/rot2.py.jinja diff --git a/symforce/codegen/python_templates/geo_package/custom_methods/rot3.py.jinja b/symforce/codegen/backends/python/templates/geo_package/custom_methods/rot3.py.jinja similarity index 100% rename from symforce/codegen/python_templates/geo_package/custom_methods/rot3.py.jinja rename to symforce/codegen/backends/python/templates/geo_package/custom_methods/rot3.py.jinja diff --git a/symforce/codegen/python_templates/ops/CLASS/__init__.py.jinja b/symforce/codegen/backends/python/templates/ops/CLASS/__init__.py.jinja similarity index 100% rename from symforce/codegen/python_templates/ops/CLASS/__init__.py.jinja rename to symforce/codegen/backends/python/templates/ops/CLASS/__init__.py.jinja diff --git a/symforce/codegen/python_templates/ops/CLASS/group_ops.py.jinja b/symforce/codegen/backends/python/templates/ops/CLASS/group_ops.py.jinja similarity index 100% rename from symforce/codegen/python_templates/ops/CLASS/group_ops.py.jinja rename to symforce/codegen/backends/python/templates/ops/CLASS/group_ops.py.jinja diff --git a/symforce/codegen/python_templates/ops/CLASS/lie_group_ops.py.jinja b/symforce/codegen/backends/python/templates/ops/CLASS/lie_group_ops.py.jinja similarity index 100% rename from symforce/codegen/python_templates/ops/CLASS/lie_group_ops.py.jinja rename to symforce/codegen/backends/python/templates/ops/CLASS/lie_group_ops.py.jinja diff --git a/symforce/codegen/backends/python/templates/ops/__init__.py.jinja b/symforce/codegen/backends/python/templates/ops/__init__.py.jinja new file mode 100644 index 000000000..e69de29bb diff --git a/symforce/codegen/python_templates/setup.py.jinja b/symforce/codegen/backends/python/templates/setup.py.jinja similarity index 100% rename from symforce/codegen/python_templates/setup.py.jinja rename to symforce/codegen/backends/python/templates/setup.py.jinja diff --git a/symforce/codegen/python_templates/tests/cam_package_python_test.py.jinja b/symforce/codegen/backends/python/templates/tests/cam_package_python_test.py.jinja similarity index 100% rename from symforce/codegen/python_templates/tests/cam_package_python_test.py.jinja rename to symforce/codegen/backends/python/templates/tests/cam_package_python_test.py.jinja diff --git a/symforce/codegen/python_templates/tests/geo_package_python_test.py.jinja b/symforce/codegen/backends/python/templates/tests/geo_package_python_test.py.jinja similarity index 100% rename from symforce/codegen/python_templates/tests/geo_package_python_test.py.jinja rename to symforce/codegen/backends/python/templates/tests/geo_package_python_test.py.jinja diff --git a/symforce/codegen/python_templates/util/util.jinja b/symforce/codegen/backends/python/templates/util/util.jinja similarity index 100% rename from symforce/codegen/python_templates/util/util.jinja rename to symforce/codegen/backends/python/templates/util/util.jinja diff --git a/symforce/codegen/cam_package_codegen.py b/symforce/codegen/cam_package_codegen.py index b2c94cda3..2c7c24b1c 100644 --- a/symforce/codegen/cam_package_codegen.py +++ b/symforce/codegen/cam_package_codegen.py @@ -257,7 +257,7 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: if isinstance(config, PythonConfig): logger.info(f'Creating Python package at: "{cam_package_dir}"') - template_dir = pathlib.Path(template_util.PYTHON_TEMPLATE_DIR) + template_dir = config.template_dir() # First generate the geo package as it's a dependency of the cam package from symforce.codegen import geo_package_codegen @@ -310,7 +310,7 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: elif isinstance(config, CppConfig): logger.info(f'Creating C++ cam package at: "{cam_package_dir}"') - template_dir = pathlib.Path(template_util.CPP_TEMPLATE_DIR, "cam_package") + template_dir = config.template_dir() / "cam_package" # First generate the geo package as it's a dependency of the cam package from symforce.codegen import geo_package_codegen @@ -331,9 +331,8 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: (".", "ops/CLASS/lie_group_ops.h"), (".", "ops/CLASS/lie_group_ops.cc"), ): - template_path = pathlib.Path( - template_util.CPP_TEMPLATE_DIR, base_dir, relative_path + ".jinja" - ) + template_path = config.template_dir() / base_dir / relative_path + ".jinja" + output_path = cam_package_dir / relative_path.replace( "CLASS", python_util.camelcase_to_snakecase(cls.__name__) ) diff --git a/symforce/codegen/codegen.py b/symforce/codegen/codegen.py index 1b02072c2..9814dd15c 100644 --- a/symforce/codegen/codegen.py +++ b/symforce/codegen/codegen.py @@ -23,13 +23,14 @@ from symforce import python_util from symforce import typing as T from symforce.values import Values +from symforce import codegen from symforce.codegen import template_util from symforce.codegen import codegen_util from symforce.codegen import codegen_config from symforce.codegen import types_package_codegen from symforce.type_helpers import symbolic_inputs -CURRENT_DIR = os.path.dirname(__file__) +CURRENT_DIR = Path(__file__).parent class LinearizationMode(enum.Enum): @@ -256,7 +257,7 @@ def common_data() -> T.Dict[str, T.Any]: data["DataBuffer"] = sm.DataBuffer data["Values"] = Values data["pathlib"] = pathlib - data["path_to_codegen"] = CURRENT_DIR + data["path_to_codegen"] = str(CURRENT_DIR) data["scalar_types"] = ("double", "float") data["camelcase_to_snakecase"] = python_util.camelcase_to_snakecase data["python_util"] = python_util @@ -393,56 +394,23 @@ def generate_function( self.namespace = namespace template_data = dict(self.common_data(), spec=self) + template_dir = self.config.template_dir() - # Generate the function - if isinstance(self.config, codegen_config.PythonConfig): - if skip_directory_nesting: - python_function_dir = output_dir - else: - python_function_dir = output_dir / "python" / "symforce" / namespace - - logger.info(f'Creating python function "{self.name}" at "{python_function_dir}"') - - templates.add( - Path(template_util.PYTHON_TEMPLATE_DIR) / "function" / "FUNCTION.py.jinja", - python_function_dir / f"{generated_file_name}.py", - template_data, - ) - templates.add( - Path(template_util.PYTHON_TEMPLATE_DIR) / "function" / "__init__.py.jinja", - python_function_dir / "__init__.py", - template_data, - ) - - out_function_dir = python_function_dir - elif isinstance(self.config, codegen_config.CppConfig): - if skip_directory_nesting: - cpp_function_dir = output_dir - else: - cpp_function_dir = output_dir / "cpp" / "symforce" / namespace - - logger.info( - f'Creating C++ function "{python_util.snakecase_to_camelcase(self.name)}" at "{cpp_function_dir}"' - ) + backend_name = self.config.backend_name() + if skip_directory_nesting: + out_function_dir = output_dir + else: + out_function_dir = output_dir / backend_name / "symforce" / namespace - templates.add( - Path(template_util.CPP_TEMPLATE_DIR) / "function" / "FUNCTION.h.jinja", - cpp_function_dir / f"{generated_file_name}.h", - template_data, - ) + logger.info(f'Creating {backend_name} function from "{self.name}" at "{out_function_dir}"') - if self.config.explicit_template_instantiation_types is not None: - templates.add( - Path(template_util.CPP_TEMPLATE_DIR) / "function" / "FUNCTION.cc.jinja", - cpp_function_dir / f"{generated_file_name}.cc", - template_data, - ) - - out_function_dir = cpp_function_dir - else: - raise NotImplementedError(f'Unknown config type: "{self.config}"') + # Get templates to render + for source, dest in self.config.templates_to_render(generated_file_name): + templates.add(template_dir / source, out_function_dir / dest, template_data) + # Render templates.render(autoformat=self.config.autoformat) + lcm_data = codegen_util.generate_lcm_types( lcm_type_dir=types_codegen_data["lcm_type_dir"], lcm_files=types_codegen_data["lcm_files"], diff --git a/symforce/codegen/codegen_config.py b/symforce/codegen/codegen_config.py index 9e03b3218..4c0e8a1e5 100644 --- a/symforce/codegen/codegen_config.py +++ b/symforce/codegen/codegen_config.py @@ -2,16 +2,20 @@ # SymForce - Copyright 2022, Skydio, Inc. # This source code is under the Apache 2.0 license found in the LICENSE file. # ---------------------------------------------------------------------------- - +from abc import abstractmethod from dataclasses import dataclass +from pathlib import Path +from sympy.printing.codeprinter import CodePrinter from symforce import typing as T +CURRENT_DIR = Path(__file__).parent + @dataclass class CodegenConfig: """ - Base class for language-specific arguments for code generation + Base class for backend-specific arguments for code generation. Args: doc_comment_line_prefix: Prefix applied to each line in a docstring, e.g. " * " for C++ @@ -20,6 +24,7 @@ class CodegenConfig: use_eigen_types: Use eigen_lcm types for vectors instead of lists autoformat: Run a code formatter on the generated code cse_optimizations: Optimizations argument to pass to sm.cse + matrix_is_1d: Whether geo.Matrix symbols get formatted as 1D """ doc_comment_line_prefix: str @@ -29,58 +34,45 @@ class CodegenConfig: cse_optimizations: T.Optional[ T.Union[T.Literal["basic"], T.Sequence[T.Tuple[T.Callable, T.Callable]]] ] = None + # TODO(hayk): Remove this parameter (by making everything 2D?) + matrix_is_1d: bool = False + @classmethod + @abstractmethod + def backend_name(cls) -> str: + """ + String name for the backend. This should match the directory name in codegen/backends + and will be used to namespace by backend in generated code. + """ + pass -@dataclass -class CppConfig(CodegenConfig): - """ - C++ Codegen configuration - - Args: - doc_comment_line_prefix: Prefix applied to each line in a docstring - line_length: Maximum allowed line length in docstrings; used for formatting docstrings. - use_eigen_types: Use eigen_lcm types for vectors instead of lists - autoformat: Run a code formatter on the generated code - cse_optimizations: Optimizations argument to pass to sm.cse - support_complex: Generate code that can work with std::complex or with regular float types - force_no_inline: Mark generated functions as `__attribute__((noinline))` - zero_initialization_sparsity_threshold: Threshold between 0 and 1 for the sparsity below - which we'll initialize an output matrix to 0, so we - don't have to generate a line to set each zero - element to 0 individually - explicit_template_instantiation_types: Explicity instantiates templated functions in a `.cc` - file for each given type. This allows the generated function to be compiled in its own - translation unit. Useful for large functions which take a long time to compile. - """ + @classmethod + @abstractmethod + def template_dir(cls) -> Path: + """ + Directory for jinja templates. + """ + pass - doc_comment_line_prefix: str = " * " - line_length: int = 100 - use_eigen_types: bool = True - support_complex: bool = False - force_no_inline: bool = False - zero_initialization_sparsity_threshold: float = 0.5 - explicit_template_instantiation_types: T.Optional[T.Sequence[str]] = None + @abstractmethod + def templates_to_render(self, generated_file_name: str) -> T.List[T.Tuple[str, str]]: + """ + Given a single symbolic function's filename, provide one or more Jinja templates to + render and the relative output paths where they should go. + """ + pass + @abstractmethod + def printer(self) -> CodePrinter: + """ + Return an instance of the code printer to use for this language. + """ + pass -@dataclass -class PythonConfig(CodegenConfig): - """ - Python Codegen configuration - - Args: - doc_comment_line_prefix: Prefix applied to each line in a docstring - line_length: Maximum allowed line length in docstrings; used for formatting docstrings. - use_eigen_types: Use eigen_lcm types for vectors instead of lists - autoformat: Run a code formatter on the generated code - cse_optimizations: Optimizations argument to pass to sm.cse - use_numba: Add the `@numba.njit` decorator to generated functions. This will greatly - speed up functions by compiling them to machine code, but has large overhead - on the first call and some overhead on subsequent calls, so it should not be - used for small functions or functions that are only called a handfull of - times. - """ - - doc_comment_line_prefix: str = "" - line_length: int = 100 - use_eigen_types: bool = True - use_numba: bool = False + # TODO(hayk): Move this into code printer. + @staticmethod + def format_data_accessor(prefix: str, index: int) -> str: + """ + Format data for accessing a data array in code. + """ + return f"{prefix}.data[{index}]" diff --git a/symforce/codegen/codegen_util.py b/symforce/codegen/codegen_util.py index 2c9d36964..2cdebfae2 100644 --- a/symforce/codegen/codegen_util.py +++ b/symforce/codegen/codegen_util.py @@ -23,7 +23,7 @@ from symforce.values import Values, IndexEntry from symforce import sympy as sm from symforce import typing as T -from symforce.codegen import printers, format_util +from symforce.codegen import format_util from symforce.codegen import codegen_config from symforce import python_util from symforce import _sympy_count_ops @@ -178,7 +178,7 @@ def count_ops(expr: sm.Expr) -> int: ) # Get printer - printer = get_code_printer(config) + printer = config.printer() # Print code intermediate_terms = [(str(var), printer.doprint(t)) for var, t in temps_formatted] @@ -341,20 +341,19 @@ def get_formatted_list( formatted_symbols = [sm.Symbol(key)] flattened_value = [value] elif issubclass(arg_cls, geo.Matrix): - if isinstance(config, codegen_config.PythonConfig): - # TODO(nathan): Not sure this works for 2D matrices + if config.matrix_is_1d: + # TODO(nathan): Not sure this works for 2D matrices. Get rid of this. formatted_symbols = [sm.Symbol(f"{key}[{j}]") for j in range(storage_dim)] - elif isinstance(config, codegen_config.CppConfig): - formatted_symbols = [] + else: # NOTE(brad): The order of the symbols must match the storage order of geo.Matrix # (as returned by geo.Matrix.to_storage). Hence, if there storage order were # changed to, say, row major, the below for loops would have to be swapped to # reflect that. + formatted_symbols = [] for j in range(value.shape[1]): for i in range(value.shape[0]): formatted_symbols.append(sm.Symbol(f"{key}({i}, {j})")) - else: - raise NotImplementedError() + flattened_value = ops.StorageOps.to_storage(value) elif issubclass(arg_cls, Values): @@ -466,12 +465,10 @@ def _get_scalar_keys_recursive( vec.extend(sm.Symbol(f"{prefix}[{i}]") for i in range(index_value.storage_dim)) else: # We have a geo/cam or other object that uses "data" to store a flat vector of scalars. - if isinstance(config, codegen_config.PythonConfig): - vec.extend(sm.Symbol(f"{prefix}.data[{i}]") for i in range(index_value.storage_dim)) - elif isinstance(config, codegen_config.CppConfig): - vec.extend(sm.Symbol(f"{prefix}.Data()[{i}]") for i in range(index_value.storage_dim)) - else: - raise NotImplementedError() + vec.extend( + sm.Symbol(config.format_data_accessor(prefix=prefix, index=i)) + for i in range(index_value.storage_dim) + ) assert len(vec) == len(set(vec)), "Non-unique keys:\n{}".format( [symbol for symbol in vec if vec.count(symbol) > 1] @@ -494,26 +491,6 @@ def get_formatted_sparse_list(sparse_outputs: Values) -> T.List[T.List[T.Scalar] return symbolic_args -def get_code_printer(config: codegen_config.CodegenConfig) -> "sm.CodePrinter": - """ - Pick a code printer for the given mode. - """ - # TODO(hayk): Consider symengine printer if this becomes slow. - - if isinstance(config, codegen_config.PythonConfig): - printer: sm.printing.codeprinter.CodePrinter = printers.PythonCodePrinter() - - elif isinstance(config, codegen_config.CppConfig): - if config.support_complex: - printer = printers.ComplexCppCodePrinter() - else: - printer = printers.CppCodePrinter() - else: - raise NotImplementedError(f"Unknown config type: {config}") - - return printer - - def _load_generated_package_internal(name: str, path: Path) -> T.Tuple[T.Any, T.List[str]]: """ Dynamically load generated package (or module). diff --git a/symforce/codegen/geo_package_codegen.py b/symforce/codegen/geo_package_codegen.py index 65abd8f5e..b3ba41ea6 100644 --- a/symforce/codegen/geo_package_codegen.py +++ b/symforce/codegen/geo_package_codegen.py @@ -145,7 +145,7 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: if isinstance(config, PythonConfig): logger.info(f'Creating Python package at: "{package_dir}"') - template_dir = Path(template_util.PYTHON_TEMPLATE_DIR) + template_dir = config.template_dir() # Build up templates for each type @@ -192,7 +192,7 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: sym_util_package_codegen.generate(config, output_dir=output_dir) logger.info(f'Creating C++ package at: "{package_dir}"') - template_dir = Path(template_util.CPP_TEMPLATE_DIR, "geo_package") + template_dir = config.template_dir() # Build up templates for each type for cls in DEFAULT_GEO_TYPES: @@ -210,14 +210,14 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: (".", "ops/CLASS/lie_group_ops.h"), (".", "ops/CLASS/lie_group_ops.cc"), ): - template_path = Path( - template_util.CPP_TEMPLATE_DIR, base_dir, relative_path + ".jinja" - ) + template_path = template_dir / base_dir / f"{relative_path}.jinja" output_path = package_dir / relative_path.replace("CLASS", cls.__name__.lower()) templates.add(template_path, output_path, data) # Render non geo type specific templates - for template_name in python_util.files_in_dir(str(template_dir / "ops"), relative=True): + for template_name in python_util.files_in_dir( + template_dir / "geo_package" / "ops", relative=True + ): if "CLASS" in template_name: continue @@ -225,15 +225,15 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: continue templates.add( - template_dir / "ops" / template_name, - package_dir / "ops" / template_name[: -len(".jinja")], + template_dir / "geo_package" / "ops" / template_name, + package_dir / "geo_package" / "ops" / template_name[: -len(".jinja")], dict(Codegen.common_data()), ) # Test example for name in ("geo_package_cpp_test.cc",): templates.add( - template_dir / ".." / "tests" / (name + ".jinja"), + template_dir / "tests" / (name + ".jinja"), Path(output_dir, "tests", name), dict( Codegen.common_data(), @@ -256,7 +256,7 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str: # LCM type_t templates.add( - Path(template_util.LCM_TEMPLATE_DIR, "symforce_types.lcm.jinja"), + template_util.LCM_TEMPLATE_DIR / "symforce_types.lcm.jinja", package_dir / ".." / "lcmtypes" / "lcmtypes" / "symforce_types.lcm", lcm_types_codegen.lcm_symforce_types_data(), ) diff --git a/symforce/codegen/printers/__init__.py b/symforce/codegen/printers/__init__.py index 8140e53a1..e69de29bb 100644 --- a/symforce/codegen/printers/__init__.py +++ b/symforce/codegen/printers/__init__.py @@ -1,7 +0,0 @@ -# ---------------------------------------------------------------------------- -# SymForce - Copyright 2022, Skydio, Inc. -# This source code is under the Apache 2.0 license found in the LICENSE file. -# ---------------------------------------------------------------------------- - -from .cpp_code_printer import CppCodePrinter, ComplexCppCodePrinter -from .python_code_printer import PythonCodePrinter diff --git a/symforce/codegen/sym_util_package_codegen.py b/symforce/codegen/sym_util_package_codegen.py index 185f36282..552ab52ae 100644 --- a/symforce/codegen/sym_util_package_codegen.py +++ b/symforce/codegen/sym_util_package_codegen.py @@ -29,14 +29,15 @@ def generate(config: codegen.CodegenConfig, output_dir: str = None) -> str: templates = template_util.TemplateList() if isinstance(config, codegen.CppConfig): + template_dir = config.template_dir() templates.add( - template_path=os.path.join(template_util.CPP_TEMPLATE_DIR, "typedefs.h.jinja"), + template_path=template_dir / "typedefs.h.jinja", output_path=os.path.join(package_dir, "typedefs.h"), data={}, ) templates.add( - template_path=os.path.join(template_util.CPP_TEMPLATE_DIR, "type_ops.h.jinja"), + template_path=template_dir / "type_ops.h.jinja", output_path=os.path.join(package_dir, "type_ops.h"), data=dict( python_util=python_util, @@ -44,17 +45,12 @@ def generate(config: codegen.CodegenConfig, output_dir: str = None) -> str: ), ) - templates.add( - template_path=os.path.join(template_util.CPP_TEMPLATE_DIR, "epsilon.h.jinja"), - output_path=os.path.join(package_dir, "epsilon.h"), - data={}, - ) - - templates.add( - template_path=os.path.join(template_util.CPP_TEMPLATE_DIR, "epsilon.cc.jinja"), - output_path=os.path.join(package_dir, "epsilon.cc"), - data={}, - ) + for filename in ("epsilon.h", "epsilon.cc"): + templates.add( + template_path=template_dir / f"{filename}.jinja", + output_path=os.path.join(package_dir, filename), + data={}, + ) else: # sym/util is currently C++ only pass diff --git a/symforce/codegen/template_util.py b/symforce/codegen/template_util.py index 3e947d322..e417cd51d 100644 --- a/symforce/codegen/template_util.py +++ b/symforce/codegen/template_util.py @@ -18,10 +18,8 @@ from symforce import typing as T from symforce.codegen import format_util -CURRENT_DIR = os.path.dirname(__file__) -CPP_TEMPLATE_DIR = os.path.join(CURRENT_DIR, "cpp_templates") -PYTHON_TEMPLATE_DIR = os.path.join(CURRENT_DIR, "python_templates") -LCM_TEMPLATE_DIR = os.path.join(CURRENT_DIR, "lcm_templates") +CURRENT_DIR = Path(__file__).parent +LCM_TEMPLATE_DIR = CURRENT_DIR / "lcm_templates" class FileType(enum.Enum): @@ -61,6 +59,45 @@ def from_template_path(template_path: Path) -> FileType: ) return FileType.from_extension(parts[-2]) + def comment_prefix(self) -> str: + """ + Return the comment prefix for this file type. + """ + if self in (FileType.CPP, FileType.CUDA, FileType.LCM): + return "//" + elif self in (FileType.PYTHON, FileType.PYTHON_INTERFACE): + return "#" + else: + raise NotImplementedError(f"Unknown comment prefix for {self}") + + def autoformat( + self, file_contents: str, template_name: T.Openable, output_path: T.Openable = None + ) -> str: + """ + Format code of this file type. + """ + # TODO(hayk): Move up to language-specific config or printer. This is quite an awkward + # place for auto-format logic, but I thought it was better centralized here than down below + # hidden in a function. We might want to somehow pass the config through to render a + # template so we can move things into the backend code. (tag=centralize-language-diffs) + if self in (FileType.CPP, FileType.CUDA): + # Come up with a fake filename to give to the formatter just for formatting purposes, + # even if this isn't being written to disk + if output_path is not None: + format_cpp_filename = os.path.basename(output_path) + else: + format_cpp_filename = os.fspath(template_name).replace(".jinja", "") + + return format_util.format_cpp( + file_contents, filename=str(CURRENT_DIR / format_cpp_filename) + ) + elif self == FileType.PYTHON: + return format_util.format_py(file_contents) + elif self == FileType.PYTHON_INTERFACE: + return format_util.format_pyi(file_contents) + else: + raise NotImplementedError(f"Unknown autoformatter for {self}") + class RelEnvironment(jinja2.Environment): """ @@ -73,21 +110,16 @@ def join_path(self, template: T.Union[jinja2.Template, str], parent: str) -> str return os.path.normpath(os.path.join(os.path.dirname(parent), str(template))) -def add_preamble(source: str, name: Path, filetype: FileType) -> str: - prefix = ( - "//" - if filetype in (FileType.CPP, FileType.CUDA, FileType.LCM, FileType.TYPESCRIPT) - else "#" - ) +def add_preamble(source: str, name: Path, comment_prefix: str) -> str: dashes = "-" * 77 return ( textwrap.dedent( f""" - {prefix} {dashes} - {prefix} This file was autogenerated by symforce from template: - {prefix} {name} - {prefix} Do NOT modify by hand. - {prefix} {dashes} + {comment_prefix} {dashes} + {comment_prefix} This file was autogenerated by symforce from template: + {comment_prefix} {name} + {comment_prefix} Do NOT modify by hand. + {comment_prefix} {dashes} """ ).lstrip() @@ -143,24 +175,14 @@ def render_template( filetype = FileType.from_template_path(Path(template_name)) template = jinja_env(template_dir).get_template(os.fspath(template_name)) - rendered_str = add_preamble(str(template.render(**data)), template_name, filetype) + rendered_str = add_preamble( + str(template.render(**data)), template_name, comment_prefix=filetype.comment_prefix() + ) if autoformat: - if filetype in (FileType.CPP, FileType.CUDA): - # Come up with a fake filename to give to the formatter just for formatting purposes, even - # if this isn't being written to disk - if output_path is not None: - format_cpp_filename = os.path.basename(output_path) - else: - format_cpp_filename = os.fspath(template_name).replace(".jinja", "") - - rendered_str = format_util.format_cpp( - rendered_str, filename=os.path.join(CURRENT_DIR, format_cpp_filename) - ) - elif filetype == FileType.PYTHON: - rendered_str = format_util.format_py(rendered_str) - elif filetype == FileType.PYTHON_INTERFACE: - rendered_str = format_util.format_pyi(rendered_str) + rendered_str = filetype.autoformat( + file_contents=rendered_str, template_name=template_name, output_path=output_path + ) if output_path: directory = os.path.dirname(output_path) diff --git a/symforce/codegen/types_package_codegen.py b/symforce/codegen/types_package_codegen.py index f998bd8c1..ac21a7b3f 100644 --- a/symforce/codegen/types_package_codegen.py +++ b/symforce/codegen/types_package_codegen.py @@ -94,7 +94,7 @@ def generate_types( lcm_files = [] if len(types_to_generate) > 0: logger.info(f'Creating LCM type at: "{lcm_type_dir}"') - lcm_template = os.path.join(template_util.LCM_TEMPLATE_DIR, "types.lcm.jinja") + lcm_template = template_util.LCM_TEMPLATE_DIR / "types.lcm.jinja" # Type definition lcm_file_name = f"{file_name}.lcm" diff --git a/symforce/codegen/values_codegen.py b/symforce/codegen/values_codegen.py index 8fecf7c05..caa94b0a9 100644 --- a/symforce/codegen/values_codegen.py +++ b/symforce/codegen/values_codegen.py @@ -7,7 +7,7 @@ import re from symforce import typing as T -from symforce.codegen import Codegen, template_util +from symforce.codegen import Codegen, CppConfig, template_util from symforce.values import generated_key_selection from symforce.values.values import Values @@ -33,6 +33,7 @@ def generate_values_keys( skip_directory_nesting: Generate the output file directly into output_dir instead of adding the usual directory structure inside output_dir """ + config = CppConfig() if not isinstance(output_dir, Path): output_dir = Path(output_dir) @@ -50,7 +51,7 @@ def generate_values_keys( cpp_function_dir = output_dir / "cpp" / "symforce" / namespace template_util.render_template( - template_path=Path(template_util.CPP_TEMPLATE_DIR) / "keys.h.jinja", + template_path=config.template_dir() / "keys.h.jinja", data=dict(Codegen.common_data(), namespace=namespace, vars=vars_to_generate), output_path=cpp_function_dir / generated_file_name, ) diff --git a/test/symforce_cpp_code_printer_codegen_test.py b/test/symforce_cpp_code_printer_codegen_test.py index cb7266626..851a31f44 100644 --- a/test/symforce_cpp_code_printer_codegen_test.py +++ b/test/symforce_cpp_code_printer_codegen_test.py @@ -7,13 +7,9 @@ import symforce -from symforce import typing as T -from symforce.test_util import TestCase, sympy_only -from symforce.values import Values - from symforce import codegen -from symforce.codegen import codegen_util from symforce import sympy as sm +from symforce.test_util import TestCase, sympy_only SYMFORCE_DIR = os.path.dirname(os.path.dirname(__file__)) TEST_DATA_DIR = os.path.join( @@ -27,7 +23,7 @@ class SymforceCppCodePrinterTest(TestCase): """ def test_max_min(self) -> None: - printer = codegen_util.get_code_printer(codegen.CppConfig()) + printer = codegen.CppConfig().printer() a = sm.Symbol("a") b = sm.Symbol("b") diff --git a/test/symforce_gen_codegen_test.py b/test/symforce_gen_codegen_test.py index 06f21822b..6aaa586d9 100644 --- a/test/symforce_gen_codegen_test.py +++ b/test/symforce_gen_codegen_test.py @@ -23,7 +23,6 @@ from symforce.codegen import sym_util_package_codegen from symforce.codegen import template_util from symforce.test_util import TestCase, symengine_only -from symforce import path_util SYMFORCE_DIR = os.path.dirname(os.path.dirname(__file__)) TEST_DATA_DIR = os.path.join( @@ -87,9 +86,10 @@ def test_gen_package_codegen_python(self) -> None: """ output_dir = self.make_output_dir("sf_gen_codegen_test_") - cam_package_codegen.generate(config=codegen.PythonConfig(), output_dir=output_dir) + config = codegen.PythonConfig() + cam_package_codegen.generate(config=config, output_dir=output_dir) template_util.render_template( - template_path=os.path.join(template_util.PYTHON_TEMPLATE_DIR, "setup.py.jinja"), + template_path=config.template_dir() / "setup.py.jinja", output_path=os.path.join(output_dir, "setup.py"), data=dict( package_name="symforce-sym", @@ -140,6 +140,7 @@ def test_gen_package_codegen_cpp(self) -> None: """ Test C++ code generation """ + config = codegen.CppConfig() output_dir = self.make_output_dir("sf_gen_codegen_test_") # Prior factors, between factors, and SLAM factors for C++. @@ -147,11 +148,11 @@ def test_gen_package_codegen_cpp(self) -> None: slam_factors_codegen.generate(os.path.join(output_dir, "sym")) # Generate typedefs.h - sym_util_package_codegen.generate(config=codegen.CppConfig(), output_dir=output_dir) + sym_util_package_codegen.generate(config=config, output_dir=output_dir) # Generate cam package, geo package, and tests # This calls geo_package_codegen.generate internally - cam_package_codegen.generate(config=codegen.CppConfig(), output_dir=output_dir) + cam_package_codegen.generate(config=config, output_dir=output_dir) # Check against existing generated package (only on SymEngine) self.compare_or_update_directory( diff --git a/test/symforce_lcm_codegen_test.py b/test/symforce_lcm_codegen_test.py index f3bd969fb..69e902b78 100644 --- a/test/symforce_lcm_codegen_test.py +++ b/test/symforce_lcm_codegen_test.py @@ -17,7 +17,7 @@ def test_generate_lcm(self) -> None: output_dir = self.make_output_dir("sf_lcm_codegen_test") template_util.render_template( - os.path.join(template_util.LCM_TEMPLATE_DIR, "symforce_types.lcm.jinja"), + template_util.LCM_TEMPLATE_DIR / "symforce_types.lcm.jinja", data=lcm_symforce_types_data(), output_path=os.path.join(output_dir, "symforce_types.lcm"), )