Skip to content

feature: make LambdaInstUseVisitor clobbering aware - #526

Open
johnbowen42 wants to merge 30 commits into
mainfrom
feature/bowen/clobbering-aware-use-analysis
Open

johnbowen42 wants to merge 30 commits into
mainfrom
feature/bowen/clobbering-aware-use-analysis

Conversation

@johnbowen42

@johnbowen42 johnbowen42 commented Sep 16, 2026 •

Copy link
Copy Markdown
Collaborator

Currently if we have LLVM IR shapes like

// Assume %0 is a closure pointer
store int32 %1, %0
store int32 %2, %0
call registered_lambda(%0)

our analysis of the lambda callsite cannot automatically tell which store to trust as the definitive source of truth for the first slot of the lambda at the callsite. Fortunately, LLVM has O(N) (N is size of the function) built-in analysis called MemorySSA, which we can "at the callsite call registered_lambda(%0), which instruction accessing %0 clobbers", where clobbering basically means is the source of truth for that memory value.

Fortunately, this LLVM IR shape is relatively exotic for lambda closure pointer. Nonetheless, we want Proteus to conservatively and correctly deal with simple cases of memory clobbering, and not trust the def-use graph as a source of ground truth for memory locations. This PR introduces a MemorySSA oracle model, that the def-use walk uses as a guide when analyzing instructions that write to lambda closure pointers. The oracle allows the analysis ask, "which instruction clobbers the lambda closure pointer at a given lambda callsite??" Now, there are many edge cases to this. This PR tries to be as conservative as possible while allowing for some complexity in LLVM IR shape.

Here are some of the edge cases where we fail to specialize
(1) What if the analysis finds that different slots in a lambda are stored from different pointers? -- fail to specialize
(2) What if we find that multiple user-defined function calls determine the value of different slots of the pointers -- fail to specialize
(3) What if we find multiple memcpys that copy different lambda slots from different source pointers? -- fail to specialize

Here are some cases where we do specialize
✅ the lambda closure pointer is clobbered at the callsite by a collection of stores, all from the same base ptr
✅ the lambda closure pointer is clobbered at the callsite by a memory intrinsic
✅ the lambda closure pointer is clobbered at the callsite by a function call, which writes to the closure pointer using one of the above patterns

We try to handle all unsupported LLVM IR patterns from the Clang frontend with #540

@johnbowen42
johnbowen42 added this pull request to stack #527 September 16, 2026 20:08
@johnbowen42
johnbowen42 force-pushed the feature/bowen/clobbering-aware-use-analysis branch from 37dfebc to 25cb26d Compare September 16, 2026 20:10
Base automatically changed from feature/bowen/mfem-raja-analysis-changes to main September 16, 2026 22:23
@johnbowen42
johnbowen42 force-pushed the feature/bowen/clobbering-aware-use-analysis branch 2 times, most recently from e1298c3 to ad2d3ee Compare September 17, 2026 18:37
@johnbowen42
johnbowen42 force-pushed the feature/bowen/clobbering-aware-use-analysis branch from ad2d3ee to 976c996 Compare September 24, 2026 00:09
johnbowen42 and others added 19 commits October 8, 2026 17:07
…iling due to recognition of various nocapture intrinsics getting marked as clobbers
         81 - lambda_call_shapes.HIP (Failed)
         82 - lambda_memory_intrinsic_shapes.HIP (Failed)
         83 - lambda_pointer_field_load.HIP (Failed)
         84 - lambda_clobber_shapes.HIP (Failed)
         85 - lambda_clobber_adversarial.HIP (Failed)
         90 - lambda_inst_visitor_adversarial.HIP (Failed)
         91 - lambda_inst_visitor_call_context.HIP (Failed)
         92 - lambda_inst_visitor_call_clobber.HIP (Failed)
@johnbowen42
johnbowen42 force-pushed the feature/bowen/clobbering-aware-use-analysis branch from c328997 to 5091dc4 Compare October 12, 2026 19:22
johnbowen42 and others added 4 commits October 12, 2026 17:11
… depth cap

visitArgument now counts only direct callsites (CB->isCallee) and declines
when the function's address escapes; previously a call merely referencing
the function was counted as its caller and the walk could follow an
unrelated operand. analyzeFunction requires exactly one value-returning
ReturnInst, matching the forward use visitor, instead of trusting the
first return of a multi-return callee. An AnalysisDepth cap bounds nested
PHI-clone and function sub-analyses, which previously recursed without
bound on cyclic PHI webs and recursive callees since each sub-analysis
starts with a fresh Seen set.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…nchored queries

Context guards: entering a callee whose formal already carries a different
remaining offset now declines at both entry points (returned-alias and
deferred); ValueOffsetMap is not call-context scoped, so a second entry
would silently retarget facts computed for the first.

