Fix batched forward constant get (#2143)

* Fix batched forward constant get * fix * Update typeunstable.jl * Update typeunstable.jl * Update typeunstable.jl * fixup
EnzymeAD · Nov 30, 2024 · c9c79ac · c9c79ac · wsmoses · Nov 30, 2024
1 parent 2a24bb5
commit c9c79ac
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 4 deletions.
diff --git a/src/errors.jl b/src/errors.jl
@@ -94,6 +94,17 @@ function julia_error(
     B::LLVM.API.LLVMBuilderRef,
 )::LLVM.API.LLVMValueRef
     msg = Base.unsafe_string(cstr)
+    julia_error(msg, val, errtype, data, data2, B)
+end
+
+function julia_error(
+    msg::String,
+    val::LLVM.API.LLVMValueRef,
+    errtype::API.ErrorType,
+    data::Ptr{Cvoid},
+    data2::LLVM.API.LLVMValueRef,
+    B::LLVM.API.LLVMBuilderRef,
+)::LLVM.API.LLVMValueRef
     bt = nothing
     ir = nothing
     if val != C_NULL
@@ -331,7 +342,9 @@ function julia_error(
                                 sres
                             end
                             shadowres = insert_value!(prevbb, shadowres, res, idx - 1)
-                            push!(created, shadowres)
+                            if shadowres isa LLVM.Instruction
+				push!(created, shadowres)
+			    end
                         end
                         return shadowres
                     end

diff --git a/src/rules/typeunstablerules.jl b/src/rules/typeunstablerules.jl
@@ -1067,6 +1067,7 @@ function common_jl_getfield_fwd(offset, B, orig, gutils, normalR, shadowR)
             shadowres = UndefValue(
                 LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(normal))),
             )
+            position!(B, LLVM.Instruction(LLVM.API.LLVMGetNextInstruction(normal)))
             for idx = 1:width
                 shadowres = insert_value!(B, shadowres, normal, idx - 1)
             end
@@ -1534,8 +1535,13 @@ end
     end
     origops = collect(operands(orig))
     width = get_width(gutils)
-    if !is_constant_value(gutils, origops[1])
-        shadowin = invert_pointer(gutils, origops[1], B)
+    if !is_constant_value(gutils, origops[1]) || !get_runtime_activity(gutils)
+        shadowin = if !is_constant_value(gutils, origops[1])
+            invert_pointer(gutils, origops[1], B)
+        else
+            estr = "Mismatched activity for: " * string(orig) * " const input " *string(origops[1]) * ", differentiable return"
+            LLVM.Value(julia_error(estr, orig.ref, API.ET_MixedActivityError, gutils.ref, origops[1].ref, B.ref))
+        end
         if width == 1
             args = LLVM.Value[
                 shadowin
@@ -1565,6 +1571,7 @@ end
             shadowres = UndefValue(
                 LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(normal))),
             )
+            position!(B, LLVM.Instruction(LLVM.API.LLVMGetNextInstruction(normal)))
             for idx = 1:width
                 shadowres = insert_value!(B, shadowres, normal, idx - 1)
             end

diff --git a/test/typeunstable.jl b/test/typeunstable.jl
@@ -101,4 +101,24 @@ end
 	res = Enzyme.autodiff(Forward, toactivepair, BatchDuplicated(2.7f0, (2.0f0, 5.0f0)), BatchDuplicated(3.1, (3.0, 7.0)))
 	@test res[1][1] ≈ 2.7f0 * 3.0 + 2.0f0 * 3.1
 	@test res[1][2] ≈ 2.7f0 * 7.0 + 5.0f0 * 3.1	
-end
+end
+
+struct InsFwdNormal1{T<:Real}
+	σ::T
+end
+
+struct InsFwdNormal2{T<:Real}
+	σ::T
+end
+
+insfwdlogpdf(d, x) = d.σ
+
+function insfwdfunc(x)
+    dists = [InsFwdNormal1{Float64}(1.0), InsFwdNormal2{Float64}(1.0)]
+    return sum(Base.Fix2(insfwdlogpdf, x), dists)
+end
+
+@testset "Forward Batch Constant insertion" begin
+    res = Enzyme.gradient(Enzyme.Forward, insfwdfunc, [0.5, 0.7])[1]
+    @test res ≈ [0.0, 0.0]
+end