feat: add a reactant.donated attr to donated args (#947)

Author: avik-pal

src/Compiler.jl

@@ -1055,6 +1055,17 @@ function compile_mlir!(
     MLIR.API.mlirOperationDestroy(compiled_f.operation)
     compiled_f.operation = MLIR.API.MlirOperation(C_NULL)
 
+    # Add a `donated` attr to the function arguments. This doesn't affect XLA, but lets us
+    # check which arguments were donated.
+    preserved_args_idx = last.(preserved_args)
+    for (i, arg) in enumerate(linear_args)
+        if i ∉ preserved_args_idx
+            MLIR.API.mlirFuncSetArgAttr(
+                func3, i - 1, "reactant.donated", MLIR.IR.UnitAttribute()
+            )
+        end
+    end
+
     return Reactant.TracedUtils.CompiledMlirFnResult(
         fnwrapped,
         func3,

test/buffer_donation.jl

@@ -17,6 +17,8 @@ end
     b = Reactant.to_rarray(3 * ones(2, 2))
     @jit(donate_fill_x_with_2(a, b))
     @test convert(Array, a) == 2 * ones(2, 2)
+    hlo = @code_hlo(donate_fill_x_with_2(a, b))
+    @test length(findall("reactant.donated", repr(hlo))) == 1
 
     (; preserved_args) = Reactant.Compiler.compile_xla(donate_fill_x_with_2, (a, b))[3]
     preserved_args_idx = last.(preserved_args)
@@ -26,6 +28,8 @@ end
     b = Reactant.to_rarray(3 * ones(2, 2))
     @jit(donate_inplace_mul(a, b))
     @test convert(Array, a) == 6 * ones(2, 2)
+    hlo = @code_hlo(donate_inplace_mul(a, b))
+    @test length(findall("reactant.donated", repr(hlo))) == 1
 
     (; preserved_args) = Reactant.Compiler.compile_xla(donate_inplace_mul, (a, b))[3]
     preserved_args_idx = last.(preserved_args)