Skip to content

fix(tileop): support CubeM32 binary slot assembly #289

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 23, 2026
Merged

fix(tileop): support CubeM32 binary slot assembly #289

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
29 changes: 29 additions & 0 deletions docs/tileop-usage/b-subview-b-assemble.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,35 @@ ScaleParent assemble_scales(ScaleBf16 (&scales)[4]) {
}
```

普通计算得到的 `CUBE_M32` fragment 也可以作为二元 producer 的输入，结果
直接写入同布局的 slot。下面的八个 `TADD` 分别携带 INIT、六个 MIDDLE 和
LAST；输入是普通 Tile，不会生成 source `B.SUBVIEW`，也不需要额外的
`TCVT`/`TMOV` 搬运：

```cpp
#include <common/pto_tileop.hpp>

#include <cstdint>
#include <utility>

using namespace pto;

using Fragment = VecTileM32<int32_t, 32, 1>;
using Parent = VecTileM32<int32_t, 32, 8>;

Parent assemble_sums(Fragment &lhs, Fragment &rhs) {
TileArray<Fragment, 1, 8> slots;
for (int col = 0; col < 8; ++col)
TADD(slots[0][col], lhs, rhs);
return TASSEMBLY<Parent>(std::move(slots));
}
```

该普通 Tile 二元路径要求 slot、`src0`、`src1` 都是 descriptor 匹配的
unboxed `CUBE_M32` Local Tile；dtype、location、physical/valid shape 和
fragment capacity 必须一致。`TADD`、`TSUB`、`TMUL`、`TDIV`、`TREM`、
`TAND`、`TOR`、`TXOR`、`TSHL`、`TSHR`、`TMAX`、`TMIN` 共用此路径。

source 侧对应物是 `TPARTVIEW`：`B.ASSEMBLE` 把 producer 的 destination 写进
parent 的某个 slot，`B.SUBVIEW` 则把 parent 的某段借出为 source fragment，
两者互为读写对偶。
Expand Down
159 changes: 151 additions & 8 deletions include/common/pto_tile_region_inline_asm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1984,14 +1984,157 @@ PTO_REGION_ALWAYS_INLINE void pto_region_binary_assemble_dispatch(
}
}

#define PTO_REGION_BINARY_DEST_WRAPPER(Name, Opcode) \
template <typename SubTile, typename Parent0, typename SubTile0, \
typename Parent1, typename SubTile1> \
PTO_REGION_ALWAYS_INLINE void Name( \
region::TileArrayOutputRef<SubTile> dst, \
region::SubTileView<Parent0, SubTile0> &src0, \
region::SubTileView<Parent1, SubTile1> &src1) { \
pto_region_binary_assemble_dispatch<Opcode>(dst, src0, src1); \
// Ordinary computed CUBE_M32 tiles can produce one fragment of a larger
// CUBE_M32 TileArray parent directly. Unlike the overload above, these
// sources are complete tiles rather than B.SUBVIEW carriers, so the producer
// binds them directly and attaches only the destination B.ASSEMBLE modifier.
//
// INIT allocates the full parent capacity. MIDDLE/LAST bind the existing
// ParentRef on the final L=1 record, while B.ASSEMBLE field 5 continues to
// describe the individual fragment writer extent in every phase.
template <int ParentSize, bool Init, bool Last, int Opcode, typename SubTile,
typename Src0, typename Src1>
requires(is_tile_data_v<Src0> && is_tile_data_v<Src1> &&
!is_subtile_view_v<Src0> && !is_subtile_view_v<Src1> &&
!is_range_v<Src0> && !is_range_v<Src1>)
PTO_REGION_ALWAYS_INLINE void pto_region_binary_tile_assemble(
region::TileArrayOutputRef<SubTile> &dst, Src0 &src0, Src1 &src1) {
static_assert(SubTile::BFractal == BLayout::CubeM32 &&
Src0::BFractal == BLayout::CubeM32 &&
Src1::BFractal == BLayout::CubeM32,
"binary TileArray ordinary-tile assembly requires CUBE_M32 "
"destination and sources");
static_assert(SubTile::SFractal == SLayout::NoneBox &&
Src0::SFractal == SLayout::NoneBox &&
Src1::SFractal == SLayout::NoneBox,
"binary TileArray ordinary-tile assembly requires unboxed "
"fragments");
static_assert(SubTile::Loc == Src0::Loc && Src0::Loc == Src1::Loc,
"binary TileArray slot and sources must have matching "
"locations");
static_assert(SubTile::Rows == Src0::Rows && SubTile::Cols == Src0::Cols &&
Src0::Rows == Src1::Rows && Src0::Cols == Src1::Cols,
"binary TileArray slot requires matching physical shapes");
static_assert(SubTile::ValidRow == Src0::ValidRow &&
SubTile::ValidCol == Src0::ValidCol &&
Src0::ValidRow == Src1::ValidRow &&
Src0::ValidCol == Src1::ValidCol,
"binary TileArray slot requires matching valid shapes");
static_assert(std::is_same_v<typename SubTile::DType, typename Src0::DType> &&
std::is_same_v<typename Src0::DType, typename Src1::DType>,
"binary TileArray slot requires matching element types");
static_assert(SubTile::LogicalTileBytes == Src0::LogicalTileBytes &&
Src0::LogicalTileBytes == Src1::LogicalTileBytes,
"binary TileArray slot and sources must have matching "
"capacities");

constexpr int writer_size =
tile_type_traits<typename SubTile::TileDType>::TilesizeCode;
static_assert(writer_size != 0,
"TileArray slot writer extent must be a nonzero size code");
const uintptr_t destination_base_units = dst.range_base_units();

#define PTO_REGION_BINARY_TILE_ASSEMBLY_PREFIX \
"BSTART.TEPL %c[Opcode], %D[Type]\n" PTO_REGION_ELEMENTWISE_LAYOUT_ASM \
"B.DIM zero, %c[ValidCol], ->lb0\n" \
"B.DIM zero, %c[ValidRow], ->lb1\n" \
"B.DIM zero, %c[Cols], ->lb2\n"
#define PTO_REGION_BINARY_TILE_ASSEMBLY_INPUTS \
[Type] "i"(type_traits<typename Src0::DType>::TypeCode), \
[Src0] "Tr"(src0.data()), [Src1] "Tr"(src1.data()), \
[ValidCol] "i"(Src0::ValidCol), [ValidRow] "i"(Src0::ValidRow), \
[Cols] "i"(Src0::Cols), [ParentSize] "i"(ParentSize), \
[Offset] "r"(destination_base_units), [Opcode] "i"(Opcode), \
[WriterSize] "i"(writer_size), [Init] "i"(Init), [Last] "i"(Last), \
[ElemLayout] "i"(local_layout_code_v<SubTile>)
if constexpr (Init) {
asm volatile(PTO_REGION_BINARY_TILE_ASSEMBLY_PREFIX
"B.IOT %[Src0], %[Src1], mask=1111, last, "
"->%[Dst]<%Z[ParentSize]>\n"
"B.ASSEMBLE %c[Init], %c[Last], %[Offset], 0, %c[WriterSize]\n"
: [Dst] "=Tr"(dst.template parent_data<ParentSize>())
: PTO_REGION_BINARY_TILE_ASSEMBLY_INPUTS
: "memory");
} else {
asm volatile(PTO_REGION_BINARY_TILE_ASSEMBLY_PREFIX
"B.IOT %[Src0], %[Src1], mask=1111\n"
"B.IOT %[Dst], mask=1111, last\n"
"B.ASSEMBLE %c[Init], %c[Last], %[Offset], 0, %c[WriterSize]\n"
:
: [Dst] "Tr"(dst.template parent_data<ParentSize>()),
PTO_REGION_BINARY_TILE_ASSEMBLY_INPUTS
: "memory");
}
#undef PTO_REGION_BINARY_TILE_ASSEMBLY_INPUTS
#undef PTO_REGION_BINARY_TILE_ASSEMBLY_PREFIX
}

template <int ParentSize, int Opcode, typename SubTile, typename Src0,
typename Src1>
requires(is_tile_data_v<Src0> && is_tile_data_v<Src1> &&
!is_subtile_view_v<Src0> && !is_subtile_view_v<Src1> &&
!is_range_v<Src0> && !is_range_v<Src1>)
PTO_REGION_ALWAYS_INLINE void pto_region_binary_tile_assemble_phase(
region::TileArrayOutputRef<SubTile> &dst, Src0 &src0, Src1 &src1) {
if (dst.slot_count() == 1)
pto_region_binary_tile_assemble<ParentSize, true, true, Opcode>(dst, src0,
src1);
else if (dst.ordinal() == 0)
pto_region_binary_tile_assemble<ParentSize, true, false, Opcode>(dst, src0,
src1);
else if (dst.ordinal() == dst.slot_count() - 1)
pto_region_binary_tile_assemble<ParentSize, false, true, Opcode>(dst, src0,
src1);
else
pto_region_binary_tile_assemble<ParentSize, false, false, Opcode>(dst, src0,
src1);
}

template <int Opcode, typename SubTile, typename Src0, typename Src1>
requires(is_tile_data_v<Src0> && is_tile_data_v<Src1> &&
!is_subtile_view_v<Src0> && !is_subtile_view_v<Src1> &&
!is_range_v<Src0> && !is_range_v<Src1>)
PTO_REGION_ALWAYS_INLINE void pto_region_binary_tile_assemble_dispatch(
region::TileArrayOutputRef<SubTile> &dst, Src0 &src0, Src1 &src1) {
switch (dst.parent_size_code()) {
#define PTO_REGION_BINARY_TILE_PARENT_CASE(N) \
case N: \
pto_region_binary_tile_assemble_phase<N, Opcode>(dst, src0, src1); \
break
PTO_REGION_BINARY_TILE_PARENT_CASE(1);
PTO_REGION_BINARY_TILE_PARENT_CASE(2);
PTO_REGION_BINARY_TILE_PARENT_CASE(3);
PTO_REGION_BINARY_TILE_PARENT_CASE(4);
PTO_REGION_BINARY_TILE_PARENT_CASE(5);
PTO_REGION_BINARY_TILE_PARENT_CASE(6);
PTO_REGION_BINARY_TILE_PARENT_CASE(7);
PTO_REGION_BINARY_TILE_PARENT_CASE(8);
PTO_REGION_BINARY_TILE_PARENT_CASE(9);
PTO_REGION_BINARY_TILE_PARENT_CASE(10);
PTO_REGION_BINARY_TILE_PARENT_CASE(11);
PTO_REGION_BINARY_TILE_PARENT_CASE(12);
#undef PTO_REGION_BINARY_TILE_PARENT_CASE
default:
__builtin_trap();
}
}

#define PTO_REGION_BINARY_DEST_WRAPPER(Name, Opcode) \
template <typename SubTile, typename Parent0, typename SubTile0, \
typename Parent1, typename SubTile1> \
PTO_REGION_ALWAYS_INLINE void Name( \
region::TileArrayOutputRef<SubTile> dst, \
region::SubTileView<Parent0, SubTile0> &src0, \
region::SubTileView<Parent1, SubTile1> &src1) { \
pto_region_binary_assemble_dispatch<Opcode>(dst, src0, src1); \
} \
template <typename SubTile, typename Src0, typename Src1> \
requires(is_tile_data_v<Src0> && is_tile_data_v<Src1> && \
!is_subtile_view_v<Src0> && !is_subtile_view_v<Src1> && \
!is_range_v<Src0> && !is_range_v<Src1>) \
PTO_REGION_ALWAYS_INLINE void Name(region::TileArrayOutputRef<SubTile> dst, \
Src0 &src0, Src1 &src1) { \
pto_region_binary_tile_assemble_dispatch<Opcode>(dst, src0, src1); \
}

PTO_REGION_BINARY_DEST_WRAPPER(TADD, 0)
Expand Down
2 changes: 1 addition & 1 deletion test/tileop_api/compile.all
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ ACTIVE_FIXTURES=(
TMatmulAccFullOptions TMatmulAccCScale TMatmulAllOptions TGEMVAllOptions TGEMVOperandOrder
PostProcessCombos SharedMatrixForms TCmpMode TColSumContract TSort TPrefetch MGatherCas MGatherAtomics MScatterAtomics
TMrgSort TQuant TImg2Col GroupMatmul TStoreShared SharedTLoad TAbs
DirectCubeLayout TileRegionBinaryMixed Issue166RmsNormB128 Issue269ReductionCapacity Issue272PackPrefix Issue274PackViews Issue275PackAssemble Issue276ReinterpretMul TCvtSameTypeRounding
DirectCubeLayout TileRegionBinaryMixed Issue166RmsNormB128 Issue269ReductionCapacity Issue272PackPrefix Issue274PackViews Issue275PackAssemble Issue276ReinterpretMul Issue288CubeM32BinaryAssembly TCvtSameTypeRounding
TAdd_mask TAdd TAdds TAnd TeplAss TeplAssSpecial TOrAssSubview TLoad TStore TCvt TDiv TDivs TExp TExpdif TMax TMaxs
RangeSubview RangeAssemble SharedRange TileRegionSharedSubview TileArray TileArrayCube TileArrayAssemblyOffsets TileArrayRegionAsm TileArrayTCVTE8M0 TileArrayWriterSize TileRegionUnaryAssembly TileRegion TileRegionUnary TileRegionBinary TileRegionBinaryAssembly TileRegionUnarySubviewAssembly TileRegionScalarAssembly TileRegionTCVTSubviewAssembly TileRegionTCVTSubview TileRegionShift TileRegionCubeSubview SubviewVecCube Issue138Subview Issue160ReductionPrefix Issue241ReductionPrefixBinary Issue187DynamicPrefix Issue178CubeTcvtNorm Issue195ConvCompose Issue194Timg2colNchw SubviewVecCubeIssue144
ValidShapeImmediate ValidShapePerDim ValidShapePerDimSweep CoopGroupMOverloads CubeInternalAcc TSELCanonical Issue200TcmpsGpr Issue203CubeTcvtTileArray IssueA2CubeMxFp4Tcvt Issue111SharedSlot
Expand Down
11 changes: 11 additions & 0 deletions test/tileop_api/run_negatives.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ TEXPDIF_CASES="type narrow n8 mixed_layout shape"
TCI2D_CASES="layout location dtype shape steps m16_rows"
I64_CASES="fp64_element unsupported_parent m16 short_m32 subview_pair assemble_pair writer_limit"
TCOLSUM_CASES="dest_rows dest_cols dtype"
ISSUE288_CASES="layout dtype shape valid_shape"
PASS=0; FAIL=0

# A negative suite is meaningless when every compile is rejected before the
Expand Down Expand Up @@ -108,6 +109,16 @@ for c in $TCOLSUM_CASES; do
expect_rejected "tcolsum_$c" "$define" \
TColSumNegatives.cpp "TCOLSUM"
done
for c in $ISSUE288_CASES; do
define=SHOULD_FAIL_ISSUE288_$(echo "$c" | tr '[:lower:]' '[:upper:]')
case "$c" in
layout) pattern='requires CUBE_M32 destination and sources' ;;
dtype) pattern='requires matching element types' ;;
shape) pattern='requires matching physical shapes' ;;
valid_shape) pattern='requires matching valid shapes' ;;
esac
expect_rejected "issue288_$c" "$define" Issue288Negatives.cpp "$pattern"
done
expect_rejected "thistogram_retired" "SHOULD_FAIL_THISTOGRAM_RETIRED" \
TColSumNegatives.cpp "THISTOGRAM is retired"
for c in $TS_CASES; do
Expand Down
34 changes: 34 additions & 0 deletions test/tileop_api/src/Issue288CubeM32BinaryAssembly.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
#include <common/pto_tileop.hpp>

#include <cstdint>
#include <utility>

using namespace pto;

using Fragment = VecTileM32<int32_t, 32, 1>;
using Parent = VecTileM32<int32_t, 32, 8>;

static_assert(Fragment::LogicalTileBytes == 128);
static_assert(Parent::LogicalTileBytes == 1024);

__attribute__((noinline)) Parent assemble_binary_cube_m32(Fragment &lhs,
Fragment &rhs) {
TileArray<Fragment, 1, 8> slots;
TADD(slots[0][0], lhs, rhs);
TADD(slots[0][1], lhs, rhs);
TADD(slots[0][2], lhs, rhs);
TADD(slots[0][3], lhs, rhs);
TADD(slots[0][4], lhs, rhs);
TADD(slots[0][5], lhs, rhs);
TADD(slots[0][6], lhs, rhs);
TADD(slots[0][7], lhs, rhs);
return TASSEMBLY<Parent>(std::move(slots));
}

int main() {
Fragment lhs;
Fragment rhs;
Parent result = assemble_binary_cube_m32(lhs, rhs);
asm volatile("" : : "Tr"(result.data()));
return 0;
}
35 changes: 35 additions & 0 deletions test/tileop_api/src/Issue288Negatives.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
#include <common/pto_tileop.hpp>

#include <cstdint>

using namespace pto;

using Fragment = VecTileM32<int32_t, 32, 1>;

#if defined(SHOULD_FAIL_ISSUE288_LAYOUT)
using BadSlot = Tile<Location::Vec, int32_t, 32, 1, BLayout::RowMajor>;
void bad_layout(Fragment &lhs, Fragment &rhs) {
TileArray<BadSlot, 1, 2> slots;
TADD(slots[0][0], lhs, rhs);
}
#elif defined(SHOULD_FAIL_ISSUE288_DTYPE)
using OtherDType = VecTileM32<uint32_t, 32, 1>;
void bad_dtype(Fragment &lhs, OtherDType &rhs) {
TileArray<Fragment, 1, 2> slots;
TADD(slots[0][0], lhs, rhs);
}
#elif defined(SHOULD_FAIL_ISSUE288_SHAPE)
using OtherShape = VecTileM32<int32_t, 32, 2>;
void bad_shape(Fragment &lhs, OtherShape &rhs) {
TileArray<Fragment, 1, 2> slots;
TADD(slots[0][0], lhs, rhs);
}
#elif defined(SHOULD_FAIL_ISSUE288_VALID_SHAPE)
using OtherValid = VecTileM32<int32_t, 32, 1, 16, 1>;
void bad_valid_shape(Fragment &lhs, OtherValid &rhs) {
TileArray<Fragment, 1, 2> slots;
TADD(slots[0][0], lhs, rhs);
}
#endif

int main() { return 0; }
50 changes: 50 additions & 0 deletions test/tileop_api/verify_target_cxx_frontend.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ for source in RangeSubview.cpp GMov.cpp TileRegionUnary.cpp \
Issue241ReductionPrefixBinary.cpp \
Issue253RowExpandDualSubview.cpp \
Issue268CooperativeRowMax.cpp \
Issue288CubeM32BinaryAssembly.cpp \
TOrAssSubview.cpp \
TileArrayTCVTE8M0.cpp \
SharedTransposeNonSquare.cpp TCI.cpp TRowExpandMul.cpp; do
Expand All @@ -62,6 +63,55 @@ for source in TileRegionUnary TileRegionCubeSubview Issue166RmsNormB128; do
"$TC_DIR/clang++" "${FLAGS[@]}" -c \
"$ROOT/test/tileop_api/src/$source.cpp" -o "$OUT/$source.o"
done
"$TC_DIR/clang++" "${FLAGS[@]}" -S -emit-llvm \
"$ROOT/test/tileop_api/src/Issue288CubeM32BinaryAssembly.cpp" \
-o "$OUT/Issue288CubeM32BinaryAssembly.ll"
"$TC_DIR/clang++" "${FLAGS[@]}" -c \
"$ROOT/test/tileop_api/src/Issue288CubeM32BinaryAssembly.cpp" \
-o "$OUT/Issue288CubeM32BinaryAssembly.o"
python3 - "$OUT/Issue288CubeM32BinaryAssembly.ll" <<'PY_ISSUE288'
import re
import sys
from pathlib import Path

text = Path(sys.argv[1]).read_text(encoding="utf-8")
body = next((part for part in re.split(r"(?=^define )", text,
flags=re.MULTILINE)
if "assemble_binary_cube_m32" in part), None)
if body is None:
raise SystemExit("Issue288: missing binary CubeM32 assembly fixture")
ops = [line for line in body.splitlines()
if "asm sideeffect" in line and "BSTART.TEPL" in line]
if len(ops) != 8:
raise SystemExit(f"Issue288: expected 8 producers, got {len(ops)}")
if any("B.DATR CUBE_M32" not in op for op in ops):
raise SystemExit("Issue288: every producer must preserve CUBE_M32 layout")
if any("B.SUBVIEW" in op or "TCVT" in op or "TMOV" in op for op in ops):
raise SystemExit("Issue288: ordinary sources must not be materialized or subviewed")
if any("B.ASSEMBLE" not in op for op in ops):
raise SystemExit("Issue288: every producer must write an assembly slot")
if "mask=1111, last, ->$0<${7:Z}>" not in ops[0]:
raise SystemExit("Issue288: INIT producer must allocate the parent")
for op in ops[1:]:
between = op.split("B.IOT", 1)[1].split("B.ASSEMBLE", 1)[0]
if between.count("B.IOT") != 1 or "mask=1111, last" not in between:
raise SystemExit("Issue288: continuation must bind the existing ParentRef last")
# The operand tail is ParentSize, Offset, Opcode, WriterSize, INIT, LAST,
# Layout. ParentSize is the 1 KiB parent (4); WriterSize is 128 bytes (1).
tails = []
for op in ops:
match = re.search(
r"i32 4, i64 ([0-7]), i32 0, i32 1, "
r"i1 (true|false), i1 (true|false), i32 29\)", op)
if match is None:
raise SystemExit("Issue288: parent capacity or writer extent was lost")
tails.append((int(match.group(1)), match.group(2), match.group(3)))
expected = [(0, "true", "false")]
expected += [(offset, "false", "false") for offset in range(1, 7)]
expected += [(7, "false", "true")]
if tails != expected:
raise SystemExit(f"Issue288: incorrect offsets or lifecycle {tails}")
PY_ISSUE288
"$TC_DIR/clang++" "${FLAGS[@]}" -S -emit-llvm \
"$ROOT/test/tileop_api/src/RangeAssemble.cpp" -o "$OUT/RangeAssemble.ll"
"$TC_DIR/clang++" "${FLAGS[@]}" -c \
Expand Down
Loading