// siso_stream_stateful_layers_toggle_df.cpp
// Stateful-across-invocations Transformer encoder pipeline for Vitis HLS,
// with TRUE steady-state overlap across layers using ping-pong (toggle) banks.
//
// Key fix vs. pointer-mux approach:
//   - NO runtime-selected pointers inside the DATAFLOW region.
//   - We instantiate two explicit "step" schedules (toggle=0 and toggle=1)
//     so HLS can prove the per-stage memories don't alias and can overlap
//     layer0(current) with layer1(prev), etc.
//
// Behavior
//  - load_parameters=1 : loads NUM_LAYERS sets of weights into local BRAM, returns.
//  - load_parameters=0 : each invocation consumes exactly 1 sequence from AXIS,
//                        produces exactly 1 sequence to AXIS.
//  - Warmup: first invocation computes full stack sequentially to initialize pipeline state.
//            After that, steady-state overlaps stages per call.
//
// Build-time knobs:
//   -DNUM_LAYERS=1..6  (default 4)
//   -DUF_J=<int>       (default 8)
//   -DUF_D=<int>       (default 8)

#include <hls_stream.h>
#include <hls_math.h>
#include <hls_half.h>
#include <ap_int.h>
#include <ap_axi_sdata.h>
#include <float.h>
#include <cstdint>
#include <cstring>

// FAT_BEGIN DATA_TYPES
typedef float data_t;
typedef float acc_t;
// FAT_END DATA_TYPES

// FAT_BEGIN MODEL_CONFIG
constexpr int SEQ_LEN = 500;
constexpr int D_MODEL = 32;
constexpr int QKV_DIM = 3 * D_MODEL;
constexpr int H = 4;
constexpr int d_k = D_MODEL / H;
constexpr acc_t LN_EPSILON = (acc_t)1e-5f;
constexpr int FF_DIM = 128;
// FAT_END MODEL_CONFIG

#ifndef NEG_LARGE
#define NEG_LARGE (-1e3f)
#endif

#ifndef UF_J
#define UF_J 8
#endif

#ifndef UF_D
#define UF_D 8
#endif

// FAT_BEGIN NUM_LAYERS
#ifndef NUM_LAYERS
#define NUM_LAYERS 1
#endif
// FAT_END NUM_LAYERS

#if (NUM_LAYERS < 1) || (NUM_LAYERS > 6)
#error "NUM_LAYERS must be 1..6"
#endif

// AXI stream type: 32-bit float payload
typedef ap_axiu<32, 0, 0, 0> axis_t;

static inline axis_t pack_f32(float x, bool last) {
#pragma HLS INLINE
  axis_t a;
  uint32_t bits;
  static_assert(sizeof(float) == 4, "Expected 32-bit float");
  std::memcpy(&bits, &x, 4);
  a.data = (ap_uint<32>)bits;
  a.keep = 0xF;
  a.strb = 0xF;
  a.last = last ? 1 : 0;
  return a;
}

static inline float unpack_f32(const axis_t& a) {
#pragma HLS INLINE
  uint32_t bits = (uint32_t)a.data;
  float x;
  std::memcpy(&x, &bits, 4);
  return x;
}

// =============================================================================
// IMPORTANT FOR DATAFLOW:
// Split ALL parameter memories per layer, so each DATAFLOW stage reads only its
// own private arrays (no multi-reader on one global object).
// =============================================================================

// ---- Layer 0 ----
static data_t w_qkv_l0[D_MODEL * (3 * D_MODEL)];
static data_t b_qkv_l0[3 * D_MODEL];
static data_t w_o_l0  [D_MODEL * D_MODEL];
static data_t b_o_l0  [D_MODEL];

static data_t norm1_gamma_l0[D_MODEL];
static data_t norm1_beta_l0 [D_MODEL];
static data_t norm2_gamma_l0[D_MODEL];
static data_t norm2_beta_l0 [D_MODEL];

static data_t ff_w1_l0[D_MODEL * FF_DIM];
static data_t ff_b1_l0[FF_DIM];
static data_t ff_w2_l0[FF_DIM * D_MODEL];
static data_t ff_b2_l0[D_MODEL];

#if (NUM_LAYERS >= 2)
// ---- Layer 1 ----
static data_t w_qkv_l1[D_MODEL * (3 * D_MODEL)];
static data_t b_qkv_l1[3 * D_MODEL];
static data_t w_o_l1  [D_MODEL * D_MODEL];
static data_t b_o_l1  [D_MODEL];

static data_t norm1_gamma_l1[D_MODEL];
static data_t norm1_beta_l1 [D_MODEL];
static data_t norm2_gamma_l1[D_MODEL];
static data_t norm2_beta_l1 [D_MODEL];

static data_t ff_w1_l1[D_MODEL * FF_DIM];
static data_t ff_b1_l1[FF_DIM];
static data_t ff_w2_l1[FF_DIM * D_MODEL];
static data_t ff_b2_l1[D_MODEL];
#endif

#if (NUM_LAYERS >= 3)
// ---- Layer 2 ----
static data_t w_qkv_l2[D_MODEL * (3 * D_MODEL)];
static data_t b_qkv_l2[3 * D_MODEL];
static data_t w_o_l2  [D_MODEL * D_MODEL];
static data_t b_o_l2  [D_MODEL];

static data_t norm1_gamma_l2[D_MODEL];
static data_t norm1_beta_l2 [D_MODEL];
static data_t norm2_gamma_l2[D_MODEL];
static data_t norm2_beta_l2 [D_MODEL];

static data_t ff_w1_l2[D_MODEL * FF_DIM];
static data_t ff_b1_l2[FF_DIM];
static data_t ff_w2_l2[FF_DIM * D_MODEL];
static data_t ff_b2_l2[D_MODEL];
#endif

#if (NUM_LAYERS >= 4)
// ---- Layer 3 ----
static data_t w_qkv_l3[D_MODEL * (3 * D_MODEL)];
static data_t b_qkv_l3[3 * D_MODEL];
static data_t w_o_l3  [D_MODEL * D_MODEL];
static data_t b_o_l3  [D_MODEL];

static data_t norm1_gamma_l3[D_MODEL];
static data_t norm1_beta_l3 [D_MODEL];
static data_t norm2_gamma_l3[D_MODEL];
static data_t norm2_beta_l3 [D_MODEL];

static data_t ff_w1_l3[D_MODEL * FF_DIM];
static data_t ff_b1_l3[FF_DIM];
static data_t ff_w2_l3[FF_DIM * D_MODEL];
static data_t ff_b2_l3[D_MODEL];
#endif

