Despecialize make_tracer (#540)

* Despecialize make_tracer

* further despec

* even more

* even more inference stop

* more testinfra

* fix

* more despec

* fixes

* fix

* fix

* fix

* fix

* fixup

* fix

* more stuff

* fix

Author: wsmoses

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>", "Avik Pal <avikpal@mit.edu>"]
-version = "0.2.18"
+version = "0.2.19"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

ext/ReactantCUDAExt.jl

@@ -722,37 +722,40 @@ Reactant.@reactant_overlay @noinline function CUDA.cufunction(
     return Core.Typeof(res)(f, res.entry)
 end
 
-function Reactant.traced_type(
-    ::Type{A}, seen::ST, ::Val{mode}, track_numbers
-) where {A<:CuTracedArray,ST,mode}
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(A::Type{<:CuTracedArray}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
+)
     return A
 end
 
-function Reactant.traced_type(
-    ::Type{A}, seen::ST, ::Val{mode}, track_numbers
-) where {T,N,A<:CUDA.CuArray{T,N},ST,mode}
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(A::Type{<:CUDA.CuArray}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
+)
+    T = eltype(A)
+    N = ndims(A)
     if mode == Reactant.ArrayToConcrete && T <: Reactant.ReactantPrimitive
         return Reactant.ConcreteRArray{T,N}
     else
-        TT = Reactant.traced_type(T, seen, Val(mode), track_numbers)
+        TT = Reactant.traced_type_inner(T, seen, mode, track_numbers)
         if TT === T
             return A
         else
-            return Array{traced_type(T, seen, Val(mode), track_numbers),N}
+            return Array{Reactant.traced_type_inner(T, seen, mode, track_numbers),N}
         end
     end
 end
 
 function Reactant.make_tracer(
-    seen, @nospecialize(prev::RT), @nospecialize(path), mode; track_numbers=(), kwargs...
-) where {RT<:CUDA.CuArray}
+    seen, @nospecialize(prev::CUDA.CuArray), @nospecialize(path), mode; @nospecialize(track_numbers::Type=Union{}), kwargs...
+)
+    RT = Core.Typeof(prev)
     if haskey(seen, prev)
         return seen[prev]
     end
     if mode == Reactant.ArrayToConcrete && eltype(RT) <: Reactant.ReactantPrimitive
         return seen[prev] = Reactant.ConcreteRArray(Array(prev))
     end
-    TT = Reactant.traced_type(eltype(RT), (), Val(mode), track_numbers)
+    TT = Reactant.traced_type(eltype(RT), Val(mode), track_numbers)
     if TT === eltype(RT)
         return prev
     end

ext/ReactantOffsetArraysExt.jl

@@ -1,12 +1,15 @@
 module ReactantOffsetArraysExt
 
+using OffsetArrays
 using OffsetArrays: OffsetArray
 using Reactant: Reactant, MLIR, Ops, TracedRArray
 
-function Reactant.traced_type(
-    ::Type{<:OffsetArray{<:Any,N,T}}, seen::ST, ::Val{mode}, track_numbers
-) where {T,N,ST,mode}
-    T2 = Reactant.traced_type(T, seen, Val(mode), track_numbers)
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(OA::Type{<:OffsetArray}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type=Union{})
+)
+    N = ndims(OA)
+    T = OffsetArrays.parenttype(OA)
+    T2 = Reactant.traced_type_inner(T, seen, mode, track_numbers)
     return OffsetArray{eltype(T2),N,T2}
 end
 

src/Reactant.jl

@@ -23,14 +23,23 @@ else
     const ReactantFloat = Union{Float16,Float32,Float64}
 end
 
-const ReactantInt = Union{Int8,UInt8,Int16,UInt16,Int32,UInt32,Int64,UInt64}
+@static if isdefined(Core, :BFloat16)
+    const ReactantComplexFloat = Union{Complex{Float16},Complex{Core.BFloat16},Complex{Float32},Complex{Float64}}
+else
+    const ReactantComplexFloat = Union{Complex{Float16},Complex{Float32},Complex{Float64}}
+end
+
+const ReactantInt = Union{Int8,UInt8,Int16,UInt16,Int32,UInt32,Int64,UInt64,Int128,UInt128}
+
+const ReactantComplexInt = Union{Complex{Int8},Complex{UInt8},Complex{Int16},Complex{UInt16},Complex{Int32},Complex{UInt32},Complex{Int64},Complex{UInt64},Complex{Int128},Complex{UInt128}}
 
 const ReactantFloatInt = Union{
     Base.uniontypes(ReactantInt)...,Base.uniontypes(ReactantFloat)...
 }
 
 const ReactantPrimitive = Union{
-    Bool,Base.uniontypes(ReactantFloatInt)...,Complex{Float32},Complex{Float64}
+    Bool,Base.uniontypes(ReactantFloatInt)...,
+    Base.uniontypes(ReactantComplexInt)...,Base.uniontypes(ReactantComplexFloat)...
 }
 
 abstract type RNumber{T<:ReactantPrimitive} <: Number end

src/TracedUtils.jl

@@ -131,7 +131,7 @@ function make_mlir_fn(
             (:args, i),
             concretein ? Reactant.ConcreteToTraced : Reactant.TracedSetPath;
             toscalar,
-            track_numbers=construct_function_without_args ? (Number,) : (),
+            track_numbers=construct_function_without_args ? Number : Union{},
         )
     end
 
@@ -201,7 +201,7 @@ function make_mlir_fn(
         result,
         (:result,),
         concretein ? Reactant.TracedTrack : Reactant.TracedSetPath;
-        track_numbers=construct_function_without_args ? (Number,) : (),
+        track_numbers=construct_function_without_args ? Number : Union{},
     )
 
     # marks buffers to be donated