Make RArray a subtype of DenseArray (#1696)

* Make `RArray` a subtype of `DenseArray`

* Make `AnyTracedRArray` a `DenseArray` subtype

* Fix permutedims errors

* Fix errors

* Simplify method definitions

* Add _parentsmatch method

Author: albertomercurio

src/ConcreteRArray.jl

@@ -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))

src/TracedRArray.jl

@@ -23,6 +23,8 @@ 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)
+
 # 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)
@@ -265,10 +267,6 @@ 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...])
-end
-
 for (jlop, hloop, hlocomp, merge) in
     ((:(Base.:(==)), :compare, "EQ", :all), (:(Base.:(!=)), :compare, "NE", :any))
     @eval function $jlop(
@@ -279,6 +277,17 @@ for (jlop, hloop, hlocomp, merge) in
     end
 end
 
+# Override _parentsmatch to avoid pointer comparisons during tracing
+# Direct TracedRArray comparisons - they don't alias unless they're the same object
+Base._parentsmatch(A::TracedRArray, B::TracedRArray) = A === B
+# ReshapedArray comparisons - check if they share the same parent (more specific than StridedArray)
+function Base._parentsmatch(
+    A::Base.ReshapedArray{<:TracedRNumber,<:Any,<:Union{TracedRArray,SubArray{<:TracedRNumber,<:Any,<:TracedRArray}}},
+    B::Base.ReshapedArray{<:TracedRNumber,<:Any,<:Union{TracedRArray,SubArray{<:TracedRNumber,<:Any,<:TracedRArray}}}
+)
+    return Base._parentsmatch(parent(A), parent(B))
+end
+
 function __default_init(
     ::Type{T}, ::Union{typeof(Base.min),typeof(Base.FastMath.min_fast)}
 ) where {T}
@@ -1348,4 +1357,15 @@ function unrolled_map(f::F, itr) where {F}
     return result
 end
 
+# permutedims for TracedRArrays and wrappers
+function Base.permutedims(A::AnyTracedRArray{T,N}, perm) where {T,N}
+    return @opcall transpose(materialize_traced_array(A), Int64[perm...])
+end
+
+function Base.permutedims!(dest::TracedRArray, src::AnyTracedRArray, perm)
+    result = @opcall transpose(materialize_traced_array(src), Int64[perm...])
+    TracedUtils.set_mlir_data!(dest, result.mlir_data)
+    return dest
+end
+
 end

src/Types.jl

@@ -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
 
@@ -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