#if (NUM_LAYERS >= 5)
// ---- Layer 4 ----
static data_t w_qkv_l4[D_MODEL * (3 * D_MODEL)];
static data_t b_qkv_l4[3 * D_MODEL];
static data_t w_o_l4  [D_MODEL * D_MODEL];
static data_t b_o_l4  [D_MODEL];

static data_t norm1_gamma_l4[D_MODEL];
static data_t norm1_beta_l4 [D_MODEL];
static data_t norm2_gamma_l4[D_MODEL];
static data_t norm2_beta_l4 [D_MODEL];

static data_t ff_w1_l4[D_MODEL * FF_DIM];
static data_t ff_b1_l4[FF_DIM];
static data_t ff_w2_l4[FF_DIM * D_MODEL];
static data_t ff_b2_l4[D_MODEL];
#endif

#if (NUM_LAYERS >= 6)
// ---- Layer 5 ----
static data_t w_qkv_l5[D_MODEL * (3 * D_MODEL)];
static data_t b_qkv_l5[3 * D_MODEL];
static data_t w_o_l5  [D_MODEL * D_MODEL];
static data_t b_o_l5  [D_MODEL];

static data_t norm1_gamma_l5[D_MODEL];
static data_t norm1_beta_l5 [D_MODEL];
static data_t norm2_gamma_l5[D_MODEL];
static data_t norm2_beta_l5 [D_MODEL];

static data_t ff_w1_l5[D_MODEL * FF_DIM];
static data_t ff_b1_l5[FF_DIM];
static data_t ff_w2_l5[FF_DIM * D_MODEL];
static data_t ff_b2_l5[D_MODEL];
#endif

// =============================================================================
// Small helper: map runtime layer index -> pointers to that layer's arrays.
// (Safe because load happens outside DATAFLOW; compute stages use templates.)
// =============================================================================
static inline data_t* ptr_w_qkv(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return w_qkv_l0;
#if (NUM_LAYERS >= 2)
    case 1: return w_qkv_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return w_qkv_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return w_qkv_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return w_qkv_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return w_qkv_l5;
#endif
    default: return w_qkv_l0;
  }
}
static inline data_t* ptr_b_qkv(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return b_qkv_l0;
#if (NUM_LAYERS >= 2)
    case 1: return b_qkv_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return b_qkv_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return b_qkv_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return b_qkv_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return b_qkv_l5;
#endif
    default: return b_qkv_l0;
  }
}
static inline data_t* ptr_w_o(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return w_o_l0;
#if (NUM_LAYERS >= 2)
    case 1: return w_o_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return w_o_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return w_o_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return w_o_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return w_o_l5;
#endif
    default: return w_o_l0;
  }
}
static inline data_t* ptr_b_o(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return b_o_l0;
#if (NUM_LAYERS >= 2)
    case 1: return b_o_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return b_o_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return b_o_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return b_o_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return b_o_l5;
#endif
    default: return b_o_l0;
  }
}
static inline data_t* ptr_norm1_gamma(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return norm1_gamma_l0;
#if (NUM_LAYERS >= 2)
    case 1: return norm1_gamma_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return norm1_gamma_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return norm1_gamma_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return norm1_gamma_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return norm1_gamma_l5;
#endif
    default: return norm1_gamma_l0;
  }
}
static inline data_t* ptr_norm1_beta(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return norm1_beta_l0;
#if (NUM_LAYERS >= 2)
    case 1: return norm1_beta_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return norm1_beta_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return norm1_beta_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return norm1_beta_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return norm1_beta_l5;
#endif
    default: return norm1_beta_l0;
  }
}
static inline data_t* ptr_norm2_gamma(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return norm2_gamma_l0;
#if (NUM_LAYERS >= 2)
    case 1: return norm2_gamma_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return norm2_gamma_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return norm2_gamma_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return norm2_gamma_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return norm2_gamma_l5;
#endif
    default: return norm2_gamma_l0;
  }
}
static inline data_t* ptr_norm2_beta(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return norm2_beta_l0;
#if (NUM_LAYERS >= 2)
    case 1: return norm2_beta_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return norm2_beta_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return norm2_beta_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return norm2_beta_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return norm2_beta_l5;
#endif
    default: return norm2_beta_l0;
  }
}
static inline data_t* ptr_ff_w1(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return ff_w1_l0;
#if (NUM_LAYERS >= 2)
    case 1: return ff_w1_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return ff_w1_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return ff_w1_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return ff_w1_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return ff_w1_l5;
#endif
    default: return ff_w1_l0;
  }
}
static inline data_t* ptr_ff_b1(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return ff_b1_l0;
#if (NUM_LAYERS >= 2)
    case 1: return ff_b1_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return ff_b1_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return ff_b1_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return ff_b1_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return ff_b1_l5;
#endif
    default: return ff_b1_l0;
  }
}
static inline data_t* ptr_ff_w2(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return ff_w2_l0;
#if (NUM_LAYERS >= 2)
    case 1: return ff_w2_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return ff_w2_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return ff_w2_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return ff_w2_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return ff_w2_l5;
#endif
    default: return ff_w2_l0;
  }
}
static inline data_t* ptr_ff_b2(int l) {
#pragma HLS INLINE
  switch (l) {
    case 0: return ff_b2_l0;
#if (NUM_LAYERS >= 2)
    case 1: return ff_b2_l1;
#endif
#if (NUM_LAYERS >= 3)
    case 2: return ff_b2_l2;
#endif
#if (NUM_LAYERS >= 4)
    case 3: return ff_b2_l3;
#endif
#if (NUM_LAYERS >= 5)
    case 4: return ff_b2_l4;
#endif
#if (NUM_LAYERS >= 6)
    case 5: return ff_b2_l5;
#endif
    default: return ff_b2_l0;
  }
}

