Skip to content

Improve detection precision of detect_unfused_reshapes by implementing physical layout verification and first-party attribution filtering: - #3339

Open
copybara-service[bot] wants to merge 1 commit into
masterfrom
test_975584585
Open

Improve detection precision of detect_unfused_reshapes by implementing physical layout verification and first-party attribution filtering:#3339
copybara-service[bot] wants to merge 1 commit into
masterfrom
test_975584585

Conversation

@copybara-service

Copy link
Copy Markdown

Improve detection precision of detect_unfused_reshapes by implementing physical layout verification and first-party attribution filtering:

  1. Physical Relayout Gate (_is_real_relayout): Drop zero-cost logical bitcast reshapes/transposes that preserve physical linearized memory ordering under identical minor-to-major layouts.
  2. FLOPs == 0 filter: Exclude misclassified compute operations reporting non-zero FLOPs.
  3. Tightened Compute Opcodes: Restrict downstream fold targets to genuinely foldable compute (dot, convolution, einsum, non-collective custom-call).
  4. Collective Custom-Call Filter: Exclude collective communications (all-gather, all-to-all, reduce-scatter) which require collective schedule optimization rather than GEMM folding.
  5. Fusion-Body Peek: Look inside consumer fusion bodies to verify presence of compute ops (dot/conv) while rejecting pure elementwise fusions.
  6. First-Party Attribution Walk: Resolve leaf framework wrapper attributions (flax, jax, etc.) via HLO stack frames to true user model code, dropping orphaned framework wrappers.
  7. Source Line Verification: Check source code on disk when available to filter passthrough statements.
  8. Deduplication and Ranking: Group duplicate operations by source callsite and result shape, sorting findings by total self-time descending with bytes accessed reported alongside.

…ing physical layout verification and first-party attribution filtering:

1. Physical Relayout Gate (_is_real_relayout): Drop zero-cost logical bitcast reshapes/transposes that preserve physical linearized memory ordering under identical minor-to-major layouts.
2. FLOPs == 0 filter: Exclude misclassified compute operations reporting non-zero FLOPs.
3. Tightened Compute Opcodes: Restrict downstream fold targets to genuinely foldable compute (dot, convolution, einsum, non-collective custom-call).
4. Collective Custom-Call Filter: Exclude collective communications (all-gather, all-to-all, reduce-scatter) which require collective schedule optimization rather than GEMM folding.
5. Fusion-Body Peek: Look inside consumer fusion bodies to verify presence of compute ops (dot/conv) while rejecting pure elementwise fusions.
6. First-Party Attribution Walk: Resolve leaf framework wrapper attributions (flax, jax, etc.) via HLO stack frames to true user model code, dropping orphaned framework wrappers.
7. Source Line Verification: Check source code on disk when available to filter passthrough statements.
8. Deduplication and Ranking: Group duplicate operations by source callsite and result shape, sorting findings by total self-time descending with bytes accessed reported alongside.

PiperOrigin-RevId: 975584585
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.

0 participants