FixedSizeArrays support (And Memory Support) #1669
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1669 +/- ##
==========================================
+ Coverage 68.16% 69.82% +1.66%
==========================================
Files 109 111 +2
Lines 11779 12157 +378
==========================================
+ Hits 8029 8489 +460
+ Misses 3750 3668 -82 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
With the latest commit, I'm outputting a It turned out to be more intricate than I initially thought, as I had to add support for In addition, I stumbled on the fact that IFRT is broken for Arrays at least: #1672 Tests are still to be done. If this is shippable I'll start writing tests for it. (and for |
Qfl3x
changed the title
|
Hmm, the builds are failing because they're running on v1.10 and Memory isn't defined there |
|
You can guard the if isdefined(Base, :Memory)
# ...
endRelatedly, note that |
|
I've found a problem with this implementation, it only works in the 1D case. the parent of a In the latest commit, I reshaped the RArray to the shape of the underlying FixedSizeArray. |
Qfl3x
force-pushed
the
FixedSizeArrays_support
branch
from
8625b1a to
658a3d7
Compare
|
Those parts have been refactored into their own functions (with |
probably should be |
I can also reproduce it with julia> using FixedSizeArrays, Reactant
julia> x = FixedSizeArray(randn(3, 4));
julia> rx = Reactant.to_rarray(x);
ERROR: MethodError: no method matching make_buffer_array(::Reactant.XLA.PJRT.Client, ::Memory{Float64}, ::Reactant.XLA.PJRT.Device)
The function `make_buffer_array` exists, but no method is defined for this combination of argument types.
Closest candidates are:
make_buffer_array(::Reactant.XLA.PJRT.Client, ::Array{T, N}, ::Reactant.XLA.PJRT.Device) where {T, N}
@ Reactant ~/.julia/packages/Reactant/eaYLR/src/xla/PJRT/Buffer.jl:9
Stacktrace: |
avik-pal
changed the title
| @inline ConcreteRArray{T}(::UndefInitializer, shape::Integer...; kwargs...) where {T} = ConcreteRArray{ | ||
| T | ||
| }( | ||
| undef, Dims(shape); kwargs... | ||
| ) |
There was a problem hiding this comment.
| @inline ConcreteRArray{T}(::UndefInitializer, shape::Integer...; kwargs...) where {T} = ConcreteRArray{ | |
| T | |
| }( | |
| undef, Dims(shape); kwargs... | |
| ) | |
| @inline ConcreteRArray{T}(::UndefInitializer, shape::Integer...; kwargs...) where {T} = | |
| ConcreteRArray{T}(undef, Dims(shape); kwargs...) |
julia> using FixedSizeArrays, Reactant
julia> M = randn(3, 4);
julia> M_fs = FixedSizeMatrix(M);
julia> rM = Reactant.to_rarray(M);
julia> rM_fs = Reactant.to_rarray(M_fs);
julia> fn(x, y) = sin.(x) .+ cos.(y)
fn (generic function with 1 method)
julia> @code_hlo fn(rM, rM)
module @reactant_fn attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<4x3xf64>) -> tensor<4x3xf64> attributes {enzymexla.memory_effects = []} {
%0 = stablehlo.sine %arg0 : tensor<4x3xf64>
%1 = stablehlo.cosine %arg0 : tensor<4x3xf64>
%2 = stablehlo.add %0, %1 : tensor<4x3xf64>
return %2 : tensor<4x3xf64>
}
}
julia> @code_hlo fn(rM_fs, rM_fs)
module @reactant_fn attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<12xf64>) -> tensor<4x3xf64> attributes {enzymexla.memory_effects = []} {
%0 = stablehlo.reshape %arg0 : (tensor<12xf64>) -> tensor<4x3xf64>
%1 = stablehlo.sine %0 : tensor<4x3xf64>
%2 = stablehlo.cosine %0 : tensor<4x3xf64>
%3 = stablehlo.add %1, %2 : tensor<4x3xf64>
return %3 : tensor<4x3xf64>
}
}Why there's an extra |
|
presumably the memory itself is being stored 1-dimensionally |
|
It is, as it the case for |
| kwargs..., | ||
| ) where {T,N} | ||
| shape = size(prev) | ||
| return reshape( |
There was a problem hiding this comment.
I guess this reshape is the culprit? How does Array work? Is it possible to construct this object directly with the right size instead of reshaping it?
There was a problem hiding this comment.
but array itself we override to keep the dimensionality and allocate everything ourselves for Array.
if the actual tracing of a fixedsizearray does the current "generic recursion into structs" it will eventually allocate a 1-dim memory, always
There was a problem hiding this comment.
The whole point of FixedSizeArray is that the size is...well...fixed. Having to reshape it all the time seems to go into the opposite direction, especially when Array doesn't have that.
There was a problem hiding this comment.
It feels to me like make_tracer should take the size as an (optional) argument. Looking at
Lines 1177 to 1196 in e4bb34f
size(prev) but could be overridden if passed explicitly.
#1649
Struggling with the Memory object. This is what I'm getting currently: