JavaScript backend core functionality

symforce/codegen/__init__.py

@@ -7,5 +7,9 @@
 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.javascript.javascript_config import JavascriptConfig
+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/backends/cpp/cpp_config.py

@@ -0,0 +1,80 @@
+# ----------------------------------------------------------------------------
+# SymForce - Copyright 2022, Skydio, Inc.
+# This source code is under the Apache 2.0 license found in the LICENSE file.
+# ----------------------------------------------------------------------------
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+
+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) -> "sympy.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}]"
+
+    @staticmethod
+    def format_matrix_accessor(key: str, i: int, j: int = None) -> str:
+        if j is None:
+            return f"{key}({i}, {0})"
+        return f"{key}({i}, {j})"

symforce/codegen/backends/javascript/javascript_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 __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+
+from symforce import typing as T
+from symforce.codegen.codegen_config import CodegenConfig
+
+
+CURRENT_DIR = Path(__file__).parent
+
+
+@dataclass
+class JavascriptConfig(CodegenConfig):
+    """
+    Code generation config for the javascript 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
+        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
+    # NOTE(hayk): Add JS autoformatter
+    autoformat: bool = False
+
+    @classmethod
+    def backend_name(cls) -> str:
+        return "javascript"
+
+    @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.js.jinja", f"{generated_file_name}.js"),
+        ]
+
+    def printer(self) -> "sm.CodePrinter":
+        from symforce.codegen.printers import javascript_code_printer
+
+        return javascript_code_printer.JavascriptCodePrinter()
+
+    @staticmethod
+    def format_matrix_accessor(key: str, i: int, j: int = None) -> str:
+        if j is None:
+            return f"{key}[{i}]"
+        return f"{key}[{i}][{j}]"

symforce/codegen/backends/javascript/templates/function/FUNCTION.js.jinja

@@ -0,0 +1,11 @@
+{# ------------------------------------------------------------------------- #}
+{# Function codegen template for Javascript                                  #}
+{# ------------------------------------------------------------------------- #}
+{%- import "../util/util.jinja" as util with context -%}
+
+{% if spec.docstring %}
+{{ util.print_docstring(spec.docstring) }}
+{% endif %}
+{{ util.function_declaration(spec) -}} {
+{{ util.expr_code(spec) }}
+}

symforce/codegen/backends/javascript/templates/util/util.jinja

@@ -0,0 +1,97 @@
+{# ------------------------------------------------------------------------- #}
+{# Utilities for Javascript code generation templates.                           #}
+{# ------------------------------------------------------------------------- #}
+
+{# ------------------------------------------------------------------------- #}
+
+ {# Format function docstring
+ #
+ # Args:
+ #     docstring (str):
+ #}
+{% macro print_docstring(docstring) %}
+{%- if docstring %}
+
+/*
+{%- for line in docstring.split('\n') %}
+*{{ ' {}'.format(line).rstrip() }}
+{% endfor -%}
+*/
+{%- endif -%}
+{% endmacro %}
+
+{# ------------------------------------------------------------------------- #}
+
+{# Generate function declaration
+ #
+ # Args:
+ #     spec (Codegen):
+ #}
+{%- macro function_declaration(spec) -%}
+function {{ camelcase_to_snakecase(spec.name) }}(
+    {%- for name in spec.inputs.keys() -%}
+    {{ name }}{% if not loop.last %}, {% endif %}
+    {%- endfor -%})
+{% endmacro -%}
+
+{# ------------------------------------------------------------------------- #}
+
+{# Generate inner code for computing the given expression.
+ #
+ # Args:
+ #     spec (Codegen):
+ #}
+{% macro expr_code(spec) %}
+    // Total ops: {{ spec.print_code_results.total_ops }}
+
+    // Input arrays
+    {% for name, type in spec.inputs.items() %}
+        {% set T = python_util.get_type(type) %}
+        {% if not issubclass(T, Values) and not issubclass(T, Matrix) and not is_symbolic(type) and not is_sequence(type) %}
+    _{{ name }} = {{ name }}.data
+        {% endif %}
+    {% endfor %}
+
+    // Intermediate terms ({{ spec.print_code_results.intermediate_terms | length }})
+    {% for lhs, rhs in spec.print_code_results.intermediate_terms %}
+    const {{ lhs }} = {{ rhs }};
+    {% endfor %}
+
+    // Output terms ({{ spec.outputs.items() | length }})
+    {% for name, type, terms in spec.print_code_results.dense_terms %}
+        {%- set T = python_util.get_type(type) -%}
+        {% if issubclass(T, Matrix) and type.shape[1] > 1 %}
+            {% set rows = type.shape[0] %}
+            {% set cols = type.shape[1] %}
+    let _{{ name }} = [...Array({{ rows }})].map(e => Array({{ cols }}));
+            {% set ns = namespace(iter=0) %}
+            {% for i in range(rows) %}
+                {% for j in range(cols) %}
+    _{{ name }}[{{ i }}][{{ j }}] = {{ terms[ns.iter][1] }};
+                    {% set ns.iter = ns.iter + 1 %}
+                {% endfor %}
+            {% endfor %}
+        {% elif not is_symbolic(type) %}
+            {% set dims = ops.StorageOps.storage_dim(type) %}
+    let _{{name}} = new Array({{ dims }});
+            {% for i in range(dims) %}
+    _{{ name }}[{{ i }}] = {{ terms[i][1] }};
+            {% endfor %}
+        {% else %}
+    const _{{name}} = {{ terms[0][1] }};
+        {% endif %}
+
+    {% endfor %}
+    return {
+    {% for name, type in spec.outputs.items() %}
+        {% set T = python_util.get_type(type) %}
+        {% if issubclass(T, (Matrix, Values)) or is_sequence(type) or is_symbolic(type) %}
+      {{ name }}: _{{name}}
+        {%- else %}
+      {{ name }}: sym.{{T.__name__}}.from_storage(_{{name}})
+        {% endif %}
+        {% if not loop.last %}, {% endif %}
+
+    {% endfor %}
+    };
+{% endmacro %}

symforce/codegen/backends/python/python_config.py

@@ -0,0 +1,63 @@
+# ----------------------------------------------------------------------------
+# SymForce - Copyright 2022, Skydio, Inc.
+# This source code is under the Apache 2.0 license found in the LICENSE file.
+# ----------------------------------------------------------------------------
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+
+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
+
+    @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) -> "sympy.CodePrinter":
+        from symforce.codegen.backends.python import python_code_printer
+
+        return python_code_printer.PythonCodePrinter()
+
+    @staticmethod
+    def format_matrix_accessor(key: str, i: int, j: int = None) -> str:
+        if j is None:
+            return f"{key}[{i}]"
+        return f"{key}[{i}, {j}]"

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