[Symforce] Allow specific explicit template instantiations

Previously we allowed users to set a flag to indicate if they wanted
generated C++ functions to be explicitly instantiated, and would
explicitly instantiate a float and a double version of the function.
However, compliling very large functions can be slow because both the
float and double versions of the function must be compiled, even if only
one is used. This PR allows the user to do explicit template
instantiation on whatever types they like, meaning they can choose to
only do explicit template instantiation on float or double if they like.

Topic: symforce_explicit_template_instantiation_list
GitOrigin-RevId: 27e5b54441a80928b85019df44fc605ecbdf75aa

Author: nathan-skydio

symforce/codegen/codegen.py

@@ -431,7 +431,7 @@ def generate_function(
                 template_data,
             )
 
-            if self.config.use_explicit_template_instantiation:
+            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",

symforce/codegen/codegen_config.py

@@ -48,9 +48,9 @@ class CppConfig(CodegenConfig):
                                                 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
-        use_explicit_template_instantiation: Explicity instantiate templated functions in a `.cc`
-            file so that generated function can be compiled in its own translation unit. Useful
-            for large functions which take a long time to compile.
+        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 = " * "
@@ -59,7 +59,7 @@ class CppConfig(CodegenConfig):
     support_complex: bool = False
     force_no_inline: bool = False
     zero_initialization_sparsity_threshold: float = 0.5
-    use_explicit_template_instantiation: bool = False
+    explicit_template_instantiation_types: T.Optional[T.Sequence[str]] = None
 
 
 @dataclass

symforce/codegen/cpp_templates/function/FUNCTION.cc.jinja

@@ -9,11 +9,11 @@
 
 namespace {{ spec.namespace }} {
 
+{% for type in spec.config.explicit_template_instantiation_types %}
 {% set name = python_util.snakecase_to_camelcase(spec.name) %}
-template {{ util.get_return_type(spec, scalar_type="double") }} {{ name }}<double>(
-    {{- util.input_args_declaration(spec, is_declaration=False, scalar_type="double") -}});
+template {{ util.get_return_type(spec, scalar_type=type) }} {{ name }}<{{ type }}>(
+    {{- util.input_args_declaration(spec, is_declaration=False, scalar_type=type) -}});
 
-template {{ util.get_return_type(spec, scalar_type="float") }} {{ name }}<float>(
-    {{- util.input_args_declaration(spec, is_declaration=False, scalar_type="float") -}});
+{% endfor -%}
 
 }  // namespace {{ spec.namespace }}

symforce/codegen/cpp_templates/function/FUNCTION.h.jinja

@@ -46,14 +46,13 @@ __attribute__((noinline))
     {{ util.expr_code(spec) -}}
 }  // NOLINT(readability/fn_size)
 
-{% if spec.config.use_explicit_template_instantiation %}
+{% if spec.config.explicit_template_instantiation_types is not none %}
+{% for type in spec.config.explicit_template_instantiation_types %}
 {% set name = python_util.snakecase_to_camelcase(spec.name) %}
-extern template {{ util.get_return_type(spec, scalar_type="double") }} {{ name }}<double>(
-    {{- util.input_args_declaration(spec, is_declaration=False, scalar_type="double") -}});
+extern template {{ util.get_return_type(spec, scalar_type=type) }} {{ name }}<{{ type }}>(
+    {{- util.input_args_declaration(spec, is_declaration=False, scalar_type=type) -}});
 
-extern template {{ util.get_return_type(spec, scalar_type="float") }} {{ name }}<float>(
-    {{- util.input_args_declaration(spec, is_declaration=False, scalar_type="float") -}});
-    
+{% endfor -%}
 {% endif -%}
 
 // NOLINTNEXTLINE(readability/fn_size)

test/symforce_codegen_test.py

@@ -837,7 +837,7 @@ def test_function_explicit_template_instantiation(self) -> None:
         cpp_func = codegen.Codegen(
             inputs,
             outputs,
-            codegen.CppConfig(use_explicit_template_instantiation=True),
+            codegen.CppConfig(explicit_template_instantiation_types=["double", "float"]),
             "codegen_explicit_template_instantiation_test",
         )
         shared_types = {