Skip to content

Conversation

@wsmoses
Copy link
Member

@wsmoses wsmoses commented Jan 16, 2025

Trying to speed up compile time


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}}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
const ReactantComplexFloat = Union{Complex{Float16},Complex{Core.BFloat16},Complex{Float32},Complex{Float64}}
const ReactantComplexFloat = Union{
Complex{Float16},Complex{Core.BFloat16},Complex{Float32},Complex{Float64}
}


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}}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
const ReactantComplexInt = Union{Complex{Int8},Complex{UInt8},Complex{Int16},Complex{UInt16},Complex{Int32},Complex{UInt32},Complex{Int64},Complex{UInt64},Complex{Int128},Complex{UInt128}}
const ReactantComplexInt = Union{
Complex{Int8},
Complex{UInt8},
Complex{Int16},
Complex{UInt16},
Complex{Int32},
Complex{UInt32},
Complex{Int64},
Complex{UInt64},
Complex{Int128},
Complex{UInt128},
}

Comment on lines +41 to +42
Bool,Base.uniontypes(ReactantFloatInt)...,
Base.uniontypes(ReactantComplexInt)...,Base.uniontypes(ReactantComplexFloat)...
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Bool,Base.uniontypes(ReactantFloatInt)...,
Base.uniontypes(ReactantComplexInt)...,Base.uniontypes(ReactantComplexFloat)...
Bool,
Base.uniontypes(ReactantFloatInt)...,
Base.uniontypes(ReactantComplexInt)...,
Base.uniontypes(ReactantComplexFloat)...,

Comment on lines +59 to +62
(ConcreteRArray{Float64,0}, TracedRArray{Float64,0}, TracedRArray{Float64, 0}),
(ConcreteRArray{Float64,1}, TracedRArray{Float64,1}, TracedRArray{Float64, 1}),
(ConcreteRArray{Float64,2}, TracedRArray{Float64,2}, TracedRArray{Float64, 2}),
(ConcreteRArray{Float64,3}, TracedRArray{Float64,3}, TracedRArray{Float64, 3}),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
(ConcreteRArray{Float64,0}, TracedRArray{Float64,0}, TracedRArray{Float64, 0}),
(ConcreteRArray{Float64,1}, TracedRArray{Float64,1}, TracedRArray{Float64, 1}),
(ConcreteRArray{Float64,2}, TracedRArray{Float64,2}, TracedRArray{Float64, 2}),
(ConcreteRArray{Float64,3}, TracedRArray{Float64,3}, TracedRArray{Float64, 3}),
(
ConcreteRArray{Float64,0},
TracedRArray{Float64,0},
TracedRArray{Float64,0},
),
(
ConcreteRArray{Float64,1},
TracedRArray{Float64,1},
TracedRArray{Float64,1},
),
(
ConcreteRArray{Float64,2},
TracedRArray{Float64,2},
TracedRArray{Float64,2},
),
(
ConcreteRArray{Float64,3},
TracedRArray{Float64,3},
TracedRArray{Float64,3},
),

Comment on lines +65 to +66
(Array{Float64,1}, Array{Float64,1}, Array{TracedRNumber{Float64}, 1}),
(Array{ConcreteRArray{Float64,2},1}, Array{TracedRArray{Float64,2},1}, Array{TracedRArray{Float64,2}, 1}),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
(Array{Float64,1}, Array{Float64,1}, Array{TracedRNumber{Float64}, 1}),
(Array{ConcreteRArray{Float64,2},1}, Array{TracedRArray{Float64,2},1}, Array{TracedRArray{Float64,2}, 1}),
(Array{Float64,1}, Array{Float64,1}, Array{TracedRNumber{Float64},1}),
(
Array{ConcreteRArray{Float64,2},1},
Array{TracedRArray{Float64,2},1},
Array{TracedRArray{Float64,2},1},
),


