Load model libraries at runtime in cython ekf (#37)

* RednoseCompileFilter tool

* Refactor sconscripts

* Move lst_sq and feature_handler to examples

* add ekf_load_code_if_needed

* Fat binary for cython ext

* Restructure. Remove common_ekf. Load libs at runtime

* Build compatibility with openpilot

* ekf_lib_init

* Fix styling issues

* Fix linker flags for mac

* Remove useless pyx imports

* Build static lib instead of shared

* newlines

* Remove lst_sq_computer and feature_handler

.gitignore

@@ -15,6 +15,7 @@ __pycache__/
 *$py.class
 
 # C extensions
+*.a
 *.o
 *.so
 

SConstruct

@@ -36,7 +36,8 @@ env = Environment(
   CFLAGS="-std=gnu11",
   CXXFLAGS="-std=c++1z",
   CPPPATH=cpppath,
-  tools=["default", "cython"],
+  REDNOSE_ROOT=Dir("#").abspath,
+  tools=["default", "cython", "rednose_filter"],
 )
 
 # Cython build enviroment
@@ -52,17 +53,7 @@ elif arch == "aarch64":
 else:
   envCython["LINKFLAGS"] = ["-pthread", "-shared"]
 
-rednose_config = {
-  'generated_folder': '#examples/generated',
-  'to_build': {
-    'live': ('#examples/live_kf.py', True, [], []),
-    'kinematic': ('#examples/kinematic_kf.py', True, [], []),
-    'compare': ('#examples/test_compare.py', True, [], []),
-    'pos_computer_4': ('#rednose/helpers/lst_sq_computer.py', False, [], []),
-    'pos_computer_5': ('#rednose/helpers/lst_sq_computer.py', False, [], []),
-    'feature_handler_5': ('#rednose/helpers/feature_handler.py', False, [], []),
-  },
-}
+Export('env', 'envCython', 'common')
 
-Export('env', 'envCython', 'arch', 'rednose_config', 'common')
 SConscript(['#rednose/SConscript'])
+SConscript(['#examples/SConscript'])

examples/SConscript

@@ -0,0 +1,19 @@
+Import('env')
+
+gen_dir = Dir('generated/').abspath
+
+env.RednoseCompileFilter(
+  target="live",
+  filter_gen_script="live_kf.py",
+  output_dir=gen_dir,
+)
+env.RednoseCompileFilter(
+  target="kinematic",
+  filter_gen_script="kinematic_kf.py",
+  output_dir=gen_dir,
+)
+env.RednoseCompileFilter(
+  target="compare",
+  filter_gen_script="test_compare.py",
+  output_dir=gen_dir,
+)

examples/test_kinematic_kf.py

@@ -38,11 +38,11 @@ def test_kinematic_kf(self):
 
       # Retrieve kf values
       state = kf.x
-      xs_kf.append(float(state[States.POSITION]))
-      vs_kf.append(float(state[States.VELOCITY]))
+      xs_kf.append(float(state[States.POSITION].item()))
+      vs_kf.append(float(state[States.VELOCITY].item()))
       std = np.sqrt(kf.P)
-      xs_kf_std.append(float(std[States.POSITION, States.POSITION]))
-      vs_kf_std.append(float(std[States.VELOCITY, States.VELOCITY]))
+      xs_kf_std.append(float(std[States.POSITION, States.POSITION].item()))
+      vs_kf_std.append(float(std[States.VELOCITY, States.VELOCITY].item()))
 
       # Update simulation
       x += v * dt

rednose/SConscript

@@ -1,45 +1,17 @@
-Import('env', 'envCython', 'arch', 'rednose_config', 'common')
+Import('env', 'envCython', 'common')
 
-generated_folder = rednose_config['generated_folder']
-
-templates = Glob('#rednose/templates/*')
-
-sympy_helpers = "#rednose/helpers/sympy_helpers.py"
-ekf_sym = "#rednose/helpers/ekf_sym.py"
-ekf_sym_pyx = "#rednose/helpers/ekf_sym_pyx.pyx"
-ekf_sym_cc = env.Object("#rednose/helpers/ekf_sym.cc")
-common_ekf = "#rednose/helpers/common_ekf.cc"
-
-found = {}
-for target, (command, *args) in rednose_config['to_build'].items():
-  if File(command).exists():
-    found[target] = (command, *args)
-
-lib_target = [common_ekf]
-for target, (command, combined_lib, extra_generated, deps) in found.items():
-  target_files = File([f'{generated_folder}/{target}.cpp', f'{generated_folder}/{target}.h'])
-  extra_generated = [File(f'{generated_folder}/{x}') for x in extra_generated]
-  command_file = File(command)
-
-  env.Command(target_files + extra_generated,
-              [templates, command_file, sympy_helpers, ekf_sym] + deps,
-              command_file.get_abspath() + " " + target + " " + Dir(generated_folder).get_abspath())
-
-  if combined_lib:
-    lib_target.append(target_files[0])
-  else:
-    env.SharedLibrary(f'{generated_folder}/' + target, [target_files[0], common_ekf])
-
-libkf = env.SharedLibrary(f'{generated_folder}/libkf', lib_target)
-
-lenv = envCython.Clone()
-lenv["LINKFLAGS"] += [libkf[0].get_labspath()]
-
-# for SWAGLOG support
+cc_sources = [
+  "helpers/ekf_load.cc",
+  "helpers/ekf_sym.cc",
+]
+libs = ["dl"]
 if common != "":
-  lenv["LIBS"] = ['zmq', common] + lenv["LIBS"]
+  # for SWAGLOG support
+  libs += [common, 'zmq']
 
-ekf_sym_so = lenv.Program('#rednose/helpers/ekf_sym_pyx.so', [ekf_sym_pyx, ekf_sym_cc, common_ekf])
-lenv.Depends(ekf_sym_so, libkf)
+ekf_objects = env.SharedObject(cc_sources)
+rednose = env.Library("helpers/ekf_sym", ekf_objects, LIBS=libs)
+rednose_python = envCython.Program("helpers/ekf_sym_pyx.so", ["helpers/ekf_sym_pyx.pyx", ekf_objects],
+                                   LIBS=libs + envCython["LIBS"])
 
-Export('libkf')
+Export('rednose', 'rednose_python')

rednose/helpers/__init__.py

@@ -13,11 +13,9 @@ def write_code(folder, name, code, header):
   open(os.path.join(folder, f"{name}.h"), 'w', encoding='utf-8').write(header)
 
 
-def load_code(folder, name, lib_name=None):
-  if lib_name is None:
-    lib_name = name
+def load_code(folder, name):
   shared_ext = "dylib" if platform.system() == "Darwin" else "so"
-  shared_fn = os.path.join(folder, f"lib{lib_name}.{shared_ext}")
+  shared_fn = os.path.join(folder, f"lib{name}.{shared_ext}")
   header_fn = os.path.join(folder, f"{name}.h")
 
   with open(header_fn, encoding='utf-8') as f:

rednose/helpers/common_ekf.h → rednose/helpers/ekf.h

@@ -32,12 +32,11 @@ struct EKF {
   std::unordered_map<std::string, extra_routine_t> extra_routines = {};
 };
 
-std::vector<const EKF*>& get_ekfs();
-const EKF* ekf_lookup(const std::string& ekf_name);
-
-void ekf_register(const EKF* ekf);
-
-#define ekf_init(ekf) \
+#define ekf_lib_init(ekf) \
+extern "C" void* ekf_get() { \
+  return (void*) &ekf; \
+} \
+extern void  __attribute__((weak)) ekf_register(const EKF* ptr); \
 static void __attribute__((constructor)) do_ekf_init_ ## ekf(void) { \
-  ekf_register(&ekf); \
+  if (ekf_register) ekf_register(&ekf); \
 }

