新增trmm-catlass-AtomGit

Aexamples/76_trmm/CMakeLists.txt+13-0

@@ -0,0 +1,13 @@
	1	+# -----------------------------------------------------------------------------------------------------------
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+# CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+# See LICENSE in the root of the software repository for the full text of the License.
	9	+# -----------------------------------------------------------------------------------------------------------
	10	+
	11	+set_source_files_properties(trmm.cpp PROPERTIES LANGUAGE ASC)
	12	+catlass_example_add_executable(76_trmm mix trmm.cpp)
	13	+target_link_libraries(76_trmm PRIVATE pthread)

Aexamples/76_trmm/README.md+35-0

@@ -0,0 +1,35 @@
	1	+# 76_trmm
	2	+
	3	+## 功能说明
	4	+
	5	+本样例演示基于 CATLASS GEMM kernel 组件实现 TRMM（triangular matrix multiply）：
	6	+
	7	+- `side = 0`：`C = alpha * op(A) * B`，其中 `A` 为三角矩阵。
	8	+- `side = 1`：`C = alpha * A * op(B)`，其中 `B` 为三角矩阵。
	9	+- `uplo = 0` 表示 lower triangular，`uplo = 1` 表示 upper triangular。
	10	+- `trans = 0` 表示不转置三角矩阵，`trans = 1` 表示转置三角矩阵。
	11	+
	12	+当前样例使用 `float32` 输入输出，`diag = 0`（non-unit diagonal）。
	13	+
	14	+## 构建
	15	+
	16	+```bash
	17	+bash scripts/build.sh 76_trmm
	18	+```
	19	+
	20	+## 运行
	21	+
	22	+命令格式：
	23	+
	24	+```bash
	25	+./output/bin/76_trmm m n side uplo trans diag alpha [device_id]
	26	+```
	27	+
	28	+示例：
	29	+
	30	+```bash
	31	+./output/bin/76_trmm 256 512 0 0 0 0 1.0
	32	+./output/bin/76_trmm 512 256 1 1 1 0 1.0
	33	+```
	34	+
	35	+程序会在 host 侧构造三角矩阵与 dense 矩阵，运行 NPU kernel，并与 CPU reference 结果做精度比对。

Aexamples/76_trmm/trmm.cpp+386-0

