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/debugging.md b/docs/src/debugging.md index 8686e304b..6b25941d5 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 @@ -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`. @@ -74,44 +75,71 @@ 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. 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 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) -C .= @. A * B + 2.0f0 +@analyze_lifetimes begin + 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 ``` -Example output: +For example, a single-use producer inside `@analyze_lifetimes` is reported as: + +```text +======================================== inter-broadcast fusion rewrite + before + begin + 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[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` +contains their replacement. No rewrite block is printed when the pass leaves +the statements unchanged. + +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`). -- If nothing prints, the expression took the unfused path (for example shape-mismatched leaves, fusion disabled, or below the min-ops threshold). +- `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/docs/src/internals.md b/docs/src/internals.md index 10f655faf..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,12 +36,32 @@ 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. - 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/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): diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index eb2d0f6b3..1c55b373d 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 @@ -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/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 2962e5a7d..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. @@ -456,18 +472,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 +494,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 +522,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 +540,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 +560,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", @@ -607,9 +646,12 @@ 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 - tree_str = BCAST_FUSION_DEBUG[] ? _bcast_tree_str(bc) : "" bc = Base.Broadcast.preprocess(dest, bc) bc = Base.Broadcast.instantiate(bc) @@ -666,11 +708,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.jl b/src/scoping.jl deleted file mode 100644 index 47ae01cea..000000000 --- a/src/scoping.jl +++ /dev/null @@ -1,469 +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 == :(=) && !(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 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). - function rewrite_bcast(e)::Tuple{Any,Vector{Expr}} - e isa Expr || return e, Expr[] - e.head == :ref && return fresh_tmp(e) - if e.head == :call && is_broadcast_op(e.args[1]) - args, hoisted = maphoist(rewrite_bcast, 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) - 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 - new_lhs, lts = rewrite(lhs) - new_rhs, rts = rewrite_bcast(rhs) - 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) - 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}() - rewritten = find_broadcast_assignments(block, assigned_vars) - result = insert_finalizers(rewritten, assigned_vars) - counter[] = 0 - return result -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 _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..c5b50a027 --- /dev/null +++ b/src/scoping/broadcast_lifetimes.jl @@ -0,0 +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 = C[2:end-1, :] +# tmp2 = A[2:end-1, :] +# tmp3 = B[2:end-1, :] +# tmp1 .= tmp2 .* tmp3 .+ 2 +# +# The destination and input slices are objects that need lifetime management; +# the `.*` and `.+` nodes are lazy and become one fused broadcast kernel. + +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_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_lazy_broadcast( + expr, slice_cache::Dict{Any,Symbol} + )::Tuple{Any,Vector{Expr}} + if !(expr isa Expr) + return expr, Expr[] + end + reference = _reference(expr) + if !isnothing(reference) + cached = get(slice_cache, expr, nothing) + if !isnothing(cached) + return cached, Expr[] + end + tmp, bind = fresh_tmp(expr) + slice_cache[expr] = tmp + return tmp, bind + 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, call.f, args...), hoisted + end + + dotcall = _dotcall(expr) + if !isnothing(dotcall) + args, hoisted = _maphoist( + arg -> rewrite_lazy_broadcast(arg, slice_cache), dotcall.args + ) + return Expr(:., dotcall.f, Expr(:tuple, args...)), hoisted + end + return rewrite_materialized(expr) + end + + function rewrite_materialized(expr)::Tuple{Any,Vector{Expr}} + if !(expr isa Expr) + 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 + 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. + broadcast_assignment = _broadcast_assignment(expr) + if !isnothing(broadcast_assignment) + (; lhs, rhs) = broadcast_assignment + # 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(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 + + reference = _reference(expr) + if !isnothing(reference) + return fresh_tmp(expr) + end + + 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 !isnothing(call) + args, hoisted = _maphoist(rewrite_materialized, call.args) + tmp, bind = fresh_tmp(Expr(:call, call.f, args...)) + return tmp, vcat(hoisted, bind) + end + + return _rewrite_children(rewrite_materialized, expr) + end + + rewritten, temps = rewrite_materialized(scope) + return _prepend_statements(rewritten, temps), assigned_vars +end + +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 new file mode 100644 index 000000000..0138c351c --- /dev/null +++ b/src/scoping/inter_broadcast_fusion.jl @@ -0,0 +1,189 @@ +module InterBroadcastFusion + +export rewrite_scope + +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}) + 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) + 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( + 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] + 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) + 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) + 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) + 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,Any}}() + for (sym, (def_idx, rhs)) in definitions + use_idx = _single_use_index(stmts, sym, def_idx) + isnothing(use_idx) && continue + dependencies = Set(walk_symbols(rhs)) + if !_safe_to_delay_broadcast(stmts, def_idx, use_idx, dependencies, lazy_defs) + continue + end + inlineable[sym] = (def_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] + 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][2], replacements) + continue + end + + source_indices = _source_indices(original_stmt, replacement_sources) + stmt = _substitute_symbols(original_stmt, replacements) + + if !isempty(source_indices) + stmt = _fuse_into_destination(stmt) + before = Expr( + :block, + (stmts[source_idx] for source_idx in source_indices)..., + original_stmt, + ) + push!(fusion_events, (; before, fused=stmt)) + end + push!(rewritten, stmt) + end + + return Expr(scope.head, rewritten...), fusion_events +end + +""" + rewrite_scope(scope; on_rewrite=nothing) -> scope + +Fuse eligible single-use broadcast producers into their consumer and return the +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; on_rewrite=nothing) + rewritten, fusion_events = _rewrite_scope(scope) + if !isnothing(on_rewrite) + for event in fusion_events + on_rewrite(event) + end + end + return rewritten +end + +function _print_expr(io::IO, expr) + clean = _strip_lines(expr) + rendered = sprint(Base.show_unquoted, clean) + for line in eachline(IOBuffer(rendered)) + println(io, " ", line) + end + return nothing +end + +function log_rewrite(event; io::IO=stdout) + println(io, "\n", "="^40, " inter-broadcast fusion rewrite") + println(io, " before") + _print_expr(io, event.before) + println(io, " fused") + _print_expr(io, event.fused) + return nothing +end + +end diff --git a/src/scoping/lifetimes.jl b/src/scoping/lifetimes.jl new file mode 100644 index 000000000..b5b88469c --- /dev/null +++ b/src/scoping/lifetimes.jl @@ -0,0 +1,75 @@ +# 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 rewrite_eager_lifetimes(scope) + assigned_vars = Set{Symbol}() + fresh_tmp(expr) = _hoist_temporary(expr, assigned_vars) + + function rewrite(expr)::Tuple{Any,Vector{Expr}} + if !(expr isa Expr) + 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 + if lhs isa Symbol + push!(assigned_vars, lhs) + end + new_rhs, temps = rewrite(rhs) + return :($lhs = $new_rhs), temps + end + + 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. + 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 + + new_rhs, rhs_temps = rewrite(rhs) + return Expr(:(.=), new_lhs, new_rhs), vcat(lhs_temps, rhs_temps) + end + + reference = _reference(expr) + if !isnothing(reference) + return fresh_tmp(expr) + end + + 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 + + return _rewrite_children(rewrite, expr) + end + + rewritten, temps = rewrite(scope) + return _prepend_statements(rewritten, temps), assigned_vars +end + +function process_lifetime_scope(scope) + return _process_lifetime_scope(scope, rewrite_eager_lifetimes) +end diff --git a/src/scoping/scoping.jl b/src/scoping/scoping.jl new file mode 100644 index 000000000..41e40d29b --- /dev/null +++ b/src/scoping/scoping.jl @@ -0,0 +1,346 @@ +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") + +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 + +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. + +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 +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)) + return esc(_lexical_scope(rewritten, bindings)) +end + +const counter = Ref(0) + +function maybe_insert_delete(var::NDArray) + return cuNumeric.destroy!(var) +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 _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[]) + 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}) + last_use = Dict{Symbol,Int}() + alias_map = Dict{Symbol,Symbol}() + + stmts = _flatten_statements(Expr(:block, exprs...)) + + # Pass 1: collect definitions and uses + for (i, stmt) in enumerate(stmts) + stmt isa Expr || continue + assignment = _assignment(stmt) + used_expr = stmt + if !isnothing(assignment) + (; lhs, rhs) = assignment + if lhs isa Symbol && rhs isa Symbol + alias_map[lhs] = rhs + end + used_expr = rhs + end + for symbol in walk_symbols(used_expr) + last_use[symbol] = i + end + end + + for (alias, src) in alias_map + 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 + # 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 + + 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_symbols(stmt) + stmt isa Symbol && return Set([stmt]) + + assignment = _assignment(stmt) + 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 + 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 + 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 in protected + return nothing + end + push!(freed, c) + 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. + aliased_source = nothing + assignment = _assignment(stmt) + if !isnothing(assignment) + (; lhs, rhs) = assignment + if lhs isa Symbol && rhs isa Symbol + aliased_source = 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 in assigned_vars && v != aliased_source + emit_delete!(v) + end + end + + if i == n + push!(out, terminal_indexed ? :nothing : res_var) + end + 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}) + 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("broadcast_lifetimes.jl") + +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; on_rewrite) + end + return process_lifetime_scope(scope) +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 + return only(call.args) +end + +function print_lifetime_analysis(block; io::IO=stdout) + rule = "-"^60 + 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 + 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) + 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) + block = _expand_dot_macros(block, __module__) + return :(print_lifetime_analysis($(QuoteNode(block)))) +end diff --git a/src/scoping/util.jl b/src/scoping/util.jl new file mode 100644 index 000000000..64f68366e --- /dev/null +++ b/src/scoping/util.jl @@ -0,0 +1,158 @@ +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, + _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 + 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 + +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 + 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 da238d01d..261b07f2e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -458,6 +458,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/broadcast_fusion_tests.jl b/test/tests/broadcast_fusion_tests.jl index 4b38fb894..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 @@ -529,6 +552,20 @@ function test_broadcast_fusion_edge_cases(; T=Float32, atol=1e-5, rtol=1e-5) a2d .= a2d .* s1 .+ b2d @allowscalar @test cuNumeric.compare(j2a .* s1 .+ j2b, a2d, 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..96ba5136e 100644 --- a/test/tests/scoping.jl +++ b/test/tests/scoping.jl @@ -86,6 +86,241 @@ 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._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( + 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) + + 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 + 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 "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)]) + 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 + + @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 + + 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 + output[1] = internal_result + internal_result + end + + @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 + + 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) A = cuNumeric.ones(T, (N, N)) B = cuNumeric.ones(T, (N, N)) @@ -106,6 +341,29 @@ 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 writes through NDArray slices" begin + out = cuNumeric.zeros(T, (N + 2, N + 2)) + @analyze_lifetimes begin + producer = A .* T(2) + out[2:(end - 1), 2:(end - 1)] = producer .+ T(1) + end + expected = zeros(T, N + 2, N + 2) + 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 function run_all_ops(FT, N)