From 2cd87e02342866a8c04fc85360f4c99796fcbd5f Mon Sep 17 00:00:00 2001 From: LinxISA Automation Date: Tue, 8 Sep 2026 10:05:47 +0800 Subject: [PATCH] TileOP: explicit-groupM overloads across the TMATMUL family Extend the cooperative Local-A/Shared-B explicit-groupM pattern (ADR-0100) from TMATMUL to the rest of the matrix family: TMATMUL_ACC(d, c, a, b, groupM) TMATMUL_BIAS(d, a, b, bias, groupM) TMATMUL_MX(d, a, sa, b, sb, groupM) TMATMUL_MX_ACC(d, c, a, sa, b, sb, groupM) TMATMUL_MX_BIAS(d, a, sa, b, sb, bias, groupM) All share one helper namespace (pto_matmul_groupm_detail) enforcing: - Local-A/Shared-B only (Shared-A derives group_M from its own shape) - D (and C) valid rows equal the per-PE A shard size - runtime group_M in 1..128 (printf + trap, mirroring the ASL dispatch) Also fixes pre-existing MX bugs surfaced by the Shared-scale paths: - basic TMATMUL_MX/MX_ACC/MX_BIAS constrained scale operands to is_tile_data_v, rejecting SharedTile scales that the options overloads already accept; relaxed to typename to match - the Local-A/Shared-B emitters bound a Shared scale with "Tr"(scale.data()) (SharedTile has no data()) and emitted B.IOT for it; now "Sr"(handle()) and B.IOS, matching the options-version convention (PTO_FIXP_MX_SB_SRC) - the 5SRC helper additionally mis-ordered a B.IOT with the Shared scale and the Extra binder; now emits B.IOS scale then B.IOT Extra Verified (fixture CoopGroupMOverloads): - TMATMUL_ACC/BIAS CubeM16 (group_M=64) and CubeM32 (group_M=128): LB0 carries the runtime group_M (B.DIM reg, 0), N/K immediate-form - TMATMUL_MX/MX_ACC/MX_BIAS with Local e4m3 A + Shared B + Shared e8m0 scale: B.IOS S1 scale binder, runtime LB0 - MXScaleVariants/GroupMatmul/SharedMatmul/TMatmulAcc*/PostProcessCombos/ HiF4X2MatmulMX/TGEMV* fixtures and unittest 40/40 all clean (TMatmulAllOptions retains 3 pre-existing match errors, identical on the unmodified baseline) --- include/jcore/template_asm.hpp | 191 ++++++++++++++++++-- test/tileop_api/compile.all | 2 +- test/tileop_api/src/CoopGroupMOverloads.cpp | 61 +++++++ 3 files changed, 241 insertions(+), 13 deletions(-) create mode 100644 test/tileop_api/src/CoopGroupMOverloads.cpp diff --git a/include/jcore/template_asm.hpp b/include/jcore/template_asm.hpp index f33e7b3..eb6410e 100644 --- a/include/jcore/template_asm.hpp +++ b/include/jcore/template_asm.hpp @@ -6054,11 +6054,11 @@ PTO_SHARED_INLINE void Name(Dst &dst, A &a, ScaleA &scale_a, B &b, \ PTO_MATMUL_HEADER(Opcode, PTO_FIXP_ATTR) \ "B.IOS %S[SharedB], mask=1111\n" \ "B.IOT %[A], mask=1111\n" ".if %c[HasScaleA]\n" "B.IOT %[ScaleA], mask=1111\n" ".endif\n" \ - "B.IOT %[ScaleB], mask=1111\n" \ + ".if %c[HasScaleB]\n" "B.IOS %S[ScaleB], mask=1111\n" ".endif\n" \ "B.IOT mask=1111, last, ->%[Dst]<%Z[TileSize]>\n" \ : [Dst] "=&Tr"(dst.data()) \ : [A] "Tr"(a.data()), [ScaleA] "Tr"(scale_a.data()), \ - [SharedB] "Sr"(b.handle()), [ScaleB] "Tr"(scale_b.data()), \ + [SharedB] "Sr"(b.handle()), [ScaleB] "Sr"(scale_b.handle()), \ PTO_FIXP_ATTR_INPUTS, PTO_MX_SCALE_INPUTS, \ PTO_MATMUL_COMMON_INPUTS(Dst, A, B, M, N, K) \ : "memory"); \ @@ -6071,7 +6071,7 @@ PTO_SHARED_INLINE void Name(Dst &dst, A &a, ScaleA &scale_a, B &b, \ "B.IOT mask=1111, last, ->%[Dst]<%Z[TileSize]>\n" \ : [Dst] "=&Tr"(dst.data()) \ : [SharedA] "Sr"(a.handle()), [ScaleA] "Tr"(scale_a.data()), \ - [SharedB] "Sr"(b.handle()), [ScaleB] "Tr"(scale_b.data()), \ + [SharedB] "Sr"(b.handle()), [ScaleB] "Sr"(scale_b.handle()), \ PTO_FIXP_ATTR_INPUTS, PTO_MX_SCALE_INPUTS, \ PTO_MATMUL_COMMON_INPUTS(Dst, A, B, M, N, K) \ : "memory"); \ @@ -6178,11 +6178,12 @@ PTO_SHARED_INLINE void Name(Dst &dst, A &a, ScaleA &scale_a, B &b, \ PTO_MATMUL_HEADER(Opcode, PTO_FIXP_ATTR) \ "B.IOS %S[SharedB], mask=1111\n" \ "B.IOT %[A], mask=1111\n" ".if %c[HasScaleA]\n" "B.IOT %[ScaleA], mask=1111\n" ".endif\n" \ - "B.IOT %[ScaleB], %[Extra], mask=1111\n" \ + "B.IOS %S[ScaleB], mask=1111\n" \ + "B.IOT %[Extra], mask=1111\n" \ "B.IOT mask=1111, last, ->%[Dst]<%Z[TileSize]>\n" \ : [Dst] "=&Tr"(dst.data()) \ : [A] "Tr"(a.data()), [ScaleA] "Tr"(scale_a.data()), \ - [SharedB] "Sr"(b.handle()), [ScaleB] "Tr"(scale_b.data()), \ + [SharedB] "Sr"(b.handle()), [ScaleB] "Sr"(scale_b.handle()), \ [Extra] "Tr"(extra.data()), \ PTO_FIXP_ATTR_INPUTS, PTO_MX_SCALE_INPUTS, \ PTO_MATMUL_COMMON_INPUTS(Dst, A, B, M, N, K) \ @@ -6197,7 +6198,7 @@ PTO_SHARED_INLINE void Name(Dst &dst, A &a, ScaleA &scale_a, B &b, \ "B.IOT mask=1111, last, ->%[Dst]<%Z[TileSize]>\n" \ : [Dst] "=&Tr"(dst.data()) \ : [SharedA] "Sr"(a.handle()), [ScaleA] "Tr"(scale_a.data()), \ - [SharedB] "Sr"(b.handle()), [ScaleB] "Tr"(scale_b.data()), \ + [SharedB] "Sr"(b.handle()), [ScaleB] "Sr"(scale_b.handle()), \ [Extra] "Tr"(extra.data()), \ PTO_FIXP_ATTR_INPUTS, PTO_MX_SCALE_INPUTS, \ PTO_MATMUL_COMMON_INPUTS(Dst, A, B, M, N, K) \ @@ -6216,6 +6217,45 @@ PTO_DEFINE_MATMUL_MX_5SRC_HELPER(matmul_mx_acc, "TMATMULMX.ACC", true) #undef PTO_MATMUL_HEADER } // namespace pto_matmul_detail +namespace pto_matmul_groupm_detail { + +// Explicit-groupM overloads exist only for the cooperative Local-A/Shared-B +// CUBE matrix forms (ADR-0100): each PE holds a [M_per_PE, K] A shard and +// LB0 must carry the core-total group_M, which a Local A tile cannot supply. +template +constexpr void validate_local_a_shared_b(const char *Name) { + (void)Name; + static_assert(!is_shared_tile_v && is_shared_tile_v, + "the explicit groupM overload is only for Local-A/Shared-B " + "cooperative CUBE matrix forms"); + static_assert(tile_role_v == Location::Left, + "cooperative matrix input A must be a Left tile"); + static_assert(tile_role_v == Location::Right, + "cooperative matrix input B must be a Right tile"); +} + +// D (and accumulator C, when present) valid rows must equal the per-PE A +// shard size: A::ValidRow already IS M_per_PE for a Local cooperative A. +template +constexpr void validate_per_pe_destination(const char *Name) { + (void)Name; + static_assert(Dst::ValidRow == A::ValidRow, + "cooperative destination valid Row must match the per-PE " + "Local-A shard (M_per_PE x N)"); +} + +// ASL dispatch: group_M in 1..128; each PE computes +// valid_M = clamp(group_M - i*M_per_PE, 0, M_per_PE). +PTO_SHARED_INLINE void validate_groupm_runtime(const char *Name, + size_t groupM) { + if (groupM < 1 || groupM > 128) { + __builtin_printf("%s: cooperative group_M must be in 1..128\n", Name); + __builtin_trap(); + } +} + +} // namespace pto_matmul_groupm_detail + // TMATMUL: C = A(M,K) * B(K,N). Supported storage combinations are // Local/Local, Local/Shared-Right, and Shared-Left/Shared-Right. A lone Shared @@ -6302,6 +6342,34 @@ PTO_SHARED_INLINE void TMATMUL_ACC(tile_shape_d &d, tile_shape_c &c, tile_shape_ pto_matmul_detail::matmul_acc(d, c, a, b, M, N, K); } +// Cooperative Local-A/Shared-B form: LB0 encodes the core-total group_M, +// which a Local A shard descriptor cannot supply, so the caller passes it +// explicitly. D/C valid rows must equal the per-PE A shard (M_per_PE). +template +PTO_SHARED_INLINE void TMATMUL_ACC(tile_shape_d &d, tile_shape_c &c, tile_shape_a &a, + tile_shape_b &b, size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_ACC"); + pto_matmul_groupm_detail::validate_per_pe_destination("TMATMUL_ACC"); + pto_matmul_groupm_detail::validate_per_pe_destination("TMATMUL_ACC"); + static_assert(is_basic_fixp_attr(Attr), + "TMATMUL_ACC supports only parameter-free FPATR options " + "(keep_acc/f16/bf16/relu); quant, PReLU, RowMax and GroupMax " + "require the overload taking fixp::Options"); + static_assert(!Attr.CScaleEn, + "CScale requires the ACC overload with a CScale tile"); + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(d, a, b); + size_t N = __shape.N; + size_t K = __shape.K; + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_ACC", groupM); + pto_matmul_detail::matmul_acc(d, c, a, b, groupM, N, K); +} + template @@ -6748,6 +6816,29 @@ PTO_SHARED_INLINE void TMATMUL_BIAS(tile_shape_c &c, tile_shape_a &a, tile_shape pto_matmul_detail::matmul_bias(c, a, b, bias, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M. +template +PTO_SHARED_INLINE void TMATMUL_BIAS(tile_shape_c &c, tile_shape_a &a, tile_shape_b &b, + tile_shape_bias &bias, size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_BIAS"); + pto_matmul_groupm_detail::validate_per_pe_destination("TMATMUL_BIAS"); + static_assert(is_basic_fixp_attr(Attr), + "TMATMUL_BIAS supports only parameter-free FPATR options " + "(keep_acc/f16/bf16/relu); quant, PReLU, RowMax and GroupMax " + "require the overload taking fixp::Options"); + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(c, a, b); + size_t N = __shape.N; + size_t K = __shape.K; + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_BIAS", groupM); + pto_matmul_detail::matmul_bias(c, a, b, bias, groupM, N, K); +} + template + typename tile_shape_bscale> PTO_SHARED_INLINE void TMATMUL_MX(tile_shape_c &c, tile_shape_a &a, tile_shape_ascale &ascale, tile_shape_b &b, tile_shape_bscale &bscale) { static_assert(is_basic_fixp_attr(Attr), @@ -6830,6 +6921,28 @@ PTO_SHARED_INLINE void TMATMUL_MX(tile_shape_c &c, tile_shape_a &a, tile_shape_a pto_matmul_detail::matmul_mx(c, a, ascale, b, bscale, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M. +template +PTO_SHARED_INLINE void TMATMUL_MX(tile_shape_c &c, tile_shape_a &a, tile_shape_ascale &ascale, + tile_shape_b &b, tile_shape_bscale &bscale, size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_MX"); + pto_matmul_groupm_detail::validate_per_pe_destination("TMATMUL_MX"); + static_assert(is_basic_fixp_attr(Attr), + "TMATMUL_MX supports only parameter-free FPATR options " + "(keep_acc/f16/bf16/relu); quant, PReLU, RowMax and GroupMax " + "require the overload taking fixp::Options"); + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(c, a, b); + size_t N = __shape.N; + size_t K = __shape.K; + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_MX", groupM); + pto_matmul_detail::matmul_mx(c, a, ascale, b, bscale, groupM, N, K); +} + template + is_local_or_shared_left tile_shape_a, typename tile_shape_sa, + is_local_or_shared_right tile_shape_b, typename tile_shape_sb> PTO_SHARED_INLINE void TMATMUL_MX_ACC(tile_shape_d &d, tile_shape_c &c, tile_shape_a &a, tile_shape_sa &scale_a, tile_shape_b &b, tile_shape_sb &scale_b) { @@ -6917,6 +7030,35 @@ PTO_SHARED_INLINE void TMATMUL_MX_ACC(tile_shape_d &d, tile_shape_c &c, tile_sha pto_matmul_detail::matmul_mx_acc(d, a, scale_a, b, scale_b, c, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M. +template +PTO_SHARED_INLINE void TMATMUL_MX_ACC(tile_shape_d &d, tile_shape_c &c, tile_shape_a &a, + tile_shape_sa &scale_a, tile_shape_b &b, + tile_shape_sb &scale_b, size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_MX_ACC"); + pto_matmul_groupm_detail::validate_per_pe_destination("TMATMUL_MX_ACC"); + pto_matmul_groupm_detail::validate_per_pe_destination("TMATMUL_MX_ACC"); + pto_matmul_detail::validate_matrix_accumulator_contract(); + static_assert(is_basic_fixp_attr(Attr), + "TMATMUL_MX_ACC supports only parameter-free FPATR options " + "(keep_acc/f16/bf16/relu); quant, PReLU, RowMax and GroupMax " + "require the overload taking fixp::Options"); + static_assert(!Attr.CScaleEn, + "CScale requires the MX_ACC overload with a CScale tile"); + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(d, a, b); + size_t N = __shape.N; + size_t K = __shape.K; + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_MX_ACC", groupM); + pto_matmul_detail::matmul_mx_acc(d, a, scale_a, b, scale_b, c, groupM, N, K); +} + template + typename tile_shape_sb, is_tile_data_v tile_shape_bias> PTO_SHARED_INLINE void TMATMUL_MX_BIAS(tile_shape_d &d, tile_shape_a &a, tile_shape_sa &scale_a, tile_shape_b &b, tile_shape_sb &scale_b, tile_shape_bias &bias) { @@ -7006,6 +7148,31 @@ PTO_SHARED_INLINE void TMATMUL_MX_BIAS(tile_shape_d &d, tile_shape_a &a, pto_matmul_detail::matmul_mx_bias(d, a, scale_a, b, scale_b, bias, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M. +template +PTO_SHARED_INLINE void TMATMUL_MX_BIAS(tile_shape_d &d, tile_shape_a &a, + tile_shape_sa &scale_a, tile_shape_b &b, + tile_shape_sb &scale_b, tile_shape_bias &bias, + size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_MX_BIAS"); + pto_matmul_groupm_detail::validate_per_pe_destination("TMATMUL_MX_BIAS"); + static_assert(is_basic_fixp_attr(Attr), + "TMATMUL_MX_BIAS supports only parameter-free FPATR options " + "(keep_acc/f16/bf16/relu); quant, PReLU, RowMax and GroupMax " + "require the overload taking fixp::Options"); + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(d, a, b); + size_t N = __shape.N; + size_t K = __shape.K; + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_MX_BIAS", groupM); + pto_matmul_detail::matmul_mx_bias(d, a, scale_a, b, scale_b, bias, groupM, N, K); +} + template 64) +// TMATMUL_MX_BIAS +#include + +using namespace pto; + +namespace fp32_forms { +// FP32 basic 系列 +using AL16 = Tile; +using AL32 = Tile; +using BS = SharedMatrixRight; +using D16 = Tile; +using D32 = Tile; +using Bias = Tile; + +__attribute__((noinline)) void acc16(D16 &d, D16 &c, AL16 &a, BS &b) { + auto sb = TMOV_L2S_INSERT(b); + TMATMUL_ACC(d, c, a, sb, 64); +} +__attribute__((noinline)) void acc32(D32 &d, D32 &c, AL32 &a, BS &b) { + auto sb = TMOV_L2S_INSERT(b); + TMATMUL_ACC(d, c, a, sb, 128); +} +__attribute__((noinline)) void bias16(D16 &d, AL16 &a, BS &b, Bias &bias) { + auto sb = TMOV_L2S_INSERT(b); + TMATMUL_BIAS(d, a, sb, bias, 64); +} + + + +} // namespace fp32_forms + +namespace mx_forms { +template using MA16 = Tile; +using SA = Tile; +template using ShB = SharedMatrixRight; +using ShSB = SharedMatrixRight<__fp8_e8m0, 1, 16, 1, 16>; +using GMF = global_tensor>; +using GM16 = global_tensor<__fp8_e4m3, RowMajor<16, 32>>; +using GMB = global_tensor<__fp8_e4m3, RowMajor<32, 16>>; +using GMSB = global_tensor<__fp8_e8m0, RowMajor<1, 16>>; +using D16 = Tile; +using Bias = Tile; + +__attribute__((noinline)) void mx(float *out, __fp8_e4m3 *ain, __fp8_e4m3 *bin, __fp8_e8m0 *sbin) { + GMF gm_d(out); GM16 gm_a(ain); GMB gm_b(bin); GMSB gm_sb(sbin); + D16 d; MA16<__fp8_e4m3> a; SA sa; + TLOAD_CUBE(d, gm_d); TLOAD_CUBE(a, gm_a); + auto b = TLOAD>(gm_b); + auto sb = TLOAD(gm_sb); + TMATMUL_MX(d, a, sa, b, sb, 64); +} + + +} // namespace mx_forms + +int main() { return 0; }