Skip to content

fix(tileop): complete atomic TSTORE and cooperative TMOV contracts #298

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
VV0003 merged 1 commit into from
Oct 29, 2026
Merged

fix(tileop): complete atomic TSTORE and cooperative TMOV contracts #298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 5 additions & 1 deletion docs/tileop-usage/concepts/local-and-shared-tile.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,13 @@ TADD(local_dst, local_lhs, local_rhs);
覆盖整个分配容量，不能只描述最小 valid CELL envelope。
- Local CUBE 到 Shared ND 使用 `TMOV_CUBE_TO_SHARED_ND`。单 PE 可一次发布完整
parent；多 PE 必须通过 `range::assemble*` 写入不越界的显式 range，并由
`LAST` 原子发布。同一 generation 不得混用 NORM、M322ND 和 M162ND。
`LAST` 原子发布。同一 generation 的有效区先保存在 pending metadata 中，按
writer offset 映射并合并，`LAST` 时才替换消费者可见的 valid rows/columns。
同一 generation 不得混用 NORM、M322ND 和 M162ND。
- Shared full store 使用 `TSTORE(gm, shared)` 和固定 mask `1111`；部分 store 使用
`TSTORE_PART<PEMask>(gm, shared)`，且 `PEMask` 必须是公开接口接受的非零集合。
- Atomic TSTORE 的 Local CUBE 形式可使用 Predicate execution mask、`PredInv`
和 PE mask；Shared Atomic TSTORE 只接受 PE mask，禁止逐元素 Predicate。

### Shared operand-role view

Expand Down
7 changes: 7 additions & 0 deletions docs/tileop-usage/tlsu/load-store-move/TMOV.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,10 @@ TMOV_CUBE_TO_SHARED_ND(assemble_writer, cube_fragment);
Shared parent capacity。多 PE 生产者必须用 `range::assemble`、
`range::assemble_middle`、`range::assemble_last` 描述明确的 128-byte range；
writer SizeCode 必须等于源 CUBE capacity，最终 `LAST` 才发布完整 parent。
每个 writer 的有效矩形会按其 `B.ASSEMBLE` 线性 offset 投影到父 RowMajor
几何，并合并为 pending valid rows/columns；`INIT`/`MIDDLE` 不替换此前已发布的
有效区，只有 `LAST` 才将完整 generation 的有效区和 payload 一起发布。因而
物理列为 4、有效列为 3 的整行 fragment 拼接后，父对象有效列仍为 3。

源、目标 backing dtype 必须相同，目标列数必须等于源 CUBE 的精确物理列数，
目标 RowMajor 几何必须完整覆盖其 SizeCode capacity，且 writer range 必须落在
Expand All @@ -76,6 +80,9 @@ NORM/M322ND/M162ND 都不合法。编码使用 Function 2 `BSTART.TMOV DTYPE_NON
`B.DATR M322ND/M162ND, DTYPE_NONE, Zero`、Shared destination `B.IOS` 和 Local
source `B.IOT`。

参与同一 cooperative generation 的 PE 必须对有效区使用一致的 SPMD 元数据。
当前 C++ carrier 不对各 PE 私有的动态 valid shape 执行额外标量归约。

操作数角色、数据类型组合、容量、PE mask 和 alias 必须符合上方约束；只能使用所选重载声明的操作数形式。

### 有效区域与 padding
Expand Down
19 changes: 19 additions & 0 deletions docs/tileop-usage/tlsu/load-store-move/TSTORE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,25 @@ TSTORE_ATOMIC_INC(dst, src);
TSTORE_ATOMIC_DEC(dst, src);
```

所有八个接口都可用模板参数选择参与 PE：

```cpp
TSTORE_ATOMIC_ADD<12>(dst, src); // PE mask 1100
TSTORE_ATOMIC_ADD<0>(dst, src); // 架构严格空操作
```

`PEMask` 可为 `0, 1, 2, 4, 8, 12, 14, 15`。Local `CUBE_M16/M32`
还可以使用 Predicate Tile execution mask；第二个模板参数控制反选：

```cpp
TSTORE_ATOMIC_ADD<15, false>(dst, local_cube, predicate);
TSTORE_ATOMIC_ADD<15, true>(dst, local_cube, predicate); // PredInv=1
```

Predicate 的 CUBE layout 和 valid shape 必须与 Local 源一致。该重载固定
`Zero=0`，不提供 zero-fill 形式。RowMajor Local 和 Shared 源不接受 Predicate
参数；Shared Atomic TSTORE 仅支持 PE participation mask。