# Union types
(Union{Nothing,Int}, Union{Nothing,Int}),
(Union{Nothing,Int}, Union{Nothing,Int}, Union{Nothing, TracedRNumber{Int}}),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
(Union{Nothing,Int}, Union{Nothing,Int}, Union{Nothing, TracedRNumber{Int}}),
(Union{Nothing,Int}, Union{Nothing,Int}, Union{Nothing,TracedRNumber{Int}}),

(
Union{Nothing,ConcreteRArray{Float64,1}},
Union{Nothing,TracedRArray{Float64,1}},
Union{Nothing, TracedRArray{Float64, 1}}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Union{Nothing, TracedRArray{Float64, 1}}
Union{Nothing,TracedRArray{Float64,1}},

Comment on lines +78 to +79
(Ptr{ConcreteRArray{Float64,1}}, Ptr{TracedRArray{Float64,1}}, Ptr{TracedRArray{Float64,1}}),
(Core.LLVMPtr{Float64}, Core.LLVMPtr{Float64}, Core.LLVMPtr{TracedRNumber{Float64}}),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
(Ptr{ConcreteRArray{Float64,1}}, Ptr{TracedRArray{Float64,1}}, Ptr{TracedRArray{Float64,1}}),
(Core.LLVMPtr{Float64}, Core.LLVMPtr{Float64}, Core.LLVMPtr{TracedRNumber{Float64}}),
(
Ptr{ConcreteRArray{Float64,1}},
Ptr{TracedRArray{Float64,1}},
Ptr{TracedRArray{Float64,1}},
),
(
Core.LLVMPtr{Float64},
Core.LLVMPtr{Float64},
Core.LLVMPtr{TracedRNumber{Float64}},
),

(
Core.LLVMPtr{ConcreteRArray{Float64,1}},
Core.LLVMPtr{TracedRArray{Float64,1}},
Core.LLVMPtr{TracedRArray{Float64,1}}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Core.LLVMPtr{TracedRArray{Float64,1}}
Core.LLVMPtr{TracedRArray{Float64,1}},
),
(
Base.RefValue{Float64},
Base.RefValue{Float64},
Base.RefValue{TracedRNumber{Float64}},

Core.LLVMPtr{TracedRArray{Float64,1}}
),
(Base.RefValue{Float64}, Base.RefValue{Float64}),
(Base.RefValue{Float64}, Base.RefValue{Float64}, Base.RefValue{TracedRNumber{Float64}}),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
(Base.RefValue{Float64}, Base.RefValue{Float64}, Base.RefValue{TracedRNumber{Float64}}),

::Type{A}, seen::ST, ::Val{mode}, track_numbers
) where {T,N,A<:CUDA.CuArray{T,N},ST,mode}
Base.@nospecializeinfer function Reactant.traced_type(
@nospecialize(A::Type{<:CUDA.CuArray}), seen::ST, ::Val{mode}, track_numbers
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would @nospecialize on ::Val{mode} do sth?

src/Tracing.jl Outdated
::Type{<:ConcreteRNumber{T}}, seen, ::Val{mode}, track_numbers
) where {T,mode}
for T in (DataType, Module, Nothing, Symbol, AbstractChar, AbstractString, AbstractFloat, Integer, RNumber)
@eval Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:$T}), seen, mode, track_numbers)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here different Val(mode) would provoke different specializations when its actually not used

Suggested change
@eval Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:$T}), seen, mode, track_numbers)
@eval Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:$T}), seen, @nospecialize(mode), track_numbers)

src/Tracing.jl Outdated
end

Base.@nospecializeinfer function traced_type(
@nospecialize(C::Type{<:Complex}), seen::ST, mode::Val{Mode}, track_numbers::TN
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think we would benefit here too because it's just passed

Suggested change
@nospecialize(C::Type{<:Complex}), seen::ST, mode::Val{Mode}, track_numbers::TN
@nospecialize(C::Type{<:Complex}), seen::ST, @nospecialize(mode::Val{Mode}), track_numbers::TN

src/Tracing.jl Outdated
end
end

Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:Function}), seen, mode, track_numbers)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:Function}), seen, mode, track_numbers)
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:Function}), seen, @nospecialize(mode), track_numbers)

