Passing "complicated" types as arguments

[nickbianco]

I am following up from my discussion board entry at #1870, where I have successfully created a project where (in some cases) I can take derivatives through Simbody with Enzyme. This seems promising, but now I've encountered some new roadblocks as I try to take derivatives through more complicated function calls which are more relevant for my intended use case.

Below is a minimum working example based on my sandbox project for testing Simbody derivatives with Enzyme. This example, like most of the calculations I care about, require calling functions that take a Simbody SimTK::State as an argument. SimTK::State primarily holds the values of state variables for a multibody system (e.g., coordinate values and speeds) at a given time point of a simulation. It also holds quantities that are the result of advancing the multibody system through all calculations needed to perform each time step of a simulation (e.g., coordinate accelerations, kinematic constraint errors, etc.).

In practice, will need to take the derivative of a value (e.g., kinetic energy, as in the example below) with respect to one or more of the quantities in the SimTK::State. For example, I might want dKEdu0, where u0 is the first coordinate speed in the SimTK::State: double u0 = getU()[0].

What is the best practice for taking such a derivative? I tried to follow this example on the website about using "complicated" types as arguments, but I'm not sure how Enzyme interprets cases where something a SimTK::State is passed as a shadow argument. I am able to run the example below, but only when passing LLDEnzymeLooseTypeFlags and the derivative does not match the expected analytical derivative value.

I've tried a few other things (e.g., wrapping the function call I needed with a custom class method that takes only the duplicated arguments I want), but for brevity's sake, I will stop here for now. I will keep digging, but any guidance would be appreciated!

#include <cstdio>
#include "Simbody.h"

using namespace SimTK;

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...);

Real wrapper(MultibodySystem* system, State* state) {
    system->realize(*state, Stage::Velocity);
    return system->calcKineticEnergy(*state);
}

void exampleSlidingMass() {
    // Define the system.
    MultibodySystem system;
    SimbodyMatterSubsystem matter(system);
    GeneralForceSubsystem forces(system);
    Force::Gravity gravity(forces, matter, -YAxis, 9.8);

    // Describe mass and visualization properties for a generic body.
    Real mass = 2.3;
    Body::Rigid bodyInfo(MassProperties(mass, Vec3(0), UnitInertia(0)));
    bodyInfo.addDecoration(Transform(), DecorativeSphere(0.1));

    // Create a slider mobilized body in the x-direction.
    MobilizedBody::Slider slider(matter.Ground(), Transform(Vec3(0)),
            bodyInfo, Transform(Vec3(1, 0, 0)));

    // Initialize the system and state.
    State state = system.realizeTopology();

    // Calculate the kinetic energy of the system.
    //
    // KE = 0.5 * m * u^2
    //
    double u = 1.7;
    state.updU()[0] = u;
    Real KE = wrapper(&system, &state);
    printf("KE_analytic = %f\n", 0.5 * mass * u * u);
    printf("KE_simbody = %f\n", KE); 

    // Calculate the derivative of the kinetic energy.
    //
    // dKEdu = m * u
    //
    // TODO: this compiles but does not produce the correct derivative
    //       (LLDEnzymeLooseTypeFlags to blame?)
    State dstate = state;
    dstate.updU()[0] = 1.0;
    Real dKEdu = __enzyme_fwddiff<Real>((void*)wrapper, 
            enzyme_const, &system, 
            enzyme_dup, &state, &dstate);
    printf("dKEdu_analytic = %f\n", mass * u);
    printf("dKEdu_enzyme = %f\n", dKEdu); // is this actually dKE_dstate? what does that mean?

    // Even with LLDEnzymeLooseTypeFlags enabled, the following error is thrown
    // when using __enzyme_autodiff:
    //
    // [build] ld.lld: error: <unknown>:0:0: in function preprocess__ZNK5SimTK6System4Guts7realizeERKNS_5StateENS_5StageE void (ptr, ptr, i32): Enzyme: No create nofree of unknown value
    // [build] ptr %0
    // [build]  at context:   %26 = tail call noundef i32 %25(ptr noundef nonnull align 8 dereferenceable(16) %0, ptr noundef nonnull align 8 dereferenceable(8) %1) #62 
    // State dstate = state;
    // dstate.updY() = 0.0;
    // __enzyme_autodiff<Real>((void*)wrapper, 
    //         enzyme_const, &system, 
    //         enzyme_dup, &state, &dstate);
}

int main () {

    exampleSlidingMass();

    return 0;
}