Add ConcreteRArray undef ctors and fix KA.allocate #1558
Conversation
oschulz
force-pushed
the
fix-ka-alloc
branch
2 times, most recently
from
f0ec5e8 to
fda9a40
Compare
| device::Union{Nothing,XLA.AbstractDevice} = nothing, | ||
| sharding::Sharding.AbstractSharding = Sharding.NoSharding(), | ||
| ) where {T} | ||
| theclient, thedevice = _select_client_and_device(client, idx, device, sharding) |
There was a problem hiding this comment.
we should just delete the similar(Concrete...) definition and have that fwd to here
There was a problem hiding this comment.
Done (I think, take a look ...).
There was a problem hiding this comment.
we still need to delete
Reactant.jl/src/ConcreteRArray.jl
Line 371 in e4e51e4
There was a problem hiding this comment.
Done.
Question: This is inline - should I make the undef-ctors inline as well?
| ) where {T} | ||
| theclient, thedevice = _select_client_and_device(client, idx, device, sharding) | ||
| # ToDo: How to avoid allocating dummy array on host? | ||
| dummy_array = Array{T}(undef, shape) |
There was a problem hiding this comment.
yeah we dont yet have an api call we can use for that I think [cc @avik-pal ]
There was a problem hiding this comment.
Should I open a seperate issue for that, for future API extensions?
oschulz
force-pushed
the
fix-ka-alloc
branch
3 times, most recently
from
fc1bd5b to
86284c0
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1558 +/- ##
===========================================
+ Coverage 21.66% 43.71% +22.04%
===========================================
Files 46 123 +77
Lines 8048 21987 +13939
===========================================
+ Hits 1744 9611 +7867
- Misses 6304 12376 +6072 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
@wsmoses Ok, tests pass now (approval-dependent ones haven't run yet). I think it's good to go from my side. This PR should be completely non-breaking. Adding client/device/sharding-into to the KernelAbstractions |
oschulz
changed the title
For some reason, using
```
something(client, XLA.default_backend())
```
within `_select_client_and_device` results in `Reactant.XLA.default_backend()`
getting stuck to `Reactant.XLA.PJRT.Client(Ptr{Nothing} @0x0000000000000000)`
if `_select_client_and_device` is called from
`ConcretePJRTArray(data::Array{T,N}; kwargs...)`, while
```
client isa Nothing ? XLA.default_backend() : client
```
works fine. Something effect in package compilation or initialization?
Instead of zero-initializing array.
Avoid array allocation on host.
|
Let's see if I fixed docgen now ... |
|
I added a docstring for |
Add function
ConcretePJRTArray(undef, ...)constructor with implementations forConcretePJRTArrayandConcreteIFRTArray.Simplify implementation of
ConcretePJRTArrayandConcreteIFRTArrayfrom-Arrayconstructors.Fix KA.allocate in ReactantKernelAbstractionsExt.
Improve implementation of
KA.zerosandKA.ones.