From d8fa273b0fce0bd6f43ebb35c4accf47b6234bb8 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Thu, 27 Aug 2026 01:14:43 +0300 Subject: [PATCH 1/2] mem_copy: the passThrough kernel is compiled for the wrong element type design.py types the line buffers bf16, but get_kernel_artifacts passes no extra_flags, so passThrough.cc's BIT_WIDTH is undefined, evaluates to 0, and the int32 branch compiles. The template moves one 64-byte vector per N elements -- 32 for int16, 16 for int32 -- so given a bf16 element count it runs width/16 iterations and writes width*4 bytes through a buffer holding width*2. Confirmed on the artifact: the object carried passThrough_aie, and the generated call passes lineWidth = 64 against memref<64xbf16>, i.e. 256 bytes through 128. mha/op.py is the same kernel done right. Setting the flag then exposes what it was hiding: at the correct N=32 mem_copy's smallest tile is two iterations, against the AIE_LOOP_MIN_ITERATION_COUNT(6) the loop declared, and the design hangs. The kernel cannot promise a trip count its caller supplies, so the assertion is removed rather than lowered. 67/67 on device, mem_copy plus mha. They also passed before, which is the point: the overrun is past the compared region. --- aie_kernels/generic/passThrough.cc | 4 +++- iron/operators/mem_copy/op.py | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/aie_kernels/generic/passThrough.cc b/aie_kernels/generic/passThrough.cc index f4a784de5..170d8b0df 100644 --- a/aie_kernels/generic/passThrough.cc +++ b/aie_kernels/generic/passThrough.cc @@ -20,7 +20,9 @@ passThrough_aie(T *restrict in, T *restrict out, const int32_t height, const int v64uint8 *restrict inPtr = (v64uint8 *)in; AIE_PREPARE_FOR_PIPELINING - AIE_LOOP_MIN_ITERATION_COUNT(6) + // No minimum trip count: the caller supplies all of height, width and N. The 6 that + // was asserted here is false for mem_copy's 64-element bf16 tile, which runs two + // iterations at N=32 and hangs on device. for (int j = 0; j < (height * width); j += N) // Nx samples per loop { *outPtr++ = *inPtr++; diff --git a/iron/operators/mem_copy/op.py b/iron/operators/mem_copy/op.py index 37058b61b..e2265ee1b 100644 --- a/iron/operators/mem_copy/op.py +++ b/iron/operators/mem_copy/op.py @@ -68,6 +68,9 @@ def get_kernel_artifacts(self): / "passThrough.cc" ) ], + # design.py types the line buffers bf16. Without this, passThrough.cc + # falls through to its int32 branch and copies twice the tile. + extra_flags=["-DBIT_WIDTH=16"], ) ] From cc7ce8aee488f252753da51a77fd14e323630350 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Fri, 4 Sep 2026 17:29:40 +0300 Subject: [PATCH 2/2] passThrough: specialize pipelined vs unhinted entry points --- aie_kernels/generic/passThrough.cc | 72 ++++++++++++++++++++++++------ iron/operators/mha/design.py | 5 ++- 2 files changed, 62 insertions(+), 15 deletions(-) diff --git a/aie_kernels/generic/passThrough.cc b/aie_kernels/generic/passThrough.cc index 170d8b0df..c78515018 100644 --- a/aie_kernels/generic/passThrough.cc +++ b/aie_kernels/generic/passThrough.cc @@ -10,7 +10,12 @@ #include #include -template +// Pipelined selects the loop's minimum-trip-count promise. The kernel cannot prove a +// trip count from height/width/N alone -- that has to come from the caller, since +// asserting >=6 for a call that runs fewer iterations hangs on device (mem_copy's +// 64-element bf16 tile runs two at N=32). Callers that can guarantee >=6 should use +// the Pipelined=true entry points below for the software-pipelined loop. +template __attribute__((noinline)) void passThrough_aie(T *restrict in, T *restrict out, const int32_t height, const int32_t width) { @@ -19,13 +24,19 @@ passThrough_aie(T *restrict in, T *restrict out, const int32_t height, const int v64uint8 *restrict outPtr = (v64uint8 *)out; v64uint8 *restrict inPtr = (v64uint8 *)in; - AIE_PREPARE_FOR_PIPELINING - // No minimum trip count: the caller supplies all of height, width and N. The 6 that - // was asserted here is false for mem_copy's 64-element bf16 tile, which runs two - // iterations at N=32 and hangs on device. - for (int j = 0; j < (height * width); j += N) // Nx samples per loop - { - *outPtr++ = *inPtr++; + if constexpr (Pipelined) { + AIE_PREPARE_FOR_PIPELINING + AIE_LOOP_MIN_ITERATION_COUNT(6) + for (int j = 0; j < (height * width); j += N) // Nx samples per loop + { + *outPtr++ = *inPtr++; + } + } else { + AIE_PREPARE_FOR_PIPELINING + for (int j = 0; j < (height * width); j += N) // Nx samples per loop + { + *outPtr++ = *inPtr++; + } } event1(); @@ -37,36 +48,69 @@ extern "C" { void passThroughLine(uint8_t *in, uint8_t *out, int32_t lineWidth) { - passThrough_aie(in, out, 1, lineWidth); + passThrough_aie(in, out, 1, lineWidth); } void passThroughTile(uint8_t *in, uint8_t *out, int32_t tileHeight, int32_t tileWidth) { - passThrough_aie(in, out, tileHeight, tileWidth); + passThrough_aie(in, out, tileHeight, tileWidth); +} + +// Trip count (height*width)/64 must be >=6, or this hangs on device. +void passThroughLinePipelined(uint8_t *in, uint8_t *out, int32_t lineWidth) +{ + passThrough_aie(in, out, 1, lineWidth); +} + +void passThroughTilePipelined(uint8_t *in, uint8_t *out, int32_t tileHeight, int32_t tileWidth) +{ + passThrough_aie(in, out, tileHeight, tileWidth); } #elif BIT_WIDTH == 16 void passThroughLine(int16_t *in, int16_t *out, int32_t lineWidth) { - passThrough_aie(in, out, 1, lineWidth); + passThrough_aie(in, out, 1, lineWidth); } void passThroughTile(int16_t *in, int16_t *out, int32_t tileHeight, int32_t tileWidth) { - passThrough_aie(in, out, tileHeight, tileWidth); + passThrough_aie(in, out, tileHeight, tileWidth); +} + +// Trip count (height*width)/32 must be >=6, or this hangs on device. +void passThroughLinePipelined(int16_t *in, int16_t *out, int32_t lineWidth) +{ + passThrough_aie(in, out, 1, lineWidth); +} + +void passThroughTilePipelined(int16_t *in, int16_t *out, int32_t tileHeight, int32_t tileWidth) +{ + passThrough_aie(in, out, tileHeight, tileWidth); } #else // 32 void passThroughLine(int32_t *in, int32_t *out, int32_t lineWidth) { - passThrough_aie(in, out, 1, lineWidth); + passThrough_aie(in, out, 1, lineWidth); } void passThroughTile(int32_t *in, int32_t *out, int32_t tileHeight, int32_t tileWidth) { - passThrough_aie(in, out, tileHeight, tileWidth); + passThrough_aie(in, out, tileHeight, tileWidth); +} + +// Trip count (height*width)/16 must be >=6, or this hangs on device. +void passThroughLinePipelined(int32_t *in, int32_t *out, int32_t lineWidth) +{ + passThrough_aie(in, out, 1, lineWidth); +} + +void passThroughTilePipelined(int32_t *in, int32_t *out, int32_t tileHeight, int32_t tileWidth) +{ + passThrough_aie(in, out, tileHeight, tileWidth); } #endif diff --git a/iron/operators/mha/design.py b/iron/operators/mha/design.py index f17d26fad..c213d4569 100644 --- a/iron/operators/mha/design.py +++ b/iron/operators/mha/design.py @@ -215,8 +215,11 @@ def fused_mha( func_type = "" if vectorized else "_scalar" zero_kernel = Kernel(f"zero_{dtype_str}", "mha.o", [qk_ty]) + # Pipelined entry point: 4*B_q/32 trips is 8 at the B_q=64 this operator is fixed + # to (op.py), safely >=6. Callers that lower B_q below 48 must switch back to + # passThroughLine. memcopy_kernel_scale = Kernel( - f"passThroughLine", "mha_passThrough.o", [s_ty, s_ty, np.int32] + f"passThroughLinePipelined", "mha_passThrough.o", [s_ty, s_ty, np.int32] ) scale_buffer_init_kernel = Kernel("init_scale_buffer", "mha.o", [s_ty, np.int32])