// =============================================================================
// Load params for NUM_LAYERS layers from GMEM to local per-layer buffers
// =============================================================================
static void load_layer_params(
    const data_t* W_QKV, const data_t* B_QKV,
    const data_t* W_O,   const data_t* B_O,
    const data_t* norm1_gamma, const data_t* norm1_beta,
    const data_t* norm2_gamma, const data_t* norm2_beta,
    const data_t* ff_b1, const data_t* ff_b2,
    const data_t* ff_w1, const data_t* ff_w2)
{
#pragma HLS INLINE off

  constexpr int WQKV_ELEMS  = D_MODEL * (3 * D_MODEL);
  constexpr int BQKV_ELEMS  = 3 * D_MODEL;
  constexpr int WO_ELEMS    = D_MODEL * D_MODEL;
  constexpr int FF_W1_ELEMS = D_MODEL * FF_DIM;
  constexpr int FF_W2_ELEMS = FF_DIM * D_MODEL;

  for (int l = 0; l < NUM_LAYERS; l++) {
    const data_t* pWQKV = W_QKV + l * WQKV_ELEMS;
    const data_t* pBQKV = B_QKV + l * BQKV_ELEMS;
    const data_t* pWO   = W_O   + l * WO_ELEMS;
    const data_t* pBO   = B_O   + l * D_MODEL;

    const data_t* pG1   = norm1_gamma + l * D_MODEL;
    const data_t* pB1n  = norm1_beta  + l * D_MODEL;
    const data_t* pG2   = norm2_gamma + l * D_MODEL;
    const data_t* pB2n  = norm2_beta  + l * D_MODEL;

    const data_t* pFFb1 = ff_b1 + l * FF_DIM;
    const data_t* pFFb2 = ff_b2 + l * D_MODEL;
    const data_t* pFFw1 = ff_w1 + l * FF_W1_ELEMS;
    const data_t* pFFw2 = ff_w2 + l * FF_W2_ELEMS;

    data_t* dWQKV = ptr_w_qkv(l);
    data_t* dBQKV = ptr_b_qkv(l);
    data_t* dWO   = ptr_w_o(l);
    data_t* dBO   = ptr_b_o(l);

    data_t* dG1   = ptr_norm1_gamma(l);
    data_t* dB1n  = ptr_norm1_beta(l);
    data_t* dG2   = ptr_norm2_gamma(l);
    data_t* dB2n  = ptr_norm2_beta(l);

    data_t* dFFb1 = ptr_ff_b1(l);
    data_t* dFFb2 = ptr_ff_b2(l);
    data_t* dFFw1 = ptr_ff_w1(l);
    data_t* dFFw2 = ptr_ff_w2(l);

    for (int i = 0; i < WQKV_ELEMS; i++) {
#pragma HLS PIPELINE II=1
      dWQKV[i] = pWQKV[i];
    }
    for (int i = 0; i < BQKV_ELEMS; i++) {
#pragma HLS PIPELINE II=1
      dBQKV[i] = pBQKV[i];
    }
    for (int i = 0; i < WO_ELEMS; i++) {
#pragma HLS PIPELINE II=1
      dWO[i] = pWO[i];
    }

    for (int i = 0; i < D_MODEL; i++) {
#pragma HLS PIPELINE II=1
      dBO[i]   = pBO[i];
      dG1[i]   = pG1[i];
      dB1n[i]  = pB1n[i];
      dG2[i]   = pG2[i];
      dB2n[i]  = pB2n[i];
      dFFb2[i] = pFFb2[i];
    }

    for (int i = 0; i < FF_W1_ELEMS; i++) {
#pragma HLS PIPELINE II=1
      dFFw1[i] = pFFw1[i];
    }
    for (int i = 0; i < FF_DIM; i++) {
#pragma HLS PIPELINE II=1
      dFFb1[i] = pFFb1[i];
    }
    for (int i = 0; i < FF_W2_ELEMS; i++) {
#pragma HLS PIPELINE II=1
      dFFw2[i] = pFFw2[i];
    }
  }
}

// =============================================================================
// Compile-time layer parameter access (NO shared 2D arrays).
// =============================================================================
template<int L> struct LayerParams;

template<> struct LayerParams<0> {
  static const data_t* WQKV() { return w_qkv_l0; }
  static const data_t* BQKV() { return b_qkv_l0; }
  static const data_t* WO()   { return w_o_l0; }
  static const data_t* BO()   { return b_o_l0; }
  static const data_t* G1()   { return norm1_gamma_l0; }
  static const data_t* B1()   { return norm1_beta_l0; }
  static const data_t* G2()   { return norm2_gamma_l0; }
  static const data_t* B2()   { return norm2_beta_l0; }
  static const data_t* FW1()  { return ff_w1_l0; }
  static const data_t* FB1()  { return ff_b1_l0; }
  static const data_t* FW2()  { return ff_w2_l0; }
  static const data_t* FB2()  { return ff_b2_l0; }
};

#if (NUM_LAYERS >= 2)
template<> struct LayerParams<1> {
  static const data_t* WQKV() { return w_qkv_l1; }
  static const data_t* BQKV() { return b_qkv_l1; }
  static const data_t* WO()   { return w_o_l1; }
  static const data_t* BO()   { return b_o_l1; }
  static const data_t* G1()   { return norm1_gamma_l1; }
  static const data_t* B1()   { return norm1_beta_l1; }
  static const data_t* G2()   { return norm2_gamma_l1; }
  static const data_t* B2()   { return norm2_beta_l1; }
  static const data_t* FW1()  { return ff_w1_l1; }
  static const data_t* FB1()  { return ff_b1_l1; }
  static const data_t* FW2()  { return ff_w2_l1; }
  static const data_t* FB2()  { return ff_b2_l1; }
};
#endif

#if (NUM_LAYERS >= 3)
template<> struct LayerParams<2> {
  static const data_t* WQKV() { return w_qkv_l2; }
  static const data_t* BQKV() { return b_qkv_l2; }
  static const data_t* WO()   { return w_o_l2; }
  static const data_t* BO()   { return b_o_l2; }
  static const data_t* G1()   { return norm1_gamma_l2; }
  static const data_t* B1()   { return norm1_beta_l2; }
  static const data_t* G2()   { return norm2_gamma_l2; }
  static const data_t* B2()   { return norm2_beta_l2; }
  static const data_t* FW1()  { return ff_w1_l2; }
  static const data_t* FB1()  { return ff_b1_l2; }
  static const data_t* FW2()  { return ff_w2_l2; }
  static const data_t* FB2()  { return ff_b2_l2; }
};
#endif

#if (NUM_LAYERS >= 4)
template<> struct LayerParams<3> {
  static const data_t* WQKV() { return w_qkv_l3; }
  static const data_t* BQKV() { return b_qkv_l3; }
  static const data_t* WO()   { return w_o_l3; }
  static const data_t* BO()   { return b_o_l3; }
  static const data_t* G1()   { return norm1_gamma_l3; }
  static const data_t* B1()   { return norm1_beta_l3; }
  static const data_t* G2()   { return norm2_gamma_l3; }
  static const data_t* B2()   { return norm2_beta_l3; }
  static const data_t* FW1()  { return ff_w1_l3; }
  static const data_t* FB1()  { return ff_b1_l3; }
  static const data_t* FW2()  { return ff_w2_l3; }
  static const data_t* FB2()  { return ff_b2_l3; }
};
#endif

