Add back support for buffer donation (#96)

* Redo call to `make_tracer` with `TracedTrack` to mark donated buffers

* Test buffer donation

* Fix `i64` to `f64` type conversion

Author: mofeing

src/utils.jl

@@ -125,12 +125,12 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
         seen_results, result, (:result,), concretein ? TracedTrack : TracedSetPath
     )
 
-    # retraced_args = ntuple(Val(N)) do i
-    #     Base.@_inline_meta
-    #     return make_tracer(
-    #         seen_results, traced_args[i], concretein ? (:resargs, i) : (), TracedTrack
-    #     )
-    # end
+    # marks buffers to be donated
+    for i in 1:N
+        make_tracer(
+            seen_results, traced_args[i], concretein ? (:resargs, i) : (), TracedTrack
+        )
+    end
 
     linear_results = TracedRArray[]
 

test/buffer_donation.jl

@@ -0,0 +1,27 @@
+using Test
+using Reactant
+
+# TODO try again with `2` to check automatic conversion from int to float
+function donate_fill_x_with_2(x, y)
+    x .= 2.0
+    return nothing
+end
+
+function donate_inplace_mul(x, y)
+    x .*= y
+    return nothing
+end
+
+@testset "buffer_donation" begin
+    a = Reactant.ConcreteRArray(ones(2, 2))
+    b = Reactant.ConcreteRArray(3 * ones(2, 2))
+    f = Reactant.compile(donate_fill_x_with_2, (a, b))
+    f(a, b)
+    @test convert(Array, a) == 2 * ones(2, 2)
+
+    a = Reactant.ConcreteRArray(2 * ones(2, 2))
+    b = Reactant.ConcreteRArray(3 * ones(2, 2))
+    f = Reactant.compile(donate_inplace_mul, (a, b))
+    f(a, b)
+    @test convert(Array, a) == 6 * ones(2, 2)
+end

test/runtests.jl

@@ -43,10 +43,11 @@ include("layout.jl")
 include("tracing.jl")
 include("basic.jl")
 include("bcast.jl")
-include("nn.jl")
 include("struct.jl")
 include("closure.jl")
 include("compile.jl")
+include("buffer_donation.jl")
+include("nn.jl")
 
 if VERSION ≥ v"1.10-" # Lux isn't supported on 1.9
     include("nn_lux.jl")