Add Escape Analysis by SaumitraLohokare · Pull Request #551 · pschachte/wybe

SaumitraLohokare · 2026-03-22T23:31:10Z

This PR implements intraprocedural escape analysis for LPVM alloc instructions. Allocations whose result does not outlive the current procedure call are now emitted as LLVM alloca (stack) instead of heap allocations, avoiding GC pressure for short-lived local structs.

How it works

Two checks are combined at each alloc site:

Mutation-chain reachability - backwards fixed-point over the body. A variable escapes if it is an output parameter, or reaches one via mutate or pass-through edges.
Alias-map check - checks whether the alloc result is connected to a global or aliased output in the incremental alias disjoint-set.

If neither fires and the size is a compile-time constant, the {stack} flag is added to the alloc. LLVM codegen then lowers it to alloca i8, i64 <size> + ptrtoint instead of a GC malloc call.

Verification

pub type point { pub point(x: int, y: int) }

# mid is never returned → should be stack-allocated
pub def print_midpoint(p1: point, p2: point) use !io {
    ?mid = point((p1^x + p2^x) / 2, (p1^y + p2^y) / 2)
    !println(mid^x)
    !println(mid^y)
}

# q is an output parameter → heap-allocated
pub def make_point(?q: point) {
    ?q = point(10, 20)
}

LLVM:

make_point: Heap-allocated.

define external fastcc i64 @"escape_ex#.make_point<0>"() {
  %"tmp#6##0" = call ccc ptr @wybe_malloc(i32 16)   ; heap allocation
  %"tmp#3##0" = ptrtoint ptr %"tmp#6##0" to i64
  %"tmp#7##0" = inttoptr i64 %"tmp#3##0" to ptr
  store i64 10, ptr %"tmp#7##0"                      ; x = 10
  %"tmp#8##0" = add i64 %"tmp#3##0", 8
  %"tmp#9##0" = inttoptr i64 %"tmp#8##0" to ptr
  store i64 20, ptr %"tmp#9##0"                      ; y = 20
  ret i64 %"tmp#3##0"
}

print_midpoint: Stack-allocated.

define external fastcc void @"escape_ex#.print_midpoint<0>"(i64 %"p1##0", i64 %"p2##0") {
  ; compute (p1.x + p2.x) / 2
  ...
  ; compute (p1.y + p2.y) / 2
  ...
  ; mid = point(...) — stack allocated
  %"tmp#44##0" = alloca i8, i64 16, align 8          ; ← alloca, not malloc
  %"tmp#25##0" = ptrtoint ptr %"tmp#44##0" to i64
  %"tmp#45##0" = inttoptr i64 %"tmp#25##0" to ptr
  store i64 %"tmp#1##0", ptr %"tmp#45##0"            ; mid.x
  %"tmp#46##0" = add i64 %"tmp#25##0", 8
  %"tmp#47##0" = inttoptr i64 %"tmp#46##0" to ptr
  store i64 %"tmp#5##0", ptr %"tmp#47##0"            ; mid.y
  ; read back and print
  ...
  ret void
}

pschachte

This looks great, and I really appreciate your documentation here. Makes the code much easier to review.

I've placed a few requests for documentation clarification inline, and one suggested improvement.

The other thing for you to fix is all the failing test cases. Before resubmitting, please run the update-exp script and scroll carefully through the differences in each file, making sure that they are all as expected (they're mostly calls to foreign lpvm alloc that have had a {stack} flag added, and the consequent changes to LLVM code). I.e., make sure that the all those stack allocations are safe. For each test case, if so, say yes, and it will replace the .exp file with the new output. Then add all the changed .exp files to the pull request branch. If not, you'll need to fix the bugs.

You should also add your new test case to the test-cases/final-dump directory and run update-exp; it will show you the output and prompt you to accept or reject it. Then also add the source and .exp files to the branch. Finally, please add a couple more tests to that test case, one that creates a fresh structure, passes it to another proc, which passes it back, and see if the escape analysis still sees that it doesn't escape (it's OK if it doesn't, but we should know), and one that does something similar, but the structure really does escape, so we can see that it doesn't optimise cases it shouldn't.

pschachte · 2026-03-23T05:53:41Z

+                -- argument cannot be a tail call (the callee would dereference
+                -- a pointer into a stack frame that tail-call would have freed).
+                -- LLVM enforces this automatically; sets doesAlloca accordingly.
+                Just sizeVal -> do


Does LLVM's alloca only work for fixed size allocations? If so, please say that here in a comment. If not, please add a comment saying this is a temporary restriction. Unless you can generalise it so it works for variable sized allocations.

alloca also works for variable sized allocations from my research, but I'm not super familiar with LLVM in general so I chose to take a safer approach. I'll comment this 👍 Thank you

I don't see why we can't support variably sized stack allocations :)

