Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix mul overload #440

Merged
merged 9 commits into from
Dec 31, 2024
Merged

Fix mul overload #440

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
873ef6c
Fix mul overload
wsmoses Dec 31, 2024
97d184a
fix
wsmoses Dec 31, 2024
411c478
fix
wsmoses Dec 31, 2024
39e07cd
Apply suggestions from code review
wsmoses Dec 31, 2024
3b2c410
fix: handle aos for mul (#441)
avik-pal Dec 31, 2024
733a901
Merge branch 'main' into mulover
wsmoses Dec 31, 2024
b1f5325
Update Reactant.jl
wsmoses Dec 31, 2024
6628582
Update ReactantArrayInterfaceExt.jl
wsmoses Dec 31, 2024
e53b003
Update ext/ReactantArrayInterfaceExt.jl
wsmoses Dec 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 11 additions & 10 deletions ext/ReactantArrayInterfaceExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,20 +2,21 @@ module ReactantArrayInterfaceExt

using ArrayInterface: ArrayInterface
using Reactant:
Reactant, RArray, ConcreteRArray, ConcreteRNumber, TracedRNumber, TracedRArray, Ops
Reactant,
RArray,
ConcreteRArray,
ConcreteRNumber,
TracedRNumber,
TracedRArray,
AnyTracedRArray,
Ops

ArrayInterface.can_setindex(::Type{<:RArray}) = false
ArrayInterface.fast_scalar_indexing(::Type{<:RArray}) = false

function ArrayInterface.aos_to_soa(x::AbstractArray{<:ConcreteRNumber{T}}) where {T}
x_c = ConcreteRArray(zeros(T, size(x)))
x_c .= x
return x_c
end

ArrayInterface.aos_to_soa(x::TracedRArray) = x
function ArrayInterface.aos_to_soa(x::AbstractArray{<:TracedRNumber{T}}) where {T}
return Ops.reshape(vcat(x...), size(x)...)
for aType in
(AbstractArray{<:ConcreteRNumber}, AbstractArray{<:TracedRNumber}, AnyTracedRArray)
@eval ArrayInterface.aos_to_soa(x::$aType) = Reactant.aos_to_soa(x)
end

end
3 changes: 2 additions & 1 deletion src/Overlay.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,10 +127,11 @@ for (cT, aT, bT) in (
@reactant_overlay @noinline function LinearAlgebra.mul!(
C::$cT, A::$aT, B::$bT, α::Number, β::Number
)
A, B = aos_to_soa(A), aos_to_soa(B)
if use_overlayed_version((C, A, B))
TracedLinearAlgebra.overloaded_mul!(C, A, B, α, β)
else
LinearAlgebra._mul!(C, A, B, α, β)
LinearAlgebra.mul!(C, A, B, α, β)
Copy link
Collaborator

Choose a reason for hiding this comment

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

this will very likely trigger the illegal instruction #393

Copy link
Member Author

Choose a reason for hiding this comment

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

There was an absint behavior that we had a while ago but was accidentally removed in a refactor which I re added with the last cuda stuff.

I hope it means we won't hit it

Copy link
Collaborator

Choose a reason for hiding this comment

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

I checked it again in the stack PR, the crash persists

Copy link
Member Author

Choose a reason for hiding this comment

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

Well darn, okay yeah if CI yells I think we need to make a separate noinline function which is called in the override, which itself calls mul

Copy link
Collaborator

Choose a reason for hiding this comment

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

lets add a test like https://github.com/EnzymeAD/Reactant.jl/blob/main/test/basic.jl#L864 to verify this

end
return C
end
Expand Down
16 changes: 16 additions & 0 deletions src/Reactant.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,22 @@ unwrapped_eltype(::RArray{T,N}) where {T,N} = T
unwrapped_eltype(::AbstractArray{T,N}) where {T,N} = unwrapped_eltype(T)
unwrapped_eltype(::AnyTracedRArray{T,N}) where {T,N} = T

aos_to_soa(x::AbstractArray) = x
aos_to_soa(x::AnyTracedRArray) = x
function aos_to_soa(x::AbstractArray{<:ConcreteRNumber{T}}) where {T}
x_c = ConcreteRArray(zeros(T, size(x)))
x_c .= x
return x_c
end
function aos_to_soa(x::AbstractArray{<:TracedRNumber{T}}) where {T}
for i in eachindex(x)
if !isassigned(x, i)
x[i] = TracedUtils.promote_to(TracedRNumber{T}, 0)
end
end
return Ops.reshape(vcat(x...), size(x)...)
end

include("Ops.jl")
include("TracedUtils.jl")

Expand Down
3 changes: 3 additions & 0 deletions src/TracedUtils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -449,6 +449,9 @@ end
new_traced_value(A::TracedRArray{T,N}) where {T,N} = TracedRArray{T,N}((), nothing, size(A))
new_traced_value(::TracedRNumber{T}) where {T} = TracedRNumber{T}((), nothing)

function broadcast_to_size(arg::AbstractArray{<:TracedRNumber}, rsize)
return broadcast_to_size(reshape(Ops.vcat(arg...), size(arg)...), rsize)
end
broadcast_to_size(arg::AbstractArray, rsize) = broadcast_to_size(Ops.constant(arg), rsize)

function broadcast_to_size(arg::Base.RefValue, rsize)
Expand Down
Loading