Forward and Reverse mode autodiff fails with Polyester.jl

[dominic-chang]

Here's a MWE where I do a reduction with a Polyester.jl. I find that both the forward pass and reverse mode fails.

using Polyester
using Enzyme

function test(arr)
    ans = 0.0
    Polyester.@batch for i in arr
        ans += i
    end
    return ans
end

arr = rand(10)

Enzyme.autodiff(Enzyme.Forward, test, Enzyme.Duplicated, Enzyme.Duplicated(arr,zeros(10)))

ERROR: Current scope: 
; Function Attrs: mustprogress willreturn
define "enzyme_type"="{[-1]:Float@double}" double @preprocess_julia_test_28814({} addrspace(10)* noundef nonnull align 16 dereferenceable(40) "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}" "enzymejl_parmtype"="4765215328" "enzymejl_parmtype_ref"="2" %0) local_unnamed_addr #78 !dbg !2019 {
top:
  %1 = alloca { { i64, [1 x i64] }, double }, align 16
  %.sub = bitcast { { i64, [1 x i64] }, double }* %1 to i8*
  %2 = call {}*** @julia.get_pgcstack() #79
  %ptls_field82 = getelementptr inbounds {}**, {}*** %2, i64 2
  %3 = bitcast {}*** %ptls_field82 to i64***
  %ptls_load8384 = load i64**, i64*** %3, align 8, !tbaa !48
  %4 = getelementptr inbounds i64*, i64** %ptls_load8384, i64 2
  %safepoint = load i64*, i64** %4, align 8, !tbaa !52
  fence syncscope("singlethread") seq_cst
  call void @julia.safepoint(i64* %safepoint) #79, !dbg !2020
  fence syncscope("singlethread") seq_cst
  %5 = call i64 @julia_nthreads_28978() #80, !dbg !2021
  %.not = icmp eq i64 %5, 1, !dbg !2023
  br i1 %.not, label %L4, label %L46, !dbg !2024

...

Stacktrace:
 [1] _atomic_xchg!
   @ ~/.julia/packages/ThreadingUtilities/3z3g0/src/atomics.jl:33
 [2] _exchange_mask!
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:67
 [3] __request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:89
 [4] _request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:61
 [5] request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:121
 [6] request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:128
 [7] batch
   @ ~/.julia/packages/Polyester/eqrC9/src/batch.jl:308
 [8] macro expansion
   @ ~/.julia/packages/Polyester/eqrC9/src/closure.jl:456
 [9] test
   @ /issue.jl:6

and on reverse mode:
Enzyme.autodiff(Enzyme.Reverse, test, Enzyme.Active, Enzyme.Duplicated(arr,zeros(10)))

ERROR: Current scope: 
; Function Attrs: mustprogress willreturn
define "enzyme_type"="{[-1]:Float@double}" double @preprocess_julia_test_28137({} addrspace(10)* noundef nonnull align 16 dereferenceable(40) "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}" "enzymejl_parmtype"="4765215328" "enzymejl_parmtype_ref"="2" %0) local_unnamed_addr #84 !dbg !2154 {
top:
  %1 = alloca { { i64, [1 x i64] }, double }, align 16
  %.sub = bitcast { { i64, [1 x i64] }, double }* %1 to i8*
  %2 = call {}*** @julia.get_pgcstack() #85
  %ptls_field82 = getelementptr inbounds {}**, {}*** %2, i64 2
  %3 = bitcast {}*** %ptls_field82 to i64***
  %ptls_load8384 = load i64**, i64*** %3, align 8, !tbaa !51
  %4 = getelementptr inbounds i64*, i64** %ptls_load8384, i64 2
  %safepoint = load i64*, i64** %4, align 8, !tbaa !55
  fence syncscope("singlethread") seq_cst
  call void @julia.safepoint(i64* %safepoint) #85, !dbg !2155
  fence syncscope("singlethread") seq_cst
  %5 = call i64 @julia_nthreads_28312() #86, !dbg !2156
  %.not = icmp eq i64 %5, 1, !dbg !2158
  br i1 %.not, label %L4, label %L46, !dbg !2159

...

Stacktrace:
 [1] _atomic_xchg!
   @ ~/.julia/packages/ThreadingUtilities/3z3g0/src/atomics.jl:33
 [2] _exchange_mask!
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:67
 [3] __request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:89
 [4] _request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:61
 [5] request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:121
 [6] request_threads
   @ ~/.julia/packages/PolyesterWeave/E9Wdf/src/request.jl:128
 [7] batch
   @ ~/.julia/packages/Polyester/eqrC9/src/batch.jl:308
 [8] macro expansion
   @ ~/.julia/packages/Polyester/eqrC9/src/closure.jl:456
 [9] test
   @ /issue.jl:6

[vchuravy]

Just as an FYI, Polyester support is not an active goal for Enzyme since Polyester deeply mucks with Julia internals.
Someone would likely need to write custom rules for Polyester.

[wsmoses]

I believe this should now be resolved on main by the jll bump, please reopen if not!