Fast pjrt copy (#1327)

* Fast pjrt copy

* number

* actually copy

* Update src/xla/PJRT/AsyncBuffer.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

src/ConcreteRArray.jl

@@ -25,8 +25,7 @@ end
 
 # deepcopy
 function Base.deepcopy(x::Union{AbstractConcreteArray,AbstractConcreteNumber})
-    fn = Reactant.compile(copy, (x,))
-    return fn(x)
+    Base.copy(x)
 end
 
 # One more reason why users shouldn't call `deepcopy`
@@ -118,6 +117,14 @@ function write_to_host_buffer!(data::Array, X::ConcretePJRTArray{T,N}) where {T,
     return nothing
 end
 
+function Base.copy(X::ConcretePJRTArray)
+    Core.Typeof(X)(Base.copy.(X.data), X.shape, X.sharding)
+end
+
+function Base.copy(X::ConcretePJRTNumber)
+    Core.Typeof(X)(Base.copy.(X.data), X.sharding)
+end
+
 function write_to_host_buffer!(data::Array, X::ConcreteIFRTArray{T,N}) where {T,N}
     XLA.to_host(X.data, data, X.sharding)
     return nothing

src/xla/PJRT/AsyncBuffer.jl

@@ -6,3 +6,8 @@ end
 const AsyncEmptyBuffer = AsyncBuffer(Buffer(C_NULL), nothing)
 
 AsyncBuffer(args...; kwargs...) = AsyncBuffer(Buffer(args...; kwargs...), nothing)
+
+function Base.copy(b::AsyncBuffer)
+    Base.wait(b)
+    return AsyncBuffer(Base.copy(b.buffer), nothing)
+end

src/xla/PJRT/Buffer.jl

@@ -87,6 +87,17 @@ function XLA.unsafe_buffer_pointer(buffer::Buffer)
     @ccall MLIR.API.mlir_c.UnsafeBufferPointer(buffer.buffer::Ptr{Cvoid})::Ptr{Cvoid}
 end
 
+function Base.copy(buffer::Buffer)
+    dev = XLA.device(buffer)
+    GC.@preserve buffer dev begin
+        Buffer(
+            @ccall MLIR.API.mlir_c.CopyBufferToDevice(
+                buffer.buffer::Ptr{Cvoid}, dev.device::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
 function XLA.copy_buffer_to_device(buffer::Buffer, dev::Device)
     XLA.device(buffer) == dev && return buffer
     GC.@preserve buffer dev begin