Separate Reactant number types

Project.toml

@@ -118,7 +118,7 @@ DLFP8Types = "0.1"
 Dates = "1.10"
 Downloads = "1.6"
 EnumX = "1"
-Enzyme = "0.13.115"
+Enzyme = "0.13.177"
 EnzymeCore = "0.8.14"
 FFTW = "1"
 FileWatching = "1.10"

docs/src/api/api.md

@@ -34,6 +34,14 @@ Reactant.to_rarray
 ```@docs
 ConcreteRArray
 ConcreteRNumber
+Reactant.RNumber
+Reactant.RReal
+RInteger
+RFloat
+RComplex
+Reactant.traced_number_type
+Reactant.pjrt_number_type
+Reactant.ifrt_number_type
 ```
 
 ## Inspect Generated HLO

ext/ReactantAbstractFFTsExt.jl

@@ -7,7 +7,21 @@ using Reactant.Ops: @opcall
 
 # FFTW.jl defines methods on StridedArrays, and we have to be more specific than FFTW to
 # catch its methods.
-const AnyStridedTracedRArray{T,N} = StridedArray{TracedRNumber{T},N}
+const AnyStridedTracedRArray{T,N} = StridedArray{<:TracedRNumber{T},N}
+
+# Traced reals hit the `fftfloat`-based conversion path in AbstractFFTs, whose
+# fallback only supports the BLAS float types and copies through a host array.
+function AbstractFFTs._fftfloat(::Type{T}) where {T<:TracedRNumber}
+    return Reactant.traced_number_type(AbstractFFTs._fftfloat(Reactant.unwrapped_eltype(T)))
+end
+
+function AbstractFFTs.complexfloat(x::AbstractArray{<:Reactant.TracedRReal{T}}) where {T}
+    return Reactant.promote_to(Reactant.TracedRArray{complex(float(T)),ndims(x)}, x)
+end
+
+function AbstractFFTs.realfloat(x::AbstractArray{<:Reactant.TracedRReal{T}}) where {T}
+    return Reactant.promote_to(Reactant.TracedRArray{float(T),ndims(x)}, x)
+end
 
 # To automatically convert FFT plans to traced versions
 # To extend a user needs to extend Reactant.reactant_fftplan for their plan type

ext/ReactantCUDAExt.jl

@@ -63,27 +63,53 @@ end
 
 Reactant.use_overlayed_version(::CuTracedArray) = true
 
-struct CuTracedRNumber{T,A} <: Number
-    ptr::Core.LLVMPtr{T,A}
+# Split by numeric kind, mirroring the traced host types, so that kernel-side
+# numbers subtype Integer and AbstractFloat where possible.
+for (CuT, AbstractT) in (
+    (:CuTracedRInteger, :(Reactant.RInteger)),
+    (:CuTracedRFloat, :(Reactant.RFloat)),
+    (:CuTracedRComplex, :(Reactant.RComplex)),
+)
+    @eval begin
+        struct $CuT{T,A} <: $AbstractT{T}
+            ptr::Core.LLVMPtr{T,A}
 
-    function CuTracedRNumber{T,A}(xs::TracedRNumber) where {T,A}
-        Reactant.Compiler.guard_from_gc_for_module(MLIR.IR.current_module(), xs)
-        ptr = Base.reinterpret(Core.LLVMPtr{T,CUDA.AS.Global}, Base.pointer_from_objref(xs))
-        return new(ptr)
-    end
-    function CuTracedRNumber{T,A}(ptr::Core.LLVMPtr{T,A}) where {T,A}
-        return new(ptr)
+            function $CuT{T,A}(xs::TracedRNumber) where {T,A}
+                Reactant.Compiler.guard_from_gc_for_module(MLIR.IR.current_module(), xs)
+                ptr = Base.reinterpret(
+                    Core.LLVMPtr{T,CUDA.AS.Global}, Base.pointer_from_objref(xs)
+                )
+                return new(ptr)
+            end
+            function $CuT{T,A}(ptr::Core.LLVMPtr{T,A}) where {T,A}
+                return new(ptr)
+            end
+        end
+
+        Reactant.use_overlayed_version(::$CuT) = true
     end
 end
 