#if (NUM_LAYERS >= 5)
template<> struct LayerParams<4> {
  static const data_t* WQKV() { return w_qkv_l4; }
  static const data_t* BQKV() { return b_qkv_l4; }
  static const data_t* WO()   { return w_o_l4; }
  static const data_t* BO()   { return b_o_l4; }
  static const data_t* G1()   { return norm1_gamma_l4; }
  static const data_t* B1()   { return norm1_beta_l4; }
  static const data_t* G2()   { return norm2_gamma_l4; }
  static const data_t* B2()   { return norm2_beta_l4; }
  static const data_t* FW1()  { return ff_w1_l4; }
  static const data_t* FB1()  { return ff_b1_l4; }
  static const data_t* FW2()  { return ff_w2_l4; }
  static const data_t* FB2()  { return ff_b2_l4; }
};
#endif

#if (NUM_LAYERS >= 6)
template<> struct LayerParams<5> {
  static const data_t* WQKV() { return w_qkv_l5; }
  static const data_t* BQKV() { return b_qkv_l5; }
  static const data_t* WO()   { return w_o_l5; }
  static const data_t* BO()   { return b_o_l5; }
  static const data_t* G1()   { return norm1_gamma_l5; }
  static const data_t* B1()   { return norm1_beta_l5; }
  static const data_t* G2()   { return norm2_gamma_l5; }
  static const data_t* B2()   { return norm2_beta_l5; }
  static const data_t* FW1()  { return ff_w1_l5; }
  static const data_t* FB1()  { return ff_b1_l5; }
  static const data_t* FW2()  { return ff_w2_l5; }
  static const data_t* FB2()  { return ff_b2_l5; }
};
#endif

