Add ijl allocations (rust-lang#560)

EnzymeAD · Mar 12, 2022 · bd9e2b3 · bd9e2b3
1 parent f951148
commit bd9e2b3
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/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 ||

diff --git a/enzyme/Enzyme/GradientUtils.h b/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") {
           }
         }

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp b/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);