Raising: expand a branch between buffers read back through a pointer by wsmoses · Pull Request #3016 · EnzymeAD/Enzyme-JAX

wsmoses · 2026-08-29T00:46:20Z

Stacked on #3027 (expandBufferBranches); merge that first.

When an affine.if yields one of two buffers as a memref and the result is only read back through enzymexla.memref2pointer → enzymexla.pointer2memref (the shape MFEM's Vector/DeviceTensor wrappers produce after Memref2Pointer2MemrefCast, #3068, has done its part but the intermediate buffer is a different rank), the branch cannot be raised: the yielded memref has no affine access to follow. This extends expandBufferBranches so that a memref result whose users are either direct affine.load/affine.stores, or such accesses through an m2p/p2m chain, is expanded per access — each access is cloned into a fresh affine.if with the arm's buffer substituted, the dead chain links are erased, and the original branch disappears.

The pointer-yield case (a branch yielding !llvm.ptr that is the same base as, or a gep of, the other side) is now handled generally in canonicalize-parallel by #3072 and is no longer part of this PR.

Test: test/lit_tests/raising/if_memref_yield.mlir.

Part of #2968.

🤖 Generated with Claude Code

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

wsmoses · 2026-08-30T13:15:32Z

Rebased: this branch now carries #3027's expand-buffer-branches commit as its prerequisite on top of pb/buffer-normalizations (the base ref), since the pointer-yield expansion edits that function and its test needs the raw-gep raising from #3006. Merge order: #3006 → #3027 → this.

🤖 Generated with Claude Code

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

MFEM's shared scratch selection yields the buffers themselves from an affine.if; the chosen one is converted back to a pointer and viewed again, and the access sits on that view. Follow the conversions from the yield to the access, rebuild them inside each arm on that arm's own buffer, and erase them once the access is gone: they are all that still names the branch, and while they stand it reaches the descent, which refuses a branch choosing between buffers that are written through. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

…es (#2949) Squash of pb/scalar-cf-raise @ 767c09d onto union2, minus the parts main already carries (#3101 fail-soft, #3102 scf.for raiser, #3103 masked while, #3104 racy stores): what remains is the masked-pick racy refinement and the base expandBufferBranches that #3027/#3016 extend.

…3016) MFEM's shared scratch selection yields the buffers themselves from an affine.if; the chosen one is converted back to a pointer and viewed again, and the access sits on that view. Follow the conversions from the yield to the access, rebuild them inside each arm on that arm's own buffer, and erase them once the access is gone: they are all that still names the branch, and while they stand it reaches the descent, which refuses a branch choosing between buffers that are written through. (cherry picked from commit 5400f00)

…xtension) An affine.if may yield the chosen buffer as a pointer rather than a memref; the views taken of that pointer are the accesses that give the branch away. Push each such access into a clone of the branch, viewing that arm's own pointer, the way the memref yields already expand.

…es (#2949) Squash of pb/scalar-cf-raise @ 767c09d onto union2, minus the parts main already carries (#3101 fail-soft, #3102 scf.for raiser, #3103 masked while, #3104 racy stores): what remains is the masked-pick racy refinement and the base expandBufferBranches that #3027/#3016 extend.

…3016) MFEM's shared scratch selection yields the buffers themselves from an affine.if; the chosen one is converted back to a pointer and viewed again, and the access sits on that view. Follow the conversions from the yield to the access, rebuild them inside each arm on that arm's own buffer, and erase them once the access is gone: they are all that still names the branch, and while they stand it reaches the descent, which refuses a branch choosing between buffers that are written through. (cherry picked from commit 5400f00)

…xtension) An affine.if may yield the chosen buffer as a pointer rather than a memref; the views taken of that pointer are the accesses that give the branch away. Push each such access into a clone of the branch, viewing that arm's own pointer, the way the memref yields already expand.

