[VMI] 256-lane E2B group broadcast: layout assignment downgrades the deinterleaved direct row to contiguous, landing on vsldb + 2x vselr (device 507035) instead of one E2B vlds

[liuzidi]

Summary

A 256-lane bf16 grouped broadcast (dist_mode="brc", group = 8, i.e. period-32
broadcast over 32 elements/group → groupSize = 256/8 = 32 elements = gb(2))
has a native E2B direct row in the layout table, but layout assignment
deliberately downgrades it, so the load ends up on the group_slot fallback:

vsldb + 2x vselr

instead of the one instruction the hardware has:

vlds {dist = "E2B_B16"}   // one packet, reused by both physical parts

vsldb requires a 32B-aligned group base and reads 256B (16x the semantic
16B footprint). With the natural packed scale layout (8 contiguous bf16 scales
per chunk, so odd rows start 16B off a 32B boundary) the emitted code passes
lowering and then traps on device with 507035.

The only thing separating the two lowerings is a single if in
getPreferredGroupBroadcastLoadLayout().

Environment

Ascend950DT_9592 (A5), CANN 9.2.0, dav-c310-vec, 72 vector cores.
PTOAS 9b0b7697f; the same behaviour reproduces on 566d6af8c (= vmi-v0.1.8)
and on MR152 (cbb478460) — the two lowerings are identical there.

Minimal repro

pto-test-opt repro.pto -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto

// A. contiguous result declared -> no E2B direct row for that layout
func.func @e2b_256_bf16_contiguous(
    %src: !pto.ptr<bf16, ub>, %off: index) -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout<contiguous>> {
  %c1 = arith.constant 1 : index
  %c8 = arith.constant 8 : index
  %aligned_off = arith.muli %off, %c8 : index
  %out = pto.vmi.vload %src[%aligned_off], %c1 {dist_mode = "brc", group = 8}
      : !pto.ptr<bf16, ub> -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout<contiguous>>
  return %out : !pto.vmi.vreg<256xbf16, #pto.vmi.layout<contiguous>>
}

// B. the deinterleaved result the direct row asks for
func.func @e2b_256_bf16_deint2(
    %src: !pto.ptr<bf16, ub>, %off: index) -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout<deinterleaved = 2>> {
  %c1 = arith.constant 1 : index
  %c8 = arith.constant 8 : index
  %aligned_off = arith.muli %off, %c8 : index
  %out = pto.vmi.vload %src[%aligned_off], %c1 {dist_mode = "brc", group = 8}
      : !pto.ptr<bf16, ub> -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout<deinterleaved = 2>>
  return %out : !pto.vmi.vreg<256xbf16, #pto.vmi.layout<deinterleaved = 2>>
}

// C. no layout declared: broadcast feeding an elementwise vmul + store
func.func @e2b_256_bf16_solver_choice(
    %src: !pto.ptr<bf16, ub>, %x: !pto.ptr<bf16, ub>, %off: index, %xoff: index,
    %dst: !pto.ptr<bf16, ub>, %doff: index, %mask: !pto.vmi.mask<256xb16>) {
  %c1 = arith.constant 1 : index
  %c8 = arith.constant 8 : index
  %aligned_off = arith.muli %off, %c8 : index
  %scale = pto.vmi.vload %src[%aligned_off], %c1 {dist_mode = "brc", group = 8}
      : !pto.ptr<bf16, ub> -> !pto.vmi.vreg<256xbf16>
  %aligned_xoff = arith.muli %xoff, %c8 : index
  %value = pto.vmi.vload %x[%aligned_xoff]
      : !pto.ptr<bf16, ub> -> !pto.vmi.vreg<256xbf16>
  %prod = pto.vmi.vmul %value, %scale, %mask
      : !pto.vmi.vreg<256xbf16>, !pto.vmi.vreg<256xbf16>, !pto.vmi.mask<256xb16> -> !pto.vmi.vreg<256xbf16>
  %aligned_doff = arith.muli %doff, %c8 : index
  pto.vmi.vstore %prod, %dst[%aligned_doff]
      : !pto.vmi.vreg<256xbf16>, !pto.ptr<bf16, ub>
  return
}

Observed lowering

A — contiguous (as declared, and as the solver picks for C):

