[Graph Optimization](revert) restore layout optimization (#1467)

[zhanwei33]

Summary

• Revert [Graph Optimization](revert) revert layout optimization (#1267) #1467 to restore [Graph Optimization](feat) Layout Optimization for main-dev #1267's GraphOptimize/layout-memory implementation on main-dev.
• Resolve post-revert conflicts while preserving the merge-concat-load-buffer pass from [TritonToGraph](feat) Add merge-concat-load-buffer pass #1494, the DiscreteMaskAccessConversion dependency from [AtomicOp](fix) Fuse expanded floating-point atomic min/max #1460, and the current pytest backend marker registration.

Validation

• Targeted pre-commit hooks: whitespace, EOF, Python AST, merge conflicts, private keys, large files, ruff, yapf, and clang-format.
• Parsed all changed, non-deleted Python sources.
• Diff whitespace check passed for all paths other than third_party/ascend/patch/, which the repository's pre-commit configuration explicitly excludes because it contains unified-diff context lines.

Not run

• Full CMake build and device tests: the worktree's two submodules are uninitialized and the host has CMake 3.27.1.

[zhanwei33]

P0 — reject this nested loop-carried pointer shape before enabling ConvertModuloToMask

This PR restores ConvertModuloToMask, but its forward use walk currently accepts an address through an unbounded number of scf.for init-arg → region-iter-arg and scf.yield → loop-result relays (collectAddressedLoads). That admits a deterministic mm_int8int2 precision regression with the default graph rule mask (511).

The important distinction is that the modulo-to-boundary-mask rewrite is algebraically valid in the TTIR below. The first semantic corruption is later, when the nested pointer state is materialized for the loop lowering. We should still fence this shape in this PR: without it, enabling the rule makes every output of the repro wrong.

Reproduction and observed boundary

Fixed configuration: M=2, N=1024, K=4096, BLOCK_M=64, BLOCK_N=128, BLOCK_K=64, fixed seed. With dynamic-CV explicitly disabled:

graph optimization	result
off	exact
on	2048 / 2048 outputs mismatch; max absolute difference 22525

Turning dynamic-CV on does not change the failure. Debug + overflow sanitization happens to be exact only because the extra use-def graph prevents this candidate from matching; debug without sanitization mismatches again. This isolates GraphOptimize rather than autotune, parallel compilation, or dynamic-CV.

Why this candidate matches

For the A coordinate, the rule sees all of the conditions in analyzeModulo:

• %offs_am_29 is a rank-1 tensor<64xi32> remainder with a tt.make_range(0, 64)-derived dividend.
• The divisor is tt.splat %M: it is runtime/dynamic, not a compile-time constant, so the constant-divisor exclusion does not apply.
• The unwrapped offset is also used by the final C-store guard, offs_am_28 < M.
• The modulo result has one tt.expand_dims axis = 1.
• Its address-only use chain is entirely allowed by collectAddressedLoads: shape/arithmetic ops, outer scf.for, inner scf.for, and finally tt.load.

The B coordinate independently matches the same pattern (tensor<128xi32>, dynamic N, axis = 0). The problematic A path is the one that crosses two nested loop-carried pointer relays.

TTIR before the rule

%offs_am_28 = arith.addi %offs_am_27, %offs_am : tensor<64xi32>
%offs_am_29 = arith.remsi %offs_am_28, %offs_am_9 : tensor<64xi32>
// %offs_am_9 = tt.splat %M

%offs_bn_32 = arith.addi %offs_bn_31, %offs_bn : tensor<128xi32>
%offs_bn_33 = arith.remsi %offs_bn_32, %offs_bn_10 : tensor<128xi32>
// %offs_bn_10 = tt.splat %N

%a_ptrs_34 = tt.expand_dims %offs_am_29 {axis = 1 : i32}
%a_ptrs_35 = arith.muli %a_ptrs_34, %a_ptrs       // %a_ptrs = splat(stride_am)
%a_ptrs_38 = tt.addptr %a_ptrs_13, %a_ptrs_37
%b_ptrs_39 = tt.expand_dims %offs_bn_33 {axis = 0 : i32}
%b_ptrs_42 = tt.addptr %b_ptrs_17, %b_ptrs_41

%outer:2 = scf.for ... iter_args(%a_outer = %a_ptrs_38, %acc = ...) {
  %inner:3 = scf.for ... iter_args(%a_inner = %a_outer, %b_inner = %b_ptrs_42, %acc2 = %acc) {
    ...
    %a_next = tt.addptr %a_inner, ...
    %b_next = tt.addptr %b_inner, ...
    scf.yield %a_next, %b_next, %acc_next
  }
  scf.yield %inner#0, %inner#2
}

%c_ptrs_43 = tt.expand_dims %offs_am_28 {axis = 1 : i32}
%c_ptrs_46 = tt.expand_dims %offs_bn_32 {axis = 0 : i32}
%c_mask_m = arith.cmpi slt, %c_ptrs_43, %c_mask    // raw row < M
%c_mask_n = arith.cmpi slt, %c_ptrs_46, %c_mask_22 // raw col < N
tt.store %c_ptrs_49, %outer#1, (broadcast(%c_mask_m) & broadcast(%c_mask_n))

TTIR after the rule

%offs_am_29 = arith.cmpi slt, %offs_am_28, %offs_am_9 : tensor<64xi1>
%offs_bn_33 = arith.cmpi slt, %offs_bn_32, %offs_bn_10 : tensor<128xi1>

%a_ptrs_34 = tt.expand_dims %offs_am_28 {axis = 1 : i32} // raw, no remsi
...
%b_ptrs_39 = tt.expand_dims %offs_bn_32 {axis = 0 : i32} // raw, no remsi
...

%a_k_mask = ...                                           // original K mask
%a_m_mask = tt.broadcast (tt.expand_dims %offs_am_29 {axis = 1 : i32})
%a_load_mask = arith.andi %a_k_mask, %a_m_mask
%a = tt.load %a_inner, %a_load_mask, %zero

%b_k_mask = ...
%b_n_mask = tt.broadcast (tt.expand_dims %offs_bn_33 {axis = 0 : i32})
%b_load_mask = arith.andi %b_k_mask, %b_n_mask
%b = tt.load %b_inner, %b_load_mask, %zero

For the first M tile, the exact lane values make the intended equivalence clear:

• r = offs_am_28 = [0, 1, 2, ..., 63]
• before: r % M = r % 2 = [0, 1, 0, 1, ..., 0, 1]
• after boundary mask: r < 2 = [true, true, false, ..., false] (2 true lanes, 62 false lanes)

Thus, before the rewrite, an A element at inner-K lane k is addressed as

A_base + (r % 2) * stride_am + k

and after it as

A_base + r * stride_am + k, guarded by (r < 2)

For rows 0 and 1, the addresses and mask remain exactly the same. For lanes 2..63, the old version reads wrapped rows 0/1 but their C stores are already masked out; the new version reads zero and the C stores are still masked out. Therefore this TTIR change is not itself the arithmetic bug.

For N=1024 and BLOCK_N=128, each of the eight N tiles has q * 128 + [0..127] < 1024, so the B boundary mask is all true for this particular shape. It is still syntactically dynamic and therefore matches the rule, but it does not explain the mismatch.

Where the pointer state is actually lost

The graph-on final TTAdapter starts the outer loop with the correct dynamic A offset and row stride, then uses them in the inner reinterpret cast:

%a_offset0 = arith.muli %offs_am, %stride_am : index
%outer:3 = scf.for ... iter_args(
    %acc = ...,
    %a_offset = %a_offset0,
    %a_row_stride = %stride_am) {
  %inner:3 = scf.for ... iter_args(
      %acc2 = %acc, %inner_a_offset = %a_offset, %b_offset = %c0) {
    %a_view = memref.reinterpret_cast %a_ptr
      to offset: [%inner_a_offset], sizes: [64, 64],
         strides: [%a_row_stride, %c1]
    ...
    %next_a_offset = arith.addi %inner_a_offset, %c64
    scf.yield %next_acc, %next_a_offset, %next_b_offset
  }
  scf.yield %inner#0, %c0, %c1
}

The final outer scf.yield should forward the inner A offset and preserve the dynamic row stride, i.e. conceptually scf.yield %inner#0, %inner#1, %a_row_stride. Instead it emits literal 0 and 1.

Each inner loop advances the A offset by 16 * 64 = 1024 bytes. The correct A address for outer iteration i, inner iteration j, row lane r, and K lane c is:

A_base + r * stride_am + (i * 16 + j) * 64 + c

After the first outer iteration, the lowered IR resets the next outer iteration to offset=0, stride=1 instead of carrying offset += 1024 and the original stride_am. Consequently outer iterations 1..3 reuse/interpret the wrong A region, which corrupts all 2 * 1024 = 2048 valid C elements.

This is the loop-state path reached by LoopConverter<scf::ForOp> / BlockDataParser::rewriteLoopOp; the latter's terminator materialization is exactly where constant offset/stride attributes become 0 / 1. The rule must not introduce this unsupported nested carried-pointer shape until that lowering contract is fixed.

Requested P0 fix

Please make collectAddressedLoads track loop-carry depth and reject the second scf.for init-arg → region-iter-arg transition. In pseudocode:

bool collectAddressedLoads(..., unsigned loopCarryHops = 0) {
  ...
  if (auto forOp = dyn_cast<scf::ForOp>(user)) {
    if (loopCarryHops >= 1)
      return false;
    return collectAddressedLoads(forOp.getRegionIterArg(index), ..., loads,
                                 loopCarryHops + 1);
  }
  if (auto yield = dyn_cast<scf::YieldOp>(user))
    return collectAddressedLoads(forOp.getResult(index), ..., loads,
                                 loopCarryHops);
}

This is deliberately candidate-local: it retains the A remsi for this unsafe two-level relay, while a one-loop candidate such as B can still be optimized. Please add (1) a positive one-loop rule test, (2) a negative nested-loop A-pointer test that keeps the remainder under mask 256/511, and (3) this M=2, N=1024, K=4096 numerical regression.

[github-actions]

✅ DCO Check Passed

Commits

Identity	Commits
zhanwei33 <22…@users.noreply.github.com>	2 (signed off)