Skip to content

B70 attention opt #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ShengchenJ merged 8 commits into from
Oct 19, 2026
Merged

B70 attention opt #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
921063c
build: detect the local GPU target, 256-GRF mode on Battlemage
ShengchenJ Oct 13, 2026
f94e4b8
base_types: bf16_2 conversion reuses the scalar converter
ShengchenJ Oct 13, 2026
2550b22
reductions: transposing butterfly for register-tile row reductions
ShengchenJ Oct 13, 2026
4fdd0b8
reductions: register-region transpose-reduce for float max/sum
ShengchenJ Oct 13, 2026
b2cf57e
gqa_forward: use wg256/seq_q8 on Battlemage
ShengchenJ Oct 13, 2026
5e921a4
B70 attention: hoisted 2D block-load payloads for K/V
ShengchenJ Oct 13, 2026
27b206a
B70 attention: cooperative K/V prefetch across subgroups
ShengchenJ Oct 13, 2026
c5c93ba
B70 attention: derive K/V block-load layout from gl strides
ShengchenJ Oct 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion include/common/base_types.dp.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,7 +477,10 @@ template <> struct convertor<sycl::float2, bf16_2> {
};
template <> struct convertor<bf16_2, sycl::float2> {
static inline bf16_2 convert(const sycl::float2 &u) {
return sycl::vec<sycl::ext::oneapi::bfloat16, 2>(u.x(), u.y());
// Reuse the scalar converter: the bf16_2 vec constructor lowers to a
// software RNE sequence (shr/and/add3 + NaN branch) on Xe2.
return bf16_2(convertor<bf16, float>::convert(u.x()),
convertor<bf16, float>::convert(u.y()));
}
};
template <> struct convertor<float, sycl::half> {
Expand Down
63 changes: 41 additions & 22 deletions include/ops/warp/memory/tile/gl_desc.dp.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,9 @@
* - Register cost: each live `gl_desc` is ~1 GRF. GEMM runs near the 256-GRF
* ceiling — share one descriptor per operand-stream and advance X/Y rather
* than holding many; re-run spillcheck after adopting.
* - The `.load()` inline asm is unvalidated in THIS tree until you build+run
* it on PVC. Confirm bit-exact against the flat load before trusting it.
* - Validated on B70 (Xe2) by the attention forward K/V loads (transposed,
* VNNI and prefetch), which match the flat path's output. PVC builds but
* has not been run; confirm against the flat load there before trusting it.
* ============================================================================
*/

Expand All @@ -60,20 +61,26 @@ namespace kittens {
* @tparam VNNI true selects the transform/VNNI read (`nt`, matches
* Subgroup2DBlockLoadTransformINTEL — the B operand);
* false selects the plain read (`nn` — the A operand).
* @tparam Transpose true selects the transposed read (`tn`, matches
* Subgroup2DBlockLoadTransposeINTEL; NumBlocks must be 1).
*
* Coordinates are ELEMENT coordinates (identical to the {col,row} the flat
* intrinsic takes): X = contiguous/column element index, Y = strided/row
* element index.
*/
template <int Bits, int BlockHeight, int BlockWidth, int NumBlocks,
bool VNNI = false>
bool VNNI = false, bool Transpose = false>
struct gl_desc {
static_assert(!(VNNI && Transpose), "VNNI and transposed reads are exclusive");
static_assert(!Transpose || NumBlocks == 1, "transposed reads take one block");

int *payload = nullptr; ///< the block2d address payload (one GRF, built once)

static constexpr int sg_size = 16;
static constexpr int atom_width = BlockWidth * NumBlocks; ///< total contiguous elems
// Per-lane storage vector matching the flat load's deposited bytes.
using elem_t = std::conditional_t<(Bits <= 16), short, int>;
// VNNI packs 32/Bits rows into each dword, so it is read back as int.
using elem_t = std::conditional_t<(Bits <= 16 && !VNNI), short, int>;
static constexpr int lane_count =
(atom_width * BlockHeight * Bits / sg_size) /
static_cast<int>(sizeof(elem_t) * 8);
Expand All @@ -92,14 +99,18 @@ struct gl_desc {

template <ducks::gl::all GL> inline void init(const GL &g) {
using U = typename GL::dtype;
const long base = reinterpret_cast<long>(g.raw_ptr);
const int width_bytes = static_cast<int>(g.template shape<3>() * sizeof(U));
const int height = static_cast<int>(g.template shape<2>());
const int pitch_bytes =
static_cast<int>(g.template stride<2>() * sizeof(U));
init(g.raw_ptr, static_cast<int>(g.template shape<3>() * sizeof(U)),
static_cast<int>(g.template shape<2>()),
static_cast<int>(g.template stride<2>() * sizeof(U)));
}

/// Build the payload for an arbitrary 2D surface (sizes in bytes / rows).
inline void init(const void *base, int width_bytes, int height,
int pitch_bytes) {
payload = __builtin_IB_subgroup_createBlock2DAddressPayload(
base, width_bytes - 1, height - 1, pitch_bytes - 1,
/*blockX=*/0, /*blockY=*/0, BlockWidth, BlockHeight, NumBlocks);
reinterpret_cast<long>(base), width_bytes - 1, height - 1,
pitch_bytes - 1, /*blockX=*/0, /*blockY=*/0, BlockWidth, BlockHeight,
NumBlocks);
}
Comment thread
ShengchenJ marked this conversation as resolved.

/// Set both surface coordinates (element units).
Expand Down Expand Up @@ -127,22 +138,29 @@ struct gl_desc {
/**
* @brief Load the block(s) into register-tile storage. `dst` is the same
* pointer the flat path writes (e.g. &reg.tiles[i][j].data[0]); the
* deposited bytes are bit-identical. Mirrors sycl-tla XE_LOAD_2D /
* XE_LOAD_2D_VNNI.
* deposited bytes are bit-identical. DX/DY are compile-time offsets
* (in `Bits`-sized elements / rows) added to the current X/Y, so a
* fixed pattern of loads around one coordinate needs no payload
* update. Mirrors sycl-tla XE_LOAD_2D / _VNNI / _TRANSPOSE.
*/
template <typename T> inline void load(T *dst) const {
template <int DX = 0, int DY = 0, typename T> inline void load(T *dst) const {
#ifdef __SYCL_DEVICE_ONLY__
auto &dv = *reinterpret_cast<store_vec_t *>(dst);
if constexpr (VNNI) {
asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nt flat[%1+(0,0)]"
if constexpr (Transpose) {
asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4tn flat[%1+(%5,%6)]"
: "=rw"(dv)
: "rw.u"(payload), "P"(Bits), "P"(BlockWidth), "P"(BlockHeight),
"P"(DX), "P"(DY));
} else if constexpr (VNNI) {
asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nt flat[%1+(%6,%7)]"
: "=rw"(dv)
: "rw.u"(payload), "P"(Bits), "P"(NumBlocks), "P"(BlockWidth),
"P"(BlockHeight));
"P"(BlockHeight), "P"(DX), "P"(DY));
} else {
asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nn flat[%1+(0,0)]"
asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nn flat[%1+(%6,%7)]"
: "=rw"(dv)
: "rw.u"(payload), "P"(Bits), "P"(NumBlocks), "P"(BlockWidth),
"P"(BlockHeight));
"P"(BlockHeight), "P"(DX), "P"(DY));
}
#else
(void)dst;
Expand All @@ -162,10 +180,11 @@ struct gl_desc {
* load coordinate is clobbered. See the migration recipe. Each extra
* descriptor is ~1 GRF; re-run spillcheck.
*/
inline void prefetch() const {
template <int DX = 0, int DY = 0> inline void prefetch() const {
#ifdef __SYCL_DEVICE_ONLY__
asm("lsc_load_block2d.ugm.ca.ca (M1, 1) %%null:d%1.%2x%3nn flat[%0+(0,0)]"
:: "rw.u"(payload), "P"(Bits), "P"(atom_width), "P"(BlockHeight));
asm("lsc_load_block2d.ugm.ca.ca (M1, 1) %%null:d%1.%2x%3nn flat[%0+(%4,%5)]"
:: "rw.u"(payload), "P"(Bits), "P"(atom_width), "P"(BlockHeight),
"P"(DX), "P"(DY));
#endif
}
};
Expand Down
214 changes: 172 additions & 42 deletions include/ops/warp/register/tile/reductions.dp.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,154 @@ inline float reduce_16x16(const sycl::float2 *first_half,
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when reset is false.
*/
#ifdef __SYCL_DEVICE_ONLY__
// Register-region transpose-reduce of 16 (or 8) per-lane row partials across the
// 16 SIMD lanes: lane L gets the reduction of row L. Every stage is a plain SIMD16
// mov/op on GRF regions, so there are no a0-indirect shuffles (which serialize on
// Xe). Same data movement as sycl-tla's hreduce16/8_float_* (BSD-3-Clause,
// Copyright 2025 Intel Corporation). For 8 rows, lanes 8-15 are left undefined.
typedef float sk_float16v __attribute__((ext_vector_type(16)));
typedef float sk_float8v __attribute__((ext_vector_type(8)));
#define SK_DEFINE_HREDUCE16(op) \
static inline float sk_hreduce16_##op(sk_float16v x) { \
float y; \
asm("{\n" \
".decl SKR_S v_type=G type=F num_elts=512 align=wordx32\n" \
".decl SKR_T v_type=G type=F num_elts=256 align=wordx32\n" \
"mov (M1_NM,16) SKR_S(0,0)<1> %1(0,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(1,0)<1> %1(0,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(2,0)<1> %1(2,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(3,0)<1> %1(2,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(4,0)<1> %1(4,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(5,0)<1> %1(4,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(6,0)<1> %1(6,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(7,0)<1> %1(6,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(8,0)<1> %1(8,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(9,0)<1> %1(8,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(10,0)<1> %1(10,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(11,0)<1> %1(10,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(12,0)<1> %1(12,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(13,0)<1> %1(12,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(14,0)<1> %1(14,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(15,0)<1> %1(14,8)<16;8,1>\n" \
#op " (M1_NM,16) SKR_T(0,0)<1> SKR_S(0,0)<1;1,0> SKR_S(1,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(1,0)<1> SKR_S(2,0)<1;1,0> SKR_S(3,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(2,0)<1> SKR_S(4,0)<1;1,0> SKR_S(5,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(3,0)<1> SKR_S(6,0)<1;1,0> SKR_S(7,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(4,0)<1> SKR_S(8,0)<1;1,0> SKR_S(9,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(5,0)<1> SKR_S(10,0)<1;1,0> SKR_S(11,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(6,0)<1> SKR_S(12,0)<1;1,0> SKR_S(13,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(7,0)<1> SKR_S(14,0)<1;1,0> SKR_S(15,0)<1;1,0>\n" \
"mov (M1_NM,16) SKR_S(16,0)<1> SKR_T(0,0)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(17,0)<1> SKR_T(0,4)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(18,0)<1> SKR_T(2,0)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(19,0)<1> SKR_T(2,4)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(20,0)<1> SKR_T(4,0)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(21,0)<1> SKR_T(4,4)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(22,0)<1> SKR_T(6,0)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(23,0)<1> SKR_T(6,4)<8;4,1>\n" \
#op " (M1_NM,16) SKR_T(8,0)<1> SKR_S(16,0)<1;1,0> SKR_S(17,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(9,0)<1> SKR_S(18,0)<1;1,0> SKR_S(19,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(10,0)<1> SKR_S(20,0)<1;1,0> SKR_S(21,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(11,0)<1> SKR_S(22,0)<1;1,0> SKR_S(23,0)<1;1,0>\n" \
"mov (M1_NM,16) SKR_S(24,0)<1> SKR_T(8,0)<4;2,1>\n" \
"mov (M1_NM,16) SKR_S(25,0)<1> SKR_T(8,2)<4;2,1>\n" \
"mov (M1_NM,16) SKR_S(26,0)<1> SKR_T(10,0)<4;2,1>\n" \
"mov (M1_NM,16) SKR_S(27,0)<1> SKR_T(10,2)<4;2,1>\n" \
#op " (M1_NM,16) SKR_T(12,0)<1> SKR_S(24,0)<1;1,0> SKR_S(25,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(13,0)<1> SKR_S(26,0)<1;1,0> SKR_S(27,0)<1;1,0>\n" \
"mov (M1_NM,16) SKR_S(28,0)<1> SKR_T(12,0)<2;1,0>\n" \
"mov (M1_NM,16) SKR_S(29,0)<1> SKR_T(12,1)<2;1,0>\n" \
#op " (M1_NM,16) %0(0,0)<1> SKR_S(28,0)<1;1,0> SKR_S(29,0)<1;1,0>\n" \
"}\n" : "=rw"(y) : "rw"(x)); \
return y; \
}

#define SK_DEFINE_HREDUCE8(op) \
static inline float sk_hreduce8_##op(sk_float8v x) { \
float y; \
asm("{\n" \
".decl SKR_S v_type=G type=F num_elts=256 align=wordx32\n" \
".decl SKR_T v_type=G type=F num_elts=128 align=wordx32\n" \
"mov (M1_NM,16) SKR_S(0,0)<1> %1(0,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(1,0)<1> %1(0,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(2,0)<1> %1(2,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(3,0)<1> %1(2,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(4,0)<1> %1(4,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(5,0)<1> %1(4,8)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(6,0)<1> %1(6,0)<16;8,1>\n" \
"mov (M1_NM,16) SKR_S(7,0)<1> %1(6,8)<16;8,1>\n" \
#op " (M1_NM,16) SKR_T(0,0)<1> SKR_S(0,0)<1;1,0> SKR_S(1,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(1,0)<1> SKR_S(2,0)<1;1,0> SKR_S(3,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(2,0)<1> SKR_S(4,0)<1;1,0> SKR_S(5,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(3,0)<1> SKR_S(6,0)<1;1,0> SKR_S(7,0)<1;1,0>\n" \
"mov (M1_NM,16) SKR_S(8,0)<1> SKR_T(0,0)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(9,0)<1> SKR_T(0,4)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(10,0)<1> SKR_T(2,0)<8;4,1>\n" \
"mov (M1_NM,16) SKR_S(11,0)<1> SKR_T(2,4)<8;4,1>\n" \
#op " (M1_NM,16) SKR_T(4,0)<1> SKR_S(8,0)<1;1,0> SKR_S(9,0)<1;1,0>\n" \
#op " (M1_NM,16) SKR_T(5,0)<1> SKR_S(10,0)<1;1,0> SKR_S(11,0)<1;1,0>\n" \
"mov (M1_NM,16) SKR_S(12,0)<1> SKR_T(4,0)<4;2,1>\n" \
"mov (M1_NM,16) SKR_S(13,0)<1> SKR_T(4,2)<4;2,1>\n" \
#op " (M1_NM,16) SKR_T(6,0)<1> SKR_S(12,0)<1;1,0> SKR_S(13,0)<1;1,0>\n" \
"mov (M1_NM,8) SKR_S(14,0)<1> SKR_T(6,0)<2;1,0>\n" \
"mov (M1_NM,8) SKR_S(15,0)<1> SKR_T(6,1)<2;1,0>\n" \
#op " (M1_NM,8) %0(0,0)<1> SKR_S(14,0)<1;1,0> SKR_S(15,0)<1;1,0>\n" \
"}\n" : "=rw"(y) : "rw"(x)); \
return y; \
}
SK_DEFINE_HREDUCE16(max)
SK_DEFINE_HREDUCE16(add)
SK_DEFINE_HREDUCE8(max)
SK_DEFINE_HREDUCE8(add)
#undef SK_DEFINE_HREDUCE16
#undef SK_DEFINE_HREDUCE8
#endif

/**
* @brief Transposing butterfly: reduce ROWS per-lane row partials across the 16
* SIMD lanes so that lane L ends with the reduction of row L (L % ROWS).
*
* Each stage swaps half of the remaining rows with the partner lane, so all rows
* are reduced together: ROWS-1 shuffles (plus one fold for ROWS=8) instead of
* 4 per row. Shuffles lower to serialized a0-indirect moves on Xe, so this
* removes ~3/4 of them from softmax row reductions.
*/
template <typename op, typename type, int ROWS>
static inline type transpose_reduce_rows(type (&a)[ROWS], int lane) {
static_assert(ROWS == 16 || ROWS == 8, "16 or 8 rows per SIMD16 tile");
#ifdef __SYCL_DEVICE_ONLY__
constexpr bool is_max = std::is_same_v<op, base_ops::max>;
constexpr bool is_sum = std::is_same_v<op, base_ops::sum>;
if constexpr (std::is_same_v<type, float> && (is_max || is_sum)) {
if constexpr (ROWS == 16) {
sk_float16v v;
#pragma unroll
for (int r = 0; r < 16; r++) v[r] = a[r];
return is_max ? sk_hreduce16_max(v) : sk_hreduce16_add(v);
} else {
sk_float8v v;
#pragma unroll
for (int r = 0; r < 8; r++) v[r] = a[r];
return is_max ? sk_hreduce8_max(v) : sk_hreduce8_add(v);
}
}
#endif
#pragma unroll
for (int s = ROWS / 2; s >= 1; s >>= 1) {
const bool hi = (lane & s) != 0;
#pragma unroll
for (int r = 0; r < s; r++) {
const type keep = hi ? a[r + s] : a[r];
const type send = hi ? a[r] : a[r + s];
a[r] = op::template op<type>(keep, packed_shfl_sync(MASK_ALL, send, lane ^ s));
}
}
if constexpr (ROWS == 8) // rows 0-7 are spread over both lane halves: fold them
a[0] = op::template op<type>(a[0], packed_shfl_sync(MASK_ALL, a[0], lane ^ 8));
return a[0];
}

template <typename op, ducks::rv::all V, ducks::rt::row_layout T, bool reset>
static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) {
// I actually like these static asserts because they give more verbose errors
Expand Down Expand Up @@ -66,56 +214,38 @@ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) {
for (int i = 0; i < src.height; i++) {
#pragma unroll
for (int j = 0; j < src.width; j++) {
type tile_result{};
// Per-lane row partials: rows[2k] = data[k].x, rows[2k+1] = data[k].y
// (the 32-column shapes first combine their two 16-column halves in-lane),
// then one transposing butterfly leaves row L's reduction in lane L.
type tile_result;
if constexpr (std::is_same_v<typename T::shape,
ducks::rt_shape::rt_8x32>) {
// 8x32: 8 rows, data[0-3]=cols0-15, data[4-7]=cols16-31 (offset +4,
// not +8). Only 4 packed pairs -> fills lanes 0-7 (rows 0-7).
// 8x32: 8 rows, data[0-3]=cols0-15, data[4-7]=cols16-31 (offset +4).
type rows[8];
#pragma unroll
for (int k = 0; k < 4; k++) {
type val_x = op::template op<type>(src.tiles[i][j].data[k].x(),
src.tiles[i][j].data[k + 4].x());
type val_y = op::template op<type>(src.tiles[i][j].data[k].y(),
src.tiles[i][j].data[k + 4].y());
// XOR butterfly reduction across 16 SIMD lanes (32 cols -> scalar)
#pragma unroll
for (int stride = 8; stride >= 1; stride >>= 1) {
val_x = op::template op<type>(
val_x, packed_shfl_sync(MASK_ALL, val_x, lane ^ stride));
val_y = op::template op<type>(
val_y, packed_shfl_sync(MASK_ALL, val_y, lane ^ stride));
}
if (lane == 2 * k) tile_result = val_x;
if (lane == 2 * k + 1) tile_result = val_y;
rows[2 * k] = op::template op<type>(src.tiles[i][j].data[k].x(),
src.tiles[i][j].data[k + 4].x());
rows[2 * k + 1] = op::template op<type>(src.tiles[i][j].data[k].y(),
src.tiles[i][j].data[k + 4].y());
}
tile_result = transpose_reduce_rows<op>(rows, lane);
} else {
type rows[16];
#pragma unroll
for (int k = 0; k < 8; k++) {
type val_x, val_y;
if constexpr (std::is_same_v<typename T::shape,
ducks::rt_shape::rt_16x32>) {
// 16x32: combine two 16-column halves within-lane first
val_x = op::template op<type>(src.tiles[i][j].data[k].x(),
src.tiles[i][j].data[k + 8].x());
val_y = op::template op<type>(src.tiles[i][j].data[k].y(),
src.tiles[i][j].data[k + 8].y());
} else {
val_x = src.tiles[i][j].data[k].x();
val_y = src.tiles[i][j].data[k].y();
}
// XOR butterfly reduction across 16 SIMD lanes
#pragma unroll
for (int stride = 8; stride >= 1; stride >>= 1) {
val_x = op::template op<type>(
val_x, packed_shfl_sync(MASK_ALL, val_x, lane ^ stride));
val_y = op::template op<type>(
val_y, packed_shfl_sync(MASK_ALL, val_y, lane ^ stride));
for (int k = 0; k < 8; k++) {
if constexpr (std::is_same_v<typename T::shape,
ducks::rt_shape::rt_16x32>) {
rows[2 * k] = op::template op<type>(src.tiles[i][j].data[k].x(),
src.tiles[i][j].data[k + 8].x());
rows[2 * k + 1] = op::template op<type>(src.tiles[i][j].data[k].y(),
src.tiles[i][j].data[k + 8].y());
} else {
rows[2 * k] = src.tiles[i][j].data[k].x();
rows[2 * k + 1] = src.tiles[i][j].data[k].y();
}
}
// All lanes now have reduce(row 2k) in val_x, reduce(row 2k+1) in val_y.
// Scatter: lane 2k stores reduce(row 2k), lane 2k+1 stores reduce(row 2k+1).
if (lane == 2 * k) tile_result = val_x;
if (lane == 2 * k + 1) tile_result = val_y;
}
tile_result = transpose_reduce_rows<op>(rows, lane);
}
if (j == 0) {
if constexpr (reset) {
Expand Down
Loading
Loading