…es (#2949) Squash of pb/scalar-cf-raise @ 767c09d onto union2, minus the parts main already carries (#3101 fail-soft, #3102 scf.for raiser, #3103 masked while, #3104 racy stores): what remains is the masked-pick racy refinement and the base expandBufferBranches that #3027/#3016 extend.

…3016) MFEM's shared scratch selection yields the buffers themselves from an affine.if; the chosen one is converted back to a pointer and viewed again, and the access sits on that view. Follow the conversions from the yield to the access, rebuild them inside each arm on that arm's own buffer, and erase them once the access is gone: they are all that still names the branch, and while they stand it reaches the descent, which refuses a branch choosing between buffers that are written through. (cherry picked from commit 5400f00)

…xtension) An affine.if may yield the chosen buffer as a pointer rather than a memref; the views taken of that pointer are the accesses that give the branch away. Push each such access into a clone of the branch, viewing that arm's own pointer, the way the memref yields already expand.

…es (#2949) Squash of pb/scalar-cf-raise @ 767c09d onto union2, minus the parts main already carries (#3101 fail-soft, #3102 scf.for raiser, #3103 masked while, #3104 racy stores): what remains is the masked-pick racy refinement and the base expandBufferBranches that #3027/#3016 extend.

…oes it #3133 splits an access through a branch between buffers inside a GPU wrapper in polygeist-mem2reg, before raising. The raiser's own copy came in with the #2949 squash (it is what #3105 carved out, and #3027/#3016 grew); with it switched off all 145 mfem objects still compile strict, so it goes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

…es (#2949) Squash of pb/scalar-cf-raise @ 767c09d onto union2, minus the parts main already carries (#3101 fail-soft, #3102 scf.for raiser, #3103 masked while, #3104 racy stores): what remains is the masked-pick racy refinement and the base expandBufferBranches that #3027/#3016 extend. (cherry picked from commit 9e18dd1)

…es (#2949) Squash of pb/scalar-cf-raise @ 767c09d onto union2, minus the parts main already carries (#3101 fail-soft, #3102 scf.for raiser, #3103 masked while, #3104 racy stores): what remains is the masked-pick racy refinement and the base expandBufferBranches that #3027/#3016 extend. (cherry picked from commit 9e18dd1) (cherry picked from commit 0f4db9b080c9c15703c6b76fefe081b97150edd6)

…es (#2949) Squash of pb/scalar-cf-raise @ 767c09d onto union2, minus the parts main already carries (#3101 fail-soft, #3102 scf.for raiser, #3103 masked while, #3104 racy stores): what remains is the masked-pick racy refinement and the base expandBufferBranches that #3027/#3016 extend. (cherry picked from commit 9e18dd1) (cherry picked from commit 008426e)

wsmoses mentioned this pull request Aug 29, 2026

Tracking: MFEM CUDA raising to xla-gpu #2968

Open

wsmoses force-pushed the pb/buffer-branches branch from 4858cb9 to 8245f4a Compare August 30, 2026 13:01

wsmoses mentioned this pull request Aug 30, 2026

Raising: expand buffer-yielding branches into their accesses #3027

Open

wsmoses changed the base branch from pb/buffer-branches to pb/buffer-normalizations August 30, 2026 13:09

wsmoses force-pushed the pb/if-ptr-yields branch from 15ed014 to cf9a46f Compare August 30, 2026 13:09

wsmoses changed the title ~~Raising: expand affine.if pointer yields into their accesses~~ Raising: expand a branch between buffers read back through a pointer Sep 2, 2026

wsmoses changed the base branch from pb/buffer-normalizations to pb/buffer-branches September 2, 2026 15:32

wsmoses force-pushed the pb/if-ptr-yields branch from 3484471 to 5400f00 Compare September 2, 2026 15:32

This was referenced Sep 2, 2026

Raising: chase register-array scratch #3015

Closed

Union of the open MFEM raising PRs, one commit per PR (rebuilt on main, #3133 merged) #3117

Draft

Navigation Menu

Raising: expand a branch between buffers read back through a pointer - #3016

Raising: expand a branch between buffers read back through a pointer#3016
wsmoses wants to merge 1 commit into
pb/buffer-branchesfrom
pb/if-ptr-yields

wsmoses commented Aug 29, 2026 •

edited

Loading

Uh oh!

wsmoses commented Aug 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

wsmoses commented Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wsmoses commented Aug 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

wsmoses commented Aug 29, 2026 •

edited

Loading