Comment on lines +176 to +177
@inline is_concrete_tuple(x::T2) where {T2} =
(x <: Tuple) && !(x === Tuple) && !(x isa UnionAll)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isn't T2 just Type? and because Type is parametric, we could do @nospecialize and maybe @nospecializeinfer too (it always returns Bool)? wdyt @gbaraldi?

Suggested change
@inline is_concrete_tuple(x::T2) where {T2} =
(x <: Tuple) && !(x === Tuple) && !(x isa UnionAll)
@inline @nospecializeinfer is_concrete_tuple(@nospecialize(T)) =
(T <: Tuple) && !(T === Tuple) && !(T isa UnionAll)

src/Tracing.jl Outdated
@inline is_concrete_tuple(x::T2) where {T2} =
(x <: Tuple) && !(x === Tuple) && !(x isa UnionAll)

Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:Tuple}), seen, mode, track_numbers)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:Tuple}), seen, mode, track_numbers)
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:Tuple}), seen, @nospecialize(mode), track_numbers)

src/Tracing.jl Outdated
return Tuple{TT...}
end

Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:NamedTuple}), seen, mode, track_numbers)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:NamedTuple}), seen, mode, track_numbers)
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:NamedTuple}), seen, @nospecialize(mode), track_numbers)

src/Tracing.jl Outdated
end

function traced_type(::Type{T}, seen, mode, track_numbers) where {T<:XLAArray}
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:XLAArray}), seen, mode, track_numbers)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:XLAArray}), seen, mode, track_numbers)
Base.@nospecializeinfer function traced_type(@nospecialize(T::Type{<:XLAArray}), seen, @nospecialize(mode), track_numbers)

src/Tracing.jl Outdated

for P in (Ptr, Core.LLVMPtr, Base.RefValue)
@eval function traced_type(::Type{P}, seen, mode, track_numbers) where {T,P<:$P{T}}
@eval Base.@nospecializeinfer function traced_type(@nospecialize(PT::Type{<:$P}), seen, mode, track_numbers)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
@eval Base.@nospecializeinfer function traced_type(@nospecialize(PT::Type{<:$P}), seen, mode, track_numbers)
@eval Base.@nospecializeinfer function traced_type(@nospecialize(PT::Type{<:$P}), seen, @nospecialize(mode), track_numbers)

src/Tracing.jl Outdated
end

function traced_type(::Type{Val{T}}, seen, mode, track_numbers) where {T}
Base.@nospecializeinfer function traced_type(@nospecialize(VT::Type{<:Val}), seen, mode, track_numbers)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Base.@nospecializeinfer function traced_type(@nospecialize(VT::Type{<:Val}), seen, mode, track_numbers)
Base.@nospecializeinfer function traced_type(@nospecialize(VT::Type{<:Val}), seen, @nospecialize(mode), track_numbers)

Comment on lines 705 to 713
@nospecialize(prev::NamedTuple),
@nospecialize(path),
mode;
track_numbers=(),
kwargs...,
) where {A,RT}
)
NT = Core.Typeof(prev)
A = NT.parameters[1]
RT = NT.parameters[2]
Copy link
Collaborator

@mofeing mofeing Jan 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what i understood from @nospecialize is that this change should be equivalent?
not 100% sure though

Copy link
Collaborator

@mofeing mofeing left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

didn't know about tracked_numbers... is that for passing the type of the result if using a TracedRArray or so? not sure what is for but probably you want to despecialize it too if you despecialize the T in traced_type

Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/Tracing.jl

Lines 353 to 354 in 5742e62

empty!(new_ci.codelocs)
resize!(new_ci.linetable, 1) # see note below


[JuliaFormatter] reported by reviewdog 🐶

new_ci.slotflags = UInt8[0x00 for i = 1:4]


[JuliaFormatter] reported by reviewdog 🐶

push!(new_ci.codelocs, 1) # see note below


[JuliaFormatter] reported by reviewdog 🐶

@eval Base.@assume_effects :removable :foldable :nothrow @inline function traced_type(T::Type, mode::Val, track_numbers::Type)


[JuliaFormatter] reported by reviewdog 🐶

