feat: add ifrt copy api (#1334)

* feat: ifrt copy api (julia bindings)

* chore: bump jll

Author: avik-pal

Project.toml

@@ -87,7 +87,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.10"
-Reactant_jll = "0.0.187"
+Reactant_jll = "0.0.188"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/ConcreteRArray.jl

@@ -23,6 +23,22 @@ function Base.copy(x::Union{AbstractConcreteArray,AbstractConcreteNumber})
     return fn(x)
 end
 
+function Base.copy(X::ConcreteIFRTArray{T,D,S,P}) where {T,D,S,P}
+    return ConcreteIFRTArray{T,D,S}(Base.copy(X.data), X.shape, X.sharding, X.padding)
+end
+
+function Base.copy(X::ConcretePJRTArray)
+    return Core.Typeof(X)(Base.copy.(X.data), X.shape, X.sharding)
+end
+
+function Base.copy(X::ConcreteIFRTNumber)
+    return Core.Typeof(X)(Base.copy(X.data), X.sharding)
+end
+
+function Base.copy(X::ConcretePJRTNumber)
+    return Core.Typeof(X)(Base.copy.(X.data), X.sharding)
+end
+
 # deepcopy
 function Base.deepcopy(x::Union{AbstractConcreteArray,AbstractConcreteNumber})
     return Base.copy(x)
@@ -117,14 +133,6 @@ function write_to_host_buffer!(data::Array, X::ConcretePJRTArray{T,N}) where {T,
     return nothing
 end
 
-function Base.copy(X::ConcretePJRTArray)
-    return Core.Typeof(X)(Base.copy.(X.data), X.shape, X.sharding)
-end
-
-function Base.copy(X::ConcretePJRTNumber)
-    return 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/IFRT/Array.jl

@@ -424,3 +424,11 @@ function copy_arrays_to_device_with_sharding(buffers::Vector{Array}, sharding::S
     end
     return dst_arrays
 end
+
+function Base.copy(b::Array)
+    GC.@preserve b begin
+        return Array(
+            @ccall MLIR.API.mlir_c.ifrt_copy_array(b.buffer::Ptr{Cvoid})::Ptr{Cvoid}
+        )
+    end
+end

src/xla/IFRT/AsyncArray.jl

@@ -25,3 +25,8 @@ function XLA.to_host(array::AsyncArray, data, reactant_sharding)
 end
 
 XLA.sharding(x::AsyncArray) = XLA.sharding(x.buffer)
+
+function Base.copy(b::AsyncArray)
+    Base.wait(b)
+    return AsyncArray(Base.copy(b.buffer), nothing)
+end