-Reactant.use_overlayed_version(::CuTracedRNumber) = true
+const CuTracedRReal{T,A} = Union{CuTracedRInteger{T,A},CuTracedRFloat{T,A}}
+const CuTracedRNumber{T,A} = Union{CuTracedRReal{T,A},CuTracedRComplex{T,A}}
+
+@inline cu_traced_number_type(::Type{T}) where {T<:Integer} = CuTracedRInteger{T}
+@inline cu_traced_number_type(::Type{T}) where {T<:AbstractFloat} = CuTracedRFloat{T}
+@inline cu_traced_number_type(::Type{T}) where {T<:Complex} = CuTracedRComplex{T}
+
+@inline function CuTracedRNumber{T,A}(xs::TracedRNumber) where {T,A}
+    return cu_traced_number_type(T){A}(xs)
+end
+@inline function CuTracedRNumber{T,A}(ptr::Core.LLVMPtr{T,A}) where {T,A}
+    return cu_traced_number_type(T){A}(ptr)
+end
 
 Base.@nospecializeinfer Reactant.is_traced_number(
     @nospecialize(T::Type{<:CuTracedRNumber})
 ) = true
 Reactant.unwrapped_eltype(::Type{<:CuTracedRNumber{T}}) where {T} = T
 
-@inline CuTracedRNumber{T,A}(val::Number) where {T,A} = convert(CuTracedRNumber{T,A}, val)
+@inline (::Type{CuT})(val::Number) where {CuT<:CuTracedRNumber} = convert(CuT, val)
 
 function Base.getindex(RN::CuTracedRNumber{T,A}) where {T,A}
     align = alignment(RN)
@@ -92,6 +118,10 @@ end
 
 Base.convert(::Type{T}, RN::CuTracedRNumber) where {T<:Number} = convert(T, getindex(RN))
 