`src` 可为 Local RowMajor、Shared RowMajor，或 Local `CUBE_M16/M32`；不接受
`CUBE_N8`。GM 与源 dtype 必须相同。支持的 dtype 如下：

Expand Down
17 changes: 17 additions & 0 deletions include/common/pto_predicate.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@ namespace predicate {

template <typename Basis> class Predicate;

namespace detail {
template <typename Basis>
__attribute__((always_inline)) inline decltype(auto)
atomic_tstore_carrier(const Predicate<Basis> &);
} // namespace detail

template <typename Basis, int Columns = Basis::ValidCol>
using LogicalBasis = Tile<Location::Vec, uint8_t,
Basis::BFractal == BLayout::CubeM16 ? 16 : 32,
Expand Down Expand Up @@ -99,11 +105,22 @@ template <typename Basis> class Predicate {
typename Global>
friend __attribute__((always_inline)) inline void
atomic_add(Global &, const Index &, const Value &, const Predicate<Source> &);
template <typename Source>
friend __attribute__((always_inline)) inline decltype(auto)
detail::atomic_tstore_carrier(const Predicate<Source> &);

public:
using BasisTile = Basis;
};

namespace detail {
template <typename Basis>
__attribute__((always_inline)) inline decltype(auto)
atomic_tstore_carrier(const Predicate<Basis> &predicate) {
return (predicate.Carrier);
}
} // namespace detail

#define PTO_PREDICATE_LAYOUT_ASM \
".if %c[Layout] == 31\nB.DATR CUBE_M16, DTYPE_NONE, Zero\n" \
".elseif %c[Layout] == 29\nB.DATR CUBE_M32, DTYPE_NONE, Zero\n.endif\n"
Expand Down
35 changes: 33 additions & 2 deletions include/common/pto_tile.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1773,11 +1773,13 @@ class SharedTile {

SharedTile()
: RowMaskInternal(ValidRow == DYNAMIC ? 0 : ValidRow),
ColMaskInternal(ValidCol == DYNAMIC ? 0 : ValidCol) {}
ColMaskInternal(ValidCol == DYNAMIC ? 0 : ValidCol),
PendingRowMaskInternal(0), PendingColMaskInternal(0) {}

explicit SharedTile(const LocalTile &local)
: RowMaskInternal(local.GetValidRow()),
ColMaskInternal(local.GetValidCol()) {}
ColMaskInternal(local.GetValidCol()), PendingRowMaskInternal(0),
PendingColMaskInternal(0) {}

int GetValidRow() const { return RowMaskInternal; }
int GetValidCol() const { return ColMaskInternal; }
Expand All @@ -1788,6 +1790,23 @@ class SharedTile {
ColMaskInternal = local.GetValidCol();
}

void BeginAssembledValidShape(int rows, int columns) {
PendingRowMaskInternal = rows;
PendingColMaskInternal = columns;
}

void MergeAssembledValidShape(int rows, int columns) {
if (rows > PendingRowMaskInternal)
PendingRowMaskInternal = rows;
if (columns > PendingColMaskInternal)
PendingColMaskInternal = columns;
}

void PublishAssembledValidShape() {
RowMaskInternal = PendingRowMaskInternal;
ColMaskInternal = PendingColMaskInternal;
}

unsigned long &handle_ref() { return Handle; }
unsigned long handle() const { return Handle; }

Expand All @@ -1799,6 +1818,8 @@ class SharedTile {
unsigned long Handle;
int RowMaskInternal;
int ColMaskInternal;
int PendingRowMaskInternal;
int PendingColMaskInternal;
};

// v5 SharedTile traits. is_tile_data_v stays scoped to ordinary Local Tiles
Expand Down Expand Up @@ -2805,6 +2826,16 @@ class Assemble {
int GetValidCol() const { return ParentValue.GetValidCol(); }
uintptr_t GetRangeBase() const { return RangeBaseValue; }

void UpdateAssembledValidShape(int rows, int columns)
requires(is_shared_tile_v<Parent>) {
if constexpr (INIT)
ParentValue.BeginAssembledValidShape(rows, columns);
else
ParentValue.MergeAssembledValidShape(rows, columns);
if constexpr (LAST)
ParentValue.PublishAssembledValidShape();
}

private:
Parent &ParentValue;
uintptr_t RangeBaseValue;
Expand Down
109 changes: 99 additions & 10 deletions include/jcore/template_asm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,15 @@

using namespace pto;

namespace pto::predicate {
template <typename Basis> class Predicate;
namespace detail {
template <typename Basis>
__attribute__((always_inline)) inline decltype(auto)
atomic_tstore_carrier(const Predicate<Basis> &);
} // namespace detail
} // namespace pto::predicate

// Must be defined before any wrapper that carries an opaque Shared value or
// tile register through inline asm; always_inline keeps such values inside
// the caller (they must never cross the ordinary C++ ABI as integer/stack).
Expand All @@ -17,7 +26,8 @@ using namespace pto;
// the compact numeric PEMask template argument. Select one canonical token
// at assembly time without emitting a decimal mask spelling.
#define PTO_PE_MASK_ASM(PREFIX, SUFFIX) \
".if %c[PEMask] == 1\n" PREFIX "0001" SUFFIX \
".if %c[PEMask] == 0\n" PREFIX "0000" SUFFIX \
".elseif %c[PEMask] == 1\n" PREFIX "0001" SUFFIX \
".elseif %c[PEMask] == 2\n" PREFIX "0010" SUFFIX \
".elseif %c[PEMask] == 3\n" PREFIX "0011" SUFFIX \
".elseif %c[PEMask] == 4\n" PREFIX "0100" SUFFIX \
Expand Down Expand Up @@ -4123,10 +4133,12 @@ inline constexpr bool type_legal_v = [] {
return TypeCode == __type_uint32;
}();

template <int Mode, is_global_data_v GM, is_local_tile_v Tile>
template <int Mode, int PEMask, is_global_data_v GM, is_local_tile_v Tile>
requires(!is_range_v<Tile>)
PTO_SHARED_INLINE void emit(GM &dst, const Tile &src) {
constexpr int TypeCode = type_traits<typename Tile::DType>::TypeCode;
static_assert(is_valid_pe_mask(PEMask),
"atomic TSTORE PEMask must be 0,1,2,4,8,12,14,or 15");
static_assert(Tile::BFractal == BLayout::RowMajor ||
Tile::BFractal == BLayout::CubeM16 ||
Tile::BFractal == BLayout::CubeM32,
Expand All @@ -4150,12 +4162,13 @@ PTO_SHARED_INLINE void emit(GM &dst, const Tile &src) {
".endif\n"
"B.DIM %[VCOL], 0, ->lb0\n"
"B.DIM %[VROW], 0, ->lb1\n"
"B.IOT %[Src], mask=1111, last\n"
PTO_PE_MASK_ASM("B.IOT %[Src], mask=", ", last\n")
"B.IOR [%[Base],%[RowStrideBytes]], []\n"
:
: [Base] "r"(dst.data()), [Src] "Tr"(src.data()),
[RowStrideBytes] "r"(dst.GetStrideBytes(3)),
[DataType] "i"(TypeCode), [Layout] "i"(Tile::CubeStoreLayout),
[PEMask] "i"(PEMask),
[Mode] "i"(Mode), [VCOL] "r"(valid_col), [VROW] "r"(valid_row)
: "memory");
} else {
Expand All @@ -4165,22 +4178,24 @@ PTO_SHARED_INLINE void emit(GM &dst, const Tile &src) {
"B.DIM %[VCOL], 0, ->lb0\n"
"B.DIM %[VROW], 0, ->lb1\n"
"B.DIM zero, %c[COL], ->lb2\n"
"B.IOT %[Src], mask=1111, last\n"
PTO_PE_MASK_ASM("B.IOT %[Src], mask=", ", last\n")
"B.IOR [%[Base],%[RowStrideBytes]], []\n"
:
: [Base] "r"(dst.data()), [Src] "Tr"(src.data()),
[RowStrideBytes] "r"(dst.GetStrideBytes(3)),
[DataType] "i"(TypeCode), [Mode] "i"(Mode),
[DataType] "i"(TypeCode), [Mode] "i"(Mode), [PEMask] "i"(PEMask),
[VCOL] "r"(valid_col), [VROW] "r"(valid_row),
[COL] "i"(Tile::PhysicalCols)
: "memory");
}
}

template <int Mode, is_global_data_v GM, is_shared_tile_v Shared>
template <int Mode, int PEMask, is_global_data_v GM, is_shared_tile_v Shared>
PTO_SHARED_INLINE void emit(GM &dst, const Shared &src) {
using Local = typename Shared::LocalTileType;
constexpr int TypeCode = type_traits<typename Local::DType>::TypeCode;
static_assert(is_valid_pe_mask(PEMask),
"atomic TSTORE PEMask must be 0,1,2,4,8,12,14,or 15");
static_assert(std::is_same_v<typename GM::DType, typename Local::DType>,
"atomic Shared TSTORE requires matching GM and source dtypes");
static_assert(Local::isRowMajor && !Local::isBoxedLayout,
Expand All @@ -4197,24 +4212,76 @@ PTO_SHARED_INLINE void emit(GM &dst, const Shared &src) {
"B.DIM %[VCOL], 0, ->lb0\n"
"B.DIM %[VROW], 0, ->lb1\n"
"B.DIM zero, %c[COL], ->lb2\n"
"B.IOS %S[Shared], mask=1111\n"
PTO_PE_MASK_ASM("B.IOS %S[Shared], mask=", "\n")
"B.IOR [%[Base],%[RowStrideBytes]], []\n"
:
: [Base] "r"(dst.data()), [Shared] "Sr"(src.handle()),
[RowStrideBytes] "r"(dst.GetStrideBytes(3)),
[DataType] "i"(TypeCode), [Mode] "i"(Mode),
[DataType] "i"(TypeCode), [Mode] "i"(Mode), [PEMask] "i"(PEMask),
[VCOL] "r"(valid_col), [VROW] "r"(valid_row),
[COL] "i"(Local::PhysicalCols)
: "memory");
}

template <int Mode, int PEMask, bool PredInv, is_global_data_v GM,
is_local_tile_v Tile, typename Basis>
requires(!is_range_v<Tile>)
PTO_SHARED_INLINE void emit(
GM &dst, const Tile &src, const predicate::Predicate<Basis> &active) {
constexpr int TypeCode = type_traits<typename Tile::DType>::TypeCode;
static_assert(is_valid_pe_mask(PEMask),
"atomic TSTORE PEMask must be 0,1,2,4,8,12,14,or 15");
static_assert(Tile::BFractal == BLayout::CubeM16 ||
Tile::BFractal == BLayout::CubeM32,
"masked atomic TSTORE requires a Local CUBE_M16 or CUBE_M32 source");
static_assert(std::is_same_v<typename GM::DType, typename Tile::DType>,
"atomic TSTORE requires matching GM and source dtypes");
static_assert(type_legal_v<Mode, TypeCode>,
"atomic TSTORE mode does not support this DataType");
static_assert(Tile::IsValidActiveSize,
"atomic TSTORE source capacity must be 128 B..256 KiB");
static_assert(Basis::BFractal == Tile::BFractal,
"atomic TSTORE Predicate layout must match the source");
static_assert(Basis::ValidRow == Tile::ValidRow &&
Basis::ValidCol == Tile::ValidCol,
"atomic TSTORE Predicate valid shape must match the source");
const size_t valid_col = src.GetValidCol();
const size_t valid_row = src.GetValidRow();
asm volatile(
"BSTART.TLSU TSTORE, %D[DataType]\n"
".if %c[Layout] == 24\n"
"B.DATR M322ND.canon, DTYPE_NONE, Null, byte0, %c[Mode], RNONE, nosat, %c[PredInv], 0\n"
".else\n"
"B.DATR M162ND.canon, DTYPE_NONE, Null, byte0, %c[Mode], RNONE, nosat, %c[PredInv], 0\n"
".endif\n"
"B.DIM %[VCOL], 0, ->lb0\n"
"B.DIM %[VROW], 0, ->lb1\n"
PTO_PE_MASK_ASM("B.IOT %[Src], %[Predicate], mask=", ", last\n")
"B.IOR [%[Base],%[RowStrideBytes]], []\n"
:
: [Base] "r"(dst.data()), [Src] "Tr"(src.data()),
[Predicate] "Tr"(predicate::detail::atomic_tstore_carrier(active)),
[RowStrideBytes] "r"(dst.GetStrideBytes(3)),
[DataType] "i"(TypeCode), [Layout] "i"(Tile::CubeStoreLayout),
[Mode] "i"(Mode), [PEMask] "i"(PEMask),
[PredInv] "i"(static_cast<int>(PredInv)), [VCOL] "r"(valid_col),
[VROW] "r"(valid_row)
: "memory");
}

} // namespace pto_tstore_atomic_detail

#define PTO_DEFINE_TSTORE_ATOMIC_API(Name, Mode) \
template <is_global_data_v GM, typename Source> \
template <int PEMask = 15, is_global_data_v GM, typename Source> \
requires(is_local_tile_v<Source> || is_shared_tile_v<Source>) \
PTO_SHARED_INLINE void Name(GM &dst, const Source &src) { \
pto_tstore_atomic_detail::emit<Mode>(dst, src); \
pto_tstore_atomic_detail::emit<Mode, PEMask>(dst, src); \
} \
template <int PEMask = 15, bool PredInv = false, is_global_data_v GM, \
is_local_tile_v Source, typename Basis> \
PTO_SHARED_INLINE void Name( \
GM &dst, const Source &src, const predicate::Predicate<Basis> &active) { \
pto_tstore_atomic_detail::emit<Mode, PEMask, PredInv>(dst, src, active); \
}

PTO_DEFINE_TSTORE_ATOMIC_API(TSTORE_ATOMIC_ADD, 0)
Expand Down Expand Up @@ -5031,6 +5098,26 @@ PTO_SHARED_INLINE void TMOV_CUBE_TO_SHARED_ND(
static_assert(type_traits<typename Cube::DType>::TypeCode != __type_rcpe6m2,
"CUBE-to-Shared-ND TMOV does not support RCPE6M2 backing");
const uintptr_t range_base = static_cast<uintptr_t>(dst.GetRangeBase());
const size_t valid_rows = static_cast<size_t>(src.GetValidRow());
const size_t valid_columns = static_cast<size_t>(src.GetValidCol());
const uintptr_t destination_offset =
((range_base + static_cast<uintptr_t>(OffsetUnits)) *
range::RangeAddressUnitBytes * 8) /
tile_element_bits_v<typename ND::DType>;
const size_t destination_start_column = destination_offset % ND::PhysicalCols;
const size_t placed_valid_rows =
valid_rows == 0 || valid_columns == 0
? 0
: (destination_offset + (valid_rows - 1) * ND::PhysicalCols +
valid_columns - 1) /
ND::PhysicalCols +
1;
const size_t placed_valid_columns =
valid_rows == 0 || valid_columns == 0
? 0
: (destination_start_column + valid_columns < ND::PhysicalCols
? destination_start_column + valid_columns
: ND::PhysicalCols);
if constexpr (INIT) {
asm volatile(
"BSTART.TMOV DTYPE_NONE\n"
Expand Down Expand Up @@ -5062,6 +5149,8 @@ PTO_SHARED_INLINE void TMOV_CUBE_TO_SHARED_ND(
[Offset] "i"(OffsetUnits), [WriterSize] "i"(WriterSizeCode)
: "memory");
}
dst.UpdateAssembledValidShape(static_cast<int>(placed_valid_rows),
static_cast<int>(placed_valid_columns));
}

template <int PEMask = 15, is_tile_data_v tile_shape_dst, typename LocalTile>
Expand Down
14 changes: 14 additions & 0 deletions test/test_issue_398_399_contract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,10 @@ def test_atomic_tstore_aliases_use_function_one_fields(self):
self.assertIn("M322ND.canon", ASM)
self.assertIn("M162ND.canon", ASM)
self.assertIn('RNONE, nosat, 0, 0', ASM)
self.assertIn('RNONE, nosat, %c[PredInv], 0', ASM)
self.assertIn('B.IOT %[Src], %[Predicate], mask=', ASM)
self.assertIn('B.IOS %S[Shared], mask=', ASM)
self.assertIn('atomic TSTORE PEMask must be 0,1,2,4,8,12,14,or 15', ASM)
self.assertNotRegex(ASM, r'BSTART\.TLSU\s+TSTORE\.ATOMIC_')

def test_cube_to_shared_nd_uses_canonical_tmov(self):
Expand All @@ -46,6 +50,12 @@ def test_cube_to_shared_nd_uses_canonical_tmov(self):
self.assertIn('[Init] "i"(1)', ASM)
self.assertIn('[Init] "i"(0)', ASM)
self.assertIn('PTO_PE_MASK_ASM("B.IOS %S[Shared], mask=", "\\n")', ASM)
self.assertIn("UpdateAssembledValidShape", ASM)
self.assertIn("BeginAssembledValidShape", TILE)
self.assertIn("PublishAssembledValidShape", TILE)
fixture = (ROOT / "test/tileop_api/src/CubeToSharedND.cpp").read_text()
self.assertIn("PartialCube::PhysicalCols == 4", fixture)
self.assertIn("TSTORE_ATOMIC_ADD<1>(gm, parent)", fixture)

def test_negative_compile_cases_are_registered(self):
driver = (ROOT / "test/tileop_api/run_negatives.sh").read_text()
Expand All @@ -54,6 +64,10 @@ def test_negative_compile_cases_are_registered(self):
"atomic_add_dtype",
"atomic_inc_dtype",
"atomic_cube_n8",
"atomic_mask_rowmajor",
"atomic_mask_shape",
"atomic_mask_shared",
"atomic_pe_mask",
"cube_shared_dtype",
"cube_shared_cols",
"cube_shared_capacity",
Expand Down
Loading
Loading