From 5e23fa322fd37321b2187d00a93ad319f7032ca7 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Sun, 3 Aug 2025 18:01:37 -0500 Subject: [PATCH 1/5] More precompile memory fixups --- ext/ReactantCUDAExt.jl | 12 +++++++++++- src/Precompile.jl | 23 +++++++++++++++++++++-- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl index 95894881bc..24781a4928 100644 --- a/ext/ReactantCUDAExt.jl +++ b/ext/ReactantCUDAExt.jl @@ -1489,7 +1489,17 @@ end end y = Reactant.ConcreteRArray([2.0]; client) Reactant.Compiler.compile_mlir(square!, (y,); optimize=false) - finalize(y) + + if y isa ConcreteIFRTArray + XLA.free_buffer(y.data.buffer) + y.data.buffer.buffer = C_NULL + else + ntuple(Val(len(y.data))) do i + Base.@_inline_meta + XLA.free_buffer(y.data[i].buffer) + y.data[i].buffer.buffer = C_NULL + end + end end end diff --git a/src/Precompile.jl b/src/Precompile.jl index aa0a278c3b..3887f21ec6 100644 --- a/src/Precompile.jl +++ b/src/Precompile.jl @@ -69,11 +69,30 @@ if Reactant_jll.is_available() @static if precompilation_supported() x = ConcreteRNumber(2.0; client) Reactant.compile(sin, (x,); client, optimize=:all) - finalize(x) + if x isa ConcreteIFRTArray + XLA.free_buffer(x.data.buffer) + x.data.buffer.buffer = C_NULL + else + ntuple(Val(len(y.data))) do i + Base.@_inline_meta + XLA.free_buffer(y.data[i].buffer) + y.data[i].buffer.buffer = C_NULL + end + end y = ConcreteRArray([2.0]; client) Reactant.compile(Base.sum, (y,); client, optimize=:all) - finalize(y) + if y isa ConcreteIFRTArray + XLA.free_buffer(y.data.buffer) + y.data.buffer.buffer = C_NULL + else + ntuple(Val(len(y.data))) do i + Base.@_inline_meta + XLA.free_buffer(y.data[i].buffer) + y.data[i].buffer.buffer = C_NULL + end + end + end end From ea934b2deb132ebb21633b10a5fdf384f3c85542 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Sun, 3 Aug 2025 18:05:31 -0500 Subject: [PATCH 2/5] fix --- ext/ReactantCUDAExt.jl | 2 +- src/Precompile.jl | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl index 24781a4928..238c68e9a1 100644 --- a/ext/ReactantCUDAExt.jl +++ b/ext/ReactantCUDAExt.jl @@ -1494,7 +1494,7 @@ end XLA.free_buffer(y.data.buffer) y.data.buffer.buffer = C_NULL else - ntuple(Val(len(y.data))) do i + ntuple(Val(length(y.data))) do i Base.@_inline_meta XLA.free_buffer(y.data[i].buffer) y.data[i].buffer.buffer = C_NULL diff --git a/src/Precompile.jl b/src/Precompile.jl index 3887f21ec6..ae4dcb526b 100644 --- a/src/Precompile.jl +++ b/src/Precompile.jl @@ -73,10 +73,10 @@ if Reactant_jll.is_available() XLA.free_buffer(x.data.buffer) x.data.buffer.buffer = C_NULL else - ntuple(Val(len(y.data))) do i + ntuple(Val(length(x.data))) do i Base.@_inline_meta - XLA.free_buffer(y.data[i].buffer) - y.data[i].buffer.buffer = C_NULL + XLA.free_buffer(x.data[i].buffer) + x.data[i].buffer.buffer = C_NULL end end @@ -86,7 +86,7 @@ if Reactant_jll.is_available() XLA.free_buffer(y.data.buffer) y.data.buffer.buffer = C_NULL else - ntuple(Val(len(y.data))) do i + ntuple(Val(length(y.data))) do i Base.@_inline_meta XLA.free_buffer(y.data[i].buffer) y.data[i].buffer.buffer = C_NULL From 54a62101327ca1d9f13a4c9668cc3af64d6fe16e Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Sun, 3 Aug 2025 18:06:36 -0500 Subject: [PATCH 3/5] fix --- src/Precompile.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Precompile.jl b/src/Precompile.jl index ae4dcb526b..80bcb944e1 100644 --- a/src/Precompile.jl +++ b/src/Precompile.jl @@ -69,7 +69,7 @@ if Reactant_jll.is_available() @static if precompilation_supported() x = ConcreteRNumber(2.0; client) Reactant.compile(sin, (x,); client, optimize=:all) - if x isa ConcreteIFRTArray + if x isa ConcreteIFRTNumber XLA.free_buffer(x.data.buffer) x.data.buffer.buffer = C_NULL else From 1b1557732723124d203889e7ef605203d9c89d3e Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Sun, 3 Aug 2025 18:10:42 -0500 Subject: [PATCH 4/5] fix --- ext/ReactantCUDAExt.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl index 238c68e9a1..893c5eb5ec 100644 --- a/ext/ReactantCUDAExt.jl +++ b/ext/ReactantCUDAExt.jl @@ -1490,13 +1490,13 @@ end y = Reactant.ConcreteRArray([2.0]; client) Reactant.Compiler.compile_mlir(square!, (y,); optimize=false) - if y isa ConcreteIFRTArray - XLA.free_buffer(y.data.buffer) + if y isa Reactant.ConcreteIFRTArray + Reactant.XLA.free_buffer(y.data.buffer) y.data.buffer.buffer = C_NULL else ntuple(Val(length(y.data))) do i Base.@_inline_meta - XLA.free_buffer(y.data[i].buffer) + Reactant.XLA.free_buffer(y.data[i].buffer) y.data[i].buffer.buffer = C_NULL end end From fcafc72cdf04e3427e1187b48f0586566d7cd041 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Sun, 3 Aug 2025 18:45:50 -0500 Subject: [PATCH 5/5] cleanup --- ext/ReactantCUDAExt.jl | 7 +++---- src/Precompile.jl | 15 ++++++--------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl index 893c5eb5ec..25ae24a5e3 100644 --- a/ext/ReactantCUDAExt.jl +++ b/ext/ReactantCUDAExt.jl @@ -1494,10 +1494,9 @@ end Reactant.XLA.free_buffer(y.data.buffer) y.data.buffer.buffer = C_NULL else - ntuple(Val(length(y.data))) do i - Base.@_inline_meta - Reactant.XLA.free_buffer(y.data[i].buffer) - y.data[i].buffer.buffer = C_NULL + for dat in y.data + Reactant.XLA.free_buffer(dat.buffer) + dat.buffer.buffer = C_NULL end end end diff --git a/src/Precompile.jl b/src/Precompile.jl index 80bcb944e1..1b8d0901fd 100644 --- a/src/Precompile.jl +++ b/src/Precompile.jl @@ -73,10 +73,9 @@ if Reactant_jll.is_available() XLA.free_buffer(x.data.buffer) x.data.buffer.buffer = C_NULL else - ntuple(Val(length(x.data))) do i - Base.@_inline_meta - XLA.free_buffer(x.data[i].buffer) - x.data[i].buffer.buffer = C_NULL + for dat in x.data + XLA.free_buffer(dat.buffer) + dat.buffer.buffer = C_NULL end end @@ -86,13 +85,11 @@ if Reactant_jll.is_available() XLA.free_buffer(y.data.buffer) y.data.buffer.buffer = C_NULL else - ntuple(Val(length(y.data))) do i - Base.@_inline_meta - XLA.free_buffer(y.data[i].buffer) - y.data[i].buffer.buffer = C_NULL + for dat in y.data + XLA.free_buffer(dat.buffer) + dat.buffer.buffer = C_NULL end end - end end