diff --git a/include/jcore/template_asm.hpp b/include/jcore/template_asm.hpp index b3ec767..8cd3240 100644 --- a/include/jcore/template_asm.hpp +++ b/include/jcore/template_asm.hpp @@ -6469,6 +6469,111 @@ PTO_SHARED_INLINE void TMATMUL_ACC(tile_shape_d &d, tile_shape_c &c, tile_shape_ quant_gpr, lrelu_gpr, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M +// (options variant; LB0 encodes group_M, which a Local A shard cannot supply). +template +PTO_SHARED_INLINE void TMATMUL_ACC(tile_shape_d &d, tile_shape_c &c, tile_shape_a &a, + tile_shape_b &b, const Options &options, + size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_ACC"); + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_ACC", groupM); + constexpr FixpAttr Attr = Options::Attr; + static_assert(tile_role_v == Location::Left && + tile_role_v == Location::Right, + "TMATMUL_ACC requires A=Left and B=Right"); + static_assert(is_valid_fixp_attr(Attr), "invalid B.FPATR configuration"); + static_assert(is_fixp_output_type(), + "TMATMUL_ACC destination dtype does not match PreQuantMode"); + + constexpr bool HasVectorQuant = + is_vector_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasScalarQuant = + is_scalar_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasRowIn = Attr.RowMaxInit; + constexpr bool HasRowOut = Attr.RowMaxEn; + constexpr bool HasGroupOut = Attr.GroupMaxEn; + constexpr bool HasPRelu = Attr.Relu == FixpReluMode::PRelu; + constexpr int SrcMask = (HasRowIn ? 1 : 0) | (HasVectorQuant ? 2 : 0) | + (HasPRelu ? 4 : 0); + constexpr int OutMask = (HasRowOut ? 1 : 0) | (HasGroupOut ? 2 : 0); + constexpr int IorMode = (HasScalarQuant ? 1 : 0) | + (Attr.Relu == FixpReluMode::LRelu ? 2 : 0); + constexpr int EffectiveM = is_shared_tile_v && Attr.TransA + ? tile_shape_a::ValidCol : tile_shape_a::ValidRow; + constexpr int EffectiveN = is_shared_tile_v && Attr.TransB + ? tile_shape_b::ValidRow : tile_shape_b::ValidCol; + // Reduction outputs (RowMax/GroupMax) reduce the per-PE D rows: ASL + // MatrixRowMaxResult iterates input.valid_rows, which for a cooperative + // TMATMUL is the per-PE clamp of group_M, not the core-total group_M. + // For cooperative forms the reduction/output rows are the per-PE M. + // Shared-A: per-PE block = rows_per_pe(group_M) derived from the core + // total. Local-A/Shared-B: A::ValidRow already IS the per-PE shard size + // (M_per_PE itself, per ADR-0100), so it must be used directly — feeding + // it back through rows_per_pe() would map M_per_PE=32 to 16 and break + // CubeM32 group_M>64 configurations. + constexpr int RedRows = + (is_shared_tile_v) + ? pto_matmul_detail::cooperative_group_m_rows_per_pe(EffectiveM) + : EffectiveM; + + static_assert(HasVectorQuant == + !std::is_same_v, + "vector PreQuant mode requires a quant parameter Tile"); + static_assert(HasPRelu == + !std::is_same_v, + "PRelu mode requires a PReLU parameter Tile"); + static_assert(HasRowIn == + !std::is_same_v, + "RowMaxInit requires a RowMaxIn Tile"); + static_assert(HasRowOut == + !std::is_same_v, + "RowMaxEn requires a RowMaxOut Tile"); + static_assert(HasGroupOut == + !std::is_same_v, + "GroupMaxEn requires a GroupMaxOut Tile"); + + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(d, a, b); + size_t M = __shape.M; + size_t N = __shape.N; + size_t K = __shape.K; + + auto &row_in = pto_matmul_detail::select_fixp_operand(options.RowIn, d); + auto &quant_tile = pto_matmul_detail::select_fixp_operand(options.Quant, d); + auto &relu_tile = pto_matmul_detail::select_fixp_operand(options.Relu, d); + auto &row_out = pto_matmul_detail::select_fixp_operand(options.RowOut, d); + auto &group_out = pto_matmul_detail::select_fixp_operand(options.GroupOut, d); + auto &cscale = pto_matmul_detail::select_fixp_operand( + options.CScale, c); + + // ASL B.FPATR: PreQuant=None and Relu!=LRelu consume no scalar + // parameter at all (BundleFPATRModeUsesScalarParameter(0)=false), so + // materialising the zero descriptors would only produce dead + // sdi/ldi round-trips (issue: keep_acc zero-descriptor dead code). + // Materialise the GPR values only when the IOR schema reads them. + uint64_t quant_gpr_storage; // addresses stable only when used + uint64_t lrelu_gpr_storage; + [[maybe_unused]] volatile uint64_t &quant_gpr_v = quant_gpr_storage; + [[maybe_unused]] volatile uint64_t &lrelu_gpr_v = lrelu_gpr_storage; + if constexpr (IorMode != 0) { + quant_gpr_storage = options.QuantDescriptor; + lrelu_gpr_storage = options.LReluDescriptor; + } + const uint64_t quant_gpr = quant_gpr_storage; + const uint64_t lrelu_gpr = lrelu_gpr_storage; + pto_matmul_detail::emit_matmul_acc_fixp( + d, c, a, b, cscale, row_in, quant_tile, relu_tile, row_out, group_out, + quant_gpr, lrelu_gpr, groupM, N, K); +} + // Cooperative Local-A/Shared-B + options form: LB0 encodes the core-total // group_M provided by the caller, while A still denotes the per-PE Local shard // and B remains the Shared KxN right operand. This keeps the existing options @@ -6901,6 +7006,74 @@ PTO_SHARED_INLINE void TMATMUL_BIAS(tile_shape_c &c, tile_shape_a &a, tile_shape pto_matmul_detail::emit_matmul_bias_fixp(c, a, b, bias, row_in, quant_tile, relu_tile, row_out, group_out, quant_gpr, lrelu_gpr, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M +// (options variant; LB0 encodes group_M, which a Local A shard cannot supply). +template +PTO_SHARED_INLINE void TMATMUL_BIAS(tile_shape_c &c, tile_shape_a &a, tile_shape_b &b, + tile_shape_bias &bias, const Options &options, + size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_BIAS"); + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_BIAS", groupM); + constexpr FixpAttr Attr = Options::Attr; + static_assert(is_valid_fixp_attr(Attr), "invalid B.FPATR configuration"); + static_assert(is_fixp_output_type(), + "TMATMUL_BIAS destination dtype does not match PreQuantMode"); + + constexpr bool HasVectorQuant = is_vector_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasScalarQuant = is_scalar_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasRowIn = Attr.RowMaxInit; + constexpr bool HasRowOut = Attr.RowMaxEn; + constexpr bool HasGroupOut = Attr.GroupMaxEn; + constexpr bool HasPRelu = Attr.Relu == FixpReluMode::PRelu; + constexpr int SrcMask = (HasRowIn ? 1 : 0) | (HasVectorQuant ? 2 : 0) | (HasPRelu ? 4 : 0); + constexpr int OutMask = (HasRowOut ? 1 : 0) | (HasGroupOut ? 2 : 0); + constexpr int IorMode = (HasScalarQuant ? 1 : 0) | (Attr.Relu == FixpReluMode::LRelu ? 2 : 0); + + static_assert(HasVectorQuant == !std::is_same_v, + "vector PreQuant mode requires a quant parameter Tile"); + static_assert(HasPRelu == !std::is_same_v, + "PRelu mode requires a PReLU parameter Tile"); + static_assert(HasRowIn == !std::is_same_v, + "RowMaxInit requires a RowMaxIn Tile"); + static_assert(HasRowOut == !std::is_same_v, + "RowMaxEn requires a RowMaxOut Tile"); + static_assert(HasGroupOut == !std::is_same_v, + "GroupMaxEn requires a GroupMaxOut Tile"); + + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(c, a, b); + size_t M = __shape.M; + size_t N = __shape.N; + size_t K = __shape.K; + + auto &row_in = pto_matmul_detail::select_fixp_operand(options.RowIn, c); + auto &quant_tile = pto_matmul_detail::select_fixp_operand(options.Quant, c); + auto &relu_tile = pto_matmul_detail::select_fixp_operand(options.Relu, c); + auto &row_out = pto_matmul_detail::select_fixp_operand(options.RowOut, c); + auto &group_out = pto_matmul_detail::select_fixp_operand(options.GroupOut, c); + + // ASL B.FPATR: PreQuant=None and Relu!=LRelu consume no scalar + // parameter at all (BundleFPATRModeUsesScalarParameter(0)=false), so + // materialising the zero descriptors would only produce dead + // sdi/ldi round-trips (issue: keep_acc zero-descriptor dead code). + // Materialise the GPR values only when the IOR schema reads them. + uint64_t quant_gpr_storage; // addresses stable only when used + uint64_t lrelu_gpr_storage; + [[maybe_unused]] volatile uint64_t &quant_gpr_v = quant_gpr_storage; + [[maybe_unused]] volatile uint64_t &lrelu_gpr_v = lrelu_gpr_storage; + if constexpr (IorMode != 0) { + quant_gpr_storage = options.QuantDescriptor; + lrelu_gpr_storage = options.LReluDescriptor; + } + const uint64_t quant_gpr = quant_gpr_storage; + const uint64_t lrelu_gpr = lrelu_gpr_storage; + pto_matmul_detail::emit_matmul_bias_fixp(c, a, b, bias, row_in, quant_tile, relu_tile, row_out, group_out, quant_gpr, lrelu_gpr, groupM, N, K); +} + // TMATMUL_MX: C = (A * aScale) * (B * bScale) (BSTART.CUBE TMATMULMX). template (c, a, ascale, b, bscale, row_in, quant_tile, relu_tile, row_out, group_out, quant_gpr, lrelu_gpr, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M +// (options variant; LB0 encodes group_M, which a Local A shard cannot supply). +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, + const Options &options, + size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_MX"); + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_MX", groupM); + constexpr FixpAttr Attr = Options::Attr; + static_assert(is_valid_fixp_attr(Attr), "invalid B.FPATR configuration"); + static_assert(is_fixp_output_type(), + "TMATMUL_MX destination dtype does not match PreQuantMode"); + + constexpr bool HasVectorQuant = is_vector_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasScalarQuant = is_scalar_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasRowIn = Attr.RowMaxInit; + constexpr bool HasRowOut = Attr.RowMaxEn; + constexpr bool HasGroupOut = Attr.GroupMaxEn; + constexpr bool HasPRelu = Attr.Relu == FixpReluMode::PRelu; + constexpr int SrcMask = (HasRowIn ? 1 : 0) | (HasVectorQuant ? 2 : 0) | (HasPRelu ? 4 : 0); + constexpr int OutMask = (HasRowOut ? 1 : 0) | (HasGroupOut ? 2 : 0); + constexpr int IorMode = (HasScalarQuant ? 1 : 0) | (Attr.Relu == FixpReluMode::LRelu ? 2 : 0); + + static_assert(HasVectorQuant == !std::is_same_v, + "vector PreQuant mode requires a quant parameter Tile"); + static_assert(HasPRelu == !std::is_same_v, + "PRelu mode requires a PReLU parameter Tile"); + static_assert(HasRowIn == !std::is_same_v, + "RowMaxInit requires a RowMaxIn Tile"); + static_assert(HasRowOut == !std::is_same_v, + "RowMaxEn requires a RowMaxOut Tile"); + static_assert(HasGroupOut == !std::is_same_v, + "GroupMaxEn requires a GroupMaxOut Tile"); + + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(c, a, b); + size_t M = __shape.M; + size_t N = __shape.N; + size_t K = __shape.K; + + auto &row_in = pto_matmul_detail::select_fixp_operand(options.RowIn, c); + auto &quant_tile = pto_matmul_detail::select_fixp_operand(options.Quant, c); + auto &relu_tile = pto_matmul_detail::select_fixp_operand(options.Relu, c); + auto &row_out = pto_matmul_detail::select_fixp_operand(options.RowOut, c); + auto &group_out = pto_matmul_detail::select_fixp_operand(options.GroupOut, c); + + // ASL B.FPATR: PreQuant=None and Relu!=LRelu consume no scalar + // parameter at all (BundleFPATRModeUsesScalarParameter(0)=false), so + // materialising the zero descriptors would only produce dead + // sdi/ldi round-trips (issue: keep_acc zero-descriptor dead code). + // Materialise the GPR values only when the IOR schema reads them. + uint64_t quant_gpr_storage; // addresses stable only when used + uint64_t lrelu_gpr_storage; + [[maybe_unused]] volatile uint64_t &quant_gpr_v = quant_gpr_storage; + [[maybe_unused]] volatile uint64_t &lrelu_gpr_v = lrelu_gpr_storage; + if constexpr (IorMode != 0) { + quant_gpr_storage = options.QuantDescriptor; + lrelu_gpr_storage = options.LReluDescriptor; + } + const uint64_t quant_gpr = quant_gpr_storage; + const uint64_t lrelu_gpr = lrelu_gpr_storage; + pto_matmul_detail::emit_matmul_mx_fixp(c, a, ascale, b, bscale, row_in, quant_tile, relu_tile, row_out, group_out, quant_gpr, lrelu_gpr, groupM, N, K); +} + 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, const Options &options, + size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_MX_ACC"); + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_MX_ACC", groupM); + constexpr FixpAttr Attr = Options::Attr; + pto_matmul_detail::validate_matrix_accumulator_contract(); + static_assert(is_valid_fixp_attr(Attr), "invalid B.FPATR configuration"); + static_assert(is_fixp_output_type(), + "TMATMUL_MX_ACC destination dtype does not match PreQuantMode"); + + constexpr bool HasVectorQuant = is_vector_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasScalarQuant = is_scalar_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasRowIn = Attr.RowMaxInit; + constexpr bool HasRowOut = Attr.RowMaxEn; + constexpr bool HasGroupOut = Attr.GroupMaxEn; + constexpr bool HasPRelu = Attr.Relu == FixpReluMode::PRelu; + constexpr int SrcMask = (HasRowIn ? 1 : 0) | (HasVectorQuant ? 2 : 0) | (HasPRelu ? 4 : 0); + constexpr int OutMask = (HasRowOut ? 1 : 0) | (HasGroupOut ? 2 : 0); + constexpr int IorMode = (HasScalarQuant ? 1 : 0) | (Attr.Relu == FixpReluMode::LRelu ? 2 : 0); + + static_assert(HasVectorQuant == !std::is_same_v, + "vector PreQuant mode requires a quant parameter Tile"); + static_assert(HasPRelu == !std::is_same_v, + "PRelu mode requires a PReLU parameter Tile"); + static_assert(HasRowIn == !std::is_same_v, + "RowMaxInit requires a RowMaxIn Tile"); + static_assert(HasRowOut == !std::is_same_v, + "RowMaxEn requires a RowMaxOut Tile"); + static_assert(HasGroupOut == !std::is_same_v, + "GroupMaxEn requires a GroupMaxOut Tile"); + + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(d, a, b); + size_t M = __shape.M; + size_t N = __shape.N; + size_t K = __shape.K; + + auto &row_in = pto_matmul_detail::select_fixp_operand(options.RowIn, d); + auto &quant_tile = pto_matmul_detail::select_fixp_operand(options.Quant, d); + auto &relu_tile = pto_matmul_detail::select_fixp_operand(options.Relu, d); + auto &row_out = pto_matmul_detail::select_fixp_operand(options.RowOut, d); + auto &group_out = pto_matmul_detail::select_fixp_operand(options.GroupOut, d); + auto &cscale = pto_matmul_detail::select_fixp_operand( + options.CScale, c); + + // ASL B.FPATR: PreQuant=None and Relu!=LRelu consume no scalar + // parameter at all (BundleFPATRModeUsesScalarParameter(0)=false), so + // materialising the zero descriptors would only produce dead + // sdi/ldi round-trips (issue: keep_acc zero-descriptor dead code). + // Materialise the GPR values only when the IOR schema reads them. + uint64_t quant_gpr_storage; // addresses stable only when used + uint64_t lrelu_gpr_storage; + [[maybe_unused]] volatile uint64_t &quant_gpr_v = quant_gpr_storage; + [[maybe_unused]] volatile uint64_t &lrelu_gpr_v = lrelu_gpr_storage; + if constexpr (IorMode != 0) { + quant_gpr_storage = options.QuantDescriptor; + lrelu_gpr_storage = options.LReluDescriptor; + } + const uint64_t quant_gpr = quant_gpr_storage; + const uint64_t lrelu_gpr = lrelu_gpr_storage; + pto_matmul_detail::emit_matmul_mx_acc_fixp( + d, c, a, scale_a, b, scale_b, cscale, row_in, quant_tile, relu_tile, + row_out, group_out, quant_gpr, lrelu_gpr, groupM, N, K); +} + template (d, a, scale_a, b, scale_b, bias, row_in, quant_tile, relu_tile, row_out, group_out, quant_gpr, lrelu_gpr, M, N, K); } +// Cooperative Local-A/Shared-B form with explicit core-total group_M +// (options variant; LB0 encodes group_M, which a Local A shard cannot supply). +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, + const Options &options, + size_t groupM) { + pto_matmul_groupm_detail::validate_local_a_shared_b("TMATMUL_MX_BIAS"); + pto_matmul_groupm_detail::validate_groupm_runtime("TMATMUL_MX_BIAS", groupM); + constexpr FixpAttr Attr = Options::Attr; + static_assert(is_valid_fixp_attr(Attr), "invalid B.FPATR configuration"); + static_assert(is_fixp_output_type(), + "TMATMUL_MX_BIAS destination dtype does not match PreQuantMode"); + + constexpr bool HasVectorQuant = is_vector_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasScalarQuant = is_scalar_fixp_pre_quant(Attr.PreQuant); + constexpr bool HasRowIn = Attr.RowMaxInit; + constexpr bool HasRowOut = Attr.RowMaxEn; + constexpr bool HasGroupOut = Attr.GroupMaxEn; + constexpr bool HasPRelu = Attr.Relu == FixpReluMode::PRelu; + constexpr int SrcMask = (HasRowIn ? 1 : 0) | (HasVectorQuant ? 2 : 0) | (HasPRelu ? 4 : 0); + constexpr int OutMask = (HasRowOut ? 1 : 0) | (HasGroupOut ? 2 : 0); + constexpr int IorMode = (HasScalarQuant ? 1 : 0) | (Attr.Relu == FixpReluMode::LRelu ? 2 : 0); + + static_assert(HasVectorQuant == !std::is_same_v, + "vector PreQuant mode requires a quant parameter Tile"); + static_assert(HasPRelu == !std::is_same_v, + "PRelu mode requires a PReLU parameter Tile"); + static_assert(HasRowIn == !std::is_same_v, + "RowMaxInit requires a RowMaxIn Tile"); + static_assert(HasRowOut == !std::is_same_v, + "RowMaxEn requires a RowMaxOut Tile"); + static_assert(HasGroupOut == !std::is_same_v, + "GroupMaxEn requires a GroupMaxOut Tile"); + + pto_matmul_detail::MatmulShape __shape = + pto_matmul_detail::resolve_matmul_shape_runtime(d, a, b); + size_t M = __shape.M; + size_t N = __shape.N; + size_t K = __shape.K; + + auto &row_in = pto_matmul_detail::select_fixp_operand(options.RowIn, d); + auto &quant_tile = pto_matmul_detail::select_fixp_operand(options.Quant, d); + auto &relu_tile = pto_matmul_detail::select_fixp_operand(options.Relu, d); + auto &row_out = pto_matmul_detail::select_fixp_operand(options.RowOut, d); + auto &group_out = pto_matmul_detail::select_fixp_operand(options.GroupOut, d); + + // ASL B.FPATR: PreQuant=None and Relu!=LRelu consume no scalar + // parameter at all (BundleFPATRModeUsesScalarParameter(0)=false), so + // materialising the zero descriptors would only produce dead + // sdi/ldi round-trips (issue: keep_acc zero-descriptor dead code). + // Materialise the GPR values only when the IOR schema reads them. + uint64_t quant_gpr_storage; // addresses stable only when used + uint64_t lrelu_gpr_storage; + [[maybe_unused]] volatile uint64_t &quant_gpr_v = quant_gpr_storage; + [[maybe_unused]] volatile uint64_t &lrelu_gpr_v = lrelu_gpr_storage; + if constexpr (IorMode != 0) { + quant_gpr_storage = options.QuantDescriptor; + lrelu_gpr_storage = options.LReluDescriptor; + } + const uint64_t quant_gpr = quant_gpr_storage; + const uint64_t lrelu_gpr = lrelu_gpr_storage; + pto_matmul_detail::emit_matmul_mx_bias_fixp(d, a, scale_a, b, scale_b, bias, row_in, quant_tile, relu_tile, row_out, group_out, quant_gpr, lrelu_gpr, groupM, N, K); +} + // MX scale presence is owned independently by each input side. The primary // itself is passed as an unused compiler operand for an absent scale; the // constant `.if` in the inline assembly removes that binder completely. diff --git a/test/tileop_api/src/CoopGroupMOverloads.cpp b/test/tileop_api/src/CoopGroupMOverloads.cpp index e707a89..b02508b 100644 --- a/test/tileop_api/src/CoopGroupMOverloads.cpp +++ b/test/tileop_api/src/CoopGroupMOverloads.cpp @@ -58,4 +58,20 @@ __attribute__((noinline)) void mx(float *out, __fp8_e4m3 *ain, __fp8_e4m3 *bin, } // namespace mx_forms + +namespace options_groupm_forms { +// 用户报错形态: TMATMUL_ACC(tO, tO, tW, tV, pvOptions, kGroupM) +using TW = Tile; +using TV = SharedMatrixRight<__bf16, 128, 16>; +using TO = Tile; +constexpr size_t kGroupM = 128; +__attribute__((noinline)) void pv(TO &tO, TW &tW, TV &tV) { + auto sv = TMOV_L2S_INSERT(tV); + auto pvOptions = fixp::Options{}; + TMATMUL_ACC(tO, tO, tW, sv, pvOptions, kGroupM); +} + + +} // namespace options_groupm_forms + int main() { return 0; }