diff --git a/NEWS.md b/NEWS.md index c65ceb45e22c7..cbf5ecb339671 100644 --- a/NEWS.md +++ b/NEWS.md @@ -18,6 +18,9 @@ Language changes Compiler/Runtime improvements ----------------------------- + - Captured variables that are assigned in all branches of an `if`/`elseif`/`else` statement + no longer allocate a `Core.Box`, reducing heap allocations in closures ([#60542]). + Command-line option changes --------------------------- diff --git a/src/julia-syntax.scm b/src/julia-syntax.scm index c9b1fd551e359..d092e0fbb8a74 100644 --- a/src/julia-syntax.scm +++ b/src/julia-syntax.scm @@ -3915,6 +3915,8 @@ f(x) = yt(x) ;; Try to identify never-undef variables, and then clear the `captured` flag for single-assigned, ;; never-undef variables to avoid allocating unnecessary `Box`es. +;; Also handles captured variables assigned in all branches of if-else, treating the if-else +;; as a single definition point when there are no other assignments and capture is after. (define (lambda-optimize-vars! lam) (assert (eq? (car lam) 'lambda)) ;; memoize all-methods-for to avoid O(n^2) behavior @@ -3933,12 +3935,23 @@ f(x) = yt(x) (decl (table)) (unused (table)) ;; variables not (yet) used (read from) in the current block (live (table)) ;; variables that have been set in the current block - (seen (table))) ;; all variables we've seen assignments to - ;; Collect candidate variables: those that are captured (and hence we want to optimize) - ;; and only assigned once. This populates the initial `unused` table. + (seen (table)) ;; all variables we've seen assignments to + ;; if-else optimization for captured multi-assigned variables: + (ifa-candidates (table)) ;; captured multi-assigned vars (potential optimization) + (assigned-outside (table)) ;; vars assigned outside any qualifying if-else + (captured-early (table)) ;; vars captured before their if-else completes + (ifa-completed (table)) ;; vars that completed a valid if-else (all branches assigned) + (in-ifa-for (table)) ;; during if-else visit: which vars we're tracking + (used-before-assign (table)) ;; vars used before assigned in some branch + (has-candidates #f)) ;; whether we have any ifa candidates + ;; Collect candidate variables: those that are captured and single-assigned go into `unused`. + ;; Captured multi-assigned vars go into ifa-candidates for if-else optimization. (for-each (lambda (v) - (if (and (vinfo:capt v) (vinfo:sa v)) - (put! unused (car v) #t))) + (if (vinfo:capt v) + (if (vinfo:sa v) + (put! unused (car v) #t) + (begin (put! ifa-candidates (car v) #t) + (set! has-candidates #t))))) vi) (define (restore old) (table.foreach (lambda (k v) @@ -3956,15 +3969,26 @@ f(x) = yt(x) (if (and (has? unused var) (not (memq var args))) (del! unused var))) (define (mark-captured var) - (if (has? unused var) - (del! unused var))) + (if (and (has? unused var) (not (has? ifa-completed var))) + (del! unused var)) + ;; For ifa candidates: if captured before if-else completes, can't optimize + (if (and (has? ifa-candidates var) + (not (has? ifa-completed var))) + (put! captured-early var #t))) (define (assign! var) (if (has? unused var) ;; When a variable is assigned, move it to the live set to protect ;; it from being removed from `unused`. (begin (put! live var #t) (put! seen var #t) - (del! unused var)))) + (del! unused var))) + ;; Track assignments for ifa candidates + (if (has? ifa-candidates var) + (begin + (put! seen var #t) + ;; If not currently inside an if-else for this var, it's assigned outside + (if (not (has? in-ifa-for var)) + (put! assigned-outside var #t))))) (define (declare! var) (if (has? unused var) (put! decl var #t))) @@ -3976,6 +4000,19 @@ f(x) = yt(x) (del! live k))) (table.keys live)) (set! decl old-decls)) + ;; Returns table of vars assigned in ALL branches + (define (intersect-branch-tables branch-tables) + (if (null? branch-tables) + (table) + (let ((result (table.clone (car branch-tables)))) + (for-each (lambda (branch-tbl) + (table.foreach + (lambda (var _) + (if (not (has? branch-tbl var)) + (del! result var))) + result)) + (cdr branch-tables)) + result))) (define (visit e) ;; returns whether e contained a symboliclabel (cond ((atom? e) (if (symbol? e) (mark-used e)) @@ -3997,7 +4034,125 @@ f(x) = yt(x) ((eq? (car e) 'symboliclabel) (kill) #t) - ((memq (car e) '(if elseif trycatch tryfinally trycatchelse)) + ((eq? (car e) 'if) + (let ((prev (table.clone live))) + (cond + ;; if-else with exactly 3 args and we have candidates + ((and (length= e 4) has-candidates) + (let ((has-label #f) + (branch-tables '()) + (prev-in-ifa (table.clone in-ifa-for)) + (branch-assigned (table))) ;; track assignments within current branch + ;; Mark eligible candidates as being inside this if-else + (table.foreach + (lambda (var _) + (if (and (not (has? ifa-completed var)) + (not (has? assigned-outside var))) + (put! in-ifa-for var #t))) + ifa-candidates) + ;; Visit condition + (if (visit (cadr e)) (set! has-label #t)) + (let ((pre-branch-live (table.clone live))) + ;; Collect assignments from each branch + (let collect-branches ((branches (list (caddr e) (cadddr e))) + (branch-assigns '())) + (if (null? branches) + (set! branch-tables (reverse branch-assigns)) + (let ((branch (car branches))) + (set! live (table.clone pre-branch-live)) + (kill) + (set! branch-assigned (table)) + (let ((branch-tbl (table))) + ;; Visit branch, tracking assignments and use-before-assign + (let visit-branch ((expr branch)) + (cond + ((atom? expr) + (if (symbol? expr) + (begin + (mark-used expr) + ;; Track use-before-assign for ifa candidates + (if (and (has? in-ifa-for expr) + (not (has? branch-assigned expr))) + (put! used-before-assign expr #t))))) + ((lambda-opt-ignored-exprs (car expr)) #f) + ((eq? (car expr) 'symboliclabel) + (set! has-label #t)) + ((eq? (car expr) '=) + (visit-branch (caddr expr)) + (let ((var (cadr expr))) + (assign! var) + (if (has? in-ifa-for var) + (begin + (put! branch-tbl var #t) + (put! branch-assigned var #t))))) + ;; Handle loops inside branches - mark vars as assigned-outside + ((or (eq? (car expr) '_while) (eq? (car expr) '_do_while)) + (let ((prev-in-ifa-loop (table.clone in-ifa-for))) + ;; Clear in-ifa-for during loop + (set! in-ifa-for (table)) + (for-each visit-branch (cdr expr)) + (set! in-ifa-for prev-in-ifa-loop))) + ;; Handle elseif chain + ((and (eq? (car expr) 'elseif) (length= expr 4)) + (visit-branch (cadr expr)) + (kill) + (set! branch-assigned (table)) + (let ((elseif-branch-tbl (table))) + (let visit-elseif-then ((e2 (caddr expr))) + (cond + ((atom? e2) + (if (symbol? e2) + (begin + (mark-used e2) + (if (and (has? in-ifa-for e2) + (not (has? branch-assigned e2))) + (put! used-before-assign e2 #t))))) + ((lambda-opt-ignored-exprs (car e2)) #f) + ((eq? (car e2) 'symboliclabel) + (set! has-label #t)) + ((eq? (car e2) '=) + (visit-elseif-then (caddr e2)) + (let ((var (cadr e2))) + (assign! var) + (if (has? in-ifa-for var) + (begin + (put! elseif-branch-tbl var #t) + (put! branch-assigned var #t))))) + ((or (eq? (car e2) '_while) (eq? (car e2) '_do_while)) + (let ((prev-in-ifa-loop (table.clone in-ifa-for))) + (set! in-ifa-for (table)) + (for-each visit-elseif-then (cdr e2)) + (set! in-ifa-for prev-in-ifa-loop))) + (else + (for-each visit-elseif-then (cdr e2))))) + (set! branch-assigns (cons elseif-branch-tbl branch-assigns))) + (set! live (table.clone pre-branch-live)) + (collect-branches (list (cadddr expr)) branch-assigns)) + (else + (for-each visit-branch (cdr expr))))) + (if (not (and (pair? branch) (eq? (car branch) 'elseif))) + (collect-branches (cdr branches) + (cons branch-tbl branch-assigns))))))) + ;; Mark vars assigned in ALL branches as completed (if not used-before-assign) + (let ((all-assigned (intersect-branch-tables branch-tables))) + (table.foreach + (lambda (var _) + (if (and (has? in-ifa-for var) + (not (has? captured-early var)) + (not (has? used-before-assign var))) + (put! ifa-completed var #t))) + all-assigned)) + (set! in-ifa-for prev-in-ifa) + (kill) + (if has-label + (begin (kill) #t) + (begin (restore prev) #f))))) + ;; Default if handling + (else + (if (eager-any (lambda (e) (begin0 (visit e) (kill))) (cdr e)) + (begin (kill) #t) + (begin (restore prev) #f)))))) + ((memq (car e) '(elseif trycatch tryfinally trycatchelse)) (let ((prev (table.clone live))) (if (eager-any (lambda (e) (begin0 (visit e) (kill))) @@ -4008,8 +4163,12 @@ f(x) = yt(x) (begin (restore prev) #f)))) ((or (eq? (car e) '_while) (eq? (car e) '_do_while)) (let ((prev (table.clone live)) - (decl- (table.clone decl))) + (decl- (table.clone decl)) + (prev-in-ifa (table.clone in-ifa-for))) + ;; Clear in-ifa-for during loop - assignments in loop are "outside" + (set! in-ifa-for (table)) (let ((result (eager-any visit (cdr e)))) + (set! in-ifa-for prev-in-ifa) (leave-loop! decl-) (if result #t @@ -4048,10 +4207,22 @@ f(x) = yt(x) (let ((vv (assq v vi))) (vinfo:set-never-undef! vv #t)))) (append (table.keys live) (table.keys unused))) + ;; Clear captured flag for single-assigned never-undef variables (for-each (lambda (v) (if (and (vinfo:sa v) (vinfo:never-undef v)) (set-car! (cddr v) (logand (caddr v) (lognot 5))))) vi) + ;; Also clear captured flag for ifa-completed vars not assigned outside or captured early + (for-each (lambda (var) + (if (and (has? ifa-completed var) + (not (has? assigned-outside var)) + (not (has? captured-early var))) + (let ((vv (assq var vi))) + (if vv + (begin + (vinfo:set-never-undef! vv #t) + (set-car! (cddr vv) (logand (caddr vv) (lognot 5)))))))) + (table.keys ifa-candidates)) lam)) (define (is-var-boxed? v lam) diff --git a/test/syntax.jl b/test/syntax.jl index 5691ca7135993..6ec7c42224cd4 100644 --- a/test/syntax.jl +++ b/test/syntax.jl @@ -4673,3 +4673,270 @@ module M59755 end @test M59755.v6 === 5 @test Base.binding_kind(M59755, :v6) == Base.PARTITION_KIND_CONST end + +# Test that `if cond; x = a; else x = b; end` doesn't introduce a Core.Box +# when x is captured, due to optimization in lambda-optimize-vars! +@testset "if-else assignment without boxing" begin + + # Tests where optimization SHOULD apply (variable not boxed) + @testset "optimized cases" begin + # Basic case: assign in both branches, capture after + function ifa_basic(cond, a, b) + if cond + x = a + else + x = b + end + return () -> x + end + closure_type = typeof(ifa_basic(true, 1, 2)) + @test fieldtype(closure_type, 1) !== Core.Box + @test ifa_basic(true, 1, 2)() == 1 + @test ifa_basic(false, 1, 2)() == 2 + + # Multiple variables assigned in both branches + function ifa_multi_vars(cond, a, b, c, d) + if cond + @noinline identity(1) + x = a + y = b + else + @noinline identity(2) + x = c + y = d + end + return () -> (x, y) + end + closure_type2 = typeof(ifa_multi_vars(true, 1, 2, 3, 4)) + @test fieldtype(closure_type2, 1) !== Core.Box + @test fieldtype(closure_type2, 2) !== Core.Box + @test ifa_multi_vars(true, 1, 2, 3, 4)() == (1, 2) + @test ifa_multi_vars(false, 1, 2, 3, 4)() == (3, 4) + + # Elseif chain: assign in all branches + function ifa_elseif(cond1, cond2, a, b, c) + if cond1 + x = a + elseif cond2 + x = b + else + x = c + end + return () -> x + end + closure_type3 = typeof(ifa_elseif(true, false, 1, 2, 3)) + @test fieldtype(closure_type3, 1) !== Core.Box + @test ifa_elseif(true, false, 1, 2, 3)() == 1 + @test ifa_elseif(false, true, 1, 2, 3)() == 2 + @test ifa_elseif(false, false, 1, 2, 3)() == 3 + + # Nested if-else: all leaf paths assign + function ifa_nested(c1, c2, a, b, c, d) + if c1 + if c2 + x = a + else + x = b + end + else + if c2 + x = c + else + x = d + end + end + return () -> x + end + closure_nested = typeof(ifa_nested(true, true, 1, 2, 3, 4)) + @test fieldtype(closure_nested, 1) !== Core.Box + @test ifa_nested(true, true, 1, 2, 3, 4)() == 1 + @test ifa_nested(true, false, 1, 2, 3, 4)() == 2 + @test ifa_nested(false, true, 1, 2, 3, 4)() == 3 + @test ifa_nested(false, false, 1, 2, 3, 4)() == 4 + + # Deeply nested if-else (3 levels): all 8 leaf paths assign + function ifa_deeply_nested(c1, c2, c3, vals...) + if c1 + if c2 + if c3 + x = vals[1] + else + x = vals[2] + end + else + if c3 + x = vals[3] + else + x = vals[4] + end + end + else + if c2 + if c3 + x = vals[5] + else + x = vals[6] + end + else + if c3 + x = vals[7] + else + x = vals[8] + end + end + end + return () -> x + end + closure_deep = typeof(ifa_deeply_nested(true, true, true, 1, 2, 3, 4, 5, 6, 7, 8)) + @test fieldtype(closure_deep, 1) !== Core.Box + @test ifa_deeply_nested(true, true, true, 1, 2, 3, 4, 5, 6, 7, 8)() == 1 + @test ifa_deeply_nested(false, false, false, 1, 2, 3, 4, 5, 6, 7, 8)() == 8 + end + + # Tests where optimization must NOT apply (variable must remain boxed) + @testset "boxed cases - incomplete branch coverage" begin + # Missing else branch (not all paths assign) + function box_missing_else(cond, a) + x = 0 + if cond + x = a + end + return () -> x + end + @test fieldtype(typeof(box_missing_else(true, 1)), 1) === Core.Box + end + + @testset "boxed cases - use/capture timing" begin + # Use before assignment in one branch + function box_use_before_assign(cond, a, b) + if cond + x = a + else + @noinline println(devnull, x) + x = b + end + return () -> x + end + @test fieldtype(typeof(box_use_before_assign(true, 1, 2)), 1) === Core.Box + + # Capture before any assignment + function box_capture_before_any_assign(cond, a, b) + f = () -> x + if cond + x = a + else + x = b + end + return f + end + @test fieldtype(typeof(box_capture_before_any_assign(true, 1, 2)), 1) === Core.Box + + # Assignment before if-else, capture before if-else + function box_assign_and_capture_before_ifelse() + x = 1 + g = () -> x + if rand() > 0.5 + x = 2 + else + x = 3 + end + return g + end + @test fieldtype(typeof(box_assign_and_capture_before_ifelse()), 1) === Core.Box + + # Capture inside branch, then assign again in same branch + function box_capture_inside_branch_then_assign() + x = 0 + if true + x = 1 + g = () -> x + x = 2 + else + x = 3 + end + return g + end + @test fieldtype(typeof(box_capture_inside_branch_then_assign()), 1) === Core.Box + end + + @testset "boxed cases - control flow" begin + # @goto can skip assignments + function box_goto_skips_assign(cond, a, b) + if cond + @goto skip + x = a + else + x = b + end + @label skip + x = 0 + return () -> x + end + @test fieldtype(typeof(box_goto_skips_assign(true, 1, 2)), 1) === Core.Box + end + + @testset "boxed cases - loops" begin + # For loop inside branch causes multiple assignments + function box_for_loop_in_branch(cond, n, a, b) + if cond + x = a + for i in 1:n + x = i + end + else + x = b + end + return () -> x + end + @test fieldtype(typeof(box_for_loop_in_branch(true, 3, 1, 2)), 1) === Core.Box + + # While loop inside branch causes multiple assignments + function box_while_loop_in_branch(cond, a, b) + if cond + x = a + i = 0 + while i < 2 + x = i + i += 1 + end + else + x = b + end + return () -> x + end + @test fieldtype(typeof(box_while_loop_in_branch(true, 1, 2)), 1) === Core.Box + + # If-else inside a loop (loop re-executes assignments) + function box_ifelse_inside_loop(n, a, b) + x = 0 + for i in 1:n + if i > n รท 2 + x = a + else + x = b + end + end + return () -> x + end + @test fieldtype(typeof(box_ifelse_inside_loop(3, 1, 2)), 1) === Core.Box + end + + @testset "boxed cases - multiple assignments" begin + # Assignment after capture (second if-else after closure) + function box_assign_after_capture(cond) + if cond + x = 1 + else + x = 2 + end + g = function () + x + end + if true + x = 123 + end + return g + end + @test fieldtype(typeof(box_assign_after_capture(true)), 1) === Core.Box + end +end