cleanup

Author: vchuravy

jlpkg/Enzyme/Project.toml

@@ -8,3 +8,13 @@ BinaryProvider = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 MCAnalyzer = "a81df072-f4bb-11e8-03d3-cfaeda626d18"
+
+[compat]
+julia = "1.0"
+
+[extras]
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+
+[targets]
+test = ["Test", "ReverseDiff"]

jlpkg/Enzyme/src/Enzyme.jl

@@ -1,86 +1,16 @@
 module Enzyme
 
-include(joinpath(@__DIR__, "..", "deps", "deps.jl"))
+export autodiff
 
 using LLVM
 using LLVM.Interop
 import MCAnalyzer: irgen
 
+include("utils.jl")
+include("ad.jl")
+include("opt.jl")
 
-@enum Diffe begin
-   Duplicate = 1
-   Output = 2
-   Constant = 3
-end
-
-function hasfieldcount(@nospecialize(dt))
-    try
-        fieldcount(dt)
-    catch
-        return false
-    end
-    return true
-end
-
-function whatType(@nospecialize(dt))
-    if <:(dt, Array)
-        sub = whatType(eltype(dt))
-        if sub == "diffe_dup"
-            return "diffe_dup"
-        elseif sub == "diffe_out"
-            return "diffe_dup"
-        else
-            @assert(sub == "diffe_const")
-            return "diffe_const"
-        end
-    end
-    if <:(dt, Real)
-        return "diffe_out"
-    end
-    if <:(dt, Int)
-        return "diffe_const"
-    end
-    if <:(dt, String)
-        return "diffe_const"
-    end
-
-    if ! hasfieldcount(dt)
-        # just be safe for now
-        return "diffe_dup"
-    end
-    @assert(hasfieldcount(dt))
-    @assert(isstructtype(dt))
-    passpointer = true
-    if passpointer
-        ty = "diffe_const"
-        for (ft, fn) in zip(fieldtypes(dt), fieldnames(dt))
-            sub = whatType(ft)
-            if sub == "diffe_dup"
-                ty = "diffe_dup"
-            elseif sub == "diffe_out"
-                ty = "diffe_dup"
-            else
-                @assert(sub == "diffe_const")
-            end
-        end
-        return ty
-    else
-        ty = "diffe_const"
-        for (ft, fn) in zip(fieldtypes(dt), fieldnames(dt))
-            sub = whatType(ft)
-            if sub == "diffe_dup"
-                ty = "diffe_dup"
-            elseif sub == "diffe_out"
-                if ty != "diffe_dup"
-                    ty = "diffe_out"
-                end
-            else
-                @assert(sub == "diffe_const")
-            end
-        end
-        return ty
-    end
-end
+using .Opt: optimize!
 
 @generated function autodiff(f, args...)
     # Obtain the function and all it's dependencies in one handy module
@@ -106,12 +36,8 @@ end
 
     i = 1
     j = 1
-    Base.println(typeof(ccf))
-    Base.println(typeof(llvmtype(ccf)))
-    Base.println(llvmtype(ccf))
     orig_params = parameters(ccf)
     for p in orig_params
-        Base.println(llvmtype(p))
         push!(argtypes2, llvmtype(p))
         if diffetypes[i] == "diffe_dup"
             push!(argtypes2, llvmtype(p))
@@ -120,8 +46,6 @@ end
             i+=1
         end
     end
-    Base.println(argtypes2)
-    
 
     # TODO get function type from ccf
     ft2  = LLVM.FunctionType(rettype, argtypes2)
@@ -148,8 +72,8 @@ end
             j+=1
         end
         push!(params, llvm_params[j])
-        j+=1
-        i+=1
+        j += 1
+        i += 1
     end
 
     Builder(ctx) do builder
@@ -168,120 +92,11 @@ end
         #end
     end
 
-    # TODO: Run pipeline and Enzyme pass
+    # Run pipeline and Enzyme pass
+    optimize!(mod)
 
     _args = (:(args[$i]) for i in 1:length(args))
     call_function(llvmf, Float64, Tuple{args...}, Expr(:tuple, _args...))
 end