@@ -0,0 +1,386 @@
	1	+/**
	2	+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+ * CANN Open Software License Agreement Version 2.0 (the "License").
	5	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+
	11	+// By setting the K_MAX_SHAPE_DIM macro, the dimension of the AscendC Tensor's ShapeInfo is configured to 0,
	12	+// optimizing stack space. If you need to use the ShapeInfo of the AscendC Tensor, please undefine this macro.
	13	+#ifndef K_MAX_SHAPE_DIM
	14	+#define K_MAX_SHAPE_DIM 0
	15	+#endif
	16	+
	17	+#include "catlass/gemm/kernel/trmm.hpp"
	18	+
	19	+#include "catlass/arch/arch.hpp"
	20	+#include "catlass/catlass.hpp"
	21	+#include "catlass/gemm/block/block_mmad.hpp"
	22	+#include "catlass/gemm/block/block_swizzle.hpp"
	23	+#include "catlass/gemm/device/device_gemm.hpp"
	24	+#include "catlass/gemm/dispatch_policy.hpp"
	25	+#include "catlass/layout/layout.hpp"
	26	+#include "catlass/status.hpp"
	27	+#include "tla/layout.hpp"
	28	+#include "tla/tensor.hpp"
	29	+
	30	+#include "golden.hpp"
	31	+#include "helper.hpp"
	32	+
	33	+#include <vector>
	34	+
	35	+using namespace Catlass;
	36	+
	37	+using Options = TrmmOptions;
	38	+
	39	+// Tile/Swizzle configurations matching TRMM delivery package
	40	+using TrmmDefaultL1 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<256>>;
	41	+using TrmmDefaultL0 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	42	+using TrmmDefaultSwizzle30 = Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
	43	+using TrmmDefaultSwizzle31 = Gemm::Block::GemmIdentityBlockSwizzle<3, 1>;
	44	+
	45	+using TrmmSRLowerL1 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	46	+using TrmmSRLowerL0 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	47	+using TrmmSRLowerSwizzle = Gemm::Block::GemmIdentityBlockSwizzle<1, 0>;
	48	+
	49	+using TrmmSRUpperL1 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	50	+using TrmmSRUpperL0 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	51	+using TrmmSRUpperSwizzle = Gemm::Block::GemmIdentityBlockSwizzle<4, 1>;
	52	+
	53	+using TrmmSRK128L1 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<128>>;
	54	+using TrmmSRK128L0 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	55	+using TrmmSRK128Swizzle = Gemm::Block::GemmIdentityBlockSwizzle<1, 1>;
	56	+
	57	+using TrmmSLK128L1 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<128>>;
	58	+using TrmmSLK128L0 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	59	+using TrmmSLK128Swizzle = Gemm::Block::GemmIdentityBlockSwizzle<1, 1>;
	60	+
	61	+constexpr uint32_t SMALL_RIGHT_N_LIMIT = 512u;
	62	+constexpr uint32_t SMALL_RIGHT_K128_N_LIMIT = 256u;
	63	+constexpr uint32_t SMALL_LEFT_M_LIMIT = 512u;
	64	+constexpr uint32_t SMALL_LEFT_K128_M_LIMIT = 256u;
	65	+
	66	+enum class TileVariant : uint64_t
	67	+{
	68	+ TILE_DEFAULT = 0,
	69	+ TILE_SMALL_RIGHT_LOWER = 1,
	70	+ TILE_SMALL_RIGHT_UPPER = 2,
	71	+ TILE_SMALL_RIGHT_K128 = 3,
	72	+ TILE_DEFAULT_SWIZZLE31 = 4,
	73	+ TILE_SMALL_LEFT_LOWER = 5,
	74	+ TILE_SMALL_LEFT_UPPER = 6,
	75	+ TILE_SMALL_LEFT_K128 = 7,
	76	+};
	77	+
	78	+static TileVariant SelectTileVariant(uint32_t side, uint32_t uplo, uint32_t trans, uint32_t m, uint32_t n)
	79	+{
	80	+ uint32_t effectiveUplo = uplo ^ trans;
	81	+ bool isUpper = (effectiveUplo == 1u);
	82	+ bool isLower = (effectiveUplo == 0u);
	83	+
	84	+ if (side == 1) {
	85	+ if (isUpper && n <= SMALL_RIGHT_K128_N_LIMIT)
	86	+ return TileVariant::TILE_SMALL_RIGHT_K128;
	87	+ if (isUpper && n <= SMALL_RIGHT_N_LIMIT)
	88	+ return TileVariant::TILE_SMALL_RIGHT_UPPER;
	89	+ if (isLower && n <= SMALL_RIGHT_N_LIMIT)
	90	+ return TileVariant::TILE_SMALL_RIGHT_LOWER;
	91	+ } else {
	92	+ if (isUpper && m <= SMALL_LEFT_K128_M_LIMIT)
	93	+ return TileVariant::TILE_SMALL_LEFT_K128;
	94	+ if (isUpper && m <= SMALL_LEFT_M_LIMIT)
	95	+ return TileVariant::TILE_SMALL_LEFT_UPPER;
	96	+ if (isLower && m <= SMALL_LEFT_M_LIMIT)
	97	+ return TileVariant::TILE_SMALL_LEFT_LOWER;
	98	+ }
	99	+
	100	+ if (m >= 2048 && m < n)
	101	+ return TileVariant::TILE_DEFAULT_SWIZZLE31;
	102	+ return TileVariant::TILE_DEFAULT;
	103	+}
	104	+
	105	+template <class LayoutTagA, class LayoutTagB, class L1TileShape, class L0TileShape, class BlockScheduler>
	106	+static bool RunTrmmDispatch(
	107	+ GemmCoord problemShape, uint8_t* deviceA, uint8_t* deviceB, uint8_t* deviceC, uint32_t side, uint32_t uplo,
	108	+ uint32_t trans, float alpha, aclrtStream stream, uint32_t aicCoreNum)
	109	+{
	110	+ using ElementA = float;
	111	+ using ElementB = float;
	112	+ using ElementC = float;
	113	+ using LayoutTagC = layout::RowMajor;
	114	+ using ArchTag = Arch::AtlasA2;
	115	+ constexpr bool enableUnitFlag = true;
	116	+ constexpr bool useHF32 = false;
	117	+ using DispatchPolicy = Gemm::MmadPingpong<ArchTag, enableUnitFlag, useHF32>;
	118	+
	119	+ using TileCopy =
	120	+ Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC>;
	121	+ using BlockMmad = Gemm::Block::BlockMmadTla<
	122	+ DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, void, TileCopy>;
	123	+ using BlockEpilogue = void;
	124	+
	125	+ using TrmmKernel = Gemm::Kernel::Trmm<BlockMmad, BlockEpilogue, BlockScheduler>;
	126	+ using TrmmAdapter = Gemm::Device::DeviceGemm<TrmmKernel>;
	127	+
	128	+ auto layoutA = tla::MakeLayout<ElementA, LayoutTagA>(problemShape.m(), problemShape.k());
	129	+ auto layoutB = tla::MakeLayout<ElementB, LayoutTagB>(problemShape.k(), problemShape.n());
	130	+ auto layoutC = tla::MakeLayout<ElementC, LayoutTagC>(problemShape.m(), problemShape.n());
	131	+
	132	+ typename TrmmKernel::Arguments arguments{problemShape, deviceA, layoutA, deviceB, layoutB, deviceC,
	133	+ layoutC, side, uplo, trans, 0u, alpha};
	134	+ TrmmAdapter trmmOp;
	135	+ if (trmmOp.CanImplement(arguments) != Status::kSuccess) {
	136	+ std::cerr << "[ERROR] TRMM kernel cannot implement the supplied arguments." << std::endl;
	137	+ return false;
	138	+ }
	139	+ size_t sizeWorkspace = trmmOp.GetWorkspaceSize(arguments);
	140	+ uint8_t* deviceWorkspace = nullptr;
	141	+ if (sizeWorkspace > 0) {
	142	+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceWorkspace), sizeWorkspace, ACL_MEM_MALLOC_HUGE_FIRST));
	143	+ }
	144	+ if (trmmOp.Initialize(arguments, deviceWorkspace) != Status::kSuccess) {
	145	+ std::cerr << "[ERROR] Failed to initialize TRMM kernel." << std::endl;
	146	+ if (deviceWorkspace != nullptr) {
	147	+ ACL_CHECK(aclrtFree(deviceWorkspace));
	148	+ }
	149	+ return false;
	150	+ }
	151	+ trmmOp(stream, aicCoreNum);
	152	+ ACL_CHECK(aclrtSynchronizeStream(stream));
	153	+
	154	+ if (sizeWorkspace > 0) {
	155	+ ACL_CHECK(aclrtFree(deviceWorkspace));
	156	+ }
	157	+ return true;
	158	+}
	159	+
	160	+#define TRMM_DISPATCH(LayoutA_t, LayoutB_t, L1_t, L0_t, Swizzle_t) \
	161	+ dispatchSucceeded = RunTrmmDispatch<LayoutA_t, LayoutB_t, L1_t, L0_t, Swizzle_t>( \
	162	+ problemShape, deviceA, deviceB, deviceC, side, uplo, trans, kernelAlpha, stream, aicCoreNum)
	163	+
	164	+static void Run(const Options& options)
	165	+{
	166	+ aclrtStream stream{nullptr};
	167	+
	168	+ ACL_CHECK(aclInit(nullptr));
	169	+ ACL_CHECK(aclrtSetDevice(options.deviceId));
	170	+ ACL_CHECK(aclrtCreateStream(&stream));
	171	+
	172	+ uint32_t m = options.problemShape.m();
	173	+ uint32_t n = options.problemShape.n();
	174	+ uint32_t k = options.problemShape.k();
	175	+ uint32_t side = options.side;
	176	+ uint32_t uplo = options.uplo;
	177	+ uint32_t trans = options.trans;
	178	+ uint32_t diag = options.diag;
	179	+ float alpha = options.alpha;
	180	+
	181	+ if (m == 0 \|\| n == 0) {
	182	+ std::cerr << "[ERROR] Dimensions M and N must be greater than 0." << std::endl;
	183	+ return;
	184	+ }
	185	+ if (side > 1) {
	186	+ std::cerr << "[ERROR] Invalid side: " << side << ". Must be 0 (left) or 1 (right)." << std::endl;
	187	+ return;
	188	+ }
	189	+ if (uplo > 1) {
	190	+ std::cerr << "[ERROR] Invalid uplo: " << uplo << ". Must be 0 (lower) or 1 (upper)." << std::endl;
	191	+ return;
	192	+ }
	193	+ if (trans > 1) {
	194	+ std::cerr << "[ERROR] Invalid trans: " << trans << ". Must be 0 (non-transpose) or 1 (transpose)." << std::endl;
	195	+ return;
	196	+ }
	197	+ if (diag != 0) {
	198	+ std::cerr << "[ERROR] Invalid diag: " << diag << ". This example currently supports only 0 (non-unit)."
	199	+ << std::endl;
	200	+ return;
	201	+ }
	202	+
	203	+ size_t lenA = static_cast<size_t>(m) * k;
	204	+ size_t lenB = static_cast<size_t>(k) * n;
	205	+ size_t lenC = static_cast<size_t>(m) * n;
	206	+
	207	+ size_t sizeA = lenA * sizeof(float);
	208	+ size_t sizeB = lenB * sizeof(float);
	209	+ size_t sizeC = lenC * sizeof(float);
	210	+
	211	+ // Layout dispatch matching TRMM:
	212	+ // trans=0: Row/Row
	213	+ // trans=1, side=left: Col/Row
	214	+ // trans=1, side=right: Row/Col
	215	+ using LayoutRowRowA = layout::RowMajor;
	216	+ using LayoutRowRowB = layout::RowMajor;
	217	+ using LayoutColRowA = layout::ColumnMajor;
	218	+ using LayoutColRowB = layout::RowMajor;
	219	+ using LayoutRowColA = layout::RowMajor;
	220	+ using LayoutRowColB = layout::ColumnMajor;
	221	+
	222	+ uint64_t triElements = static_cast<uint64_t>(k) * static_cast<uint64_t>(k);
	223	+ uint64_t outputElements = static_cast<uint64_t>(m) * static_cast<uint64_t>(n);
	224	+ bool fuseAlphaInPrepare = (triElements <= outputElements);
	225	+ float prepareAlpha = fuseAlphaInPrepare ? alpha : 1.0f;
	226	+ float kernelAlpha = fuseAlphaInPrepare ? 1.0f : alpha;
	227	+
	228	+ std::vector<float> hostA(lenA);
	229	+ std::vector<float> hostB(lenB);
	230	+ if (side == 0) {
	231	+ golden::FillTriangularData(hostA, m, k, uplo, diag, 0.5f, 1.5f, prepareAlpha);
	232	+ golden::FillRandomData(hostB, 0.5f, 1.5f);
	233	+ } else {
	234	+ golden::FillRandomData(hostA, 0.5f, 1.5f);
	235	+ golden::FillTriangularData(hostB, k, n, uplo, diag, 0.5f, 1.5f, prepareAlpha);
	236	+ }
	237	+
	238	+ uint8_t* deviceA{nullptr};
	239	+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceA), sizeA, ACL_MEM_MALLOC_HUGE_FIRST));
	240	+ ACL_CHECK(aclrtMemcpy(deviceA, sizeA, hostA.data(), sizeA, ACL_MEMCPY_HOST_TO_DEVICE));
	241	+
	242	+ uint8_t* deviceB{nullptr};
	243	+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceB), sizeB, ACL_MEM_MALLOC_HUGE_FIRST));
	244	+ ACL_CHECK(aclrtMemcpy(deviceB, sizeB, hostB.data(), sizeB, ACL_MEMCPY_HOST_TO_DEVICE));
	245	+
	246	+ uint8_t* deviceC{nullptr};
	247	+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceC), sizeC, ACL_MEM_MALLOC_HUGE_FIRST));
	248	+
	249	+ auto aicCoreNum = platform_ascendc::PlatformAscendCManager::GetInstance()->GetCoreNumAic();
	250	+
	251	+ GemmCoord problemShape{m, n, k};
	252	+ TileVariant tileVariant = SelectTileVariant(side, uplo, trans, m, n);
	253	+ bool dispatchSucceeded = false;
	254	+
	255	+ // Dispatch matching TRMM TilingKey logic
	256	+ if (trans == 0) {
	257	+ // === layout = Row/Row (trans=0) ===
	258	+ switch (tileVariant) {
	259	+ case TileVariant::TILE_DEFAULT:
	260	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	261	+ break;
	262	+ case TileVariant::TILE_DEFAULT_SWIZZLE31:
	263	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle31);
	264	+ break;
	265	+ case TileVariant::TILE_SMALL_RIGHT_LOWER:
	266	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmSRLowerL1, TrmmSRLowerL0, TrmmSRLowerSwizzle);
	267	+ break;
	268	+ case TileVariant::TILE_SMALL_RIGHT_UPPER:
	269	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmSRUpperL1, TrmmSRUpperL0, TrmmSRUpperSwizzle);
	270	+ break;
	271	+ case TileVariant::TILE_SMALL_RIGHT_K128:
	272	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmSRK128L1, TrmmSRK128L0, TrmmSRK128Swizzle);
	273	+ break;
	274	+ case TileVariant::TILE_SMALL_LEFT_LOWER:
	275	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	276	+ break;
	277	+ case TileVariant::TILE_SMALL_LEFT_UPPER:
	278	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	279	+ break;
	280	+ case TileVariant::TILE_SMALL_LEFT_K128:
	281	+ TRMM_DISPATCH(LayoutRowRowA, LayoutRowRowB, TrmmSLK128L1, TrmmSLK128L0, TrmmSLK128Swizzle);
	282	+ break;
	283	+ }
	284	+ } else if (trans == 1 && side == 0) {
	285	+ // === layout = Col/Row (trans=1, side=left) ===
	286	+ switch (tileVariant) {
	287	+ case TileVariant::TILE_DEFAULT:
	288	+ TRMM_DISPATCH(LayoutColRowA, LayoutColRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	289	+ break;
	290	+ case TileVariant::TILE_DEFAULT_SWIZZLE31:
	291	+ TRMM_DISPATCH(LayoutColRowA, LayoutColRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle31);
	292	+ break;
	293	+ case TileVariant::TILE_SMALL_LEFT_LOWER:
	294	+ TRMM_DISPATCH(LayoutColRowA, LayoutColRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	295	+ break;
	296	+ case TileVariant::TILE_SMALL_LEFT_UPPER:
	297	+ TRMM_DISPATCH(LayoutColRowA, LayoutColRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	298	+ break;
	299	+ case TileVariant::TILE_SMALL_LEFT_K128:
	300	+ TRMM_DISPATCH(LayoutColRowA, LayoutColRowB, TrmmSLK128L1, TrmmSLK128L0, TrmmSLK128Swizzle);
	301	+ break;
	302	+ default:
	303	+ TRMM_DISPATCH(LayoutColRowA, LayoutColRowB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	304	+ break;
	305	+ }
	306	+ } else {
	307	+ // === layout = Row/Col (trans=1, side=right) ===
	308	+ switch (tileVariant) {
	309	+ case TileVariant::TILE_DEFAULT:
	310	+ TRMM_DISPATCH(LayoutRowColA, LayoutRowColB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	311	+ break;
	312	+ case TileVariant::TILE_DEFAULT_SWIZZLE31:
	313	+ TRMM_DISPATCH(LayoutRowColA, LayoutRowColB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle31);
	314	+ break;
	315	+ case TileVariant::TILE_SMALL_RIGHT_LOWER:
	316	+ TRMM_DISPATCH(LayoutRowColA, LayoutRowColB, TrmmSRLowerL1, TrmmSRLowerL0, TrmmSRLowerSwizzle);
	317	+ break;
	318	+ case TileVariant::TILE_SMALL_RIGHT_UPPER:
	319	+ TRMM_DISPATCH(LayoutRowColA, LayoutRowColB, TrmmSRUpperL1, TrmmSRUpperL0, TrmmSRUpperSwizzle);
	320	+ break;
	321	+ case TileVariant::TILE_SMALL_RIGHT_K128:
	322	+ TRMM_DISPATCH(LayoutRowColA, LayoutRowColB, TrmmSRK128L1, TrmmSRK128L0, TrmmSRK128Swizzle);
	323	+ break;
	324	+ default:
	325	+ TRMM_DISPATCH(LayoutRowColA, LayoutRowColB, TrmmDefaultL1, TrmmDefaultL0, TrmmDefaultSwizzle30);
	326	+ break;
	327	+ }
	328	+ }
	329	+
	330	+ if (!dispatchSucceeded) {
	331	+ ACL_CHECK(aclrtFree(deviceA));
	332	+ ACL_CHECK(aclrtFree(deviceB));
	333	+ ACL_CHECK(aclrtFree(deviceC));
	334	+ ACL_CHECK(aclrtDestroyStream(stream));
	335	+ ACL_CHECK(aclrtResetDevice(options.deviceId));
	336	+ ACL_CHECK(aclFinalize());
	337	+ return;
	338	+ }
	339	+
	340	+ std::vector<float> hostC(lenC);
	341	+ ACL_CHECK(aclrtMemcpy(hostC.data(), sizeC, deviceC, sizeC, ACL_MEMCPY_DEVICE_TO_HOST));
	342	+
	343	+ std::vector<float> hostGolden(lenC);
	344	+ if (trans == 0) {
	345	+ golden::ComputeMatmul(
	346	+ problemShape, hostA, layout::RowMajor(m, k), hostB, layout::RowMajor(k, n), hostGolden,
	347	+ layout::RowMajor(m, n));
	348	+ } else if (side == 0) {
	349	+ golden::ComputeMatmul(
	350	+ problemShape, hostA, layout::ColumnMajor(m, k), hostB, layout::RowMajor(k, n), hostGolden,
	351	+ layout::RowMajor(m, n));
	352	+ } else {
	353	+ golden::ComputeMatmul(
	354	+ problemShape, hostA, layout::RowMajor(m, k), hostB, layout::ColumnMajor(k, n), hostGolden,
	355	+ layout::RowMajor(m, n));
	356	+ }
	357	+ if (kernelAlpha != 1.0f) {
	358	+ for (auto& v : hostGolden) {
	359	+ v *= kernelAlpha;
	360	+ }
	361	+ }
	362	+ std::vector<uint64_t> errorIndices = golden::CompareData(hostC, hostGolden, k);
	363	+ if (errorIndices.empty()) {
	364	+ std::cout << "Compare success." << std::endl;
	365	+ } else {
	366	+ std::cerr << "Compare failed. Error count: " << errorIndices.size() << std::endl;
	367	+ }
	368	+
	369	+ ACL_CHECK(aclrtFree(deviceA));
	370	+ ACL_CHECK(aclrtFree(deviceB));
	371	+ ACL_CHECK(aclrtFree(deviceC));
	372	+
	373	+ ACL_CHECK(aclrtDestroyStream(stream));
	374	+ ACL_CHECK(aclrtResetDevice(options.deviceId));
	375	+ ACL_CHECK(aclFinalize());
	376	+}
	377	+
	378	+int main(int argc, const char** argv)
	379	+{
	380	+ Options options;
	381	+ if (options.Parse(argc, argv) != 0) {
	382	+ return -1;
	383	+ }
	384	+ Run(options);
	385	+ return 0;
	386	+}

Mexamples/CMakeLists.txt+1-0

	@@ -142,6 +142,7 @@ set(EXAMPLE_ATLASA2
142	45_strided_batched_matmul_tla	142	45_strided_batched_matmul_tla
143	52_quant_multi_core_splitk_matmul_tla	143	52_quant_multi_core_splitk_matmul_tla
144	75_symm	144	75_symm
		145	+ 76_trmm
145	102_dynamic_optimized_matmul	146	102_dynamic_optimized_matmul
146	103_dynamic_optimized_quant_matmul_per_token_basic	147	103_dynamic_optimized_quant_matmul_per_token_basic
147	)	148	)

Mexamples/common/golden/fill_data.hpp+27-2

	@@ -11,10 +11,11 @@
11	#ifndef EXAMPLES_COMMON_GOLDEN_FILL_DATA_HPP	11	#ifndef EXAMPLES_COMMON_GOLDEN_FILL_DATA_HPP
12	#define EXAMPLES_COMMON_GOLDEN_FILL_DATA_HPP	12	#define EXAMPLES_COMMON_GOLDEN_FILL_DATA_HPP
13		13
14	-#include <stack>	14	+#include <cstdint>
15	-#include <vector>
16	#include <cstdlib>	15	#include <cstdlib>
17	#include <ctime>	16	#include <ctime>
		17	+#include <stack>
		18	+#include <vector>
18		19
19	namespace Catlass::golden {	20	namespace Catlass::golden {
20		21
	@@ -45,6 +46,30 @@ void FillRandomData<int8_t, int>(std::vector<int8_t>& data, int low, int high)
45	}	46	}
46	}	47	}
47		48
		49	+template <class Element, class ElementRandom>
		50	+void FillTriangularData(
		51	+ std::vector<Element>& data, uint32_t rows, uint32_t columns, uint32_t uplo, uint32_t diag, ElementRandom low,
		52	+ ElementRandom high, Element alpha = static_cast<Element>(1))
		53	+{
		54	+ data.resize(static_cast<std::size_t>(rows) * columns);
		55	+ FillRandomData(data, low, high);
		56	+
		57	+ for (uint32_t i = 0; i < rows; ++i) {
		58	+ for (uint32_t j = 0; j < columns; ++j) {
		59	+ bool keep = (uplo == 0) ? (j <= i) : (j >= i);
		60	+ auto index = static_cast<std::size_t>(i) * columns + j;
		61	+ if (!keep) {
		62	+ data[index] = static_cast<Element>(0);
		63	+ continue;
		64	+ }
		65	+ data[index] = static_cast<Element>(data[index] * alpha);
		66	+ if (diag != 0 && i == j) {
		67	+ data[index] = alpha;
		68	+ }
		69	+ }
		70	+ }
		71	+}
		72	+
48	template <class Element, class ElementRandom>	73	template <class Element, class ElementRandom>
49	void FillSymmRandomData(	74	void FillSymmRandomData(
50	std::vector<Element>& data, uint32_t dim, bool upperStorage, ElementRandom low, ElementRandom high)	75	std::vector<Element>& data, uint32_t dim, bool upperStorage, ElementRandom low, ElementRandom high)

Mexamples/common/options.hpp+58-0

	@@ -11,6 +11,8 @@
11	#ifndef EXAMPLES_COMMON_OPTIONS_HPP	11	#ifndef EXAMPLES_COMMON_OPTIONS_HPP
12	#define EXAMPLES_COMMON_OPTIONS_HPP	12	#define EXAMPLES_COMMON_OPTIONS_HPP
13		13
		14	+#include <cstdint>
		15	+#include <cstdlib>
14	#include <iostream>	16	#include <iostream>
15	#include <string>	17	#include <string>
16		18
	@@ -144,4 +146,60 @@ struct GroupedGemmOptions {
144	}	146	}
145	};	147	};
146		148
		149	+/**
		150	+ * @struct TrmmOptions
		151	+ * @brief Options structure for TRMM examples.
		152	+ * @brief Arguments: `example_name m n side uplo trans diag alpha [device_id]`
		153	+ */
		154	+struct TrmmOptions {
		155	+ const std::string HELPER = "m n side uplo trans diag alpha [device_id]";
		156	+
		157	+ Catlass::GemmCoord problemShape{128, 128, 128};
		158	+ uint32_t side{0};
		159	+ uint32_t uplo{0};
		160	+ uint32_t trans{0};
		161	+ uint32_t diag{0};
		162	+ float alpha{1.0f};
		163	+ int32_t deviceId{0};
		164	+
		165	+ TrmmOptions() = default;
		166	+
		167	+ int Parse(int argc, const char** argv)
		168	+ {
		169	+ enum class ArgsIndex
		170	+ {
		171	+ M_INDEX = 1,
		172	+ N_INDEX,
		173	+ SIDE_INDEX,
		174	+ UPLO_INDEX,
		175	+ TRANS_INDEX,
		176	+ DIAG_INDEX,
		177	+ ALPHA_INDEX,
		178	+ DEVICE_ID_INDEX,
		179	+ ARGS_MAX
		180	+ };
		181	+
		182	+ if (argc > static_cast<uint32_t>(ArgsIndex::ARGS_MAX) \|\|
		183	+ argc < static_cast<uint32_t>(ArgsIndex::ALPHA_INDEX) + 1) {
		184	+ std::cerr << TOSTRING(CATLASS_EXAMPLE_NAME) << " " << HELPER << std::endl;
		185	+ return -1;
		186	+ }
		187	+
		188	+ problemShape.m() = std::atoi(argv[static_cast<uint32_t>(ArgsIndex::M_INDEX)]);
		189	+ problemShape.n() = std::atoi(argv[static_cast<uint32_t>(ArgsIndex::N_INDEX)]);
		190	+ side = static_cast<uint32_t>(std::atoi(argv[static_cast<uint32_t>(ArgsIndex::SIDE_INDEX)]));
		191	+ uplo = static_cast<uint32_t>(std::atoi(argv[static_cast<uint32_t>(ArgsIndex::UPLO_INDEX)]));
		192	+ trans = static_cast<uint32_t>(std::atoi(argv[static_cast<uint32_t>(ArgsIndex::TRANS_INDEX)]));
		193	+ diag = static_cast<uint32_t>(std::atoi(argv[static_cast<uint32_t>(ArgsIndex::DIAG_INDEX)]));
		194	+ alpha = std::atof(argv[static_cast<uint32_t>(ArgsIndex::ALPHA_INDEX)]);
		195	+
		196	+ problemShape.k() = (side == 0) ? problemShape.m() : problemShape.n();
		197	+
		198	+ if (argc == static_cast<uint32_t>(ArgsIndex::ARGS_MAX)) {
		199	+ deviceId = std::atoi(argv[static_cast<uint32_t>(ArgsIndex::DEVICE_ID_INDEX)]);
		200	+ }
		201	+ return 0;
		202	+ }
		203	+};
		204	+
