Skip to content

Fix composite pass missing captured values in ops with nested regions#9760

Open
mstojkovicTT wants to merge 2 commits intopytorch:masterfrom
tenstorrent:mstojkovic/composite_nested_regions_upstream
Open

Fix composite pass missing captured values in ops with nested regions#9760
mstojkovicTT wants to merge 2 commits intopytorch:masterfrom
tenstorrent:mstojkovic/composite_nested_regions_upstream

Conversation

@mstojkovicTT
Copy link
Copy Markdown

Problem

The BuildStableHLOCompositePass performs a reverse graph traversal from output boundary markers to input boundary markers to collect all ops that belong to a composite. The traversal followed op->getOperands() to discover dependencies, but this only captures values passed directly to an op, and it misses values that are used inside nested regions (e.g., the body of a stablehlo.reduce) but defined in the parent function.

This caused the generated composite impl function to reference values that were never included as arguments or cloned ops, leading to invalid MLIR.

Fix

After processing each op's direct operands, the traversal now also inspects all nested regions using mlir::getUsedValuesDefinedAbove() (mlir doc) to find captured values. These are handled with the same logic as direct operands:

  • Block arguments added as composite arguments
  • Constants included in the impl function body
  • Other ops added to the traversal worklist

@mstojkovicTT
Copy link
Copy Markdown
Author

Tagging @qihqi and @ysiraichi for visibility

@nvukobratTT nvukobratTT requested review from qihqi and ysiraichi March 30, 2026 15:09
@ysiraichi
Copy link
Copy Markdown
Collaborator

Could you add some tests that show the failure? The build failure is detailed in this comment.
Again, since I'm not familiar with this module, I can't really review it.

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.

3 participants