rednose/helpers/ekf_load.cc

@@ -0,0 +1,39 @@
+#include "ekf_load.h"
+#include <dlfcn.h>
+
+std::vector<const EKF*>& ekf_get_all() {
+  static std::vector<const EKF*> vec;
+  return vec;
+}
+
+void ekf_register(const EKF* ekf) {
+  ekf_get_all().push_back(ekf);
+}
+
+const EKF* ekf_lookup(const std::string& ekf_name) {
+  for (const auto& ekfi : ekf_get_all()) {
+    if (ekf_name == ekfi->name) {
+      return ekfi;
+    }
+  }
+  return NULL;
+}
+
+void ekf_load_and_register(const std::string& ekf_directory, const std::string& ekf_name) {
+  if (ekf_lookup(ekf_name)) {
+    return;
+  }
+
+#ifdef __APPLE__
+  std::string dylib_ext = ".dylib";
+#else
+  std::string dylib_ext = ".so";
+#endif
+  std::string ekf_path = ekf_directory + "/lib" + ekf_name + dylib_ext;
+  void* handle = dlopen(ekf_path.c_str(), RTLD_NOW);
+  assert(handle);
+  void* (*ekf_get)() = (void*(*)())dlsym(handle, "ekf_get");
+  assert(ekf_get != NULL);
+  const EKF* ekf = (const EKF*)ekf_get();
+  ekf_register(ekf);
+}

