Inter-broadcast expr kernel fusion and MacroTools refactor

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!))
 ```

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.
 

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
 

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):

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.