[Graph Optimization](feat) Layout Optimization for main-dev

[zhanwei33]

Summary

• Port the graph/layout optimization work from [Graph Optimization](feat) Layout Optimization #1262 onto main-dev.
• Integrate load/store transpose, transpose-pointwise reorder, store
  coalescing, and migrated legacy layout-memory rules.
• Auto-configure the graph UB budget from the Ascend architecture while
  retaining explicit capacity bounds.

Why

main and main-dev have diverged, so opening #1262's original head directly
against main-dev would include dozens of unrelated main commits. This PR
replays only the feature commits on top of main-dev.

main-dev integration

• Preserve main-dev's runtime Ascend patch initialization.
• Fold changes to upstream-managed Triton files into
  third_party/ascend/patch/triton-ascend-3.6.0.patch.
• Expose the migrated graph-memory headers to all patched Triton tools.

Validation

• Built an aarch64 CPython 3.11 triton-ascend wheel with LLVM/Clang 22 and
  CANN 9.0.
• Layout-memory compiler contracts: 30 passed, 4 hardware-gated tests skipped.
• Graph optimizer load/store FileCheck: passed.
• Strided-axis FileCheck with variable scoping enabled: passed.
• NPU vector-add: passed with maximum difference 0.0.

Related: #1262

Graph Optimize : Layout Optimization

This document covers only:

1. registration of the Graph Optimize pass and its rules;
2. the main optimization path, non-match exits, rewriting, and rewrite-failure exits.

1. Registration Flow

flowchart TD
  A["Passes.td defines<br/>GraphOptimize<br/>constructor=<br/>createGraph<br/>OptimizePass"]
  B["C++/pybind exposes<br/>add_graph_optimize"]
  C["Python make_ttir"]
  D{"enable_graph_<br/>optimize?"}
  E["Do not add<br/>GraphOptimizePass<br/>to the PassManager"]
  F["Build<br/>GraphOptimization<br/>Options<br/>rule mask /<br/>rewrite cap /<br/>UB / remarks /<br/>force SIMT"]
  G["createGraph<br/>OptimizePass<br/>(options)"]
  H["Add<br/>GraphOptimizePass<br/>to the PassManager"]
  I["pm.run invokes<br/>runOnOperation"]
  J["populateBuiltin<br/>GraphOptimization<br/>Rules"]
  K["Construct and<br/>enable rules<br/>according to<br/>the rule mask"]

  A --> B
  C --> D
  D -->|No| E
  D -->|Yes| B
  B --> F --> G --> H --> I --> J --> K

Loading

Registered rules:

Order	Rule	Enable condition
1	LoadStoreTransposeRule	Enabled in the rule mask
2	TransposePointwiseReorderRule	Enabled in the rule mask
3	StoreCoalescingRule	Enabled in the rule mask; produces no candidates when the UB capacity is 0
4	RowCoalescingRule	Enabled in the rule mask and forceSimtOnly=true

The first three rules run in a fixed phase order. RowCoalescingRule is attempted once after those phases and is not limited by the per-function rewrite budget.

2. Main Graph Optimization Flow

2.1 Main Optimization and Non-Match Exit Flow

flowchart TD
  A["GraphOptimizePass<br/>::runOnOperation"]
  B["Validate and<br/>stabilize options"]
  C["Construct enabled<br/>builtin rules"]
  D["For each tt.func:<br/>create context;<br/>rewriteCount = 0"]
  E{"Any enabled phase<br/>remaining among<br/>the first three<br/>rules?"}
  F["Select next phase<br/>in fixed order:<br/>LoadStoreTranspose<br/>Rule →<br/>TransposePointwise<br/>ReorderRule →<br/>StoreCoalescing<br/>Rule"]
  G{"Per-function<br/>rewrite budget<br/>remaining?"}
  H["Prepare analyses<br/>needed by pattern<br/>matching and proof"]
  I["findCandidates:<br/>run pattern match<br/>and prove<br/>applicability;<br/>do not modify IR"]
  J{"Any candidate<br/>plans?"}
  K["Sort candidates<br/>benefit →<br/>program order →<br/>rule id"]
  L["Discard stale<br/>plans;<br/>revalidate the<br/>remaining plans"]
  M{"Any usable plan<br/>remaining?"}
  N["Enter the<br/>rewrite flow"]
  O["Rewrite succeeds"]
  P["Destroy all<br/>stale plans;<br/>invalidate<br/>analyses; epoch++"]
  Q{"Was this<br/>RowCoalescing?"}
  R["rewriteCount++;<br/>continue current<br/>phase"]
  S["Non-match:<br/>current attempt<br/>does not<br/>modify IR;<br/>no pass failure"]
  T{"Processing<br/>RowCoalescing?"}
  U["End the current<br/>phase"]
  V{"RowCoalescing<br/>enabled by mask<br/>and forceSimtOnly?"}
  W["Attempt<br/>RowCoalescingRule<br/>once after phases;<br/>not limited by the<br/>per-function<br/>rewrite budget"]
  X["Finish current<br/>tt.func;<br/>process the next"]

  A --> B --> C --> D --> E
  E -->|Yes| F --> G
  E -->|No| V
  G -->|Yes| H --> I --> J
  G -->|No| V
  J -->|Yes| K --> L --> M
  J -->|No| S
  M -->|Yes| N --> O --> P --> Q
  M -->|No| S
  Q -->|No| R --> G
  Q -->|Yes| X
  S --> T
  T -->|No| U --> E
  T -->|Yes| X
  V -->|No| X
  V -->|Yes| W --> H

