Differentiating across multiple libraries?

[Nick Bianco]

Hello,

I am developer for OpenSim, a musculoskeletal modeling and simulation software project, and I am interested in using Enzyme to support automatic differentiation in trajectory optimization tools we develop. OpenSim depends on Simbody, a multibody dynamics library, and essentially all functions that need be differentiated will require differentiating through Simbody function calls. 

I made a lengthier post about this on the Enzyme discussion board, but posting here to ask more directly my main question: is it possible to differentiate across multiple libraries? I would assume that it is not possible to differentiate through a pre-compiled library, but I'm wondering if it's somehow possible to perform AD through other means. For example, would it be possible apply Enzyme to dependent functions in Simbody first, recompile Simbody, and then use Enzyme again in OpenSim to get the derivatives we need. 

Any thoughts or suggestions would be greatly appreciated! If this is something that is worth discussing in the open developer calls, I would be happy to join those as well.

Best,

Nick

[William Moses]

Hey Nick,

Yes please join the next design meetings (which we're presently in the process of rescheduling), and apologies for not noticing the post.

It is indeed possible to differentiate through multiple libraries and there's a couple of different ways. You could compile the library with either LTO or embedded bitcode (which ensures that future compilers can access the necessary information to differentiate), you could tell Enzyme export derivatives for various functions which may be needed to differentiate, among other options.

Maybe it's worth coming to the next design meeting and chatting through constraints/pros/cons of different methods?

--
You received this message because you are subscribed to the Google Groups "Enzyme AD" group.
To unsubscribe from this group and stop receiving emails from it, send an email to enzyme-dev+...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/enzyme-dev/c4eb73de-9663-4496-97f2-691a57a34f62n%40googlegroups.com.

[Nick Bianco]

Hi William,

That sounds great, I will definitely join the next design meeting. Great to hear that this is possible!

Best,

Nick

[Nick Bianco]

Hi William,

I gave your first suggestion a shot and compiled OpenSim's dependencies using LTO (e.g., added "-flto" to CMAKE_C_FLAGS and CMAKE_CXX_FLAGS and replaced my linker with lld). I ran the following example:

#include <cstdio>
#include <OpenSim/Common/PolynomialFunction.h>
#include <OpenSim/Actuators/DeGrooteFregly2016Muscle.h>
#include <OpenSim/Common/Sine.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);
}

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

However, I run into the following linker error:

[build] [100%] Linking CXX executable ../../../sandboxEnzyme
[build] ld.lld: error: OpenSim/Common/PropertyTable.h:105:40: in function preprocess__Z13wrapperMuscleIN7OpenSim24DeGrooteFregly2016MuscleEJdEEDaT_DpOT0_ double (ptr, ptr): Enzyme: No forward mode derivative found for _ZNK7OpenSim13PropertyTable26getAbstractPropertyByIndexEi
[build]  at context:   %6 = tail call noundef nonnull align 8 dereferenceable(84) ptr @_ZNK7OpenSim13PropertyTable26getAbstractPropertyByIndexEi(ptr noundef nonnull align 8 dereferenceable(80) %4, i32 noundef %5) #39, !dbg !26244
[build]

This example requires derefencing a custom smart pointer implemented by the dependency Simbody, which I believe explains the "ptr noundef nonnull align 8 dereferenceable(80)" message. (The pointer is used to store "property" values in OpenSim, which are serializable static parameters in OpenSim models, via the "AbstractProperty" class.)

I'm planning to join the ODM tomorrow to discuss the options you suggested above, but figured I'd post my progress here beforehand.

Best,

Nick

[Nick Bianco]

Hi Billy,

Thanks for helping me get started debugging this issue on the Zoom call today! Following your suggestions, I started by trying to generate LLVM IR and looking at that directly (based on the instructions at https://enzyme.mit.edu/getting_started/UsingEnzyme/).  The "input.ll" file with the emitted LLVM IR is quite large, but I find the mangled method names where the pointer deference happens:

invoke.cont:                                      ; preds = %entry
  call void @llvm.lifetime.start.p0(i64 8, ptr nonnull %y) #29
  store double 1.100000e+00, ptr %y, align 8, !tbaa !78
  call void @llvm.lifetime.start.p0(i64 8, ptr nonnull %dy) #29
  store double 1.000000e+00, ptr %dy, align 8, !tbaa !78
  %PropertyIndex_tendon_strain_at_one_norm_force.i.i.i.i = getelementptr inbounds %"class.OpenSim::DeGrooteFregly2016Muscle", ptr %muscle, i64 0, i32 9
  %_propertyTable.i.i.i.i.i = getelementptr inbounds %"class.OpenSim::Object", ptr %muscle, i64 0, i32 6
  %0 = load i32, ptr %PropertyIndex_tendon_strain_at_one_norm_force.i.i.i.i, align 8, !tbaa !82
  %call.i.i.i.i.i.i34 = invoke noundef nonnull align 8 dereferenceable(84) ptr @_ZNK7OpenSim13PropertyTable26getAbstractPropertyByIndexEi(ptr noundef nonnull align 8 dereferenceable(80) %_propertyTable.i.i.i.i.i, i32 noundef %0)
          to label %call.i.i.i.i.i.i.noexc unwind label %lpad1

Applying AD with the LLVMEnzyme plugin (and adding the -enzyme-print) flag yields the expected error, along with some extra information I'm not sure how to parse yet:

prefn:
; Function Attrs: mustprogress uwtable
define linkonce_odr dso_local noundef double @_Z13wrapperMuscleIN7OpenSim24DeGrooteFregly2016MuscleEJdEEDaT_DpOT0_(ptr noundef %obj, ptr noundef nonnull align 8 dereferenceable(8) %args) #5 comdat {
entry:
  %PropertyIndex_tendon_strain_at_one_norm_force.i.i.i.i = getelementptr inbounds %"class.OpenSim::DeGrooteFregly2016Muscle", ptr %obj, i64 0, i32 9
  %_propertyTable.i.i.i.i.i = getelementptr inbounds %"class.OpenSim::Object", ptr %obj, i64 0, i32 6
  %0 = load i32, ptr %PropertyIndex_tendon_strain_at_one_norm_force.i.i.i.i, align 4, !tbaa !82
  %call.i.i.i.i.i.i = tail call noundef nonnull align 8 dereferenceable(84) ptr @_ZNK7OpenSim13PropertyTable26getAbstractPropertyByIndexEi(ptr noundef nonnull align 8 dereferenceable(80) %_propertyTable.i.i.i.i.i, i32 noundef %0)
  %call2.i.i.i.i.i.i = tail call noundef nonnull align 8 dereferenceable(84) ptr @_ZN7OpenSim8PropertyIdE5getAsERKNS_16AbstractPropertyE(ptr noundef nonnull align 8 dereferenceable(84) %call.i.i.i.i.i.i)
  %call2.i.i.i = tail call noundef nonnull align 8 dereferenceable(8) ptr @_ZNK7OpenSim8PropertyIdE8getValueEi(ptr noundef nonnull align 8 dereferenceable(84) %call2.i.i.i.i.i.i, i32 noundef -1)
  %1 = load double, ptr %call2.i.i.i, align 8, !tbaa !78
  %add.i.i = fadd double %1, 1.000000e+00
  %sub.i.i = fadd double %add.i.i, -1.000000e+00
  %div.i.i = fdiv double 0x3FFCAB0BFA2A2002, %sub.i.i
  %2 = load double, ptr %args, align 8, !tbaa !78
  %sub.i = fadd double %2, -1.000000e+00
  %mul.i = fmul double %sub.i, %div.i.i
  %call2.i = tail call double @exp(double noundef %mul.i) #30
  %3 = tail call noundef double @llvm.fmuladd.f64(double %call2.i, double 2.000000e-01, double -2.000000e-01)
  ret double %3
}

after simplification :
; Function Attrs: mustprogress willreturn uwtable
define linkonce_odr dso_local noundef double @preprocess__Z13wrapperMuscleIN7OpenSim24DeGrooteFregly2016MuscleEJdEEDaT_DpOT0_(ptr noundef %obj, ptr noundef nonnull align 8 dereferenceable(8) %args) #30 {
entry:
  %PropertyIndex_tendon_strain_at_one_norm_force.i.i.i.i = getelementptr inbounds %"class.OpenSim::DeGrooteFregly2016Muscle", ptr %obj, i64 0, i32 9
  %_propertyTable.i.i.i.i.i = getelementptr inbounds %"class.OpenSim::Object", ptr %obj, i64 0, i32 6
  %0 = load i32, ptr %PropertyIndex_tendon_strain_at_one_norm_force.i.i.i.i, align 4, !tbaa !82
  %call.i.i.i.i.i.i = tail call noundef nonnull align 8 dereferenceable(84) ptr @_ZNK7OpenSim13PropertyTable26getAbstractPropertyByIndexEi(ptr noundef nonnull align 8 dereferenceable(80) %_propertyTable.i.i.i.i.i, i32 noundef %0) #31
  %call2.i.i.i.i.i.i = tail call noundef nonnull align 8 dereferenceable(84) ptr @_ZN7OpenSim8PropertyIdE5getAsERKNS_16AbstractPropertyE(ptr noundef nonnull align 8 dereferenceable(84) %call.i.i.i.i.i.i) #31
  %call2.i.i.i = tail call noundef nonnull align 8 dereferenceable(8) ptr @_ZNK7OpenSim8PropertyIdE8getValueEi(ptr noundef nonnull align 8 dereferenceable(84) %call2.i.i.i.i.i.i, i32 noundef -1) #31
  %1 = load double, ptr %call2.i.i.i, align 8, !tbaa !78
  %add.i.i = fadd double %1, 1.000000e+00
  %sub.i.i = fadd double %add.i.i, -1.000000e+00
  %div.i.i = fdiv double 0x3FFCAB0BFA2A2002, %sub.i.i
  %2 = load double, ptr %args, align 8, !tbaa !78
  %sub.i = fadd double %2, -1.000000e+00
  %mul.i = fmul double %sub.i, %div.i.i
  %call2.i = tail call double @exp(double noundef %mul.i) #32
  %3 = tail call noundef double @llvm.fmuladd.f64(double %call2.i, double 2.000000e-01, double -2.000000e-01) #31
  ret double %3
}

error: <unknown>:0:0: in function preprocess__Z13wrapperMuscleIN7OpenSim24DeGrooteFregly2016MuscleEJdEEDaT_DpOT0_ double (ptr, ptr): Enzyme: No forward mode derivative found for _ZNK7OpenSim13PropertyTable26getAbstractPropertyByIndexEi
 at context:   %call.i.i.i.i.i.i = tail call noundef nonnull align 8 dereferenceable(84) ptr @_ZNK7OpenSim13PropertyTable26getAbstractPropertyByIndexEi(ptr noundef nonnull align 8 dereferenceable(80) %_propertyTable.i.i.i.i.i, i32 noundef %0) #31

I believe you had some other suggestions, but I don't think I caught everything on the call. I would be happy to jump on a quick call to parse through the rest of the Enzyme AD process for this example.

Best,

Nick