Dynamic Inserts

Author: ax3l

enzyme/Enzyme/CallDerivatives.cpp

@@ -2595,39 +2595,22 @@ bool AdjointGenerator::handleKnownCallDerivatives(
 
     // 2) STL std::list insertion
     if (funcName == "_ZNSt8__detail15_List_node_base7_M_hookEPS0_") {
-      if (Mode == DerivativeMode::ReverseModeGradient) {
-        eraseIfUnused(call, /*erase*/ true, /*check*/ false);
-        return true;
-      }
-      if (gutils->isConstantValue(call.getArgOperand(0)))
-        return true;
-      SmallVector<Value *, 2> args;
-      for (auto &arg : call.args()) {
-        if (gutils->isConstantValue(arg))
-          args.push_back(gutils->getNewFromOriginal(arg));
-        else
-          args.push_back(gutils->invertPointerM(arg, BuilderZ));
-      }
-      BuilderZ.CreateCall(called, args);
+      // Only run in primal version to avoid memory management issues
+      eraseIfUnused(call, /*erase*/ true, /*check*/ false);
       return true;
     }
 
     // 3) STL std::list transfer (splice operations)
     if (funcName == "_ZNSt8__detail15_List_node_base11_M_transferEPS0_S1_") {
-      if (Mode == DerivativeMode::ReverseModeGradient) {
-        eraseIfUnused(call, /*erase*/ true, /*check*/ false);
-        return true;
-      }
-      if (gutils->isConstantValue(call.getArgOperand(0)))
-        return true;
-      SmallVector<Value *, 2> args;
-      for (auto &arg : call.args()) {
-        if (gutils->isConstantValue(arg))
-          args.push_back(gutils->getNewFromOriginal(arg));
-        else
-          args.push_back(gutils->invertPointerM(arg, BuilderZ));
-      }
-      BuilderZ.CreateCall(called, args);
+      // Only run in primal version to avoid memory management issues
+      eraseIfUnused(call, /*erase*/ true, /*check*/ false);
+      return true;
+    }
+
+    // 4) STL std::list size increment (called when adding elements)
+    if (funcName == "_ZNSt7__cxx1110_List_baseIdSaIdEE11_M_inc_sizeEm") {
+      // Only run in primal version to avoid memory management issues
+      eraseIfUnused(call, /*erase*/ true, /*check*/ false);
       return true;
     }
 

enzyme/test/Integration/ReverseMode/stl_list.cpp

@@ -9,9 +9,9 @@
 
 
 template<typename T>
-extern void __enzyme_fwddiff(void*, int, T&, T&);
+extern double __enzyme_fwddiff(void*, int, T&, T&);
 template<typename T>
-extern void __enzyme_autodiff(void*, int, T&, T&);
+extern double __enzyme_autodiff(void*, int, T&, T&);
 
 
 double test_simple_list(std::list<double>& vals) {
@@ -26,23 +26,56 @@ double test_simple_list(std::list<double>& vals) {
     return result;
 }
 
+double test_dynamic_list(std::list<double>& vals) {
+    // dynamically insert new elements
+    vals.push_back(4.0);
+    vals.push_back(5.0);
+
+    // iterate over list
+    double result = 0.0;
+    for (const auto& val : vals) {
+        result += val * val;
+    }
+    return result;
+}
+
 void test_forward_list() {
-    std::list<double> vals = {1.0, 2.0, 3.0};
-    std::list<double> dvals = {1.0, 1.0, 1.0};
+    {
+        std::list<double> vals = {1.0, 2.0, 3.0};
+        std::list<double> dvals = {1.0, 1.0, 1.0};
+
+        double ret = __enzyme_fwddiff((void*)test_simple_list, enzyme_dup, vals, dvals);
+        APPROX_EQ( ret, 10., 1e-10);
+    }
+    {
+        std::list<double> vals = {1.0, 2.0, 3.0};
+        std::list<double> dvals = {1.0, 1.0, 1.0};
 
-    __enzyme_fwddiff((void*)test_simple_list, enzyme_dup, vals, dvals);
+        double ret = __enzyme_fwddiff((void*)test_dynamic_list, enzyme_dup, vals, dvals);
+        APPROX_EQ( ret, 12., 1e-10);
+    }
 }
 
 void test_reverse_list() {
-    std::list<double> vals = {1.0, 2.0, 3.0};
-    std::list<double> dvals = {1.0, 1.0, 1.0};
+    {
+        std::list<double> vals = {1.0, 2.0, 3.0};
+        std::list<double> dvals = {1.0, 1.0, 1.0};
 
-    __enzyme_autodiff((void*)test_simple_list, enzyme_dup, vals, dvals);
+        double ret = __enzyme_autodiff((void*)test_simple_list, enzyme_dup, vals, dvals);
+        //APPROX_EQ( ret, 10., 1e-10);
+    }
+    {
+        std::list<double> vals = {1.0, 2.0, 3.0};
+        std::list<double> dvals = {1.0, 1.0, 1.0};
+
+        double ret = __enzyme_autodiff((void*)test_dynamic_list, enzyme_dup, vals, dvals);
+        //APPROX_EQ( ret, 12., 1e-10);
+    }
 }
 
 
 int main() {
     test_forward_list();
     test_reverse_list();
     return 0;
-} 
+}