mem_copy: the passThrough kernel is compiled for the wrong element type

aie_kernels/generic/passThrough.cc

@@ -10,7 +10,12 @@
 #include <stdint.h>
 #include <stdlib.h>
 
-template <typename T, int N>
+// 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 <typename T, int N, bool Pipelined>
 __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<uint8_t, 64>(in, out, 1, lineWidth);
+    passThrough_aie<uint8_t, 64, false>(in, out, 1, lineWidth);
 }
 
 void passThroughTile(uint8_t *in, uint8_t *out, int32_t tileHeight, int32_t tileWidth)
 {
-    passThrough_aie<uint8_t, 64>(in, out, tileHeight, tileWidth);
+    passThrough_aie<uint8_t, 64, false>(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<uint8_t, 64, true>(in, out, 1, lineWidth);
+}
+
+void passThroughTilePipelined(uint8_t *in, uint8_t *out, int32_t tileHeight, int32_t tileWidth)
+{
+    passThrough_aie<uint8_t, 64, true>(in, out, tileHeight, tileWidth);
 }
 
 #elif BIT_WIDTH == 16
 
 void passThroughLine(int16_t *in, int16_t *out, int32_t lineWidth)
 {
-    passThrough_aie<int16_t, 32>(in, out, 1, lineWidth);
+    passThrough_aie<int16_t, 32, false>(in, out, 1, lineWidth);
 }
 
 void passThroughTile(int16_t *in, int16_t *out, int32_t tileHeight, int32_t tileWidth)
 {
-    passThrough_aie<int16_t, 32>(in, out, tileHeight, tileWidth);
+    passThrough_aie<int16_t, 32, false>(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<int16_t, 32, true>(in, out, 1, lineWidth);
+}
+
+void passThroughTilePipelined(int16_t *in, int16_t *out, int32_t tileHeight, int32_t tileWidth)
+{
+    passThrough_aie<int16_t, 32, true>(in, out, tileHeight, tileWidth);
 }
 
 #else // 32
 
 void passThroughLine(int32_t *in, int32_t *out, int32_t lineWidth)
 {
-    passThrough_aie<int32_t, 16>(in, out, 1, lineWidth);
+    passThrough_aie<int32_t, 16, false>(in, out, 1, lineWidth);
 }
 
 void passThroughTile(int32_t *in, int32_t *out, int32_t tileHeight, int32_t tileWidth)
 {
-    passThrough_aie<int32_t, 16>(in, out, tileHeight, tileWidth);
+    passThrough_aie<int32_t, 16, false>(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<int32_t, 16, true>(in, out, 1, lineWidth);
+}
+
+void passThroughTilePipelined(int32_t *in, int32_t *out, int32_t tileHeight, int32_t tileWidth)
+{
+    passThrough_aie<int32_t, 16, true>(in, out, tileHeight, tileWidth);
 }
 
 #endif

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"],
             )
         ]
 

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])