Hello, sorry for the late reply. Been a bit busy with proposal and presentations. But I looked into variable sized alloca and I think it isn't worth the effort. Mainly for two reasons:

Variable-sized alloca forces LLVM to keep a frame pointer because the stack frame size can't be computed statically. Wybe actively uses TCO (TailCallModCons is an optimization flag). These two features are fundamentally incompatible.

malloc can return null; a stack overflow from a large VLA just segfaults or corrupts the stack silently. For objects that are large or runtime-sized, heap allocation is the safer default.

I think we should support it in LLVM generation, but never generate it in the escape analysis. There are cases where we alloc a variable length array, and if the user wanted to supply their own stack flag I don't see why we shouldn't support that

Since tail recursion optimisation only works for certain calling conventions, and from what you've told me, you can't use a calling convention that supports TRO together with variable-sized alloca, I think we should call this Not Yet Implemented for now (there's a nyi function to report such missing features). If we get around to doing our own TRO during LLVM generation, we should revisit this.

pschachte · 2026-03-23T05:57:35Z

+isAllocOrMutate (PrimForeign "lpvm" "mutate" _ _) = True
+isAllocOrMutate _                                  = False
+
+-- | Compute the set of variables whose allocated memory must survive the


I think you mean "may" rather than "must". This is an important distinction: this optimisation can only be done if it's impossible for the allocation to survive the proc body it appears in. You have to be conservative in approximating this (and you have to approximate; you can't always get it exactly right).

Oh, I think I understand what you mean. OK, it's difficult to express this clearly, but I think it's clearer to write of what will happen when the program runs, so think of a value as escaping a proc body if it is possible that the value could be referred to after the proc returns.

pschachte · 2026-03-23T06:18:23Z

+        prims = List.map content $ collectAllBodyPrims body
+        -- For mutate(fIn, fOut, offset, destr, size, startOff, member, ...),
+        -- if fOut escapes then:
+        --   (a) fIn escapes (same struct, just a new version)


fIn might not escape in this case, if the destructive flag is not set in the mutate. In that case, the mutate instruction makes a fresh copy, so fOut escaping doesn't imply that fIn does, unless the structure contains pointers.

If your analysis/transformation is happening before the destructive update transformation, it would be useful to document that this limits the stack allocation transformation here. Otherwise, I think you should be able to do better here.

SaumitraLohokare · 2026-03-23T10:34:57Z

I also had the thought of adding a flag to disable/enable escape analysis. Is this a good idea or should we just leave it as always enabled?

pschachte · 2026-03-24T00:16:07Z

Adding a command line flag is a good idea. It is stack allocation, not the escape analysis, that would be relevant to programmers, so I would name the switch in terms of stack allocation rather than escape analysis. So far, we optimise by default, so I think the switch should disable automatic stack allocation.

SaumitraLohokare · 2026-03-31T11:51:23Z

I added the flag for stack-alloc under --opt. It is on by default. And also added a new test case for stack allocation.
Ran update-exp as well, all tests passing.

jimbxb · 2026-03-31T11:52:07Z

+                -- Side-effect: any call that receives this address as an i64
+                -- argument cannot be a tail call (the callee would dereference
+                -- a pointer into a stack frame that tail-call would have freed).
+                -- LLVM enforces this automatically; sets doesAlloca accordingly.


We should be able to do a tail call if the stack allocated space never enters the tail call.

Something like

def foo use !io { ?p = position(1, 2) !print(p) !foo }

Should be able to stack allocate p.

I'm thinking this case could be fairly common, say if something is used for an intermediate computation before a tail call.

Good point! Fixed - instead of blocking all tail calls once any alloca is emitted, we now track which variables hold stack-allocated addresses and only block tail calls that directly receive one of those variables as an argument. Since escape analysis already guarantees stack-allocated addresses can't reach a call through any other mechanism, checking direct arguments is sufficient.
This adds back a lot of tail calls.

