From 214010d1e57d8baab06180d23d30b2366d1073ed Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Tue, 26 Sep 2023 16:54:47 -0500
Subject: [PATCH] Fix custom rule with batch fwd pointer (#1083)

---
 src/compiler.jl | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/compiler.jl b/src/compiler.jl
index 8745a6aa6c..69a92ca393 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -4642,7 +4642,14 @@ function enzyme_custom_common_rev(forward::Bool, B, orig::LLVM.CallInst, gutils,
             shadowV = extract_value!(B, res, idx)
             if get_return_info(RealRt)[2] !== nothing
                 dval = invert_pointer(gutils, operands(orig)[1], B)
-                store!(B, shadowV, dval)
+
+                for idx in 1:width
+                    to_store = (width == 1) ? shadowV : extract_value!(B, shadowV, idx-1)
+
+                    store_ptr = (width == 1) ? dval : extract_value!(B, dval, idx-1)
+
+                    store!(B, to_store, store_ptr)
+                end
                 shadowV = C_NULL
             else
                 @assert value_type(shadowV) == shadowType