JavaScript backend core functionality

symforce/codegen/__init__.py

@@ -7,5 +7,11 @@
 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
+
+
+# TODO(hayk): Do we want to explicitly expose all configs here? (tag=centralize-language-diffs)
+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 aims to make it very straightforward to add new backends.
+
+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 C++ and Python examples.

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 __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) -> "sm.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/backends/javascript/javascript_config.py

@@ -0,0 +1,52 @@
+# ----------------------------------------------------------------------------
+# 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
+    matrix_is_1d: bool = True
+
+    @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()

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

@@ -0,0 +1,23 @@
+{# ------------------------------------------------------------------------- #}
+{# Function codegen template for Javascript                                  #}
+{# ------------------------------------------------------------------------- #}
+{%- import "../util/util.jinja" as util with context -%}
+
+{% for typename in ('Rot3', 'Rot2', 'Pose3', 'Pose2') %}
+{% if typename in spec.types_included %}
+{# #include <sym/{{ python_util.camelcase_to_snakecase(typename) }}.h> #}
+{% endif %}
+{% endfor %}
+
+{%- for name in spec.inputs.keys() | list + spec.outputs.keys() | list %}
+{% if name in spec.typenames_dict %}
+{# #include <lcmtypes/{{ spec.namespaces_dict[name] }}/{{ spec.typenames_dict[name] }}.hpp> #}
+{% endif %}
+{% endfor %}
+
+{% 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,177 @@
+{# ------------------------------------------------------------------------- #}
+{# Utilities for Javascript code generation templates.                           #}
+{# ------------------------------------------------------------------------- #}
+
+{# Vector type of a given dimension
+ #
+ # Args:
+ #     dim (int):
+ #}
+{%- macro vector_type(dim) -%}
+Eigen::Matrix<Scalar, {{ dim }}, 1>
+{%- endmacro -%}
+
+{# Matrix type of a given dimension
+ #
+ # Args:
+ #     rows (int):
+ #     cols (int):
+ #     is_input (bool): Is this an input argument or return value?
+ #}
+{%- macro matrix_type(rows, cols, is_input) -%}
+    {%- if cols == 1 -%}
+        T.{%- if is_input -%}Sequence{%- else -%}List{%- endif -%}[float]
+    {%- else -%}
+        numpy.ndarray
+    {%- endif -%}
+{%- endmacro -%}
+
+{# ------------------------------------------------------------------------- #}
+
+{# Convert a class to the emitted string
+ #
+ # Args:
+ #     T_or_value (type or Element):
+ #     name (str): Name in case type is a generated struct
+ #     is_input (bool): Is this an input argument or return value?
+ #}
+{%- macro format_typename(T_or_value, name, is_input) %}
+    {%- set T = python_util.get_type(T_or_value) -%}
+    {%- if T.__name__ == 'Symbol' or is_symbolic(T_or_value) -%}
+        float
+    {%- elif T.__name__ == 'NoneType' -%}
+        None
+    {%- elif issubclass(T, Matrix) -%}
+        {{ matrix_type(T_or_value.shape[0], T_or_value.shape[1], is_input) }}
+    {%- elif issubclass(T, Values) -%}
+        {#- TODO(aaron): We don't currently know where to import lcmtypes from or what they should be
+         # called, at some point we should fix this and do something like
+         # {{ spec.namespaces_dict[name] }}.{{ spec.typenames_dict[name] }}
+        -#}
+        T.Any
+    {%- elif is_sequence(T_or_value) -%}
+        {%- if is_input -%}
+            T.Sequence[{{ format_typename(T_or_value[0], name, is_input) }}]
+        {%- else -%}
+            T.List[float]
+        {%- endif -%}
+    {%- else -%}
+        {%- if "geo" in T.__module__ or "cam" in T.__module__ -%}
+            sym.
+        {%- endif -%}
+        {{- T.__name__ -}}
+    {%- endif -%}
+{% endmacro -%}
+
+{# ------------------------------------------------------------------------- #}
+
+{# Get the type of the object in the ouput Values with key given by spec.return_key
+ #
+ # Args:
+ #     spec (Codegen):
+ #}
+{%- macro get_return_type(spec) %}
+    {%- if spec.outputs.keys() | length == 1 -%}
+        {%- set name, type = spec.outputs.items() | first -%}
+        {{ format_typename(type, name, is_input=False) }}
+    {%- elif spec.outputs -%}
+        T.Tuple[
+        {%- for name, type in spec.outputs.items() -%}
+        {{ format_typename(type, name, is_input=False) }}{% if not loop.last %}, {% endif %}
+        {%- endfor -%}]
+    {%- else -%}
+        None
+    {%- endif -%}
+{% endmacro -%}
+
+{# ------------------------------------------------------------------------- #}
+
+ {# 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] %}
+    _{{ name }} = numpy.zeros(({{ rows }}, {{ 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}} = 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}}
+        {%- else %}
+ sym.{{T.__name__}}.from_storage(_{{name}})
+        {%- endif %}
+        {%- if not loop.last %}, {% endif %}
+    {%- endfor -%}
+{% endmacro %}