// =============================================================================
// One layer core: in_act -> out_act using compile-time layer L
// =============================================================================
template<int L>
static void transformer_layer_core(
    const data_t in_act [SEQ_LEN][D_MODEL],
          data_t out_act[SEQ_LEN][D_MODEL],
    int mask_count,
    data_t qkv_out   [SEQ_LEN][QKV_DIM],
    data_t mha_out   [SEQ_LEN][D_MODEL],
    data_t res1      [SEQ_LEN][D_MODEL],
    data_t norm1_out [SEQ_LEN][D_MODEL],
    data_t h_buf     [SEQ_LEN][FF_DIM],
    data_t ff_out    [SEQ_LEN][D_MODEL],
    data_t res2      [SEQ_LEN][D_MODEL])
{
#pragma HLS INLINE off

  const data_t* WQKV = LayerParams<L>::WQKV();
  const data_t* BQKV = LayerParams<L>::BQKV();
  const data_t* WO   = LayerParams<L>::WO();
  const data_t* BO   = LayerParams<L>::BO();

  const data_t* G1   = LayerParams<L>::G1();
  const data_t* B1n  = LayerParams<L>::B1();
  const data_t* G2   = LayerParams<L>::G2();
  const data_t* B2n  = LayerParams<L>::B2();

  const data_t* FW1  = LayerParams<L>::FW1();
  const data_t* FB1  = LayerParams<L>::FB1();
  const data_t* FW2  = LayerParams<L>::FW2();
  const data_t* FB2  = LayerParams<L>::FB2();

  const int WQKV_STRIDE = QKV_DIM;
  const int W1_STRIDE   = FF_DIM;
  const int W2_STRIDE   = D_MODEL;

  // 1) QKV projection
  for (int i = 0; i < SEQ_LEN; i++) {
    for (int j0 = 0; j0 < QKV_DIM; j0 += UF_J) {
      acc_t acc[UF_J];
#pragma HLS ARRAY_PARTITION variable=acc complete

      for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
        int j = j0 + u;
        acc[u] = (j < QKV_DIM) ? (acc_t)BQKV[j] : (acc_t)0.0f;
      }

      for (int k = 0; k < D_MODEL; k++) {
#pragma HLS PIPELINE II=1
        acc_t a = (acc_t)in_act[i][k];
        for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
          int j = j0 + u;
          if (j < QKV_DIM) {
            acc_t w = (acc_t)WQKV[k * WQKV_STRIDE + j];
            acc[u] += a * w;
          }
        }
      }

      for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
        int j = j0 + u;
        if (j < QKV_DIM) qkv_out[i][j] = (data_t)acc[u];
      }
    }
  }

  // 2) Q/K scaling
  {
    acc_t scale = hls::powf((acc_t)d_k, (acc_t)-0.25f);
    for (int t = 0; t < SEQ_LEN; t++) {
#pragma HLS PIPELINE II=1
      for (int d = 0; d < QKV_DIM; d++) {
        acc_t val = (acc_t)qkv_out[t][d];
        if (d < D_MODEL)        qkv_out[t][d] = (data_t)(val * scale);
        else if (d < 2*D_MODEL) qkv_out[t][d] = (data_t)(val * scale);
      }
    }
  }

  // 3) Multi-Head Attention + W_O
  for (int i = 0; i < SEQ_LEN; i++) {

    acc_t out_acc[D_MODEL];
#pragma HLS ARRAY_PARTITION variable=out_acc cyclic factor=UF_J
    for (int oj = 0; oj < D_MODEL; oj++) {
#pragma HLS PIPELINE II=1
      out_acc[oj] = (acc_t)BO[oj];
    }

    for (int hh = 0; hh < H; hh++) {
      acc_t scores[SEQ_LEN];
      acc_t ebuf  [SEQ_LEN];
#pragma HLS BIND_STORAGE variable=scores type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=ebuf   type=ram_2p impl=bram

      acc_t maxv = -FLT_MAX;

      // scores + max
      for (int j = 0; j < SEQ_LEN; j++) {
#pragma HLS PIPELINE II=1
        acc_t dot = 0.f;

        for (int d0 = 0; d0 < d_k; d0 += UF_D) {
#pragma HLS UNROLL factor=1
          for (int u = 0; u < UF_D; u++) {
#pragma HLS UNROLL
            int d = d0 + u;
            if (d < d_k) {
              int col = hh * d_k + d;
              dot += (acc_t)qkv_out[i][col] * (acc_t)qkv_out[j][D_MODEL + col];
            }
          }
        }

        if (j >= mask_count) dot = (acc_t)NEG_LARGE;
        scores[j] = dot;
        if (dot > maxv) maxv = dot;
      }

      // exp + sum
      acc_t sum = 0.f;
      for (int j = 0; j < SEQ_LEN; j++) {
#pragma HLS PIPELINE II=1
        acc_t e = hls::expf(scores[j] - maxv);
        ebuf[j] = e;
        sum += e;
      }
      acc_t inv_sum = (sum > 0.f) ? (acc_t)(1.0f / sum) : (acc_t)0.f;

      // softmax*V -> head_vec
      acc_t head_vec[d_k];
#pragma HLS ARRAY_PARTITION variable=head_vec cyclic factor=UF_D
      for (int d = 0; d < d_k; d++) {
#pragma HLS PIPELINE II=1
        head_vec[d] = 0.f;
      }

      for (int j = 0; j < SEQ_LEN; j++) {
#pragma HLS PIPELINE II=1
        acc_t w = ebuf[j] * inv_sum;
        for (int d0 = 0; d0 < d_k; d0 += UF_D) {
#pragma HLS UNROLL factor=1
          for (int u = 0; u < UF_D; u++) {
#pragma HLS UNROLL
            int d = d0 + u;
            if (d < d_k) {
              int col = hh * d_k + d;
              head_vec[d] += w * (acc_t)qkv_out[j][2*D_MODEL + col];
            }
          }
        }
      }

      // project head slice
      const int base = hh * d_k;
      for (int oj0 = 0; oj0 < D_MODEL; oj0 += UF_J) {
        acc_t part[UF_J];
#pragma HLS ARRAY_PARTITION variable=part complete
        for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
          part[u] = 0.f;
        }

        for (int d = 0; d < d_k; d++) {
#pragma HLS PIPELINE II=1
          acc_t a = head_vec[d];
          for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
            int oj = oj0 + u;
            if (oj < D_MODEL) {
              part[u] += a * (acc_t)WO[oj * D_MODEL + (base + d)];
            }
          }
        }

        for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
          int oj = oj0 + u;
          if (oj < D_MODEL) out_acc[oj] += part[u];
        }
      }
    } // heads

    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      mha_out[i][j] = (data_t)out_acc[j];
    }
  }

  // 4) Residual + LN1  (reuse out_act as LN1 output to reduce buffers)
  for (int i = 0; i < SEQ_LEN; i++) {
    acc_t sum = 0.f, sumsq = 0.f;
    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      acc_t v = (acc_t)in_act[i][j] + (acc_t)mha_out[i][j];
      res1[i][j] = (data_t)v;
      sum   += v;
      sumsq += v * v;
    }
    acc_t mu  = sum / (acc_t)D_MODEL;
    acc_t var = sumsq / (acc_t)D_MODEL - mu * mu;
    acc_t inv = hls::powf(var + LN_EPSILON, (acc_t)-0.5f);

    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      acc_t v = (acc_t)res1[i][j];
      out_act[i][j] = (data_t)((acc_t)G1[j] * ((v - mu) * inv) + (acc_t)B1n[j]);
    }
  }

  // 5) FFN: FC1 + ReLU (uses out_act as LN1 output)
  for (int i = 0; i < SEQ_LEN; i++) {
    for (int j0 = 0; j0 < FF_DIM; j0 += UF_J) {
      acc_t acc[UF_J];
#pragma HLS ARRAY_PARTITION variable=acc complete

      for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
        int j = j0 + u;
        acc[u] = (j < FF_DIM) ? (acc_t)FB1[j] : (acc_t)0.f;
      }

      for (int k = 0; k < D_MODEL; k++) {
#pragma HLS PIPELINE II=1
        acc_t a = (acc_t)out_act[i][k];
        for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
          int j = j0 + u;
          if (j < FF_DIM) {
            acc[u] += a * (acc_t)FW1[k * W1_STRIDE + j];
          }
        }
      }

      for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
        int j = j0 + u;
        if (j < FF_DIM) {
          acc_t v = acc[u];
          h_buf[i][j] = (data_t)((v > 0.f) ? v : 0.f);
        }
      }
    }
  }

  // FC2
  for (int i = 0; i < SEQ_LEN; i++) {
    for (int j0 = 0; j0 < D_MODEL; j0 += UF_J) {
      acc_t acc[UF_J];
#pragma HLS ARRAY_PARTITION variable=acc complete

      for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
        int j = j0 + u;
        acc[u] = (j < D_MODEL) ? (acc_t)FB2[j] : (acc_t)0.f;
      }

      for (int k = 0; k < FF_DIM; k++) {
#pragma HLS PIPELINE II=1
        acc_t a = (acc_t)h_buf[i][k];
        for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
          int j = j0 + u;
          if (j < D_MODEL) {
            acc[u] += a * (acc_t)FW2[k * W2_STRIDE + j];
          }
        }
      }

      for (int u = 0; u < UF_J; u++) {
#pragma HLS UNROLL
        int j = j0 + u;
        if (j < D_MODEL) ff_out[i][j] = (data_t)acc[u];
      }
    }
  }

  // 6) Residual + LN2 -> out_act
  for (int i = 0; i < SEQ_LEN; i++) {
    acc_t sum = 0.f, sumsq = 0.f;
    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      acc_t v = (acc_t)out_act[i][j] + (acc_t)ff_out[i][j];
      res2[i][j] = (data_t)v;
      sum   += v;
      sumsq += v * v;
    }
    acc_t mu  = sum / (acc_t)D_MODEL;
    acc_t var = sumsq / (acc_t)D_MODEL - mu * mu;
    acc_t inv = hls::powf(var + LN_EPSILON, (acc_t)-0.5f);

    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      acc_t v = (acc_t)res2[i][j];
      out_act[i][j] = (data_t)((acc_t)G2[j] * ((v - mu) * inv) + (acc_t)B2n[j]);
    }
  }
}