%result = pto.vsldb %2, %c0_i16, %c0_i16, %1 : !pto.ptr<bf16, ub>, i16, i16, !pto.mask<b16> -> !pto.vreg<128xbf16>
%6 = pto.vselr %result, %5 : !pto.vreg<128xbf16>, !pto.vreg<128xui16> -> !pto.vreg<128xbf16>
%8 = pto.vselr %result, %7 : !pto.vreg<128xbf16>, !pto.vreg<128xui16> -> !pto.vreg<128xbf16>
return %6, %8 : !pto.vreg<128xbf16>, !pto.vreg<128xbf16>

B — deinterleaved = 2 declared (this is the whole difference):

%result = pto.vlds %arg0[%0] {dist = "E2B_B16"} : !pto.ptr<bf16, ub> -> !pto.vreg<128xbf16>
return %result, %result : !pto.vreg<128xbf16>, !pto.vreg<128xbf16>   // one packet, both parts

C — no layout declared: the solver repeats A, then pays for the mismatch at the
store (note there is no d(2) in the consumer chain, so the assigned broadcast
result is contiguous):

%result = pto.vsldb %2, %c0_i16, %c0_i16, %1 : !pto.ptr<bf16, ub>, i16, i16, !pto.mask<b16> -> !pto.vreg<128xbf16>
%6 = pto.vselr %result, %5 : ...
%8 = pto.vselr %result, %7 : ...
%13 = pto.vmul %result_0, %6, %12 : ...
%15 = pto.vmul %result_1, %8, %14 : ...
pto.vstus ... %13 ...    // + pto.vstus ... %15 ...

For reference, the AscendC form of the same operation is exactly one E2B packet
plus the broadcast split, with no slot load and no register gather:

1x vlds_brc_elem2datablock (E2B_B16)    // 16B group base alignment is enough
1x vintlv
2x vlds UNPK4
2x vmul
2x vsts NORM_B16

Root cause

kGroupBroadcastLoadDirectPatterns (lib/PTO/Transforms/VMILayoutSupportTables.inc:510)
already carries the row for this shape — groupSize = 32 elements is gb(2):

{VMIGroupBroadcastLoadDirectKind::E2B, G<8>(), gb(2), bits<16, 32>(),
 memContiguous(), d(2)},                                    // line 513

But getPreferredGroupBroadcastLoadLayout()
(lib/PTO/Transforms/VMILayoutAssignment.cpp:403, introduced by
0d7da8454 "vmi: narrow FP4 layout and physical part lowering", 2026-09-07)
throws that layout away whenever the result is unassigned:

VMILayoutAttr directLayout = fact->layout.resultLayout;
// A direct E2B packet fills exactly one physical part. When the
// contiguous form of this broadcast spans multiple physical chunks, the
// direct table can only offer a deinterleaved split layout. Prefer the
// generic contiguous lowering (group_slots -> contiguous) so consumers
// such as plain elementwise vmul can stay contiguous and avoid
// vldsx2/vintlv-style deinterleave materialization.
if (fact->kind == VMIGroupBroadcastLoadDirectKind::E2B &&
    directLayout.isDeinterleaved()) {            // line 424/425
  ...
  return contiguous;                             // no E2B direct row exists for c()
}
return directLayout;

c() has no E2B row (only gb(1)→c(), gb(2)→d(2), gb(4)→d(4) exist), so the
value falls through to the generic group_broadcast → gs(8) layout table, and
the lowering takes lowerGroupSlotFallback
(VMIToVPTOPatternInternals3.cpp:1680) instead of lowerDirectE2B (:1545).

The comment states the trade-off honestly (keep vmul contiguous), but the
comparison is one-sided: it measures the cost of the d(2) path
(vdintlv/vintlv materialization) and ignores the cost of the contiguous path
(vsldb + vselr, plus a 32B base-alignment requirement and a 256B read
footprint for a 16B semantic load). On this shape the contiguous path is the
more expensive one, on every axis.

Proposal

The d(2) row is not the problem — the downgrade is. Two possible directions:

1. Make the choice cost-aware instead of unconditional: keep directLayout
   when the E2B row exists and the consumer chain can absorb d(2), and only
   fall back to contiguous when the cost actually favors it. At minimum, the
   vsldb path's 32B alignment requirement (and its 256B read envelope) should
   count against it — a slot path that is known-misaligned today is a device
   trap, not a slower-but-valid alternative.
2. Let d(2) reach the store. With directLayout kept, what remains is a
   vdintlv/vintlv pair around the elementwise consumer (see the local
   experiment below). If elementwise ops could stay in d(2) end-to-end and
   only the store did the interleave, the 256-lane form would generate exactly
   the AscendC instruction sequence. This looks like the missing piece — the
   lowering exists, the layout negotiation policy is what stops it.

