Fix hierarchical reduced predicates in copy layout lowering

[HydraQYH]

Motivation

Fix copy layout lowering for reduced predicates when ATOM_REST has a hierarchical shape.

Technical Details

A reduced copy predicate contains one boolean value per copy atom, so its shape intentionally omits ATOM_V. The source and predicate shapes follow this contract:

source:    ((ATOM_V, ATOM_REST), REST...)
predicate: (ATOM_REST, REST...)

The previous lowering detected reduced predicates using:

if (pred && predLayoutAttr.rank() == srcRank - 1) {
  // Prepend a unit mode.
}

This assumes that removing ATOM_V always reduces the rank by one. That is not true when ATOM_REST is hierarchical.

For example:

• source V-mode: (1, (2, 2))
• predicate V-mode: (2, 2)

Both shapes have rank 2, so the old condition does not recognize the predicate as reduced.

The subsequent recursive copy expansion slices the predicate according to the source hierarchy and eventually generates an invalid operation:

vector.extract_strided_slice offset=[4], size=[2] from vector<4xi1>

Changes

• derives the expected reduced predicate shape from the source shape structure;
• normalizes (ATOM_REST, REST...) to ((1, ATOM_REST), REST...);
• removes the rank-based reduced-predicate heuristic;
• adds a regression test covering hierarchical ATOM_REST=(2,2).

Minimal reproduction

Starting from examples/01-vectorAdd.py, make the following changes:

-    copy_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Float32)
+    copy_atom = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32)

Apply the same copy-atom change in the JIT function and change the value layout:

-    copy_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Float32)
+    copy_atom = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32)

     tiled_copy = fx.make_tiled_copy_tv(
         copy_atom,
         fx.make_ordered_layout((8, 16), order=(1, 0)),
-        fx.make_ordered_layout((1, 4), order=(0, 1)),
+        fx.make_ordered_layout((2, 2), order=(1, 0)),
     )

The existing predicate construction remains unchanged:

thr_cC = thr_copy.partition_S(cC)[(0, None), None, None]
thr_pC = fx.make_fragment_like(thr_cC, dtype=fx.Boolean)

for i in fx.range_constexpr(fx.size(thr_pC.shape).unpack()):
    thr_pC[i] = fx.elem_less(thr_cC[i], (M, N))

fx.copy(copy_atom, thr_gA, thr_rA, pred=thr_pC)

This produces:

• source shape: ((1,(2,2)),1,1)
• predicate shape: ((2,2),1,1)

Before this fix, compiling the example fails with the out-of-bounds vector.extract_strided_slice described above.

Fix

The lowering now derives the expected reduced predicate shape structurally:

((ATOM_V, ATOM_REST), REST...)
↓
(ATOM_REST, REST...)

When the actual predicate shape matches this structure, its layout is normalized before recursive copy expansion:

(ATOM_REST, REST...)
-> (1, ATOM_REST, REST...)
-> ((1, ATOM_REST), REST...)

The inserted unit mode uses zero stride:

1:0

This preserves one predicate per copy atom without allocating or copying predicate data. It only aligns the predicate layout hierarchy with the source hierarchy.

Test Plan

• Full FlyDSL C++ and Python build completed successfully.
• New hierarchical reduced-predicate layout-lowering test passed.
• New full SSA-lowering regression test passed.
• Existing expand_copy_reduced_pred.mlir test passed.

Test Result

The reproducer ran successfully on gfx950/MI355X and reported: PASS

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.