jimbxb · 2026-03-31T11:54:28Z

+                -- argument cannot be a tail call (the callee would dereference
+                -- a pointer into a stack frame that tail-call would have freed).
+                -- LLVM enforces this automatically; sets doesAlloca accordingly.
+                Just sizeVal -> do


I don't see why we can't support variably sized stack allocations :)

jimbxb · 2026-03-31T11:58:44Z

+        -- For mutate(fIn, fOut, offset, destr, size, startOff, member, ...),
+        -- if fOut escapes then:
+        --   (a) fIn escapes (same struct, just a new version).  Note: this
+        --       analysis runs before the destructive-update transformation sets


Why can't we run it afterwards?

technically we could, but it would require splitting the current single-pass design. The destr flags are set incrementally during transformBody using an alias map that's built as the body is traversed; computeEscapedVars runs before that traversal starts, so the updated flags don't exist yet. To use post-destr escape analysis we'd need a preliminary sub-pass just to accumulate the alias map and annotate destr flags, then re-run escape analysis on the result. Happy to do that as a follow-up if the precision loss is worth fixing.

Sure. I will leave this unresolved as a reminder (assuming it doesn't block the approval)

jimbxb · 2026-03-31T12:01:13Z

+        --     the same type escapes.  This is sound because these ops can't
+        --     embed a value into a result of a different type.
+        --   • user-defined calls (PrimCall) and foreign calls to other
+        --     languages (e.g. C): we have no visibility into the callee, so


I would hope for Wybe procs we do.

Is there a way we can annotate Params as saying "this input escapes iff some output also escapes"?

Not sure that makes sense, but I'm sure we can do better here. For (co-recursive) price we'd have to be conservative or attempt to find a fixed point.

For instance,

def foo(p0, ?p1, q0, ?q1) { p0 = ?p1 q0 = ?q1 }

For a call to foo, p0 should only escape if p1 does, not if q1 does.

Unless I'm misunderstanding, I believe the current analysis would have p0 and q0 escape if p1 or q1 escape.

You are correct. I did try to implement a fix for this, but my fix is causing ./complex/exp/testcase_multi_specz-int_list.out to fail with a segmentation fault.
Since this is more of a precision issue, rather than correctness, I'd like to move this to my next PR so I can take a closer look at this.

Again, I'll leave this open as a reminder

jimbxb · 2026-05-19T23:21:06Z

+# escaped), so p also does not escape.
+def {noinline} test_pass_through use !io {
+    ?p = position(1, 2)
+    !pass_back(p, ?q)


Why does this have a !? It shouldn't need it

SaumitraLohokare added 3 commits February 24, 2026 22:17

Initial plan commit

4bfc851

Merge branch 'master' of github.com:SaumitraLohokare/wybe

e58568e

Add old docs

739c578

pschachte requested changes Mar 23, 2026

View reviewed changes

jimbxb reviewed Mar 23, 2026

View reviewed changes

Comment thread src/Transform.hs Outdated

Comment thread src/Transform.hs

jimbxb requested changes Mar 31, 2026

View reviewed changes

SaumitraLohokare marked this pull request as draft April 1, 2026 00:50

SaumitraLohokare and others added 8 commits May 20, 2026 06:46

Merge branch 'pschachte:master' into master

e814f02

Added escape analysis

61b9469

Removed temp docs

42f818e

Fixed comments and updated test cases

562bb65

Added new test case for stack allocation

c645606

Added stack frame size limit and compiler flag

bc286c9

Added note on variable size alloca

fe3b1c0

Updated tests after merging master

25900bd

SaumitraLohokare force-pushed the feat/escape-analysis branch from ab5bf6e to 25900bd Compare May 19, 2026 21:26

Fixed tail calls

19b00b8

SaumitraLohokare marked this pull request as ready for review May 19, 2026 23:07

jimbxb reviewed May 19, 2026

View reviewed changes

Conversation

SaumitraLohokare commented Mar 22, 2026

How it works

Verification

Uh oh!

pschachte left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SaumitraLohokare commented Mar 23, 2026

Uh oh!

pschachte commented Mar 24, 2026

Uh oh!

SaumitraLohokare commented Mar 31, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jimbxb Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jimbxb Mar 31, 2026 •

edited

Loading