// =============================================================================
// Per-layer instance wrapper with private intermediates.
// Template instantiations ensure distinct statics => DATAFLOW safe.
// =============================================================================
template<int L>
static void transformer_layer_inst(
    const data_t in_act [SEQ_LEN][D_MODEL],
          data_t out_act[SEQ_LEN][D_MODEL],
    int mask_count)
{
#pragma HLS INLINE off

  static data_t qkv_out   [SEQ_LEN][QKV_DIM];
  static data_t mha_out   [SEQ_LEN][D_MODEL];
  static data_t res1      [SEQ_LEN][D_MODEL];
  static data_t h_buf     [SEQ_LEN][FF_DIM];
  static data_t ff_out    [SEQ_LEN][D_MODEL];
  static data_t res2      [SEQ_LEN][D_MODEL];

#pragma HLS BIND_STORAGE variable=qkv_out   type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=mha_out   type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=res1      type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=h_buf     type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=ff_out    type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=res2      type=ram_2p impl=bram

  transformer_layer_core<L>(in_act, out_act, mask_count,
                            qkv_out, mha_out, res1, /*norm1_out reused*/ out_act,
                            h_buf, ff_out, res2);
}

// =============================================================================
// Small tasks (functions) so DATAFLOW can schedule them as separate processes
// =============================================================================
static void task_read_input_axis(hls::stream<axis_t>& s_axis_in,
                                data_t in_local[SEQ_LEN][D_MODEL])
{
#pragma HLS INLINE off
  for (int i = 0; i < SEQ_LEN; i++) {
    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      in_local[i][j] = (data_t)unpack_f32(s_axis_in.read());
    }
  }
}

template<int L>
static void task_run_layer(const data_t in_act [SEQ_LEN][D_MODEL],
                           data_t out_act[SEQ_LEN][D_MODEL],
                           int mask_count)
{
#pragma HLS INLINE off
  transformer_layer_inst<L>(in_act, out_act, mask_count);
}

static void task_stream_output_axis(const data_t out_local[SEQ_LEN][D_MODEL],
                                    hls::stream<axis_t>& m_axis_out)
{
#pragma HLS INLINE off
  for (int i = 0; i < SEQ_LEN; i++) {
    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      bool last = (i == SEQ_LEN - 1) && (j == D_MODEL - 1);
      m_axis_out.write(pack_f32((float)out_local[i][j], last));
    }
  }
}

template<int L>
static void task_run_last_layer_and_stream(
    const data_t in_act[SEQ_LEN][D_MODEL],
    hls::stream<axis_t>& m_axis_out,
    int mask_count)
{
#pragma HLS INLINE off

  // Local output buffer stays *inside* this one process
  static data_t out_local[SEQ_LEN][D_MODEL];
#pragma HLS BIND_STORAGE variable=out_local type=ram_2p impl=bram

  transformer_layer_inst<L>(in_act, out_local, mask_count);

  // stream out here (no separate DATAFLOW reader)
  for (int i = 0; i < SEQ_LEN; i++) {
    for (int j = 0; j < D_MODEL; j++) {
#pragma HLS PIPELINE II=1
      bool last = (i == SEQ_LEN - 1) && (j == D_MODEL - 1);
      m_axis_out.write(pack_f32((float)out_local[i][j], last));
    }
  }
}

// =============================================================================
// Steady-state step schedules (toggle=0 and toggle=1)
// IMPORTANT: arrays are explicit arguments (no runtime pointer mux inside DATAFLOW)
// =============================================================================

#if (NUM_LAYERS == 1)

static void step_single(
    hls::stream<axis_t>& s_axis_in,
    hls::stream<axis_t>& m_axis_out,
    int mask_count)
{
#pragma HLS INLINE off
  static data_t in_local[SEQ_LEN][D_MODEL];
#pragma HLS BIND_STORAGE variable=in_local type=ram_2p impl=bram

#pragma HLS DATAFLOW
  task_read_input_axis(s_axis_in, in_local);
  task_run_last_layer_and_stream<0>(in_local, m_axis_out, mask_count);
}

#else // NUM_LAYERS >= 2

static void step_2layers(
    hls::stream<axis_t>& s_axis_in,
    hls::stream<axis_t>& m_axis_out,
    int mc_cur,
    data_t prev0[SEQ_LEN][D_MODEL],
    data_t next0[SEQ_LEN][D_MODEL],
    int &mc_prev0,
    int &mc_next0)
{
#pragma HLS INLINE off
  static data_t in_local[SEQ_LEN][D_MODEL];
#pragma HLS BIND_STORAGE variable=in_local type=ram_2p impl=bram

#pragma HLS DATAFLOW
  task_read_input_axis(s_axis_in, in_local);

  // layer0(current) -> next0
  mc_next0 = mc_cur;
  task_run_layer<0>(in_local, next0, mc_cur);

  // layer1(prev0) -> out
  task_run_last_layer_and_stream<1>(prev0, m_axis_out, mc_prev0);
}

#if (NUM_LAYERS >= 3)
static void step_3layers(
    hls::stream<axis_t>& s_axis_in,
    hls::stream<axis_t>& m_axis_out,
    int mc_cur,
    data_t prev0[SEQ_LEN][D_MODEL], data_t next0[SEQ_LEN][D_MODEL],
    data_t prev1[SEQ_LEN][D_MODEL], data_t next1[SEQ_LEN][D_MODEL],
    int &mc_prev0, int &mc_next0,
    int &mc_prev1, int &mc_next1)
{
#pragma HLS INLINE off
  static data_t in_local[SEQ_LEN][D_MODEL];
#pragma HLS BIND_STORAGE variable=in_local type=ram_2p impl=bram

#pragma HLS DATAFLOW
  task_read_input_axis(s_axis_in, in_local);

  // layer0(current) -> next0
  mc_next0 = mc_cur;
  task_run_layer<0>(in_local, next0, mc_cur);

  // layer1(prev0) -> next1
  mc_next1 = mc_prev0;
  task_run_layer<1>(prev0, next1, mc_prev0);

  // layer2(prev1) -> out
  task_run_last_layer_and_stream<2>(prev1, m_axis_out, mc_prev1);
}
#endif

#if (NUM_LAYERS >= 4)
static void step_4layers(
    hls::stream<axis_t>& s_axis_in,
    hls::stream<axis_t>& m_axis_out,
    int mc_cur,
    data_t prev0[SEQ_LEN][D_MODEL], data_t next0[SEQ_LEN][D_MODEL],
    data_t prev1[SEQ_LEN][D_MODEL], data_t next1[SEQ_LEN][D_MODEL],
    data_t prev2[SEQ_LEN][D_MODEL], data_t next2[SEQ_LEN][D_MODEL],
    int &mc_prev0, int &mc_next0,
    int &mc_prev1, int &mc_next1,
    int &mc_prev2, int &mc_next2)
{
#pragma HLS INLINE off
  static data_t in_local[SEQ_LEN][D_MODEL];
#pragma HLS BIND_STORAGE variable=in_local type=ram_2p impl=bram

#pragma HLS DATAFLOW
  task_read_input_axis(s_axis_in, in_local);

  // layer0(current) -> next0
  mc_next0 = mc_cur;
  task_run_layer<0>(in_local, next0, mc_cur);

  // layer1(prev0) -> next1
  mc_next1 = mc_prev0;
  task_run_layer<1>(prev0, next1, mc_prev0);

  // layer2(prev1) -> next2
  mc_next2 = mc_prev1;
  task_run_layer<2>(prev1, next2, mc_prev1);

  // layer3(prev2) -> out
  task_run_last_layer_and_stream<3>(prev2, m_axis_out, mc_prev2);
}
#endif

