Prefix sum (scan) implementation (C2V sync).

examples/jit_cpp/scan/CustomTSync.hpp

@@ -0,0 +1,62 @@
+// CustomTSync.hpp — bisheng-safe C2V/V2C cross-core sync wrapper
+//
+// Workaround for two bugs in pto-isa TSync primitives (see issue_report.md):
+//   • TSync_Custom::record() — flag_id is a runtime uint16_t member, passed as
+//     the second arg of ffts_cross_core_sync.  Bisheng requires a compile-time
+//     literal there; a runtime value silently produces wrong output.
+//   • Event::Init()          — srcPipe is a computed constexpr pipe_t, not a
+//     literal; bisheng rejects it as the first arg (compile error).
+//
+// Fix: FlagID is a template non-type parameter.  The message constant kMsg is
+// computed entirely from template arguments at instantiation time, so bisheng
+// sees it as a literal in both ffts_cross_core_sync arguments.
+//
+// Usage (C2V, flag slot 0):
+//
+//   CustomTSync<0> c2v_sync;
+//
+//   // Cube side — after TSTORE completes:
+//   c2v_sync.record();
+//
+//   // Vector side — before TLOAD begins:
+//   c2v_sync.wait();
+//
+// Requires: included after <pto/pto-inst.hpp> (provides AICORE, PIPE_MTE3,
+//           ffts_cross_core_sync, wait_flag_dev).
+
+#pragma once
+
+enum SyncDir { CubeToVec, VecToCube };
+
+// CustomTSync<FlagID, Dir>
+//
+//   FlagID — FFTS flag slot index [0..11]; compile-time template param.
+//   Dir    — CubeToVec (default): cube calls record(), vector calls wait().
+//            VecToCube:          vector calls record(), cube calls wait().
+template <uint8_t FlagID, SyncDir Dir = CubeToVec>
+struct CustomTSync {
+  // FFTS message word (encoding from TSyncCVID.hpp / _getFFTSMsg):
+  //   bits [3:0]   base    = 1
+  //   bits [5:4]   mode    = CV_CORE_SYNC = 2
+  //   bits [11:8]  flag_id = FlagID
+  // Every term is a template/compile-time constant → kMsg is a literal.
+  static constexpr uint16_t kBase = 1;
+  static constexpr uint16_t kMode = 2;  // CV_CORE_SYNC
+  static constexpr uint64_t kMsg =
+      kBase | (kMode << 4) | ((uint64_t)FlagID << 8);
+
+  // Producer side: signal that data written to GM is visible.
+  //
+  // C2V: emits ffts_cross_core_sync(PIPE_FIX, kMsg).
+  // V2C: emits ffts_cross_core_sync(PIPE_MTE3, kMsg).
+  AICORE inline void record() const {
+    if (Dir == CubeToVec) {
+      ffts_cross_core_sync(PIPE_FIX, kMsg);
+    } else {
+      ffts_cross_core_sync(PIPE_MTE3, kMsg);
+    }
+  }
+
+  // Consumer side: stall until the producer's record() signal arrives.
+  AICORE inline void wait() const { wait_flag_dev(FlagID); }
+};

examples/jit_cpp/scan/README.md