147	#endif	205	#endif

Ainclude/catlass/gemm/kernel/trmm.hpp+314-0

@@ -0,0 +1,314 @@
	1	+/**
	2	+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+ * CANN Open Software License Agreement Version 2.0 (the "License").
	5	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+
	11	+#ifndef CATLASS_GEMM_KERNEL_TRMM_HPP
	12	+#define CATLASS_GEMM_KERNEL_TRMM_HPP
	13	+
	14	+#include "catlass/catlass.hpp"
	15	+#include "catlass/arch/resource.hpp"
	16	+#include "catlass/arch/cross_core_sync.hpp"
	17	+#include "catlass/coord.hpp"
	18	+#include "catlass/epilogue/tile/copy_gm_to_ub_tla.hpp"
	19	+#include "catlass/epilogue/tile/copy_ub_to_gm_tla.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/matrix_coord.hpp"
	22	+#include "tla/layout.hpp"
	23	+#include "tla/tensor.hpp"
	24	+
	25	+namespace Catlass::Gemm::Kernel {
	26	+
	27	+// Compute TRMM C = alpha * op(T) * B or C = alpha * A * op(T), where T is triangular.
	28	+// Supported kernel parameters are side={left,right}, uplo={lower,upper}, trans={no,trans}, and diag={nonunit}.
	29	+// The kernel reduces work by clipping each output tile to the valid triangular K range and expects inactive triangular
	30	+// elements to be zeroed by the caller-side wrapper.
	31	+template <class BlockMmad_, class BlockEpilogue_, class BlockScheduler_>
	32	+class Trmm {
	33	+public:
	34	+ using BlockMmad = BlockMmad_;
	35	+ using ArchTag = typename BlockMmad::ArchTag;
	36	+ using L1TileShape = typename BlockMmad::L1TileShape;
	37	+ using ElementA = typename BlockMmad::ElementA;
	38	+ using LayoutA = typename BlockMmad::LayoutA;
	39	+ using ElementB = typename BlockMmad::ElementB;
	40	+ using LayoutB = typename BlockMmad::LayoutB;
	41	+ using ElementC = typename BlockMmad::ElementC;
	42	+ using LayoutC = typename BlockMmad::LayoutC;
	43	+ using ElementAccumulator = typename BlockMmad::ElementAccumulator;
	44	+
	45	+ using BlockEpilogue = BlockEpilogue_;
	46	+ using BlockScheduler = BlockScheduler_;
	47	+
	48	+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});
	49	+ static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});
	50	+ static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{});
	51	+
	52	+ struct Params {
	53	+ GemmCoord problemShape;
	54	+ GM_ADDR ptrA;
	55	+ LayoutA layoutA;
	56	+ GM_ADDR ptrB;
	57	+ LayoutB layoutB;
	58	+ GM_ADDR ptrC;
	59	+ LayoutC layoutC;
	60	+ uint32_t side;
	61	+ uint32_t uplo;
	62	+ uint32_t trans;
	63	+ uint32_t diag;
	64	+ float alpha;
	65	+
	66	+ CATLASS_HOST_DEVICE
	67	+ Params()
	68	+ {}
	69	+
	70	+ CATLASS_HOST_DEVICE
	71	+ Params(
	72	+ GemmCoord const& problemShape_, GM_ADDR ptrA_, LayoutA layoutA_, GM_ADDR ptrB_, LayoutB layoutB_,
	73	+ GM_ADDR ptrC_, LayoutC layoutC_, uint32_t side_, uint32_t uplo_, uint32_t trans_, uint32_t diag_,
	74	+ float alpha_)
	75	+ : problemShape(problemShape_),
	76	+ ptrA(ptrA_),
	77	+ layoutA(layoutA_),
	78	+ ptrB(ptrB_),
	79	+ layoutB(layoutB_),
	80	+ ptrC(ptrC_),
	81	+ layoutC(layoutC_),
	82	+ side(side_),
	83	+ uplo(uplo_),
	84	+ trans(trans_),
	85	+ diag(diag_),
	86	+ alpha(alpha_)
	87	+ {}
	88	+ };
	89	+
	90	+ struct Arguments {
	91	+ GemmCoord problemShape;
	92	+ GM_ADDR ptrA;
	93	+ LayoutA layoutA;
	94	+ GM_ADDR ptrB;
	95	+ LayoutB layoutB;
	96	+ GM_ADDR ptrC;
	97	+ LayoutC layoutC;
	98	+ uint32_t side{0};
	99	+ uint32_t uplo{0};
	100	+ uint32_t trans{0};
	101	+ uint32_t diag{0};
	102	+ float alpha{1.0f};
	103	+ };
	104	+
	105	+ static bool CanImplement(const Arguments& args)
	106	+ {
	107	+ if (args.ptrA == nullptr \|\| args.ptrB == nullptr \|\| args.ptrC == nullptr) {
	108	+ return false;
	109	+ }
	110	+ if (args.problemShape.m() == 0 \|\| args.problemShape.n() == 0 \|\| args.problemShape.k() == 0) {
	111	+ return false;
	112	+ }
	113	+ if (args.side > 1 \|\| args.uplo > 1 \|\| args.trans > 1 \|\| args.diag != 0) {
	114	+ return false;
	115	+ }
	116	+ if (args.side == 0) {
	117	+ return args.problemShape.m() == args.problemShape.k();
	118	+ }
	119	+ return args.problemShape.n() == args.problemShape.k();
	120	+ }
	121	+
	122	+ static size_t GetWorkspaceSize(const Arguments& args)
	123	+ {
	124	+ return 0;
	125	+ }
	126	+
	127	+ static Params ToUnderlyingArguments(const Arguments& args, uint8_t* workspace)
	128	+ {
	129	+ Params params{args.problemShape, args.ptrA, args.layoutA, args.ptrB, args.layoutB, args.ptrC,
	130	+ args.layoutC, args.side, args.uplo, args.trans, args.diag, args.alpha};
	131	+ return params;
	132	+ }
	133	+
	134	+ CATLASS_DEVICE
	135	+ Trmm()
	136	+ {}
	137	+
	138	+ template <int32_t CORE_TYPE = g_coreType>
	139	+ CATLASS_DEVICE void operator()(Params const& params);
	140	+
	141	+ template <>
	142	+ CATLASS_DEVICE void operator()<AscendC::AIC>(Params const& params)
	143	+ {
	144	+ if constexpr (std::is_same_v<ElementA, float> && std::is_same_v<ElementB, float>) {
	145	+ AscendC::SetHF32Mode(false);
	146	+ }
	147	+ BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N));
	148	+ uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();
	149	+
	150	+ Arch::Resource<ArchTag> resource;
	151	+ BlockMmad blockMmad(resource);
	152	+
	153	+ AscendC::GlobalTensor<ElementA> gmA;
	154	+ gmA.SetGlobalBuffer((__gm__ ElementA*)params.ptrA);
	155	+ AscendC::GlobalTensor<ElementB> gmB;
	156	+ gmB.SetGlobalBuffer((__gm__ ElementB*)params.ptrB);
	157	+ AscendC::GlobalTensor<ElementC> gmC;
	158	+ gmC.SetGlobalBuffer((__gm__ ElementC*)params.ptrC);
	159	+
	160	+ auto tensorA = tla::MakeTensor(gmA, params.layoutA, Arch::PositionGM{});
	161	+ auto tensorB = tla::MakeTensor(gmB, params.layoutB, Arch::PositionGM{});
	162	+ auto tensorC = tla::MakeTensor(gmC, params.layoutC, Arch::PositionGM{});
	163	+
	164	+ for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) {
	165	+ GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
	166	+ GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);
	167	+
	168	+ uint32_t mStart = blockCoord.m() * L1_TILE_M;
	169	+ uint32_t nStart = blockCoord.n() * L1_TILE_N;
	170	+ uint32_t mEnd = mStart + actualBlockShape.m();
	171	+ uint32_t nEnd = nStart + actualBlockShape.n();
	172	+ uint32_t kStart = 0;
	173	+ uint32_t kEnd = params.problemShape.k();
	174	+ uint32_t effectiveUplo = params.uplo ^ params.trans;
	175	+ if (params.side == 0) {
	176	+ if (effectiveUplo == 0) {
	177	+ kEnd = mEnd < kEnd ? mEnd : kEnd;
	178	+ } else {
	179	+ kStart = mStart < kEnd ? mStart : kEnd;
	180	+ }
	181	+ } else {
	182	+ if (effectiveUplo == 0) {
	183	+ kStart = nStart < kEnd ? nStart : kEnd;
	184	+ } else {
	185	+ kEnd = nEnd < kEnd ? nEnd : kEnd;
	186	+ }
	187	+ }
	188	+ if (kStart >= kEnd) {
	189	+ continue;
	190	+ }
	191	+ actualBlockShape.k() = kEnd - kStart;
	192	+
	193	+ auto tensorBlockA = GetTile(
	194	+ tensorA, tla::MakeCoord(mStart, kStart), tla::MakeShape(actualBlockShape.m(), actualBlockShape.k()));
	195	+ auto tensorBlockB = GetTile(
	196	+ tensorB, tla::MakeCoord(kStart, nStart), tla::MakeShape(actualBlockShape.k(), actualBlockShape.n()));
	197	+ auto tensorBlockC = GetTile(
	198	+ tensorC, tla::MakeCoord(mStart, nStart), tla::MakeShape(actualBlockShape.m(), actualBlockShape.n()));
	199	+
	200	+ blockMmad(tensorBlockA, tensorBlockB, tensorBlockC, actualBlockShape);
	201	+
	202	+ if (params.alpha != 1.0f) {
	203	+ Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(flagAicFinishStore);
	204	+ }
	205	+ }
	206	+
	207	+ AscendC::PipeBarrier<PIPE_ALL>();
	208	+ }
	209	+
	210	+ template <>
	211	+ CATLASS_DEVICE void operator()<AscendC::AIV>(Params const& params)
	212	+ {
	213	+ if (params.alpha == 1.0f) {
	214	+ return;
	215	+ }
	216	+
	217	+ BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N));
	218	+ uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();
	219	+
	220	+ Arch::Resource<ArchTag> resource;
	221	+ auto ubTensor = resource.ubBuf.template GetBufferByByte<ElementC>(0);
	222	+
	223	+ AscendC::GlobalTensor<ElementC> gmC;
	224	+ gmC.SetGlobalBuffer((__gm__ ElementC*)params.ptrC);
	225	+ auto tensorC = tla::MakeTensor(gmC, params.layoutC, Arch::PositionGM{});
	226	+
	227	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	228	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	229	+ uint32_t aicoreIndex = AscendC::GetBlockIdx() / subBlockNum;
	230	+ uint32_t aicoreNum = AscendC::GetBlockNum() / subBlockNum;
	231	+
	232	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
	233	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
	234	+
	235	+ for (uint32_t loopIdx = aicoreIndex; loopIdx < coreLoops; loopIdx += aicoreNum) {
	236	+ GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
	237	+ GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);
	238	+ uint32_t mStart = blockCoord.m() * L1_TILE_M;
	239	+ uint32_t nStart = blockCoord.n() * L1_TILE_N;
	240	+ uint32_t mEnd = mStart + actualBlockShape.m();
	241	+ uint32_t nEnd = nStart + actualBlockShape.n();
	242	+ uint32_t kStart = 0;
	243	+ uint32_t kEnd = params.problemShape.k();
	244	+ uint32_t effectiveUplo = params.uplo ^ params.trans;
	245	+ if (params.side == 0) {
	246	+ if (effectiveUplo == 0) {
	247	+ kEnd = mEnd < kEnd ? mEnd : kEnd;
	248	+ } else {
	249	+ kStart = mStart < kEnd ? mStart : kEnd;
	250	+ }
	251	+ } else {
	252	+ if (effectiveUplo == 0) {
	253	+ kStart = nStart < kEnd ? nStart : kEnd;
	254	+ } else {
	255	+ kEnd = nEnd < kEnd ? nEnd : kEnd;
	256	+ }
	257	+ }
	258	+ if (kStart >= kEnd) {
	259	+ continue;
	260	+ }
	261	+ Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagAicFinishStore);
	262	+
	263	+ uint32_t subBlockM = CeilDiv(actualBlockShape.m(), subBlockNum);
	264	+ uint32_t subBlockMOffset = subBlockIdx * subBlockM;
	265	+ if (subBlockMOffset >= actualBlockShape.m()) {
	266	+ continue;
	267	+ }
	268	+ uint32_t actualSubBlockM = actualBlockShape.m() - subBlockMOffset;
	269	+ actualSubBlockM = actualSubBlockM < subBlockM ? actualSubBlockM : subBlockM;
	270	+
	271	+ auto tensorBlockC = GetTile(
	272	+ tensorC, tla::MakeCoord(mStart + subBlockMOffset, nStart),
	273	+ tla::MakeShape(actualSubBlockM, actualBlockShape.n()));
	274	+ auto layoutUbC = tla::MakeLayout(
	275	+ tla::MakeShape(actualSubBlockM, actualBlockShape.n()),
	276	+ tla::MakeStride(actualBlockShape.n(), tla::Int<1>{}));
	277	+ auto tensorUbC = tla::MakeTensor(ubTensor, layoutUbC, Arch::PositionUB{});
	278	+
	279	+ using CopyGmToUbC = Epilogue::Tile::CopyGm2UbTla<ArchTag, decltype(tensorBlockC), decltype(tensorUbC)>;
	280	+ using CopyUbToGmC = Epilogue::Tile::CopyUb2GmTla<ArchTag, decltype(tensorUbC), decltype(tensorBlockC)>;
	281	+ CopyGmToUbC copyGmToUbC;
	282	+ CopyUbToGmC copyUbToGmC;
	283	+
	284	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
	285	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
	286	+ copyGmToUbC(tensorUbC, tensorBlockC);
	287	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	288	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	289	+ AscendC::Muls(
	290	+ tensorUbC.data(), tensorUbC.data(), static_cast<ElementC>(params.alpha),
	291	+ actualSubBlockM * actualBlockShape.n());
	292	+ AscendC::PipeBarrier<PIPE_V>();
	293	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
	294	+
	295	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	296	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	297	+ copyUbToGmC(tensorBlockC, tensorUbC);
	298	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
	299	+ }
	300	+
	301	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
	302	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
	303	+ AscendC::PipeBarrier<PIPE_ALL>();
	304	+ }
	305	+
	306	+private:
	307	+ static constexpr Arch::FlagID FLAG_AIC_FINISH_STORE = 0;
	308	+ static constexpr Arch::FlagID RV_FLAG_AIC_FINISH_STORE = 1;
	309	+ Arch::CrossCoreFlagWithReverse<> flagAicFinishStore{FLAG_AIC_FINISH_STORE, RV_FLAG_AIC_FINISH_STORE};
	310	+};
	311	+
	312	+} // namespace Catlass::Gemm::Kernel
	313	+
	314	+#endif // CATLASS_GEMM_KERNEL_TRMM_HPP

Mtests/optest/include/catlass_kernel_jit.h+16-0

	@@ -115,6 +115,17 @@ struct GemmParams : public MatmulParams {
115	float beta = 0.0f; ///< Beta scale in D = alpha * A * B + beta * C.	115	float beta = 0.0f; ///< Beta scale in D = alpha * A * B + beta * C.
116	};	116	};
117		117
		118	+/**
		119	+ * @brief Runtime parameters for example 76_trmm.
		120	+ */
		121	+struct TrmmParams : public MatmulParams {
		122	+ uint32_t side = 0; ///< 0: left triangular input, 1: right triangular input.
		123	+ uint32_t uplo = 0; ///< 0: lower triangular, 1: upper triangular.
		124	+ uint32_t trans = 0; ///< 0: no transpose, 1: transpose triangular input.
		125	+ uint32_t diag = 0; ///< 0: non-unit diagonal. Unit diagonal is not supported by the kernel.
		126	+ float alpha = 1.0f; ///< Output scaling factor.
		127	+};
		128	+