#if (NUM_LAYERS >= 5)
static void step_5layers(
    hls::stream<axis_t>& s_axis_in,
    hls::stream<axis_t>& m_axis_out,
    int mc_cur,
    data_t prev0[SEQ_LEN][D_MODEL], data_t next0[SEQ_LEN][D_MODEL],
    data_t prev1[SEQ_LEN][D_MODEL], data_t next1[SEQ_LEN][D_MODEL],
    data_t prev2[SEQ_LEN][D_MODEL], data_t next2[SEQ_LEN][D_MODEL],
    data_t prev3[SEQ_LEN][D_MODEL], data_t next3[SEQ_LEN][D_MODEL],
    int &mc_prev0, int &mc_next0,
    int &mc_prev1, int &mc_next1,
    int &mc_prev2, int &mc_next2,
    int &mc_prev3, int &mc_next3)
{
#pragma HLS INLINE off
  static data_t in_local[SEQ_LEN][D_MODEL];
#pragma HLS BIND_STORAGE variable=in_local type=ram_2p impl=bram

#pragma HLS DATAFLOW
  task_read_input_axis(s_axis_in, in_local);

  // layer0(current) -> next0
  mc_next0 = mc_cur;
  task_run_layer<0>(in_local, next0, mc_cur);

  // layer1(prev0) -> next1
  mc_next1 = mc_prev0;
  task_run_layer<1>(prev0, next1, mc_prev0);

  // layer2(prev1) -> next2
  mc_next2 = mc_prev1;
  task_run_layer<2>(prev1, next2, mc_prev1);

  // layer3(prev2) -> next3
  mc_next3 = mc_prev2;
  task_run_layer<3>(prev2, next3, mc_prev2);

  // layer4(prev3) -> out
  task_run_last_layer_and_stream<4>(prev3, m_axis_out, mc_prev3);
}
#endif

#if (NUM_LAYERS >= 6)
static void step_6layers(
    hls::stream<axis_t>& s_axis_in,
    hls::stream<axis_t>& m_axis_out,
    int mc_cur,
    data_t prev0[SEQ_LEN][D_MODEL], data_t next0[SEQ_LEN][D_MODEL],
    data_t prev1[SEQ_LEN][D_MODEL], data_t next1[SEQ_LEN][D_MODEL],
    data_t prev2[SEQ_LEN][D_MODEL], data_t next2[SEQ_LEN][D_MODEL],
    data_t prev3[SEQ_LEN][D_MODEL], data_t next3[SEQ_LEN][D_MODEL],
    data_t prev4[SEQ_LEN][D_MODEL], data_t next4[SEQ_LEN][D_MODEL],
    int &mc_prev0, int &mc_next0,
    int &mc_prev1, int &mc_next1,
    int &mc_prev2, int &mc_next2,
    int &mc_prev3, int &mc_next3,
    int &mc_prev4, int &mc_next4)
{
#pragma HLS INLINE off
  static data_t in_local[SEQ_LEN][D_MODEL];
#pragma HLS BIND_STORAGE variable=in_local type=ram_2p impl=bram

#pragma HLS DATAFLOW
  task_read_input_axis(s_axis_in, in_local);

  // layer0(current) -> next0
  mc_next0 = mc_cur;
  task_run_layer<0>(in_local, next0, mc_cur);

  // layer1(prev0) -> next1
  mc_next1 = mc_prev0;
  task_run_layer<1>(prev0, next1, mc_prev0);

  // layer2(prev1) -> next2
  mc_next2 = mc_prev1;
  task_run_layer<2>(prev1, next2, mc_prev1);

  // layer3(prev2) -> next3
  mc_next3 = mc_prev2;
  task_run_layer<3>(prev2, next3, mc_prev2);

  // layer4(prev3) -> next4
  mc_next4 = mc_prev3;
  task_run_layer<4>(prev3, next4, mc_prev3);

  // layer5(prev4) -> out
  task_run_last_layer_and_stream<5>(prev4, m_axis_out, mc_prev4);
}
#endif

#endif // NUM_LAYERS