+# Defined per numeric kind: a single method on the `CuTracedRNumber` union
+# would be ambiguous with Base's methods on `Real`/`Integer`/`AbstractFloat`.
+const CU_TRACED_NUMBER_KINDS = (CuTracedRInteger, CuTracedRFloat, CuTracedRComplex)
+
 for jlop in (
     :(Base.min),
     :(Base.mod),
@@ -106,10 +136,32 @@ for jlop in (
     :(Base.:(==)),
     :(Base.:(!=)),
 )
-    @eval begin
-        @inline $jlop(a::CuTracedRNumber, b::CuTracedRNumber) = $jlop(a[], b[])
-        @inline $jlop(a::CuTracedRNumber{T,A}, b::Number) where {T,A} = $jlop(a[], b)
-        @inline $jlop(a::Number, b::CuTracedRNumber{T,A}) where {T,A} = $jlop(a, b[])
+    for K1 in CU_TRACED_NUMBER_KINDS
+        @eval begin
+            @inline $jlop(a::$K1, b::Number) = $jlop(a[], b)
+            @inline $jlop(a::Number, b::$K1) = $jlop(a, b[])
+        end
+        for K2 in CU_TRACED_NUMBER_KINDS
+            @eval @inline $jlop(a::$K1, b::$K2) = $jlop(a[], b[])
+        end
+        for X in (
+            :Real,
+            :Integer,
+            :Bool,
+            :AbstractFloat,
+            :Complex,
+            :(Complex{Bool}),
+            :Rational,
+            :BigInt,
+            :BigFloat,
+            :AbstractIrrational,
+        )
+            jlop == :(Base.:^) && X === :Integer && continue
+            @eval begin
+                @inline $jlop(a::$X, b::$K1) = $jlop(a, b[])
+                @inline $jlop(a::$K1, b::$X) = $jlop(a[], b)
+            end
+        end
     end
 end
 
@@ -123,28 +175,28 @@ end
 @inline Base.ifelse(cond::CuTracedRNumber, a::CuTracedRNumber, b::CuTracedRNumber) =
     ifelse(cond[], a[], b[])
 
-Base.@constprop :aggressive @inline Base.:^(
-    a::CuTracedRNumber{T,A}, b::Integer
-) where {T,A} = ^(a[], b)
+for K in CU_TRACED_NUMBER_KINDS
+    @eval Base.@constprop :aggressive @inline Base.:^(a::$K, b::Integer) = ^(a[], b)
+end
 
 @inline Base.unsafe_trunc(::Type{T}, a::CuTracedRNumber) where {T} =
     Base.unsafe_trunc(T, a[])
 
-for jlop in (:(Base.:+), :(Base.:-), :(Base.isnan), :(Base.isfinite), :(Base.isinf))
-    @eval begin
-        @inline $jlop(a::CuTracedRNumber) = $jlop(a[])
-    end
+for jlop in (:(Base.:+), :(Base.:-), :(Base.isnan), :(Base.isfinite), :(Base.isinf)),
+    K in CU_TRACED_NUMBER_KINDS
+
+    @eval @inline $jlop(a::$K) = $jlop(a[])
 end
 
-Base.OneTo(x::CuTracedRNumber{<:Integer}) = Base.OneTo(x[])
+Base.OneTo(x::CuTracedRInteger) = Base.OneTo(x[])
 
 @static if isdefined(Base, :unchecked_oneto)
     function Base.unchecked_oneto(x::CuTracedRNumber{<:Integer})
         return Base.unchecked_oneto(x[])
     end
 end
 
-@inline function Base.convert(CT::Type{CuTracedRNumber{Float64,1}}, x::Number)
+@inline function Base.convert(CT::Type{<:CuTracedRNumber{Float64,1}}, x::Number)
     return CT(
         Base.reinterpret(
             Core.LLVMPtr{Float64,1},
@@ -168,7 +220,7 @@ end
     )
 end
 
-@inline function Base.convert(CT::Type{CuTracedRNumber{Float32,1}}, x::Number)
+@inline function Base.convert(CT::Type{<:CuTracedRNumber{Float32,1}}, x::Number)
     return CT(
         Base.reinterpret(
             Core.LLVMPtr{Float32,1},
@@ -273,8 +325,12 @@ function Base.show(io::IO, a::AT) where {AT<:CuTracedArray}
     Printf.@printf(io, "%s cu traced array at %p", join(size(a), '×'), Int(pointer(a)))
 end
 
-function Base.show(io::IO, a::AT) where {AT<:CuTracedRNumber}
-    Printf.@printf(io, "%s cu traced rnumber at %p", join(size(a), '×'), Int(pointer(a)))
+for K in CU_TRACED_NUMBER_KINDS
+    @eval function Base.show(io::IO, a::AT) where {AT<:$K}
+        Printf.@printf(
+            io, "%s cu traced rnumber at %p", join(size(a), '×'), Int(pointer(a))
+        )
+    end
 end
 
 ## array interface
@@ -811,7 +867,7 @@ function compile(job)
             MLIR.IR.setattr!(MLIR.IR.Operation(cur_module), dl_attr_name, MLIR.IR.Attribute(dl))
         end
 
-        String(Reactant.TracedUtils.get_attribute_by_name(linkRes, "sym_name"))
+        return String(Reactant.TracedUtils.get_attribute_by_name(linkRes, "sym_name"))
     end
 
     return LLVMFunc{job.source.specTypes.parameters[1],job.source.specTypes}(nothing, entry)
@@ -1595,15 +1651,17 @@ Base.@nospecializeinfer function Reactant.traced_type_inner(
     return A
 end
 
-Base.@nospecializeinfer function Reactant.traced_type_inner(
-    @nospecialize(A::Type{<:CuTracedRNumber}),
-    seen,
-    @nospecialize(mode::Reactant.TraceMode),
-    @nospecialize(track_numbers::Type),
-    @nospecialize(ndevices),
-    @nospecialize(runtime)
-)
-    return A
+for CuT in (:CuTracedRNumber, :CuTracedRInteger, :CuTracedRFloat, :CuTracedRComplex)
+    @eval Base.@nospecializeinfer function Reactant.traced_type_inner(
+        @nospecialize(A::Type{<:$CuT}),
+        seen,
+        @nospecialize(mode::Reactant.TraceMode),
+        @nospecialize(track_numbers::Type),
+        @nospecialize(ndevices),
+        @nospecialize(runtime)
+    )
+        return A
+    end
 end
 
 Base.@nospecializeinfer function Reactant.traced_type_inner(

ext/ReactantOneHotArraysExt.jl

@@ -7,11 +7,15 @@ using ReactantCore: ReactantCore
 using Reactant.Ops: @opcall
 
 __compatible_eltype(::Type{T}, ::Type{U}) where {T,U} = T
-function __compatible_eltype(::Type{TracedRNumber{T}}, ::Type{TracedRNumber{U}}) where {T,U}
-    return TracedRNumber{T}
+function __compatible_eltype(
+    ::Type{<:TracedRNumber{T}}, ::Type{<:TracedRNumber{U}}
+) where {T,U}
+    return Reactant.traced_number_type(T)
+end
+__compatible_eltype(::Type{<:TracedRNumber{T}}, ::Type{U}) where {T,U} = T
+function __compatible_eltype(::Type{T}, ::Type{<:TracedRNumber{U}}) where {T,U}
+    return Reactant.traced_number_type(T)
 end
-__compatible_eltype(::Type{TracedRNumber{T}}, ::Type{U}) where {T,U} = T
-__compatible_eltype(::Type{T}, ::Type{TracedRNumber{U}}) where {T,U} = TracedRNumber{T}
 
 function Reactant.traced_type_inner(
     @nospecialize(_::Type{OneHotArray{T,N,Np1,I}}),
@@ -60,7 +64,7 @@ function OneHotArrays.onehotbatch(data::AnyTracedRArray{<:Any,N}, labels) where
     )
     data = ReactantCore.materialize_traced_array(reshape(data, 1, size(data)...))
     indices = UInt32.(@opcall(findfirst(data .== labels_expanded; dimension=1)))
-    return OneHotArray{TracedRNumber{UInt32},N,N + 1,typeof(indices)}(
+    return OneHotArray{Reactant.TracedRInteger{UInt32},N,N + 1,typeof(indices)}(
         indices, length(labels)
     )
 end
@@ -73,7 +77,7 @@ function OneHotArrays.onehotbatch(
         TracedRNumber{UInt32} ∘ Base.Fix2(+, 1 - first(labels)),
         ReactantCore.materialize_traced_array(data),
     )
-    return OneHotArray{TracedRNumber{UInt32},N,N + 1,typeof(indices)}(
+    return OneHotArray{Reactant.TracedRInteger{UInt32},N,N + 1,typeof(indices)}(
         indices, length(labels)
     )
 end

ext/ReactantSparseArraysExt/Errors.jl

@@ -1,15 +1,15 @@
 # We don't yet support sparse arrays. Throwing errors on this dispatches to avoid
 # ambiguities.
-function Base.getindex(::AbstractSparseArray{TracedRNumber{T},N,1}, ::Int64) where {T,N}
+function Base.getindex(::AbstractSparseArray{<:TracedRNumber{T},N,1}, ::Int64) where {T,N}
     return error("Sparse arrays are not supported by reactant yet.")
 end
 
 function Base.getindex(
-    ::AbstractSparseMatrixCSC{TracedRNumber{T}}, ::Int64, ::Int64
+    ::AbstractSparseMatrixCSC{<:TracedRNumber{T}}, ::Int64, ::Int64
 ) where {T}
     return error("Sparse arrays are not supported by reactant yet.")
 end
 
-function Base.getindex(::CHOLMOD.Sparse{TracedRNumber{T}}, ::Int64, ::Int64) where {T}
+function Base.getindex(::CHOLMOD.Sparse{<:TracedRNumber{T}}, ::Int64, ::Int64) where {T}
     return error("Sparse arrays are not supported by reactant yet.")
 end

ext/ReactantSparseArraysExt/ReactantSparseArraysExt.jl

@@ -1,10 +1,39 @@
 module ReactantSparseArraysExt
 
-using Reactant: Reactant, TracedRNumber
+using Reactant: Reactant, TracedRInteger, TracedRNumber
 using SparseArrays:
-    SparseArrays, ReadOnly, AbstractSparseArray, CHOLMOD, AbstractSparseMatrixCSC
+    SparseArrays,
+    ReadOnly,
+    AbstractSparseArray,
+    AbstractSparseVector,
+    CHOLMOD,
+    AbstractSparseMatrixCSC
 
 include("Errors.jl")
 include("ReadOnly.jl")
 
+# Scalar indexing disambiguators against the generic traced-array `getindex`;
+# they delegate to the SparseArrays implementations.
+function Base.getindex(
+    x::AbstractSparseMatrixCSC{<:TracedRNumber},
+    i::Union{Int,TracedRInteger{Int}},
+    j::Union{Int,TracedRInteger{Int}},
+)
+    return Base.invoke(getindex, Tuple{AbstractSparseMatrixCSC,Integer,Integer}, x, i, j)
+end
+
+function Base.getindex(
+    x::AbstractSparseVector{<:TracedRNumber}, i::Union{Int,TracedRInteger{Int}}
+)
+    return Base.invoke(getindex, Tuple{AbstractSparseVector,Integer}, x, i)
+end
+
+function Base.getindex(
+    x::CHOLMOD.Sparse{<:TracedRNumber},
+    i::Union{Int,TracedRInteger{Int}},
+    j::Union{Int,TracedRInteger{Int}},
+)
+    return Base.invoke(getindex, Tuple{CHOLMOD.Sparse,Integer,Integer}, x, i, j)
+end
+
 end