$(Expr(:meta, :generated, traced_type_generator))


[JuliaFormatter] reported by reviewdog 🐶

@nospecialize(prev::Union{Base.ExceptionStack, Core.MethodInstance}),


[JuliaFormatter] reported by reviewdog 🐶

function make_tracer(seen, prev::ConcreteRNumber{T}, @nospecialize(path), mode; kwargs...) where {T}


[JuliaFormatter] reported by reviewdog 🐶

seen, @nospecialize(prev::Number), @nospecialize(path), mode; @nospecialize(track_numbers::Type), kwargs...


[JuliaFormatter] reported by reviewdog 🐶

seen, @nospecialize(prev::Array), @nospecialize(path), mode; @nospecialize(track_numbers::Type), kwargs...


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/Tracing.jl

Lines 774 to 776 in 5742e62

function make_tracer(
seen, @nospecialize(prev::Tuple), @nospecialize(path), mode; kwargs...
)


[JuliaFormatter] reported by reviewdog 🐶

function to_rarray_internal(@nospecialize(::TracedRArray), @nospecialize(track_numbers::Type))


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/Tracing.jl

Lines 837 to 838 in 5742e62

@inline to_rarray_internal(@nospecialize(x::ConcreteRArray), @nospecialize(track_numbers::Type)) = x
@inline function to_rarray_internal(@nospecialize(x::Array{<:ReactantPrimitive}), @nospecialize(track_numbers::Type))


[JuliaFormatter] reported by reviewdog 🐶

@inline to_rarray_internal(@nospecialize(x::ConcreteRNumber), @nospecialize(track_numbers::Type)) = x


[JuliaFormatter] reported by reviewdog 🐶

tracedty = traced_type(
origty, Val(ConcreteToTraced), ()
)


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/test/tracing.jl

Lines 102 to 104 in 5742e62

tracedty2 = traced_type(
origty, Val(ConcreteToTraced), ReactantPrimitive
)


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/test/tracing.jl

Lines 128 to 130 in 5742e62

@test_throws NoFieldMatchError traced_type(
Node, Val(ArrayToConcrete), ()
)

@inline getmap(::Val{T}, ::Val{T}, ::Val{T2}, args...) where {T,T2} = T2

function traced_type(::Type{T}, seen, mode, track_numbers) where {T}
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type), seen, mode::TraceMode, @nospecialize(track_numbers::Type)
)

Comment on lines +131 to +132
for T in (DataType, Module, Nothing, Symbol, AbstractChar, AbstractString, AbstractFloat, Integer, RNumber)
@eval Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:$T}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
for T in (DataType, Module, Nothing, Symbol, AbstractChar, AbstractString, AbstractFloat, Integer, RNumber)
@eval Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:$T}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
for T in (
DataType,
Module,
Nothing,
Symbol,
AbstractChar,
AbstractString,
AbstractFloat,
Integer,
RNumber,
)
@eval Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type{<:$T}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)
)

end

Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type{<:ReactantPrimitive}), seen, @nospecialize(mode::TraceMode), @nospecialize(track_numbers::Type)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@nospecialize(T::Type{<:ReactantPrimitive}), seen, @nospecialize(mode::TraceMode), @nospecialize(track_numbers::Type)
@nospecialize(T::Type{<:ReactantPrimitive}),
seen,
@nospecialize(mode::TraceMode),
@nospecialize(track_numbers::Type)

end

