feat: add specialization of copyto! from julia array to concreterarray (#942)

* add specialization of copyto! from julia array to concreterarray

* support subarray et. al. as well

* Update src/ConcreteRArray.jl

Co-authored-by: Gregory L. Wagner <[email protected]>

* Update src/ConcreteRArray.jl

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

* fix: proper copyto

* fix: restrict types

---------

Co-authored-by: Gregory L. Wagner <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Avik Pal <[email protected]>

Author: 4 people

lib/ReactantCore/src/ReactantCore.jl

@@ -231,7 +231,11 @@ function trace_for(mod, expr; track_numbers)
                 end
 
             $(ReactantCore).traced_while(
-                cond_fn, body_fn, args; track_numbers=$(track_numbers), verify_arg_names=$(QuoteNode(args_init))
+                cond_fn,
+                body_fn,
+                args;
+                track_numbers=$(track_numbers),
+                verify_arg_names=$(QuoteNode(args_init)),
             )
         end
     end

src/ConcreteRArray.jl

@@ -398,18 +398,25 @@ function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteI
     return fn(bc.args...)
 end
 
-# XXX: This is not necessarily correct. We need to check for sharding and also device
-#      compatibility.
-function Base.copyto!(dest::AbstractConcreteArray, src::AbstractConcreteArray)
-    dest.data = src.data
-    return dest
-end
-
 function mycopyto!(dest, src)
     dest .= src # use broadcasting instead of copyto!
     return nothing
 end
 
+for aType in (:ConcretePJRTArray, :ConcreteIFRTArray)
+    @eval function Base.copyto!(dest::$(aType), src::$(aType))
+        if dest.sharding == src.sharding &&
+            XLA.device(dest) == XLA.device(src) &&
+            XLA.client(dest) == XLA.client(src)
+            dest.data = src.data
+        else
+            fn = compile(mycopyto!, (dest, src))
+            fn(dest, src)
+        end
+        return dest
+    end
+end
+
 function Base.copyto!(
     dest::Union{AnyConcreteIFRTArray,AnyConcretePJRTArray}, src::AbstractConcreteArray
 )
@@ -418,6 +425,22 @@ function Base.copyto!(
     return dest
 end
 
+for aType in (:ConcretePJRTArray, :ConcreteIFRTArray)
+    anyaType = Symbol(:Any, aType)
+    @eval function Base.copyto!(dest::$(anyaType), src::Array{<:ReactantPrimitive})
+        ancestor_dest = ancestor(dest)
+        return copyto!(
+            dest,
+            $(aType)(
+                src;
+                sharding=ancestor_dest.sharding,
+                client=XLA.client(ancestor_dest),
+                device=XLA.device(ancestor_dest),
+            ),
+        )
+    end
+end
+
 for aType in (:ConcretePJRTArray, :ConcreteIFRTArray)
     @eval begin
         function Base.copyto!(

src/Ops.jl

@@ -1791,7 +1791,7 @@ end
             return_dialect=:stablehlo,
             args_in_result=:none,
             do_transpose=false,
-            verify_arg_names
+            verify_arg_names,
         ).f
 
     cond_reg = Reactant.TracedUtils.__take_region(cond_fn_compiled)

src/TracedUtils.jl

@@ -356,9 +356,11 @@ function make_mlir_fn(
         conflicts = setdiff(resis, argis)
         @assert !isempty(conflicts) "Expected to have some conflicts, but none were found."
 
-        error("""Types do not match between function arguments and results.
-        The following arguments should be traced: $(join(verify_arg_names.args[collect(conflicts)], ", "))
-        """)
+        error(
+            """Types do not match between function arguments and results.
+      The following arguments should be traced: $(join(verify_arg_names.args[collect(conflicts)], ", "))
+      """,
+        )
     end
 
     out_tys = if do_transpose