Add ijl allocations (rust-lang#560)

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -9574,7 +9574,11 @@ class AdjointGenerator
             }
 
             // No need to free GC.
-            if (funcName == "jl_alloc_array_1d" ||
+            if (funcName == "ijl_alloc_array_1d" ||
+                funcName == "ijl_alloc_array_2d" ||
+                funcName == "ijl_alloc_array_3d" ||
+                funcName == "ijl_array_copy" ||
+                funcName == "jl_alloc_array_1d" ||
                 funcName == "jl_alloc_array_2d" ||
                 funcName == "jl_alloc_array_3d" ||
                 funcName == "jl_array_copy" || funcName == "julia.gc_alloc_obj")
@@ -9669,6 +9673,9 @@ class AdjointGenerator
       // Thus it only needs caching if there is a need for it in the reverse.
       if (funcName == "jl_alloc_array_1d" || funcName == "jl_alloc_array_2d" ||
           funcName == "jl_alloc_array_3d" || funcName == "jl_array_copy" ||
+          funcName == "ijl_alloc_array_1d" ||
+          funcName == "ijl_alloc_array_2d" ||
+          funcName == "ijl_alloc_array_3d" || funcName == "ijl_array_copy" ||
           funcName == "julia.gc_alloc_obj") {
         if (!primalNeededInReverse) {
           if (Mode == DerivativeMode::ReverseModeGradient ||

enzyme/Enzyme/GradientUtils.h

@@ -627,6 +627,10 @@ class GradientUtils : public CacheUtility {
               funcName == "jl_alloc_array_1d" ||
               funcName == "jl_alloc_array_2d" ||
               funcName == "jl_alloc_array_3d" || funcName == "jl_array_copy" ||
+              funcName == "ijl_alloc_array_1d" ||
+              funcName == "ijl_alloc_array_2d" ||
+              funcName == "ijl_alloc_array_3d" ||
+              funcName == "ijl_array_copy" ||
               funcName == "julia.gc_alloc_obj") {
             primalInitializationOfShadow = true;
           } else {
@@ -758,6 +762,10 @@ class GradientUtils : public CacheUtility {
                   funcName == "jl_alloc_array_2d" ||
                   funcName == "jl_alloc_array_3d" ||
                   funcName == "jl_array_copy" ||
+                  funcName == "ijl_alloc_array_1d" ||
+                  funcName == "ijl_alloc_array_2d" ||
+                  funcName == "ijl_alloc_array_3d" ||
+                  funcName == "ijl_array_copy" ||
                   funcName == "julia.gc_alloc_obj") {
                 primalInitializationOfShadow = true;
               } else {
@@ -876,6 +884,10 @@ class GradientUtils : public CacheUtility {
           if (funcName == "jl_alloc_array_1d" ||
               funcName == "jl_alloc_array_2d" ||
               funcName == "jl_alloc_array_3d" || funcName == "jl_array_copy" ||
+              funcName == "ijl_alloc_array_1d" ||
+              funcName == "ijl_alloc_array_2d" ||
+              funcName == "ijl_alloc_array_3d" ||
+              funcName == "ijl_array_copy" ||
               funcName == "julia.gc_alloc_obj") {
           }
         }

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -804,10 +804,9 @@ void TypeAnalyzer::considerTBAA() {
         }
         if (F) {
           std::set<std::string> JuliaKnownTypes = {
-              "julia.gc_alloc_obj",
-              "jl_alloc_array_1d",
-              "jl_alloc_array_2d",
-              "jl_alloc_array_3d",
+              "julia.gc_alloc_obj", "jl_alloc_array_1d",  "jl_alloc_array_2d",
+              "jl_alloc_array_3d",  "ijl_alloc_array_1d", "ijl_alloc_array_2d",
+              "ijl_alloc_array_3d",
           };
           if (JuliaKnownTypes.count(F->getName().str())) {
             visitCallInst(*call);