Base.@nospecializeinfer function traced_type_inner(
@nospecialize(C::Type{<:Complex}), seen, @nospecialize(mode::TraceMode), @nospecialize(track_numbers::Type)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@nospecialize(C::Type{<:Complex}), seen, @nospecialize(mode::TraceMode), @nospecialize(track_numbers::Type)
@nospecialize(C::Type{<:Complex}),
seen,
@nospecialize(mode::TraceMode),
@nospecialize(track_numbers::Type)

Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter

[JuliaFormatter] reported by reviewdog 🐶

function make_tracer(seen, prev::ConcreteRNumber{T}, @nospecialize(path), mode; kwargs...) where {T}


[JuliaFormatter] reported by reviewdog 🐶

seen, @nospecialize(prev::Number), @nospecialize(path), mode; @nospecialize(track_numbers::Type=Union{}), kwargs...


[JuliaFormatter] reported by reviewdog 🐶

seen, @nospecialize(prev::Array), @nospecialize(path), mode; @nospecialize(track_numbers::Type), kwargs...


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/Tracing.jl

Lines 776 to 778 in c61d2e1

function make_tracer(
seen, @nospecialize(prev::Tuple), @nospecialize(path), mode; kwargs...
)


[JuliaFormatter] reported by reviewdog 🐶

@nospecialize(track_numbers::Type=Union{}),


[JuliaFormatter] reported by reviewdog 🐶

function to_rarray_internal(@nospecialize(::TracedRArray), @nospecialize(track_numbers::Type))


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/Tracing.jl

Lines 839 to 840 in c61d2e1

@inline to_rarray_internal(@nospecialize(x::ConcreteRArray), @nospecialize(track_numbers::Type)) = x
@inline function to_rarray_internal(@nospecialize(x::Array{<:ReactantPrimitive}), @nospecialize(track_numbers::Type))


[JuliaFormatter] reported by reviewdog 🐶

@inline to_rarray_internal(@nospecialize(x::ConcreteRNumber), @nospecialize(track_numbers::Type)) = x


[JuliaFormatter] reported by reviewdog 🐶

tracedty = traced_type(
origty, Val(ConcreteToTraced), Union{}
)


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/test/tracing.jl

Lines 102 to 104 in c61d2e1

tracedty2 = traced_type(
origty, Val(ConcreteToTraced), ReactantPrimitive
)


[JuliaFormatter] reported by reviewdog 🐶

type, Val(ConcreteToTraced), Union{}


[JuliaFormatter] reported by reviewdog 🐶

Real, Val(Reactant.ArrayToConcrete), Union{}


[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/test/tracing.jl

Lines 128 to 130 in c61d2e1

@test_throws NoFieldMatchError traced_type(
Node, Val(ArrayToConcrete), Union{}
)

end
end

Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:Function}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:Function}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type{<:Function}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)
)

@inline is_concrete_tuple(x::T2) where {T2} =
(x <: Tuple) && !(x === Tuple) && !(x isa UnionAll)

Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:Tuple}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:Tuple}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type{<:Tuple}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)
)

return Tuple{TT...}
end

Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:NamedTuple}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:NamedTuple}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type{<:NamedTuple}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)
)

return NamedTuple{N,traced_type_inner(V, seen, mode, track_numbers)}
end

Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:AbstractDict}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:AbstractDict}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type{<:AbstractDict}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)
)

end

Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T0::Type{<:ConcreteRNumber}), seen, mode::TraceMode, @nospecialize(track_numbers::Type)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@nospecialize(T0::Type{<:ConcreteRNumber}), seen, mode::TraceMode, @nospecialize(track_numbers::Type)
@nospecialize(T0::Type{<:ConcreteRNumber}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)

end

function traced_type(::Type{T}, seen, mode, track_numbers) where {T<:XLAArray}
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:XLAArray}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function traced_type_inner(@nospecialize(T::Type{<:XLAArray}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(T::Type{<:XLAArray}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)
)

::Type{A}, seen::ST, ::Val{mode}, track_numbers
) where {T,N,A<:Array{T,N},ST,mode}
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(A::Type{<:Array}), seen, mode::TraceMode, @nospecialize(track_numbers::Type)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@nospecialize(A::Type{<:Array}), seen, mode::TraceMode, @nospecialize(track_numbers::Type)
@nospecialize(A::Type{<:Array}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)