Loading

Preparing analyses is only a prerequisite for matching. The actual pattern match and applicability proof run inside findCandidates(), which does not modify the IR. The optimization takes effect only after candidate matching and revalidation succeed and apply() commits the new IR.

The maxRewritesPerFunction budget is initialized once per tt.func and shared by LoadStoreTransposeRule, TransposePointwiseReorderRule, and StoreCoalescingRule in that fixed phase order. Reaching the budget skips the remaining phases among these three rules, but it does not suppress the one-shot RowCoalescingRule attempt.

The following cases take the normal non-match exit:

• findCandidates() produces no candidate plan;
• the candidate's epoch is stale;
• revalidate() fails;
• no usable plan remains after filtering.

These cases only discard a candidate, end the current phase among the first three rules, or finish the one-shot Row attempt; they do not fail the compilation.
Only an analysis-construction failure or findCandidates() itself returning failure() is an internal processing error that triggers signalPassFailure().

2.2 Rewrite and Rewrite-Failure Exit Flow

flowchart TD
  A["Select a<br/>successfully<br/>revalidated plan"]
  B["Create IRRewriter"]
  C["plan.apply<br/>(rewriter)"]
  D["Attempt to create<br/>replacement IR<br/>or a detached<br/>sandbox"]
  E{"Did creation<br/>succeed?"}
  F["Verify new ops,<br/>types, attributes,<br/>and additional<br/>proofs"]
  G{"Did verification<br/>or a secondary<br/>proof fail?"}
  H["Replace old uses /<br/>IR / module attrs"]
  I["Erase old ops"]
  J["apply returns<br/>success"]
  K["Driver destroys<br/>plans and<br/>invalidates<br/>analyses"]
  L["Erase IR created<br/>by this attempt<br/>or discard the<br/>sandbox"]
  M["Do not commit<br/>or erase<br/>the old IR"]
  N["apply returns<br/>failure"]
  O["Driver clears<br/>selectedPlan<br/>and the remaining<br/>plans"]
  P["function.emitError"]
  Q["signalPassFailure<br/>+ return"]

  A --> B --> C --> D --> E
  E -->|Yes| F --> G
  G -->|No| H --> I --> J --> K
  E -->|No| L
  G -->|Yes| L
  L --> M --> N --> O --> P --> Q

Loading

Rewriting uses transactional commit semantics: construct and verify the new IR first, then replace the old IR only after every check passes. After a successful rewrite, all analyses are invalidated, and the next match must be computed against a new epoch. If creation, verification, or a secondary proof fails, the attempt cleans up its new IR, preserves the old IR, and terminates the Graph Optimize pass.

A rewrite failure differs from a non-match:

• Non-match: a normal optimization bypass; preserve the original IR and continue with a later phase or function.
• Rewrite failure: a candidate was selected, but transactional construction or verification failed; roll back the new IR, preserve the old IR, and terminate the Graph Optimize pass.

3. Summary

The core Graph Optimize path is:

Register the pass and rules → match candidates per function → exit normally when no candidate exists → revalidate a candidate → transactionally construct and verify new IR → commit and invalidate analyses on success; on rewrite failure, clean up the new IR, preserve the old IR, and terminate the pass.

[zhanwei33]

retry

[zhanwei33]

/retry

[zhanwei33]

/retry

[zhanwei33]

/retry

[zhanwei33]

/retry

[zhanwei33]

/retry

[zhanwei33]

/retry