-    
-
-function jlpasses!(pm)
-    ccall(:jl_add_optimization_passes, Nothing,
-          (LLVM.API.LLVMPassManagerRef, Cint),
-           LLVM.ref(pm), optlevel[])
-end
-
-function enzyme_pass!(pm)
-    ccall((libenzyme, :AddEnzymePass), Nothing, (LLVM.API.LLVMPassManagerRef,) LLVM.ref(pm))
-end
-
-function optimize!(mod, opt_level=2)
-    # everying except unroll, slpvec, loop-vec
-    # then finish Julia GC
-    ModulePassManager() do pm
-        if opt_level < 2
-            cfgsimplification!(pm)
-            if opt_level == 1
-                scalar_repl_aggregates!(pm) # SSA variant?
-                instruction_combinining!(pm)
-                early_cse!(pm)
-            end
-            mem_cpy_opt!(pm)
-            always_inliner!(pm)
-            lower_simdloop!(pm)
-
-            # GC passes
-            barrier_noop!(pm)
-            lower_exc_handlers!(pm)
-            gc_invariant_verifier(pm, false)
-            late_lower_gc_frame!(pm)
-            # TODO: FinalLowerGCPass
-            lower_ptls!(pm, #=dump_native=# false)
-
-            # Enzyme pass
-            barrier_noop!(pm)
-            enzyme_pass!(pm)
-        else
-            propagate_julia_addrsp!(pm)
-            scoped_no_alias_aa!(pm)
-            type_based_alias_analysis!(pm)
-            if opt_level >= 2
-                basic_alias_analysis!(pm)
-            end
-            cfgsimplification!(pm)
-            # TODO: DCE
-            scalar_repl_aggregates!(pm) # SSA variant?
-            mem_cpy_opt!(pm)
-            always_inliner!(pm)
-            alloc_opt!(pm)
-            instruction_combinining!(pm)
-            cfgsimplification!(pm)
-            scalar_repl_aggregates!(pm) # SSA variant?
-            instruction_combinining!(pm)
-            jump_threading!(pm)
-            instruction_combinining!(pm)
-            reassociate!(pm)
-            early_cse!(pm)
-            alloc_opt!(pm)
-            loop_idiom!(pm)
-            loop_rotate!(pm)
-            lower_simdloop!(pm)
-            licm!(pm)
-            loop_unswitch!(pm)
-            instruction_combinining!(pm)
-            ind_var_simplify!(pm)
-            loop_deletion!(pm)
-            # SimpleLoopUnroll -- not for Enzyme
-            alloc_opt!(pm)
-            scalar_repl_aggregates!(pm) # SSA variant?
-            instruction_combinining!(pm)
-            gvn!(pm)
-            mem_cpy_opt!(pm)
-            sccp!(pm)
-            # TODO: Sinking Pass
-            # TODO: LLVM <7 InstructionSimplifier
-            instruction_combinining!(pm)
-            jump_threading!(pm)
-            dead_store_elimination!(pm)
-            alloc_opt!(pm)
-            cfgsimplification!(pm)
-            loop_idiom!(pm)
-            loop_deletion!(pm)
-            jump_threading!(pm)
-            # SLP_Vectorizer -- not for Enzyme
-            aggressive_dce!(pm)
-            instruction_combinining!(pm)
-            # Loop Vectorize -- not for Enzyme
-            # InstCombine
-
-            # GC passes
-            barrier_noop!(pm)
-            lower_exc_handlers!(pm)
-            gc_invariant_verifier(pm, false)
-            late_lower_gc_frame!(pm)
-            # TODO: FinalLowerGCPass
-            # TODO: DCE
-            lower_ptls!(pm, #=dump_native=# false)
-            cfgsimplification!(pm)
-            instruction_combinining!(pm) # Extra for Enzyme
-            # CombineMulAddPass will run on second pass
-
-            # Enzyme pass
-            barrier_noop!(pm)
-            enzyme_pass!(pm)
-        end
-        run!(pm, mod)
-    end
-end
-
-end
+   
+end # module

jlpkg/Enzyme/src/ad.jl

@@ -0,0 +1,66 @@
+@enum Diffe begin
+   Duplicate = 1
+   Output = 2
+   Constant = 3
+end
+
+function whatType(@nospecialize(dt))
+    if <:(dt, Array)
+        sub = whatType(eltype(dt))
+        if sub == "diffe_dup"
+            return "diffe_dup"
+        elseif sub == "diffe_out"
+            return "diffe_dup"
+        else
+            @assert(sub == "diffe_const")
+            return "diffe_const"
+        end
+    end
+    if <:(dt, Real)
+        return "diffe_out"
+    end
+    if <:(dt, Int)
+        return "diffe_const"
+    end
+    if <:(dt, String)
+        return "diffe_const"
+    end
+
+    if !hasfieldcount(dt)
+        # just be safe for now
+        return "diffe_dup"
+    end
+
+    @assert(hasfieldcount(dt))
+    @assert(isstructtype(dt))
+    passpointer = true
+    if passpointer
+        ty = "diffe_const"
+        for (ft, fn) in zip(fieldtypes(dt), fieldnames(dt))
+            sub = whatType(ft)
+            if sub == "diffe_dup"
+                ty = "diffe_dup"
+            elseif sub == "diffe_out"
+                ty = "diffe_dup"
+            else
+                @assert(sub == "diffe_const")
+            end
+        end
+        return ty
+    else
+        ty = "diffe_const"
+        for (ft, fn) in zip(fieldtypes(dt), fieldnames(dt))
+            sub = whatType(ft)
+            if sub == "diffe_dup"
+                ty = "diffe_dup"
+            elseif sub == "diffe_out"
+                if ty != "diffe_dup"
+                    ty = "diffe_out"
+                end
+            else
+                @assert(sub == "diffe_const")
+            end
+        end
+        return ty
+    end
+end