for P in (Ptr, Core.LLVMPtr, Base.RefValue)
@eval function traced_type(::Type{P}, seen, mode, track_numbers) where {T,P<:$P{T}}
return $P{traced_type(T, seen, mode, track_numbers)}
@eval Base.@nospecializeinfer function traced_type_inner(@nospecialize(PT::Type{<:$P}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@eval Base.@nospecializeinfer function traced_type_inner(@nospecialize(PT::Type{<:$P}), seen, mode::TraceMode, @nospecialize(track_numbers::Type))
@eval Base.@nospecializeinfer function traced_type_inner(
@nospecialize(PT::Type{<:$P}),
seen,
mode::TraceMode,
@nospecialize(track_numbers::Type)
)


function traced_type(::Type{Val{T}}, seen, mode, track_numbers) where {T}
if traced_type(typeof(T), seen, mode, track_numbers) == typeof(T)
Base.@nospecializeinfer function traced_type_inner(@nospecialize(VT::Type{<:Val}), seen, @nospecialize(mode::TraceMode), @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function traced_type_inner(@nospecialize(VT::Type{<:Val}), seen, @nospecialize(mode::TraceMode), @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function traced_type_inner(
@nospecialize(VT::Type{<:Val}),
seen,
@nospecialize(mode::TraceMode),
@nospecialize(track_numbers::Type)
)

throw("Val type $(Val{T}) cannot be traced")
end

const traced_type_cache = Dict{Tuple{TraceMode, Type}, Dict{Type, Type}}()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
const traced_type_cache = Dict{Tuple{TraceMode, Type}, Dict{Type, Type}}()
const traced_type_cache = Dict{Tuple{TraceMode,Type},Dict{Type,Type}}()

Comment on lines 97 to 99
tracedty = traced_type(
origty, Reactant.OrderedIdDict(), Val(ConcreteToTraced), ()
origty, Val(ConcreteToTraced), Union{}
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
tracedty = traced_type(
origty, Reactant.OrderedIdDict(), Val(ConcreteToTraced), ()
origty, Val(ConcreteToTraced), Union{}
)
tracedty = traced_type(origty, Val(ConcreteToTraced), Union{})

Comment on lines +102 to +104
tracedty2 = traced_type(
origty, Val(ConcreteToTraced), ReactantPrimitive
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
tracedty2 = traced_type(
origty, Val(ConcreteToTraced), ReactantPrimitive
)
tracedty2 = traced_type(origty, Val(ConcreteToTraced), ReactantPrimitive)

]
@test_throws Union{ErrorException,String} traced_type(
type, Reactant.OrderedIdDict(), Val(ConcreteToTraced), ()
type, Val(ConcreteToTraced), Union{}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
type, Val(ConcreteToTraced), Union{}
type, Val(ConcreteToTraced), Union{}

@testset "traced_type exceptions" begin
@test_throws TracedTypeError Reactant.traced_type(
Real, Reactant.OrderedIdDict(), Val(Reactant.ArrayToConcrete), ()
Real, Val(Reactant.ArrayToConcrete), Union{}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Real, Val(Reactant.ArrayToConcrete), Union{}
Real, Val(Reactant.ArrayToConcrete), Union{}

Comment on lines 128 to 130
@test_throws NoFieldMatchError traced_type(
Node, Reactant.OrderedIdDict(), Val(ArrayToConcrete), ()
Node, Val(ArrayToConcrete), Union{}
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@test_throws NoFieldMatchError traced_type(
Node, Reactant.OrderedIdDict(), Val(ArrayToConcrete), ()
Node, Val(ArrayToConcrete), Union{}
)
@test_throws NoFieldMatchError traced_type(Node, Val(ArrayToConcrete), Union{})

::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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@nospecialize(A::Type{<:CuTracedArray}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
@nospecialize(A::Type{<:CuTracedArray}),
seen,
mode::Reactant.TraceMode,
@nospecialize(track_numbers::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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@nospecialize(A::Type{<:CUDA.CuArray}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
@nospecialize(A::Type{<:CUDA.CuArray}),
seen,
mode::Reactant.TraceMode,
@nospecialize(track_numbers::Type)

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...
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
seen, @nospecialize(prev::CUDA.CuArray), @nospecialize(path), mode; @nospecialize(track_numbers::Type=Union{}), kwargs...
seen,
@nospecialize(prev::CUDA.CuArray),
@nospecialize(path),
mode;
@nospecialize(track_numbers::Type = Union{}),
kwargs...,

) 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{})
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@nospecialize(OA::Type{<:OffsetArray}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type=Union{})
@nospecialize(OA::Type{<:OffsetArray}),
seen,
mode::Reactant.TraceMode,
@nospecialize(track_numbers::Type = Union{})

# $(Expr(:meta, :generated, traced_type_generator))
# end

Base.@assume_effects :total @inline function traced_type(T::Type, ::Val{mode}, track_numbers::Type) where mode
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@assume_effects :total @inline function traced_type(T::Type, ::Val{mode}, track_numbers::Type) where mode
Base.@assume_effects :total @inline function traced_type(
T::Type, ::Val{mode}, track_numbers::Type
) where {mode}

if haskey(traced_type_cache, cache_key)
cache = traced_type_cache[cache_key]
else
cache = Dict{Type, Type}()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
cache = Dict{Type, Type}()
cache = Dict{Type,Type}()

cache = Dict{Type, Type}()
traced_type_cache[cache_key] = cache
end
res1 = traced_type_inner(T, cache, mode, track_numbers)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
res1 = traced_type_inner(T, cache, mode, track_numbers)
return res1 = traced_type_inner(T, cache, mode, track_numbers)

end

function to_rarray_internal(@nospecialize(::TracedRArray), ::Tuple)
function to_rarray_internal(@nospecialize(::TracedRArray), @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
function to_rarray_internal(@nospecialize(::TracedRArray), @nospecialize(track_numbers::Type))
function to_rarray_internal(
@nospecialize(::TracedRArray), @nospecialize(track_numbers::Type)
)

Comment on lines +859 to +860
@inline to_rarray_internal(@nospecialize(x::ConcreteRArray), @nospecialize(track_numbers::Type)) = x
@inline function to_rarray_internal(@nospecialize(x::Array{<:ReactantPrimitive}), @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@inline to_rarray_internal(@nospecialize(x::ConcreteRArray), @nospecialize(track_numbers::Type)) = x
@inline function to_rarray_internal(@nospecialize(x::Array{<:ReactantPrimitive}), @nospecialize(track_numbers::Type))
@inline to_rarray_internal(
@nospecialize(x::ConcreteRArray), @nospecialize(track_numbers::Type)
) = x
@inline function to_rarray_internal(
@nospecialize(x::Array{<:ReactantPrimitive}), @nospecialize(track_numbers::Type)
)

end

@inline to_rarray_internal(@nospecialize(x::ConcreteRNumber), ::Tuple) = x
@inline to_rarray_internal(@nospecialize(x::ConcreteRNumber), @nospecialize(track_numbers::Type)) = x
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@inline to_rarray_internal(@nospecialize(x::ConcreteRNumber), @nospecialize(track_numbers::Type)) = x
@inline to_rarray_internal(
@nospecialize(x::ConcreteRNumber), @nospecialize(track_numbers::Type)
) = x

inds::Vector{Symbol}
end

Base.@nospecializeinfer function Reactant.traced_type_inner(@nospecialize(A::Type{<:MockTensor}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
Base.@nospecializeinfer function Reactant.traced_type_inner(@nospecialize(A::Type{<:MockTensor}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type))
Base.@nospecializeinfer function Reactant.traced_type_inner(
@nospecialize(A::Type{<:MockTensor}),
seen,
mode::Reactant.TraceMode,
@nospecialize(track_numbers::Type)
)


Base.@nospecializeinfer function Reactant.traced_type_inner(@nospecialize(A::Type{<:MockTensor}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type))
T2 = Reactant.traced_type_inner(A.parameters[3], seen, mode, track_numbers)
MT = MockTensor{eltype(T2), ndims(A), T2}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
MT = MockTensor{eltype(T2), ndims(A), T2}
MT = MockTensor{eltype(T2),ndims(A),T2}

@wsmoses wsmoses merged commit 015a748 into main Jan 16, 2025
32 of 39 checks passed
@wsmoses wsmoses deleted the despec branch January 16, 2025 23:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants