Added another example passing parameters to function

Author: meister

tlm3_dihedral/Dihed.cc

@@ -76,7 +76,7 @@ struct Dihedral {
  extern double __enzyme_autodiff(void*, ...);
 
 void  dump_vec(const std::string& msg, float* deriv, size_t num) {
-  std::cout << std::setw(10) << msg << " ";
+  std::cout << std::setw(20) << msg << " ";
   for ( int i=0; i<num; i++ ) {
     std::cout << std::setprecision(4) << std::fixed << std::setw(8) << deriv[i] << " ";
   }
@@ -87,6 +87,74 @@ void  dump_vec(const std::string& msg, float* deriv, size_t num) {
 #define mysqrt(VAR) (std::sqrt(VAR))
 #define reciprocal(VAR) (1.0/VAR)
 
+double old_simple_dihedral_energy_gradient( double sinPhase, double cosPhase, double V, double DN, int IN, int I1, int I2, int I3, int I4, float* pos, float* deriv ) {
+#define USE_EXPLICIT_DECLARES 1
+#define DECLARE_FLOAT(VAR) double VAR;
+#define Real double
+  double EraseLinearDihedral;
+  double x1, y1, z1;
+  double x2, y2, z2;
+  double x3, y3, z3;
+  double x4, y4, z4;
+#include "_Dihedral_termDeclares.cc"
+#define DIHEDRAL_SET_PARAMETER(VAR) {}
+#define DIHEDRAL_SET_POSITION(VAR,IDX,OFFSET) { VAR = pos[IDX+OFFSET]; }
+#define DIHEDRAL_ENERGY_ACCUMULATE(VAR) { result += VAR; }
+#define DIHEDRAL_FORCE_ACCUMULATE(IDX,OFFSET,VAR) { deriv[IDX+OFFSET] -= VAR; }
+  double result = 0.0;
+  double SinNPhi;
+  double CosNPhi;
+  bool calcForce = true;
+#define DIHEDRAL_CALC_FORCE 1
+#include "_Dihedral_termCode.cc"
+#undef Real
+#undef DECLARE_FLOAT
+#undef DIHEDRAL_SET_PARAMETER
+#undef DIHEDRAL_SET_POSITION
+#undef DIHEDRAL_ENERGY_ACCUMULATE
+#undef DIHEDRAL_FORCE_ACCUMULATE
+  return result;
+}
+
+double simple_dihedral_energy( double sinPhase, double cosPhase, double V, double DN, int IN, int I1, int I2, int I3, int I4, float* pos ) {
+#define USE_EXPLICIT_DECLARES 1
+#define DECLARE_FLOAT(VAR) double VAR;
+#define Real double
+  double EraseLinearDihedral;
+  double x1, y1, z1;
+  double x2, y2, z2;
+  double x3, y3, z3;
+  double x4, y4, z4;
+#include "_DihedralEnergy_termDeclares.cc"
+#define DIHEDRAL_SET_PARAMETER(VAR) {}
+#define DIHEDRAL_SET_POSITION(VAR,IDX,OFFSET) { VAR = pos[IDX+OFFSET]; }
+#define DIHEDRAL_ENERGY_ACCUMULATE(VAR) { result += VAR; }
+  double result = 0.0;
+  double SinNPhi;
+  double CosNPhi;
+#include "_DihedralEnergy_termCode.cc"
+#undef Real
+#undef DECLARE_FLOAT
+#undef DIHEDRAL_SET_PARAMETER
+#undef DIHEDRAL_SET_POSITION
+#undef DIHEDRAL_ENERGY_ACCUMULATE
+  return result;
+}
+
+void simple_dihedral_gradient( double sinPhase, double cosPhase, double V, double DN, int IN, int I1, int I2, int I3, int I4, float* pos, float* grad ) {
+  __enzyme_autodiff( (void*)simple_dihedral_energy,
+                     enzyme_const, sinPhase,
+                     enzyme_const, cosPhase,
+                     enzyme_const, V,
+                     enzyme_const, DN,
+                     enzyme_const, IN,
+                     enzyme_const, I1,
+                     enzyme_const, I2,
+                     enzyme_const, I3,
+                     enzyme_const, I4,
+                     enzyme_dup, pos, grad );
+}
+
 template <typename Real>
 float dihedral_energy(Dihedral* dihedral_begin, Dihedral* dihedral_end, float* pos) {
 #define USE_EXPLICIT_DECLARES 1
@@ -123,7 +191,7 @@ float dihedral_energy(Dihedral* dihedral_begin, Dihedral* dihedral_end, float* p
 }
 
 template <typename Real>
-void grad_dihedral_energy(Dihedral* dihedral_begin, Dihedral* dihedral_end, float* pos, float* deriv ) {
+__attribute__((noinline)) void grad_dihedral_energy(Dihedral* dihedral_begin, Dihedral* dihedral_end, float* pos, float* deriv ) {
   __enzyme_autodiff( (void*)dihedral_energy<Real>,
                      enzyme_const, dihedral_begin,
                      enzyme_const, dihedral_end,
@@ -174,7 +242,7 @@ Real old_dihedral_energy(Dihedral* dihedral_begin, Dihedral*dihedral_end, float*
 #define DIHEDRAL_SET_PARAMETER(VAR) { VAR = dihedral->VAR; }
 #define DIHEDRAL_SET_POSITION(VAR,IDX,OFFSET) { VAR = pos[IDX+OFFSET]; }
 #define DIHEDRAL_ENERGY_ACCUMULATE(VAR) { result += VAR; }
-#define DIHEDRAL_FORCE_ACCUMULATE(IDX,OFFSET,VAR) { deriv[IDX+OFFSET] += VAR; }
+#define DIHEDRAL_FORCE_ACCUMULATE(IDX,OFFSET,VAR) { deriv[IDX+OFFSET] -= VAR; }
   int I1;
   int I2;
   int I3;
@@ -228,18 +296,21 @@ int main( int argc, const char* argv[] ) {
       zeroVec(deriv,12);
       grad_dihedral_energy<float>( &dihedral[0], &dihedral[1], pos, deriv);
     }
+    dump_vec("new float deriv", deriv, 12);
   } else if (arg1 == "new-double") {
     std::cout << "New method" << "\n";
     // energy = dihedral_energy( &dihedral[0], &dihedral[1], pos );
     for ( size_t nn = 0; nn<num; nn++ ) {
       zeroVec(deriv,12);
       grad_dihedral_energy<double>( &dihedral[0], &dihedral[1], pos, deriv);
     }
+    dump_vec("new double deriv", deriv, 12);
   } else if ( arg1 == "fin") {
       for ( size_t nn = 0; nn<num; nn++ ) {
         zeroVec(deriv,12);
         finite_diff_dihedral_energy(&dihedral[0], &dihedral[1], pos, deriv);
       }
+    dump_vec("fin deriv", deriv, 12);
   } else if (arg1 == "old-float") {
     float fenergy;
     float fderiv[12];
@@ -264,5 +335,20 @@ int main( int argc, const char* argv[] ) {
     }
     std::cout << "Energy = " << fenergy << "\n";
     dump_vec("deriv", fderiv, 12);
+  } else if (arg1 == "new-simple") {
+    float fderiv[12];
+    for ( size_t nn = 0; nn<num; nn++ ) {
+      zeroVec(fderiv,12);
+      simple_dihedral_gradient( 0.0, 1.0, 10.0, 2.0, 2, 0, 3, 6, 9, pos, fderiv );
+    }
+    dump_vec("simple deriv via Enzyme", fderiv, 12 );
+  } else if (arg1 == "old-simple") {
+    float fderiv[12];
+    for ( size_t nn = 0; nn<num; nn++ ) {
+      zeroVec(fderiv,12);
+      old_simple_dihedral_energy_gradient( 0.0, 1.0, 10.0, 2.0, 2, 0, 3, 6, 9, pos, fderiv );
+    }
+    dump_vec("simple deriv old code  ", fderiv, 12 );
   }
+
 }