Support more Julia 1.11 functions (#2121)

* Support more Julia 1.11 functions * Consider allocating functions
EnzymeAD · Oct 17, 2024 · f085ac8 · f085ac8
1 parent 8416a68
commit f085ac8
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 30 deletions.
diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp
@@ -1708,7 +1708,9 @@ bool ActivityAnalyzer::isConstantValue(TypeResults const &TR, Value *Val) {
         }
         if (funcName == "jl_array_copy" || funcName == "ijl_array_copy" ||
             funcName == "jl_idtable_rehash" ||
-            funcName == "ijl_idtable_rehash") {
+            funcName == "ijl_idtable_rehash" ||
+            funcName == "jl_genericmemory_copy_slice" ||
+            funcName == "ijl_genericmemory_copy_slice") {
           // This pointer is inactive if it is either not actively stored to
           // and not actively loaded from and the copied input is inactive.
           if (directions & DOWN && directions & UP) {

diff --git a/enzyme/Enzyme/CallDerivatives.cpp b/enzyme/Enzyme/CallDerivatives.cpp
@@ -3386,15 +3386,8 @@ bool AdjointGenerator::handleKnownCallDerivatives(
           }
 
           // No need to free GC.
-          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" ||
-              funcName == "jl_gc_alloc_typed" ||
-              funcName == "ijl_gc_alloc_typed") {
+          if (EnzymeJuliaAddrLoad && isa<PointerType>(call.getType()) &&
+              cast<PointerType>(call.getType())->getAddressSpace() == 10) {
             if (Mode == DerivativeMode::ReverseModeGradient && AllocationLoop)
               gutils->rematerializedPrimalOrShadowAllocations.push_back(
                   newCall);
@@ -3459,12 +3452,8 @@ bool AdjointGenerator::handleKnownCallDerivatives(
 
     // If an object is managed by the GC do not preserve it for later free,
     // 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" || funcName == "jl_gc_alloc_typed" ||
-        funcName == "ijl_gc_alloc_typed") {
+    if (EnzymeJuliaAddrLoad && isa<PointerType>(call.getType()) &&
+        cast<PointerType>(call.getType())->getAddressSpace() == 10) {
       if (!subretused) {
         eraseIfUnused(call, /*erase*/ true, /*check*/ false);
         return true;

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -9267,15 +9267,10 @@ void GradientUtils::computeGuaranteedFrees() {
         if (hasMetadata(CI, "enzyme_fromstack")) {
           allocationsWithGuaranteedFree[CI].insert(CI);
         }
-        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" ||
-            funcName == "jl_gc_alloc_typed" ||
-            funcName == "ijl_gc_alloc_typed") {
+        // TODO: special case object managed by the GC as it is automatically
+        // freed.
+        if (EnzymeJuliaAddrLoad && isa<PointerType>(CI->getType()) &&
+            cast<PointerType>(CI->getType())->getAddressSpace() == 10) {
         }
       }
     }

diff --git a/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp b/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp
@@ -1826,7 +1826,11 @@ bool mlir::enzyme::ActivityAnalyzer::isConstantValue(MTypeResults const &TR,
         //     }
         //   }
         // }
-        if (funcName == "jl_array_copy" || funcName == "ijl_array_copy") {
+        if (funcName == "jl_array_copy" || funcName == "ijl_array_copy" ||
+            funcName == "jl_idtable_rehash" ||
+            funcName == "ijl_idtable_rehash" ||
+            funcName == "jl_genericmemory_copy_slice" ||
+            funcName == "ijl_genericmemory_copy_slice") {
           // This pointer is inactive if it is either not actively stored to
           // and not actively loaded from.
           if (directions & DOWN && directions & UP) {

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
@@ -1346,10 +1346,19 @@ void TypeAnalyzer::considerTBAA() {
           updateAnalysis(call->getOperand(0), TT.Only(-1, call), call);
         }
         if (F) {
-          StringSet<> JuliaKnownTypes = {
-              "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", "jl_gc_alloc_typed",  "ijl_gc_alloc_typed"};
+          StringSet<> JuliaKnownTypes = {"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",
+                                         "jl_gc_alloc_typed",
+                                         "ijl_gc_alloc_typed",
+                                         "jl_alloc_genericmemory",
+                                         "ijl_alloc_genericmemory",
+                                         "jl_new_array",
+                                         "ijl_new_array"};
           if (JuliaKnownTypes.count(F->getName())) {
             visitCallBase(*call);
             continue;

diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp
@@ -2331,6 +2331,10 @@ bool writesToMemoryReadBy(const TypeResults *TR, llvm::AAResults &AA,
     if (funcName == "jl_array_copy" || funcName == "ijl_array_copy")
       return false;
 
+    if (funcName == "jl_genericmemory_copy_slice" ||
+        funcName == "ijl_genericmemory_copy_slice")
+      return false;
+
     if (funcName == "jl_new_array" || funcName == "ijl_new_array")
       return false;
 
@@ -2494,6 +2498,10 @@ bool writesToMemoryReadBy(const TypeResults *TR, llvm::AAResults &AA,
     if (funcName == "jl_array_copy" || funcName == "ijl_array_copy")
       return false;
 
+    if (funcName == "jl_genericmemory_copy_slice" ||
+        funcName == "ijl_genericmemory_copy_slice")
+      return false;
+
     if (funcName == "jl_idtable_rehash" || funcName == "ijl_idtable_rehash")
       return false;