rednose/helpers/ekf_load.h

@@ -0,0 +1,9 @@
+#include <vector>
+#include <string>
+
+#include "ekf.h"
+
+std::vector<const EKF*>& ekf_get_all();
+const EKF* ekf_lookup(const std::string& ekf_name);
+void ekf_register(const EKF* ekf);
+void ekf_load_and_register(const std::string& ekf_directory, const std::string& ekf_name);

rednose/helpers/ekf_sym.cc

@@ -9,7 +9,6 @@ EKFSym::EKFSym(std::string name, Map<MatrixXdr> Q, Map<VectorXd> x_initial, Map<
     std::vector<int> quaternion_idxs, std::vector<std::string> global_vars, double max_rewind_age)
 {
   // TODO: add logger
-
   this->ekf = ekf_lookup(name);
   assert(this->ekf);
 

rednose/helpers/ekf_sym.h

@@ -12,7 +12,8 @@
 
 #include <eigen3/Eigen/Dense>
 
-#include "common_ekf.h"
+#include "ekf.h"
+#include "ekf_load.h"
 
 #define REWIND_TO_KEEP 512
 

rednose/helpers/ekf_sym.py

@@ -116,7 +116,7 @@ def gen_code(folder, name, f_sym, dt_sym, x_sym, obs_eqs, dim_x, dim_err, eskf_p
   sympy_header, code = sympy_into_c(sympy_functions, global_vars)
 
   header = "#pragma once\n"
-  header += "#include \"rednose/helpers/common_ekf.h\"\n"
+  header += "#include \"rednose/helpers/ekf.h\"\n"
   header += "extern \"C\" {\n"
 
   pre_code = f"#include \"{name}.h\"\n"
@@ -200,7 +200,7 @@ def gen_code(folder, name, f_sym, dt_sym, x_sym, obs_eqs, dim_x, dim_err, eskf_p
     post_code += f"    {{ \"{f}\", {name}_{f} }},\n"
   post_code += "  },\n"
   post_code += "};\n\n"
-  post_code += f"ekf_init({name});\n"
+  post_code += f"ekf_lib_init({name})\n"
 
   # merge code blocks
   header += "}"
@@ -252,7 +252,7 @@ def __init__(self, folder, name, Q, x_initial, P_initial, dim_main, dim_main_err
     self.rewind_obscache = []
     self.init_state(x_initial, P_initial, None)
 
-    ffi, lib = load_code(folder, name, "kf")
+    ffi, lib = load_code(folder, name)
     kinds, self.feature_track_kinds = [], []
     for func in dir(lib):
       if func[:len(name) + 3] == f'{name}_h_':

rednose/helpers/ekf_sym_pyx.pyx

@@ -11,12 +11,16 @@ cimport numpy as np
 
 import numpy as np
 
+
 cdef extern from "<optional>" namespace "std" nogil:
   cdef cppclass optional[T]:
     ctypedef T value_type
     bool has_value()
     T& value()
 
+cdef extern from "rednose/helpers/ekf_load.h":
+  cdef void ekf_load_and_register(string directory, string name)
+
 cdef extern from "rednose/helpers/ekf_sym.h" namespace "EKFS":
   cdef cppclass MapVectorXd "Eigen::Map<Eigen::VectorXd>":
     MapVectorXd(double*, int)
@@ -85,6 +89,7 @@ cdef class EKF_sym_pyx:
       int dim_main_err, int N=0, int dim_augment=0, int dim_augment_err=0, list maha_test_kinds=[],
       list quaternion_idxs=[], list global_vars=[], double max_rewind_age=1.0, logger=None):
     # TODO logger
+    ekf_load_and_register(gen_dir.encode('utf8'), name.encode('utf8'))
 
     cdef np.ndarray[np.float64_t, ndim=2, mode='c'] Q_b = np.ascontiguousarray(Q, dtype=np.double)
     cdef np.ndarray[np.float64_t, ndim=1, mode='c'] x_initial_b = np.ascontiguousarray(x_initial, dtype=np.double)