[D2M] Implement graph coloring-based DST register allocation

[brnorris03]

Ticket

N/A

Summary

New pass d2m-insert-dst-register-gc implements graph coloring register allocation for DST (destination registers). Reduces DST usage by reusing slices for non-interfering values.
Note that while this patch appears large, more than 80% of it is tests.

Motivation

Current InsertDstRegisterAccess uses sequential allocation - each value gets a new DST slice. Graph coloring identifies values with non-overlapping lifetimes and assigns them to the same slice, reducing DST pressure.

Example: Diamond pattern (4 values, 2 interference edges) needs only 2 slices instead of 4 (50% reduction).

     dst0
    /    \
  dst1   dst2
    \    /
     dst3

Interference edges: dst0 - dst1, dst1 - dst2, dst2 - dst3. No edges between dst0 - dst2 or dst1 - dst3 (early releases enable reuse). Graph coloring assigns: dst0 and dst2 → slice 0, dst1 and dst3 → slice 1.

Changes

Overview of design

The design includes the defintion of simple abstract interfaces for DstAnalysis and ColoringStrategy which enable multiple implementations outside of any specific pass defiition.

                          DstAnalysis
                             ^   ^
                   inherits /     \ inherits
                           /       \
               DstAnalysisBasic   DstAnalysisGraphColoring
                                      | composes
                                      v
                              ColoringStrategy
                                 ^       ^
                     implements /         \ implements
                               /           \
               ChaitinBriggsColoring   GreedyColoring
                              |            |
                              | uses       | uses
                              v            v
                           InterferenceGraph utils

The analysis and transformation passes use the above strategies:

D2MDstRequirementAnalysisPass
    ├─ reads option 'strategy'
    ├─ instantiates DstAnalysis implementation
    │     ├─ basic  -> DstAnalysisBasic
    │     ├─ greedy -> DstAnalysisGraphColoring + GreedyColoring
    │     └─ graph-coloring -> DstAnalysisGraphColoring + ChaitinBriggsColoring
    └─ invokes analysis->analyze(funcOp) to report required slices

D2MInsertDstRegisterGCPass
    ├─ invokes DstCapacityAnalysis(funcOp) for capacity limit
    ├─ instantiates runtime strategy based on option
    │     ├─ basic  -> DstAnalysisBasic (equivalent to D2MInsertDstRegisterAccess)
    │     ├─ greedy -> DstAnalysisGraphColoring + GreedyColoring
    │     └─ graph-coloring -> DstAnalysisGraphColoring + ChaitinBriggsColoring
    ├─ performs pre-check via selected DstAnalysis->analyze(funcOp)
    └─ on success, uses DstAnalysisGraphColoring + chosen ColoringStrategy
       to allocate slices and rewrite operations

New Operations

d2m.release_dst - Marks end of DST value lifetime:

%dst = d2m.acquire_dst() : memref<1x!ttcore.tile<32x32, f32>, #dst>
// ... use dst ...
d2m.release_dst %dst : memref<1x!ttcore.tile<32x32, f32>, #dst>

Enables precise liveness analysis and early reuse. Verifiers ensure pairing with acquire_dst and prevent use-after-release.

Pass Implementation

D2MInsertDstRegisterGCPass (lib/Dialect/D2M/Transforms/InsertDstRegisterGC.cpp):

• Collects affine.load operations for DST-consuming ops (via OperandLoadStoreRegisterOpInterface)
• Builds interference graph using SSA liveness analysis
• Applies graph coloring to assign DST slices
• Generates L1↔DST data copy loops
• Inserts acquire_dst/release_dst pairs

Coloring Strategies

Two algorithms implemented in GraphColoringStrategy.{h,cpp}:

1. Chaitin-Briggs (default): Simplification-based coloring
2. Greedy: Sequential first-available-color assignment

Strategy selection via pass option:

ttmlir-opt --d2m-insert-dst-register-gc="coloring-strategy=greedy" input.mlir

Pipeline Integration

Not yet integrated into main pipeline. Current d2m-insert-dst-register-access performs multiple unrelated transformations (most notably linalg to affine conversion). These need refactoring into separate passes that would be applied before this pass.

Future Work

1. Keep accumulators in DST during reduction loops: For operations like matmul that accumulate results (e.g., C += A * B in a loop), the current pass loads the accumulator from L1 memory and writes it back on every iteration. This is inefficient. A better approach (used by the existing D2MInsertDSTAccess pass) keeps the accumulator in fast DST registers throughout the entire reduction loop and only writes back to L1 once at the end. This requires splitting the loop into three phases: (1) load initial values into DST, (2) compute with accumulator staying in DST, (3) write final results back to L1. This optimization could be implemented as a separate pass that runs before register allocation. (high priority)
2. Precise loop dependence analysis using MLIR affine utilities (medium priority)
3. Spill code generation for insufficient DST capacity (high priority)
4. PBQP-based allocation for cost modeling (may be overkill)
5. Move coalescing to eliminate redundant copies (medium priority)
6. Performance benchmarking on real workloads (medium priority)

[vmilosevic]

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

[codecov-commenter]

Codecov Report

❌ Patch coverage is 80.78406% with 299 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.77%. Comparing base (becbe0a) to head (8a9827c).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
lib/Dialect/D2M/Utils/TileMatmulUtils.cpp	18.77%	173 Missing ⚠️
lib/Dialect/D2M/Transforms/InsertDstRegisterGC.cpp	79.57%	106 Missing ⚠️
lib/Dialect/D2M/Analysis/DstAnalysisPass.cpp	78.12%	7 Missing ⚠️
lib/Dialect/D2M/IR/D2MOps.cpp	87.87%	4 Missing ⚠️
...b/Dialect/D2M/Transforms/GraphColoringStrategy.cpp	98.17%	3 Missing ⚠️
lib/Conversion/D2MToTTKernel/D2MToTTKernel.cpp	94.59%	2 Missing ⚠️
lib/Dialect/D2M/Analysis/DstAnalysisBasic.cpp	94.59%	2 Missing ⚠️
lib/Dialect/D2M/Analysis/DstCapacityAnalysis.cpp	95.65%	1 Missing ⚠️
...Dialect/D2M/Transforms/InsertDstRegisterAccess.cpp	97.36%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5879      +/-   ##
==========================================
+ Coverage   69.34%   69.77%   +0.42%     
==========================================
  Files         334      347      +13     
  Lines       50999    52483    +1484     
==========================================
+ Hits        35367    36621    +1254     
- Misses      15632    15862     +230     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[vmilosevic]

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

[vmilosevic]

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

[vmilosevic]

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

[vmilosevic]

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

[vmilosevic]

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

[vmilosevic]

⚠️ cppcoreguidelines-use-default-member-init ⚠️
use default member initializer for totalIterations

Suggested change

	LoopContext() : totalIterations(1) {}
	LoopContext() : {}