Recursion/resources: pointer-returning helpers that recurse (directly or
through the continuation parent chain) are rejected -- each entry mints a
fresh ReturnContinuation and Seen is keyed on (Value, continuation), so
recursion previously re-enqueued work without bound. popBack enforces a
128K work-item budget as a backstop.

Deterministic reconciliation: ClobberQuery now carries TrackedByteOffset
-- the byte the query resolves -- as the query identity (cache key); the
unused SizeOfWriteToBasePtr is removed. Every MemorySSA query routes
through getQueryLocation (exact tracked-byte location when a visited
pointer denotes it, covering superset otherwise) and the
proven-returned-alias skip now applies on every query path. Candidate
consistency (frame/tracked-ptr/boundary) is checked before any query, and
the sole-memcpy precise retry is subsumed by precise-first querying.
Soundness does not depend on the anchor choice: acceptance requires the
MemorySSA selection to be a collected covering candidate.

Verified via the pseudo-GPU harness on LLVM 19.1.7 and 22.1.0: per-kernel
decisions identical across 15 lambda suites, lambda_mfem_style_launch
specializes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…letion, stable Seen keys

Deferred work items are now pushed through the same accounting popBack
undoes, instead of WorkList.append which decremented PendingWorkItems
without a matching increment and wrapped the unsigned counter.
finishCurrentWorkItem now also runs when a popped duplicate is skipped,
since that duplicate may be the final pending item of a continuation.
The visitor owns every ReturnContinuation it creates so the raw addresses
used as Seen keys cannot be freed and reused within one analysis.

Verified via the pseudo-GPU harness: decisions identical across 15 lambda
suites on LLVM 19 and 22.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
johnbowen42 and others added 7 commits October 12, 2026 18:25
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
lambda_interior_slot_overwrite re-stores a two-slot functor's first slot
from the authoritative source and overwrites the second slot from a donor
functor before the call. The analysis anchors its reconciliation query at
the narrow slot-0 store, so the slot-1 write is invisible: it does not
cover the tracked byte (candidate collection skips its GEP) and does not
alias the 4-byte anchor once SROA gives both stores one provable base. The
analysis concludes the whole functor is authoritative from Initial and
specializes; slot 1 executes with Donor's value, so both FileCheck
assertions fail (a LambdaSpec line appears against CHECK-NOT, and Y prints
757 instead of 999).

This FAILS until functor-extent tracking lands (gate item 8 in
docs/dev/reviews/clobbering-aware-use-analysis.md): the query and candidate
coverage must reason about [offset, offset + sizeof(functor)), not byte 0.
Note the optnone variant of this shape declines only by accident -- spill
reloads give the two slot stores unrelatable bases, so AA conservatively
aliases the anchor -- which is the anchor-luck failure mode the review
describes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ysis

The wrapper operator's proteus.wrapper_call annotation gains a second
payload: sizeof the registered closure, folded per-target by Sema at
template instantiation, immune to IR type erasure. Both compilations parse
it and stamp proteus.functor_size metadata on the wrapper operator; on the
host, attachFunctorSizeToCallsiteMetadata recomputes the size from the
registered functor type, so setU64FunctionMetadata's conflict fatal doubles
as a host/device layout cross-check. The device manifest requires the
metadata and passes the extent to analyzeLambdaUses.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Three size-based checks, each pinned by a test, none touching witness
resolution: (a) the fallback MemorySSA query location is bounded to the
prefix [0, tracked offset + size), so writes past the functor (envelope
suffixes, enclosure trailing captures) can no longer be selected as the
clobber and spuriously decline correct shapes; (b) a transfer that covers
the tracked byte must cover the functor's full extent or decline, in both
visitors -- a shorter copy (the classic wrong-sizeof bug, legal C++ with no
casts) splices two sources into one functor and the byte-0 witness would
lie for the uncovered slots; (c) memset declines on any overlap with the
extent instead of ignoring interior-only hits.

The precise zero-remaining-GEP location is deliberately NOT widened to the
extent: interior slot stores are legitimate parallel copies in slot-wise
assembly (the mfem shape), and extent-wide visibility requires per-piece
reconciliation to accept them. Interior-store soundness is a frontend
contract; see gate 8 in docs/dev/reviews/clobbering-aware-use-analysis.md.

Verified via the pseudo-GPU harness on LLVM 19: all 16 lambda suites
decision-identical to the pre-change baseline, and lambda_size_test's
envelope-suffix case flips from spurious decline to specialize.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
lambda_size_test pins the two extent behaviors: the enclosure shape (slot-wise assembly with a trailing capture past the functor) must specialize, and the envelope-suffix shape must specialize rather than spuriously decline when zero-index GEPs fold at -O1 and the suffix store would otherwise be selected as the clobber. The review doc records the resolved gate-8 design and the remaining frontend-diagnostic work item.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@johnbowen42
johnbowen42 force-pushed the feature/bowen/clobbering-aware-use-analysis branch 2 times, most recently from c465b25 to 3a50d90 Compare October 15, 2026 02:45

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant