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)