fix: codegen to fix aliasing issues (#1115)

* fix: codegen to fix aliasing issues

* fix: incorrect logic

* test: more tests

* fix: overload deepcopy_internal

Author: avik-pal

src/Compiler.jl

@@ -1262,8 +1262,10 @@ function compile_mlir!(
     nresults = MLIR.IR.Value[]
     linear_results2 = TracedType[]
     results_mask = falses(length(results))
+
     for (i, op) in enumerate(results)
-        if !MLIR.IR.is_block_arg(op)
+        if !MLIR.IR.is_block_arg(op) ||
+            !Reactant.TracedUtils.has_idx(linear_results[i], :args) # new buffer
             push!(nresults, op)
             push!(linear_results2, linear_results[i])
             results_mask[i] = true

src/ConcreteRArray.jl

@@ -17,6 +17,28 @@ for runtime in (:PJRT, :IFRT)
     end
 end
 
+# copy
+function Base.copy(x::Union{AbstractConcreteArray,AbstractConcreteNumber})
+    fn = Reactant.compile(copy, (x,))
+    return fn(x)
+end
+
+# deepcopy
+function Base.deepcopy(x::Union{AbstractConcreteArray,AbstractConcreteNumber})
+    fn = Reactant.compile(copy, (x,))
+    return fn(x)
+end
+
+# One more reason why users shouldn't call `deepcopy`
+function Base.deepcopy_internal(
+    x::Union{AbstractConcreteArray,AbstractConcreteNumber}, stackdict::IdDict
+)
+    if haskey(stackdict, x)
+        return stackdict[x]::typeof(x)
+    end
+    return deepcopy(x)
+end
+
 Base.size(::AbstractConcreteNumber) = ()
 Base.real(x::AbstractConcreteNumber{<:Real}) = x
 function Base.rtoldefault(T::Type{<:AbstractConcreteNumber})
@@ -410,14 +432,10 @@ 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
+        # We can't directly set the data field. it will alias the inner buffers without
+        # actually copying them.
+        fn = compile(mycopyto!, (dest, src))
+        fn(dest, src)
         return dest
     end
 end

test/basic.jl

@@ -1117,3 +1117,37 @@ end
         @test_throws Reactant.XLA.ReactantInternalError Reactant.XLA.allocatorstats()
     end
 end
+
+@testset "copy/deepcopy" begin
+    for op in (copy, deepcopy)
+        x = Reactant.to_rarray(ones(4, 4))
+        if x isa Reactant.ConcretePJRTArray
+            orig_ptr = only(x.data).buffer.buffer
+            y = op(x)
+            @test y isa Reactant.ConcretePJRTArray
+            @test only(y.data).buffer.buffer != orig_ptr
+            @test only(x.data).buffer.buffer == orig_ptr
+        else
+            orig_ptr = x.data.buffer.buffer
+            y = op(x)
+            @test y isa Reactant.ConcreteIFRTArray
+            @test y.data.buffer.buffer != orig_ptr
+            @test x.data.buffer.buffer == orig_ptr
+        end
+
+        x = Reactant.to_rarray(4.0; track_numbers=Number)
+        if x isa Reactant.ConcretePJRTNumber
+            orig_ptr = only(x.data).buffer.buffer
+            y = op(x)
+            @test y isa Reactant.ConcretePJRTNumber
+            @test only(y.data).buffer.buffer != orig_ptr
+            @test only(x.data).buffer.buffer == orig_ptr
+        else
+            orig_ptr = x.data.buffer.buffer
+            y = op(x)
+            @test y isa Reactant.ConcreteIFRTNumber
+            @test y.data.buffer.buffer != orig_ptr
+            @test x.data.buffer.buffer == orig_ptr
+        end
+    end
+end