A local 5-line experiment (gate the downgrade behind an env var, nothing else
changed) confirms the direction. Same .pto input:

	generated broadcast load	on-device
default	vsldb + 2x vselr	507035 (traps)
downgrade gated off	1x vlds {dist="E2B_B16"} + vdintlv + 2x vmul + vintlv	maxdiff 0

     if (fact->kind == VMIGroupBroadcastLoadDirectKind::E2B &&
-        directLayout.isDeinterleaved()) {
+        directLayout.isDeinterleaved() &&
+        std::getenv("PTOAS_KEEP_E2B_D2") == nullptr) {

Measured on a real kernel (fp4 e2m1 → bf16 dequant, m=512 h=512 bk=512, 72
cores, msprof warmup=0 rep=20 x 3 rounds, medians of two devices):

variant	device 0	device 1	vs ASC
AscendC reference	2.5256 us	2.5521 us	1.000
128-lane E2B form (two carriers)	2.5474 us	2.5338 us	+0.9% / +0.7%
256-lane ideal form + PTOAS_KEEP_E2B_D2=1	2.6927 us	2.6737 us	+6.6% / +4.8%

(I am not proposing this env-var gate as a patch — global removal of the
downgrade would slow down kernels whose consumers genuinely want contiguous,
which is what the comment is protecting. It is reported as an experiment to
characterize the gap. The remaining +5–7% is exactly the vdintlv/vintlv pair
that a "keep d(2) through the consumer" policy would remove.)

Note: the default path's trap is also reachable without any d(2)
consideration — a group-slot slots=8 load with a 16B-only-aligned base emits
vsldb and traps. I filed a related guard for that case previously (see
#1502 for a sibling problem in the same lowering family); the request here is
about the policy that sends this shape there in the first place.

Happy to contribute the repro as a lit test or to try the cost-aware variant if
that direction is acceptable.

[liuzidi]

End-to-end repro, in two parts. Part 1 needs no hardware (just pto-test-opt)
and shows the lowering difference directly. Part 2 is the real tilelang kernel on
one A5 board, and shows the 507035.

All files referenced below are attached to this comment / available on request.
SHA256 is listed so you can confirm you have the same bytes.

---

Part 1 — compiler-only repro (no board)

1.1 Input

repro.pto (sha256 d1b5716d…827ee) — three functions, only the declared layout varies:

// A. result declared contiguous -> no E2B direct row for that layout
func.func @e2b_256_bf16_contiguous(%src: !pto.ptr<bf16, ub>, %off: index)
    -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout<contiguous>> {
  %c1 = arith.constant 1 : index
  %c8 = arith.constant 8 : index
  %aligned_off = arith.muli %off, %c8 : index
  %out = pto.vmi.vload %src[%aligned_off], %c1 {dist_mode = "brc", group = 8}
      : !pto.ptr<bf16, ub> -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout<contiguous>>
  return %out : !pto.vmi.vreg<256xbf16, #pto.vmi.layout<contiguous>>
}

// B. the deinterleaved result the direct row asks for
func.func @e2b_256_bf16_deint2(%src: !pto.ptr<bf16, ub>, %off: index)
    -> !pto.vmi.vreg<256xbf16, #pto.vmi.layout<deinterleaved = 2>> { ... same load, d(2) declared ... }

// C. no layout declared: broadcast -> vmul -> vstore, the real kernel shape
func.func @e2b_256_bf16_solver_choice(%src, %x, %off, %xoff, %dst, %doff, %mask) { ... }

1.2 Run

pto-test-opt repro.pto \
  -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment \
  -vmi-layout-assignment -vmi-to-vpto \
  -o repro_default.mlir                 # sha256 95a69c14…7573f

1.3 Observed

• A (contiguous declared) and C (solver-chosen, which also picks contiguous)
  → vsldb + 2× vselr:

%1 = pto.pset_b16 "PAT_VL8" : !pto.mask<b16>
%2 = pto.addptr %arg0, %0 : <bf16, ub> -> <bf16, ub>
%result = pto.vsldb %2, %c0_i16, %c0_i16, %1 : !pto.ptr<bf16, ub>, i16, i16, !pto.mask<b16> -> !pto.vreg<128xbf16>
%4 = pto.vci %c0_i16 : i16 -> !pto.vreg<128xui16>
%5 = pto.vshrs %4, %c5_i16, %3 : ...
%6 = pto.vselr %result, %5 : !pto.vreg<128xbf16>, !pto.vreg<128xui16> -> !pto.vreg<128xbf16>
%7 = pto.vadds %5, %c4_i16, %3 : ...
%8 = pto.vselr %result, %7 : !pto.vreg<128xbf16>, !pto.vreg<128xui16> -> !pto.vreg<128xbf16>

• B (d(2) explicitly declared) → the one-instruction form:

%0 = arith.muli %arg2, %c8 : index
%result = pto.vlds %arg0[%0] {dist = "E2B_B16"} : !pto.ptr<bf16, ub> -> !pto.vreg<128xbf16>

So the native direct row is reachable — the solver just never chooses it for the
shapes A and C produce.

1.4 The 5-line experiment (d2_layout_patch.diff, sha256 aab68833…f7a05d)

The only thing separating A/C from B is the guard in
getPreferredGroupBroadcastLoadLayout() (lib/PTO/Transforms/VMILayoutAssignment.cpp:424),
which downgrades the deinterleaved direct layout to contiguous.
Gating that if off for this experiment:

@@ -18,6 +18,8 @@
 #include "PTO/Transforms/VMILayoutSupport.h"
 
+#include <cstdlib>
+
@@ -422,7 +424,8 @@ struct LayoutSolver {
     if (fact->kind == VMIGroupBroadcastLoadDirectKind::E2B &&
-        directLayout.isDeinterleaved()) {
+        directLayout.isDeinterleaved() &&
+        std::getenv("PTOAS_KEEP_E2B_D2") == nullptr) {
       VMILayoutAttr contiguous = getContiguousLayout();

Build (PTOAS 9b0b7697f, LLVM 19 from /opt/ptoas/llvm):

source /opt/ptoas/llvm/env-19.sh
cmake -G Ninja -S third_party/ptoas -B third_party/ptoas/build-d2 \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_DIR=/opt/ptoas/llvm/current-19/lib/cmake/llvm \
  -DMLIR_DIR=/opt/ptoas/llvm/current-19/lib/cmake/mlir \
  -DPython_EXECUTABLE=$(which python) \
  -DCMAKE_CXX_FLAGS="-D_FORTIFY_SOURCE=2 -fstack-protector-strong -ftrapv -Wno-error=restrict" \
  -DPTOAS_ENABLE_WERROR=OFF
ninja -C third_party/ptoas/build-d2 -j 100

1.5 Re-run with the toggle

PTOAS_KEEP_E2B_D2=1 pto-test-opt repro.pto \
  -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment \
  -vmi-layout-assignment -vmi-to-vpto \
  -o repro_keepd2.mlir                  # sha256 9815a66b…5c183b

Function C now reads:

%result  = pto.vlds %arg0[%0] {dist = "E2B_B16"} : !pto.ptr<bf16, ub> -> !pto.vreg<128xbf16>   // ONE packet, no vselr
%2       = pto.addptr %arg1, %1 : <bf16, ub> -> <bf16, ub>
%result_0, ... = pto.vldus %2, %3, %c128 : ...          // first 128 lanes
%result_1, ... = pto.vldus %updated_base, %updated_align, %c128 : ...   // second 128 lanes
%low, %high = pto.vdintlv %result_0, %result_1 : ... -> ...
%5 = pto.vmul %low,  %result, %4 : ...
%7 = pto.vmul %high, %result, %6 : ...
%low_4, %high_5 = pto.vintlv %5, %7 : ... -> ...
pto.vstus ... x2

i.e. 1× vlds {dist="E2B_B16"} (the same packet feeds both physical parts) + one
vdintlv + 2× vmul + one vintlv, instead of vsldb + 2× vselr.
The only remaining cost is the vdintlv/vintlv pair at the elementwise boundary.

(The PTOAS_KEEP_E2B_D2 env var here is only an experiment toggle so both paths
can be produced from one binary; I am not proposing it as a patch.)

---

Part 2 — end-to-end tilelang kernel on the board

2.1 Environment

• Ascend950DT_9592 (A5), CANN 9.2.0, dav-c310-vec, 72 vector cores
• PTOAS 9b0b7697f (same result on 566d6af8c / vmi-v0.1.8; MR152 cbb478460 is also identical here)
• tilelang with the VMI frontend, torch_npu

2.2 The kernel

e2b_256_e2e_repro.py (sha256 2d5f7150…b4ef5a) — standalone, ~150 lines, no other
repo needed. The whole compute is:

with T.SimdVF():
    mask256 = V.create_mask(256, size=256)
    for row in T.serial(block_m):                 # block_m = 56
        for chunk in T.serial(block_k // 256):    # block_k = 512 -> 2 chunks
            col = chunk * 256
            scale256 = V.vload(sf_ub[row, chunk * 8], size=256,
                               stride=1, dist_mode='brc', group=8)
            x256 = V.vcvt(V.vload(x_ub[row, col], size=256), 'bfloat16')
            V.vstore(V.vmul(x256, scale256, mask256), out_ub[row, col], mask256)

Geometry: num_tokens=512, hidden=512, block_k=512, block_m=56, e2m1 → bf16,
72 cores, 3 stages, persistent scheduling. sf_ub is (56, 16) bf16 = 8 packed
scales per 256-column chunk, row stride 32 B.

The input staging is what makes the alignment fault deterministic: 8 bf16 scales =
16 B per chunk, so chunk 1 of every row starts 16 B off a 32 B boundary, while
vsldb needs a 32 B-aligned group base.

2.3 Run — default compiler

python e2b_256_e2e_repro.py --build            # -> BUILD_OK
python e2b_256_e2e_repro.py --run --device 0

Observed:

BUILD_OK
LAUNCH_OK
RuntimeError: npuSynchronizeDevice:.../NPUStream.cpp:798 NPU function error:
              device error type 3, error code is 507035

i.e. it compiles, launches, and then traps on the vector core at synchronize
(ACL_ERROR_RT_VECTOR_CORE_EXCEPTION). This matches the lowering in Part 1:
source/dst/offset are all UB with the compact scale layout, so the vsldb
group base is only 16 B aligned.

Corresponding emitted VPTO (ptoas kernel.pto --pto-backend=vpto --pto-level=level3 --emit-vpto):

%result   = pto.vsldb %96, %c0_i16, %c0_i16, %85 : !pto.ptr<bf16, ub>, i16, i16, !pto.mask<b16> -> !pto.vreg<128xbf16>
%result_0 = pto.vlds %38[%98] {dist = "UNPK4"} : ...
%result_1 = pto.vlds %38[%99] {dist = "UNPK4"} : ...
%105 = pto.vselr %result, %87 : ...
%106 = pto.vselr %result, %88 : ...

2.4 Run — with the toggle (same kernel source, no change)

PTOAS_KEEP_E2B_D2=1 ptoas kernel.pto --pto-backend=vpto --pto-level=level3 --emit-vpto

%result_1 = pto.vlds %60[%94] {dist = "E2B_B16"} : !pto.ptr<bf16, ub> -> !pto.vreg<128xbf16>   // no vsldb, no vselr
%low, %high = pto.vdintlv %102, %103 : ...
%low_2, %high_3 = pto.vintlv %104, %105 : ...

On device:

BUILD_OK
LAUNCH_OK
SYNC_OK
maxdiff 0.0 mismatch 0

2.5 Numbers

msprof device time, median of 3 rounds (do_bench warmup=0 rep=20),
M=512 H=512 e2m1→bf16, 72 cores:

variant	device 0	device 1
256-lane E2B, default lowering	507035	507035
256-lane E2B, deinterleaved row kept	2.6927 µs	2.6660 / 2.6737 / 2.6740 µs
128-lane E2B (equivalent ASC-shaped spelling, works today)	2.5474 µs	2.5338 µs
ASC reference (S.*)	2.5256 µs	2.5521 µs

So the shape is not a performance question: the default path faults, and once the
direct row is kept it is ~6% off ASC with a correct result.

---

Files

file	sha256	what
repro.pto	d1b5716d…827ee	Part 1 input, 3 variants
repro_default.mlir	95a69c14…7573f	Part 1 output, default
repro_keepd2.mlir	9815a66b…5c183b	Part 1 output, toggle on
d2_layout_patch.diff	aab68833…f7a05d	the 5-line experiment
e2b_256_e2e_repro.py	2d5f7150…b4ef5a	Part 2 standalone kernel
build_host_only.sh	4f3d6a23…79a7cb	env preamble for Part 2
cast_back_vmi_ideal256.py	bd149ef8…e5117	the same shape inside a full production-shaped cast-back kernel