Skip to content

Make RArray a subtype of DenseArray #1696

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

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Make RArray a subtype of DenseArray #1696

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3f715d1
Make `RArray` a subtype of `DenseArray`
albertomercurio Sep 21, 2025
8988f3f
Make `AnyTracedRArray` a `DenseArray` subtype
albertomercurio Sep 22, 2025
04e3d42
Fix permutedims errors
albertomercurio Sep 22, 2025
68a6d96
Fix errors
albertomercurio Sep 22, 2025
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
2 changes: 1 addition & 1 deletion src/ConcreteRArray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ for jlop in (
:(Base.:^),
:(Base.:(==)),
),
T in (AbstractConcreteNumber, AbstractConcreteArray{<:Any,0})
T in (AbstractConcreteNumber, AbstractConcreteArray{<:Number,0})

@eval begin
$(jlop)(x::$(T), y::$(T)) = $(jlop)(to_number(x), to_number(y))
Expand Down
38 changes: 36 additions & 2 deletions src/TracedRArray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,33 @@ Base.strides(x::TracedRArray) = Base.size_to_strides(1, size(x)...)

Base.IndexStyle(::Type{<:TracedRArray}) = Base.IndexLinear()

Base.elsize(::Type{TracedRArray{T,N}}) where {T,N} = sizeof(T)

const ArrayTypesAlias = (
:(TracedRArray),
:(SubArray{<:TracedRNumber,<:Any,<:TracedRArray}),
:(Base.ReshapedArray{<:TracedRNumber,<:Any,<:TracedRArray}),
:(SubArray{<:TracedRNumber,<:Any,<:Base.ReshapedArray{<:TracedRNumber,<:Any,<:TracedRArray}}),
:(Base.ReshapedArray{<:TracedRNumber,<:Any,<:SubArray{<:TracedRNumber,<:Any,<:TracedRArray}}),
)
for ArrayType1 in ArrayTypesAlias
for ArrayType2 in ArrayTypesAlias
@eval Base.mightalias(::$ArrayType1, ::$ArrayType2) = false
Copy link
Member

Choose a reason for hiding this comment

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

this definitely seems wrong since if they are equal they certainly may alias

Copy link
Author

Choose a reason for hiding this comment

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

I'm not familiar with the internal structure of TracedRArray. I could in principle see the implementation for Base arrays if I know where to look in TracedRArray.

end
end

# Base.mightalias(::TracedRArray, ::TracedRArray) = false
# Base.mightalias(
# ::SubArray{<:TracedRNumber,<:Any,<:TracedRArray},
# ::SubArray{<:TracedRNumber,<:Any,<:TracedRArray},
# ) = false
# Base.mightalias(
# ::SubArray{<:TracedRNumber,<:Any,<:TracedRArray}, ::TracedRArray
# ) = false
# Base.mightalias(
# ::TracedRArray, ::SubArray{<:TracedRNumber,<:Any,<:TracedRArray}
# ) = false

# This is required otherwise we will copy a tracedrarray each time
# we use it
Base.convert(T::Type{<:TracedRArray}, x::AbstractArray) = Reactant.promote_to(T, x)
Expand Down Expand Up @@ -265,8 +292,15 @@ function Base.show(io::IOty, X::TracedRArray{T,N}) where {T,N,IOty<:Union{IO,IOC
return print(io, "TracedRArray{", T, ",", N, "N}(", X.paths, ", size=", size(X), ")")
end

function Base.permutedims(A::AnyTracedRArray{T,N}, perm) where {T,N}
return @opcall transpose(materialize_traced_array(A), Int64[perm...])
for ArrayType in (
:(AnyTracedRArray{T,N}),
:(TracedRArray{T,N}),
:(SubArray{<:TracedRNumber{T},N,<:TracedRArray}),
:(Base.ReshapedArray{<:TracedRNumber{T},N,<:TracedRArray})
)
@eval function Base.permutedims(A::$ArrayType, perm) where {T,N}
return @opcall transpose(materialize_traced_array(A), Int64[perm...])
Copy link
Member

Choose a reason for hiding this comment

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

why doesn't anytacedrarray work?

Copy link
Author

Choose a reason for hiding this comment

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

Apparently SubArray and ReshapedArray are not subtypes of AnyTracedRArray

julia> SubArray{<:TracedRNumber,1,<:TracedRArray} <: AnyTracedRArray
false

julia> SubArray{<:TracedRNumber,1,<:TracedRArray} |> supertypes
(SubArray{<:TracedRNumber, 1, <:TracedRArray}, AbstractVector{<:TracedRNumber}, Any)

Copy link
Member

Choose a reason for hiding this comment

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

Cc @avik-pal i thought anytracedrarray was just an abstractarray of a traced number?

end
end

for (jlop, hloop, hlocomp, merge) in
Expand Down
7 changes: 6 additions & 1 deletion src/Types.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ abstract type RNumber{T<:ReactantPrimitive} <: Number end

abstract type AbstractConcreteNumber{T} <: RNumber{T} end

abstract type RArray{T,N} <: AbstractArray{T,N} end
abstract type RArray{T,N} <: DenseArray{T,N} end

abstract type AbstractConcreteArray{T,N} <: RArray{T,N} end

Expand Down Expand Up @@ -52,6 +52,11 @@ mutable struct TracedRNumber{T} <: RNumber{T}
end
end

Base.elsize(::Type{TracedRNumber{T}}) where {T} = sizeof(T)
Base.elsize(::Type{RNumber{T}}) where {T} = sizeof(T)
Base.elsize(::Type{<:AbstractConcreteNumber{T}}) where {T} = sizeof(T)
Base.elsize(::Type{<:AbstractConcreteArray{T}}) where {T} = sizeof(T)

function repath(x::TracedRNumber{T}, paths) where {T}
return TracedRNumber{T}(paths, x.mlir_data)
end
Expand Down
Loading