118	/**	129	/**
119	* @brief Runtime parameters for example 75_symm.	130	* @brief Runtime parameters for example 75_symm.
120	*/	131	*/
	@@ -539,6 +550,11 @@ void Ascend950MultiCoreSplitkMatmul(
539	void Ascend950TailMultiCoreSplitkMatmul(	550	void Ascend950TailMultiCoreSplitkMatmul(
540	const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const MatmulParams& params);	551	const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const MatmulParams& params);
541		552
		553	+/**
		554	+ * @brief JIT interface for example 76_trmm.
		555	+ */
		556	+void Trmm(const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const TrmmParams& params);
		557	+
542	/**	558	/**
543	* @brief JIT interface for example 80_grouped_matmul_slice_m_gelu.	559	* @brief JIT interface for example 80_grouped_matmul_slice_m_gelu.
544	*/	560	*/

Atests/optest/kernels/76_trmm/CMakeLists.txt+15-0

@@ -0,0 +1,15 @@
	1	+# -----------------------------------------------------------------------------------------------------------
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+# CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+# See LICENSE in the root of the software repository for the full text of the License.
	9	+# -----------------------------------------------------------------------------------------------------------
	10	+
	11	+add_kernel(NAME trmm
	12	+ NPU_ARCH_LIST 2201
	13	+ KERNEL_TYPE jit
	14	+ ${CMAKE_CURRENT_SOURCE_DIR}/trmm.cpp
	15	+ TEMPLATE ${CMAKE_CURRENT_SOURCE_DIR}/trmm_impl.cpp)

Atests/optest/kernels/76_trmm/trmm.cpp+105-0

@@ -0,0 +1,105 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#include "catlass_kernel.h"
	13	+#include "jit_compiler.h"
	14	+#include "jit_macro_generator.h"
	15	+
	16	+#include <iostream>
	17	+
	18	+namespace CatlassKernel {
	19	+
	20	+namespace {
	21	+constexpr uint32_t SMALL_RIGHT_N_LIMIT = 512u;
	22	+constexpr uint32_t SMALL_RIGHT_K128_N_LIMIT = 256u;
	23	+constexpr uint32_t SMALL_LEFT_M_LIMIT = 512u;
	24	+constexpr uint32_t SMALL_LEFT_K128_M_LIMIT = 256u;
	25	+
	26	+enum class TileVariant : uint32_t
	27	+{
	28	+ TILE_DEFAULT = 0,
	29	+ TILE_SMALL_RIGHT_LOWER = 1,
	30	+ TILE_SMALL_RIGHT_UPPER = 2,
	31	+ TILE_SMALL_RIGHT_K128 = 3,
	32	+ TILE_DEFAULT_SWIZZLE31 = 4,
	33	+ TILE_SMALL_LEFT_LOWER = 5,
	34	+ TILE_SMALL_LEFT_UPPER = 6,
	35	+ TILE_SMALL_LEFT_K128 = 7,
	36	+};
	37	+
	38	+TileVariant SelectTileVariant(const TrmmParams& params)
	39	+{
	40	+ uint32_t effectiveUplo = params.uplo ^ params.trans;
	41	+ bool isUpper = (effectiveUplo == 1u);
	42	+ bool isLower = (effectiveUplo == 0u);
	43	+
	44	+ if (params.side == 1) {
	45	+ if (isUpper && params.n <= SMALL_RIGHT_K128_N_LIMIT)
	46	+ return TileVariant::TILE_SMALL_RIGHT_K128;
	47	+ if (isUpper && params.n <= SMALL_RIGHT_N_LIMIT)
	48	+ return TileVariant::TILE_SMALL_RIGHT_UPPER;
	49	+ if (isLower && params.n <= SMALL_RIGHT_N_LIMIT)
	50	+ return TileVariant::TILE_SMALL_RIGHT_LOWER;
	51	+ } else {
	52	+ if (isUpper && params.m <= SMALL_LEFT_K128_M_LIMIT)
	53	+ return TileVariant::TILE_SMALL_LEFT_K128;
	54	+ if (isUpper && params.m <= SMALL_LEFT_M_LIMIT)
	55	+ return TileVariant::TILE_SMALL_LEFT_UPPER;
	56	+ if (isLower && params.m <= SMALL_LEFT_M_LIMIT)
	57	+ return TileVariant::TILE_SMALL_LEFT_LOWER;
	58	+ }
	59	+
	60	+ if (params.m >= 2048 && params.m < params.n)
	61	+ return TileVariant::TILE_DEFAULT_SWIZZLE31;
	62	+ return TileVariant::TILE_DEFAULT;
	63	+}
	64	+
	65	+const char* SchedulerMacro(TileVariant variant)
	66	+{
	67	+ switch (variant) {
	68	+ case TileVariant::TILE_SMALL_RIGHT_LOWER:
	69	+ return "10";
	70	+ case TileVariant::TILE_SMALL_RIGHT_UPPER:
	71	+ return "41";
	72	+ case TileVariant::TILE_SMALL_RIGHT_K128:
	73	+ case TileVariant::TILE_SMALL_LEFT_K128:
	74	+ return "11";
	75	+ case TileVariant::TILE_DEFAULT_SWIZZLE31:
	76	+ return "31";
	77	+ case TileVariant::TILE_DEFAULT:
	78	+ case TileVariant::TILE_SMALL_LEFT_LOWER:
	79	+ case TileVariant::TILE_SMALL_LEFT_UPPER:
	80	+ default:
	81	+ return "30";
	82	+ }
	83	+}
	84	+} // namespace
	85	+
	86	+/**
	87	+ * @brief example 76_trmm: Resolve and launch the JIT-specialized TRMM implementation.
	88	+ */
	89	+extern "C" void Trmm(const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const TrmmParams& params)
	90	+{
	91	+ auto macros = JitMacroGenerator<TParams>::generate("trmm", tParams);
	92	+ auto variant = SelectTileVariant(params);
	93	+ macros["CATLASS_JIT_TRMM_TILE_VARIANT"] = std::to_string(static_cast<uint32_t>(variant));
	94	+ macros["CATLASS_JIT_BLOCK_SCHEDULER"] = SchedulerMacro(variant);
	95	+
	96	+ auto* entry = JitCompiler::instance().getKernel("trmm_impl.cpp", macros, JitKernelType::MIX);
	97	+ if (entry == nullptr) {
	98	+ std::cerr << "[ERROR] Failed to resolve JIT kernel: trmm_impl.cpp" << std::endl;
	99	+ return;
	100	+ }
	101	+ entry(blockNum, stream, &params);
	102	+ aclrtSynchronizeStream(stream);
	103	+}
	104	+
	105	+} // namespace CatlassKernel

Atests/optest/kernels/76_trmm/trmm_impl.cpp+117-0

@@ -0,0 +1,117 @@
	1	+#ifndef K_MAX_SHAPE_DIM
	2	+#define K_MAX_SHAPE_DIM 0
	3	+#endif
	4	+
	5	+#include "catlass/arch/arch.hpp"
	6	+#include "catlass/catlass.hpp"
	7	+#include "catlass/gemm/block/block_mmad.hpp"
	8	+#include "catlass/gemm/block/block_swizzle.hpp"
	9	+#include "catlass/gemm/dispatch_policy.hpp"
	10	+#include "catlass/gemm/kernel/trmm.hpp"
	11	+#include "catlass/gemm_coord.hpp"
	12	+#include "catlass/layout/layout.hpp"
	13	+#include "catlass/matrix_coord.hpp"
	14	+#include "catlass/status.hpp"
	15	+#include "tla/layout.hpp"
	16	+#include "tla/tensor.hpp"
	17	+
	18	+#include "catlass_kernel.h"
	19	+#include "common/kernel_runner.h"
	20	+#include "common/tile_shape_scaler_tla.h"
	21	+
	22	+#include <type_traits>
	23	+
	24	+using namespace Catlass;
	25	+
	26	+#ifdef CATLASS_JIT_ELEMENT_A
	27	+using ElementA = CATLASS_JIT_ELEMENT_A;
	28	+#else
	29	+using ElementA = float;
	30	+#endif
	31	+#ifdef CATLASS_JIT_ELEMENT_B
	32	+using ElementB = CATLASS_JIT_ELEMENT_B;
	33	+#else
	34	+using ElementB = float;
	35	+#endif
	36	+#ifdef CATLASS_JIT_ELEMENT_C
	37	+using ElementC = CATLASS_JIT_ELEMENT_C;
	38	+#else
	39	+using ElementC = float;
	40	+#endif
	41	+#ifdef CATLASS_JIT_LAYOUT_A
	42	+using LayoutTagA = layout::CATLASS_JIT_LAYOUT_A;
	43	+#else
	44	+using LayoutTagA = layout::RowMajor;
	45	+#endif
	46	+#ifdef CATLASS_JIT_LAYOUT_B
	47	+using LayoutTagB = layout::CATLASS_JIT_LAYOUT_B;
	48	+#else
	49	+using LayoutTagB = layout::RowMajor;
	50	+#endif
	51	+#ifdef CATLASS_JIT_LAYOUT_C
	52	+using LayoutTagC = layout::CATLASS_JIT_LAYOUT_C;
	53	+#else
	54	+using LayoutTagC = layout::RowMajor;
	55	+#endif
	56	+
	57	+#ifdef CATLASS_JIT_TRMM_TILE_VARIANT
	58	+constexpr uint32_t trmmTileVariant = CATLASS_JIT_TRMM_TILE_VARIANT;
	59	+#else
	60	+constexpr uint32_t trmmTileVariant = 0;
	61	+#endif
	62	+
	63	+#ifdef CATLASS_JIT_BLOCK_SCHEDULER
	64	+constexpr uint32_t blockSchedulerCode = CATLASS_JIT_BLOCK_SCHEDULER;
	65	+#else
	66	+constexpr uint32_t blockSchedulerCode = 30;
	67	+#endif
	68	+
	69	+using BaseL1 = std::conditional_t<
	70	+ (trmmTileVariant == 1 \|\| trmmTileVariant == 2), tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>,
	71	+ std::conditional_t<
	72	+ (trmmTileVariant == 3 \|\| trmmTileVariant == 7), tla::tuple<tla::C<128>, tla::C<128>, tla::C<128>>,
	73	+ tla::tuple<tla::C<128>, tla::C<128>, tla::C<256>>>>;
	74	+using BaseL0 = tla::tuple<tla::C<128>, tla::C<128>, tla::C<64>>;
	75	+
	76	+using L1TileShape = typename CatlassKernel::TileShapeScalerTLA<ElementA, float, BaseL1>::type;
	77	+using L0TileShape = typename CatlassKernel::TileShapeScalerTLA<ElementA, float, BaseL0>::type;
	78	+using BlockScheduler = Gemm::Block::GemmIdentityBlockSwizzle<(blockSchedulerCode / 10), (blockSchedulerCode % 10)>;
	79	+
	80	+using ArchTag = Arch::AtlasA2;
	81	+constexpr bool enableUnitFlag = true;
	82	+constexpr bool useHF32 = false;
	83	+using DispatchPolicy = Gemm::MmadPingpong<ArchTag, enableUnitFlag, useHF32>;
	84	+using TileCopy =
	85	+ Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC>;
	86	+using BlockMmad =
	87	+ Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, void, TileCopy>;
	88	+using BlockEpilogue = void;
	89	+
	90	+using TrmmKernel = Gemm::Kernel::Trmm<BlockMmad, BlockEpilogue, BlockScheduler>;
	91	+
	92	+extern "C" void run(uint32_t blockNum, aclrtStream stream, const CatlassKernel::TrmmParams* params)
	93	+{
	94	+ if (params == nullptr) {
	95	+ return;
	96	+ }
	97	+
	98	+ auto layoutA = tla::MakeLayout<ElementA, LayoutTagA>(params->m, params->k);
	99	+ auto layoutB = tla::MakeLayout<ElementB, LayoutTagB>(params->k, params->n);
	100	+ auto layoutC = tla::MakeLayout<ElementC, LayoutTagC>(params->m, params->n);
	101	+
	102	+ typename TrmmKernel::Arguments arguments{
	103	+ GemmCoord{params->m, params->n, params->k},
	104	+ params->inputAddr[0],
	105	+ layoutA,
	106	+ params->inputAddr[1],
	107	+ layoutB,
	108	+ params->outputAddr[0],
	109	+ layoutC,
	110	+ params->side,
	111	+ params->uplo,
	112	+ params->trans,
	113	+ params->diag,
	114	+ params->alpha};
	115	+
	116	+ Catlass::RunKernel<TrmmKernel>(arguments, stream, blockNum);
	117	+}

Mtests/optest/kernels/CMakeLists.txt+1-0

	@@ -278,6 +278,7 @@ add_subdirectory(74_ascend950_weight_quant_a8w4_grouped_mx_matmul)
278	add_subdirectory(56_ascend950_basic_conv2d_tla)	278	add_subdirectory(56_ascend950_basic_conv2d_tla)
279	add_subdirectory(72_ascend950_fp8_mx_flash_attention_infer)	279	add_subdirectory(72_ascend950_fp8_mx_flash_attention_infer)
280	add_subdirectory(75_symm)	280	add_subdirectory(75_symm)
		281	+add_subdirectory(76_trmm)
281	add_subdirectory(80_grouped_matmul_slice_m_gelu)	282	add_subdirectory(80_grouped_matmul_slice_m_gelu)
282		283
283	if(_JIT_KERNEL_ENTRIES)	284	if(_JIT_KERNEL_ENTRIES)

Mtests/optest/src/catlass_torch.cpp+4-0

	@@ -48,6 +48,7 @@
48	#include "template/a8w4_mx_matmul.h"	48	#include "template/a8w4_mx_matmul.h"
49	#include "template/a8w4_grouped_mx_matmul.h"	49	#include "template/a8w4_grouped_mx_matmul.h"
50	#include "template/svd_quant_matmul.h"	50	#include "template/svd_quant_matmul.h"
		51	+#include "template/trmm.h"
51	#include "template/conv_bias.h"	52	#include "template/conv_bias.h"
52	#include "template/symm.h"	53	#include "template/symm.h"
53	// ── Workspace allocator bridge ──	54	// ── Workspace allocator bridge ──
	@@ -251,6 +252,9 @@ using A2Fp8E4M3MatmulOp = MatmulLike<CatlassKernel::A2Fp8E4M3Matmul>;
251	static auto& a2_fp8_e4m3_matmul = A2Fp8E4M3MatmulOp::Run;	252	static auto& a2_fp8_e4m3_matmul = A2Fp8E4M3MatmulOp::Run;
252	REGISTER_TORCH_FUNC(a2_fp8_e4m3_matmul);	253	REGISTER_TORCH_FUNC(a2_fp8_e4m3_matmul);
253		254
		255	+using TrmmOp = TrmmLike<CatlassKernel::Trmm>;
		256	+static auto& trmm = TrmmOp::Run;
		257	+REGISTER_TORCH_FUNC(trmm);
254	static auto& basic_conv2d = BasicConv2dOp::Run;	258	static auto& basic_conv2d = BasicConv2dOp::Run;
255	REGISTER_TORCH_FUNC(basic_conv2d);	259	REGISTER_TORCH_FUNC(basic_conv2d);
256		260

Atests/optest/src/include/template/trmm.h+93-0

@@ -0,0 +1,93 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef OPTEST_TRMM_H
	13	+#define OPTEST_TRMM_H
	14	+
	15	+#include <torch/torch.h>
	16	+#include <tiling/platform/platform_ascendc.h>
	17	+
	18	+#include "catlass_kernel_jit.h"
	19	+#include "common/run_npu_func.h"
	20	+#include "torch_utils.h"
	21	+
	22	+namespace CatlassKernelWrapper {
	23	+
	24	+using TrmmKernelFn =
	25	+ void (*)(const uint32_t, aclrtStream, const CatlassKernel::TParams&, const CatlassKernel::TrmmParams&);
	26	+
	27	+template <TrmmKernelFn KernelFunc>
	28	+struct TrmmLike {
	29	+ using OutputType = at::Tensor;
	30	+
	31	+ static OutputType Run(
	32	+ const at::Tensor& mat1, const at::Tensor& mat2, const c10::ScalarType& outDType, int64_t side, int64_t uplo,
	33	+ int64_t trans, int64_t diag, double alpha)
	34	+ {
	35	+ TORCH_CHECK(mat1.dim() == 2 && mat2.dim() == 2, "trmm expects two 2-D tensors");
	36	+ TORCH_CHECK(side == 0 \|\| side == 1, "trmm side must be 0 (left) or 1 (right)");
	37	+ TORCH_CHECK(uplo == 0 \|\| uplo == 1, "trmm uplo must be 0 (lower) or 1 (upper)");
	38	+ TORCH_CHECK(trans == 0 \|\| trans == 1, "trmm trans must be 0 (no transpose) or 1 (transpose)");
	39	+ TORCH_CHECK(diag == 0, "trmm currently supports only diag=0 (non-unit diagonal)");
	40	+ TORCH_CHECK(
	41	+ mat1.scalar_type() == at::kFloat && mat2.scalar_type() == at::kFloat,
	42	+ "trmm currently supports float32 inputs");
	43	+ TORCH_CHECK(outDType == at::kFloat, "trmm currently supports float32 output");
	44	+
	45	+ CatlassKernel::TParams tParams;
	46	+ CatlassKernel::TrmmParams params;
	47	+
	48	+ tParams.element["A"] = TorchDtypeToAclDtype(mat1.scalar_type());
	49	+ tParams.element["B"] = TorchDtypeToAclDtype(mat2.scalar_type());
	50	+ tParams.element["C"] = TorchDtypeToAclDtype(outDType);
	51	+ tParams.transpose["A"] = (side == 0 && trans == 1);
	52	+ tParams.transpose["B"] = (side == 1 && trans == 1);
	53	+ tParams.useNz["A"] = false;
	54	+ tParams.useNz["B"] = false;
	55	+ tParams.useNz["C"] = false;
	56	+
	57	+ if (side == 0) {
	58	+ TORCH_CHECK(mat1.size(0) == mat1.size(1), "left trmm expects mat1 to be square");
	59	+ TORCH_CHECK(mat1.size(1) == mat2.size(0), "left trmm shape mismatch");
	60	+ params.m = static_cast<uint32_t>(mat1.size(0));
	61	+ params.k = static_cast<uint32_t>(mat1.size(1));
	62	+ params.n = static_cast<uint32_t>(mat2.size(1));
	63	+ } else {
	64	+ TORCH_CHECK(mat2.size(0) == mat2.size(1), "right trmm expects mat2 to be square");
	65	+ TORCH_CHECK(mat1.size(1) == mat2.size(0), "right trmm shape mismatch");
	66	+ params.m = static_cast<uint32_t>(mat1.size(0));
	67	+ params.k = static_cast<uint32_t>(mat2.size(0));
	68	+ params.n = static_cast<uint32_t>(mat2.size(1));
	69	+ }
	70	+
	71	+ params.side = static_cast<uint32_t>(side);
	72	+ params.uplo = static_cast<uint32_t>(uplo);
	73	+ params.trans = static_cast<uint32_t>(trans);
	74	+ params.diag = static_cast<uint32_t>(diag);
	75	+ params.alpha = static_cast<float>(alpha);
	76	+ params.inputAddr.resize(2);
	77	+ params.inputAddr[0] = static_cast<uint8_t>(const_cast<void>(mat1.storage().data()));
	78	+ params.inputAddr[1] = static_cast<uint8_t>(const_cast<void>(mat2.storage().data()));
	79	+
	80	+ OutputType output = GetOutputTensor({params.m, params.n}, AclDtypeToTorchDtype(tParams.elem("C")));
	81	+ params.outputAddr.resize(1);
	82	+ params.outputAddr[0] = static_cast<uint8_t>(const_cast<void>(output.storage().data()));
	83	+
	84	+ aclrtStream stream = c10_npu::getCurrentNPUStream().stream(false);
	85	+ uint32_t coreNum = platform_ascendc::PlatformAscendCManager::GetInstance()->GetCoreNumAic();
	86	+ RUN_NPU_FUNC(KernelFunc, coreNum, stream, tParams, params);
	87	+ return output;
	88	+ }
	89	+};
	90	+
	91	+} // namespace CatlassKernelWrapper
	92	+
	93	+#endif // OPTEST_TRMM_H

Atests/optest/tests/test_76_trmm.py+70-0

@@ -0,0 +1,70 @@
	1	+import pytest
	2	+import torch
	3	+import torch_npu
	4	+
	5	+from common import only_on_2201
	6	+
	7	+pytestmark = [
	8	+ only_on_2201,
	9	+ pytest.mark.skipif(
	10	+ torch_npu.npu.device_count() <= 0,
	11	+ reason="torch-catlass integration tests require an available Ascend NPU",
	12	+ ),
	13	+]
	14	+
	15	+
	16	+def _triangular(x, uplo):
	17	+ return torch.tril(x) if uplo == 0 else torch.triu(x)
	18	+
	19	+
	20	+def _rand_triangular(size, dtype, uplo):
	21	+ return _triangular(torch.randn(size, size, dtype=dtype), uplo)
	22	+
	23	+
	24	+@pytest.mark.parametrize(
	25	+ "side,uplo,trans,m,n",
	26	+ [
	27	+ (0, 0, 0, 128, 96),
	28	+ (0, 1, 1, 128, 64),
	29	+ (1, 0, 0, 96, 128),
	30	+ (1, 1, 1, 64, 128),
	31	+ (0, 0, 1, 128, 96),
	32	+ (0, 1, 0, 128, 64),
	33	+ (1, 0, 1, 96, 128),
	34	+ (1, 1, 0, 64, 128),
	35	+ ],
	36	+)
	atomgit-bot7月10日 🟡 Medium Priority `tests/optest/tests/test_74_trmm_tla.py` 第 19–27 行的参数化组合全部满足 `uplo == trans`（即 `effectiveUplo = uplo ^ trans = 0`）：缺失的组合（uplo≠trans → effectiveUplo=1）：当 `effectiveUplo==1` 时，kernel AIC/AIV 路径会触发 `kStart >= kEnd` 的 `continue` 分支（`trmm_tla.hpp` 第 177–178 行、第 254–256 行），该路径涉及 K 范围裁剪和核间同步跳过逻辑。不测试此路径意味着 K 范围裁剪的正确性以及 Continue 路径下 SetFlag/WaitFlag 同步对的一致性均未验证。建议：在 parametrize 列表中补充 effectiveUplo==1 的组合，例如 (0, 0, 1, 128, 96)、(0, 1, 0, 128, 64)、(1, 0, 1, 96, 128)、(1, 1, 0, 64, 128)，确保 K 范围裁剪和 continue 路径均被覆盖。不准确？
	37	+@pytest.mark.parametrize("alpha", [1.0, 0.5])
	38	+def test_trmm(side, uplo, trans, m, n, alpha):
	39	+ import torch_catlass
	40	+
	41	+ dtype = torch.float32
	42	+
	43	+ if side == 0:
	44	+ tri_cpu = _rand_triangular(m, dtype, uplo)
	45	+ dense_cpu = torch.randn(m, n, dtype=dtype)
	46	+ tri = tri_cpu.to("npu")
	47	+ dense = dense_cpu.to("npu")
	48	+ result = torch_catlass.trmm(tri, dense, dtype, side, uplo, trans, 0, alpha)
	49	+ tri_op = tri_cpu.T if trans else tri_cpu
	50	+ expected = alpha * torch.matmul(tri_op, dense_cpu)
	51	+ else:
	52	+ dense_cpu = torch.randn(m, n, dtype=dtype)
	53	+ tri_cpu = _rand_triangular(n, dtype, uplo)
	54	+ dense = dense_cpu.to("npu")
	55	+ tri = tri_cpu.to("npu")
	56	+ result = torch_catlass.trmm(dense, tri, dtype, side, uplo, trans, 0, alpha)
	57	+ tri_op = tri_cpu.T if trans else tri_cpu
	58	+ expected = alpha * torch.matmul(dense_cpu, tri_op)
	59	+
	60	+ assert result.shape == (m, n)
	61	+ assert result.dtype == dtype
	62	+ assert result.device.type == "npu"
	63	+ result_cpu = result.cpu()
	64	+ assert torch.allclose(result_cpu, expected, rtol=1e-2, atol=1e-2), (
	65	+ f"Results not close: max diff = {(result_cpu - expected).abs().max().item()}"
	66	+ )
	67	+
	68	+
	69	+if __name__ == "__main__":
	70	+ pytest.main([__file__, "-v", "-s"])

Mtests/optest/torch_catlass/ops/__init__.py+2-0

	@@ -106,6 +106,7 @@ from .streamk_matmul import streamk_matmul # example 37
106	from .strided_batched_matmul_tla import strided_batched_matmul_tla # example 45	106	from .strided_batched_matmul_tla import strided_batched_matmul_tla # example 45
107	from .svd_quant_matmul import ascend950_svd_quant_matmul # example 61	107	from .svd_quant_matmul import ascend950_svd_quant_matmul # example 61
108	from .symm import symm # example 75	108	from .symm import symm # example 75
		109	+from .trmm import trmm # example 76
109	from .tail_multi_core_splitk_matmul import ascend950_tail_multi_core_splitk_matmul # example 69	110	from .tail_multi_core_splitk_matmul import ascend950_tail_multi_core_splitk_matmul # example 69
110	from .w4a4_matmul_per_token_per_channel_dequant import (	111	from .w4a4_matmul_per_token_per_channel_dequant import (
111	w4a4_matmul_per_token_per_channel_dequant, # example 38	112	w4a4_matmul_per_token_per_channel_dequant, # example 38
	@@ -195,5 +196,6 @@ __all__ = [
195	"ascend950_basic_conv2d_tla", # example 56	196	"ascend950_basic_conv2d_tla", # example 56
196	"ascend950_matmul_full_loadA", # example 73	197	"ascend950_matmul_full_loadA", # example 73
197	"symm", # example 75	198	"symm", # example 75
		199	+ "trmm", # example 76
198	"grouped_matmul_slice_m_gelu", # example 80	200	"grouped_matmul_slice_m_gelu", # example 80
199	]	201	]

Atests/optest/torch_catlass/ops/trmm.py+49-0

@@ -0,0 +1,49 @@
	1	+import torch
	2	+from torch import Tensor
	3	+
	4	+_DTYPE_MAP = {
	5	+ "float": torch.float32,
	6	+ "float32": torch.float32,
	7	+ "torch.float32": torch.float32,
	8	+}
	9	+
	10	+
	11	+def _normalize_dtype(dtype: str \| torch.dtype) -> torch.dtype:
	12	+ if isinstance(dtype, torch.dtype):
	13	+ return dtype
	14	+ try:
	15	+ return _DTYPE_MAP[dtype.lower()]
	16	+ except KeyError as exc:
	17	+ raise ValueError(f"{dtype} is not supported by torch_catlass.trmm") from exc
	18	+
	19	+
	20	+def trmm(
	21	+ mat1: Tensor,
	22	+ mat2: Tensor,
	23	+ out_dtype: str \| torch.dtype = torch.float32,
	24	+ side: int = 0,
	25	+ uplo: int = 0,
	26	+ trans: int = 0,
	27	+ diag: int = 0,
	28	+ alpha: float = 1.0,
	29	+) -> Tensor:
	30	+ """Run CATLASS TRMM on NPU tensors.
	31	+
	32	+ Source: example 76_trmm.
	33	+
	34	+ Args:
	35	+ mat1: Left input. When ``side=0`` this is the triangular matrix.
	36	+ mat2: Right input. When ``side=1`` this is the triangular matrix.
	37	+ out_dtype: Output dtype. The current TRMM kernel supports ``float32``.
	38	+ side: ``0`` computes ``op(mat1) @ mat2``; ``1`` computes ``mat1 @ op(mat2)``.
	39	+ uplo: ``0`` keeps the lower triangle; ``1`` keeps the upper triangle.
	40	+ trans: Whether to transpose the triangular input.
	41	+ diag: Diagonal mode. The current TRMM kernel supports only ``0``.
	42	+ alpha: Output scaling factor.
	43	+
	44	+ Returns:
	45	+ Output tensor with shape ``(M, N)`` on the active NPU device.
	46	+ """
	47	+ return torch.ops.catlass.trmm(
	48	+ mat1, mat2, _normalize_dtype(out_dtype), side, uplo, trans, diag, alpha
	49	+ )