diff --git a/aie_kernels/generic/passThrough.cc b/aie_kernels/generic/passThrough.cc index f4a784de5..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,11 +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 - AIE_LOOP_MIN_ITERATION_COUNT(6) - 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(); @@ -35,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/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"], ) ] 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])