@@ -0,0 +1,13 @@
+# Single core prefix sum (scan)
+
+An implementation of prefix sum (scan) algorithm, based on https://arxiv.org/abs/2505.15112v1. Only single core algorithm is implemented (ScanU from the paper).
+
+Usage:
+
+```bash
+
+# Optional
+export PTO_LIB_PATH=${ASCEND_TOOLKIT_HOME}  # reuse CANN 8.5.0 headers
+
+# Run scan tests
+python ./run_scan.py

examples/jit_cpp/scan/jit_util_scan.py

@@ -0,0 +1,109 @@
+import os
+import subprocess
+import ctypes
+import torch
+
+ASCEND_TOOLKIT_HOME = os.environ.get("ASCEND_TOOLKIT_HOME", "")
+PTO_LIB_PATH = os.environ.get("PTO_LIB_PATH", ASCEND_TOOLKIT_HOME)
+
+
+def _get_lib_path(kernel_cpp: str, tile_size: int) -> str:
+    basename = kernel_cpp.replace(".cpp", "")
+    return f"{basename}_{tile_size}_jit.so"
+
+
+def compile_cpp(
+    kernel_cpp: str, tile_size=64, verbose: bool = False, timeout: int = 120
+) -> str:
+    lib_path = _get_lib_path(kernel_cpp, tile_size)
+
+    flags = [
+        "-fPIC",
+        "-shared",
+        "-xcce",
+        "--npu-arch=dav-2201",
+        "-DMEMORY_BASE",
+        "-DSCAN_TILE_SIZE=" + str(tile_size),
+        "-O2",
+        "-std=c++17",
+        f"-I{PTO_LIB_PATH}/include",
+        f"-I{ASCEND_TOOLKIT_HOME}/include",
+        f"-I{ASCEND_TOOLKIT_HOME}/pkg_inc",
+        f"-I{ASCEND_TOOLKIT_HOME}/pkg_inc/runtime",
+        f"-I{ASCEND_TOOLKIT_HOME}/pkg_inc/profiling",
+    ]
+
+    command = ["bisheng", *flags, kernel_cpp, "-o", lib_path]
+    if verbose:
+        print(f"compile {kernel_cpp} with command: \n", command)
+
+    try:
+        subprocess.run(
+            command,
+            timeout=timeout,
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+        )
+    except Exception as e:
+        output = (
+            e.stdout.decode("utf-8", errors="replace")
+            if hasattr(e, "stdout") and e.stdout
+            else ""
+        )
+        raise RuntimeError(
+            f"Compile failed with exit code {e.returncode}:\n{output}"
+        ) from e
+
+    if verbose:
+        print(f"generated {lib_path}")
+    return lib_path
+
+
+def torch_to_ctypes(tensor):
+    return ctypes.c_void_p(tensor.data_ptr())
+
+
+def load_lib(lib_path, check_type=True):
+    lib_path = os.path.abspath(lib_path)
+    lib = ctypes.CDLL(lib_path)
+
+    default_block_dim = 1  # single core scan
+
+    if check_type:
+        lib.scan_fp32.argtypes = [
+            ctypes.c_uint32,  # blockDim
+            ctypes.c_void_p,  # stream
+            ctypes.c_void_p,  # x
+            ctypes.c_void_p,  # y
+            ctypes.c_void_p,  # u
+            ctypes.c_uint,  # total_len
+        ]
+        lib.scan_fp32.restype = None
+
+    def scan_func(x, y, u, total_len, block_dim=default_block_dim, stream_ptr=None):
+        if stream_ptr is None:
+            stream_ptr = torch.npu.current_stream()._as_parameter_  # noqa
+
+        lib.scan_fp32(
+            block_dim,
+            stream_ptr,
+            torch_to_ctypes(x),
+            torch_to_ctypes(y),
+            torch_to_ctypes(u),
+            total_len,
+        )
+
+    return scan_func
+
+
+def jit_compile(src_path, tile_size=64):
+    lib_path = compile_cpp(src_path, tile_size=tile_size, verbose=False)
+    func = load_lib(lib_path, check_type=True)
+    return func
+
+
+def clean_up(kernel_cpp: str, tile_size: int):
+    lib_path = _get_lib_path(kernel_cpp, tile_size)
+    if os.path.exists(lib_path):
+        os.remove(lib_path)

examples/jit_cpp/scan/kernel_scan_single_core.cpp

@@ -0,0 +1,212 @@
+#include <pto/pto-inst.hpp>
+
+#include "CustomTSync.hpp"
+#include "runtime/rt.h"
+
+using namespace pto;
+
+#ifndef SCAN_TILE_SIZE
+#define SCAN_TILE_SIZE 64
+#endif
+
+constexpr uint32_t TILE_SIZE = SCAN_TILE_SIZE;
+
+template <typename InputT>
+struct ScanAccType {};
+
+template <>
+struct ScanAccType<int8_t> {
+  using type = int32_t;
+};
+
+template <>
+struct ScanAccType<half> {
+  using type = float;
+};
+
+template <>
+struct ScanAccType<float> {
+  using type = float;
+};
+
+template <>
+struct ScanAccType<bfloat16_t> {
+  using type = float;
+};
+
+template <typename InputT>
+AICORE void run_scan_single_core_pto(
+    __gm__ InputT *x, __gm__ typename ScanAccType<InputT>::type *y,
+    __gm__ InputT *u, __gm__ uint8_t *ffts_addr, uint32_t total_len) {
+  using AccT = typename ScanAccType<InputT>::type;
+
+#if (__CHECK_FEATURE_AT_PRECOMPILE) || \
+    (__CCE_AICORE__ == 220 && defined(__DAV_C220_CUBE__))
+
+  // Cube code path
+
+  set_ffts_base_addr((uint64_t)ffts_addr);
+  set_padding(0);
+  set_atomic_none();
+
+  if (get_block_idx() != 0) return;  // Only process on a single core
+
+  using TileL1 =
+      Tile<TileType::Mat, InputT, TILE_SIZE, TILE_SIZE, BLayout::ColMajor,
+           TILE_SIZE, TILE_SIZE, SLayout::RowMajor, 512>;
+  TileL1 aTileL1;
+  TileL1 uTileL1;
+  TASSIGN(aTileL1, 0x0);
+  TASSIGN(uTileL1, 0x0 + TILE_SIZE * TILE_SIZE * sizeof(InputT));
+
+  using TensorShape = TileShape2D<InputT, TILE_SIZE, TILE_SIZE, Layout::ND>;
+  using TensorStrides = BaseShape2D<InputT, TILE_SIZE, TILE_SIZE, Layout::ND>;
+  using GlobalDataIn =
+      GlobalTensor<InputT, TensorShape, TensorStrides, Layout::ND>;
+
+  GlobalDataIn uGM(u);
+
+  using NDValidShapeC = TileShape2D<AccT, TILE_SIZE, TILE_SIZE, Layout::ND>;
+  using NDWholeShapeC = BaseShape2D<AccT, TILE_SIZE, TILE_SIZE, Layout::ND>;
+  using GlobalDataOut =
+      GlobalTensor<AccT, NDValidShapeC, NDWholeShapeC, Layout::ND>;
+
+  using TileA = TileLeft<InputT, TILE_SIZE, TILE_SIZE>;
+  using TileU = TileRight<InputT, TILE_SIZE, TILE_SIZE>;
+  using TileC = TileAcc<AccT, TILE_SIZE, TILE_SIZE>;
+
+  TileA aTile;
+  TASSIGN(aTile, 0x0);
+  TileU uTile;
+  TASSIGN(uTile, 0x0);
+  TileC cTile;
+  TASSIGN(cTile, 0x0);  // L0C bound
+
+  // Load U
+  TLOAD(uTileL1, uGM);
+
+  set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
+  wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
+
+  TMOV(uTile, uTileL1);
+
+  set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
+  wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
+
+  for (uint32_t offset = 0; offset < total_len;
+       offset += TILE_SIZE * TILE_SIZE) {
+    GlobalDataIn gm_x_mat(x + offset);
+    GlobalDataOut gm_out_mat(y + offset);
+
+    // 1. Load A
+    TLOAD(aTileL1, gm_x_mat);
+
+    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
+    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
+
+    TMOV(aTile, aTileL1);
+
+    set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1);
+    wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1);
+
+    // 2. Compute C = A @ U
+    TMATMUL(cTile, aTile, uTile);
+
+    set_flag(PIPE_M, PIPE_FIX, EVENT_ID1);
+    wait_flag(PIPE_M, PIPE_FIX, EVENT_ID1);
+
+    // Store C to global memory
+    TSTORE(gm_out_mat, cTile);
+
+    set_flag(PIPE_FIX, PIPE_MTE2, EVENT_ID2);
+    wait_flag(PIPE_FIX, PIPE_MTE2, EVENT_ID2);
+
+    // Sync with AIV: wait for GM to settle, then notify AIV.
+    // pipe_barrier(PIPE_MTE3); // It was suggested that this is required, but
+    // it doesn't seem to be the case.
+    CustomTSync<0, CubeToVec>().record();
+    // Wait for AIV to finish processing this chunk before advancing.
+    CustomTSync<1, VecToCube>().wait();
+  }
+
+#elif (__CHECK_FEATURE_AT_PRECOMPILE) || \
+    (__CCE_AICORE__ == 220 && defined(__DAV_C220_VEC__))
+  // AIV code path
+
+  set_ffts_base_addr((uint64_t)ffts_addr);
+  set_mask_norm();
+  set_vector_mask(-1, -1);
+
+  if (get_block_idx() != 0) return;
+
+  uint32_t addr = 0;
+  const uint32_t UB_R = addr;
+  addr += TILE_SIZE * sizeof(AccT);
+
+  using VectorRow = Tile<TileType::Vec, AccT, 1, TILE_SIZE>;
+  VectorRow rowTile;
+  TASSIGN(rowTile, UB_R);
+
+  AccT running_sum = 0;
+
+  using GlobalDataRow = GlobalTensor<AccT, Shape<1, 1, 1, TILE_SIZE, TILE_SIZE>,
+                                     Stride<1, 1, 1, TILE_SIZE, 1>>;
+
+  for (uint32_t offset = 0; offset < total_len;
+       offset += TILE_SIZE * TILE_SIZE) {
+    // Wait for AIC to finish generating this chunk.
+    CustomTSync<0, CubeToVec>().wait();
+
+    // 3. Vector phase: process row-by-row, only one vector core works, but both
+    // sync
+    if (get_subblockid() == 0) {
+      for (uint32_t i = 0; i < TILE_SIZE; ++i) {
+        GlobalDataRow gm_row(y + offset + i * TILE_SIZE);
+
+        TLOAD(rowTile, gm_row);
+
+        set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);
+        wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);
+
+        TADDS(rowTile, rowTile, running_sum);
+
+        // Sync before extracting value (simulating scalar fallback)
+        set_flag(PIPE_V, PIPE_S, EVENT_ID2);
+        wait_flag(PIPE_V, PIPE_S, EVENT_ID2);
+        running_sum = rowTile.GetValue(TILE_SIZE - 1);
+        set_flag(PIPE_S, PIPE_V, EVENT_ID2);
+        wait_flag(PIPE_S, PIPE_V, EVENT_ID2);
+
+        TSTORE(gm_row, rowTile);
+
+        set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3);
+        wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3);
+      }
+
+      set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3);
+      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3);
+    }
+    // Notify AIC that chunk processing is fully completed.
+    CustomTSync<1, VecToCube>().record();
+  }
+#endif
+}
+
+__global__ AICORE void kernel_scan_single_core(__gm__ void *x, __gm__ void *out,
+                                               __gm__ void *u_s,
+                                               __gm__ uint8_t *ffts_addr,
+                                               uint32_t total_len) {
+  run_scan_single_core_pto<float>((__gm__ float *)x, (__gm__ float *)out,
+                                  (__gm__ float *)u_s, ffts_addr, total_len);
+}
+
+extern "C" void scan_fp32(uint32_t blockDim, void *stream, void *x, void *out,
+                          void *u_s, const uint32_t total_len) {
+  void *ffts_addr;
+  uint32_t ffts_len;
+  rtGetC2cCtrlAddr((uint64_t *)&ffts_addr, &ffts_len);
+
+  kernel_scan_single_core<<<blockDim, nullptr, stream>>>(
+      (float *)x, (float *)out, (float *)u_s, (__gm__ uint8_t *)ffts_addr,
+      total_len);
+}