fix: force finalize the buffers after precompile (#1522)

Author: avik-pal

ext/ReactantCUDAExt.jl

@@ -1489,6 +1489,7 @@ end
                 end
                 y = Reactant.ConcreteRArray([2.0]; client)
                 Reactant.Compiler.compile_mlir(square!, (y,); optimize=false)
+                finalize(y)
             end
         end
 

src/Precompile.jl

@@ -69,9 +69,11 @@ if Reactant_jll.is_available()
             @static if precompilation_supported()
                 x = ConcreteRNumber(2.0; client)
                 Reactant.compile(sin, (x,); client, optimize=:all)
+                finalize(x)
 
                 y = ConcreteRArray([2.0]; client)
                 Reactant.compile(Base.sum, (y,); client, optimize=:all)
+                finalize(y)
             end
         end
 

src/xla/Buffer.jl

@@ -1,5 +1,6 @@
 abstract type AbstractBuffer end
 
+function free_buffer end
 function synced_buffer end
 function buffer_on_cpu end
 function to_host end

src/xla/IFRT/Array.jl

@@ -3,7 +3,7 @@ mutable struct Array <: XLA.AbstractBuffer
 
     function Array(buffer::Ptr{Cvoid}, owned::Bool=true)
         !owned && return new(buffer)
-        return finalizer(free_ifrt_array, new(buffer))
+        return finalizer(XLA.free_buffer, new(buffer))
     end
 end
 
@@ -158,7 +158,7 @@ function Array(
     return Array(client, array, ifrt_sharding)
 end
 
-@inline function free_ifrt_array(buffer::Array)
+@inline function XLA.free_buffer(buffer::Array)
     if buffer.buffer != C_NULL
         @ccall MLIR.API.mlir_c.ifrt_free_array(buffer.buffer::Ptr{Cvoid})::Cvoid
     end

src/xla/PJRT/Buffer.jl

@@ -2,7 +2,7 @@ mutable struct Buffer <: XLA.AbstractBuffer
     buffer::Ptr{Cvoid}
 
     function Buffer(buffer::Ptr{Cvoid})
-        return finalizer(free_buffer, new(buffer))
+        return finalizer(XLA.free_buffer, new(buffer))
     end
 end
 
@@ -114,7 +114,7 @@ function Base.similar(a::Buffer, S::Type, dims::Dims)
     return Base.similar(Buffer, S, dims; client=XLA.client(a), device=XLA.device(a))
 end
 
-@inline function free_buffer(buffer::Buffer)
+@inline function XLA.free_buffer(buffer::Buffer)
     sbuffer = buffer.buffer
     if sbuffer != C_NULL
         @ccall MLIR.API.mlir_c.PjRtBufferFree(sbuffer::Ptr{Cvoid})::Cvoid