TSort32¶

Tile Operation Diagram¶

Introduction¶

Sort each 32-element block of src together with the corresponding indices from idx, and write the sorted value-index pairs into dst. The underlying SFU instruction is VBS32 (vbitsort), which sorts one or more independent 32-element lists in a single call.

Hardware: VBS32 (vbitsort)¶

VBS32 runs on the SFU (not the vector pipeline). One invocation sorts repeat consecutive 32-element blocks, each consisting of 32 values + 32 indices packed as value-index pairs:

void vbitsort(__ubuf__ T *dst,        // sorted value-index pairs out
              __ubuf__ T *src0,        // 32 values per block × repeat
              __ubuf__ uint32_t *src1, // 32 indices per block × repeat
              uint8_t repeat);         // number of 32-element blocks (1..255)

• repeat (cap REPEAT_MAX = 255) is packed into config[63:56].
• Blocks are contiguous, strided by 32 elements: block b reads src0[b*32 : b*32+32] and src1[b*32 : b*32+32], writes dst[b*32*coef : ...] where coef = 2 (float) or 4 (half) — the value-index pair expansion factor.
• Sort order: descending by value; ties broken by smaller index first.

Math Interpretation¶

For each row r, src is processed in independent 32-element blocks. Let block b cover columns 32b … 32b+31, and n_b = min(32, C - 32b) be its valid count.

Sort the pairs by value (descending); output the reordered sequence:

where π is the sort permutation for that block.

Notes: - idx is an input tile (indices permuted with values), not an output. - dst stores sorted value-index pairs, not just sorted values.

C++ Intrinsic¶

Declared in include/pto/common/pto_instr.hpp:

// 3-arg: src must be 32-aligned (validCol % 32 == 0)
template <typename DstTileData, typename SrcTileData, typename IdxTileData>
PTO_INST RecordEvent TSort32(DstTileData &dst, SrcTileData &src, IdxTileData &idx);

// 4-arg: supports non-32-aligned tails (validCol % 32 != 0) via tmp padding
template <typename DstTileData, typename SrcTileData, typename IdxTileData, typename TmpTileData>
PTO_INST RecordEvent TSort32(DstTileData &dst, SrcTileData &src, IdxTileData &idx, TmpTileData &tmp);

Tile Sizes & Data Types¶

For src of shape \(R \times C\) (valid region), block size 32:

dst expansion factor (typeCoef): each input element becomes an 8-byte tuple [value (4 B), index (4 B)] — for float the value fills 4 B; for half the 2-B value is zero-padded to 4 B. So dst is always C × 8 bytes.

Constraints¶

tmp size equation (4-arg)¶

Let \(C\) = validCol, \(b\) = sizeof(T) bytes, \(G\) = 32 (block size). The implementation branches on whether the whole row fits MAX_UB_TMP = 8160; the threshold unit differs between targets:

• ceil_G(C) = \(C\) rounded up to the next multiple of 32.
• A2A3: the threshold is in elements (srcShapeBytesPerRow / sizeof(T) <= MAX_UB_TMP), so \(C \le 8160\) regardless of dtype (float → \(C \le 8160\), half → \(C \le 8160\)).
• A5: the threshold is in bytes (srcShapeBytesPerRow <= MAX_UB_TMP), so \(C \cdot b \le 8160\) (float → \(C \le 2040\), half → \(C \le 4080\)). This is the pto_copy_ubuf_to_ubuf (MOV_UB_TO_UB) repeat cap = 255 blocks × 32 B.
• Tail block = \(t = C \bmod G\) elements (the trailing partial block), extended to \(G\) with \(-\infty\).
• Path A (\(C \cdot b \le 8160\)) copies the entire row from its start into tmp, then pads the last 32 elements in place.
• Path B (\(C \cdot b > 8160\)) copies only the tail block into tmp; full blocks are sorted directly from src.
• VBS32 hard cap: repeat ≤ REPEAT_MAX = 255 blocks per call (≤ 8160 elements); rows longer than 255 blocks are split across multiple vbitsort calls.
• UB placement: tmp should be placed right after dst (32-B aligned), sized ceil(C·b, 32) bytes (equivalently ceil(ceil(C, 32)·b, 32) since \(b \in \{2,4\}\) divides 32) — not at a fixed 8KB offset, since Path A (A2A3) needs up to ~32KB for float near the threshold (\(C \le 8160\) elements = 32 KB for float).

4-arg tail handling¶

When validCol % 32 != 0, the trailing partial block (\(t = C \bmod 32\) elements) must be padded to a full 32-element block before vbitsort. Two paths:

• A2A3: \(C \le 8160\) (elements) / A5: \(C \cdot b \le 8160\) (bytes) (small row): the entire row is copied to tmp, then the last 32 elements are overwritten in place with \(-\infty\) padding via vdup; the row is sorted from tmp.
• A2A3: \(C > 8160\) (elements) / A5: \(C \cdot b > 8160\) (bytes) (large row): only the tail block is copied to tmp and padded; full blocks are sorted directly from src, only the tail is sorted from tmp.

Padding values (\(-\infty\) = -(0.0/0.0)) land at the bottom of the descending order. If validCol > 32 × 255, the row is chunked into REPEAT_MAX-sized groups, each sorted via a separate vbitsort call.

Assembly Syntax¶

AS Level 1 (SSA)¶

%dst = pto.tsort32 %src, %idx : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>

AS Level 2 (DPS)¶

pto.tsort32 ins(%src, %idx : !pto.tile_buf<...>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>)

Examples¶

#include <pto/pto-inst.hpp>
using namespace pto;

// 32-aligned: single block per row
using SrcT = Tile<TileType::Vec, float, 1, 32>;
using IdxT = Tile<TileType::Vec, uint32_t, 1, 32>;
using DstT = Tile<TileType::Vec, float, 1, 64>;   // 2× src cols (float)
SrcT src; IdxT idx; DstT dst;
TSort32(dst, src, idx);

// Non-32-aligned tail: 4-arg with tmp
using SrcT2 = Tile<TileType::Vec, half, 1, 100>;
using IdxT2 = Tile<TileType::Vec, uint32_t, 1, 100>;
using DstT2 = Tile<TileType::Vec, half, 1, 400>;  // 4× src cols (half)
using TmpT  = Tile<TileType::Vec, half, 1, 128>;  // ≥ ceil32(100)=128
TSort32(dst2, src2, idx2, tmp);

ASM Form Examples¶

Auto Mode¶

%dst = pto.tsort32 %src, %idx : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>

Manual Mode¶

# pto.tassign %arg0, @tile(0x1000)
# pto.tassign %arg1, @tile(0x2000)
# pto.tassign %arg2, @tile(0x3000)
%dst = pto.tsort32 %src, %idx : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>

PTO Assembly Form¶

%dst = tsort32 %src, %idx : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>
# AS Level 2 (DPS)
pto.tsort32 ins(%src, %idx : !pto.tile_buf<...>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>)