Modularize backend-specific code

Move almost all backend-specific code into a sub-package of symforce/codegen/backends, rather than being scattered around multiple core files. This includes the code printer, codegen config, and a bunch of switch statements around the codegen machinery. The goal is to centralize and make it cleaner to add new codegen backends. These changes were made in parallel with adding a javascript backend, which will come in a follow-on review.

See codegen/backends/README.md for the steps to add a backend after this review.

There are still several ways to improve code quality and reorganize backend-specific code, but this is a step forward.

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

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`.

symforce/codegen/printers/cpp_code_printer.py → 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):
     """

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}]"

symforce/codegen/printers/python_code_printer.py → 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):

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()

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__)
                 )

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"],