Merge branch 'main' into vc/blocking_ring

Author: wsmoses

Project.toml

@@ -1,7 +1,7 @@
 name = "Enzyme"
 uuid = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>"]
-version = "0.13.96"
+version = "0.13.97"
 
 [deps]
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
@@ -43,9 +43,9 @@ ADTypes = "1"
 BFloat16s = "0.2, 0.3, 0.4, 0.5, 0.6"
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
-DynamicPPL = "0.35, 0.36, 0.37, 0.38"
+DynamicPPL = "0.35 - 0.39"
 EnzymeCore = "0.8.15"
-Enzyme_jll = "0.0.207"
+Enzyme_jll = "0.0.208"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.6.2"
 LLVM = "9.1"

src/typetree.jl

@@ -305,7 +305,7 @@ else
         dl,
         seen::TypeTreeTable,
     ) where {kind,T}
-        tt = copy(typetree(Ptr{T}, ctx, dl, seen))
+        tt = copy(typetree(typed_fieldtype(AT, 1), ctx, dl, seen))
         shift!(tt, dl, 0, sizeof(Int), 0)
 
         for f = 2:fieldcount(AT)

src/utils.jl

@@ -462,10 +462,23 @@ end
 
 else
 
+function is_memory_ref_field2_an_offset(@nospecialize(T::Type{<:GenericMemoryRef}))
+    ET = eltype(T)
+
+    # 0 = inlinealloc
+    # 1 = isboxed
+    # 2 = isbitsunion
+    return (Base.datatype_arrayelem(T.types[2]) == 2) || Base.datatype_layoutsize(T.types[2]) == 0
+end
+
 @inline function typed_fieldtype(@nospecialize(T::Type), i::Int)::Type
     if T <: GenericMemoryRef && i == 1 || T <: GenericMemory && i == 2
-        eT = eltype(T)
-        Ptr{eT}
+        if T <: GenericMemoryRef && i == 1 && is_memory_ref_field2_an_offset(T)
+            Int
+        else
+            eT = eltype(T)
+            Ptr{eT}
+        end
     else
         fieldtype(T, i)
     end

test/typetree.jl

@@ -48,6 +48,9 @@ struct UnionStruct1{T}
     y::Any
 end
 
+struct EmptyStruct
+end
+
 @testset "TypeTree" begin
     @test tt(Float16) == "{[-1]:Float@half}"
     @test tt(Float32) == "{[-1]:Float@float}"
@@ -99,6 +102,14 @@ end
               "{[0]:Pointer, [0,0]:Pointer, [0,4]:Float@float, [0,8]:Float@double, [4]:Integer, [8]:Pointer, [8,0]:Pointer, [8,4]:Float@float, [8,8]:Float@double, [12]:Integer, [16]:Pointer, [16,0]:Pointer, [16,4]:Float@float, [16,8]:Float@double, [20]:Integer, [24]:Pointer, [24,0]:Pointer, [24,4]:Float@float, [24,8]:Float@double}"
     end
 
+    @static if VERSION >= v"1.11-"
+        if Sys.WORD_SIZE == 64
+            @test tt(MemoryRef{EmptyStruct}) == "{[0]:Integer, [1]:Integer, [2]:Integer, [3]:Integer, [4]:Integer, [5]:Integer, [6]:Integer, [7]:Integer, [8]:Pointer, [8,0]:Integer, [8,1]:Integer, [8,2]:Integer, [8,3]:Integer, [8,4]:Integer, [8,5]:Integer, [8,6]:Integer, [8,7]:Integer, [8,8]:Pointer}"
+            @test tt(MemoryRef{Union{Float64, Int64}}) == "{[0]:Integer, [1]:Integer, [2]:Integer, [3]:Integer, [4]:Integer, [5]:Integer, [6]:Integer, [7]:Integer, [8]:Pointer, [8,0]:Integer, [8,1]:Integer, [8,2]:Integer, [8,3]:Integer, [8,4]:Integer, [8,5]:Integer, [8,6]:Integer, [8,7]:Integer, [8,8]:Pointer}"
+            @test tt(MemoryRef{Float64}) == "{[-1]:Pointer, [0,-1]:Float@double, [8,0]:Integer, [8,1]:Integer, [8,2]:Integer, [8,3]:Integer, [8,4]:Integer, [8,5]:Integer, [8,6]:Integer, [8,7]:Integer, [8,8]:Pointer, [8,8,-1]:Float@double}"
+        end
+    end
+
     @test tt(UnionStruct1{Float32}) == "{[0]:Float@float, [4]:Integer, [8]:Pointer}"
 
     if Sys.WORD_SIZE == 64