feat: implement a separate TracedRNumber (#161)

* feat: TracedRScalar

* feat: partial progress on getting scalars to work

* refactor: Scalar --> Number

* fix: batching

* fix: promote_rule and introduce union over primitive types

* chore: apply formatting

* feat: type-restrict arrays

* refactor: move scalar ops to a separate file

* feat: support Base.float

* fix: import ordering

* feat: handle `broadcast_preserving_zero_d` in a generic fashion

* refactor: move code a bit

* test: more test fixes

* chore: apply formatting

* fix: setindex with scalars

* fix: scalar broadcasting case

* feat: support BFloat16 from Core (if available)

* test: more native lux functionality unblocked

* refactor: use a union type for traced types

* fix: check for reactant primitives

* fix: missing import

* fix: correct semantics for Colon mapreduce

* fix: trace_type

* fix: minor fixes

* feat: support logsoftmax

* fix: bool promote rule

* fix: broadcasting of closures

* refactor: use TracedTypes

* Fix type of `preserved_args`

* Rename `TracedTypes` to `TracedType`

* small testset rename

* fix: special handling for concatenation of numbers

* Reenable tests

* Rename `ReactantPrimitives` to `ReactantPrimitive`

---------

Co-authored-by: Sergio Sánchez Ramírez <[email protected]>
Co-authored-by: Sergio Sánchez Ramírez <[email protected]>

ext/ReactantNNlibExt.jl

@@ -1,31 +1,24 @@
 module ReactantNNlibExt
 
 using NNlib
-using Reactant: Reactant, TracedRArray, AnyTracedRArray, materialize_traced_array, MLIR
+using Reactant:
+    Reactant, TracedRArray, AnyTracedRArray, materialize_traced_array, MLIR, TracedRNumber
 
 for (jlop, hloop) in (
     (:(NNlib.tanh_fast), :tanh),
     (:(NNlib.sigmoid_fast), :logistic),
     (:(NNlib.sigmoid), :logistic),
 )
-    @eval function $(jlop)(x::TracedRArray{T,0}) where {T}
-        return TracedRArray{T,0}(
+    @eval function $(jlop)(x::TracedRNumber{T}) where {T}
+        return TracedRNumber{T}(
             (),
             Reactant.MLIR.IR.result(
                 Reactant.MLIR.Dialects.stablehlo.$(hloop)(x.mlir_data), 1
             ),
-            (),
         )
     end
 end
 
-# Don't confuse our poor scalar arrays, we no like numbers we like 0D arrays
-for nnlib_op in setdiff(Tuple(NNlib.ACTIVATIONS), (:tanh_fast, :sigmoid_fast, :sigmoid, :σ))
-    @eval function NNlib.$(nnlib_op)(x::TracedRArray{T,0}) where {T}
-        return invoke(NNlib.$(nnlib_op), Tuple{Any}, x)
-    end
-end
-
 # TODO handle non finite cases
 function NNlib.softmax!(out::TracedRArray{T,N}, x::AbstractArray; dims=1) where {T,N}
     max_ = NNlib.fast_maximum(x; dims)
@@ -39,6 +32,20 @@ function NNlib.softmax!(out::TracedRArray{T,N}, x::AbstractArray; dims=1) where
     return out ./= tmp
 end
 
+function NNlib.logsoftmax!(out::TracedRArray{T}, x::AbstractArray; dims=1) where {T}
+    max_ = NNlib.fast_maximum(x; dims)
+    # if all(isfinite, max_)
+    @fastmath out .= x .- max_
+    # else
+    #     _zero, _minf, _inf = T(0), T(-Inf), T(Inf)
+    #     @. out = ifelse(
+    #         isequal(max_, _inf), ifelse(isequal(x, _inf), _zero, _minf), x - max_
+    #     )
+    # end
+    @fastmath log_ = log.(sum(exp, out; dims))
+    return out .-= log_
+end
+
 function NNlib.conv(
     x::AnyTracedRArray{T,N}, W::AnyTracedRArray{T}, cdims::DenseConvDims
 ) where {T,N}

src/Compiler.jl

@@ -6,10 +6,12 @@ import ..Reactant:
     XLA,
     ConcreteRArray,
     TracedRArray,
+    TracedRNumber,
     OrderedIdDict,
     make_tracer,
     TracedToConcrete,
-    append_path
+    append_path,
+    TracedType
 
 @inline traced_getfield(@nospecialize(obj), field) = Base.getfield(obj, field)
 
@@ -286,10 +288,10 @@ function compile_mlir!(mod, f, args; optimize=true)
         )
     end
 
-    preserved_args = Tuple{TracedRArray,Int}[]
+    preserved_args = Tuple{TracedType,Int}[]
     results = [MLIR.IR.operand(ret, i) for i in 1:MLIR.IR.noperands(ret)]
     nresults = MLIR.IR.Value[]
-    linear_results2 = TracedRArray[]
+    linear_results2 = TracedType[]
     for (i, op) in enumerate(results)
         if !MLIR.IR.is_block_arg(op)
             push!(nresults, op)
@@ -573,7 +575,7 @@ end
 function compile_xla(f, args; client=nothing)
     # register MLIR dialects
     ctx = MLIR.IR.Context()
-    Base.append!(Reactant.registry[]; context=ctx)
+    append!(Reactant.registry[]; context=ctx)
     @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
 
     return MLIR.IR.context!(ctx) do

src/ConcreteRArray.jl

@@ -74,10 +74,6 @@ function Base.convert(::Type{T}, x::ConcreteRArray{T,0}) where {T}
     return to_float(x)
 end
 
-function Base.promote_rule(::Type{<:RArray{T1,0}}, ::Type{T2}) where {T1,T2}
-    return Base.promote_rule(T1, T2)
-end
-
 for jlop in (:(Base.isless), :(Base.:+), :(Base.:-), :(Base.:*), :(Base.:/), :(Base.:^))
     @eval begin
         function $jlop(x::ConcreteRArray{T,0}, y::ConcreteRArray{U,0}) where {T,U}
@@ -158,7 +154,7 @@ function Base.getindex(a::ConcreteRArray{T}, args::Vararg{Int,N}) where {T,N}
 end
 
 function mysetindex!(a, v, args::Vararg{Int,N}) where {N}
-    Base.setindex!(a, v, args...)
+    setindex!(a, v, args...)
     return nothing
 end
 

src/Reactant.jl

@@ -7,7 +7,45 @@ include("OrderedIdDict.jl")
 
 using Enzyme
 
-abstract type RArray{T,N} <: AbstractArray{T,N} end
+@static if isdefined(Core, :BFloat16)
+    const ReactantPrimitive = Union{
+        Bool,
+        Int8,
+        UInt8,
+        Int16,
+        UInt16,
+        Int32,
+        UInt32,
+        Int64,
+        UInt64,
+        Float16,
+        Core.BFloat16,
+        Float32,
+        Float64,
+        Complex{Float32},
+        Complex{Float64},
+    }
+else
+    const ReactantPrimitive = Union{
+        Bool,
+        Int8,
+        UInt8,
+        Int16,
+        UInt16,
+        Int32,
+        UInt32,
+        Int64,
+        UInt64,
+        Float16,
+        Float32,
+        Float64,
+        Complex{Float32},
+        Complex{Float64},
+    }
+end
+
+abstract type RArray{T<:ReactantPrimitive,N} <: AbstractArray{T,N} end
+abstract type RNumber{T<:ReactantPrimitive} <: Number end
 
 function Base.reshape(A::RArray, dims::Tuple{Vararg{Union{Int,Colon}}})
     return reshape(A, Base._reshape_uncolon(A, dims))
@@ -45,8 +83,13 @@ include("mlir/MLIR.jl")
 include("XLA.jl")
 include("Interpreter.jl")
 include("utils.jl")
+
 include("ConcreteRArray.jl")
+include("TracedRNumber.jl")
 include("TracedRArray.jl")
+
+const TracedType = Union{TracedRArray,TracedRNumber}
+
 include("Tracing.jl")
 include("Compiler.jl")