dsl: support l0c2l1 zN-catlass-AtomGit

Minclude/catlass/gemm/tile/ascend950/copy_l0c_to_dst.hpp+12-0

	@@ -273,6 +273,18 @@ struct CopyL0CToGmTla {
273	static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l0c to gm, can not find the specialization.");	273	static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l0c to gm, can not find the specialization.");
274	};	274	};
275		275
		276	+template <
		277	+ class ArchTag,
		278	+ class TensorSrc,
		279	+ class TensorDst,
		280	+ ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT,
		281	+ bool ReluEnable = false,
		282	+ class Enable = void
		283	+>
		284	+struct CopyL0CToL1Tla {
		285	+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l0c to l1, can not find the specialization.");
		286	+};
		287	+
276	/////////////////////////////////////////////////////////////////////////////////////////////////////////////	288	/////////////////////////////////////////////////////////////////////////////////////////////////////////////
277		289
278	enum class CopyL0CToUBMode {	290	enum class CopyL0CToUBMode {

Ainclude/catlass/gemm/tile/ascend950/copy_l0c_to_l1.hpp+67-0

@@ -0,0 +1,67 @@
	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	+#ifndef CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_L1_HPP
	12	+#define CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_L1_HPP
	13	+
	14	+#include "catlass/arch/arch.hpp"
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/gemm/tile/ascend950/copy_l0c_to_dst.hpp"
	17	+#include "tla/tensor.hpp"
	18	+
	19	+namespace Catlass::Gemm::Tile {
	20	+
	21	+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
	22	+struct CopyL0CToL1Tla<
	23	+ Catlass::Arch::Ascend950, TensorSrc_,
	24	+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::A1>,
	25	+ ScaleGranularity::NO_QUANT, ReluEnable_, std::enable_if_t<tla::detail::iszN<ElementDst_, LayoutDst_>::value>> {
	26	+ using ArchTag = Catlass::Arch::Ascend950;
	27	+ using ElementDst = ElementDst_;
	28	+ using ElementSrc = typename TensorSrc_::Element;
	29	+ static constexpr auto quantPre =
	30	+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE;
	31	+ static constexpr auto reluEn = ReluEnable_;
	32	+
	33	+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementDst); // L1
	34	+
	35	+ template <class TensorDst, class TensorSrc>
	36	+ CATLASS_DEVICE void operator()(TensorDst const& dstTensor, TensorSrc const& srcTensor, uint8_t unitFlag = 0)
	37	+ {
	38	+ static_assert(
	39	+ tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value &&
	40	+ TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::A1,
	41	+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be L1 and zN");
	42	+
	43	+ AscendC::DataCopyCO12DstParams intriParams;
	44	+
	45	+ intriParams.nSize = RoundUp<ELE_NUM_PER_C0>(tla::get<1>(dstTensor.originShape()));
	46	+ intriParams.mSize = tla::get<0>(dstTensor.originShape());
	47	+ intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0;
	48	+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
	49	+ intriParams.quantPre = quantPre;
	50	+ intriParams.nz2ndEn = false;
	51	+ intriParams.reluPre = reluEn;
	52	+ intriParams.unitFlag = unitFlag;
	53	+
	54	+ if constexpr (std::is_same_v<ElementSrc, float> && std::is_same_v<ElementDst, float>) {
	55	+ intriParams.channelSplit = true;
	56	+ }
	57	+
	58	+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
	59	+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
	60	+
	61	+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
	62	+ }
	63	+};
	64	+
	65	+} // namespace Catlass::Gemm::Tile
	66	+
	67	+#endif // CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_L1_HPP

Ainclude/catlass/gemm/tile/copy_l0c_to_l1.hpp+18-0

@@ -0,0 +1,18 @@
	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	+#ifndef CATLASS_GEMM_TILE_COPY_L0C_TO_L1_HPP
	12	+#define CATLASS_GEMM_TILE_COPY_L0C_TO_L1_HPP
	13	+
	14	+#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510)
	15	+#include "catlass/gemm/tile/ascend950/copy_l0c_to_l1.hpp"
	16	+#endif
	17	+
	18	+#endif

Mpython/tla_dsl/catlass/core_api.py+8-6

	@@ -4231,7 +4231,7 @@ Description:
4231		4231
4232	# (src addrspace, dst addrspace) routes and the region each must be nested in.	4232	# (src addrspace, dst addrspace) routes and the region each must be nested in.
4233	_COPY_CUBE_ROUTES = {	4233	_COPY_CUBE_ROUTES = {
4234	- ("gm", "l1"), ("l1", "l0a"), ("l1", "l0b"), ("l0c", "gm"),	4234	+ ("gm", "l1"), ("l1", "l0a"), ("l1", "l0b"), ("l0c", "gm"), ("l0c", "l1"),
		tianxinghuiAug 18 新路径未同步到 dialect 约束。 [TlaOps.cpp (line 769)] [Tla.td (line 654)] [core_api.py (line 4245)] 仍遗漏 L0C → L1
4235	("l0c", "ub"), ("l1", "ub"),	4235	("l0c", "ub"), ("l1", "ub"),
4236	}	4236	}
4237	_COPY_VECTOR_ROUTES = {("gm", "ub"), ("ub", "gm"), ("ub", "l1")}	4237	_COPY_VECTOR_ROUTES = {("gm", "ub"), ("ub", "gm"), ("ub", "l1")}
	@@ -4242,7 +4242,7 @@ def copy(dst: Tensor, src: Tensor, params: CopyParams \| None = None, *, loc: mli
4242	"""Directory: Data Movement	4242	"""Directory: Data Movement
4243	Description:	4243	Description:
4244	Copy data between tiles. The hardware path follows `src`/`dst` address spaces	4244	Copy data between tiles. The hardware path follows `src`/`dst` address spaces
4245	- (vector: GM↔UB, UB→L1; cube: GM→L1, L1→L0A/L0B, L0C→GM\|UB, L1→UB).	4245	+ (vector: GM↔UB, UB→L1; cube: GM→L1, L1→L0A/L0B, L0C→GM\|UB\|L1, L1→UB).
4246	Layout tags on the tiles select format conversion (for example ND→zN).	4246	Layout tags on the tiles select format conversion (for example ND→zN).
4247		4247
4248	Copy / tiling sizes follow each tile's logical `origin_shape`	4248	Copy / tiling sizes follow each tile's logical `origin_shape`
	@@ -4339,12 +4339,12 @@ Description:
4339	src_dtype = src_desc.element_type.lower()	4339	src_dtype = src_desc.element_type.lower()
4340	dst_dtype = dst_desc.element_type.lower()	4340	dst_dtype = dst_desc.element_type.lower()
4341	same_dtype = (src_dtype == dst_dtype)	4341	same_dtype = (src_dtype == dst_dtype)
4342	- src_layout = src_desc.layout_tag.lower()	4342	+ src_layout = src_desc.layout_tag
4343	- dst_layout = dst_desc.layout_tag.lower()	4343	+ dst_layout = dst_desc.layout_tag
4344		4344
4345	if _route[0] == "l0c":	4345	if _route[0] == "l0c":
4346	- if src_layout != "l0clayout":	4346	+ if src_layout != "L0Clayout":
4347	- raise TlaLoweringError(f"L0C layout_tag only support l0clayout, got {src_layout}")	4347	+ raise TlaLoweringError(f"L0C layout_tag only support L0Clayout, got {src_layout}")
4348	if src_dtype != "f32":	4348	if src_dtype != "f32":
4349	raise NotImplementedError(f"currently l0c dtype only support f32, got {src_dtype}")	4349	raise NotImplementedError(f"currently l0c dtype only support f32, got {src_dtype}")
4350	if src_dtype == "f32" and dst_dtype not in ("f32", "f16", "bf16"):	4350	if src_dtype == "f32" and dst_dtype not in ("f32", "f16", "bf16"):
	@@ -4356,6 +4356,8 @@ Description:
4356	"currently copy l0c to ub only support dst [row_major, column_major],"	4356	"currently copy l0c to ub only support dst [row_major, column_major],"
4357	f" got {dst_layout}"	4357	f" got {dst_layout}"
4358	)	4358	)
		4359	+ if _route[1] == "l1" and dst_layout not in ("zN",):
		4360	+ raise TlaLoweringError(f"l0c2l1 dst layout shoud be zN, got {dst_layout}")
4359		4361
4360	if params is None:	4362	if params is None:
4361	params = CopyL0C2DstParams() # use default	4363	params = CopyL0C2DstParams() # use default

Mpython/tla_dsl/csrc/mlir/bc/Cube/dma.cpp+26-0

	@@ -3,6 +3,7 @@
3	#include "catlass/gemm/tile/copy_gm_to_l1.hpp"	3	#include "catlass/gemm/tile/copy_gm_to_l1.hpp"
4	#include "catlass/gemm/tile/copy_l0c_to_gm.hpp"	4	#include "catlass/gemm/tile/copy_l0c_to_gm.hpp"
5	#include "catlass/gemm/tile/copy_l0c_to_ub.hpp"	5	#include "catlass/gemm/tile/copy_l0c_to_ub.hpp"
		6	+#include "catlass/gemm/tile/copy_l0c_to_l1.hpp"
6	#include "catlass/gemm/tile/copy_l1_to_l0a.hpp"	7	#include "catlass/gemm/tile/copy_l1_to_l0a.hpp"
7	#include "catlass/gemm/tile/copy_l1_to_l0b.hpp"	8	#include "catlass/gemm/tile/copy_l1_to_l0b.hpp"
8	#include "catlass/epilogue/tile/tile_copy.hpp"	9	#include "catlass/epilogue/tile/tile_copy.hpp"
	@@ -52,6 +53,18 @@ CATLASS_DEVICE void copyL0CToGM(
52	srcTensor, unitFlag);	53	srcTensor, unitFlag);
53	}	54	}
54		55
		56	+// L0C holds fp32 MMAD accumulator; ElementDst may be f32/f16/bf16 (Ascend950 fixpipe).
		57	+template <class ArchTag, LayoutTag DstLayout, typename ElementSrc, typename ElementDst>
		58	+CATLASS_DEVICE void copyL0CToL1(
		59	+ memref_t<__cc__ ElementSrc, 1> src, memref_t<__cbuf__ ElementDst, 1> dst,
		60	+ const TensorDesc &srcDesc, const TensorDesc &dstDesc, uint8_t unitFlag) {
		61	+ auto srcTensor = makeL0CTensor<LayoutTag::L0C, ElementSrc>(src, srcDesc);
		62	+ auto dstTensor = makeL1Tensor<DstLayout, ElementDst>(dst, dstDesc);
		63	+ Catlass::Gemm::Tile::CopyL0CToL1Tla<ArchTag, decltype(srcTensor),
		64	+ decltype(dstTensor)>{}(dstTensor,
		65	+ srcTensor, unitFlag);
		66	+}
		67	+