// =============================================================================
// Top kernel
// =============================================================================
extern "C" {
void siso(
    const ap_uint<1> load_parameters,
    const int mask_count,

    hls::stream<axis_t>& s_axis_in,
    hls::stream<axis_t>& m_axis_out,

    const data_t* W_QKV, const data_t* B_QKV,
    const data_t* W_O,   const data_t* B_O,

    const data_t* norm1_gamma, const data_t* norm1_beta,
    const data_t* norm2_gamma, const data_t* norm2_beta,
    const data_t* ff_b1, const data_t* ff_b2,
    const data_t* ff_w1, const data_t* ff_w2)
{
#pragma HLS INTERFACE s_axilite port=load_parameters
#pragma HLS INTERFACE s_axilite port=mask_count
#pragma HLS INTERFACE s_axilite port=return

#pragma HLS INTERFACE axis port=s_axis_in
#pragma HLS INTERFACE axis port=m_axis_out

#pragma HLS INTERFACE m_axi port=W_QKV       offset=slave
#pragma HLS INTERFACE m_axi port=B_QKV       offset=slave
#pragma HLS INTERFACE m_axi port=W_O         offset=slave
#pragma HLS INTERFACE m_axi port=B_O         offset=slave
#pragma HLS INTERFACE m_axi port=norm1_gamma offset=slave
#pragma HLS INTERFACE m_axi port=norm1_beta  offset=slave
#pragma HLS INTERFACE m_axi port=norm2_gamma offset=slave
#pragma HLS INTERFACE m_axi port=norm2_beta  offset=slave
#pragma HLS INTERFACE m_axi port=ff_b1       offset=slave
#pragma HLS INTERFACE m_axi port=ff_b2       offset=slave
#pragma HLS INTERFACE m_axi port=ff_w1       offset=slave
#pragma HLS INTERFACE m_axi port=ff_w2       offset=slave

  // --------------------------------------------------------------------------
  // Parameter load mode
  // --------------------------------------------------------------------------
  if (load_parameters) {
    load_layer_params(W_QKV, B_QKV, W_O, B_O,
                      norm1_gamma, norm1_beta, norm2_gamma, norm2_beta,
                      ff_b1, ff_b2, ff_w1, ff_w2);
    return;
  }

  // Clamp mask_count (0..SEQ_LEN) once per call
  int mc_cur = mask_count;
  if (mc_cur < 0) mc_cur = 0;
  if (mc_cur > SEQ_LEN) mc_cur = SEQ_LEN;

#if (NUM_LAYERS == 1)
  // Simple single-layer run (no state)
  step_single(s_axis_in, m_axis_out, mc_cur);
  return;

#else
  // --------------------------------------------------------------------------
  // Persistent pipeline state (ping-pong banks)
  // --------------------------------------------------------------------------
  static ap_uint<1> toggle = 0;

  // Intermediate states after layers 0..N-2
  static data_t stateA0[SEQ_LEN][D_MODEL], stateB0[SEQ_LEN][D_MODEL];
#if (NUM_LAYERS >= 3)
  static data_t stateA1[SEQ_LEN][D_MODEL], stateB1[SEQ_LEN][D_MODEL];
#endif
#if (NUM_LAYERS >= 4)
  static data_t stateA2[SEQ_LEN][D_MODEL], stateB2[SEQ_LEN][D_MODEL];
#endif
#if (NUM_LAYERS >= 5)
  static data_t stateA3[SEQ_LEN][D_MODEL], stateB3[SEQ_LEN][D_MODEL];
#endif
#if (NUM_LAYERS >= 6)
  static data_t stateA4[SEQ_LEN][D_MODEL], stateB4[SEQ_LEN][D_MODEL];
#endif

#pragma HLS BIND_STORAGE variable=stateA0 type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=stateB0 type=ram_2p impl=bram
#if (NUM_LAYERS >= 3)
#pragma HLS BIND_STORAGE variable=stateA1 type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=stateB1 type=ram_2p impl=bram
#endif
#if (NUM_LAYERS >= 4)
#pragma HLS BIND_STORAGE variable=stateA2 type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=stateB2 type=ram_2p impl=bram
#endif
#if (NUM_LAYERS >= 5)
#pragma HLS BIND_STORAGE variable=stateA3 type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=stateB3 type=ram_2p impl=bram
#endif
#if (NUM_LAYERS >= 6)
#pragma HLS BIND_STORAGE variable=stateA4 type=ram_2p impl=bram
#pragma HLS BIND_STORAGE variable=stateB4 type=ram_2p impl=bram
#endif

  // Ping-pong state for mask_count aligned with state* buffers above
  static int mcA0 = 0, mcB0 = 0;
#if (NUM_LAYERS >= 3)
  static int mcA1 = 0, mcB1 = 0;
#endif
#if (NUM_LAYERS >= 4)
  static int mcA2 = 0, mcB2 = 0;
#endif
#if (NUM_LAYERS >= 5)
  static int mcA3 = 0, mcB3 = 0;
#endif
#if (NUM_LAYERS >= 6)
  static int mcA4 = 0, mcB4 = 0;
#endif

  // --------------------------------------------------------------------------
  // Steady-state: overlap stage0(current) with later stages(previous)
  // via explicit toggle schedules (no pointer mux in DATAFLOW).
  // --------------------------------------------------------------------------
  if (toggle == 0) {
    // prev = A, next = B
#if (NUM_LAYERS == 2)
    step_2layers(s_axis_in, m_axis_out, mc_cur,
                 stateA0, stateB0,
                 mcA0, mcB0);
#elif (NUM_LAYERS == 3)
    step_3layers(s_axis_in, m_axis_out, mc_cur,
                 stateA0, stateB0, stateA1, stateB1,
                 mcA0, mcB0, mcA1, mcB1);
#elif (NUM_LAYERS == 4)
    step_4layers(s_axis_in, m_axis_out, mc_cur,
                 stateA0, stateB0, stateA1, stateB1, stateA2, stateB2,
                 mcA0, mcB0, mcA1, mcB1, mcA2, mcB2);
#elif (NUM_LAYERS == 5)
    step_5layers(s_axis_in, m_axis_out, mc_cur,
                 stateA0, stateB0, stateA1, stateB1, stateA2, stateB2, stateA3, stateB3,
                 mcA0, mcB0, mcA1, mcB1, mcA2, mcB2, mcA3, mcB3);
#elif (NUM_LAYERS == 6)
    step_6layers(s_axis_in, m_axis_out, mc_cur,
                 stateA0, stateB0, stateA1, stateB1, stateA2, stateB2, stateA3, stateB3, stateA4, stateB4,
                 mcA0, mcB0, mcA1, mcB1, mcA2, mcB2, mcA3, mcB3, mcA4, mcB4);
#endif
  } else {
    // prev = B, next = A
#if (NUM_LAYERS == 2)
    step_2layers(s_axis_in, m_axis_out, mc_cur,
                 stateB0, stateA0,
                 mcB0, mcA0);
#elif (NUM_LAYERS == 3)
    step_3layers(s_axis_in, m_axis_out, mc_cur,
                 stateB0, stateA0, stateB1, stateA1,
                 mcB0, mcA0, mcB1, mcA1);
#elif (NUM_LAYERS == 4)
    step_4layers(s_axis_in, m_axis_out, mc_cur,
                 stateB0, stateA0, stateB1, stateA1, stateB2, stateA2,
                 mcB0, mcA0, mcB1, mcA1, mcB2, mcA2);
#elif (NUM_LAYERS == 5)
    step_5layers(s_axis_in, m_axis_out, mc_cur,
                 stateB0, stateA0, stateB1, stateA1, stateB2, stateA2, stateB3, stateA3,
                 mcB0, mcA0, mcB1, mcA1, mcB2, mcA2, mcB3, mcA3);
#elif (NUM_LAYERS == 6)
    step_6layers(s_axis_in, m_axis_out, mc_cur,
                 stateB0, stateA0, stateB1, stateA1, stateB2, stateA2, stateB3, stateA3, stateB4, stateA4,
                 mcB0, mcA0, mcB1, mcA1, mcB2, mcA2, mcB3, mcA3, mcB4, mcA4);
#endif
  }

  toggle ^= 1;
  return;
#endif // NUM_LAYERS
}
}
