From 8846a8ce4984db28c91f9e812f3e3526581ae322 Mon Sep 17 00:00:00 2001 From: krasow Date: Tue, 28 Jul 2026 11:35:58 -0500 Subject: [PATCH 01/12] inter-broadcast kernel fusion. Leverage lifetime analysis to rewrite expr blocks to maximize fusion and reducing temps --- src/scoping.jl | 176 +++++++++++++++++++++++++-- test/tests/broadcast_fusion_tests.jl | 14 +++ test/tests/scoping.jl | 14 +++ 3 files changed, 196 insertions(+), 8 deletions(-) diff --git a/src/scoping.jl b/src/scoping.jl index 47ae01cea..707253a49 100644 --- a/src/scoping.jl +++ b/src/scoping.jl @@ -114,7 +114,8 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) return v end - is_indexed_assign(s) = s isa Expr && s.head == :(=) && !(s.args[1] isa Symbol) + is_indexed_assign(s) = + s isa Expr && s.head in (:(=), :(.=)) && !(s.args[1] isa Symbol) result_symbol(s) = if s isa Symbol s @@ -312,6 +313,129 @@ end is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") +function _is_broadcast_syntax(e) + return e isa Expr && + ( + (e.head == :call && !isempty(e.args) && is_broadcast_op(e.args[1])) || + (e.head == :. && length(e.args) == 2) + ) +end + +function _symbol_occurrences(x, target::Symbol) + x === target && return 1 + x isa Expr || return 0 + return sum(arg -> _symbol_occurrences(arg, target), x.args; init=0) +end + +function _substitute_symbols(x, replacements::Dict{Symbol,Any}) + x isa Symbol && return get(replacements, x, x) + x isa Expr || return x + + # A symbol assignment introduces/redefines its LHS; only substitute in RHS. + if x.head == :(=) && x.args[1] isa Symbol + return Expr(:(=), x.args[1], _substitute_symbols(x.args[2], replacements)) + end + return Expr(x.head, (_substitute_symbols(arg, replacements) for arg in x.args)...) +end + +function _indexed_assignment_base(stmt) + if stmt isa Expr && stmt.head == :(=) && + stmt.args[1] isa Expr && stmt.args[1].head == :ref + base = stmt.args[1].args[1] + return base isa Symbol ? base : nothing + end + return nothing +end + +function _safe_to_delay_broadcast( + stmts, def_idx::Int, use_idx::Int, dependencies::Set{Symbol}, lazy_defs::Set{Int} +) + for i in (def_idx + 1):(use_idx - 1) + stmt = stmts[i] + stmt isa LineNumberNode && continue + i in lazy_defs && continue + + # An indexed write to an unrelated array does not invalidate the lazy + # producer. Any other intervening statement is conservatively a barrier. + mutated = _indexed_assignment_base(stmt) + mutated !== nothing && !(mutated in dependencies) && continue + return false + end + return true +end + +""" +Inline single-use broadcast temporaries across statements before lifetime +hoisting. This lets a stencil such as + + lap = ... + out[slice] = scale .* lap .+ ... + +become one broadcast tree writing directly to `out[slice]`. + +Only broadcast definitions are delayed, and only across other lazy definitions +or indexed writes to arrays not referenced by the producer. +""" +function inline_single_use_broadcasts(block) + Meta.isexpr(block, (:block, :begin)) || return block + stmts = collect(block.args) + + definitions = Dict{Symbol,Tuple{Int,Any}}() + lazy_defs = Set{Int}() + for (i, stmt) in enumerate(stmts) + if stmt isa Expr && stmt.head == :(=) && stmt.args[1] isa Symbol && + _is_broadcast_syntax(stmt.args[2]) + definitions[stmt.args[1]] = (i, stmt.args[2]) + push!(lazy_defs, i) + end + end + + inlineable = Dict{Symbol,Tuple{Int,Int,Any}}() + for (sym, (def_idx, rhs)) in definitions + use_indices = Int[] + occurrences = 0 + for i in (def_idx + 1):length(stmts) + n = _symbol_occurrences(stmts[i], sym) + if n > 0 + occurrences += n + push!(use_indices, i) + end + end + occurrences == 1 || continue + use_idx = only(use_indices) + deps = Set(walk_symbols(rhs)) + _safe_to_delay_broadcast(stmts, def_idx, use_idx, deps, lazy_defs) || continue + inlineable[sym] = (def_idx, use_idx, rhs) + end + + replacements = Dict{Symbol,Any}() + removed = Set(first(info) for info in values(inlineable)) + out = Any[] + for (i, stmt) in enumerate(stmts) + i in removed && begin + sym = only(sym for (sym, info) in inlineable if first(info) == i) + replacements[sym] = _substitute_symbols(inlineable[sym][3], replacements) + continue + end + + stmt = _substitute_symbols(stmt, replacements) + + # A materialized indexed assignment with a broadcast RHS normally + # allocates a temporary and copies it into the slice. If the destination + # base is absent from the RHS, write the fused tree directly to the view. + if stmt isa Expr && stmt.head == :(=) && + stmt.args[1] isa Expr && stmt.args[1].head == :ref && + _is_broadcast_syntax(stmt.args[2]) + base = stmt.args[1].args[1] + if base isa Symbol && !(base in Set(walk_symbols(stmt.args[2]))) + stmt = Expr(:(.=), stmt.args[1], stmt.args[2]) + end + end + push!(out, stmt) + end + return Expr(block.head, out...) +end + function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) local_assigned = Set{Symbol}() @@ -335,16 +459,32 @@ function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) # Inside a broadcast tree: hoist slices, keep dotted ops/f.(…) lazy, and # delegate anything else to rewrite() (it breaks the tree → real NDArray). - function rewrite_bcast(e)::Tuple{Any,Vector{Expr}} + # + # The slice cache is deliberately scoped to one fused tree. Repeated views + # then become one task argument and are still destroyed immediately after + # that task's last use. Sharing the cache across trees would keep transformed + # stores alive across task submissions, which can force runtime copies and + # delay device-memory reclamation. + function rewrite_bcast(e, slice_cache::Dict{Any,Symbol})::Tuple{Any,Vector{Expr}} e isa Expr || return e, Expr[] - e.head == :ref && return fresh_tmp(e) + if e.head == :ref + cached = get(slice_cache, e, nothing) + cached === nothing || return cached, Expr[] + tmp, bind = fresh_tmp(e) + slice_cache[e] = tmp + return tmp, bind + end if e.head == :call && is_broadcast_op(e.args[1]) - args, hoisted = maphoist(rewrite_bcast, e.args[2:end]) + args, hoisted = maphoist( + x -> rewrite_bcast(x, slice_cache), e.args[2:end] + ) return Expr(:call, e.args[1], args...), hoisted end if e.head == :. && length(e.args) == 2 && e.args[2] isa Expr && e.args[2].head == :tuple - args, hoisted = maphoist(rewrite_bcast, e.args[2].args) + args, hoisted = maphoist( + x -> rewrite_bcast(x, slice_cache), e.args[2].args + ) return Expr(:., e.args[1], Expr(:tuple, args...)), hoisted end return rewrite(e) @@ -363,8 +503,27 @@ function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) # .= RHS is a broadcast tree: only the slices inside it are hoisted. if e.head == :(.=) lhs, rhs = e.args - new_lhs, lts = rewrite(lhs) - new_rhs, rts = rewrite_bcast(rhs) + # Preserve Julia's indexed `.=` semantics. Hoisting `A[inds...]` + # directly is valid for NDArray (whose getindex returns a view), but + # produces a detached copy for Array. An explicit view works for + # both, while still giving the lifetime pass a concrete NDArray + # wrapper to destroy immediately after task submission. + new_lhs, lts = + if lhs isa Expr && lhs.head == :ref + view_expr = Base.macroexpand( + @__MODULE__, + Expr( + :macrocall, + GlobalRef(Base, Symbol("@view")), + LineNumberNode(0), + lhs, + ), + ) + fresh_tmp(view_expr) + else + rewrite(lhs) + end + new_rhs, rts = rewrite_bcast(rhs, Dict{Any,Symbol}()) return Expr(:(.=), new_lhs, new_rhs), vcat(lts, rts) end @@ -372,7 +531,7 @@ function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) # Broadcast root: hoist the fused result as one temp. if e.head == :call && is_broadcast_op(e.args[1]) - inner, hoisted = rewrite_bcast(e) + inner, hoisted = rewrite_bcast(e, Dict{Any,Symbol}()) tmp, bind = fresh_tmp(inner) return tmp, vcat(hoisted, bind) end @@ -411,6 +570,7 @@ end function process_broadcast_scope(block) assigned_vars = Set{Symbol}() + block = inline_single_use_broadcasts(block) rewritten = find_broadcast_assignments(block, assigned_vars) result = insert_finalizers(rewritten, assigned_vars) counter[] = 0 diff --git a/test/tests/broadcast_fusion_tests.jl b/test/tests/broadcast_fusion_tests.jl index ceaa84593..a4d0716c4 100644 --- a/test/tests/broadcast_fusion_tests.jl +++ b/test/tests/broadcast_fusion_tests.jl @@ -430,6 +430,20 @@ function test_broadcast_fusion_edge_cases(; T=Float32, atol=1e-5, rtol=1e-5) a2 .= a2 .* s1 .+ b @allowscalar @test cuNumeric.compare(ja .* s1 .+ jb, a2, atol, rtol) end + + @testset "cross-statement fusion into a slice" begin + N = 16 + ja = reshape(T.(1:(N * N)), N, N) + a = @allowscalar NDArray(ja) + out = cuNumeric.zeros(T, (N + 2, N + 2)) + @analyze_lifetimes begin + producer = a .* s1 + out[2:(end - 1), 2:(end - 1)] = producer .+ s2 + end + expected = zeros(T, N + 2, N + 2) + expected[2:(end - 1), 2:(end - 1)] = ja .* s1 .+ s2 + @allowscalar @test cuNumeric.compare(expected, out, atol, rtol) + end end #= Broadcast fusion PTX compilation cache. diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index fc600d67a..9b7ffbcf4 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -106,6 +106,20 @@ function test_scoping_regressions(T, N) @test res isa cuNumeric.NDArray @test all(Array(res) .== T(4.0)) end + + if cuNumeric.FUSE_BROADCAST_EXPRS + @testset "Indexed fused assignment preserves Array view semantics" begin + src = reshape(T.(1:20), 4, 5) + out = fill(T(-1), 6, 7) + @analyze_lifetimes begin + producer = src .* T(2) + out[2:(end - 1), 2:(end - 1)] = producer .+ T(1) + end + @test out[2:(end - 1), 2:(end - 1)] == src .* T(2) .+ T(1) + @test all(out[[1, end], :] .== T(-1)) + @test all(out[:, [1, end]] .== T(-1)) + end + end end function run_all_ops(FT, N) From 86b964a588023b417a87e9acc12959693892b8bc Mon Sep 17 00:00:00 2001 From: krasow Date: Fri, 31 Jul 2026 17:09:18 -0500 Subject: [PATCH 02/12] scoping debug printer --- benchmark/benchmarks.toml | 59 -------------------------- docs/src/debugging.md | 21 ++++++++-- src/scoping.jl | 87 ++++++++++++++++++++++++++++++++++++--- test/tests/scoping.jl | 31 ++++++++++++++ 4 files changed, 130 insertions(+), 68 deletions(-) delete mode 100644 benchmark/benchmarks.toml diff --git a/benchmark/benchmarks.toml b/benchmark/benchmarks.toml deleted file mode 100644 index 70d5a47a6..000000000 --- a/benchmark/benchmarks.toml +++ /dev/null @@ -1,59 +0,0 @@ -[Global] -n_warmup = 5 -n_iter = 1000 -n_trial = 5 -cupynumeric = true # (needs install_cupynumeric.sh) -cuda = false # compare against CUDA.jl (single-GPU configs only) -# One CPU-reference check per config (not per timed iter). Written to CSV. -check_correctness = true -n_correctness_iter = 5 - -#################################### -# GEMM # -# Weak scaling, restricted to NxN. # -# Work ~ 2*N^2*M = 2*N^3. # -# N^3 / P --> constant. # -# N = baseline * P^(1/3). # -#################################### - -[[gemm]] -T = ["Float32"] -gpus = [1, 2, 4, 8] -cpus = 16 -N = [20000, 25200, 31752, 40000] -M = [20000, 25200, 31752, 40000] - -################################# -# Gray-Scott # -# Weak scaling, square NxN grid.# -# Work ~ N*M = N^2. # -# N^2 / P --> constant. # -# N = baseline * P^(1/2). # -################################# - -[[grayscott_baseline]] -T = "Float32" -gpus = [1, 2, 4, 8] -cpus = 16 -fusion = [true, false] -N = [2000, 2832, 4000, 5656] -M = [2000, 2832, 4000, 5656] - -[[grayscott_lifetimes]] -T = "Float32" -gpus = [1, 2, 4, 8] -cpus = 16 -fusion = false -N = [2000, 2832, 4000, 5656] -M = [2000, 2832, 4000, 5656] - -################################# -# Monte-Carlo Integration # -# Work ~ N. Scale N linearly # -################################# - -[[montecarlo]] -T = "Float32" -gpus = [1, 2, 4, 8] -cpus = 16 -N = [1_000_000, 2_000_000, 4_000_000, 8_000_000] diff --git a/docs/src/debugging.md b/docs/src/debugging.md index 8686e304b..49f4eb0d7 100644 --- a/docs/src/debugging.md +++ b/docs/src/debugging.md @@ -5,7 +5,7 @@ Debug the layer that matches the problem: | Question | Tool | |---|---| | Which operations did Legate submit, and when did they run? | [Legate logs and profiles](#trace-legate-runtime-work) | -| Did a broadcast become one fused kernel? | [`BCAST_FUSION_DEBUG`](#inspect-fused-broadcasts-with-bcast_fusion_debug) | +| How were broadcasts fused? | [`BCAST_FUSION_DEBUG`](#inspect-fused-broadcasts-with-bcast_fusion_debug) | | Where does `@analyze_lifetimes` free temporaries? | [`@show_lifetimes`](#inspect-lifetime-rewrites-with-show_lifetimes) | ## Trace Legate runtime work @@ -74,7 +74,9 @@ cuNumeric already supplies names for individual operations when task-scope namin ## Inspect fused broadcasts with `BCAST_FUSION_DEBUG` -When broadcast fusion is on, set `cuNumeric.BCAST_FUSION_DEBUG[] = true` to print each fused kernel before launch: the expression tree, inputs, scalars, arg map, and launch geometry. +When broadcast fusion is on, set `cuNumeric.BCAST_FUSION_DEBUG[] = true` to +print inter-statement rewrites and each fused kernel's expression tree, +arguments, and launch geometry. ```julia using cuNumeric @@ -91,7 +93,20 @@ C .= @. A * B + 2.0f0 cuNumeric.BCAST_FUSION_DEBUG[] = false ``` -Example output: +For example, a single-use producer inside `@analyze_lifetimes` is reported as: + +```text +======================================== inter-broadcast fusion rewrite + before + begin + product = A .* B + C[:, :] = product .+ 2.0f0 + end + fused + C[:, :] .= A .* B .+ 2.0f0 +``` + +The fused kernel is reported separately: ```text ======================================== fused broadcast kernel diff --git a/src/scoping.jl b/src/scoping.jl index 707253a49..75d72ad65 100644 --- a/src/scoping.jl +++ b/src/scoping.jl @@ -376,8 +376,8 @@ become one broadcast tree writing directly to `out[slice]`. Only broadcast definitions are delayed, and only across other lazy definitions or indexed writes to arrays not referenced by the producer. """ -function inline_single_use_broadcasts(block) - Meta.isexpr(block, (:block, :begin)) || return block +function _inline_single_use_broadcasts(block) + Meta.isexpr(block, (:block, :begin)) || return block, NamedTuple[] stmts = collect(block.args) definitions = Dict{Symbol,Tuple{Int,Any}}() @@ -409,33 +409,67 @@ function inline_single_use_broadcasts(block) end replacements = Dict{Symbol,Any}() + replacement_sources = Dict{Symbol,Vector{Int}}() removed = Set(first(info) for info in values(inlineable)) + def_symbols = Dict(info[1] => sym for (sym, info) in inlineable) + fusion_events = NamedTuple[] out = Any[] for (i, stmt) in enumerate(stmts) i in removed && begin - sym = only(sym for (sym, info) in inlineable if first(info) == i) + sym = def_symbols[i] + source_indices = Int[] + for dependency in walk_symbols(stmt.args[2]) + haskey(replacement_sources, dependency) || continue + append!(source_indices, replacement_sources[dependency]) + end + unique!(source_indices) + sort!(source_indices) + push!(source_indices, i) + replacement_sources[sym] = source_indices replacements[sym] = _substitute_symbols(inlineable[sym][3], replacements) continue end + original_stmt = stmt + source_indices = Int[] + for dependency in walk_symbols(original_stmt) + haskey(replacement_sources, dependency) || continue + append!(source_indices, replacement_sources[dependency]) + end + unique!(source_indices) + sort!(source_indices) + stmt = _substitute_symbols(stmt, replacements) # A materialized indexed assignment with a broadcast RHS normally # allocates a temporary and copies it into the slice. If the destination # base is absent from the RHS, write the fused tree directly to the view. + direct_write = false if stmt isa Expr && stmt.head == :(=) && stmt.args[1] isa Expr && stmt.args[1].head == :ref && _is_broadcast_syntax(stmt.args[2]) base = stmt.args[1].args[1] if base isa Symbol && !(base in Set(walk_symbols(stmt.args[2]))) stmt = Expr(:(.=), stmt.args[1], stmt.args[2]) + direct_write = true end end + + if !isempty(source_indices) || direct_write + before = Expr( + :block, + (deepcopy(stmts[source_idx]) for source_idx in source_indices)..., + deepcopy(original_stmt), + ) + push!(fusion_events, (; before, fused=deepcopy(stmt))) + end push!(out, stmt) end - return Expr(block.head, out...) + return Expr(block.head, out...), fusion_events end +inline_single_use_broadcasts(block) = first(_inline_single_use_broadcasts(block)) + function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) local_assigned = Set{Symbol}() @@ -570,15 +604,54 @@ end function process_broadcast_scope(block) assigned_vars = Set{Symbol}() - block = inline_single_use_broadcasts(block) + block, fusion_events = _inline_single_use_broadcasts(block) rewritten = find_broadcast_assignments(block, assigned_vars) result = insert_finalizers(rewritten, assigned_vars) + if !isempty(fusion_events) + log_calls = Expr[] + for event in fusion_events + before = QuoteNode(event.before) + fused = QuoteNode(event.fused) + push!( + log_calls, + :(cuNumeric._maybe_log_inter_broadcast_fusion($before, $fused)), + ) + end + result = Expr(:block, log_calls..., result.args...) + end counter[] = 0 return result end +# Log the source statements that were recombined by inter-statement broadcast +# fusion and the fused expression they become. The call is injected into the +# analyzed block so `BCAST_FUSION_DEBUG[]` is checked when the block executes, +# rather than earlier when its macro is expanded. +function _print_fusion_expr(io::IO, expr) + clean = expr isa Expr ? Base.remove_linenums!(deepcopy(expr)) : expr + rendered = sprint(Base.show_unquoted, clean) + for line in eachline(IOBuffer(rendered)) + println(io, " ", line) + end + return nothing +end + +function _maybe_log_inter_broadcast_fusion(before, fused; io::IO=stdout) + BCAST_FUSION_DEBUG[] || return nothing + println(io, "\n", "="^40, " inter-broadcast fusion rewrite") + println(io, " before") + _print_fusion_expr(io, before) + println(io, " fused") + _print_fusion_expr(io, fused) + return nothing +end + # Pretty-print the @analyze_lifetimes rewrite (see @show_lifetimes below). _is_delete_call(s) = Meta.isexpr(s, :call) && s.args[1] == :(cuNumeric.maybe_insert_delete) +function _is_fusion_log_call(s) + return Meta.isexpr(s, :call) && + s.args[1] == :(cuNumeric._maybe_log_inter_broadcast_fusion) +end # Flatten nested begin/blocks into a linear statement list, dropping line nodes. function _flatten_stmts(x) @@ -603,7 +676,9 @@ function print_lifetime_analysis(block; io::IO=stdout) n = 0 for s in stmts - if _is_delete_call(s) + if _is_fusion_log_call(s) + continue + elseif _is_delete_call(s) printstyled(io, lpad("✗ free ", 11), s.args[2], "\n"; color=:red) else n += 1 diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index 9b7ffbcf4..6f31dcede 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -108,6 +108,37 @@ function test_scoping_regressions(T, N) end if cuNumeric.FUSE_BROADCAST_EXPRS + @testset "Inter-broadcast fusion logger" begin + source = quote + product = A .* B + shifted = product .+ T(2) + C[:, :] = shifted ./ T(3) + end + _, events = cuNumeric._inline_single_use_broadcasts(source) + + @test length(events) == 1 + fused = sprint(Base.show_unquoted, events[1].fused) + @test !occursin("product", fused) + @test !occursin("shifted", fused) + @test occursin(".=", fused) + + old_debug = cuNumeric.BCAST_FUSION_DEBUG[] + try + cuNumeric.BCAST_FUSION_DEBUG[] = true + io = IOBuffer() + cuNumeric._maybe_log_inter_broadcast_fusion( + events[1].before, events[1].fused; io + ) + log = String(take!(io)) + @test occursin("inter-broadcast fusion rewrite", log) + @test occursin("product =", log) + @test occursin("shifted =", log) + @test occursin("fused", log) + finally + cuNumeric.BCAST_FUSION_DEBUG[] = old_debug + end + end + @testset "Indexed fused assignment preserves Array view semantics" begin src = reshape(T.(1:20), 4, 5) out = fill(T(-1), 6, 7) From 66b780bfb4e772859c6b549313f2dc6947f6760f Mon Sep 17 00:00:00 2001 From: krasow Date: Fri, 31 Jul 2026 17:43:58 -0500 Subject: [PATCH 03/12] refactor scoping analysis --- docs/src/api.md | 2 +- docs/src/internals.md | 2 +- src/cuNumeric.jl | 2 +- src/scoping.jl | 704 -------------------------- src/scoping/broadcast_lifetimes.jl | 135 +++++ src/scoping/inter_broadcast_fusion.jl | 205 ++++++++ src/scoping/lifetimes.jl | 89 ++++ src/scoping/scoping.jl | 258 ++++++++++ test/tests/scoping.jl | 33 +- 9 files changed, 704 insertions(+), 726 deletions(-) delete mode 100644 src/scoping.jl create mode 100644 src/scoping/broadcast_lifetimes.jl create mode 100644 src/scoping/inter_broadcast_fusion.jl create mode 100644 src/scoping/lifetimes.jl create mode 100644 src/scoping/scoping.jl diff --git a/docs/src/api.md b/docs/src/api.md index fd7ba459e..82345f271 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -4,6 +4,6 @@ Indexing, reshaping, reductions, comparisons, memory helpers, lifetime macros, a ```@autodocs Modules = [cuNumeric] -Pages = ["ndarray/ndarray.jl", "ndarray/linalg.jl", "cuNumeric.jl", "warnings.jl", "util.jl", "memory.jl", "scoping.jl"] +Pages = ["ndarray/ndarray.jl", "ndarray/linalg.jl", "cuNumeric.jl", "warnings.jl", "util.jl", "memory.jl", "scoping/scoping.jl"] Filter = t -> !(t isa Function && nameof(t) in (:zeros, :ones, :fill, :trues, :falses, :eye, :rand, :rand!)) ``` diff --git a/docs/src/internals.md b/docs/src/internals.md index 10f655faf..6fa90b799 100644 --- a/docs/src/internals.md +++ b/docs/src/internals.md @@ -41,7 +41,7 @@ Limits to keep in mind: - Analysis is statement-linear. It is not a full control-flow graph pass. Wrap hot loop bodies, not entire programs. - Some paths free eagerly outside the macro (for example LHS slice views created during indexed assignment). -Relevant source: `src/scoping.jl`. +Relevant source: `src/scoping/`. ### Allocation-driven GC heuristics diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index eb2d0f6b3..9eb9d5b56 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -174,7 +174,7 @@ include("ndarray/ndarray.jl") include("ndarray/unary.jl") include("ndarray/binary.jl") include("ndarray/linalg.jl") -include("scoping.jl") +include("scoping/scoping.jl") # From https://github.com/JuliaGraphics/QML.jl/blob/dca239404135d85fe5d4afe34ed3dc5f61736c63/src/QML.jl#L147 mutable struct ArgcArgv diff --git a/src/scoping.jl b/src/scoping.jl deleted file mode 100644 index 75d72ad65..000000000 --- a/src/scoping.jl +++ /dev/null @@ -1,704 +0,0 @@ -export @analyze_lifetimes, @show_lifetimes - -@doc""" - @analyze_lifetimes expr - -Wraps a block of code so that all temporary `NDArray` allocations -(e.g. from slicing or function calls) are tracked and safely freed -at the end of the block. Ensures proper cleanup of GPU memory by -inserting `maybe_insert_delete` calls automatically. - -When broadcast fusion is enabled (`FUSE_BROADCAST_EXPRS`), dotted operators -(`.+`, `.*`, etc.) form a lazy `Base.Broadcast.Broadcasted` tree compiled into -a single PTX kernel; intermediate nodes are not real `NDArray` allocations and -are not individually hoisted. The macro automatically selects the -broadcast-aware analysis in that case and the plain analysis otherwise. -""" -macro analyze_lifetimes(block) - return esc(process_ndarray_scope(block)) -end - -const counter = Ref(0) - -function maybe_insert_delete(var::NDArray) - return cuNumeric.destroy!(var) -end - -maybe_insert_delete(x) = x - -""" - walk_symbols(x) -> Vector{Symbol} - -Recursively collect all symbols that appear inside expression `x`. -""" -function walk_symbols(x) - syms = Symbol[] - if x isa Symbol - push!(syms, x) - elseif x isa Expr - for a in x.args - append!(syms, walk_symbols(a)) - end - elseif x isa AbstractArray - for a in x - append!(syms, walk_symbols(a)) - end - end - return syms -end - -""" - insert_finalizers(stmts::Vector) -Insert `cuNumeric.maybe_insert_delete(var)` after the last use of each temporary variable. -""" -function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) - uses = Dict{Symbol,Vector{Int}}() - defs = Dict{Symbol,Int}() - alias_map = Dict{Symbol,Symbol}() - - # Collect all statements, flattening blocks and skipping LineNumberNodes - stmts = Any[] - for expr in exprs - if expr isa LineNumberNode - continue - elseif expr isa Expr && expr.head == :block - for arg in expr.args - arg isa LineNumberNode || push!(stmts, arg) - end - else - push!(stmts, expr) - end - end - - # Pass 1: collect definitions and uses - for (i, stmt) in enumerate(stmts) - stmt isa Expr || continue - stmt.head == :line && continue - - if stmt.head == :(=) - lhs, rhs = stmt.args - if lhs isa Symbol - defs[lhs] = i - end - if lhs isa Symbol && rhs isa Symbol - alias_map[lhs] = rhs - end - for s in walk_symbols(rhs) - push!(get!(uses, s, Int[]), i) - end - else - for s in walk_symbols(stmt) - push!(get!(uses, s, Int[]), i) - end - end - end - - for (alias, src) in alias_map - append!(get!(uses, src, Int[]), get(uses, alias, Int[])) - end - - # Compute last usage index per variable - last_use = Dict{Symbol,Int}() - for (v, idxs) in uses - last_use[v] = maximum(idxs) - end - - # `F_u = tmp1` aliases one NDArray under two names; resolve to a canonical rep - # so it's freed once (double-free is masked only by the ptr=0 null-out). - function canon(v) - seen = Set{Symbol}() - while haskey(alias_map, v) && !(v in seen) - push!(seen, v) - v = alias_map[v] - end - return v - end - - is_indexed_assign(s) = - s isa Expr && s.head in (:(=), :(.=)) && !(s.args[1] isa Symbol) - result_symbol(s) = - if s isa Symbol - s - else - (s isa Expr && s.head == :(=) && s.args[1] isa Symbol ? s.args[1] : nothing) - end - - # Pass 2: insert finalizers - out = Any[] - n = length(stmts) - freed = Set{Symbol}() - - # The block's value escapes to the caller, except for `A[...] = rhs`: Julia - # returns `rhs` there, but that's a dead temp nobody consumes — free it and - # return `nothing` rather than leak it or hand back a dangling handle. - terminal_indexed = n > 0 && is_indexed_assign(stmts[n]) - - protected = Set{Symbol}() - if n > 0 && !terminal_indexed - rs = result_symbol(stmts[n]) - rs isa Symbol && push!(protected, canon(rs)) - end - - function emit_delete!(v) - c = canon(v) - (c in freed || c in protected) && return nothing - push!(freed, c) - return push!(out, :(cuNumeric.maybe_insert_delete($v))) - end - - for (i, stmt) in enumerate(stmts) - # `v = w` aliases w, so don't finalize w at this statement. - skip_finalize = Set{Symbol}() - if stmt isa Expr && stmt.head == :(=) - lhs, rhs = stmt.args - if lhs isa Symbol && rhs isa Symbol - push!(skip_finalize, rhs) - end - end - - if i == n && !terminal_indexed - res_var = Symbol(:res, counter[]) - counter[] += 1 - push!(out, :($res_var = $stmt)) - else - push!(out, stmt) - end - - for (v, lasti) in last_use - if lasti == i && v ∈ assigned_vars && !(v ∈ skip_finalize) - emit_delete!(v) - end - end - - i == n && push!(out, terminal_indexed ? :nothing : res_var) - end - - return out -end - -""" - insert_finalizers(block::Expr) -Apply finalizer insertion to a `begin ... end` or `:block` expression. -""" -function insert_finalizers(block::Expr, assigned_vars::Set{Symbol}) - if block.head == :block || block.head == :begin - # Filter out LineNumberNodes before processing - stmts = [s for s in block.args if !(s isa LineNumberNode)] - new_stmts = insert_finalizers(stmts, assigned_vars) - return Expr(:block, new_stmts...) - else - error("Expected a begin/block expression") - end -end - -function process_ndarray_scope(block) - # Broadcast trees need the fusion-aware hoisting; otherwise the - # plain analysis treats every call (dotted or not) as a real allocation. - @static if FUSE_BROADCAST_EXPRS - return process_broadcast_scope(block) - end - assigned_vars = Set{Symbol}() - # Process the entire block at once so lifetimes are tracked across statements - rewritten = find_ndarray_assignments(block, assigned_vars) - result = insert_finalizers(rewritten, assigned_vars) - counter[] = 0 - return result -end - -function find_ndarray_assignments(ex, assigned_vars::Set{Symbol}) - cache = Dict{Any,Symbol}() # expression → temp mapping - local_assigned = Set{Symbol}() # track all assigned symbols - - # --- create a fresh temp for any expression --- - function fresh_tmp(expr) - counter[] += 1 - tmp = Symbol(:tmp, counter[]) - cache[expr] = tmp - push!(local_assigned, tmp) - return tmp, [:($tmp = $expr)] - end - - # --- recursive rewrite --- - function rewrite(e)::Tuple{Any,Vector{Expr}} - if !(e isa Expr) - return e, Expr[] - end - - # --- assignment: leave LHS intact --- - if e.head == :(=) - lhs, rhs = e.args - if lhs isa Symbol - push!(local_assigned, lhs) - end - new_rhs, temps = rewrite(rhs) - return :($lhs = $new_rhs), temps - end - - # --- broadcasted assignment: preserve fusion --- - if e.head == :(.=) - lhs, rhs = e.args - new_lhs, lhs_temps = rewrite(lhs) - # Do not hoist the top-level call of the RHS to preserve fusion - if rhs isa Expr && rhs.head == :call - op = rhs.args[1] - new_rhs_args, rhs_temps = Any[], Expr[] - for arg in rhs.args[2:end] - new_arg, t = rewrite(arg) - push!(new_rhs_args, new_arg) - append!(rhs_temps, t) - end - new_rhs = Expr(:call, op, new_rhs_args...) - return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) - else - new_rhs, rhs_temps = rewrite(rhs) - return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) - end - end - - # --- array slice reference --- - if e.head == :ref - return fresh_tmp(e) - end - - # --- function calls --- - if e.head == :call - op = e.args[1] - new_args, hoisted = Any[], Expr[] - - # recursively rewrite arguments first - for arg in e.args[2:end] - new_arg, temps = rewrite(arg) - push!(new_args, new_arg) - append!(hoisted, temps) - end - - # recreate call with rewritten args - new_expr = Expr(:call, op, new_args...) - - # always hoist calls (arrays and non-arrays) - tmp, bind = fresh_tmp(new_expr) - return tmp, vcat(hoisted, bind) - end - - # --- fallback for other Expr types --- - new_args, hoisted = Any[], Expr[] - is_block = e.head == :block || e.head == :begin - for arg in e.args - new_arg, temps = rewrite(arg) - if is_block && !(arg isa LineNumberNode) - append!(new_args, temps) - push!(new_args, new_arg) - else - push!(new_args, new_arg) - append!(hoisted, temps) - end - end - return Expr(e.head, new_args...), hoisted - end - - new_ex, temps = rewrite(ex) - union!(assigned_vars, local_assigned) - - if new_ex isa Expr && new_ex.head == :block - return Expr(:block, temps..., new_ex.args...) - else - return Expr(:block, temps..., new_ex) - end -end - -# Broadcast-fusion-aware lifetime analysis. Under fusion, dotted operators -# (.+, .*, …) form a lazy Broadcasted tree lowered to one PTX kernel, so their -# intermediate nodes are not real NDArrays — only slices and the tree root are. -# Non-broadcast sub-expressions break the tree and are hoisted like any call. - -is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") - -function _is_broadcast_syntax(e) - return e isa Expr && - ( - (e.head == :call && !isempty(e.args) && is_broadcast_op(e.args[1])) || - (e.head == :. && length(e.args) == 2) - ) -end - -function _symbol_occurrences(x, target::Symbol) - x === target && return 1 - x isa Expr || return 0 - return sum(arg -> _symbol_occurrences(arg, target), x.args; init=0) -end - -function _substitute_symbols(x, replacements::Dict{Symbol,Any}) - x isa Symbol && return get(replacements, x, x) - x isa Expr || return x - - # A symbol assignment introduces/redefines its LHS; only substitute in RHS. - if x.head == :(=) && x.args[1] isa Symbol - return Expr(:(=), x.args[1], _substitute_symbols(x.args[2], replacements)) - end - return Expr(x.head, (_substitute_symbols(arg, replacements) for arg in x.args)...) -end - -function _indexed_assignment_base(stmt) - if stmt isa Expr && stmt.head == :(=) && - stmt.args[1] isa Expr && stmt.args[1].head == :ref - base = stmt.args[1].args[1] - return base isa Symbol ? base : nothing - end - return nothing -end - -function _safe_to_delay_broadcast( - stmts, def_idx::Int, use_idx::Int, dependencies::Set{Symbol}, lazy_defs::Set{Int} -) - for i in (def_idx + 1):(use_idx - 1) - stmt = stmts[i] - stmt isa LineNumberNode && continue - i in lazy_defs && continue - - # An indexed write to an unrelated array does not invalidate the lazy - # producer. Any other intervening statement is conservatively a barrier. - mutated = _indexed_assignment_base(stmt) - mutated !== nothing && !(mutated in dependencies) && continue - return false - end - return true -end - -""" -Inline single-use broadcast temporaries across statements before lifetime -hoisting. This lets a stencil such as - - lap = ... - out[slice] = scale .* lap .+ ... - -become one broadcast tree writing directly to `out[slice]`. - -Only broadcast definitions are delayed, and only across other lazy definitions -or indexed writes to arrays not referenced by the producer. -""" -function _inline_single_use_broadcasts(block) - Meta.isexpr(block, (:block, :begin)) || return block, NamedTuple[] - stmts = collect(block.args) - - definitions = Dict{Symbol,Tuple{Int,Any}}() - lazy_defs = Set{Int}() - for (i, stmt) in enumerate(stmts) - if stmt isa Expr && stmt.head == :(=) && stmt.args[1] isa Symbol && - _is_broadcast_syntax(stmt.args[2]) - definitions[stmt.args[1]] = (i, stmt.args[2]) - push!(lazy_defs, i) - end - end - - inlineable = Dict{Symbol,Tuple{Int,Int,Any}}() - for (sym, (def_idx, rhs)) in definitions - use_indices = Int[] - occurrences = 0 - for i in (def_idx + 1):length(stmts) - n = _symbol_occurrences(stmts[i], sym) - if n > 0 - occurrences += n - push!(use_indices, i) - end - end - occurrences == 1 || continue - use_idx = only(use_indices) - deps = Set(walk_symbols(rhs)) - _safe_to_delay_broadcast(stmts, def_idx, use_idx, deps, lazy_defs) || continue - inlineable[sym] = (def_idx, use_idx, rhs) - end - - replacements = Dict{Symbol,Any}() - replacement_sources = Dict{Symbol,Vector{Int}}() - removed = Set(first(info) for info in values(inlineable)) - def_symbols = Dict(info[1] => sym for (sym, info) in inlineable) - fusion_events = NamedTuple[] - out = Any[] - for (i, stmt) in enumerate(stmts) - i in removed && begin - sym = def_symbols[i] - source_indices = Int[] - for dependency in walk_symbols(stmt.args[2]) - haskey(replacement_sources, dependency) || continue - append!(source_indices, replacement_sources[dependency]) - end - unique!(source_indices) - sort!(source_indices) - push!(source_indices, i) - replacement_sources[sym] = source_indices - replacements[sym] = _substitute_symbols(inlineable[sym][3], replacements) - continue - end - - original_stmt = stmt - source_indices = Int[] - for dependency in walk_symbols(original_stmt) - haskey(replacement_sources, dependency) || continue - append!(source_indices, replacement_sources[dependency]) - end - unique!(source_indices) - sort!(source_indices) - - stmt = _substitute_symbols(stmt, replacements) - - # A materialized indexed assignment with a broadcast RHS normally - # allocates a temporary and copies it into the slice. If the destination - # base is absent from the RHS, write the fused tree directly to the view. - direct_write = false - if stmt isa Expr && stmt.head == :(=) && - stmt.args[1] isa Expr && stmt.args[1].head == :ref && - _is_broadcast_syntax(stmt.args[2]) - base = stmt.args[1].args[1] - if base isa Symbol && !(base in Set(walk_symbols(stmt.args[2]))) - stmt = Expr(:(.=), stmt.args[1], stmt.args[2]) - direct_write = true - end - end - - if !isempty(source_indices) || direct_write - before = Expr( - :block, - (deepcopy(stmts[source_idx]) for source_idx in source_indices)..., - deepcopy(original_stmt), - ) - push!(fusion_events, (; before, fused=deepcopy(stmt))) - end - push!(out, stmt) - end - return Expr(block.head, out...), fusion_events -end - -inline_single_use_broadcasts(block) = first(_inline_single_use_broadcasts(block)) - -function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) - local_assigned = Set{Symbol}() - - function fresh_tmp(expr) - counter[] += 1 - tmp = Symbol(:tmp, counter[]) - push!(local_assigned, tmp) - return tmp, [:($tmp = $expr)] - end - - # Rewrite each arg with `f`, collecting the temps each one hoists. - function maphoist(f, args) - new_args, hoisted = Any[], Expr[] - for a in args - na, ts = f(a) - push!(new_args, na) - append!(hoisted, ts) - end - return new_args, hoisted - end - - # Inside a broadcast tree: hoist slices, keep dotted ops/f.(…) lazy, and - # delegate anything else to rewrite() (it breaks the tree → real NDArray). - # - # The slice cache is deliberately scoped to one fused tree. Repeated views - # then become one task argument and are still destroyed immediately after - # that task's last use. Sharing the cache across trees would keep transformed - # stores alive across task submissions, which can force runtime copies and - # delay device-memory reclamation. - function rewrite_bcast(e, slice_cache::Dict{Any,Symbol})::Tuple{Any,Vector{Expr}} - e isa Expr || return e, Expr[] - if e.head == :ref - cached = get(slice_cache, e, nothing) - cached === nothing || return cached, Expr[] - tmp, bind = fresh_tmp(e) - slice_cache[e] = tmp - return tmp, bind - end - if e.head == :call && is_broadcast_op(e.args[1]) - args, hoisted = maphoist( - x -> rewrite_bcast(x, slice_cache), e.args[2:end] - ) - return Expr(:call, e.args[1], args...), hoisted - end - if e.head == :. && length(e.args) == 2 && - e.args[2] isa Expr && e.args[2].head == :tuple - args, hoisted = maphoist( - x -> rewrite_bcast(x, slice_cache), e.args[2].args - ) - return Expr(:., e.args[1], Expr(:tuple, args...)), hoisted - end - return rewrite(e) - end - - function rewrite(e)::Tuple{Any,Vector{Expr}} - e isa Expr || return e, Expr[] - - if e.head == :(=) - lhs, rhs = e.args - lhs isa Symbol && push!(local_assigned, lhs) - new_rhs, temps = rewrite(rhs) - return :($lhs = $new_rhs), temps - end - - # .= RHS is a broadcast tree: only the slices inside it are hoisted. - if e.head == :(.=) - lhs, rhs = e.args - # Preserve Julia's indexed `.=` semantics. Hoisting `A[inds...]` - # directly is valid for NDArray (whose getindex returns a view), but - # produces a detached copy for Array. An explicit view works for - # both, while still giving the lifetime pass a concrete NDArray - # wrapper to destroy immediately after task submission. - new_lhs, lts = - if lhs isa Expr && lhs.head == :ref - view_expr = Base.macroexpand( - @__MODULE__, - Expr( - :macrocall, - GlobalRef(Base, Symbol("@view")), - LineNumberNode(0), - lhs, - ), - ) - fresh_tmp(view_expr) - else - rewrite(lhs) - end - new_rhs, rts = rewrite_bcast(rhs, Dict{Any,Symbol}()) - return Expr(:(.=), new_lhs, new_rhs), vcat(lts, rts) - end - - e.head == :ref && return fresh_tmp(e) - - # Broadcast root: hoist the fused result as one temp. - if e.head == :call && is_broadcast_op(e.args[1]) - inner, hoisted = rewrite_bcast(e, Dict{Any,Symbol}()) - tmp, bind = fresh_tmp(inner) - return tmp, vcat(hoisted, bind) - end - - # Regular call: hoist it and its args. - if e.head == :call - args, hoisted = maphoist(rewrite, e.args[2:end]) - tmp, bind = fresh_tmp(Expr(:call, e.args[1], args...)) - return tmp, vcat(hoisted, bind) - end - - # Other exprs: recurse, splicing hoisted temps inline within blocks. - new_args, hoisted = Any[], Expr[] - is_block = e.head == :block || e.head == :begin - for a in e.args - na, ts = rewrite(a) - if is_block && !(a isa LineNumberNode) - append!(new_args, ts) - push!(new_args, na) - else - push!(new_args, na) - append!(hoisted, ts) - end - end - return Expr(e.head, new_args...), hoisted - end - - new_ex, temps = rewrite(ex) - union!(assigned_vars, local_assigned) - if new_ex isa Expr && new_ex.head == :block - return Expr(:block, temps..., new_ex.args...) - else - return Expr(:block, temps..., new_ex) - end -end - -function process_broadcast_scope(block) - assigned_vars = Set{Symbol}() - block, fusion_events = _inline_single_use_broadcasts(block) - rewritten = find_broadcast_assignments(block, assigned_vars) - result = insert_finalizers(rewritten, assigned_vars) - if !isempty(fusion_events) - log_calls = Expr[] - for event in fusion_events - before = QuoteNode(event.before) - fused = QuoteNode(event.fused) - push!( - log_calls, - :(cuNumeric._maybe_log_inter_broadcast_fusion($before, $fused)), - ) - end - result = Expr(:block, log_calls..., result.args...) - end - counter[] = 0 - return result -end - -# Log the source statements that were recombined by inter-statement broadcast -# fusion and the fused expression they become. The call is injected into the -# analyzed block so `BCAST_FUSION_DEBUG[]` is checked when the block executes, -# rather than earlier when its macro is expanded. -function _print_fusion_expr(io::IO, expr) - clean = expr isa Expr ? Base.remove_linenums!(deepcopy(expr)) : expr - rendered = sprint(Base.show_unquoted, clean) - for line in eachline(IOBuffer(rendered)) - println(io, " ", line) - end - return nothing -end - -function _maybe_log_inter_broadcast_fusion(before, fused; io::IO=stdout) - BCAST_FUSION_DEBUG[] || return nothing - println(io, "\n", "="^40, " inter-broadcast fusion rewrite") - println(io, " before") - _print_fusion_expr(io, before) - println(io, " fused") - _print_fusion_expr(io, fused) - return nothing -end - -# Pretty-print the @analyze_lifetimes rewrite (see @show_lifetimes below). -_is_delete_call(s) = Meta.isexpr(s, :call) && s.args[1] == :(cuNumeric.maybe_insert_delete) -function _is_fusion_log_call(s) - return Meta.isexpr(s, :call) && - s.args[1] == :(cuNumeric._maybe_log_inter_broadcast_fusion) -end - -# Flatten nested begin/blocks into a linear statement list, dropping line nodes. -function _flatten_stmts(x) - stmts = Any[] - function walk(e) - if Meta.isexpr(e, (:block, :begin)) - foreach(walk, e.args) - elseif !(e isa LineNumberNode) - push!(stmts, e) - end - end - walk(x) - return stmts -end - -function print_lifetime_analysis(block; io::IO=stdout) - rule = "-"^60 - stmts = _flatten_stmts(process_ndarray_scope(block)) - mode = FUSE_BROADCAST_EXPRS ? "fusion-aware" : "plain" - - println(io, "@analyze_lifetimes expansion ($mode analysis)\n", rule) - - n = 0 - for s in stmts - if _is_fusion_log_call(s) - continue - elseif _is_delete_call(s) - printstyled(io, lpad("✗ free ", 11), s.args[2], "\n"; color=:red) - else - n += 1 - println(io, lpad(n, 4), " ", s) - end - end - - println(io, rule) - return nothing -end - -@doc""" - @show_lifetimes expr - -Print the lifetime-analysis rewrite of `expr` — the same transformation -[`@analyze_lifetimes`](@ref) applies — without running it. Every statement is -shown in source order and each inserted `maybe_insert_delete` is highlighted so -you can see exactly where each temporary is freed. Pure AST work, so it runs on -CPU-only checkouts. -""" -macro show_lifetimes(block) - return :(print_lifetime_analysis($(QuoteNode(block)))) -end diff --git a/src/scoping/broadcast_lifetimes.jl b/src/scoping/broadcast_lifetimes.jl new file mode 100644 index 000000000..a6f377aec --- /dev/null +++ b/src/scoping/broadcast_lifetimes.jl @@ -0,0 +1,135 @@ +# Lifetime analysis for lazy broadcast expression trees. + +is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") + +function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) + local_assigned = Set{Symbol}() + + function fresh_tmp(expr) + counter[] += 1 + tmp = Symbol(:tmp, counter[]) + push!(local_assigned, tmp) + return tmp, [:($tmp = $expr)] + end + + function maphoist(f, args) + new_args, hoisted = Any[], Expr[] + for arg in args + new_arg, temps = f(arg) + push!(new_args, new_arg) + append!(hoisted, temps) + end + return new_args, hoisted + end + + # Inside a broadcast tree: hoist slices, keep dotted ops/f.(…) lazy, and + # delegate anything else to rewrite() (it breaks the tree → real NDArray). + # The slice cache is scoped to one fused tree so repeated views become one + # task argument without extending their lifetime across task submissions. + function rewrite_broadcast( + expr, slice_cache::Dict{Any,Symbol} + )::Tuple{Any,Vector{Expr}} + expr isa Expr || return expr, Expr[] + if expr.head == :ref + cached = get(slice_cache, expr, nothing) + cached === nothing || return cached, Expr[] + tmp, bind = fresh_tmp(expr) + slice_cache[expr] = tmp + return tmp, bind + end + if expr.head == :call && is_broadcast_op(expr.args[1]) + args, hoisted = maphoist( + arg -> rewrite_broadcast(arg, slice_cache), expr.args[2:end] + ) + return Expr(:call, expr.args[1], args...), hoisted + end + if expr.head == :. && length(expr.args) == 2 && + expr.args[2] isa Expr && expr.args[2].head == :tuple + args, hoisted = maphoist( + arg -> rewrite_broadcast(arg, slice_cache), expr.args[2].args + ) + return Expr(:., expr.args[1], Expr(:tuple, args...)), hoisted + end + return rewrite(expr) + end + + function rewrite(expr)::Tuple{Any,Vector{Expr}} + expr isa Expr || return expr, Expr[] + InterBroadcastFusion.is_log_call(expr) && return expr, Expr[] + + if expr.head == :(=) + lhs, rhs = expr.args + lhs isa Symbol && push!(local_assigned, lhs) + new_rhs, temps = rewrite(rhs) + return :($lhs = $new_rhs), temps + end + + # A `.=` RHS is a broadcast tree: only its slices are hoisted. + if expr.head == :(.=) + lhs, rhs = expr.args + # An explicit view preserves indexed `.=` semantics for both Array + # and NDArray while giving the lifetime pass a wrapper to destroy. + new_lhs, lhs_temps = + if lhs isa Expr && lhs.head == :ref + view_expr = Base.macroexpand( + @__MODULE__, + Expr( + :macrocall, + GlobalRef(Base, Symbol("@view")), + LineNumberNode(0), + lhs, + ), + ) + fresh_tmp(view_expr) + else + rewrite(lhs) + end + new_rhs, rhs_temps = rewrite_broadcast(rhs, Dict{Any,Symbol}()) + return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) + end + + expr.head == :ref && return fresh_tmp(expr) + + if expr.head == :call && is_broadcast_op(expr.args[1]) + inner, hoisted = rewrite_broadcast(expr, Dict{Any,Symbol}()) + tmp, bind = fresh_tmp(inner) + return tmp, vcat(hoisted, bind) + end + + if expr.head == :call + args, hoisted = maphoist(rewrite, expr.args[2:end]) + tmp, bind = fresh_tmp(Expr(:call, expr.args[1], args...)) + return tmp, vcat(hoisted, bind) + end + + new_args, hoisted = Any[], Expr[] + is_block = expr.head == :block || expr.head == :begin + for arg in expr.args + new_arg, temps = rewrite(arg) + if is_block && !(arg isa LineNumberNode) + append!(new_args, temps) + push!(new_args, new_arg) + else + push!(new_args, new_arg) + append!(hoisted, temps) + end + end + return Expr(expr.head, new_args...), hoisted + end + + new_ex, temps = rewrite(ex) + union!(assigned_vars, local_assigned) + if new_ex isa Expr && new_ex.head == :block + return Expr(:block, temps..., new_ex.args...) + end + return Expr(:block, temps..., new_ex) +end + +function process_broadcast_lifetime_scope(scope) + assigned_vars = Set{Symbol}() + scope = InterBroadcastFusion.rewrite_scope(scope) + rewritten = find_broadcast_assignments(scope, assigned_vars) + result = insert_finalizers(rewritten, assigned_vars) + counter[] = 0 + return result +end diff --git a/src/scoping/inter_broadcast_fusion.jl b/src/scoping/inter_broadcast_fusion.jl new file mode 100644 index 000000000..e6175940a --- /dev/null +++ b/src/scoping/inter_broadcast_fusion.jl @@ -0,0 +1,205 @@ +module InterBroadcastFusion + +export rewrite_scope + +import ..cuNumeric: walk_symbols + +_is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") + +function _is_broadcast_syntax(expr) + return expr isa Expr && + ( + (expr.head == :call && !isempty(expr.args) && _is_broadcast_op(expr.args[1])) || + (expr.head == :. && length(expr.args) == 2) + ) +end + +function _symbol_occurrences(expr, target::Symbol) + expr === target && return 1 + expr isa Expr || return 0 + return sum(arg -> _symbol_occurrences(arg, target), expr.args; init=0) +end + +function _substitute_symbols(expr, replacements::Dict{Symbol,Any}) + expr isa Symbol && return get(replacements, expr, expr) + expr isa Expr || return expr + + # A symbol assignment introduces/redefines its LHS; only substitute in RHS. + if expr.head == :(=) && expr.args[1] isa Symbol + return Expr( + :(=), expr.args[1], _substitute_symbols(expr.args[2], replacements) + ) + end + return Expr( + expr.head, (_substitute_symbols(arg, replacements) for arg in expr.args)... + ) +end + +function _indexed_assignment_base(stmt) + if stmt isa Expr && stmt.head == :(=) && + stmt.args[1] isa Expr && stmt.args[1].head == :ref + base = stmt.args[1].args[1] + return base isa Symbol ? base : nothing + end + return nothing +end + +function _safe_to_delay_broadcast( + stmts, def_idx::Int, use_idx::Int, dependencies::Set{Symbol}, lazy_defs::Set{Int} +) + for i in (def_idx + 1):(use_idx - 1) + stmt = stmts[i] + stmt isa LineNumberNode && continue + i in lazy_defs && continue + + # An indexed write to an unrelated array does not invalidate the lazy + # producer. Any other intervening statement is conservatively a barrier. + mutated = _indexed_assignment_base(stmt) + mutated !== nothing && !(mutated in dependencies) && continue + return false + end + return true +end + +function _rewrite_scope(scope) + Meta.isexpr(scope, (:block, :begin)) || return scope, NamedTuple[] + stmts = collect(scope.args) + + definitions = Dict{Symbol,Tuple{Int,Any}}() + lazy_defs = Set{Int}() + for (i, stmt) in enumerate(stmts) + if stmt isa Expr && stmt.head == :(=) && stmt.args[1] isa Symbol && + _is_broadcast_syntax(stmt.args[2]) + definitions[stmt.args[1]] = (i, stmt.args[2]) + push!(lazy_defs, i) + end + end + + inlineable = Dict{Symbol,Tuple{Int,Int,Any}}() + for (sym, (def_idx, rhs)) in definitions + use_indices = Int[] + occurrences = 0 + for i in (def_idx + 1):length(stmts) + count = _symbol_occurrences(stmts[i], sym) + if count > 0 + occurrences += count + push!(use_indices, i) + end + end + occurrences == 1 || continue + use_idx = only(use_indices) + dependencies = Set(walk_symbols(rhs)) + _safe_to_delay_broadcast(stmts, def_idx, use_idx, dependencies, lazy_defs) || + continue + inlineable[sym] = (def_idx, use_idx, rhs) + end + + replacements = Dict{Symbol,Any}() + replacement_sources = Dict{Symbol,Vector{Int}}() + removed = Set(first(info) for info in values(inlineable)) + def_symbols = Dict(info[1] => sym for (sym, info) in inlineable) + fusion_events = NamedTuple[] + rewritten = Any[] + + for (i, original_stmt) in enumerate(stmts) + if i in removed + sym = def_symbols[i] + source_indices = Int[] + for dependency in walk_symbols(original_stmt.args[2]) + haskey(replacement_sources, dependency) || continue + append!(source_indices, replacement_sources[dependency]) + end + unique!(source_indices) + sort!(source_indices) + push!(source_indices, i) + replacement_sources[sym] = source_indices + replacements[sym] = _substitute_symbols(inlineable[sym][3], replacements) + continue + end + + source_indices = Int[] + for dependency in walk_symbols(original_stmt) + haskey(replacement_sources, dependency) || continue + append!(source_indices, replacement_sources[dependency]) + end + unique!(source_indices) + sort!(source_indices) + + stmt = _substitute_symbols(original_stmt, replacements) + + # A materialized indexed assignment with a broadcast RHS normally + # allocates a temporary and copies it into the slice. If the destination + # base is absent from the RHS, write the fused tree directly to the view. + direct_write = false + if stmt isa Expr && stmt.head == :(=) && + stmt.args[1] isa Expr && stmt.args[1].head == :ref && + _is_broadcast_syntax(stmt.args[2]) + base = stmt.args[1].args[1] + if base isa Symbol && !(base in Set(walk_symbols(stmt.args[2]))) + stmt = Expr(:(.=), stmt.args[1], stmt.args[2]) + direct_write = true + end + end + + if !isempty(source_indices) || direct_write + before = Expr( + :block, + (deepcopy(stmts[source_idx]) for source_idx in source_indices)..., + deepcopy(original_stmt), + ) + push!(fusion_events, (; before, fused=deepcopy(stmt))) + end + push!(rewritten, stmt) + end + + return Expr(scope.head, rewritten...), fusion_events +end + +const _LOG_FUNCTION = GlobalRef(@__MODULE__, :maybe_log_rewrite) +const _DEBUG_FLAG = GlobalRef(parentmodule(@__MODULE__), :BCAST_FUSION_DEBUG) + +function _log_call(event) + before = QuoteNode(event.before) + fused = QuoteNode(event.fused) + enabled = Expr(:ref, _DEBUG_FLAG) + return Expr(:call, _LOG_FUNCTION, enabled, before, fused) +end + +""" + rewrite_scope(scope) -> scope + +Fuse eligible single-use broadcast producers into their consumer and return the +rewritten scope. Runtime log hooks are included for rewrites and are controlled +by `cuNumeric.BCAST_FUSION_DEBUG[]`. +""" +function rewrite_scope(scope) + rewritten, fusion_events = _rewrite_scope(scope) + isempty(fusion_events) && return rewritten + log_calls = (_log_call(event) for event in fusion_events) + return Expr(rewritten.head, log_calls..., rewritten.args...) +end + +function is_log_call(expr) + return Meta.isexpr(expr, :call) && expr.args[1] == _LOG_FUNCTION +end + +function _print_expr(io::IO, expr) + clean = expr isa Expr ? Base.remove_linenums!(deepcopy(expr)) : expr + rendered = sprint(Base.show_unquoted, clean) + for line in eachline(IOBuffer(rendered)) + println(io, " ", line) + end + return nothing +end + +function maybe_log_rewrite(enabled::Bool, before, fused; io::IO=stdout) + enabled || return nothing + println(io, "\n", "="^40, " inter-broadcast fusion rewrite") + println(io, " before") + _print_expr(io, before) + println(io, " fused") + _print_expr(io, fused) + return nothing +end + +end diff --git a/src/scoping/lifetimes.jl b/src/scoping/lifetimes.jl new file mode 100644 index 000000000..730a5f5b9 --- /dev/null +++ b/src/scoping/lifetimes.jl @@ -0,0 +1,89 @@ +# Lifetime analysis when broadcast expressions are evaluated eagerly. + +function find_ndarray_assignments(ex, assigned_vars::Set{Symbol}) + cache = Dict{Any,Symbol}() # expression → temp mapping + local_assigned = Set{Symbol}() # track all assigned symbols + + function fresh_tmp(expr) + counter[] += 1 + tmp = Symbol(:tmp, counter[]) + cache[expr] = tmp + push!(local_assigned, tmp) + return tmp, [:($tmp = $expr)] + end + + function rewrite(e)::Tuple{Any,Vector{Expr}} + e isa Expr || return e, Expr[] + + if e.head == :(=) + lhs, rhs = e.args + lhs isa Symbol && push!(local_assigned, lhs) + new_rhs, temps = rewrite(rhs) + return :($lhs = $new_rhs), temps + end + + if e.head == :(.=) + lhs, rhs = e.args + new_lhs, lhs_temps = rewrite(lhs) + # Do not hoist the top-level call of the RHS to preserve fusion. + if rhs isa Expr && rhs.head == :call + op = rhs.args[1] + new_rhs_args, rhs_temps = Any[], Expr[] + for arg in rhs.args[2:end] + new_arg, temps = rewrite(arg) + push!(new_rhs_args, new_arg) + append!(rhs_temps, temps) + end + new_rhs = Expr(:call, op, new_rhs_args...) + return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) + end + + new_rhs, rhs_temps = rewrite(rhs) + return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) + end + + e.head == :ref && return fresh_tmp(e) + + if e.head == :call + op = e.args[1] + new_args, hoisted = Any[], Expr[] + for arg in e.args[2:end] + new_arg, temps = rewrite(arg) + push!(new_args, new_arg) + append!(hoisted, temps) + end + tmp, bind = fresh_tmp(Expr(:call, op, new_args...)) + return tmp, vcat(hoisted, bind) + end + + new_args, hoisted = Any[], Expr[] + is_block = e.head == :block || e.head == :begin + for arg in e.args + new_arg, temps = rewrite(arg) + if is_block && !(arg isa LineNumberNode) + append!(new_args, temps) + push!(new_args, new_arg) + else + push!(new_args, new_arg) + append!(hoisted, temps) + end + end + return Expr(e.head, new_args...), hoisted + end + + new_ex, temps = rewrite(ex) + union!(assigned_vars, local_assigned) + + if new_ex isa Expr && new_ex.head == :block + return Expr(:block, temps..., new_ex.args...) + end + return Expr(:block, temps..., new_ex) +end + +function process_lifetime_scope(scope) + assigned_vars = Set{Symbol}() + rewritten = find_ndarray_assignments(scope, assigned_vars) + result = insert_finalizers(rewritten, assigned_vars) + counter[] = 0 + return result +end diff --git a/src/scoping/scoping.jl b/src/scoping/scoping.jl new file mode 100644 index 000000000..ce6cf9a79 --- /dev/null +++ b/src/scoping/scoping.jl @@ -0,0 +1,258 @@ +export @analyze_lifetimes, @show_lifetimes + +@doc""" + @analyze_lifetimes expr + +Wraps a block of code so that all temporary `NDArray` allocations +(e.g. from slicing or function calls) are tracked and safely freed +at the end of the block. Ensures proper cleanup of GPU memory by +inserting `maybe_insert_delete` calls automatically. + +When broadcast fusion is enabled (`FUSE_BROADCAST_EXPRS`), dotted operators +(`.+`, `.*`, etc.) form a lazy `Base.Broadcast.Broadcasted` tree compiled into +a single PTX kernel; intermediate nodes are not real `NDArray` allocations and +are not individually hoisted. The macro automatically selects the +broadcast-aware analysis in that case and the plain analysis otherwise. +""" +macro analyze_lifetimes(block) + return esc(process_ndarray_scope(block)) +end + +const counter = Ref(0) + +function maybe_insert_delete(var::NDArray) + return cuNumeric.destroy!(var) +end + +maybe_insert_delete(x) = x + +""" + walk_symbols(x) -> Vector{Symbol} + +Recursively collect all symbols that appear inside expression `x`. +""" +function walk_symbols(x) + syms = Symbol[] + if x isa Symbol + push!(syms, x) + elseif x isa Expr + for a in x.args + append!(syms, walk_symbols(a)) + end + elseif x isa AbstractArray + for a in x + append!(syms, walk_symbols(a)) + end + end + return syms +end + +""" + insert_finalizers(stmts::Vector) +Insert `cuNumeric.maybe_insert_delete(var)` after the last use of each temporary variable. +""" +function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) + uses = Dict{Symbol,Vector{Int}}() + defs = Dict{Symbol,Int}() + alias_map = Dict{Symbol,Symbol}() + + # Collect all statements, flattening blocks and skipping LineNumberNodes + stmts = Any[] + for expr in exprs + if expr isa LineNumberNode + continue + elseif expr isa Expr && expr.head == :block + for arg in expr.args + arg isa LineNumberNode || push!(stmts, arg) + end + else + push!(stmts, expr) + end + end + + # Pass 1: collect definitions and uses + for (i, stmt) in enumerate(stmts) + stmt isa Expr || continue + stmt.head == :line && continue + + if stmt.head == :(=) + lhs, rhs = stmt.args + if lhs isa Symbol + defs[lhs] = i + end + if lhs isa Symbol && rhs isa Symbol + alias_map[lhs] = rhs + end + for s in walk_symbols(rhs) + push!(get!(uses, s, Int[]), i) + end + else + for s in walk_symbols(stmt) + push!(get!(uses, s, Int[]), i) + end + end + end + + for (alias, src) in alias_map + append!(get!(uses, src, Int[]), get(uses, alias, Int[])) + end + + # Compute last usage index per variable + last_use = Dict{Symbol,Int}() + for (v, idxs) in uses + last_use[v] = maximum(idxs) + end + + # `F_u = tmp1` aliases one NDArray under two names; resolve to a canonical rep + # so it's freed once (double-free is masked only by the ptr=0 null-out). + function canon(v) + seen = Set{Symbol}() + while haskey(alias_map, v) && !(v in seen) + push!(seen, v) + v = alias_map[v] + end + return v + end + + is_indexed_assign(s) = + s isa Expr && s.head in (:(=), :(.=)) && !(s.args[1] isa Symbol) + result_symbol(s) = + if s isa Symbol + s + else + (s isa Expr && s.head == :(=) && s.args[1] isa Symbol ? s.args[1] : nothing) + end + + # Pass 2: insert finalizers + out = Any[] + n = length(stmts) + freed = Set{Symbol}() + + # The block's value escapes to the caller, except for `A[...] = rhs`: Julia + # returns `rhs` there, but that's a dead temp nobody consumes — free it and + # return `nothing` rather than leak it or hand back a dangling handle. + terminal_indexed = n > 0 && is_indexed_assign(stmts[n]) + + protected = Set{Symbol}() + if n > 0 && !terminal_indexed + rs = result_symbol(stmts[n]) + rs isa Symbol && push!(protected, canon(rs)) + end + + function emit_delete!(v) + c = canon(v) + (c in freed || c in protected) && return nothing + push!(freed, c) + return push!(out, :(cuNumeric.maybe_insert_delete($v))) + end + + for (i, stmt) in enumerate(stmts) + # `v = w` aliases w, so don't finalize w at this statement. + skip_finalize = Set{Symbol}() + if stmt isa Expr && stmt.head == :(=) + lhs, rhs = stmt.args + if lhs isa Symbol && rhs isa Symbol + push!(skip_finalize, rhs) + end + end + + if i == n && !terminal_indexed + res_var = Symbol(:res, counter[]) + counter[] += 1 + push!(out, :($res_var = $stmt)) + else + push!(out, stmt) + end + + for (v, lasti) in last_use + if lasti == i && v ∈ assigned_vars && !(v ∈ skip_finalize) + emit_delete!(v) + end + end + + i == n && push!(out, terminal_indexed ? :nothing : res_var) + end + + return out +end + +""" + insert_finalizers(block::Expr) +Apply finalizer insertion to a `begin ... end` or `:block` expression. +""" +function insert_finalizers(block::Expr, assigned_vars::Set{Symbol}) + if block.head == :block || block.head == :begin + # Filter out LineNumberNodes before processing + stmts = [s for s in block.args if !(s isa LineNumberNode)] + new_stmts = insert_finalizers(stmts, assigned_vars) + return Expr(:block, new_stmts...) + else + error("Expected a begin/block expression") + end +end + +include("lifetimes.jl") +include("inter_broadcast_fusion.jl") +include("broadcast_lifetimes.jl") + +function process_ndarray_scope(scope) + # Broadcast expressions stay lazy only when fusion is enabled; otherwise + # every call is analyzed as an eager allocation. + @static if FUSE_BROADCAST_EXPRS + return process_broadcast_lifetime_scope(scope) + end + return process_lifetime_scope(scope) +end + +# Pretty-print the @analyze_lifetimes rewrite (see @show_lifetimes below). +_is_delete_call(s) = Meta.isexpr(s, :call) && s.args[1] == :(cuNumeric.maybe_insert_delete) + +# Flatten nested begin/blocks into a linear statement list, dropping line nodes. +function _flatten_stmts(x) + stmts = Any[] + function walk(e) + if Meta.isexpr(e, (:block, :begin)) + foreach(walk, e.args) + elseif !(e isa LineNumberNode) + push!(stmts, e) + end + end + walk(x) + return stmts +end + +function print_lifetime_analysis(block; io::IO=stdout) + rule = "-"^60 + stmts = _flatten_stmts(process_ndarray_scope(block)) + mode = FUSE_BROADCAST_EXPRS ? "fusion-aware" : "plain" + + println(io, "@analyze_lifetimes expansion ($mode analysis)\n", rule) + + n = 0 + for s in stmts + if InterBroadcastFusion.is_log_call(s) + continue + elseif _is_delete_call(s) + printstyled(io, lpad("✗ free ", 11), s.args[2], "\n"; color=:red) + else + n += 1 + println(io, lpad(n, 4), " ", s) + end + end + + println(io, rule) + return nothing +end + +@doc""" + @show_lifetimes expr + +Print the lifetime-analysis rewrite of `expr` — the same transformation +[`@analyze_lifetimes`](@ref) applies — without running it. Every statement is +shown in source order and each inserted `maybe_insert_delete` is highlighted so +you can see exactly where each temporary is freed. Pure AST work, so it runs on +CPU-only checkouts. +""" +macro show_lifetimes(block) + return :(print_lifetime_analysis($(QuoteNode(block)))) +end diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index 6f31dcede..367be7b53 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -114,29 +114,24 @@ function test_scoping_regressions(T, N) shifted = product .+ T(2) C[:, :] = shifted ./ T(3) end - _, events = cuNumeric._inline_single_use_broadcasts(source) - - @test length(events) == 1 - fused = sprint(Base.show_unquoted, events[1].fused) + rewritten = cuNumeric.InterBroadcastFusion.rewrite_scope(source) + stmts = cuNumeric._flatten_stmts(rewritten) + log_call = only(filter(cuNumeric.InterBroadcastFusion.is_log_call, stmts)) + fused_stmt = only(filter(!cuNumeric.InterBroadcastFusion.is_log_call, stmts)) + fused = sprint(Base.show_unquoted, fused_stmt) @test !occursin("product", fused) @test !occursin("shifted", fused) @test occursin(".=", fused) - old_debug = cuNumeric.BCAST_FUSION_DEBUG[] - try - cuNumeric.BCAST_FUSION_DEBUG[] = true - io = IOBuffer() - cuNumeric._maybe_log_inter_broadcast_fusion( - events[1].before, events[1].fused; io - ) - log = String(take!(io)) - @test occursin("inter-broadcast fusion rewrite", log) - @test occursin("product =", log) - @test occursin("shifted =", log) - @test occursin("fused", log) - finally - cuNumeric.BCAST_FUSION_DEBUG[] = old_debug - end + io = IOBuffer() + cuNumeric.InterBroadcastFusion.maybe_log_rewrite( + true, log_call.args[3].value, log_call.args[4].value; io + ) + log = String(take!(io)) + @test occursin("inter-broadcast fusion rewrite", log) + @test occursin("product =", log) + @test occursin("shifted =", log) + @test occursin("fused", log) end @testset "Indexed fused assignment preserves Array view semantics" begin From 4c640046ca8a0e70b4f68d56fd00cc3a3802b5a8 Mon Sep 17 00:00:00 2001 From: David Krasowska Date: Sun, 2 Aug 2026 15:08:31 -0500 Subject: [PATCH 04/12] Refactor lifetime analysis w/ MacroTools. (#157) --- docs/src/debugging.md | 15 +- docs/src/internals.md | 22 ++- lib/cunumeric_jl_wrapper/src/wrapper.cpp | 3 - src/scoping/broadcast_lifetimes.jl | 170 ++++++++---------- src/scoping/inter_broadcast_fusion.jl | 216 +++++++++++------------ src/scoping/lifetimes.jl | 107 +++++------ src/scoping/scoping.jl | 195 ++++++++++---------- src/scoping/util.jl | 138 +++++++++++++++ test/runtests.jl | 2 + test/tests/scoping.jl | 133 +++++++++++--- 10 files changed, 596 insertions(+), 405 deletions(-) create mode 100644 src/scoping/util.jl diff --git a/docs/src/debugging.md b/docs/src/debugging.md index 49f4eb0d7..05860988e 100644 --- a/docs/src/debugging.md +++ b/docs/src/debugging.md @@ -76,7 +76,9 @@ cuNumeric already supplies names for individual operations when task-scope namin When broadcast fusion is on, set `cuNumeric.BCAST_FUSION_DEBUG[] = true` to print inter-statement rewrites and each fused kernel's expression tree, -arguments, and launch geometry. +arguments, and launch geometry. Inter-statement rewrites are reported when +`@analyze_lifetimes` expands, so enable the flag before defining or evaluating +the expression you want to inspect. Kernel details are reported at runtime. ```julia using cuNumeric @@ -88,7 +90,10 @@ A = cuNumeric.ones(Float32, N, N) B = cuNumeric.ones(Float32, N, N) C = cuNumeric.zeros(Float32, N, N) -C .= @. A * B + 2.0f0 +@analyze_lifetimes begin + product = A .* B + C[:, :] = product .+ 2.0f0 +end cuNumeric.BCAST_FUSION_DEBUG[] = false ``` @@ -106,6 +111,10 @@ For example, a single-use producer inside `@analyze_lifetimes` is reported as: C[:, :] .= A .* B .+ 2.0f0 ``` +`before` contains exactly the statements that were recombined, and `fused` +contains their replacement. No rewrite block is printed when the pass leaves +the statements unchanged. + The fused kernel is reported separately: ```text @@ -126,7 +135,7 @@ How to read it: - `expr` is the fused op tree. - `inputs` / `scalars` are the runtime arguments passed into the kernel. - `arg_map` encodes how kernel slots map to the output (`0`), inputs (`>= 1`), and scalars (`< 0`). -- If nothing prints, the expression took the unfused path (for example shape-mismatched leaves, fusion disabled, or below the min-ops threshold). +- If no kernel block prints, the expression took the unfused path (for example shape-mismatched leaves, fusion disabled, or below the min-ops threshold). Turn the flag off when you are done. It prints on every fused launch and can be noisy in loops. diff --git a/docs/src/internals.md b/docs/src/internals.md index 6fa90b799..1afa20c03 100644 --- a/docs/src/internals.md +++ b/docs/src/internals.md @@ -15,7 +15,7 @@ We compile nested Julia broadcast expressions on `NDArray` to a single CUDA kern ## Lifetimes and GC -Julia's GC sees an `NDArray` as a small handle. The actual data is owned by the Legate runtime. This means that to Julia's GC `NDArrays` do not create memory pressure and GC is never executed. To avoid out-of-memory errors we created the `@analyze_lifetiems` macro so users can manually manage the lifetimes of a code block and also manually track device memory to automatically invoke GC. +Julia's GC sees an `NDArray` as a small handle. The actual data is owned by the Legate runtime. This means that to Julia's GC `NDArrays` do not create memory pressure and GC is never executed. To avoid out-of-memory errors we created the `@analyze_lifetimes` macro so users can manually manage the lifetimes of a code block and also manually track device memory to automatically invoke GC. ### Eager last-use freeing with `@analyze_lifetimes` @@ -36,6 +36,26 @@ end Use `@show_lifetimes` to print the rewritten block and the free sites without running the code. That is pure AST work and works without a GPU. +The implementation uses separate passes for inter-statement broadcast fusion, +allocation hoisting, and finalizer insertion. The top-level scoping pass selects +the appropriate lifetime analysis based on whether broadcast fusion is enabled. + +- **Inter-statement broadcast fusion** (`rewrite_scope`): + Merges single-use broadcast statements into their consumer, e.g.: + `p = A .* B; C .= p .+ 1` → `C .= A .* B .+ 1` + This pass operates only on syntax and does not depend on cuNumeric types. + +- **Lifetime analysis**: + - With fusion enabled (`rewrite_broadcast_lifetimes`): keeps dotted trees lazy + and hoists only materialized values (slices, non-broadcast calls, etc.) + - With fusion disabled (`rewrite_eager_lifetimes`): hoists all allocating calls + Both passes rename temps and insert free calls after static last uses. + +- **Finalizer insertion** (`insert_finalizers`): + Traverses the AST generated by either lifetime pass and inserts `delete` calls + at the computed last-use sites. It preserves the final return value of the + block by not freeing it. + Limits to keep in mind: - Analysis is statement-linear. It is not a full control-flow graph pass. Wrap hot loop bodies, not entire programs. diff --git a/lib/cunumeric_jl_wrapper/src/wrapper.cpp b/lib/cunumeric_jl_wrapper/src/wrapper.cpp index 03434a514..562ad7af7 100644 --- a/lib/cunumeric_jl_wrapper/src/wrapper.cpp +++ b/lib/cunumeric_jl_wrapper/src/wrapper.cpp @@ -73,9 +73,6 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { using jlcxx::TypeVar; using legate_util::HalfType; - // Map C++ complex types to Julia complex types - mod.map_type>("ComplexF64"); - mod.map_type>("ComplexF32"); mod.map_type("Float16"); // These are the types/dims used to generate templated functions diff --git a/src/scoping/broadcast_lifetimes.jl b/src/scoping/broadcast_lifetimes.jl index a6f377aec..62f2eb2e0 100644 --- a/src/scoping/broadcast_lifetimes.jl +++ b/src/scoping/broadcast_lifetimes.jl @@ -1,135 +1,121 @@ # Lifetime analysis for lazy broadcast expression trees. +# +# C[2:end-1, :] .= A[2:end-1, :] .* B[2:end-1, :] .+ 2 +# +# hoists only materialized values while leaving the dotted tree intact: +# +# tmp1 = @view C[2:end-1, :] +# tmp2 = A[2:end-1, :] +# tmp3 = B[2:end-1, :] +# tmp1 .= tmp2 .* tmp3 .+ 2 +# +# The destination view and input slices are objects that need lifetime management; +# the `.*` and `.+` nodes are lazy and become one fused broadcast kernel. -is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") - -function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) - local_assigned = Set{Symbol}() - - function fresh_tmp(expr) - counter[] += 1 - tmp = Symbol(:tmp, counter[]) - push!(local_assigned, tmp) - return tmp, [:($tmp = $expr)] - end +function _expand_view(expr) + return Base.macroexpand(@__MODULE__, :(@view $expr)) +end - function maphoist(f, args) - new_args, hoisted = Any[], Expr[] - for arg in args - new_arg, temps = f(arg) - push!(new_args, new_arg) - append!(hoisted, temps) - end - return new_args, hoisted - end +function rewrite_broadcast_lifetimes(scope) + assigned_vars = Set{Symbol}() + fresh_tmp(expr) = _hoist_temporary(expr, assigned_vars) # Inside a broadcast tree: hoist slices, keep dotted ops/f.(…) lazy, and - # delegate anything else to rewrite() (it breaks the tree → real NDArray). + # delegate anything else to rewrite_materialized() because it breaks the + # tree and produces a real NDArray. # The slice cache is scoped to one fused tree so repeated views become one # task argument without extending their lifetime across task submissions. - function rewrite_broadcast( + function rewrite_lazy_broadcast( expr, slice_cache::Dict{Any,Symbol} )::Tuple{Any,Vector{Expr}} - expr isa Expr || return expr, Expr[] - if expr.head == :ref + if !(expr isa Expr) + return expr, Expr[] + end + reference = _reference(expr) + if !isnothing(reference) cached = get(slice_cache, expr, nothing) - cached === nothing || return cached, Expr[] + if !isnothing(cached) + return cached, Expr[] + end tmp, bind = fresh_tmp(expr) slice_cache[expr] = tmp return tmp, bind end - if expr.head == :call && is_broadcast_op(expr.args[1]) - args, hoisted = maphoist( - arg -> rewrite_broadcast(arg, slice_cache), expr.args[2:end] + call = _call(expr) + if !isnothing(call) && _is_broadcast_op(call.f) + args, hoisted = _maphoist( + arg -> rewrite_lazy_broadcast(arg, slice_cache), call.args ) - return Expr(:call, expr.args[1], args...), hoisted + return Expr(:call, call.f, args...), hoisted end - if expr.head == :. && length(expr.args) == 2 && - expr.args[2] isa Expr && expr.args[2].head == :tuple - args, hoisted = maphoist( - arg -> rewrite_broadcast(arg, slice_cache), expr.args[2].args + + dotcall = _dotcall(expr) + if !isnothing(dotcall) + args, hoisted = _maphoist( + arg -> rewrite_lazy_broadcast(arg, slice_cache), dotcall.args ) - return Expr(:., expr.args[1], Expr(:tuple, args...)), hoisted + return Expr(:., dotcall.f, Expr(:tuple, args...)), hoisted end - return rewrite(expr) + return rewrite_materialized(expr) end - function rewrite(expr)::Tuple{Any,Vector{Expr}} - expr isa Expr || return expr, Expr[] - InterBroadcastFusion.is_log_call(expr) && return expr, Expr[] + function rewrite_materialized(expr)::Tuple{Any,Vector{Expr}} + if !(expr isa Expr) + return expr, Expr[] + end - if expr.head == :(=) - lhs, rhs = expr.args - lhs isa Symbol && push!(local_assigned, lhs) - new_rhs, temps = rewrite(rhs) + assignment = _assignment(expr) + if !isnothing(assignment) + (; lhs, rhs) = assignment + if lhs isa Symbol + push!(assigned_vars, lhs) + end + new_rhs, temps = rewrite_materialized(rhs) return :($lhs = $new_rhs), temps end # A `.=` RHS is a broadcast tree: only its slices are hoisted. - if expr.head == :(.=) - lhs, rhs = expr.args + broadcast_assignment = _broadcast_assignment(expr) + if !isnothing(broadcast_assignment) + (; lhs, rhs) = broadcast_assignment # An explicit view preserves indexed `.=` semantics for both Array # and NDArray while giving the lifetime pass a wrapper to destroy. - new_lhs, lhs_temps = - if lhs isa Expr && lhs.head == :ref - view_expr = Base.macroexpand( - @__MODULE__, - Expr( - :macrocall, - GlobalRef(Base, Symbol("@view")), - LineNumberNode(0), - lhs, - ), - ) - fresh_tmp(view_expr) - else - rewrite(lhs) - end - new_rhs, rhs_temps = rewrite_broadcast(rhs, Dict{Any,Symbol}()) + lhs_reference = _reference(lhs) + if isnothing(lhs_reference) + new_lhs, lhs_temps = rewrite_materialized(lhs) + else + new_lhs, lhs_temps = fresh_tmp(_expand_view(lhs)) + end + new_rhs, rhs_temps = rewrite_lazy_broadcast(rhs, Dict{Any,Symbol}()) return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) end - expr.head == :ref && return fresh_tmp(expr) + reference = _reference(expr) + if !isnothing(reference) + return fresh_tmp(expr) + end - if expr.head == :call && is_broadcast_op(expr.args[1]) - inner, hoisted = rewrite_broadcast(expr, Dict{Any,Symbol}()) + call = _call(expr) + if !isnothing(call) && _is_broadcast_op(call.f) + inner, hoisted = rewrite_lazy_broadcast(expr, Dict{Any,Symbol}()) tmp, bind = fresh_tmp(inner) return tmp, vcat(hoisted, bind) end - if expr.head == :call - args, hoisted = maphoist(rewrite, expr.args[2:end]) - tmp, bind = fresh_tmp(Expr(:call, expr.args[1], args...)) + if !isnothing(call) + args, hoisted = _maphoist(rewrite_materialized, call.args) + tmp, bind = fresh_tmp(Expr(:call, call.f, args...)) return tmp, vcat(hoisted, bind) end - new_args, hoisted = Any[], Expr[] - is_block = expr.head == :block || expr.head == :begin - for arg in expr.args - new_arg, temps = rewrite(arg) - if is_block && !(arg isa LineNumberNode) - append!(new_args, temps) - push!(new_args, new_arg) - else - push!(new_args, new_arg) - append!(hoisted, temps) - end - end - return Expr(expr.head, new_args...), hoisted + return _rewrite_children(rewrite_materialized, expr) end - new_ex, temps = rewrite(ex) - union!(assigned_vars, local_assigned) - if new_ex isa Expr && new_ex.head == :block - return Expr(:block, temps..., new_ex.args...) - end - return Expr(:block, temps..., new_ex) + rewritten, temps = rewrite_materialized(scope) + return _prepend_statements(rewritten, temps), assigned_vars end -function process_broadcast_lifetime_scope(scope) - assigned_vars = Set{Symbol}() - scope = InterBroadcastFusion.rewrite_scope(scope) - rewritten = find_broadcast_assignments(scope, assigned_vars) - result = insert_finalizers(rewritten, assigned_vars) - counter[] = 0 - return result +function process_broadcast_lifetime_scope(scope; on_rewrite=nothing) + scope = InterBroadcastFusion.rewrite_scope(scope; on_rewrite) + return _process_lifetime_scope(scope, rewrite_broadcast_lifetimes) end diff --git a/src/scoping/inter_broadcast_fusion.jl b/src/scoping/inter_broadcast_fusion.jl index e6175940a..0138c351c 100644 --- a/src/scoping/inter_broadcast_fusion.jl +++ b/src/scoping/inter_broadcast_fusion.jl @@ -2,46 +2,34 @@ module InterBroadcastFusion export rewrite_scope -import ..cuNumeric: walk_symbols - -_is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") - -function _is_broadcast_syntax(expr) - return expr isa Expr && - ( - (expr.head == :call && !isempty(expr.args) && _is_broadcast_op(expr.args[1])) || - (expr.head == :. && length(expr.args) == 2) - ) -end - -function _symbol_occurrences(expr, target::Symbol) - expr === target && return 1 - expr isa Expr || return 0 - return sum(arg -> _symbol_occurrences(arg, target), expr.args; init=0) -end +using ..ScopingUtils + +# Recombine single-use broadcast statements before lifetime analysis: +# +# product = A .* B +# C[:, :] = product .+ 2 +# +# becomes: +# +# C[:, :] .= A .* B .+ 2 +# +# The pass is syntax-only and has no NDArray or cuNumeric dependencies. function _substitute_symbols(expr, replacements::Dict{Symbol,Any}) - expr isa Symbol && return get(replacements, expr, expr) - expr isa Expr || return expr - - # A symbol assignment introduces/redefines its LHS; only substitute in RHS. - if expr.head == :(=) && expr.args[1] isa Symbol - return Expr( - :(=), expr.args[1], _substitute_symbols(expr.args[2], replacements) - ) - end - return Expr( - expr.head, (_substitute_symbols(arg, replacements) for arg in expr.args)... - ) + assignment = _assignment(expr) + isnothing(assignment) && return _replace_symbols(expr, replacements) + assignment.lhs isa Symbol || return _replace_symbols(expr, replacements) + rhs = _replace_symbols(assignment.rhs, replacements) + return :($(assignment.lhs) = $rhs) end function _indexed_assignment_base(stmt) - if stmt isa Expr && stmt.head == :(=) && - stmt.args[1] isa Expr && stmt.args[1].head == :ref - base = stmt.args[1].args[1] - return base isa Symbol ? base : nothing - end - return nothing + assignment = _assignment(stmt) + isnothing(assignment) && return nothing + reference = _reference(assignment.lhs) + isnothing(reference) && return nothing + reference.array isa Symbol || return nothing + return reference.array end function _safe_to_delay_broadcast( @@ -49,49 +37,82 @@ function _safe_to_delay_broadcast( ) for i in (def_idx + 1):(use_idx - 1) stmt = stmts[i] - stmt isa LineNumberNode && continue i in lazy_defs && continue # An indexed write to an unrelated array does not invalidate the lazy # producer. Any other intervening statement is conservatively a barrier. mutated = _indexed_assignment_base(stmt) - mutated !== nothing && !(mutated in dependencies) && continue + if !isnothing(mutated) && !(mutated in dependencies) + continue + end return false end return true end +function _single_use_index(stmts, symbol::Symbol, def_idx::Int) + use_idx = nothing + for i in (def_idx + 1):length(stmts) + symbols = walk_symbols(stmts[i]) + occurrences = count(candidate -> candidate == symbol, symbols) + occurrences == 0 && continue + if occurrences != 1 || !isnothing(use_idx) + return nothing + end + use_idx = i + end + return use_idx +end + +function _source_indices(expr, replacement_sources) + indices = Int[] + for symbol in walk_symbols(expr) + append!(indices, get(replacement_sources, symbol, Int[])) + end + unique!(indices) + sort!(indices) + return indices +end + +function _fuse_into_destination(stmt) + assignment = _assignment(stmt) + isnothing(assignment) && return stmt + _is_broadcast_syntax(assignment.rhs) || return stmt + + reference = _reference(assignment.lhs) + isnothing(reference) && return stmt + reference.array isa Symbol || return stmt + + if reference.array in walk_symbols(assignment.rhs) + return stmt + end + return Expr(:(.=), assignment.lhs, assignment.rhs) +end + function _rewrite_scope(scope) - Meta.isexpr(scope, (:block, :begin)) || return scope, NamedTuple[] - stmts = collect(scope.args) + stmts = _scope_statements(scope) + isnothing(stmts) && return scope, NamedTuple[] definitions = Dict{Symbol,Tuple{Int,Any}}() lazy_defs = Set{Int}() for (i, stmt) in enumerate(stmts) - if stmt isa Expr && stmt.head == :(=) && stmt.args[1] isa Symbol && - _is_broadcast_syntax(stmt.args[2]) - definitions[stmt.args[1]] = (i, stmt.args[2]) + assignment = _assignment(stmt) + if !isnothing(assignment) && assignment.lhs isa Symbol && + _is_broadcast_syntax(assignment.rhs) + definitions[assignment.lhs] = (i, assignment.rhs) push!(lazy_defs, i) end end - inlineable = Dict{Symbol,Tuple{Int,Int,Any}}() + inlineable = Dict{Symbol,Tuple{Int,Any}}() for (sym, (def_idx, rhs)) in definitions - use_indices = Int[] - occurrences = 0 - for i in (def_idx + 1):length(stmts) - count = _symbol_occurrences(stmts[i], sym) - if count > 0 - occurrences += count - push!(use_indices, i) - end - end - occurrences == 1 || continue - use_idx = only(use_indices) + use_idx = _single_use_index(stmts, sym, def_idx) + isnothing(use_idx) && continue dependencies = Set(walk_symbols(rhs)) - _safe_to_delay_broadcast(stmts, def_idx, use_idx, dependencies, lazy_defs) || + if !_safe_to_delay_broadcast(stmts, def_idx, use_idx, dependencies, lazy_defs) continue - inlineable[sym] = (def_idx, use_idx, rhs) + end + inlineable[sym] = (def_idx, rhs) end replacements = Dict{Symbol,Any}() @@ -104,50 +125,25 @@ function _rewrite_scope(scope) for (i, original_stmt) in enumerate(stmts) if i in removed sym = def_symbols[i] - source_indices = Int[] - for dependency in walk_symbols(original_stmt.args[2]) - haskey(replacement_sources, dependency) || continue - append!(source_indices, replacement_sources[dependency]) - end - unique!(source_indices) - sort!(source_indices) + assignment = _assignment(original_stmt) + source_indices = _source_indices(assignment.rhs, replacement_sources) push!(source_indices, i) replacement_sources[sym] = source_indices - replacements[sym] = _substitute_symbols(inlineable[sym][3], replacements) + replacements[sym] = _substitute_symbols(inlineable[sym][2], replacements) continue end - source_indices = Int[] - for dependency in walk_symbols(original_stmt) - haskey(replacement_sources, dependency) || continue - append!(source_indices, replacement_sources[dependency]) - end - unique!(source_indices) - sort!(source_indices) - + source_indices = _source_indices(original_stmt, replacement_sources) stmt = _substitute_symbols(original_stmt, replacements) - # A materialized indexed assignment with a broadcast RHS normally - # allocates a temporary and copies it into the slice. If the destination - # base is absent from the RHS, write the fused tree directly to the view. - direct_write = false - if stmt isa Expr && stmt.head == :(=) && - stmt.args[1] isa Expr && stmt.args[1].head == :ref && - _is_broadcast_syntax(stmt.args[2]) - base = stmt.args[1].args[1] - if base isa Symbol && !(base in Set(walk_symbols(stmt.args[2]))) - stmt = Expr(:(.=), stmt.args[1], stmt.args[2]) - direct_write = true - end - end - - if !isempty(source_indices) || direct_write + if !isempty(source_indices) + stmt = _fuse_into_destination(stmt) before = Expr( :block, - (deepcopy(stmts[source_idx]) for source_idx in source_indices)..., - deepcopy(original_stmt), + (stmts[source_idx] for source_idx in source_indices)..., + original_stmt, ) - push!(fusion_events, (; before, fused=deepcopy(stmt))) + push!(fusion_events, (; before, fused=stmt)) end push!(rewritten, stmt) end @@ -155,36 +151,25 @@ function _rewrite_scope(scope) return Expr(scope.head, rewritten...), fusion_events end -const _LOG_FUNCTION = GlobalRef(@__MODULE__, :maybe_log_rewrite) -const _DEBUG_FLAG = GlobalRef(parentmodule(@__MODULE__), :BCAST_FUSION_DEBUG) - -function _log_call(event) - before = QuoteNode(event.before) - fused = QuoteNode(event.fused) - enabled = Expr(:ref, _DEBUG_FLAG) - return Expr(:call, _LOG_FUNCTION, enabled, before, fused) -end - """ - rewrite_scope(scope) -> scope + rewrite_scope(scope; on_rewrite=nothing) -> scope Fuse eligible single-use broadcast producers into their consumer and return the -rewritten scope. Runtime log hooks are included for rewrites and are controlled -by `cuNumeric.BCAST_FUSION_DEBUG[]`. +rewritten scope. When provided, `on_rewrite` is called with a named tuple +containing the `before` and `fused` expressions for each rewrite. """ -function rewrite_scope(scope) +function rewrite_scope(scope; on_rewrite=nothing) rewritten, fusion_events = _rewrite_scope(scope) - isempty(fusion_events) && return rewritten - log_calls = (_log_call(event) for event in fusion_events) - return Expr(rewritten.head, log_calls..., rewritten.args...) -end - -function is_log_call(expr) - return Meta.isexpr(expr, :call) && expr.args[1] == _LOG_FUNCTION + if !isnothing(on_rewrite) + for event in fusion_events + on_rewrite(event) + end + end + return rewritten end function _print_expr(io::IO, expr) - clean = expr isa Expr ? Base.remove_linenums!(deepcopy(expr)) : expr + clean = _strip_lines(expr) rendered = sprint(Base.show_unquoted, clean) for line in eachline(IOBuffer(rendered)) println(io, " ", line) @@ -192,13 +177,12 @@ function _print_expr(io::IO, expr) return nothing end -function maybe_log_rewrite(enabled::Bool, before, fused; io::IO=stdout) - enabled || return nothing +function log_rewrite(event; io::IO=stdout) println(io, "\n", "="^40, " inter-broadcast fusion rewrite") println(io, " before") - _print_expr(io, before) + _print_expr(io, event.before) println(io, " fused") - _print_expr(io, fused) + _print_expr(io, event.fused) return nothing end diff --git a/src/scoping/lifetimes.jl b/src/scoping/lifetimes.jl index 730a5f5b9..b80c577d1 100644 --- a/src/scoping/lifetimes.jl +++ b/src/scoping/lifetimes.jl @@ -1,40 +1,46 @@ # Lifetime analysis when broadcast expressions are evaluated eagerly. +# +# result = f(A[2:end-1, :]) +# consume(result) +# +# becomes a linear sequence of named allocations: +# +# tmp1 = A[2:end-1, :] +# tmp2 = f(tmp1) +# result = tmp2 +# tmp3 = consume(result) +# tmp3 +# +# Finalizer insertion is a separate pass in scoping.jl. -function find_ndarray_assignments(ex, assigned_vars::Set{Symbol}) - cache = Dict{Any,Symbol}() # expression → temp mapping - local_assigned = Set{Symbol}() # track all assigned symbols - - function fresh_tmp(expr) - counter[] += 1 - tmp = Symbol(:tmp, counter[]) - cache[expr] = tmp - push!(local_assigned, tmp) - return tmp, [:($tmp = $expr)] - end +function rewrite_eager_lifetimes(scope) + assigned_vars = Set{Symbol}() + fresh_tmp(expr) = _hoist_temporary(expr, assigned_vars) - function rewrite(e)::Tuple{Any,Vector{Expr}} - e isa Expr || return e, Expr[] + function rewrite(expr)::Tuple{Any,Vector{Expr}} + if !(expr isa Expr) + return expr, Expr[] + end - if e.head == :(=) - lhs, rhs = e.args - lhs isa Symbol && push!(local_assigned, lhs) + assignment = _assignment(expr) + if !isnothing(assignment) + (; lhs, rhs) = assignment + if lhs isa Symbol + push!(assigned_vars, lhs) + end new_rhs, temps = rewrite(rhs) return :($lhs = $new_rhs), temps end - if e.head == :(.=) - lhs, rhs = e.args + broadcast_assignment = _broadcast_assignment(expr) + if !isnothing(broadcast_assignment) + (; lhs, rhs) = broadcast_assignment new_lhs, lhs_temps = rewrite(lhs) # Do not hoist the top-level call of the RHS to preserve fusion. - if rhs isa Expr && rhs.head == :call - op = rhs.args[1] - new_rhs_args, rhs_temps = Any[], Expr[] - for arg in rhs.args[2:end] - new_arg, temps = rewrite(arg) - push!(new_rhs_args, new_arg) - append!(rhs_temps, temps) - end - new_rhs = Expr(:call, op, new_rhs_args...) + call = _call(rhs) + if !isnothing(call) + new_rhs_args, rhs_temps = _maphoist(rewrite, call.args) + new_rhs = Expr(:call, call.f, new_rhs_args...) return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) end @@ -42,48 +48,25 @@ function find_ndarray_assignments(ex, assigned_vars::Set{Symbol}) return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) end - e.head == :ref && return fresh_tmp(e) + reference = _reference(expr) + if !isnothing(reference) + return fresh_tmp(expr) + end - if e.head == :call - op = e.args[1] - new_args, hoisted = Any[], Expr[] - for arg in e.args[2:end] - new_arg, temps = rewrite(arg) - push!(new_args, new_arg) - append!(hoisted, temps) - end - tmp, bind = fresh_tmp(Expr(:call, op, new_args...)) + call = _call(expr) + if !isnothing(call) + new_args, hoisted = _maphoist(rewrite, call.args) + tmp, bind = fresh_tmp(Expr(:call, call.f, new_args...)) return tmp, vcat(hoisted, bind) end - new_args, hoisted = Any[], Expr[] - is_block = e.head == :block || e.head == :begin - for arg in e.args - new_arg, temps = rewrite(arg) - if is_block && !(arg isa LineNumberNode) - append!(new_args, temps) - push!(new_args, new_arg) - else - push!(new_args, new_arg) - append!(hoisted, temps) - end - end - return Expr(e.head, new_args...), hoisted + return _rewrite_children(rewrite, expr) end - new_ex, temps = rewrite(ex) - union!(assigned_vars, local_assigned) - - if new_ex isa Expr && new_ex.head == :block - return Expr(:block, temps..., new_ex.args...) - end - return Expr(:block, temps..., new_ex) + rewritten, temps = rewrite(scope) + return _prepend_statements(rewritten, temps), assigned_vars end function process_lifetime_scope(scope) - assigned_vars = Set{Symbol}() - rewritten = find_ndarray_assignments(scope, assigned_vars) - result = insert_finalizers(rewritten, assigned_vars) - counter[] = 0 - return result + return _process_lifetime_scope(scope, rewrite_eager_lifetimes) end diff --git a/src/scoping/scoping.jl b/src/scoping/scoping.jl index ce6cf9a79..46494aec8 100644 --- a/src/scoping/scoping.jl +++ b/src/scoping/scoping.jl @@ -1,5 +1,11 @@ export @analyze_lifetimes, @show_lifetimes +# Include generic syntax layers before the cuNumeric-specific lifetime passes. +include("util.jl") +using .ScopingUtils + +include("inter_broadcast_fusion.jl") + @doc""" @analyze_lifetimes expr @@ -15,7 +21,8 @@ are not individually hoisted. The macro automatically selects the broadcast-aware analysis in that case and the plain analysis otherwise. """ macro analyze_lifetimes(block) - return esc(process_ndarray_scope(block)) + on_rewrite = BCAST_FUSION_DEBUG[] ? InterBroadcastFusion.log_rewrite : nothing + return esc(process_ndarray_scope(block; on_rewrite)) end const counter = Ref(0) @@ -26,81 +33,51 @@ end maybe_insert_delete(x) = x -""" - walk_symbols(x) -> Vector{Symbol} - -Recursively collect all symbols that appear inside expression `x`. -""" -function walk_symbols(x) - syms = Symbol[] - if x isa Symbol - push!(syms, x) - elseif x isa Expr - for a in x.args - append!(syms, walk_symbols(a)) - end - elseif x isa AbstractArray - for a in x - append!(syms, walk_symbols(a)) - end - end - return syms +function _hoist_temporary(expr, assigned_vars) + counter[] += 1 + temporary = Symbol(:tmp, counter[]) + push!(assigned_vars, temporary) + return temporary, [:($temporary = $expr)] end +# Turn the named allocations produced by either lifetime rewriter into an +# executable scope. For example, if tmp1 and tmp2 are last used by statement 3: +# +# 3 tmp3 = f(tmp1, tmp2) +# maybe_insert_delete(tmp1) +# maybe_insert_delete(tmp2) +# +# The final value of the block is protected because it escapes to the caller. """ insert_finalizers(stmts::Vector) Insert `cuNumeric.maybe_insert_delete(var)` after the last use of each temporary variable. """ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) - uses = Dict{Symbol,Vector{Int}}() - defs = Dict{Symbol,Int}() + last_use = Dict{Symbol,Int}() alias_map = Dict{Symbol,Symbol}() - # Collect all statements, flattening blocks and skipping LineNumberNodes - stmts = Any[] - for expr in exprs - if expr isa LineNumberNode - continue - elseif expr isa Expr && expr.head == :block - for arg in expr.args - arg isa LineNumberNode || push!(stmts, arg) - end - else - push!(stmts, expr) - end - end + stmts = _flatten_statements(Expr(:block, exprs...)) # Pass 1: collect definitions and uses for (i, stmt) in enumerate(stmts) stmt isa Expr || continue - stmt.head == :line && continue - - if stmt.head == :(=) - lhs, rhs = stmt.args - if lhs isa Symbol - defs[lhs] = i - end + assignment = _assignment(stmt) + used_expr = stmt + if !isnothing(assignment) + (; lhs, rhs) = assignment if lhs isa Symbol && rhs isa Symbol alias_map[lhs] = rhs end - for s in walk_symbols(rhs) - push!(get!(uses, s, Int[]), i) - end - else - for s in walk_symbols(stmt) - push!(get!(uses, s, Int[]), i) - end + used_expr = rhs + end + for symbol in walk_symbols(used_expr) + last_use[symbol] = i end end for (alias, src) in alias_map - append!(get!(uses, src, Int[]), get(uses, alias, Int[])) - end - - # Compute last usage index per variable - last_use = Dict{Symbol,Int}() - for (v, idxs) in uses - last_use[v] = maximum(idxs) + alias_last_use = get(last_use, alias, 0) + last_use[src] = max(get(last_use, src, 0), alias_last_use) end # `F_u = tmp1` aliases one NDArray under two names; resolve to a canonical rep @@ -114,14 +91,23 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) return v end - is_indexed_assign(s) = - s isa Expr && s.head in (:(=), :(.=)) && !(s.args[1] isa Symbol) - result_symbol(s) = - if s isa Symbol - s - else - (s isa Expr && s.head == :(=) && s.args[1] isa Symbol ? s.args[1] : nothing) + function is_indexed_assign(stmt) + assignment = _assignment(stmt) + if isnothing(assignment) + assignment = _broadcast_assignment(stmt) end + if !isnothing(assignment) + return !(assignment.lhs isa Symbol) + end + return false + end + function result_symbol(stmt) + stmt isa Symbol && return stmt + assignment = _assignment(stmt) + isnothing(assignment) && return nothing + assignment.lhs isa Symbol || return nothing + return assignment.lhs + end # Pass 2: insert finalizers out = Any[] @@ -133,26 +119,32 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) # return `nothing` rather than leak it or hand back a dangling handle. terminal_indexed = n > 0 && is_indexed_assign(stmts[n]) - protected = Set{Symbol}() + protected = nothing if n > 0 && !terminal_indexed rs = result_symbol(stmts[n]) - rs isa Symbol && push!(protected, canon(rs)) + if rs isa Symbol + protected = canon(rs) + end end function emit_delete!(v) c = canon(v) - (c in freed || c in protected) && return nothing + if c in freed || c == protected + return nothing + end push!(freed, c) - return push!(out, :(cuNumeric.maybe_insert_delete($v))) + push!(out, :(cuNumeric.maybe_insert_delete($v))) + return nothing end for (i, stmt) in enumerate(stmts) # `v = w` aliases w, so don't finalize w at this statement. - skip_finalize = Set{Symbol}() - if stmt isa Expr && stmt.head == :(=) - lhs, rhs = stmt.args + aliased_source = nothing + assignment = _assignment(stmt) + if !isnothing(assignment) + (; lhs, rhs) = assignment if lhs isa Symbol && rhs isa Symbol - push!(skip_finalize, rhs) + aliased_source = rhs end end @@ -165,12 +157,14 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) end for (v, lasti) in last_use - if lasti == i && v ∈ assigned_vars && !(v ∈ skip_finalize) + if lasti == i && v in assigned_vars && v != aliased_source emit_delete!(v) end end - i == n && push!(out, terminal_indexed ? :nothing : res_var) + if i == n + push!(out, terminal_indexed ? :nothing : res_var) + end end return out @@ -181,59 +175,56 @@ end Apply finalizer insertion to a `begin ... end` or `:block` expression. """ function insert_finalizers(block::Expr, assigned_vars::Set{Symbol}) - if block.head == :block || block.head == :begin - # Filter out LineNumberNodes before processing - stmts = [s for s in block.args if !(s isa LineNumberNode)] - new_stmts = insert_finalizers(stmts, assigned_vars) - return Expr(:block, new_stmts...) - else - error("Expected a begin/block expression") + stmts = _scope_statements(block) + isnothing(stmts) && error("Expected a begin/block expression") + return Expr(:block, insert_finalizers(stmts, assigned_vars)...) +end + +function _process_lifetime_scope(scope, rewrite_lifetimes) + try + rewritten, assigned_vars = rewrite_lifetimes(scope) + return insert_finalizers(rewritten, assigned_vars) + finally + counter[] = 0 end end +# Package-specific passes. The broadcast-aware pass also consumes the generic +# inter-broadcast fusion module included above. include("lifetimes.jl") -include("inter_broadcast_fusion.jl") include("broadcast_lifetimes.jl") -function process_ndarray_scope(scope) +function process_ndarray_scope(scope; on_rewrite=nothing) # Broadcast expressions stay lazy only when fusion is enabled; otherwise # every call is analyzed as an eager allocation. @static if FUSE_BROADCAST_EXPRS - return process_broadcast_lifetime_scope(scope) + return process_broadcast_lifetime_scope(scope; on_rewrite) end return process_lifetime_scope(scope) end -# Pretty-print the @analyze_lifetimes rewrite (see @show_lifetimes below). -_is_delete_call(s) = Meta.isexpr(s, :call) && s.args[1] == :(cuNumeric.maybe_insert_delete) - -# Flatten nested begin/blocks into a linear statement list, dropping line nodes. -function _flatten_stmts(x) - stmts = Any[] - function walk(e) - if Meta.isexpr(e, (:block, :begin)) - foreach(walk, e.args) - elseif !(e isa LineNumberNode) - push!(stmts, e) - end +# Return the deleted value for a generated finalizer call. +function _delete_argument(expr) + call = _call(expr) + isnothing(call) && return nothing + if call.f != :(cuNumeric.maybe_insert_delete) + return nothing end - walk(x) - return stmts + return only(call.args) end function print_lifetime_analysis(block; io::IO=stdout) rule = "-"^60 - stmts = _flatten_stmts(process_ndarray_scope(block)) + stmts = _flatten_statements(process_ndarray_scope(block)) mode = FUSE_BROADCAST_EXPRS ? "fusion-aware" : "plain" println(io, "@analyze_lifetimes expansion ($mode analysis)\n", rule) n = 0 for s in stmts - if InterBroadcastFusion.is_log_call(s) - continue - elseif _is_delete_call(s) - printstyled(io, lpad("✗ free ", 11), s.args[2], "\n"; color=:red) + deleted = _delete_argument(s) + if !isnothing(deleted) + printstyled(io, lpad("✗ free ", 11), deleted, "\n"; color=:red) else n += 1 println(io, lpad(n, 4), " ", s) diff --git a/src/scoping/util.jl b/src/scoping/util.jl new file mode 100644 index 000000000..416e57757 --- /dev/null +++ b/src/scoping/util.jl @@ -0,0 +1,138 @@ +module ScopingUtils + +using MacroTools: MacroTools + +# Shared syntax vocabulary for every scoping pass. For example: +# +# _assignment(:(result = A .+ B)) +# -> (lhs=:result, rhs=:(A .+ B)) +# _reference(:(A[2:end-1, :])) +# -> (array=:A, indices=Any[:(2:end-1), :(:)]) +# +# Keeping these matches here means the transformation files describe policy +# instead of repeating Expr head/argument indexing. + +export _assignment, _broadcast_assignment, _call, _dotcall, + _flatten_statements, _is_broadcast_op, _is_broadcast_syntax, _maphoist, + _reference, _prepend_statements, _replace_symbols, _rewrite_children, + _scope_statements, _strip_lines, walk_symbols + +function _assignment(expr) + MacroTools.isexpr(expr, :(=)) || return nothing + MacroTools.@capture(expr, lhs_ = rhs_) || return nothing + return (; lhs, rhs) +end + +function _broadcast_assignment(expr) + MacroTools.isexpr(expr, :(.=)) || return nothing + MacroTools.@capture(expr, lhs_ .= rhs_) || return nothing + return (; lhs, rhs) +end + +function _call(expr) + MacroTools.isexpr(expr, :call) || return nothing + MacroTools.@capture(expr, f_(args__)) || return nothing + return (; f, args) +end + +function _dotcall(expr) + MacroTools.isexpr(expr, :.) || return nothing + MacroTools.@capture(expr, f_.(args__)) || return nothing + return (; f, args) +end + +_is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") + +function _is_broadcast_syntax(expr) + call = _call(expr) + if !isnothing(call) && _is_broadcast_op(call.f) + return true + end + return !isnothing(_dotcall(expr)) +end + +function _reference(expr) + MacroTools.isexpr(expr, :ref) || return nothing + MacroTools.@capture(expr, array_[indices__]) || return nothing + return (; array, indices) +end + +_is_line(expr) = MacroTools.isline(expr) +_is_scope(expr) = MacroTools.isexpr(expr, :block, :begin) +_strip_lines(expr) = MacroTools.striplines(expr) + +function _scope_statements(scope) + _is_scope(scope) || return nothing + return MacroTools.rmlines(scope).args +end + +function _flatten_statements(scope) + flattened = MacroTools.flatten(scope) + _is_line(flattened) && return Any[] + statements = _scope_statements(flattened) + isnothing(statements) && return Any[flattened] + return Any[statements...] +end + +function _maphoist(transform, expressions) + rewritten = Any[] + hoisted = Expr[] + for expr in expressions + new_expr, temps = transform(expr) + push!(rewritten, new_expr) + append!(hoisted, temps) + end + return rewritten, hoisted +end + +function _rewrite_children(transform, expr::Expr) + rewritten = Any[] + hoisted = Expr[] + is_scope = _is_scope(expr) + for arg in expr.args + new_arg, temps = transform(arg) + if is_scope && !_is_line(arg) + append!(rewritten, temps) + push!(rewritten, new_arg) + else + push!(rewritten, new_arg) + append!(hoisted, temps) + end + end + return Expr(expr.head, rewritten...), hoisted +end + +function _prepend_statements(expr, statements) + if _is_scope(expr) + return Expr(:block, statements..., expr.args...) + end + return Expr(:block, statements..., expr) +end + +function _replace_symbols(expr, replacements::AbstractDict{Symbol}) + return MacroTools.postwalk(expr) do node + node isa Symbol || return node + return get(replacements, node, node) + end +end + +""" + walk_symbols(x) -> Vector{Symbol} + +Recursively collect all symbols that appear inside expression `x`. +""" +function walk_symbols(x) + syms = Symbol[] + MacroTools.postwalk(x) do node + node isa Symbol && push!(syms, node) + if node isa AbstractArray + for element in node + append!(syms, walk_symbols(element)) + end + end + return node + end + return syms +end + +end diff --git a/test/runtests.jl b/test/runtests.jl index 7ec03491e..b2dba0f7b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -446,6 +446,8 @@ end @testset verbose = true "Scoping" begin N = 100 + @testset "Rewrite pipeline" test_scoping_rewrite_pipeline() + @testset verbose = true for T in Base.uniontypes(cuNumeric.SUPPORTED_FLOAT_TYPES) allowscalar() do results = run_all_ops(T, N) diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index 367be7b53..f3671589d 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -86,6 +86,113 @@ const SLICE_OPS = Dict( ), ) +function test_scoping_rewrite_pipeline() + utils = cuNumeric.ScopingUtils + + @testset "Syntax helpers" begin + @test utils._assignment(:(x = y)) == (lhs=:x, rhs=:y) + @test utils._broadcast_assignment(:(A[:] .= x)).rhs == :x + @test utils._call(:(f(x, y))) == (f=:f, args=Any[:x, :y]) + @test utils._dotcall(:(f.(x, y))) == (f=:f, args=Any[:x, :y]) + @test utils._reference(:(A[i, j])) == (array=:A, indices=Any[:i, :j]) + @test utils._is_broadcast_syntax(:(x .* y)) + @test utils._replace_symbols(:(x .+ y), Dict(:x => :(a .* b))) == + :((a .* b) .+ y) + @test isnothing( + utils._assignment(quote + x = y + end), + ) + end + + @testset "Inter-broadcast fusion" begin + source = quote + product = A .* B + shifted = product .+ 2 + C[:, :] = shifted ./ 3 + end + events = NamedTuple[] + rewritten = cuNumeric.InterBroadcastFusion.rewrite_scope( + source; on_rewrite=event -> push!(events, event) + ) + stmts = utils._flatten_statements(rewritten) + fused = sprint(Base.show_unquoted, only(stmts)) + @test !occursin("product", fused) + @test !occursin("shifted", fused) + @test occursin(".=", fused) + + io = IOBuffer() + cuNumeric.InterBroadcastFusion.log_rewrite(only(events); io) + log = String(take!(io)) + @test occursin("product =", log) + @test occursin("shifted =", log) + @test occursin("fused", log) + + untouched = quote + C[:, :] = A .* B + end + empty!(events) + rewritten = cuNumeric.InterBroadcastFusion.rewrite_scope( + untouched; on_rewrite=event -> push!(events, event) + ) + @test isempty(events) + @test utils._strip_lines(rewritten) == utils._strip_lines(untouched) + end + + @testset "Fusion-aware lifetime rewrite" begin + source = quote + C[2:(end - 1), 2:(end - 1)] .= + A[2:(end - 1), 2:(end - 1)] .* B[2:(end - 1), 2:(end - 1)] .+ 2 + end + + cuNumeric.counter[] = 0 + try + rewritten, assigned = cuNumeric.rewrite_broadcast_lifetimes(source) + stmts = utils._flatten_statements(rewritten) + + @test assigned == Set([:tmp1, :tmp2, :tmp3]) + @test utils._assignment(stmts[1]).lhs == :tmp1 + @test utils._assignment(stmts[2]) == + (lhs=:tmp2, rhs=:(A[2:(end - 1), 2:(end - 1)])) + @test utils._assignment(stmts[3]) == + (lhs=:tmp3, rhs=:(B[2:(end - 1), 2:(end - 1)])) + @test !isnothing(utils._broadcast_assignment(stmts[4])) + + finalized = cuNumeric.insert_finalizers(rewritten, assigned) + freed = Set{Symbol}() + for stmt in utils._flatten_statements(finalized) + argument = cuNumeric._delete_argument(stmt) + isnothing(argument) || push!(freed, argument) + end + @test freed == assigned + finally + cuNumeric.counter[] = 0 + end + end + + @testset "Eager lifetime rewrite" begin + source = quote + result = f(A[2:(end - 1), 2:(end - 1)]) + consume(result) + end + + cuNumeric.counter[] = 0 + try + rewritten, assigned = cuNumeric.rewrite_eager_lifetimes(source) + stmts = utils._flatten_statements(rewritten) + + @test assigned == Set([:result, :tmp1, :tmp2, :tmp3]) + @test utils._assignment(stmts[1]) == + (lhs=:tmp1, rhs=:(A[2:(end - 1), 2:(end - 1)])) + @test utils._call(utils._assignment(stmts[2]).rhs).f == :f + @test utils._assignment(stmts[3]) == (lhs=:result, rhs=:tmp2) + @test utils._call(utils._assignment(stmts[4]).rhs).f == :consume + finally + cuNumeric.counter[] = 0 + end + end +end + function test_scoping_regressions(T, N) A = cuNumeric.ones(T, (N, N)) B = cuNumeric.ones(T, (N, N)) @@ -108,32 +215,6 @@ function test_scoping_regressions(T, N) end if cuNumeric.FUSE_BROADCAST_EXPRS - @testset "Inter-broadcast fusion logger" begin - source = quote - product = A .* B - shifted = product .+ T(2) - C[:, :] = shifted ./ T(3) - end - rewritten = cuNumeric.InterBroadcastFusion.rewrite_scope(source) - stmts = cuNumeric._flatten_stmts(rewritten) - log_call = only(filter(cuNumeric.InterBroadcastFusion.is_log_call, stmts)) - fused_stmt = only(filter(!cuNumeric.InterBroadcastFusion.is_log_call, stmts)) - fused = sprint(Base.show_unquoted, fused_stmt) - @test !occursin("product", fused) - @test !occursin("shifted", fused) - @test occursin(".=", fused) - - io = IOBuffer() - cuNumeric.InterBroadcastFusion.maybe_log_rewrite( - true, log_call.args[3].value, log_call.args[4].value; io - ) - log = String(take!(io)) - @test occursin("inter-broadcast fusion rewrite", log) - @test occursin("product =", log) - @test occursin("shifted =", log) - @test occursin("fused", log) - end - @testset "Indexed fused assignment preserves Array view semantics" begin src = reshape(T.(1:20), 4, 5) out = fill(T(-1), 6, 7) From bca198f5d0f5841531954e9da50ffa7548eec876 Mon Sep 17 00:00:00 2001 From: krasow Date: Sun, 2 Aug 2026 15:13:25 -0500 Subject: [PATCH 05/12] add back in benchmarks.toml --- benchmark/benchmarks.toml | 59 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 benchmark/benchmarks.toml diff --git a/benchmark/benchmarks.toml b/benchmark/benchmarks.toml new file mode 100644 index 000000000..70d5a47a6 --- /dev/null +++ b/benchmark/benchmarks.toml @@ -0,0 +1,59 @@ +[Global] +n_warmup = 5 +n_iter = 1000 +n_trial = 5 +cupynumeric = true # (needs install_cupynumeric.sh) +cuda = false # compare against CUDA.jl (single-GPU configs only) +# One CPU-reference check per config (not per timed iter). Written to CSV. +check_correctness = true +n_correctness_iter = 5 + +#################################### +# GEMM # +# Weak scaling, restricted to NxN. # +# Work ~ 2*N^2*M = 2*N^3. # +# N^3 / P --> constant. # +# N = baseline * P^(1/3). # +#################################### + +[[gemm]] +T = ["Float32"] +gpus = [1, 2, 4, 8] +cpus = 16 +N = [20000, 25200, 31752, 40000] +M = [20000, 25200, 31752, 40000] + +################################# +# Gray-Scott # +# Weak scaling, square NxN grid.# +# Work ~ N*M = N^2. # +# N^2 / P --> constant. # +# N = baseline * P^(1/2). # +################################# + +[[grayscott_baseline]] +T = "Float32" +gpus = [1, 2, 4, 8] +cpus = 16 +fusion = [true, false] +N = [2000, 2832, 4000, 5656] +M = [2000, 2832, 4000, 5656] + +[[grayscott_lifetimes]] +T = "Float32" +gpus = [1, 2, 4, 8] +cpus = 16 +fusion = false +N = [2000, 2832, 4000, 5656] +M = [2000, 2832, 4000, 5656] + +################################# +# Monte-Carlo Integration # +# Work ~ N. Scale N linearly # +################################# + +[[montecarlo]] +T = "Float32" +gpus = [1, 2, 4, 8] +cpus = 16 +N = [1_000_000, 2_000_000, 4_000_000, 8_000_000] From 554251d579780624189f8fed3b1efa8210831832 Mon Sep 17 00:00:00 2001 From: krasow Date: Mon, 3 Aug 2026 16:05:14 -0500 Subject: [PATCH 06/12] encapsulate scoped expr in a let block --- src/scoping/scoping.jl | 77 ++++++++++++++++++++++++++++++++++++------ test/tests/scoping.jl | 75 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 11 deletions(-) diff --git a/src/scoping/scoping.jl b/src/scoping/scoping.jl index 46494aec8..f729f538d 100644 --- a/src/scoping/scoping.jl +++ b/src/scoping/scoping.jl @@ -14,6 +14,10 @@ Wraps a block of code so that all temporary `NDArray` allocations at the end of the block. Ensures proper cleanup of GPU memory by inserting `maybe_insert_delete` calls automatically. +Assignments created inside the macro are scoped to its lexical region. Existing +arrays can still be mutated in place, and the final value of the block is +returned, but internal bindings do not leak into the surrounding scope. + When broadcast fusion is enabled (`FUSE_BROADCAST_EXPRS`), dotted operators (`.+`, `.*`, etc.) form a lazy `Base.Broadcast.Broadcasted` tree compiled into a single PTX kernel; intermediate nodes are not real `NDArray` allocations and @@ -22,7 +26,9 @@ broadcast-aware analysis in that case and the plain analysis otherwise. """ macro analyze_lifetimes(block) on_rewrite = BCAST_FUSION_DEBUG[] ? InterBroadcastFusion.log_rewrite : nothing - return esc(process_ndarray_scope(block; on_rewrite)) + rewritten = process_ndarray_scope(block; on_rewrite) + bindings = union(_assigned_symbols(block), _assigned_symbols(rewritten)) + return esc(_lexical_scope(rewritten, bindings)) end const counter = Ref(0) @@ -33,6 +39,43 @@ end maybe_insert_delete(x) = x +# `@analyze_lifetimes` is an ownership region, analogous to a C++ `{ ... }` +# block. Bind every source and generated assignment explicitly so it cannot +# accidentally reuse or leak a caller local with the same name. Indexed and +# broadcast assignments are mutations, not new bindings, and remain visible. +function _assigned_symbols(expr) + assigned = Set{Symbol}() + + function collect_binding(lhs) + if lhs isa Symbol + push!(assigned, lhs) + elseif lhs isa Expr && lhs.head in (:tuple, :parameters) + foreach(collect_binding, lhs.args) + elseif lhs isa Expr && lhs.head in (:(::), :(...)) + collect_binding(first(lhs.args)) + end + return nothing + end + + function visit(node) + node isa Expr || return nothing + assignment = _assignment(node) + if !isnothing(assignment) + collect_binding(assignment.lhs) + end + foreach(visit, node.args) + return nothing + end + + visit(expr) + return assigned +end + +function _lexical_scope(body, bindings::Set{Symbol}) + ordered = sort!(collect(bindings); by=string) + return Expr(:let, Expr(:block, ordered...), body) +end + function _hoist_temporary(expr, assigned_vars) counter[] += 1 temporary = Symbol(:tmp, counter[]) @@ -101,12 +144,25 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) end return false end - function result_symbol(stmt) - stmt isa Symbol && return stmt + function result_symbols(stmt) + stmt isa Symbol && return Set([stmt]) + assignment = _assignment(stmt) - isnothing(assignment) && return nothing - assignment.lhs isa Symbol || return nothing - return assignment.lhs + if !isnothing(assignment) + return result_symbols(assignment.rhs) + end + + if stmt isa Expr && stmt.head in (:tuple, :parameters) + return mapreduce(result_symbols, union, stmt.args; init=Set{Symbol}()) + end + if stmt isa Expr && stmt.head == :kw + return result_symbols(last(stmt.args)) + end + if stmt isa Expr && stmt.head == :(::) + return result_symbols(first(stmt.args)) + end + + return Set{Symbol}() end # Pass 2: insert finalizers @@ -119,17 +175,16 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) # return `nothing` rather than leak it or hand back a dangling handle. terminal_indexed = n > 0 && is_indexed_assign(stmts[n]) - protected = nothing + protected = Set{Symbol}() if n > 0 && !terminal_indexed - rs = result_symbol(stmts[n]) - if rs isa Symbol - protected = canon(rs) + for result in result_symbols(stmts[n]) + push!(protected, canon(result)) end end function emit_delete!(v) c = canon(v) - if c in freed || c == protected + if c in freed || c in protected return nothing end push!(freed, c) diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index f3671589d..54e0c2931 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -191,6 +191,81 @@ function test_scoping_rewrite_pipeline() cuNumeric.counter[] = 0 end end + + @testset "Multiple returned bindings" begin + source = quote + tmp1 = f(A) + first_result = tmp1 + tmp2 = g(A) + second_result = tmp2 + (first_result, second_result) + end + assigned = Set([:tmp1, :first_result, :tmp2, :second_result]) + finalized = cuNumeric.insert_finalizers(source, assigned) + freed = Set( + filter( + !isnothing, map(cuNumeric._delete_argument, utils._flatten_statements(finalized)) + ), + ) + + @test isempty(freed) + end + + @testset "Lexical lifetime scope" begin + function hidden_binding() + @analyze_lifetimes begin + internal_result = 41 + nothing + end + return internal_result + end + + function shadowed_binding() + internal_result = :outer + @analyze_lifetimes begin + internal_result = :inner + nothing + end + return internal_result + end + + function hidden_destructured_bindings() + @analyze_lifetimes begin + internal_first, internal_second = (1, 2) + nothing + end + return internal_first, internal_second + end + + output = [0] + returned = @analyze_lifetimes begin + internal_result = 42 + output[1] = internal_result + internal_result + end + + @test_throws UndefVarError hidden_binding() + @test_throws UndefVarError hidden_destructured_bindings() + @test shadowed_binding() == :outer + @test output == [42] + @test returned == 42 + + if cuNumeric.FUSE_BROADCAST_EXPRS + function hidden_fused_binding(a, b, destination) + @analyze_lifetimes begin + fused_result = a .* b + destination .= fused_result .+ 1 + end + return fused_result + end + + destination = zeros(Int, 2) + @test_throws UndefVarError hidden_fused_binding( + [2, 3], [4, 5], destination + ) + @test destination == [9, 16] + end + end end function test_scoping_regressions(T, N) From 3f0e342d2133c8a109f3f9ca97700b8c783ce4ee Mon Sep 17 00:00:00 2001 From: krasow Date: Mon, 3 Aug 2026 18:56:15 -0500 Subject: [PATCH 07/12] lifetime scoping: scalars don't need to become hoisted temps. And updated pretty printers --- docs/src/debugging.md | 42 ++++++++------ src/ndarray/broadcast_fusion.jl | 86 +++++++++++++++++++--------- src/scoping/broadcast_lifetimes.jl | 18 +++--- src/scoping/lifetimes.jl | 3 + src/scoping/util.jl | 26 ++++++++- test/tests/broadcast_fusion_tests.jl | 23 ++++++++ test/tests/scoping.jl | 34 ++++++++--- 7 files changed, 168 insertions(+), 64 deletions(-) diff --git a/docs/src/debugging.md b/docs/src/debugging.md index 05860988e..6b25941d5 100644 --- a/docs/src/debugging.md +++ b/docs/src/debugging.md @@ -38,7 +38,8 @@ export LEGATE_CONFIG="--gpus 1 --cpus 4 --logging legate=debug --log-to-file" julia --project=. workload.jl ``` -cuNumeric operations now carry short labels such as `zeros`, `matmul`, and `broadcast.+(*(input0, input1), scalar0)`. Search for those labels in `legate_*.log` to connect runtime messages to calls in `workload.jl`. +With task-scope naming enabled and Julia restarted, cuNumeric operations carry short labels such as `zeros`, `matmul`, and +`broadcast.+(*(input{0}, input{1}), 2.0f0)`. Search for those labels in `legate_*.log` to connect runtime messages to calls in `workload.jl`. For a timeline instead, replace the logging flags with `--profile`, run the same workload, and process the resulting `legate_*.prof` files with `legate_prof`. @@ -85,14 +86,14 @@ using cuNumeric cuNumeric.BCAST_FUSION_DEBUG[] = true -N = 8 +N = 10 A = cuNumeric.ones(Float32, N, N) B = cuNumeric.ones(Float32, N, N) C = cuNumeric.zeros(Float32, N, N) @analyze_lifetimes begin - product = A .* B - C[:, :] = product .+ 2.0f0 + product = A[2:end-1, 2:end-1] .* B[2:end-1, 2:end-1] + C[2:end-1, 2:end-1] = product .+ 2.0f0 end cuNumeric.BCAST_FUSION_DEBUG[] = false @@ -104,11 +105,11 @@ For example, a single-use producer inside `@analyze_lifetimes` is reported as: ======================================== inter-broadcast fusion rewrite before begin - product = A .* B - C[:, :] = product .+ 2.0f0 + product = A[2:end - 1, 2:end - 1] .* B[2:end - 1, 2:end - 1] + C[2:end - 1, 2:end - 1] = product .+ 2.0f0 end fused - C[:, :] .= A .* B .+ 2.0f0 + C[2:end - 1, 2:end - 1] .= A[2:end - 1, 2:end - 1] .* B[2:end - 1, 2:end - 1] .+ 2.0f0 ``` `before` contains exactly the statements that were recombined, and `fused` @@ -119,22 +120,25 @@ The fused kernel is reported separately: ```text ======================================== fused broadcast kernel - expr +(*(NDArray, NDArray), 2.0f0) - output NDArray{Float32, 2, false, Nothing} (8, 8) - inputs 2 unique NDArray(s) - [0] NDArray{Float32, 2, false, Nothing} (8, 8) - [1] NDArray{Float32, 2, false, Nothing} (8, 8) - scalars 2.0f0 - arg_map [0, 1, 2, -1] (0=output, >=1=input idx+1, <0=scalar) - launch host thread budget=1024, indexing=linear, blocks=device(local tile), global_ndrange=(8, 8) - call broadcast.gpu_broadcast_kernel_linear_splat(input0, input1, scalar0) + expr +(*(input{0}, input{1}), 2.0f0) + output NDArray{Float32, 2} (8, 8) slice, parent (10, 10) + inputs input{N} (2 unique) + N value + 0 NDArray{Float32, 2} (8, 8) slice, parent (10, 10) + 1 NDArray{Float32, 2} (8, 8) slice, parent (10, 10) + launch host thread budget=1024, indexing=cartesian, blocks=device(local tile), global_ndrange=(8, 8) + call broadcast.gpu_broadcast_kernel_cartesian_splat(input{0}, input{1}, 2.0f0) ``` How to read it: -- `expr` is the fused op tree. -- `inputs` / `scalars` are the runtime arguments passed into the kernel. -- `arg_map` encodes how kernel slots map to the output (`0`), inputs (`>= 1`), and scalars (`< 0`). +- `expr` is the fused op tree. Its `input{N}` names index the `inputs` table; + scalar arguments appear directly as their runtime values. +- `inputs` summarizes each runtime array without exposing internal `NDArray` + type parameters. A slice includes its parent shape. Source names such as `A` + are available in the static lifetime rewrite, but Julia does not retain those + binding names on the runtime array object. +- `call` shows the packed kernel argument order, including reused inputs. - If no kernel block prints, the expression took the unfused path (for example shape-mismatched leaves, fusion disabled, or below the min-ops threshold). Turn the flag off when you are done. It prints on every fused launch and can be noisy in loops. diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 2962e5a7d..cc6c634ec 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -401,7 +401,7 @@ function get_ptx( # dump_module=true keeps only_entry=false so linked libdevice helpers (e.g. cos # slowpaths) are not emptied into unresolved .externs before cuModuleLoad. CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, arg_types...); - raw=false, dump_module=true, kernel=true, ptx=v"9.0") + raw=false, dump_module=true, kernel=true, ptx=v"7.8") return String(take!(buf)), threads, ctx end @@ -456,18 +456,21 @@ function _bcast_tree_str(bc::Base.Broadcast.Broadcasted) end end -function _bcast_scope_name(bc::Base.Broadcast.Broadcasted, ndarray_to_input_idx) +function _bcast_runtime_tree_str( + bc::Base.Broadcast.Broadcasted, ndarray_to_input_idx, actual_scalars=nothing +) scalar_idx = 0 - tree = _bcast_tree_str(bc) do x + return _bcast_tree_str(bc) do x if x isa NDArray input_idx = get(ndarray_to_input_idx, objectid(x), nothing) - return input_idx === nothing ? "NDArray" : string("input", input_idx) + return input_idx === nothing ? "NDArray" : "input{$input_idx}" end if x isa Number idx = scalar_idx scalar_idx += 1 - return string("scalar", idx) + value = isnothing(actual_scalars) ? x : actual_scalars[idx + 1] + return repr(value) end if x isa Base.RefValue @@ -475,14 +478,24 @@ function _bcast_scope_name(bc::Base.Broadcast.Broadcasted, ndarray_to_input_idx) if value isa Number idx = scalar_idx scalar_idx += 1 - return string("scalar", idx) + runtime_value = + isnothing(actual_scalars) ? value : actual_scalars[idx + 1] + return repr(runtime_value) end return repr(value) end return string("<", typeof(x), ">") end - return string("broadcast.", tree) +end + +function _bcast_scope_name( + bc::Base.Broadcast.Broadcasted, ndarray_to_input_idx, actual_scalars=nothing +) + return string( + "broadcast.", + _bcast_runtime_tree_str(bc, ndarray_to_input_idx, actual_scalars), + ) end # Recover the kernel's plain name @@ -493,7 +506,17 @@ function _demangle_head(s::AbstractString) end # `arg_map` records the kernel's argument order (0=output, ≥1=input idx+1, -# <0=scalar idx); reconstruct it as a readable `output -> name(args...)` call. +# <0=scalar idx). Scalar slots are replaced with their values in the signature. +function _kernel_arg_name(a::Integer) + return if a == 0 + "output" + elseif a > 0 + "input{$(a - 1)}" + else + "scalar{$(-a - 1)}" + end +end + function _kernel_signature( dest::NDArray, unique_ndarrays::AbstractVector{<:NDArray}, @@ -501,18 +524,19 @@ function _kernel_signature( arg_map::AbstractVector{<:Integer}, func::AbstractString, ) - token(a::Integer) = - if a == 0 - "output" - elseif a > 0 - string("input", a - 1) - else - string("scalar", -a - 1) - end + token(a) = a < 0 ? repr(actual_scalars[-a]) : _kernel_arg_name(a) call_args = [token(a) for a in arg_map if a != 0] return string("broadcast.", _demangle_head(func), "(", join(call_args, ", "), ")") end +function _ndarray_debug_summary(nd::NDArray) + summary = "NDArray{$(eltype(nd)), $(ndims(nd))} $(size(nd))" + if _is_ndarray_slice(nd) + return "$summary slice, parent $(size(nd.parent))" + end + return summary +end + function _describe_fused_broadcast( dest, tree_str, unique_ndarrays, actual_scalars, static_args, arg_map, fkm, ndrange ) @@ -520,15 +544,14 @@ function _describe_fused_broadcast( field(k, v) = println(io, " ", rpad(k, 8), v) println(io, "\n", "="^40, " fused broadcast kernel") field("expr", tree_str) - field("output", "$(typeof(dest)) $(size(dest))") - field("inputs", "$(length(unique_ndarrays)) unique NDArray(s)") + field("output", _ndarray_debug_summary(dest)) + field("inputs", "input{N} ($(length(unique_ndarrays)) unique)") + println(io, " ", rpad("N", 4), "value") for (i, nd) in enumerate(unique_ndarrays) alias = objectid(nd) == objectid(dest) ? " (aliases output)" : "" - println(io, " [", i - 1, "] ", typeof(nd), " ", size(nd), alias) + println(io, " ", rpad(string(i - 1), 4), _ndarray_debug_summary(nd), alias) end - isempty(actual_scalars) || field("scalars", join(repr.(actual_scalars), ", ")) isempty(static_args) || field("static", join(repr.(static_args), ", ")) - field("arg_map", "$(Int.(arg_map)) (0=output, >=1=input idx+1, <0=scalar)") indexing = ndims(dest) in (2, 3) ? "cartesian" : "linear" field( "launch", @@ -609,7 +632,6 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas # Capture the readable tree before flatten collapses the nesting. bc_scope = bc - tree_str = BCAST_FUSION_DEBUG[] ? _bcast_tree_str(bc) : "" bc = Base.Broadcast.preprocess(dest, bc) bc = Base.Broadcast.instantiate(bc) @@ -666,11 +688,23 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas input_ndarrays = tuple(unique_ndarrays...) - BCAST_FUSION_DEBUG[] && _describe_fused_broadcast( - dest, tree_str, unique_ndarrays, actual_scalars, static_args, arg_map, fkm, ndrange - ) + if BCAST_FUSION_DEBUG[] + tree_str = _bcast_runtime_tree_str( + bc_scope, ndarray_to_input_idx, actual_scalars + ) + _describe_fused_broadcast( + dest, + tree_str, + unique_ndarrays, + actual_scalars, + static_args, + arg_map, + fkm, + ndrange, + ) + end - @task_scope _bcast_scope_name(bc_scope, ndarray_to_input_idx) begin + @task_scope _bcast_scope_name(bc_scope, ndarray_to_input_idx, actual_scalars) begin # `blocks=1` is a placeholder; RunPTXBroadcastTask overwrites grid dims # from the local PhysicalArray. `threads` is only the occupancy budget (tx). # Scalars after ctx: num_kernel_args, arg_map... diff --git a/src/scoping/broadcast_lifetimes.jl b/src/scoping/broadcast_lifetimes.jl index 62f2eb2e0..c5b50a027 100644 --- a/src/scoping/broadcast_lifetimes.jl +++ b/src/scoping/broadcast_lifetimes.jl @@ -4,18 +4,14 @@ # # hoists only materialized values while leaving the dotted tree intact: # -# tmp1 = @view C[2:end-1, :] +# tmp1 = C[2:end-1, :] # tmp2 = A[2:end-1, :] # tmp3 = B[2:end-1, :] # tmp1 .= tmp2 .* tmp3 .+ 2 # -# The destination view and input slices are objects that need lifetime management; +# The destination and input slices are objects that need lifetime management; # the `.*` and `.+` nodes are lazy and become one fused broadcast kernel. -function _expand_view(expr) - return Base.macroexpand(@__MODULE__, :(@view $expr)) -end - function rewrite_broadcast_lifetimes(scope) assigned_vars = Set{Symbol}() fresh_tmp(expr) = _hoist_temporary(expr, assigned_vars) @@ -64,6 +60,10 @@ function rewrite_broadcast_lifetimes(scope) return expr, Expr[] end + # Scalar arithmetic is evaluated while the broadcast tree is built; + # it does not create an NDArray whose lifetime needs to be tracked. + _is_scalar_expression(expr) && return expr, Expr[] + assignment = _assignment(expr) if !isnothing(assignment) (; lhs, rhs) = assignment @@ -78,13 +78,13 @@ function rewrite_broadcast_lifetimes(scope) broadcast_assignment = _broadcast_assignment(expr) if !isnothing(broadcast_assignment) (; lhs, rhs) = broadcast_assignment - # An explicit view preserves indexed `.=` semantics for both Array - # and NDArray while giving the lifetime pass a wrapper to destroy. + # NDArray slices are writable views. Hoist the destination slice so + # the fused broadcast writes through it, then destroy its handle. lhs_reference = _reference(lhs) if isnothing(lhs_reference) new_lhs, lhs_temps = rewrite_materialized(lhs) else - new_lhs, lhs_temps = fresh_tmp(_expand_view(lhs)) + new_lhs, lhs_temps = fresh_tmp(lhs) end new_rhs, rhs_temps = rewrite_lazy_broadcast(rhs, Dict{Any,Symbol}()) return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) diff --git a/src/scoping/lifetimes.jl b/src/scoping/lifetimes.jl index b80c577d1..b5b88469c 100644 --- a/src/scoping/lifetimes.jl +++ b/src/scoping/lifetimes.jl @@ -22,6 +22,9 @@ function rewrite_eager_lifetimes(scope) return expr, Expr[] end + # Scalar arithmetic cannot allocate an NDArray, so keep it inline. + _is_scalar_expression(expr) && return expr, Expr[] + assignment = _assignment(expr) if !isnothing(assignment) (; lhs, rhs) = assignment diff --git a/src/scoping/util.jl b/src/scoping/util.jl index 416e57757..64f68366e 100644 --- a/src/scoping/util.jl +++ b/src/scoping/util.jl @@ -13,9 +13,10 @@ using MacroTools: MacroTools # instead of repeating Expr head/argument indexing. export _assignment, _broadcast_assignment, _call, _dotcall, - _flatten_statements, _is_broadcast_op, _is_broadcast_syntax, _maphoist, - _reference, _prepend_statements, _replace_symbols, _rewrite_children, - _scope_statements, _strip_lines, walk_symbols + _flatten_statements, _is_broadcast_op, _is_broadcast_syntax, + _is_scalar_expression, _maphoist, _reference, _prepend_statements, + _replace_symbols, _rewrite_children, _scope_statements, _strip_lines, + walk_symbols function _assignment(expr) MacroTools.isexpr(expr, :(=)) || return nothing @@ -51,6 +52,25 @@ function _is_broadcast_syntax(expr) return !isnothing(_dotcall(expr)) end +const _SCALAR_ARITHMETIC = (:+, :-, :*, :/, :^, :%, :fld, :cld, :mod, :rem) + +function _is_property_access(expr) + return MacroTools.isexpr(expr, :.) && length(expr.args) == 2 && + expr.args[2] isa QuoteNode +end + +_is_scalar_expression(::Number) = true +_is_scalar_expression(expr::QuoteNode) = _is_scalar_expression(expr.value) +_is_scalar_expression(::Any) = false + +function _is_scalar_expression(expr::Expr) + _is_property_access(expr) && return true + + call = _call(expr) + return !isnothing(call) && call.f in _SCALAR_ARITHMETIC && + all(_is_scalar_expression, call.args) +end + function _reference(expr) MacroTools.isexpr(expr, :ref) || return nothing MacroTools.@capture(expr, array_[indices__]) || return nothing diff --git a/test/tests/broadcast_fusion_tests.jl b/test/tests/broadcast_fusion_tests.jl index 579a869f6..48c16cd9c 100644 --- a/test/tests/broadcast_fusion_tests.jl +++ b/test/tests/broadcast_fusion_tests.jl @@ -45,6 +45,29 @@ function test_broadcast_fusion(; T=Float32, N=100, atol=1e-5, rtol=1e-5) s2 = T(1.0) s3 = T(0.5) + @testset "Debug formatting" begin + input_indices = Dict(objectid(a) => 0, objectid(b) => 1) + tree = Base.broadcasted(+, Base.broadcasted(*, a, b), s1) + + @test cuNumeric._bcast_runtime_tree_str(tree, input_indices, Any[s1]) == + "+(*(input{0}, input{1}), $(repr(s1)))" + @test cuNumeric._kernel_signature( + a, + [a, b], + Any[s1], + Int32[0, 1, 2, -1], + "gpu_broadcast_kernel_linear_splat", + ) == + "broadcast.gpu_broadcast_kernel_linear_splat(input{0}, input{1}, $(repr(s1)))" + @test cuNumeric._ndarray_debug_summary(a) == + "NDArray{$T, 1} ($(size(a, 1)),)" + + slice = a[2:(end - 1)] + @test cuNumeric._ndarray_debug_summary(slice) == + "NDArray{$T, 1} ($(size(slice, 1)),) slice, parent ($(size(a, 1)),)" + cuNumeric.destroy!(slice) + end + # two different arrays @testset "A + B (two different arrays)" begin expected = julia_a .+ julia_b diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index 54e0c2931..c4fb9cb06 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -96,6 +96,9 @@ function test_scoping_rewrite_pipeline() @test utils._dotcall(:(f.(x, y))) == (f=:f, args=Any[:x, :y]) @test utils._reference(:(A[i, j])) == (array=:A, indices=Any[:i, :j]) @test utils._is_broadcast_syntax(:(x .* y)) + @test utils._is_scalar_expression(:(args.dx ^ 2)) + @test utils._is_scalar_expression(:(args.f + args.k)) + @test !utils._is_scalar_expression(:(A ^ 2)) @test utils._replace_symbols(:(x .+ y), Dict(:x => :(a .* b))) == :((a .* b) .+ y) @test isnothing( @@ -170,6 +173,24 @@ function test_scoping_rewrite_pipeline() end end + @testset "Scalar arithmetic stays inline" begin + source = quote + C .= A ./ args.dx^2 .+ (args.f + args.k) + end + + cuNumeric.counter[] = 0 + try + rewritten, assigned = cuNumeric.rewrite_broadcast_lifetimes(source) + rendered = sprint(Base.show_unquoted, utils._strip_lines(rewritten)) + + @test isempty(assigned) + @test occursin("args.dx ^ 2", rendered) + @test occursin("args.f + args.k", rendered) + finally + cuNumeric.counter[] = 0 + end + end + @testset "Eager lifetime rewrite" begin source = quote result = f(A[2:(end - 1), 2:(end - 1)]) @@ -290,16 +311,15 @@ function test_scoping_regressions(T, N) end if cuNumeric.FUSE_BROADCAST_EXPRS - @testset "Indexed fused assignment preserves Array view semantics" begin - src = reshape(T.(1:20), 4, 5) - out = fill(T(-1), 6, 7) + @testset "Indexed fused assignment writes through NDArray slices" begin + out = cuNumeric.zeros(T, (N + 2, N + 2)) @analyze_lifetimes begin - producer = src .* T(2) + producer = A .* T(2) out[2:(end - 1), 2:(end - 1)] = producer .+ T(1) end - @test out[2:(end - 1), 2:(end - 1)] == src .* T(2) .+ T(1) - @test all(out[[1, end], :] .== T(-1)) - @test all(out[:, [1, end]] .== T(-1)) + expected = zeros(T, N + 2, N + 2) + expected[2:(end - 1), 2:(end - 1)] .= T(3) + @test Array(out) == expected end end end From 4f195f76af48657731a2bb8d23da763218fc7863 Mon Sep 17 00:00:00 2001 From: krasow Date: Mon, 3 Aug 2026 19:09:12 -0500 Subject: [PATCH 08/12] hints for undef --- src/cuNumeric.jl | 2 ++ src/scoping/scoping.jl | 14 ++++++++++++++ test/tests/scoping.jl | 20 ++++++++++++++++++++ 3 files changed, 36 insertions(+) diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 9eb9d5b56..1c55b373d 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -252,6 +252,8 @@ function __init__() _is_precompiling() && return nothing + _register_scoping_error_hint!() + # Cannot set LEGATE_CONFIG on CI machines used # to register packages. So we will just skip starting # legate/cunumeric when using registry CI machines. diff --git a/src/scoping/scoping.jl b/src/scoping/scoping.jl index f729f538d..ae148f2ae 100644 --- a/src/scoping/scoping.jl +++ b/src/scoping/scoping.jl @@ -76,6 +76,20 @@ function _lexical_scope(body, bindings::Set{Symbol}) return Expr(:let, Expr(:block, ordered...), body) end +function _register_scoping_error_hint!() + isdefined(Base.Experimental, :register_error_hint) || return nothing + Base.Experimental.register_error_hint(UndefVarError) do io, exc + return print( + io, + "\nHint: bindings assigned inside `@analyze_lifetimes` are local to its " * + "block. If `", + exc.var, + "` was created there, return it from the block to use it afterward.", + ) + end + return nothing +end + function _hoist_temporary(expr, assigned_vars) counter[] += 1 temporary = Symbol(:tmp, counter[]) diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index c4fb9cb06..03a2d1a22 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -258,6 +258,19 @@ function test_scoping_rewrite_pipeline() return internal_first, internal_second end + function unrelated_undefined_binding() + return unrelated_result + end + + function rendered_error(f) + try + f() + catch exc + return sprint(io -> showerror(io, exc, catch_backtrace())) + end + return "" + end + output = [0] returned = @analyze_lifetimes begin internal_result = 42 @@ -267,6 +280,13 @@ function test_scoping_rewrite_pipeline() @test_throws UndefVarError hidden_binding() @test_throws UndefVarError hidden_destructured_bindings() + @test occursin( + "If `internal_result` was created there", rendered_error(hidden_binding) + ) + @test occursin( + "If `unrelated_result` was created there", + rendered_error(unrelated_undefined_binding), + ) @test shadowed_binding() == :outer @test output == [42] @test returned == 42 From 47dd2feec090b16f1be62915c996b1de93a6ae2f Mon Sep 17 00:00:00 2001 From: krasow Date: Mon, 3 Aug 2026 19:16:19 -0500 Subject: [PATCH 09/12] add documentation about our inter-broadcast fusion anti pattern regarding pre-allocations --- docs/src/perf/kernel_fusion.md | 41 ++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/docs/src/perf/kernel_fusion.md b/docs/src/perf/kernel_fusion.md index daed74284..04c437742 100644 --- a/docs/src/perf/kernel_fusion.md +++ b/docs/src/perf/kernel_fusion.md @@ -12,6 +12,47 @@ y .= .-a .+ b .* c y .= @. -a + b * c ``` +## Avoid preallocated intermediate broadcast buffers + +Preallocation is useful for a final output or a buffer that must persist across +iterations. It can be counterproductive for a single-use intermediate inside +`@analyze_lifetimes`, however. An in-place `.=` assignment is an observable +mutation, so inter-statement broadcast fusion treats it as a kernel boundary: + +```julia +tmp = cuNumeric.zeros(Float32, N, N) +result = cuNumeric.zeros(Float32, N, N) + +@analyze_lifetimes begin + tmp .= @. A + B + result .= @. tmp * C + 1.0f0 +end +``` + +This materializes `tmp` before the second expression and requires separate +kernel launches. Instead, use an ordinary assignment for a single-use +intermediate and keep `.=` for the final destination: + +```julia +result = cuNumeric.zeros(Float32, N, N) + +@analyze_lifetimes begin + tmp = @. A + B + result .= @. tmp * C + 1.0f0 +end +``` + +The inter-statement pass can substitute `tmp` into its only consumer, producing +the equivalent of `result .= @. (A + B) * C + 1.0f0`. The intermediate is never +materialized, so the full expression can run as one fused kernel. + +This rewrite is intentionally conservative: the intermediate must have one +use, no intervening statement may invalidate its inputs, and all normal fusion +requirements still apply. Keep preallocation when an intermediate is reused, +must preserve mutation semantics, or cannot be fused. Use +[`BCAST_FUSION_DEBUG`](../debugging.md#inspect-fused-broadcasts-with-bcast_fusion_debug) +to confirm whether the rewrite occurred. + Fusion applies when CUDA is available, the array leaves share the same shape, and the expression has at least `FUSE_BROADCAST_MIN_OPS` ops (default 2). Otherwise cuNumeric falls back to evaluating one op at a time. Shape-mismatched broadcasts such as `matrix .+ vector` use the unfused path. Toggle fusion through `CNPreferences` (restart Julia after changing these): From 3290ac9a65a78b4fb1118b435375324eee1255b5 Mon Sep 17 00:00:00 2001 From: David Krasowska Date: Mon, 3 Aug 2026 20:31:45 -0500 Subject: [PATCH 10/12] revert PTX version from 7.8 to 9.0 --- src/ndarray/broadcast_fusion.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index cc6c634ec..8867059d8 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -401,7 +401,7 @@ function get_ptx( # dump_module=true keeps only_entry=false so linked libdevice helpers (e.g. cos # slowpaths) are not emptied into unresolved .externs before cuModuleLoad. CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, arg_types...); - raw=false, dump_module=true, kernel=true, ptx=v"7.8") + raw=false, dump_module=true, kernel=true, ptx=v"9.0") return String(take!(buf)), threads, ctx end From b2e3cae025df8dfc7f4ecdd7110a0bb06172a7d4 Mon Sep 17 00:00:00 2001 From: krasow Date: Mon, 3 Aug 2026 21:30:44 -0500 Subject: [PATCH 11/12] handle dot macro expansion properly --- src/scoping/scoping.jl | 28 ++++++++++++++++++++++++++++ test/tests/scoping.jl | 22 ++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/scoping/scoping.jl b/src/scoping/scoping.jl index ae148f2ae..41e40d29b 100644 --- a/src/scoping/scoping.jl +++ b/src/scoping/scoping.jl @@ -6,6 +6,32 @@ using .ScopingUtils include("inter_broadcast_fusion.jl") +const _DOT_MACRO_NAME = Symbol("@__dot__") + +_is_dot_macro(::Any) = false +_is_dot_macro(name::Symbol) = name === _DOT_MACRO_NAME +_is_dot_macro(name::GlobalRef) = name.name === _DOT_MACRO_NAME + +function _is_dot_macro(name::Expr) + name.head === :. || return false + quoted_name = last(name.args) + return quoted_name isa QuoteNode && quoted_name.value === _DOT_MACRO_NAME +end + +_expand_dot_macros(value, ::Module) = value + +function _expand_dot_macros(expr::Expr, caller::Module) + expr.head in (:quote, :inert) && return expr + + if expr.head === :macrocall && _is_dot_macro(first(expr.args)) + expanded = Base.macroexpand(caller, expr; recursive=false) + return _expand_dot_macros(expanded, caller) + end + + args = map(arg -> _expand_dot_macros(arg, caller), expr.args) + return Expr(expr.head, args...) +end + @doc""" @analyze_lifetimes expr @@ -25,6 +51,7 @@ are not individually hoisted. The macro automatically selects the broadcast-aware analysis in that case and the plain analysis otherwise. """ macro analyze_lifetimes(block) + block = _expand_dot_macros(block, __module__) on_rewrite = BCAST_FUSION_DEBUG[] ? InterBroadcastFusion.log_rewrite : nothing rewritten = process_ndarray_scope(block; on_rewrite) bindings = union(_assigned_symbols(block), _assigned_symbols(rewritten)) @@ -314,5 +341,6 @@ you can see exactly where each temporary is freed. Pure AST work, so it runs on CPU-only checkouts. """ macro show_lifetimes(block) + block = _expand_dot_macros(block, __module__) return :(print_lifetime_analysis($(QuoteNode(block)))) end diff --git a/test/tests/scoping.jl b/test/tests/scoping.jl index 03a2d1a22..96ba5136e 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -140,6 +140,18 @@ function test_scoping_rewrite_pipeline() ) @test isempty(events) @test utils._strip_lines(rewritten) == utils._strip_lines(untouched) + + dotted = quote + tmp = @. A + B + result .= @. tmp * C + Float64(1.0) + end + expanded = cuNumeric._expand_dot_macros(dotted, @__MODULE__) + @test !occursin("@__dot__", sprint(Base.show_unquoted, expanded)) + + rewritten = cuNumeric.InterBroadcastFusion.rewrite_scope(expanded) + rendered = sprint(Base.show_unquoted, utils._strip_lines(rewritten)) + @test !occursin("tmp =", rendered) + @test occursin("Float64.(1.0)", rendered) end @testset "Fusion-aware lifetime rewrite" begin @@ -341,6 +353,16 @@ function test_scoping_regressions(T, N) expected[2:(end - 1), 2:(end - 1)] .= T(3) @test Array(out) == expected end + + @testset "Nested @. macros fuse before lifetime analysis" begin + multiplier = cuNumeric.ones(T, (N, N)) + result = cuNumeric.zeros(T, (N, N)) + @analyze_lifetimes begin + tmp = @. A + B + result .= @. tmp * multiplier + T(1.0) + end + @test Array(result) == fill(T(3), N, N) + end end end From c4b4b14ba2eff67140c8211d5bfdebd9a3e78fd9 Mon Sep 17 00:00:00 2001 From: krasow Date: Mon, 3 Aug 2026 22:16:52 -0500 Subject: [PATCH 12/12] materialize scalar broadcasts to fix test case: @analyze_lifetimes begin tmp = @. A + B result .= @. tmp * multiplier + T(1.0) end --- src/ndarray/broadcast_fusion.jl | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 8867059d8..6ee83e225 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -284,6 +284,22 @@ end @inline _unwrap_fusion_arg(x) = x @inline _unwrap_fusion_args(args::Tuple) = _unwrap_fusion_arg.(args) +# `Broadcast.flatten` discards the result types of nested scalar-only nodes. +# Materialize those nodes first so an explicit conversion such as +# `Float32.(1.0)` reaches runtime-argument alignment as a `Float32` scalar. +@inline _fold_fused_scalar_broadcasts(x) = x +@inline function _fold_fused_scalar_broadcasts( + bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.DefaultArrayStyle{0}} +) + args = map(_fold_fused_scalar_broadcasts, bc.args) + scalar_bc = Base.Broadcast.Broadcasted(bc.style, bc.f, args, bc.axes) + return Base.Broadcast.materialize(scalar_bc) +end +@inline function _fold_fused_scalar_broadcasts(bc::Base.Broadcast.Broadcasted) + args = map(_fold_fused_scalar_broadcasts, bc.args) + return Base.Broadcast.Broadcasted(bc.style, bc.f, args, bc.axes) +end + # Host-side scalar alignment for fusion — same as unfused # `T_IN = __my_promote_type(...); unchecked_promote_arr.(args, T_IN)`, but # `unchecked_promote_scalar` keeps Numbers as scalars for the PTX arg buffer. @@ -630,6 +646,10 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas # Promotion checks use the pre-flatten tree (same shape as unfused unravel). _assert_fused_broadcast_promotion(dest, bc) + # Preserve the values and types produced by scalar-only subtrees before + # flattening turns their inputs into top-level runtime arguments. + bc = _fold_fused_scalar_broadcasts(bc) + # Capture the readable tree before flatten collapses the nesting. bc_scope = bc