Refactor lifetime analysis w/ MacroTools.

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.
 

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.

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<std::complex<double>>("ComplexF64");
-  mod.map_type<std::complex<float>>("ComplexF32");
   mod.map_type<HalfType>("Float16");
 
   // These are the types/dims used to generate templated functions

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