Add additional support for svecref (#1246)

src/compiler.jl

@@ -2740,7 +2740,7 @@ function annotate!(mod, mode)
         end
     end
 
-    for fname in ("jl_f_getfield","ijl_f_getfield","jl_get_nth_field_checked","ijl_get_nth_field_checked")
+    for fname in ("jl_f_getfield","ijl_f_getfield","jl_get_nth_field_checked","ijl_get_nth_field_checked", "jl_f__svec_ref", "ijl_f__svec_ref")
         if haskey(fns, fname)
             fn = fns[fname]
             push!(function_attributes(fn), LLVM.EnumAttribute("readonly", 0))

src/compiler/optimize.jl

@@ -1,3 +1,29 @@
+mutable struct PipelineConfig
+    Speedup::Cint
+    Size::Cint
+    lower_intrinsics::Cint
+    dump_native::Cint
+    external_use::Cint
+    llvm_only::Cint
+    always_inline::Cint
+    enable_early_simplifications::Cint
+    enable_early_optimizations::Cint
+    enable_scalar_optimizations::Cint
+    enable_loop_optimizations::Cint
+    enable_vector_pipeline::Cint
+    remove_ni::Cint
+    cleanup::Cint
+end
+
+function pipeline_options(; lower_intrinsics=true, dump_native=false, external_use=false, llvm_only=false, always_inline=true, enalbe_early_simplifications=true,
+       enable_scalar_optimizations=true,
+       enable_loop_optimizations=true,
+       enable_vector_pipeline=true,
+       remove_ni=true,
+       cleanup=true, Size=0, Speedup=3)
+    return PipelineConfig(Speedup, Size, lower_intrinsics, dump_native, external_use, llvm_only, always_inline, enable_early_simplifications, enable_early_optimizations, enable_scalar_optimizations, enable_loop_optimizations, enable_vector_pipeline, remove_ni, cleanup)
+end
+
 function addNA(inst, node::LLVM.Metadata, MD)
     md = metadata(inst)
     next = nothing 

src/rules/typeunstablerules.jl

@@ -702,24 +702,73 @@ function common_f_svec_ref_fwd(offset, B, orig, gutils, normalR, shadowR)
     return false
 end
 
+function error_if_differentiable(::Type{T}) where T
+    seen = ()
+    areg = active_reg_inner(T, seen, nothing, #=justActive=#Val(true))
+    if areg != AnyState
+        throw(AssertionError("Found unhandled differentiable variable in jl_f_svec_ref $T"))
+    end
+    nothing
+end
+
 function common_f_svec_ref_augfwd(offset, B, orig, gutils, normalR, shadowR, tapeR)
-    if is_constant_value(gutils, orig) && is_constant_inst(gutils, orig)
+    if is_constant_value(gutils, orig)
         return true
     end
-    emit_error(B, orig, "Enzyme: Not yet implemented augmented forward for jl_f__svec_ref")
 
-    normal = (unsafe_load(normalR) != C_NULL) ? LLVM.Instruction(unsafe_load(normalR)) : nothing
-    if shadowR != C_NULL && normal !== nothing
-        unsafe_store!(shadowR, normal.ref)
+    width = get_width(gutils)
+
+    origmi, origh, origkey = operands(orig)[offset:end-1]
+
+    shadowh = invert_pointer(gutils, origh, B)
+
+    newvals = API.CValueType[API.VT_Primal, API.VT_Shadow, API.VT_Primal]
+
+    if offset != 1
+        pushfirst!(newvals, API.VT_Primal)
     end
+
+    errfn = if is_constant_value(gutils, origh)
+        error_if_differentiable
+    else
+        error_if_active
+    end
+
+    mi = new_from_original(gutils, origmi)
+
+    shadowres = if width == 1
+        newops = LLVM.Value[mi, shadowh, new_from_original(gutils, origkey)]
+        if offset != 1
+            pushfirst!(newops, operands(orig)[1])
+        end
+        cal = call_samefunc_with_inverted_bundles!(B, gutils, orig, newops, newvals, #=lookup=#false)
+        callconv!(cal, callconv(orig))
+
+
+        emit_apply_generic!(B, LLVM.Value[unsafe_to_llvm(errfn), emit_jltypeof!(B, cal)])
+        cal
+    else
+        ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(orig)))
+        shadow = LLVM.UndefValue(ST)
+        for j in 1:width
+            newops = LLVM.Value[mi, extract_value!(B, shadowh, j-1), new_from_original(gutils, origkey)]
+            if offset != 1
+                pushfirst!(newops, operands(orig)[1])
+            end
+            cal = call_samefunc_with_inverted_bundles!(B, gutils, orig, newops, newvals, #=lookup=#false)
+            callconv!(cal, callconv(orig))
+            emit_apply_generic!(B, LLVM.Value[unsafe_to_llvm(errfn), emit_jltypeof!(B, cal)])
+            shadow = insert_value!(B, shadow, cal, j-1)
+        end
+        shadow
+    end
+
+    unsafe_store!(shadowR, shadowres.ref)
 
     return false
 end
 
 function common_f_svec_ref_rev(offset, B, orig, gutils, tape)
-    if !is_constant_value(gutils, orig) || !is_constant_inst(gutils, orig)
-        emit_error(B, orig, "Enzyme: Not yet implemented reverse for jl_f__svec_ref")
-    end
     return nothing
 end