diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index b079e9fd8ae18..f31e621633d3e 100644 --- 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 index 53bb00f5d0291..36fdba8cd5eb1 100644 --- 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 index 3a297ed04ed80..682bc282fde80 100644 --- a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp +++ b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp @@ -804,10 +804,9 @@ void TypeAnalyzer::considerTBAA() { } if (F) { std::set 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);