55	template <class ArchTag, LayoutTag DstLayout,	68	template <class ArchTag, LayoutTag DstLayout,
56	Catlass::Gemm::Tile::CopyL0CToUBMode Mode, typename ElementSrc, typename ElementDst>	69	Catlass::Gemm::Tile::CopyL0CToUBMode Mode, typename ElementSrc, typename ElementDst>
57	CATLASS_DEVICE void copyL0CToUB(	70	CATLASS_DEVICE void copyL0CToUB(
	@@ -135,6 +148,19 @@ REGISTER_L0C_TO_GM(float, float)
135	REGISTER_L0C_TO_GM(float, half)	148	REGISTER_L0C_TO_GM(float, half)
136	REGISTER_L0C_TO_GM(float, bf16)	149	REGISTER_L0C_TO_GM(float, bf16)
137		150
		151	+#define REGISTER_L0C_TO_L1(DTypeSrc, DTypeDst) \
		152	+[aicore] __attribute__((always_inline)) \
		153	+void _mlir_ciface_copy_l0c_to_l1_zN_##DTypeDst( \
		154	+ memref_t<__cc__ DTypeSrc, 1> src, memref_t<__cbuf__ DTypeDst, 1> dst, \
		155	+ DESC_ABI_PARAMS(src), DESC_ABI_PARAMS(dst), uint8_t unitFlag) { \
		156	+ copyL0CToL1<Catlass::Arch::Ascend950, LayoutTag::zN, DTypeSrc, DTypeDst>( \
		157	+ src, dst, TENSOR_DESC_12(src), TENSOR_DESC_12(dst), unitFlag); \
		158	+}
		159	+
		160	+REGISTER_L0C_TO_L1(float, float)
		161	+REGISTER_L0C_TO_L1(float, half)
		162	+REGISTER_L0C_TO_L1(float, bf16)
		163	+
138	#define REGISTER_L0C_TO_UB(NameDst, EnumDst, mode, MODE, DTypeSrc, DTypeDst) \	164	#define REGISTER_L0C_TO_UB(NameDst, EnumDst, mode, MODE, DTypeSrc, DTypeDst) \
139	[aicore] __attribute__((always_inline)) \	165	[aicore] __attribute__((always_inline)) \
140	void _mlir_ciface_copy_l0c_to_ub_##NameDst##_##mode##_##DTypeDst( \	166	void _mlir_ciface_copy_l0c_to_ub_##NameDst##_##mode##_##DTypeDst( \

Mpython/tla_dsl/csrc/mlir/include/Dialect/Tla/IR/Tla.td+1-1

	@@ -651,7 +651,7 @@ def Tla_CopyOp : Tla_Op<"copy", []> {
651	let summary = "Copy";	651	let summary = "Copy";
652	let description = [{	652	let description = [{
653	Copy between tensors/views in different address spaces. Cube data-path	653	Copy between tensors/views in different address spaces. Cube data-path
654	- copies (GM->L1, L1->L0A, L1->L0B, L0C->GM, L0C->UB, L1->UB) must be nested	654	+ copies (GM->L1, L1->L0A, L1->L0B, L0C->GM, L0C->UB, L0C->L1, L1->UB) must be nested
655	inside a ``tla.cube`` region; vector staging copies (GM->UB, UB->GM, UB->L1)	655	inside a ``tla.cube`` region; vector staging copies (GM->UB, UB->GM, UB->L1)
656	must be nested inside a ``tla.vector`` region.	656	must be nested inside a ``tla.vector`` region.
657		657

Mpython/tla_dsl/csrc/mlir/lib/Dialect/Tla/IR/TlaOps.cpp+2-1

	@@ -766,12 +766,13 @@ mlir::LogicalResult CopyOp::verify() {
766	AddressSpace src = srcTy.getPtr().getAddrspace();	766	AddressSpace src = srcTy.getPtr().getAddrspace();
767	AddressSpace dst = dstTy.getPtr().getAddrspace();	767	AddressSpace dst = dstTy.getPtr().getAddrspace();
768		768
769	- // Cube data-path copies: GM->L1, L1->L0A, L1->L0B, L0C->GM, L0C->UB, L1->UB.	769	+ // Cube data-path copies: GM->L1, L1->L0A, L1->L0B, L0C->GM, L0C->UB, L0C->L1, L1->UB.
770	bool cubeRoute = (src == AddressSpace::gm && dst == AddressSpace::l1) \|\|	770	bool cubeRoute = (src == AddressSpace::gm && dst == AddressSpace::l1) \|\|
771	(src == AddressSpace::l1 && dst == AddressSpace::l0a) \|\|	771	(src == AddressSpace::l1 && dst == AddressSpace::l0a) \|\|
772	(src == AddressSpace::l1 && dst == AddressSpace::l0b) \|\|	772	(src == AddressSpace::l1 && dst == AddressSpace::l0b) \|\|
773	(src == AddressSpace::l0c && dst == AddressSpace::gm) \|\|	773	(src == AddressSpace::l0c && dst == AddressSpace::gm) \|\|
774	(src == AddressSpace::l0c && dst == AddressSpace::ub) \|\|	774	(src == AddressSpace::l0c && dst == AddressSpace::ub) \|\|
		775	+ (src == AddressSpace::l0c && dst == AddressSpace::l1) \|\|
775	(src == AddressSpace::l1 && dst == AddressSpace::ub);	776	(src == AddressSpace::l1 && dst == AddressSpace::ub);
776	// Vector staging copies: GM->UB, UB->GM, UB->L1.	777	// Vector staging copies: GM->UB, UB->GM, UB->L1.
777	bool vectorRoute = (src == AddressSpace::gm && dst == AddressSpace::ub) \|\|	778	bool vectorRoute = (src == AddressSpace::gm && dst == AddressSpace::ub) \|\|

Mpython/tla_dsl/csrc/mlir/lib/Passes/TlaTensorToMemref.cpp+12-1

	@@ -741,6 +741,16 @@ std::string getCopyRouteCallee(MLIRContext *ctx, StringRef srcAddrspace, StringR
741	return {};	741	return {};
742	return Twine("copy_l0c_to_ub_column_major_").concat(extraDesc).concat("_").concat(suffix).str();	742	return Twine("copy_l0c_to_ub_column_major_").concat(extraDesc).concat("_").concat(suffix).str();
743	}	743	}
		744	+ // L0C (fp32 MMAD acc) -> L1 zN: dst may be f32 / f16 / bf16 (narrowing on fixpipe).
		745	+ if (srcSpace == hivm::AddressSpace::L0C && dstSpace == hivm::AddressSpace::L1 &&
		746	+ srcLayout == TensorLayoutTag::L0C && dstLayout == TensorLayoutTag::zN) {
		747	+ if (srcElementType != "f32")
		748	+ return {};
		749	+ StringRef suffix = copyRuntimeElemSuffix(dstElem);
		750	+ if (suffix.empty())
		751	+ return {};
		752	+ return Twine("copy_l0c_to_l1_zN_").concat(suffix).str();
		753	+ }
744	return {};	754	return {};
745	}	755	}
746		756
	@@ -792,7 +802,8 @@ static bool isAicTemplateRuntimeCall(StringRef name) {
792	name.starts_with("copy_l1_nZ_to_l0b_nZ_") \|\|	802	name.starts_with("copy_l1_nZ_to_l0b_nZ_") \|\|
793	name.starts_with("copy_l0c_to_ub_row_major_") \|\|	803	name.starts_with("copy_l0c_to_ub_row_major_") \|\|
794	name.starts_with("copy_l0c_to_gm_row_major_") \|\|	804	name.starts_with("copy_l0c_to_gm_row_major_") \|\|
795	- name.starts_with("copy_l0c_to_ub_column_major_");	805	+ name.starts_with("copy_l0c_to_ub_column_major_") \|\|
		806	+ name.starts_with("copy_l0c_to_l1_zN_");
796	}	807	}
797		808
798	static bool isAivTemplateRuntimeCall(StringRef name) {	809	static bool isAivTemplateRuntimeCall(StringRef name) {

Mpython/tla_dsl/docs/en/kernel_api_reference.md+48-48

	@@ -356,7 +356,7 @@ dst = tla.make_tensor_like(ptr, like=src_tile, layoutTag=tla.arch.RowMajor)
356		356
357	### `make_ptr`	357	### `make_ptr`
358		358
359	-Source: [`catlass.core_api.make_ptr`](../../catlass/core_api.py#L6977)	359	+Source: [`catlass.core_api.make_ptr`](../../catlass/core_api.py#L6979)
360		360
361	Description:	361	Description:
362		362
	@@ -391,7 +391,7 @@ ptr = tla.make_ptr(tla.Float16, addr, mem_space=tla.AddressSpace.gm)
391		391
392	### `recast_ptr`	392	### `recast_ptr`
393		393
394	-Source: [`catlass.core_api.recast_ptr`](../../catlass/core_api.py#L7031)	394	+Source: [`catlass.core_api.recast_ptr`](../../catlass/core_api.py#L7033)
395		395
396	Description:	396	Description:
397		397
	@@ -432,7 +432,7 @@ Tensor copies between on-chip and global memory, and UB register load/store.
432	Description:	432	Description:
433		433
434	Copy data between tiles. The hardware path follows `src`/`dst` address spaces	434	Copy data between tiles. The hardware path follows `src`/`dst` address spaces
435	-(vector: GM↔UB, UB→L1; cube: GM→L1, L1→L0A/L0B, L0C→GM\|UB, L1→UB).	435	+(vector: GM↔UB, UB→L1; cube: GM→L1, L1→L0A/L0B, L0C→GM\|UB\|L1, L1→UB).
436	Layout tags on the tiles select format conversion (for example ND→zN).	436	Layout tags on the tiles select format conversion (for example ND→zN).
437		437
438	Copy / tiling sizes follow each tile's logical `origin_shape`	438	Copy / tiling sizes follow each tile's logical `origin_shape`
	@@ -592,7 +592,7 @@ Cube-side matrix multiply-accumulate (`tla.mmad`).
592		592
593	### `mmad`	593	### `mmad`
594		594
595	-Source: [`catlass.core_api.mmad`](../../catlass/core_api.py#L5123)	595	+Source: [`catlass.core_api.mmad`](../../catlass/core_api.py#L5125)
596		596
597	Description:	597	Description:
598		598
	@@ -643,7 +643,7 @@ Mask creation and tail-mask updates.
643		643
644	#### `create_mask`	644	#### `create_mask`
645		645
646	-Source: [`catlass.core_api.create_mask`](../../catlass/core_api.py#L7289)	646	+Source: [`catlass.core_api.create_mask`](../../catlass/core_api.py#L7291)
647		647
648	Description:	648	Description:
649		649
	@@ -699,7 +699,7 @@ with tla.vec.func(mode="simd"):
699		699
700	#### `update_mask`	700	#### `update_mask`
701		701
702	-Source: [`catlass.core_api.update_mask`](../../catlass/core_api.py#L7355)	702	+Source: [`catlass.core_api.update_mask`](../../catlass/core_api.py#L7357)
703		703
704	Description:	704	Description:
705		705
	@@ -736,7 +736,7 @@ Element-wise arithmetic and unary math ops. `VectorSSA` overloads `+` / `-` / `*
736		736
737	#### `exp`	737	#### `exp`
738		738
739	-Source: [`catlass.core_api.exp`](../../catlass/core_api.py#L5708)	739	+Source: [`catlass.core_api.exp`](../../catlass/core_api.py#L5710)
740		740
741	Description:	741	Description:
742		742
	@@ -769,7 +769,7 @@ with tla.vec.func(mode="simd"):
769		769
770	#### `log`	770	#### `log`
771		771
772	-Source: [`catlass.core_api.log`](../../catlass/core_api.py#L5730)	772	+Source: [`catlass.core_api.log`](../../catlass/core_api.py#L5732)
773		773
774	Description:	774	Description:
775		775
	@@ -802,7 +802,7 @@ with tla.vec.func(mode="simd"):
802		802
803	#### `sqrt`	803	#### `sqrt`
804		804
805	-Source: [`catlass.core_api.sqrt`](../../catlass/core_api.py#L5752)	805	+Source: [`catlass.core_api.sqrt`](../../catlass/core_api.py#L5754)
806		806
807	Description:	807	Description:
808		808
	@@ -835,7 +835,7 @@ with tla.vec.func(mode="simd"):
835		835
836	#### `abs`	836	#### `abs`
837		837
838	-Source: [`catlass.core_api.abs`](../../catlass/core_api.py#L5774)	838	+Source: [`catlass.core_api.abs`](../../catlass/core_api.py#L5776)
839		839
840	Description:	840	Description:
841		841
	@@ -868,7 +868,7 @@ with tla.vec.func(mode="simd"):
868		868
869	#### `neg`	869	#### `neg`
870		870
871	-Source: [`catlass.core_api.neg`](../../catlass/core_api.py#L5796)	871	+Source: [`catlass.core_api.neg`](../../catlass/core_api.py#L5798)
872		872
873	Description:	873	Description:
874		874
	@@ -901,7 +901,7 @@ with tla.vec.func(mode="simd"):
901		901
902	#### `add`	902	#### `add`
903		903
904	-Source: [`catlass.core_api.add`](../../catlass/core_api.py#L5976)	904	+Source: [`catlass.core_api.add`](../../catlass/core_api.py#L5978)
905		905
906	Description:	906	Description:
907		907
	@@ -942,7 +942,7 @@ with tla.vec.func(mode="simd"):
942		942
943	#### `sub`	943	#### `sub`
944		944
945	-Source: [`catlass.core_api.sub`](../../catlass/core_api.py#L6022)	945	+Source: [`catlass.core_api.sub`](../../catlass/core_api.py#L6024)
946		946
947	Description:	947	Description:
948		948
	@@ -980,7 +980,7 @@ with tla.vec.func(mode="simd"):
980		980
981	#### `mul`	981	#### `mul`
982		982
983	-Source: [`catlass.core_api.mul`](../../catlass/core_api.py#L6059)	983	+Source: [`catlass.core_api.mul`](../../catlass/core_api.py#L6061)
984		984
985	Description:	985	Description:
986		986
	@@ -1019,7 +1019,7 @@ with tla.vec.func(mode="simd"):
1019		1019
1020	#### `max`	1020	#### `max`
1021		1021
1022	-Source: [`catlass.core_api.max`](../../catlass/core_api.py#L6104)	1022	+Source: [`catlass.core_api.max`](../../catlass/core_api.py#L6106)
1023		1023
1024	Description:	1024	Description:
1025		1025
	@@ -1053,7 +1053,7 @@ with tla.vec.func(mode="simd"):
1053		1053
1054	#### `min`	1054	#### `min`
1055		1055
1056	-Source: [`catlass.core_api.min`](../../catlass/core_api.py#L6144)	1056	+Source: [`catlass.core_api.min`](../../catlass/core_api.py#L6146)
1057		1057
1058	Description:	1058	Description:
1059		1059
	@@ -1087,7 +1087,7 @@ with tla.vec.func(mode="simd"):
1087		1087
1088	#### `div`	1088	#### `div`
1089		1089
1090	-Source: [`catlass.core_api.div`](../../catlass/core_api.py#L6184)	1090	+Source: [`catlass.core_api.div`](../../catlass/core_api.py#L6186)
1091		1091
1092	Description:	1092	Description:
1093		1093
	@@ -1129,7 +1129,7 @@ Bitwise and logical ops on Mask / Vector.
1129		1129
1130	#### `bitwise_not`	1130	#### `bitwise_not`
1131		1131
1132	-Source: [`catlass.core_api.bitwise_not`](../../catlass/core_api.py#L5942)	1132	+Source: [`catlass.core_api.bitwise_not`](../../catlass/core_api.py#L5944)
1133		1133
1134	Description:	1134	Description:
1135		1135
	@@ -1162,7 +1162,7 @@ with tla.vec.func(mode="simd"):
1162		1162
1163	#### `bitwise_and`	1163	#### `bitwise_and`
1164		1164
1165	-Source: [`catlass.core_api.bitwise_and`](../../catlass/core_api.py#L6553)	1165	+Source: [`catlass.core_api.bitwise_and`](../../catlass/core_api.py#L6555)
1166		1166
1167	Description:	1167	Description:
1168		1168
	@@ -1196,7 +1196,7 @@ with tla.vec.func(mode="simd"):
1196		1196
1197	#### `bitwise_or`	1197	#### `bitwise_or`
1198		1198
1199	-Source: [`catlass.core_api.bitwise_or`](../../catlass/core_api.py#L6591)	1199	+Source: [`catlass.core_api.bitwise_or`](../../catlass/core_api.py#L6593)
1200		1200
1201	Description:	1201	Description:
1202		1202
	@@ -1230,7 +1230,7 @@ with tla.vec.func(mode="simd"):
1230		1230
1231	#### `bitwise_xor`	1231	#### `bitwise_xor`
1232		1232
1233	-Source: [`catlass.core_api.bitwise_xor`](../../catlass/core_api.py#L6629)	1233	+Source: [`catlass.core_api.bitwise_xor`](../../catlass/core_api.py#L6631)
1234		1234
1235	Description:	1235	Description:
1236		1236
	@@ -1268,7 +1268,7 @@ Vector compares that produce masks, and masked select.
1268		1268
1269	#### `where`	1269	#### `where`
1270		1270
1271	-Source: [`catlass.core_api.where`](../../catlass/core_api.py#L6268)	1271	+Source: [`catlass.core_api.where`](../../catlass/core_api.py#L6270)
1272		1272
1273	Description:	1273	Description:
1274		1274
	@@ -1302,7 +1302,7 @@ with tla.vec.func(mode="simd"):
1302		1302
1303	#### `cmp`	1303	#### `cmp`
1304		1304
1305	-Source: [`catlass.core_api.cmp`](../../catlass/core_api.py#L6475)	1305	+Source: [`catlass.core_api.cmp`](../../catlass/core_api.py#L6477)
1306		1306
1307	Description:	1307	Description:
1308		1308
	@@ -1341,7 +1341,7 @@ Constant fill and lane-index sequence construction.
1341		1341
1342	#### `full`	1342	#### `full`
1343		1343
1344	-Source: [`catlass.core_api.full`](../../catlass/core_api.py#L5220)	1344	+Source: [`catlass.core_api.full`](../../catlass/core_api.py#L5222)
1345		1345
1346	Description:	1346	Description:
1347		1347
	@@ -1374,7 +1374,7 @@ with tla.vec.func(mode="simd"):
1374		1374
1375	#### `arange`	1375	#### `arange`
1376		1376
1377	-Source: [`catlass.core_api.arange`](../../catlass/core_api.py#L5293)	1377	+Source: [`catlass.core_api.arange`](../../catlass/core_api.py#L5295)
1378		1378
1379	Description:	1379	Description:
1380		1380
	@@ -1412,7 +1412,7 @@ Gather elements from a UB tensor by index.
1412		1412
1413	#### `gather`	1413	#### `gather`
1414		1414
1415	-Source: [`catlass.core_api.gather`](../../catlass/core_api.py#L6667)	1415	+Source: [`catlass.core_api.gather`](../../catlass/core_api.py#L6669)
1416		1416
1417	Description:	1417	Description:
1418		1418
	@@ -1450,7 +1450,7 @@ Interleave / deinterleave and related lane reshuffles.
1450		1450
1451	#### `interleave`	1451	#### `interleave`
1452		1452
1453	-Source: [`catlass.core_api.interleave`](../../catlass/core_api.py#L5833)	1453	+Source: [`catlass.core_api.interleave`](../../catlass/core_api.py#L5835)
1454		1454
1455	Description:	1455	Description:
1456		1456
	@@ -1483,7 +1483,7 @@ with tla.vec.func(mode="simd"):
1483		1483
1484	#### `deinterleave`	1484	#### `deinterleave`
1485		1485
1486	-Source: [`catlass.core_api.deinterleave`](../../catlass/core_api.py#L5887)	1486	+Source: [`catlass.core_api.deinterleave`](../../catlass/core_api.py#L5889)
1487		1487
1488	Description:	1488	Description:
1489		1489
	@@ -1520,7 +1520,7 @@ Compress valid lanes under a mask.
1520		1520
1521	#### `squeeze`	1521	#### `squeeze`
1522		1522
1523	-Source: [`catlass.core_api.squeeze`](../../catlass/core_api.py#L6322)	1523	+Source: [`catlass.core_api.squeeze`](../../catlass/core_api.py#L6324)
1524		1524
1525	Description:	1525	Description:
1526		1526
	@@ -1557,7 +1557,7 @@ In-core / cross-core flags, pipe barriers, mutexes, and local-memory barriers.
1557		1557
1558	### `flag`	1558	### `flag`
1559		1559
1560	-Source: [`catlass.core_api.flag`](../../catlass/core_api.py#L4431)	1560	+Source: [`catlass.core_api.flag`](../../catlass/core_api.py#L4433)
1561		1561
1562	Description:	1562	Description:
1563		1563
	@@ -1595,7 +1595,7 @@ with tla.vector():
1595		1595
1596	### `cross_flag`	1596	### `cross_flag`
1597		1597
1598	-Source: [`catlass.core_api.cross_flag`](../../catlass/core_api.py#L4484)	1598	+Source: [`catlass.core_api.cross_flag`](../../catlass/core_api.py#L4486)
1599		1599
1600	Description:	1600	Description:
1601		1601
	@@ -1627,7 +1627,7 @@ cf = tla.cross_flag("aic_aiv", mode=2)
1627		1627
1628	### `cross_core_set_flag`	1628	### `cross_core_set_flag`
1629		1629
1630	-Source: [`catlass.core_api.cross_core_set_flag`](../../catlass/core_api.py#L4556)	1630	+Source: [`catlass.core_api.cross_core_set_flag`](../../catlass/core_api.py#L4558)
1631		1631
1632	Description:	1632	Description:
1633		1633
	@@ -1662,7 +1662,7 @@ with tla.cube():
1662		1662
1663	### `cross_core_wait_flag`	1663	### `cross_core_wait_flag`
1664		1664
1665	-Source: [`catlass.core_api.cross_core_wait_flag`](../../catlass/core_api.py#L4602)	1665	+Source: [`catlass.core_api.cross_core_wait_flag`](../../catlass/core_api.py#L4604)
1666		1666
1667	Description:	1667	Description:
1668		1668
	@@ -1696,7 +1696,7 @@ with tla.vector():
1696		1696
1697	### `set_flag`	1697	### `set_flag`
1698		1698
1699	-Source: [`catlass.core_api.set_flag`](../../catlass/core_api.py#L4647)	1699	+Source: [`catlass.core_api.set_flag`](../../catlass/core_api.py#L4649)
1700		1700
1701	Description:	1701	Description:
1702		1702
	@@ -1728,7 +1728,7 @@ with tla.vector():
1728		1728
1729	### `wait_flag`	1729	### `wait_flag`
1730		1730
1731	-Source: [`catlass.core_api.wait_flag`](../../catlass/core_api.py#L4673)	1731	+Source: [`catlass.core_api.wait_flag`](../../catlass/core_api.py#L4675)
1732		1732
1733	Description:	1733	Description:
1734		1734
	@@ -1760,7 +1760,7 @@ with tla.vector():
1760		1760
1761	### `pipe_barrier`	1761	### `pipe_barrier`
1762		1762
1763	-Source: [`catlass.core_api.pipe_barrier`](../../catlass/core_api.py#L4699)	1763	+Source: [`catlass.core_api.pipe_barrier`](../../catlass/core_api.py#L4701)
1764		1764
1765	Description:	1765	Description:
1766		1766
	@@ -1792,7 +1792,7 @@ with tla.vector():
1792		1792
1793	### `mutex`	1793	### `mutex`
1794		1794
1795	-Source: [`catlass.core_api.mutex`](../../catlass/core_api.py#L4734)	1795	+Source: [`catlass.core_api.mutex`](../../catlass/core_api.py#L4736)
1796		1796
1797	Description:	1797	Description:
1798		1798
	@@ -1824,7 +1824,7 @@ mtx = tla.mutex("l1_buf", id=0)
1824		1824
1825	### `mutex_guard`	1825	### `mutex_guard`
1826		1826
1827	-Source: [`catlass.core_api.mutex_guard`](../../catlass/core_api.py#L4782)	1827	+Source: [`catlass.core_api.mutex_guard`](../../catlass/core_api.py#L4784)
1828		1828
1829	Description:	1829	Description:
1830		1830
	@@ -1856,7 +1856,7 @@ with tla.mutex_guard(mtx):
1856		1856
1857	### `mutex_lock`	1857	### `mutex_lock`
1858		1858
1859	-Source: [`catlass.core_api.mutex_lock`](../../catlass/core_api.py#L4823)	1859	+Source: [`catlass.core_api.mutex_lock`](../../catlass/core_api.py#L4825)
1860		1860
1861	Description:	1861	Description:
1862		1862
	@@ -1888,7 +1888,7 @@ tla.mutex_lock(mtx, pipe=tla.arch.MTE2)
1888		1888
1889	### `mutex_unlock`	1889	### `mutex_unlock`
1890		1890
1891	-Source: [`catlass.core_api.mutex_unlock`](../../catlass/core_api.py#L4853)	1891	+Source: [`catlass.core_api.mutex_unlock`](../../catlass/core_api.py#L4855)
1892		1892
1893	Description:	1893	Description:
1894		1894
	@@ -1920,7 +1920,7 @@ tla.mutex_unlock(mtx, pipe=tla.arch.MTE2)
1920		1920
1921	### `local_mem_bar`	1921	### `local_mem_bar`
1922		1922
1923	-Source: [`catlass.core_api.local_mem_bar`](../../catlass/core_api.py#L4882)	1923	+Source: [`catlass.core_api.local_mem_bar`](../../catlass/core_api.py#L4884)
1924		1924
1925	Description:	1925	Description:
1926		1926
	@@ -1957,7 +1957,7 @@ Architecture attributes on `tla.arch` (layout tags, pipe identifiers, block help
1957		1957
1958	### `arch`	1958	### `arch`
1959		1959
1960	-Source: [`catlass.core_api.arch`](../../catlass/core_api.py#L7147)	1960	+Source: [`catlass.core_api.arch`](../../catlass/core_api.py#L7149)
1961		1961
1962	Description:	1962	Description:
1963		1963
	@@ -2028,7 +2028,7 @@ On-chip scratch allocation via `allocate`.
2028		2028
2029	### `allocate`	2029	### `allocate`
2030		2030
2031	-Source: [`catlass.core_api.allocate`](../../catlass/core_api.py#L6915)	2031	+Source: [`catlass.core_api.allocate`](../../catlass/core_api.py#L6917)
2032		2032
2033	Description:	2033	Description:
2034		2034
	@@ -2109,7 +2109,7 @@ Cube / Vector / `vec.func` regions and kernel-side loop ranges.
2109		2109
2110	### `range`	2110	### `range`
2111		2111
2112	-Source: [`catlass.core_api.range`](../../catlass/core_api.py#L4932)	2112	+Source: [`catlass.core_api.range`](../../catlass/core_api.py#L4934)
2113		2113
2114	Description:	2114	Description:
2115		2115
	@@ -2143,7 +2143,7 @@ for i in tla.range(0, n, 1):
2143		2143
2144	### `range_constexpr`	2144	### `range_constexpr`
2145		2145
2146	-Source: [`catlass.core_api.range_constexpr`](../../catlass/core_api.py#L4982)	2146	+Source: [`catlass.core_api.range_constexpr`](../../catlass/core_api.py#L4984)
2147		2147
2148	Description:	2148	Description:
2149		2149
	@@ -2177,7 +2177,7 @@ for k in tla.range_constexpr(0, 4):
2177		2177
2178	### `cube`	2178	### `cube`
2179		2179
2180	-Source: [`catlass.core_api.cube`](../../catlass/core_api.py#L5028)	2180	+Source: [`catlass.core_api.cube`](../../catlass/core_api.py#L5030)
2181		2181
2182	Description:	2182	Description:
2183		2183
	@@ -2209,7 +2209,7 @@ with tla.cube():
2209		2209
2210	### `vector`	2210	### `vector`
2211		2211
2212	-Source: [`catlass.core_api.vector`](../../catlass/core_api.py#L5050)	2212	+Source: [`catlass.core_api.vector`](../../catlass/core_api.py#L5052)
2213		2213
2214	Description:	2214	Description:
2215		2215
	@@ -2241,7 +2241,7 @@ with tla.vector():
2241		2241
2242	### `vec.func`	2242	### `vec.func`
2243		2243
2244	-Source: [`catlass.core_api._vec_func`](../../catlass/core_api.py#L5084)	2244	+Source: [`catlass.core_api._vec_func`](../../catlass/core_api.py#L5086)
2245		2245
2246	Description:	2246	Description:
2247		2247

Apython/tla_dsl/examples/end_to_end/basic_mmad/basic_matmul_l0c2l1.py+206-0

@@ -0,0 +1,206 @@
	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	+from __future__ import annotations
	12	+
	13	+import argparse
	14	+
	15	+import catlass.tla as tla
	16	+
	17	+M_DIM = 50
	18	+N_DIM = 60
	19	+K_DIM = 64
	20	+
	21	+M_L1 = 64
	22	+N_L1 = 64
	23	+K_L1 = 64
	24	+
	25	+# First mmad: C = A(M,K) @ B(K,N)
	26	+L1A_SIZE = M_L1 * K_L1
	27	+L1B_SIZE = K_L1 * N_L1
	28	+L0A1_SIZE = M_L1 * K_L1
	29	+L0B1_SIZE = K_L1 * N_L1
	30	+
	31	+# L0C accumulator, reused by both mmads (holds C, then E)
	32	+L0C_SIZE = M_L1 * N_L1
	33	+
	34	+# Second mmad: E = C(M,N) @ D(N,N)
	35	+L1C_SIZE = M_L1 * N_L1
	36	+L1D_SIZE = N_L1 * N_L1
	37	+L0A2_SIZE = M_L1 * N_L1
	38	+L0B2_SIZE = N_L1 * N_L1
	39	+
	40	+L0A_SIZE = max(L0A1_SIZE, L0A2_SIZE)
	41	+L0B_SIZE = max(L0B1_SIZE, L0B2_SIZE)
	42	+
	43	+DESCRIPTION = "Basic MMAD L0C->L1 pipeline: E = (A@B)@D; f32, cube-only."
	44	+
	45	+
	46	+# ---------------------------------------------------------------------------
	47	+# Kernel
	48	+# ---------------------------------------------------------------------------
	49	+
	50	+@tla.kernel
	51	+def basic_matmul_l0c2l1(
	52	+ mem_a: tla.Tensor,
	53	+ mem_b: tla.Tensor,
	54	+ mem_d: tla.Tensor,
	55	+ mem_e: tla.Tensor,
	56	+) -> None:
	57	+ l1ab_done = tla.flag("l1ab_done", tla.arch.MTE2, tla.arch.MTE1)
	58	+ l0_loaded = tla.flag("l0_loaded", tla.arch.MTE1, tla.arch.CUBE)
	59	+ mmad_done = tla.flag("mmad_done", tla.arch.CUBE, tla.arch.FIX)
	60	+
	61	+ fixl1_done = tla.flag("fixl1_done", tla.arch.FIX, tla.arch.MTE1)
	62	+
	63	+ l1d_done = tla.flag("l1d_done", tla.arch.MTE2, tla.arch.MTE1)
	64	+
	65	+ l1a_ptr = tla.allocate(L1A_SIZE, tla.Float32, tla.AddressSpace.l1, 512)
	66	+ l1b_ptr = tla.allocate(L1B_SIZE, tla.Float32, tla.AddressSpace.l1, 512)
	67	+ l1c_ptr = tla.allocate(L1C_SIZE, tla.Float32, tla.AddressSpace.l1, 512)
	68	+ l1d_ptr = tla.allocate(L1D_SIZE, tla.Float32, tla.AddressSpace.l1, 512)
	69	+
	70	+ # the two-mmads shares the same l0a/l0b/l0c space
	71	+ l0a_ptr = tla.allocate(L0A_SIZE, tla.Float32, tla.AddressSpace.l0a, 512)
	72	+ l0b_ptr = tla.allocate(L0B_SIZE, tla.Float32, tla.AddressSpace.l0b, 512)
	73	+ l0c_ptr = tla.allocate(L0C_SIZE, tla.Float32, tla.AddressSpace.l0c, 512)
	74	+
	75	+ with tla.cube():
	76	+ # ---- C = A @ B -------------------------------------------------
	77	+ gm_a = tla.tile_view(mem_a, tla.make_shape(M_L1, K_L1), tla.make_coord(0, 0))
	78	+ gm_b = tla.tile_view(mem_b, tla.make_shape(K_L1, N_L1), tla.make_coord(0, 0))
	79	+ l1_a = tla.make_tensor_like(l1a_ptr, gm_a, tla.arch.zN)
	80	+ l1_b = tla.make_tensor_like(l1b_ptr, gm_b, tla.arch.zN)
	81	+
	82	+ tla.copy(l1_a, gm_a)
	83	+ tla.copy(l1_b, gm_b)
	84	+ tla.set_flag(l1ab_done)
	85	+ tla.wait_flag(l1ab_done)
	86	+
	87	+ l1_a_l0 = tla.tile_view(l1_a, tla.make_shape(M_L1, K_L1), tla.make_coord(0, 0))
	88	+ l1_b_l0 = tla.tile_view(l1_b, tla.make_shape(K_L1, N_L1), tla.make_coord(0, 0))
	89	+ l0a_a = tla.make_tensor_like(l0a_ptr, l1_a_l0, tla.arch.zN)
	90	+ l0b_b = tla.make_tensor_like(l0b_ptr, l1_b_l0, tla.arch.nZ)
	91	+
	92	+ tla.copy(l0a_a, l1_a_l0)
	93	+ tla.copy(l0b_b, l1_b_l0)
	94	+ tla.set_flag(l0_loaded)
	95	+ tla.wait_flag(l0_loaded)
	96	+
	97	+ gm_e = tla.tile_view(mem_e, tla.make_shape(M_L1, M_L1), tla.make_coord(0, 0))
	98	+ l0c_c = tla.make_tensor_like(l0c_ptr, gm_e, tla.arch.L0Clayout)
	99	+ tla.mmad(l0c_c, l0a_a, l0b_b, init_c=True)
	100	+ tla.set_flag(mmad_done)
	101	+ tla.wait_flag(mmad_done)
	102	+
	103	+ # copy l0c->l1c, L0Clayout->zN
	104	+ l1_c = tla.make_tensor_like(l1c_ptr, l0c_c, tla.arch.zN)
	105	+ tla.copy(l1_c, l0c_c)
	106	+ tla.set_flag(fixl1_done)
	107	+ tla.wait_flag(fixl1_done)
	108	+
	109	+ # ---- E = C @ D -------------------------------------------------
	110	+ gm_d = tla.tile_view(mem_d, tla.make_shape(N_L1, N_L1), tla.make_coord(0, 0))
	111	+ l1_d = tla.make_tensor_like(l1d_ptr, gm_d, tla.arch.zN)
	112	+ tla.copy(l1_d, gm_d)
	113	+ tla.set_flag(l1d_done)
	114	+ tla.wait_flag(l1d_done)
	115	+
	116	+ l1_c_l0 = tla.tile_view(l1_c, tla.make_shape(M_L1, N_L1), tla.make_coord(0, 0))
	117	+ l1_d_l0 = tla.tile_view(l1_d, tla.make_shape(N_L1, N_L1), tla.make_coord(0, 0))
	118	+ l0a_c = tla.make_tensor_like(l0a_ptr, l1_c_l0, tla.arch.zN)
	119	+ l0b_d = tla.make_tensor_like(l0b_ptr, l1_d_l0, tla.arch.nZ)
	120	+ tla.copy(l0a_c, l1_c)
	121	+ tla.copy(l0b_d, l1_d)
	122	+ tla.set_flag(l0_loaded)
	123	+ tla.wait_flag(l0_loaded)
	124	+
	125	+ l0c_e = tla.make_tensor_like(l0c_ptr, l0c_c, tla.arch.L0Clayout) # two-mmad out shapes are same
	126	+ tla.mmad(l0c_e, l0a_c, l0b_d, init_c=True)
	127	+ tla.set_flag(mmad_done)
	128	+ tla.wait_flag(mmad_done)
	129	+
	130	+ tla.copy(gm_e, l0c_e)
	131	+
	132	+ tla.pipe_barrier(tla.pipes.ALL)
	133	+
	134	+
	135	+# ---------------------------------------------------------------------------
	136	+# Host
	137	+# ---------------------------------------------------------------------------
	138	+
	139	+
	140	+def golden(a, b, d):
	141	+ import torch
	142	+
	143	+ a = a.to(torch.float32)
	144	+ b = b.to(torch.float32)
	145	+ d = d.to(torch.float32)
	146	+ return a @ b @ d
	147	+
	148	+
	149	+def prepare_npu(buf, layout: str):
	150	+ storage = buf.contiguous() if layout == "row" else buf.permute(1, 0).contiguous()
	151	+ return storage.npu()
	152	+
	153	+
	154	+def run(args: argparse.Namespace) -> int:
	155	+ import torch
	156	+ import torch_npu
	157	+
	158	+ from examples.end_to_end.common import (
	159	+ create_tla_tensor,
	160	+ compare,
	161	+ )
	162	+
	163	+ torch.npu.set_device(args.device)
	164	+ # e = a @ b @ d
	165	+ a = torch.rand((M_DIM, K_DIM), dtype=torch.float32, device="cpu") * 10.0 - 5.0
	166	+ b = torch.rand((K_DIM, N_DIM), dtype=torch.float32, device="cpu") * 10.0 - 5.0
	167	+ d = torch.rand((N_DIM, N_DIM), dtype=torch.float32, device="cpu") * 10.0 - 5.0
	168	+ e = torch.full((M_DIM, N_DIM), args.sentinel, dtype=torch.float32, device="cpu")
	169	+ ref = golden(a, b, d)
	170	+
	171	+ a = prepare_npu(a, "row")
	172	+ b = prepare_npu(b, "row")
	173	+ d = prepare_npu(d, "row")
	174	+ e = prepare_npu(e, "row")
	175	+ a_tensor = create_tla_tensor(a, "row")
	176	+ b_tensor = create_tla_tensor(b, "row")
	177	+ d_tensor = create_tla_tensor(d, "row")
	178	+ e_tensor = create_tla_tensor(e, "row")
	179	+
	180	+ artifact = tla.compile(
	181	+ basic_matmul_l0c2l1,
	182	+ a_tensor,
	183	+ b_tensor,
	184	+ d_tensor,
	185	+ e_tensor,
	186	+ options="--npu-arch 3510"
	187	+ )
	188	+ artifact(a_tensor, b_tensor, d_tensor, e_tensor, block_num=args.block_num)
	189	+ torch.npu.synchronize()
	190	+
	191	+ passed = compare(e.detach().cpu(), ref, max(K_DIM, N_DIM))
	192	+ print(f"passed={passed} cache_key={artifact.cache_key}")
	193	+ print(f"kernel.o={artifact.kernel_binary_path}")
	194	+ return 0 if passed else 1
	195	+
	196	+
	197	+def main() -> int:
	198	+ parser = argparse.ArgumentParser(description=DESCRIPTION)
	199	+ parser.add_argument("--device", type=int, default=0)
	200	+ parser.add_argument("--block-num", type=int, default=1)
	201	+ parser.add_argument("--sentinel", type=float, default=-9.0)
	202	+ return run(parser.parse_args())
	203	+
	204	+
	205	+if __name__ == "__main__":
	206	+ raise SystemExit(main())

Mtests/dsl_battery/test_battery.py+4-0

	@@ -134,6 +134,10 @@ def _cases(device: int) -> Iterator[tuple[str, list[list[str]]]]:
134		134
135	yield ("mmad-ptr", [["basic_mmad/basic_mmad_ptr.py", *dev]])	135	yield ("mmad-ptr", [["basic_mmad/basic_mmad_ptr.py", *dev]])
136		136
		137	+ # basic_mmad_l0c2l1: cube-only E=(A@B)@D, L0C->L1 (fixpipe) staging reused
		138	+ # as the second mmad's LHS; single static tile (M=50/N=60/K=64).
		139	+ yield ("mmad-l0c2l1", [["basic_mmad/basic_matmul_l0c2l1.py", *dev]])
		140	+
137	# --- basic_vadd: every dtype under each sync mode ---	141	# --- basic_vadd: every dtype under each sync mode ---
138	for mode in VADD_MODES:	142	for mode in VADD_MODES:
139	for dtype in VADD_DTYPES:	143	for dtype in VADD_DTYPES:

Mtests/run_dsl_test.sh+8-1

	@@ -9,7 +9,8 @@
9	# See LICENSE in the root of the software repository for the full text of the License.	9	# See LICENSE in the root of the software repository for the full text of the License.
10	# -----------------------------------------------------------------------------------------------------------	10	# -----------------------------------------------------------------------------------------------------------
11	#	11	#
12	-# End-to-end validation for python/tla_dsl/examples/end_to_end/basic_mmad (basic_matmul*.py, basic_mmad_ptr.py),	12	+# End-to-end validation for python/tla_dsl/examples/end_to_end/basic_mmad (basic_matmul*.py, basic_mmad_ptr.py,
		13	+# basic_matmul_l0c2l1.py),
13	# python/tla_dsl/examples/end_to_end/basic_vadd (basic_vadd.py),	14	# python/tla_dsl/examples/end_to_end/basic_vadd (basic_vadd.py),
14	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed.py), python/tla_dsl/examples/end_to_end/basic_mixed_mutex (basic_mixed_mutex.py) and	15	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed.py), python/tla_dsl/examples/end_to_end/basic_mixed_mutex (basic_mixed_mutex.py) and
15	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed_ub2l1.py, basic_mixed_store_zN.py,	16	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed_ub2l1.py, basic_mixed_store_zN.py,
	@@ -96,6 +97,7 @@ export CATLASS_DSL_FORCE_RECOMPILE="${CATLASS_DSL_FORCE_RECOMPILE:-1}"
96	BASIC_MMAD_REL="examples/end_to_end/basic_mmad/basic_matmul.py"	97	BASIC_MMAD_REL="examples/end_to_end/basic_mmad/basic_matmul.py"
97	BASIC_MMAD_AUTO_SYNC_REL="examples/end_to_end/basic_mmad/basic_matmul_auto_sync.py"	98	BASIC_MMAD_AUTO_SYNC_REL="examples/end_to_end/basic_mmad/basic_matmul_auto_sync.py"
98	BASIC_MMAD_PTR_REL="examples/end_to_end/basic_mmad/basic_mmad_ptr.py"	99	BASIC_MMAD_PTR_REL="examples/end_to_end/basic_mmad/basic_mmad_ptr.py"
		100	+BASIC_MMAD_L0C2L1_REL="examples/end_to_end/basic_mmad/basic_matmul_l0c2l1.py"
99	BASIC_VADD_REL="examples/end_to_end/basic_vadd/basic_vadd.py"	101	BASIC_VADD_REL="examples/end_to_end/basic_vadd/basic_vadd.py"
100	BASIC_MIXED_REL="examples/end_to_end/basic_mixed/basic_mixed.py"	102	BASIC_MIXED_REL="examples/end_to_end/basic_mixed/basic_mixed.py"
101	BASIC_MIXED_MUTEX_REL="examples/end_to_end/basic_mixed/basic_mixed_mutex.py"	103	BASIC_MIXED_MUTEX_REL="examples/end_to_end/basic_mixed/basic_mixed_mutex.py"
	@@ -164,6 +166,7 @@ Run end-to-end validation for:
164	- basic_mmad (full flag-sync mnk/layout/dtype matrix; representative manual and automatic mutex cases;	166	- basic_mmad (full flag-sync mnk/layout/dtype matrix; representative manual and automatic mutex cases;
165	atomic-add coverage for each supported input dtype)	167	atomic-add coverage for each supported input dtype)
166	- basic_mmad_ptr (basic_mmad_ptr.py)	168	- basic_mmad_ptr (basic_mmad_ptr.py)
		169	+ - basic_mmad_l0c2l1 (basic_matmul_l0c2l1.py: cube-only E=(A@B)@D with L0C->L1 staging)
167	- basic_vadd (basic_vadd.py with per-dtype CLI invocations, plus mutex variants)	170	- basic_vadd (basic_vadd.py with per-dtype CLI invocations, plus mutex variants)
168	- basic_mixed (basic_mixed.py with dynamic GM mnk list, including --use-mutex; basic_mixed_ub2l1.py,	171	- basic_mixed (basic_mixed.py with dynamic GM mnk list, including --use-mutex; basic_mixed_ub2l1.py,
169	basic_mixed_store_zN.py, basic_mixed_store_zNUnAlign.py for m=64/m=50)	172	basic_mixed_store_zN.py, basic_mixed_store_zNUnAlign.py for m=64/m=50)
	@@ -401,6 +404,10 @@ if [[ ! -f "${CATLASS_DSL_DIR}/${BASIC_MMAD_PTR_REL}" ]]; then
401	echo "error: missing ${BASIC_MMAD_PTR_REL} under ${CATLASS_DSL_DIR}" >&2	404	echo "error: missing ${BASIC_MMAD_PTR_REL} under ${CATLASS_DSL_DIR}" >&2
402	exit 1	405	exit 1
403	fi	406	fi
		407	+if [[ ! -f "${CATLASS_DSL_DIR}/${BASIC_MMAD_L0C2L1_REL}" ]]; then
		408	+ echo "error: missing ${BASIC_MMAD_L0C2L1_REL} under ${CATLASS_DSL_DIR}" >&2
		409	+ exit 1
		410	+fi
404	if [[ ! -f "${CATLASS_DSL_DIR}/${BASIC_VADD_REL}" ]]; then	411	if [[ ! -f "${CATLASS_DSL_DIR}/${BASIC_VADD_REL}" ]]; then
405	echo "error: missing ${BASIC_VADD_REL} under ${CATLASS_DSL_DIR}" >&2	412	echo "error: missing ${BASIC_VADD_REL} under ${CATLASS_DSL_DIR}" >&2
406	exit 1	413	exit 1