Differentiating OpenSim models with Enzyme

[nickbianco]

Hello,

I work on OpenSim, a musculoskeletal modeling and simulation library used in biomechanics research applications. (A "musculoskeletal" model is essentially a multibody dynamics model with custom actuators to represent muscles and other biological components). OpenSim includes a trajectory optimization library called Moco, and I'm hoping to use Enzyme to differentiate OpenSim models to open the door to solving optimal control problems using automatic differentiation through CasADi.

OpenSim is built on top of a separate library, Simbody, a Featherstone-like multibody dynamics engine. Therefore I'm investigating the feasibility of using LLDEnzyme to differentiate functions that bridge the OpenSim and Simbody libraries. I'm also trying to understand if it's possible to differentiate functions with custom data types from OpenSim/Simbody. To test these, I've written the following example:

#include <cstdio>
#include <OpenSim/Common/PolynomialFunction.h>
#include <OpenSim/Actuators/DeGrooteFregly2016Muscle.h>

using namespace OpenSim;

extern int enzyme_dup;
extern int enzyme_dupnoneed;
extern int enzyme_out;
extern int enzyme_const;

template < typename return_type, typename ... T >
return_type __enzyme_fwddiff(void*, T ... );

template < typename return_type, typename ... T >
return_type __enzyme_autodiff(void*, T ... );

template < typename T, typename ... arg_types >
auto wrapperMuscle(T obj, arg_types && ... args) {
    return obj.calcTendonForceMultiplier(args ... );
}

void testDeGrooteFreglyMuscle() {
    DeGrooteFregly2016Muscle muscle;
    muscle.finalizeFromProperties();

    SimTK::Real y = 1.1;
    SimTK::Real dy = 1.0;
    SimTK::Real f = muscle.calcTendonForceMultiplier(y);
    SimTK::Real dfdy = muscle.calcTendonForceMultiplierDerivative(y);
    printf("Muscle analytic derivative\n");
    printf("f(y) = %f, f'(y) = %f\n\n", f, dfdy);

    dfdy = __enzyme_fwddiff<SimTK::Real>(
            (void*)wrapperMuscle<DeGrooteFregly2016Muscle, SimTK::Real>, 
            enzyme_const, muscle, 
            enzyme_dup, &y, &dy);
    printf("Muscle derivative w/ Enzyme\n");
    printf("f(y) = %f, f'(y) = %f\n\n", wrapperMuscle(muscle, y), dfdy);
}

template < typename T, typename ... arg_types >
auto wrapperPolynomial(T obj, arg_types && ... args) {
    return obj.calcValue(args ... );
}

void testPolynomialFunction() {
    // f = x^2
    SimTK::Vector coefficients(3);
    coefficients[0] = 1;
    coefficients[1] = 0;
    coefficients[2] = 0;
    PolynomialFunction func(coefficients);

    SimTK::Vector y(1, 5.0);
    SimTK::Vector dy(1, 1.0);
    SimTK::Real dfdy = __enzyme_fwddiff<SimTK::Real>(
            (void*)wrapperPolynomial<PolynomialFunction, SimTK::Vector>, 
            enzyme_const, func, 
            enzyme_dup, &y, &dy);
    printf("PolynomialFunction\n");
    printf("f(y) = %f, f'(y) = %f\n\n", wrapperPolynomial(func, y), dfdy);
}

int main() {
    testDeGrooteFreglyMuscle();
    testPolynomialFunction();
    return 0;
}

testDeGrooteFreglyMuscle() tests differentiating calcTendonForceMultiplier(), a member function of the class DeGrooteFregly2016Muscle which models a muscle-like actuator. SimTK::Real is a basic scalar type provided by Simbody. The definition for calcTendonForceMultiplier() lives in the header for DeGrooteFregly2016Muscle and only requires evaluating analytical expressions based on member variables of the muscle class. We happen to have the analytical derivative of this function in calcTendonForceMultiplierDerivative(), and comparing shows that we are getting the correct derivative from Enzyme:

Muscle analytic derivative
f(y) = 7.546290, f'(y) = 283.254880

Muscle derivative w/ Enzyme
f(y) = 7.546290, f'(y) = 283.254880

The second case, testPolynomialFunction(), is where things get more complicated. Here, I implement a PolynomialFunction object representing the expression f = x^2. Since Functions in OpenSim can depend on multiple arguments, the interface to evaluate functions, calcValue(), requires a vector input. The input type here is SimTK::Vector, a vector-type provided by Simbody. Under the hood, PolynomialFunction is just a thin wrapper around SimTK::Function::Polynomial, a polynomial function class also provided by Simbody. Building the example with testPolynomialFunction() include results in the following error during linking:

[build] [100%] Linking CXX executable ../../../enzyme_test_2
[build] ld.lld: error: OpenSim/Sandbox/Enzyme/enzyme_test.cpp:44:16: in function preprocess__Z17wrapperPolynomialIN7OpenSim18PolynomialFunctionEJN5SimTK7Vector_IdEEEEDaT_DpOT0_ double (ptr, ptr): Enzyme: No forward mode derivative found for _ZNK7OpenSim8Function9calcValueERKN5SimTK7Vector_IdEE
[build]  at context:   %3 = tail call noundef double @_ZNK7OpenSim8Function9calcValueERKN5SimTK7Vector_IdEE(ptr noundef nonnull align 8 dereferenceable(296) %0, ptr noundef nonnull align 8 dereferenceable(8) %1) #34, !dbg !30674
[build] 
[build] clang-16: error: linker command failed with exit code 1 (use -v to see invocation)

It looks like the name mangling of PolynomialFunction::calcValue() here is preventing the forward mode derivative, but I'm not totally sure. I also thought that I had not enabled LTO, but I pass LLDEnzymeFlags to my executable target via target_link_libraries which (I think) enables LTO. My next though to try generating the LLVM IR directly and taking a look at that, but any suggestions on how to focus my debugging effort here would be appreciated!