a2a3 mgather & mscatter kernels-pto-isa-AtomGit

已合并

a2a3 mgather & mscatter kernels #935

Sathi Sarveswara Reddy创建于 5月17日从ssreddy/pto-isa-a2a3-mgather-mscatter-may15： master 合入到 cann/pto-isa： master

a2a3 mgather & mscatter kernels #935

已合并

Sathi Sarveswara Reddy创建于 5月17日

从ssreddy/pto-isa-a2a3-mgather-mscatter-may15 master合入到cann/pto-isa master

共 16 个文件变更+4368-0

Minclude/pto/common/pto_instr.hpp+47-0

	@@ -1822,6 +1822,25 @@ PTO_INST RecordEvent MGATHER(TileDst &dst, GlobalData &src, TileInd &indexes, Wa
1822	}	1822	}
1823	#endif	1823	#endif
1824		1824
		1825	+#ifdef PTO_NPU_ARCH_A2A3
		1826	+template <Coalesce CMode, typename TileDst, typename GlobalData, typename TileInd, typename... WaitEvents>
		1827	+PTO_INST RecordEvent MGATHER(TileDst &dst, GlobalData &src, TileInd &indexes, WaitEvents &...events)
		1828	+{
		1829	+ TSYNC(events...);
		1830	+ MGATHER_IMPL<CMode>(dst, src, indexes);
		1831	+ return {};
		1832	+}
		1833	+
		1834	+template <Coalesce CMode, GatherOOB Mode, typename TileDst, typename GlobalData, typename TileInd,
		1835	+ typename... WaitEvents>
		1836	+PTO_INST RecordEvent MGATHER(TileDst &dst, GlobalData &src, TileInd &indexes, WaitEvents &...events)
		1837	+{
		1838	+ TSYNC(events...);
		1839	+ MGATHER_IMPL<CMode, Mode>(dst, src, indexes);
		1840	+ return {};
		1841	+}
		1842	+#endif
		1843	+
1825	template <typename GlobalData, typename TileSrc, typename TileInd, typename... WaitEvents>	1844	template <typename GlobalData, typename TileSrc, typename TileInd, typename... WaitEvents>
1826	PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, WaitEvents &...events)	1845	PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, WaitEvents &...events)
1827	{	1846	{
	@@ -1867,6 +1886,34 @@ PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, W
1867	}	1886	}
1868	#endif	1887	#endif
1869		1888
		1889	+#ifdef PTO_NPU_ARCH_A2A3
		1890	+template <Coalesce Mode, typename GlobalData, typename TileSrc, typename TileInd, typename... WaitEvents>
		1891	+PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, WaitEvents &...events)
		1892	+{
		1893	+ TSYNC(events...);
		1894	+ MSCATTER_IMPL<Mode>(dst, src, indexes);
		1895	+ return {};
		1896	+}
		1897	+
		1898	+template <Coalesce Mode, ScatterAtomicOp Atomic, typename GlobalData, typename TileSrc, typename TileInd,
		1899	+ typename... WaitEvents>
		1900	+PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, WaitEvents &...events)
		1901	+{
		1902	+ TSYNC(events...);
		1903	+ MSCATTER_IMPL<Mode, Atomic>(dst, src, indexes);
		1904	+ return {};
		1905	+}
		1906	+
		1907	+template <Coalesce Mode, ScatterAtomicOp Atomic, ScatterOOB Oob, typename GlobalData, typename TileSrc,
		1908	+ typename TileInd, typename... WaitEvents>
		1909	+PTO_INST RecordEvent MSCATTER(GlobalData &dst, TileSrc &src, TileInd &indexes, WaitEvents &...events)
		1910	+{
		1911	+ TSYNC(events...);
		1912	+ MSCATTER_IMPL<Mode, Atomic, Oob>(dst, src, indexes);
		1913	+ return {};
		1914	+}
		1915	+#endif
		1916	+
1870	template <typename TileDataDst, typename TileDataSrc, typename... WaitEvents>	1917	template <typename TileDataDst, typename TileDataSrc, typename... WaitEvents>
1871	PTO_INST RecordEvent TNEG(TileDataDst &dst, TileDataSrc &src, WaitEvents &...events)	1918	PTO_INST RecordEvent TNEG(TileDataDst &dst, TileDataSrc &src, WaitEvents &...events)
1872	{	1919	{

Minclude/pto/common/pto_instr_impl.hpp+2-0

	@@ -143,6 +143,8 @@ See LICENSE in the root of the software repository for the full text of the Lice
143	#include "pto/npu/a2a3/TPrefetch.hpp"	143	#include "pto/npu/a2a3/TPrefetch.hpp"
144	#include "pto/npu/a2a3/TPrelu.hpp"	144	#include "pto/npu/a2a3/TPrelu.hpp"
145	#include "pto/npu/a2a3/TInsert.hpp"	145	#include "pto/npu/a2a3/TInsert.hpp"
		146	+#include "pto/npu/a2a3/MGather.hpp"
		147	+#include "pto/npu/a2a3/MScatter.hpp"
146	#include "pto/npu/a2a3/TRowExpandExpdif.hpp"	148	#include "pto/npu/a2a3/TRowExpandExpdif.hpp"
147	#include "pto/npu/a2a3/TColExpandAdd.hpp"	149	#include "pto/npu/a2a3/TColExpandAdd.hpp"
148	#include "pto/npu/a2a3/TColExpandMax.hpp"	150	#include "pto/npu/a2a3/TColExpandMax.hpp"

Ainclude/pto/npu/a2a3/MGather.hpp+431-0

@@ -0,0 +1,431 @@
	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 MGATHER_HPP
	12	+#define MGATHER_HPP
	13	+
	14	+#include <pto/common/utils.hpp>
	15	+#include <pto/common/constants.hpp>
	16	+#include "common.hpp"
	17	+
	18	+namespace pto {
	19	+
	20	+enum class GatherOOB : uint8_t
	21	+{
	22	+ Undefined = 0,
	23	+ Clamp = 1,
	24	+ Wrap = 2,
	25	+ Zero = 3
	26	+};
	27	+
	28	+#ifndef PTO_COALESCE_ENUM_DEFINED
	29	+#define PTO_COALESCE_ENUM_DEFINED
	30	+enum class Coalesce : uint8_t
	31	+{
	32	+ Row = 0,
	33	+ Elem = 1
	34	+};
	35	+#endif
	36	+
	37	+template <typename T>
	38	+struct IsValidMGatherDType {
	39	+ static constexpr bool value = std::is_same_v<T, int8_t> \|\| std::is_same_v<T, uint8_t> \|\|
	40	+ std::is_same_v<T, int16_t> \|\| std::is_same_v<T, uint16_t> \|\|
	41	+ std::is_same_v<T, int32_t> \|\| std::is_same_v<T, uint32_t> \|\|
	42	+ std::is_same_v<T, half> \|\| std::is_same_v<T, bfloat16_t> \|\| std::is_same_v<T, float>;
	43	+};
	44	+
	45	+template <typename Tile>
	46	+struct IsMGatherNDTile {
	47	+ static constexpr bool value = Tile::isRowMajor && (Tile::SFractal == SLayout::NoneBox);
	48	+};
	49	+
	50	+template <typename Tile>
	51	+struct IsMGatherNZTile {
	52	+ static constexpr bool value =
	53	+ !Tile::isRowMajor && (Tile::SFractal == SLayout::RowMajor) && (Tile::SFractalSize == TileConfig::fractalABSize);
	54	+};
	55	+
	56	+template <GatherOOB Oob>
	57	+AICORE PTO_INLINE uint32_t mgather_remap(uint32_t idx, uint32_t cap, uint32_t &doRead)
	58	+{
	59	+ if constexpr (Oob == GatherOOB::Undefined) {
	60	+ doRead = 1u;
	61	+ return idx;
	62	+ } else if constexpr (Oob == GatherOOB::Clamp) {
	63	+ doRead = 1u;
	64	+ return (idx >= cap) ? (cap - 1u) : idx;
	65	+ } else if constexpr (Oob == GatherOOB::Wrap) {
	66	+ doRead = 1u;
	67	+ return idx % cap;
	68	+ } else {
	69	+ doRead = (idx < cap) ? 1u : 0u;
	70	+ return idx;
	71	+ }
	72	+}
	73	+
	74	+template <typename T>
	75	+AICORE PTO_INLINE void MGatherRowDma(__ubuf__ T dst, __gm__ T src, uint32_t lenBytes)
	76	+{
	77	+ if constexpr (sizeof(T) == 1) {
	78	+ copy_gm_to_ubuf_align_b8(dst, src, 0, 1, lenBytes, 0, 0, 0, 0);
	79	+ } else if constexpr (sizeof(T) == 2) {
	80	+ copy_gm_to_ubuf_align_b16(dst, src, 0, 1, lenBytes, 0, 0, 0, 0);
	81	+ } else if constexpr (sizeof(T) == 4) {
	82	+ copy_gm_to_ubuf_align_b32(dst, src, 0, 1, lenBytes, 0, 0, 0, 0);
	83	+ }
	84	+}
	85	+
	86	+template <typename T>
	87	+AICORE PTO_INLINE void MGatherRowMultiDma(__ubuf__ T dst, __gm__ T src, uint16_t nBurst, uint32_t lenBytes,
	88	+ uint32_t gmGapBytes, uint32_t ubGapBlocks)
	89	+{
	90	+ if constexpr (sizeof(T) == 1) {
	91	+ copy_gm_to_ubuf_align_b8(dst, src, 0, nBurst, lenBytes, 0, 0, gmGapBytes, ubGapBlocks);
	92	+ } else if constexpr (sizeof(T) == 2) {
	93	+ copy_gm_to_ubuf_align_b16(dst, src, 0, nBurst, lenBytes, 0, 0, gmGapBytes, ubGapBlocks);
	94	+ } else if constexpr (sizeof(T) == 4) {
	95	+ copy_gm_to_ubuf_align_b32(dst, src, 0, nBurst, lenBytes, 0, 0, gmGapBytes, ubGapBlocks);
	96	+ }
	97	+}
	98	+
	99	+template <typename T>
	100	+AICORE PTO_INLINE uint64_t MGatherNZGmOffset(uint32_t logicalRow, uint32_t logicalCol, int gShape0, int gShape1,
	101	+ int gStride0, int gStride1, int gStride2, int gStride3, int gStride4)
	102	+{
	103	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	104	+ constexpr uint32_t kFRow = FRACTAL_NZ_ROW;
	105	+ const uint32_t blockColCombined = logicalCol / kC0;
	106	+ const uint32_t colInBlock = logicalCol - blockColCombined * kC0;
	107	+ const uint32_t blockRow = logicalRow / kFRow;
	108	+ const uint32_t rowInBlock = logicalRow - blockRow * kFRow;
	109	+ const uint32_t blockColOuter0 = (gShape0 == 1) ? 0u : (blockColCombined / (uint32_t)gShape1);
	110	+ const uint32_t blockColOuter1 = (gShape0 == 1) ? blockColCombined : (blockColCombined - blockColOuter0 * gShape1);
	111	+ return (uint64_t)blockColOuter0 * (uint64_t)gStride0 + (uint64_t)blockColOuter1 * (uint64_t)gStride1 +
	112	+ (uint64_t)blockRow * (uint64_t)gStride2 + (uint64_t)rowInBlock * (uint64_t)gStride3 +
	113	+ (uint64_t)colInBlock * (uint64_t)gStride4;
	114	+}
	115	+
	116	+template <GatherOOB Oob, typename T, typename TIdx, typename DstTile, typename IdxTile>
	117	+__tf__ AICORE void MGatherRowImpl(typename DstTile::TileDType __out__ dst, __gm__ T *tablePtr,
	118	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, uint32_t validCol,
	119	+ uint32_t tableRows, uint32_t tableRowStride)
	120	+{
	121	+ __ubuf__ T dstPtr = (__ubuf__ T )__cce_get_tile_ptr(dst);
	122	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	123	+
	124	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	125	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	126	+
	127	+ const uint32_t lenBytes = validCol * sizeof(T);
	128	+ constexpr uint32_t kRowStride = DstTile::RowStride;
	129	+
	130	+ if constexpr (Oob == GatherOOB::Zero) {
	131	+ for (uint32_t r = 0; r < validRow; r++) {
	132	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[r]);
	133	+ uint32_t doRead;
	134	+ uint32_t safeIdx = mgather_remap<Oob>(rawIdx, tableRows, doRead);
	135	+ __ubuf__ T dstRow = dstPtr + r kRowStride;
	136	+ if (doRead) {
	137	+ __gm__ T srcRow = tablePtr + static_cast<uint64_t>(safeIdx) tableRowStride;
	138	+ MGatherRowDma<T>(dstRow, srcRow, lenBytes);
	139	+ } else {
	140	+ for (uint32_t c = 0; c < validCol; c++) {
	141	+ dstRow[c] = static_cast<T>(0);
	142	+ }
	143	+ }
	144	+ }
	145	+ } else {
	146	+ for (uint32_t r = 0; r < validRow; r++) {
	147	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[r]);
	148	+ uint32_t doRead;
	149	+ uint32_t safeIdx = mgather_remap<Oob>(rawIdx, tableRows, doRead);
	150	+ if (doRead) {
	151	+ __gm__ T srcRow = tablePtr + static_cast<uint64_t>(safeIdx) tableRowStride;
	152	+ __ubuf__ T dstRow = dstPtr + r kRowStride;
	153	+ MGatherRowDma<T>(dstRow, srcRow, lenBytes);
	154	+ }
	155	+ }
	156	+ }
	157	+
	158	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	159	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_V>();
	160	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_MTE3>();
	161	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	162	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	163	+}
	164	+
	165	+template <GatherOOB Oob, typename T, typename TIdx, typename DstTile, typename IdxTile>
	166	+__tf__ AICORE void MGatherRowNzImpl(typename DstTile::TileDType __out__ dst, __gm__ T *tablePtr,
	167	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, int gShape0,
	168	+ int gShape1, int gShape2, int gStride0, int gStride1, int gStride2, int gStride3)
	169	+{
	170	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	171	+ constexpr uint32_t kFRow = FRACTAL_NZ_ROW;
	172	+ constexpr uint32_t kFractalRowBytes = kC0 * sizeof(T);
	173	+
	174	+ __ubuf__ T dstPtr = (__ubuf__ T )__cce_get_tile_ptr(dst);
	175	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	176	+
	177	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	178	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	179	+
	180	+ const uint32_t tableLogicalRows = (uint32_t)gShape2 * kFRow;
	181	+ const uint32_t gmGapBytes = ((uint32_t)gStride1 - kC0) * (uint32_t)sizeof(T);
	182	+ constexpr uint32_t ubGapBlocks = (uint32_t)DstTile::Rows - 1u;
	183	+ const int64_t tileOuterStrideElem = (int64_t)gShape1 * (int64_t)DstTile::Rows * (int64_t)kC0;
	184	+
	185	+ if constexpr (Oob == GatherOOB::Zero) {
	186	+ constexpr uint32_t kDstNumel = (uint32_t)DstTile::Rows * (uint32_t)DstTile::Cols;
	187	+ for (uint32_t i = 0; i < kDstNumel; i++) {
	188	+ dstPtr[i] = static_cast<T>(0);
	189	+ }
	190	+ }
	191	+
	192	+ for (uint32_t r = 0; r < validRow; r++) {
	193	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[r]);
	194	+ uint32_t doRead;
	195	+ uint32_t safeIdx = mgather_remap<Oob>(rawIdx, tableLogicalRows, doRead);
	196	+ if (doRead) {
	197	+ const uint32_t srcBlockRow = safeIdx / kFRow;
	198	+ const uint32_t srcRowInBlock = safeIdx - srcBlockRow * kFRow;
	199	+ const uint32_t dstBlockRow = r / kFRow;
	200	+ const uint32_t dstRowInBlock = r - dstBlockRow * kFRow;
	201	+
	202	+ for (uint32_t i = 0; i < (uint32_t)gShape0; i++) {
	203	+ __gm__ T srcAddr = tablePtr + (int64_t)i (int64_t)gStride0 +
	204	+ (int64_t)srcBlockRow * (int64_t)gStride2 +
	205	+ (int64_t)srcRowInBlock * (int64_t)gStride3;
	206	+ __ubuf__ T dstAddr = dstPtr + (int64_t)i tileOuterStrideElem +
	207	+ (int64_t)dstBlockRow * (int64_t)kFRow * (int64_t)kC0 +
	208	+ (int64_t)dstRowInBlock * (int64_t)kC0;
	209	+ MGatherRowMultiDma<T>(dstAddr, srcAddr, (uint16_t)gShape1, kFractalRowBytes, gmGapBytes, ubGapBlocks);
	210	+ }
	211	+ }
	212	+ }
	213	+
	214	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	215	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_V>();
	216	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_MTE3>();
	217	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	218	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	219	+}
	220	+
	221	+template <GatherOOB Oob, typename T, typename TIdx, typename DstTile, typename IdxTile>
	222	+__tf__ AICORE void MGatherElemImpl(typename DstTile::TileDType __out__ dst, __gm__ T *tablePtr,
	223	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, uint32_t validCol,
	224	+ uint32_t tableSize)
	225	+{
	226	+ __ubuf__ T dstPtr = (__ubuf__ T )__cce_get_tile_ptr(dst);
	227	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	228	+
	229	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	230	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	231	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_S>();
	232	+
	233	+ constexpr uint32_t kDstRowStride = DstTile::RowStride;
	234	+ constexpr uint32_t kIdxRowStride = IdxTile::RowStride;
	235	+
	236	+ for (uint32_t r = 0; r < validRow; r++) {
	237	+ const uint32_t idxRowOff = r * kIdxRowStride;
	238	+ const uint32_t dstRowOff = r * kDstRowStride;
	239	+ for (uint32_t c = 0; c < validCol; c++) {
	240	+ const uint32_t idxOff = idxRowOff + c;
	241	+ const uint32_t dstOff = dstRowOff + c;
	242	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[idxOff]);
	243	+ uint32_t doRead;
	244	+ uint32_t safeIdx = mgather_remap<Oob>(rawIdx, tableSize, doRead);
	245	+ if (doRead) {
	246	+ dstPtr[dstOff] = tablePtr[safeIdx];
	247	+ } else if constexpr (Oob == GatherOOB::Zero) {
	248	+ dstPtr[dstOff] = static_cast<T>(0);
	249	+ }
	250	+ }
	251	+ }
	252	+
	253	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	254	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	255	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	256	+}
	257	+
	258	+template <GatherOOB Oob, typename T, typename TIdx, typename DstTile, typename IdxTile>
	259	+__tf__ AICORE void MGatherElemNzImpl(typename DstTile::TileDType __out__ dst, __gm__ T *tablePtr,
	260	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, uint32_t validCol,
	261	+ uint32_t tableSize, int gShape0, int gShape1, int gStride0, int gStride1,
	262	+ int gStride2, int gStride3, int gStride4, uint32_t nLogicalCols)
	263	+{
	264	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	265	+ __ubuf__ T dstPtr = (__ubuf__ T )__cce_get_tile_ptr(dst);
	266	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	267	+
	268	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	269	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	270	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_S>();
	271	+
	272	+ constexpr uint32_t kIdxRowStride = IdxTile::RowStride;
	273	+ const uint32_t nColBlocks = (validCol + kC0 - 1u) / kC0;
	274	+ const uint32_t kDstColBlockStride = (uint32_t)DstTile::Rows * kC0;
	275	+
	276	+ for (uint32_t bcol = 0; bcol < nColBlocks; bcol++) {
	277	+ const uint32_t cBase = bcol * kC0;
	278	+ const uint32_t cLimit = (cBase + kC0 < validCol) ? (cBase + kC0) : validCol;
	279	+ const uint32_t kInBlock = cLimit - cBase;
	280	+ __ubuf__ T dstBlockBase = dstPtr + (uint64_t)bcol (uint64_t)kDstColBlockStride;
	281	+ for (uint32_t r = 0; r < validRow; r++) {
	282	+ const uint32_t idxRowOff = r * kIdxRowStride;
	283	+ __ubuf__ T dstRowBase = dstBlockBase + (uint64_t)r (uint64_t)kC0;
	284	+ for (uint32_t cInner = 0; cInner < kInBlock; cInner++) {
	285	+ const uint32_t c = cBase + cInner;
	286	+ const uint32_t idxOff = idxRowOff + c;
	287	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[idxOff]);
	288	+ uint32_t doRead;
	289	+ uint32_t safeIdx = mgather_remap<Oob>(rawIdx, tableSize, doRead);
	290	+ if (doRead) {
	291	+ const uint32_t logicalRow = safeIdx / nLogicalCols;
	292	+ const uint32_t logicalCol = safeIdx - logicalRow * nLogicalCols;
	293	+ const uint64_t srcOff = MGatherNZGmOffset<T>(logicalRow, logicalCol, gShape0, gShape1, gStride0,
	294	+ gStride1, gStride2, gStride3, gStride4);
	295	+ dstRowBase[cInner] = tablePtr[srcOff];
	296	+ } else if constexpr (Oob == GatherOOB::Zero) {
	297	+ dstRowBase[cInner] = static_cast<T>(0);
	298	+ }
	299	+ }
	300	+ }
	301	+ }
	302	+
	303	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	304	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	305	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	306	+}
	307	+
	308	+template <Coalesce Mode, GatherOOB Oob, typename DstTile, typename GlobalTable, typename IdxTile>
	309	+PTO_INTERNAL void MGatherCheck()
	310	+{
	311	+ using T = typename DstTile::DType;
	312	+ using TIdx = typename IdxTile::DType;
	313	+
	314	+ static_assert(IsValidMGatherDType<T>::value,
	315	+ "MGATHER A2/A3 data type must be int8/uint8/int16/uint16/int32/uint32/half/bfloat16/float.");
	316	+ static_assert(std::is_same_v<TIdx, int32_t> \|\| std::is_same_v<TIdx, uint32_t>,
	317	+ "MGATHER A2/A3 index type must be int32_t or uint32_t.");
	318	+ static_assert(std::is_same_v<typename GlobalTable::DType, __gm__ T>,
	319	+ "MGATHER A2/A3 source table must be a GM GlobalTensor with element type matching the destination.");
	320	+ static_assert(DstTile::Loc == TileType::Vec, "MGATHER A2/A3 destination must be a Vec tile (UB).");
	321	+ static_assert(IdxTile::Loc == TileType::Vec, "MGATHER A2/A3 indices must be a Vec tile (UB).");
	322	+
	323	+ static_assert(IdxTile::isRowMajor, "MGATHER A2/A3 index tile must be BLayout::RowMajor.");
	324	+ static_assert(IdxTile::SFractal == SLayout::NoneBox, "MGATHER A2/A3 index tile must be ND (SLayout::NoneBox).");
	325	+
	326	+ constexpr bool kIsTableND = (GlobalTable::layout == Layout::ND);
	327	+ constexpr bool kIsTableNZ = (GlobalTable::layout == Layout::NZ);
	328	+ constexpr bool kIsDstND = IsMGatherNDTile<DstTile>::value;
	329	+ constexpr bool kIsDstNZ = IsMGatherNZTile<DstTile>::value;
	330	+
	331	+ static_assert(kIsTableND \|\| kIsTableNZ, "MGATHER A2/A3 table must use Layout::ND or Layout::NZ.");
	332	+ static_assert((kIsTableND && kIsDstND) \|\| (kIsTableNZ && kIsDstNZ),
	333	+ "MGATHER A2/A3 layout pairing must be either:\n"
	334	+ " (a) GM Layout::ND + UB tile (BLayout::RowMajor + SLayout::NoneBox), or\n"
	335	+ " (b) GM Layout::NZ + UB tile (BLayout::ColMajor + SLayout::RowMajor + SFractalSize=512).");
	336	+
	337	+ static_assert(DstTile::Cols * sizeof(T) % BLOCK_BYTE_SIZE == 0,
	338	+ "MGATHER A2/A3 destination tile padded Cols*sizeof(T) must be 32-byte aligned.");
	339	+
	340	+ if constexpr (kIsTableNZ) {
	341	+ static_assert(GlobalTable::staticShape[3] == FRACTAL_NZ_ROW,
	342	+ "MGATHER A2/A3 NZ table requires staticShape[3] == FRACTAL_NZ_ROW (16).");
	343	+ static_assert(GlobalTable::staticShape[4] == C0_SIZE_BYTE / sizeof(T),
	344	+ "MGATHER A2/A3 NZ table requires staticShape[4] == 32 / sizeof(T).");
	345	+ static_assert(DstTile::Cols % (C0_SIZE_BYTE / sizeof(T)) == 0,
	346	+ "MGATHER A2/A3 NZ destination tile Cols must be a multiple of C0 (= 32 / sizeof(T)).");
	347	+ static_assert(DstTile::Rows % FRACTAL_NZ_ROW == 0,
	348	+ "MGATHER A2/A3 NZ destination tile Rows must be a multiple of FRACTAL_NZ_ROW (16).");
	349	+ }
	350	+
	351	+ constexpr int kDstValidR = DstTile::ValidRow;
	352	+ constexpr int kDstValidC = DstTile::ValidCol;
	353	+ constexpr int kIdxValidR = IdxTile::ValidRow;
	354	+ constexpr int kIdxValidC = IdxTile::ValidCol;
	355	+
	356	+ if constexpr (Mode == Coalesce::Row) {
	357	+ if constexpr (kDstValidR > 0 && kIdxValidR > 0 && kIdxValidC > 0) {
	358	+ static_assert(kIdxValidR == 1 && kIdxValidC == kDstValidR,
	359	+ "MGATHER A2/A3 Coalesce::Row requires index tile valid shape [1, R].");
	360	+ }
	361	+ } else {
	362	+ if constexpr (kDstValidR > 0 && kIdxValidR > 0) {
	363	+ static_assert(kIdxValidR == kDstValidR,
	364	+ "MGATHER A2/A3 Coalesce::Elem requires index tile ValidRow == destination ValidRow.");
	365	+ }
	366	+ if constexpr (kDstValidC > 0 && kIdxValidC > 0) {
	367	+ static_assert(kIdxValidC == kDstValidC,
	368	+ "MGATHER A2/A3 Coalesce::Elem requires index tile ValidCol == destination ValidCol.");
	369	+ }
	370	+ }
	371	+}
	372	+
	373	+template <Coalesce Mode = Coalesce::Row, GatherOOB Oob = GatherOOB::Undefined, typename DstTile, typename GlobalTable,
	374	+ typename IdxTile>
	375	+PTO_INTERNAL void MGATHER_IMPL(DstTile &dst, GlobalTable &table, IdxTile &indices)
	376	+{
	377	+ using T = typename DstTile::DType;
	378	+ using TIdx = typename IdxTile::DType;
	379	+
	380	+ MGatherCheck<Mode, Oob, DstTile, GlobalTable, IdxTile>();
	381	+
	382	+ __gm__ T tablePtr = reinterpret_cast<__gm__ T >(table.data());
	383	+
	384	+ const uint32_t validRow = dst.GetValidRow();
	385	+ const uint32_t validCol = dst.GetValidCol();
	386	+
	387	+ constexpr bool kIsTableNZ = (GlobalTable::layout == Layout::NZ);
	388	+
	389	+ if constexpr (kIsTableNZ) {
	390	+ const int gShape0 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_0));
	391	+ const int gShape1 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_1));
	392	+ const int gShape2 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_2));
	393	+ const int gStride0 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_0));
	394	+ const int gStride1 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_1));
	395	+ const int gStride2 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_2));
	396	+ const int gStride3 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_3));
	397	+ const int gStride4 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_4));
	398	+
	399	+ if constexpr (Mode == Coalesce::Row) {
	400	+ MGatherRowNzImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, gShape0,
	401	+ gShape1, gShape2, gStride0, gStride1, gStride2, gStride3);
	402	+ } else {
	403	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	404	+ const uint32_t nLogicalCols = static_cast<uint32_t>(gShape0 * gShape1) * kC0;
	405	+ const uint32_t tableSize = static_cast<uint32_t>(gShape2 * FRACTAL_NZ_ROW) * nLogicalCols;
	406	+ MGatherElemNzImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, validCol,
	407	+ tableSize, gShape0, gShape1, gStride0, gStride1, gStride2,
	408	+ gStride3, gStride4, nLogicalCols);
	409	+ }
	410	+ } else {
	411	+ if constexpr (Mode == Coalesce::Row) {
	412	+ const uint32_t tableRows =
	413	+ static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) *
	414	+ table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3));
	415	+ const uint32_t tableRowStride = static_cast<uint32_t>(table.GetStride(GlobalTensorDim::DIM_3));
	416	+ MGatherRowImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, validCol,
	417	+ tableRows, tableRowStride);
	418	+ } else {
	419	+ const uint32_t tableSize =
	420	+ static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) *
	421	+ table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) *
	422	+ table.GetShape(GlobalTensorDim::DIM_4));
	423	+ MGatherElemImpl<Oob, T, TIdx, DstTile, IdxTile>(dst.data(), tablePtr, indices.data(), validRow, validCol,
	424	+ tableSize);
	425	+ }
	426	+ }
	427	+}
	428	+
	429	+} // namespace pto
	430	+
	431	+#endif // MGATHER_HPP

Ainclude/pto/npu/a2a3/MScatter.hpp+493-0

@@ -0,0 +1,493 @@
	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 MSCATTER_HPP
	12	+#define MSCATTER_HPP
	13	+
	14	+#include <pto/common/utils.hpp>
	15	+#include <pto/common/constants.hpp>
	16	+#include "common.hpp"
	17	+
	18	+namespace pto {
	19	+
	20	+enum class ScatterAtomicOp : uint8_t
	21	+{
	22	+ None = 0,
	23	+ Add = 1,
	24	+ Max = 2,
	25	+ Min = 3
	26	+};
	27	+
	28	+enum class ScatterOOB : uint8_t
	29	+{
	30	+ Undefined = 0,
	31	+ Skip = 1,
	32	+ Clamp = 2,
	33	+ Wrap = 3
	34	+};
	35	+
	36	+#ifndef PTO_COALESCE_ENUM_DEFINED
	37	+#define PTO_COALESCE_ENUM_DEFINED
	38	+enum class Coalesce : uint8_t
	39	+{
	40	+ Row = 0,
	41	+ Elem = 1
	42	+};
	43	+#endif
	44	+
	45	+template <typename T>
	46	+struct IsValidMScatterDType {
	47	+ static constexpr bool value = std::is_same_v<T, int8_t> \|\| std::is_same_v<T, uint8_t> \|\|
	48	+ std::is_same_v<T, int16_t> \|\| std::is_same_v<T, uint16_t> \|\|
	49	+ std::is_same_v<T, int32_t> \|\| std::is_same_v<T, uint32_t> \|\|
	50	+ std::is_same_v<T, half> \|\| std::is_same_v<T, bfloat16_t> \|\| std::is_same_v<T, float>;
	51	+};
	52	+
	53	+template <typename T, ScatterAtomicOp Atomic, Coalesce Mode>
	54	+struct IsValidMScatterAtomic {
	55	+ static constexpr bool value =
	56	+ (Atomic == ScatterAtomicOp::None) \|\|
	57	+ ((Atomic == ScatterAtomicOp::Add) &&
	58	+ (std::is_same_v<T, float> \|\| std::is_same_v<T, int32_t> \|\| std::is_same_v<T, half> \|\|
	59	+ std::is_same_v<T, bfloat16_t> \|\| std::is_same_v<T, int16_t> \|\| std::is_same_v<T, int8_t>));
	60	+};
	61	+
	62	+template <typename Tile>
	63	+struct IsMScatterNDTile {
	64	+ static constexpr bool value = Tile::isRowMajor && (Tile::SFractal == SLayout::NoneBox);
	65	+};
	66	+
	67	+template <typename Tile>
	68	+struct IsMScatterNZTile {
	69	+ static constexpr bool value =
	70	+ !Tile::isRowMajor && (Tile::SFractal == SLayout::RowMajor) && (Tile::SFractalSize == TileConfig::fractalABSize);
	71	+};
	72	+
	73	+template <ScatterOOB Oob>
	74	+AICORE PTO_INLINE uint32_t mscatter_remap(uint32_t idx, uint32_t cap, uint32_t &doWrite)
	75	+{
	76	+ if constexpr (Oob == ScatterOOB::Undefined) {
	77	+ doWrite = 1u;
	78	+ return idx;
	79	+ } else if constexpr (Oob == ScatterOOB::Skip) {
	80	+ doWrite = (idx < cap) ? 1u : 0u;
	81	+ return idx;
	82	+ } else if constexpr (Oob == ScatterOOB::Clamp) {
	83	+ doWrite = 1u;
	84	+ return (idx >= cap) ? (cap - 1u) : idx;
	85	+ } else {
	86	+ doWrite = 1u;
	87	+ return idx % cap;
	88	+ }
	89	+}
	90	+
	91	+template <typename T>
	92	+AICORE PTO_INLINE void MScatterRowDma(__gm__ T dst, __ubuf__ T src, uint32_t lenBytes)
	93	+{
	94	+ if constexpr (sizeof(T) == 1) {
	95	+ copy_ubuf_to_gm_align_b8(dst, src, 0, 1, lenBytes, 0, 0, 0, 0);
	96	+ } else if constexpr (sizeof(T) == 2) {
	97	+ copy_ubuf_to_gm_align_b16(dst, src, 0, 1, lenBytes, 0, 0, 0, 0);
	98	+ } else if constexpr (sizeof(T) == 4) {
	99	+ copy_ubuf_to_gm_align_b32(dst, src, 0, 1, lenBytes, 0, 0, 0, 0);
	100	+ }
	101	+}
	102	+
	103	+template <typename T>
	104	+AICORE PTO_INLINE void MScatterRowMultiDma(__gm__ T dst, __ubuf__ T src, uint16_t nBurst, uint32_t lenBytes,
	105	+ uint32_t ubGapBlocks, uint32_t gmGapBytes)
	106	+{
	107	+ if constexpr (sizeof(T) == 1) {
	108	+ copy_ubuf_to_gm_align_b8(dst, src, 0, nBurst, lenBytes, 0, 0, ubGapBlocks, gmGapBytes);
	109	+ } else if constexpr (sizeof(T) == 2) {
	110	+ copy_ubuf_to_gm_align_b16(dst, src, 0, nBurst, lenBytes, 0, 0, ubGapBlocks, gmGapBytes);
	111	+ } else if constexpr (sizeof(T) == 4) {
	112	+ copy_ubuf_to_gm_align_b32(dst, src, 0, nBurst, lenBytes, 0, 0, ubGapBlocks, gmGapBytes);
	113	+ }
	114	+}
	115	+
	116	+template <typename T>
	117	+AICORE PTO_INLINE void MScatterAtomicAddSet()
	118	+{
	119	+ if constexpr (std::is_same_v<T, float>) {
	120	+ set_atomic_f32();
	121	+ } else if constexpr (std::is_same_v<T, half>) {
	122	+ set_atomic_f16();
	123	+ } else if constexpr (std::is_same_v<T, bfloat16_t>) {
	124	+ set_atomic_bf16();
	125	+ } else if constexpr (std::is_same_v<T, int32_t>) {
	126	+ set_atomic_s32();
	127	+ } else if constexpr (std::is_same_v<T, int16_t>) {
	128	+ set_atomic_s16();
	129	+ } else if constexpr (std::is_same_v<T, int8_t>) {
	130	+ set_atomic_s8();
	131	+ }
	132	+ set_atomic_add();
	133	+}
	134	+
	135	+AICORE PTO_INLINE void MScatterAtomicNone()
	136	+{
	137	+ set_atomic_none();
	138	+}
	139	+
	140	+template <typename T>
	141	+AICORE PTO_INLINE uint64_t MScatterNZGmOffset(uint32_t logicalRow, uint32_t logicalCol, int gShape0, int gShape1,
	142	+ int gStride0, int gStride1, int gStride2, int gStride3, int gStride4)
	143	+{
	144	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	145	+ constexpr uint32_t kFRow = FRACTAL_NZ_ROW;
	146	+ const uint32_t blockColCombined = logicalCol / kC0;
	147	+ const uint32_t colInBlock = logicalCol - blockColCombined * kC0;
	148	+ const uint32_t blockRow = logicalRow / kFRow;
	149	+ const uint32_t rowInBlock = logicalRow - blockRow * kFRow;
	150	+ const uint32_t blockColOuter0 = (gShape0 == 1) ? 0u : (blockColCombined / (uint32_t)gShape1);
	151	+ const uint32_t blockColOuter1 = (gShape0 == 1) ? blockColCombined : (blockColCombined - blockColOuter0 * gShape1);
	152	+ return (uint64_t)blockColOuter0 * (uint64_t)gStride0 + (uint64_t)blockColOuter1 * (uint64_t)gStride1 +
	153	+ (uint64_t)blockRow * (uint64_t)gStride2 + (uint64_t)rowInBlock * (uint64_t)gStride3 +
	154	+ (uint64_t)colInBlock * (uint64_t)gStride4;
	155	+}
	156	+
	157	+template <ScatterAtomicOp Atomic, typename T>
	158	+AICORE PTO_INLINE void MScatterScalarStore(__gm__ T *dst, T value)
	159	+{
	160	+ if constexpr (Atomic == ScatterAtomicOp::Add) {
	161	+ if constexpr (std::is_same_v<T, half> \|\| std::is_same_v<T, bfloat16_t>) {
	162	+ float prev = static_cast<float>(*dst);
	163	+ float vsrc = static_cast<float>(value);
	164	+ *dst = static_cast<T>(prev + vsrc);
	165	+ } else if constexpr (std::is_same_v<T, int8_t>) {
	166	+ int32_t prev = static_cast<int32_t>(*dst);
	167	+ int32_t vsrc = static_cast<int32_t>(value);
	168	+ *dst = static_cast<T>(prev + vsrc);
	169	+ } else if constexpr (std::is_same_v<T, int16_t>) {
	170	+ int32_t prev = static_cast<int32_t>(*dst);
	171	+ int32_t vsrc = static_cast<int32_t>(value);
	172	+ *dst = static_cast<T>(prev + vsrc);
	173	+ } else {
	174	+ dst = static_cast<T>(dst + value);
	175	+ }
	176	+ } else {
	177	+ *dst = value;
	178	+ }
	179	+}
	180	+
	181	+template <ScatterAtomicOp Atomic, ScatterOOB Oob, typename T, typename TIdx, typename SrcTile, typename IdxTile>
	182	+__tf__ AICORE void MScatterRowImpl(__gm__ T *tablePtr, typename SrcTile::TileDType __in__ src,
	183	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, uint32_t validCol,
	184	+ uint32_t tableRows, uint32_t tableRowStride)
	185	+{
	186	+ __ubuf__ T srcPtr = (__ubuf__ T )__cce_get_tile_ptr(src);
	187	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	188	+
	189	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	190	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_S>();
	191	+
	192	+ if constexpr (Atomic == ScatterAtomicOp::Add) {
	193	+ MScatterAtomicAddSet<T>();
	194	+ }
	195	+
	196	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	197	+
	198	+ const uint32_t lenBytes = validCol * sizeof(T);
	199	+ constexpr uint32_t kRowStride = SrcTile::RowStride;
	200	+
	201	+ for (uint32_t r = 0; r < validRow; r++) {
	202	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[r]);
	203	+ uint32_t doWrite;
	204	+ uint32_t safeIdx = mscatter_remap<Oob>(rawIdx, tableRows, doWrite);
	205	+ if (doWrite) {
	206	+ __gm__ T dstRow = tablePtr + static_cast<uint64_t>(safeIdx) tableRowStride;
	207	+ __ubuf__ T srcRow = srcPtr + r kRowStride;
	208	+ MScatterRowDma<T>(dstRow, srcRow, lenBytes);
	209	+ }
	210	+ }
	211	+
	212	+ if constexpr (Atomic == ScatterAtomicOp::Add) {
	213	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	214	+ MScatterAtomicNone();
	215	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	216	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	217	+ }
	218	+
	219	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_V>();
	220	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_MTE2>();
	221	+}
	222	+
	223	+template <ScatterAtomicOp Atomic, ScatterOOB Oob, typename T, typename TIdx, typename SrcTile, typename IdxTile>
	224	+__tf__ AICORE void MScatterRowNzImpl(__gm__ T *tablePtr, typename SrcTile::TileDType __in__ src,
	225	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, int gShape0,
	226	+ int gShape1, int gShape2, int gStride0, int gStride1, int gStride2, int gStride3)
	227	+{
	228	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	229	+ constexpr uint32_t kFRow = FRACTAL_NZ_ROW;
	230	+ constexpr uint32_t kFractalRowBytes = kC0 * sizeof(T);
	231	+
	232	+ __ubuf__ T srcPtr = (__ubuf__ T )__cce_get_tile_ptr(src);
	233	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	234	+
	235	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	236	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_S>();
	237	+
	238	+ if constexpr (Atomic == ScatterAtomicOp::Add) {
	239	+ MScatterAtomicAddSet<T>();
	240	+ }
	241	+
	242	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	243	+
	244	+ const uint32_t tableLogicalRows = (uint32_t)gShape2 * kFRow;
	245	+ const uint32_t gmGapBytes = ((uint32_t)gStride1 - kC0) * (uint32_t)sizeof(T);
	246	+ constexpr uint32_t ubGapBlocks = (uint32_t)SrcTile::Rows - 1u;
	247	+ const int64_t tileOuterStrideElem = (int64_t)gShape1 * (int64_t)SrcTile::Rows * (int64_t)kC0;
	248	+
	249	+ for (uint32_t r = 0; r < validRow; r++) {
	250	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[r]);
	251	+ uint32_t doWrite;
	252	+ uint32_t safeIdx = mscatter_remap<Oob>(rawIdx, tableLogicalRows, doWrite);
	253	+ if (doWrite) {
	254	+ const uint32_t dstBlockRow = safeIdx / kFRow;
	255	+ const uint32_t dstRowInBlock = safeIdx - dstBlockRow * kFRow;
	256	+ const uint32_t srcBlockRow = r / kFRow;
	257	+ const uint32_t srcRowInBlock = r - srcBlockRow * kFRow;
	258	+
	259	+ for (uint32_t i = 0; i < (uint32_t)gShape0; i++) {
	260	+ __gm__ T dstAddr = tablePtr + (int64_t)i (int64_t)gStride0 +
	261	+ (int64_t)dstBlockRow * (int64_t)gStride2 +
	262	+ (int64_t)dstRowInBlock * (int64_t)gStride3;
	263	+ __ubuf__ T srcAddr = srcPtr + (int64_t)i tileOuterStrideElem +
	264	+ (int64_t)srcBlockRow * (int64_t)kFRow * (int64_t)kC0 +
	265	+ (int64_t)srcRowInBlock * (int64_t)kC0;
	266	+ MScatterRowMultiDma<T>(dstAddr, srcAddr, (uint16_t)gShape1, kFractalRowBytes, ubGapBlocks, gmGapBytes);
	267	+ }
	268	+ }
	269	+ }
	270	+
	271	+ if constexpr (Atomic == ScatterAtomicOp::Add) {
	272	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	273	+ MScatterAtomicNone();
	274	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	275	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	276	+ }
	277	+
	278	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_V>();
	279	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_MTE2>();
	280	+}
	281	+
	282	+template <ScatterAtomicOp Atomic, ScatterOOB Oob, typename T, typename TIdx, typename SrcTile, typename IdxTile>
	283	+__tf__ AICORE void MScatterElemImpl(__gm__ T *tablePtr, typename SrcTile::TileDType __in__ src,
	284	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, uint32_t validCol,
	285	+ uint32_t tableSize)
	286	+{
	287	+ __ubuf__ T srcPtr = (__ubuf__ T )__cce_get_tile_ptr(src);
	288	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	289	+
	290	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	291	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_S>();
	292	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	293	+
	294	+ constexpr uint32_t kSrcRowStride = SrcTile::RowStride;
	295	+ constexpr uint32_t kIdxRowStride = IdxTile::RowStride;
	296	+
	297	+ for (uint32_t r = 0; r < validRow; r++) {
	298	+ const uint32_t srcRowOff = r * kSrcRowStride;
	299	+ const uint32_t idxRowOff = r * kIdxRowStride;
	300	+ for (uint32_t c = 0; c < validCol; c++) {
	301	+ const uint32_t idxOff = idxRowOff + c;
	302	+ const uint32_t srcOff = srcRowOff + c;
	303	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[idxOff]);
	304	+ uint32_t doWrite;
	305	+ uint32_t safeIdx = mscatter_remap<Oob>(rawIdx, tableSize, doWrite);
	306	+ if (doWrite) {
	307	+ MScatterScalarStore<Atomic, T>(tablePtr + safeIdx, srcPtr[srcOff]);
	308	+ }
	309	+ }
	310	+ }
	311	+
	312	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	313	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	314	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	315	+}
	316	+
	317	+template <ScatterAtomicOp Atomic, ScatterOOB Oob, typename T, typename TIdx, typename SrcTile, typename IdxTile>
	318	+__tf__ AICORE void MScatterElemNzImpl(__gm__ T *tablePtr, typename SrcTile::TileDType __in__ src,
	319	+ typename IdxTile::TileDType __in__ indices, uint32_t validRow, uint32_t validCol,
	320	+ uint32_t tableSize, int gShape0, int gShape1, int gStride0, int gStride1,
	321	+ int gStride2, int gStride3, int gStride4, uint32_t nLogicalCols)
	322	+{
	323	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	324	+ __ubuf__ T srcPtr = (__ubuf__ T )__cce_get_tile_ptr(src);
	325	+ __ubuf__ TIdx idxPtr = (__ubuf__ TIdx )__cce_get_tile_ptr(indices);
	326	+
	327	+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
	328	+ PtoSetWaitFlag<PIPE_MTE2, PIPE_S>();
	329	+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
	330	+
	331	+ constexpr uint32_t kIdxRowStride = IdxTile::RowStride;
	332	+ const uint32_t nColBlocks = (validCol + kC0 - 1u) / kC0;
	333	+ const uint32_t kSrcColBlockStride = (uint32_t)SrcTile::Rows * kC0;
	334	+
	335	+ for (uint32_t bcol = 0; bcol < nColBlocks; bcol++) {
	336	+ const uint32_t cBase = bcol * kC0;
	337	+ const uint32_t cLimit = (cBase + kC0 < validCol) ? (cBase + kC0) : validCol;
	338	+ const uint32_t kInBlock = cLimit - cBase;
	339	+ __ubuf__ T srcBlockBase = srcPtr + (uint64_t)bcol (uint64_t)kSrcColBlockStride;
	340	+ for (uint32_t r = 0; r < validRow; r++) {
	341	+ const uint32_t idxRowOff = r * kIdxRowStride;
	342	+ __ubuf__ T srcRowBase = srcBlockBase + (uint64_t)r (uint64_t)kC0;
	343	+ for (uint32_t cInner = 0; cInner < kInBlock; cInner++) {
	344	+ const uint32_t c = cBase + cInner;
	345	+ const uint32_t idxOff = idxRowOff + c;
	346	+ uint32_t rawIdx = static_cast<uint32_t>(idxPtr[idxOff]);
	347	+ uint32_t doWrite;
	348	+ uint32_t safeIdx = mscatter_remap<Oob>(rawIdx, tableSize, doWrite);
	349	+ if (doWrite) {
	350	+ const uint32_t logicalRow = safeIdx / nLogicalCols;
	351	+ const uint32_t logicalCol = safeIdx - logicalRow * nLogicalCols;
	352	+ const uint64_t dstOff = MScatterNZGmOffset<T>(logicalRow, logicalCol, gShape0, gShape1, gStride0,
	353	+ gStride1, gStride2, gStride3, gStride4);
	354	+ MScatterScalarStore<Atomic, T>(tablePtr + dstOff, srcRowBase[cInner]);
	355	+ }
	356	+ }
	357	+ }
	358	+ }
	359	+
	360	+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
	361	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE2>();
	362	+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
	363	+}
	364	+
	365	+template <Coalesce Mode, ScatterAtomicOp Atomic, typename GlobalTable, typename SrcTile, typename IdxTile>
	366	+PTO_INTERNAL void MScatterCheck()
	367	+{
	368	+ using T = typename SrcTile::DType;
	369	+ using TIdx = typename IdxTile::DType;
	370	+
	371	+ static_assert(IsValidMScatterDType<T>::value,
	372	+ "MSCATTER A2/A3 data type must be int8/uint8/int16/uint16/int32/uint32/half/bfloat16/float.");
	373	+ static_assert(std::is_same_v<TIdx, int32_t> \|\| std::is_same_v<TIdx, uint32_t>,
	374	+ "MSCATTER A2/A3 index type must be int32_t or uint32_t.");
	375	+ static_assert(std::is_same_v<typename GlobalTable::DType, __gm__ T>,
	376	+ "MSCATTER A2/A3 destination table must be a GM GlobalTensor with element type matching the source.");
	377	+ static_assert(SrcTile::Loc == TileType::Vec, "MSCATTER A2/A3 source must be a Vec tile (UB).");
	378	+ static_assert(IdxTile::Loc == TileType::Vec, "MSCATTER A2/A3 indices must be a Vec tile (UB).");
	379	+
	380	+ static_assert(IdxTile::isRowMajor, "MSCATTER A2/A3 index tile must be BLayout::RowMajor.");
	381	+ static_assert(IdxTile::SFractal == SLayout::NoneBox, "MSCATTER A2/A3 index tile must be ND (SLayout::NoneBox).");
	382	+
	383	+ constexpr bool kIsTableND = (GlobalTable::layout == Layout::ND);
	384	+ constexpr bool kIsTableNZ = (GlobalTable::layout == Layout::NZ);
	385	+ constexpr bool kIsSrcND = IsMScatterNDTile<SrcTile>::value;
	386	+ constexpr bool kIsSrcNZ = IsMScatterNZTile<SrcTile>::value;
	387	+
	388	+ static_assert(kIsTableND \|\| kIsTableNZ, "MSCATTER A2/A3 table must use Layout::ND or Layout::NZ.");
	389	+ static_assert((kIsTableND && kIsSrcND) \|\| (kIsTableNZ && kIsSrcNZ),
	390	+ "MSCATTER A2/A3 layout pairing must be either:\n"
	391	+ " (a) GM Layout::ND + UB tile (BLayout::RowMajor + SLayout::NoneBox), or\n"
	392	+ " (b) GM Layout::NZ + UB tile (BLayout::ColMajor + SLayout::RowMajor + SFractalSize=512).");
	393	+
	394	+ static_assert(SrcTile::Cols * sizeof(T) % BLOCK_BYTE_SIZE == 0,
	395	+ "MSCATTER A2/A3 source tile padded Cols*sizeof(T) must be 32-byte aligned.");
	396	+
	397	+ if constexpr (kIsTableNZ) {
	398	+ static_assert(GlobalTable::staticShape[3] == FRACTAL_NZ_ROW,
	399	+ "MSCATTER A2/A3 NZ table requires staticShape[3] == FRACTAL_NZ_ROW (16).");
	400	+ static_assert(GlobalTable::staticShape[4] == C0_SIZE_BYTE / sizeof(T),
	401	+ "MSCATTER A2/A3 NZ table requires staticShape[4] == 32 / sizeof(T).");
	402	+ static_assert(SrcTile::Cols % (C0_SIZE_BYTE / sizeof(T)) == 0,
	403	+ "MSCATTER A2/A3 NZ source tile Cols must be a multiple of C0 (= 32 / sizeof(T)).");
	404	+ static_assert(SrcTile::Rows % FRACTAL_NZ_ROW == 0,
	405	+ "MSCATTER A2/A3 NZ source tile Rows must be a multiple of FRACTAL_NZ_ROW (16).");
	406	+ }
	407	+
	408	+ static_assert(IsValidMScatterAtomic<T, Atomic, Mode>::value,
	409	+ "MSCATTER A2/A3 atomic operation: Add only valid for int8/int16/int32/half/bfloat16/float; "
	410	+ "Max/Min not supported.");
	411	+
	412	+ constexpr int kSrcValidR = SrcTile::ValidRow;
	413	+ constexpr int kSrcValidC = SrcTile::ValidCol;
	414	+ constexpr int kIdxValidR = IdxTile::ValidRow;
	415	+ constexpr int kIdxValidC = IdxTile::ValidCol;
	416	+
	417	+ if constexpr (Mode == Coalesce::Row) {
	418	+ if constexpr (kSrcValidR > 0 && kIdxValidR > 0 && kIdxValidC > 0) {
	419	+ static_assert(kIdxValidR == 1 && kIdxValidC == kSrcValidR,
	420	+ "MSCATTER A2/A3 Coalesce::Row requires index tile valid shape [1, R].");
	421	+ }
	422	+ } else {
	423	+ if constexpr (kSrcValidR > 0 && kIdxValidR > 0) {
	424	+ static_assert(kIdxValidR == kSrcValidR,
	425	+ "MSCATTER A2/A3 Coalesce::Elem requires index tile ValidRow == source ValidRow.");
	426	+ }
	427	+ if constexpr (kSrcValidC > 0 && kIdxValidC > 0) {
	428	+ static_assert(kIdxValidC == kSrcValidC,
	429	+ "MSCATTER A2/A3 Coalesce::Elem requires index tile ValidCol == source ValidCol.");
	430	+ }
	431	+ }
	432	+}
	433	+
	434	+template <Coalesce Mode = Coalesce::Row, ScatterAtomicOp Atomic = ScatterAtomicOp::None,
	435	+ ScatterOOB Oob = ScatterOOB::Undefined, typename GlobalTable, typename SrcTile, typename IdxTile>
	436	+PTO_INTERNAL void MSCATTER_IMPL(GlobalTable &table, SrcTile &src, IdxTile &indices)
	437	+{
	438	+ using T = typename SrcTile::DType;
	439	+ using TIdx = typename IdxTile::DType;
	440	+
	441	+ MScatterCheck<Mode, Atomic, GlobalTable, SrcTile, IdxTile>();
	442	+
	443	+ __gm__ T tablePtr = reinterpret_cast<__gm__ T >(table.data());
	444	+
	445	+ const uint32_t validRow = src.GetValidRow();
	446	+ const uint32_t validCol = src.GetValidCol();
	447	+
	448	+ constexpr bool kIsTableNZ = (GlobalTable::layout == Layout::NZ);
	449	+
	450	+ if constexpr (kIsTableNZ) {
	451	+ const int gShape0 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_0));
	452	+ const int gShape1 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_1));
	453	+ const int gShape2 = static_cast<int>(table.GetShape(GlobalTensorDim::DIM_2));
	454	+ const int gStride0 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_0));
	455	+ const int gStride1 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_1));
	456	+ const int gStride2 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_2));
	457	+ const int gStride3 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_3));
	458	+ const int gStride4 = static_cast<int>(table.GetStride(GlobalTensorDim::DIM_4));
	459	+
	460	+ if constexpr (Mode == Coalesce::Row) {
	461	+ MScatterRowNzImpl<Atomic, Oob, T, TIdx, SrcTile, IdxTile>(tablePtr, src.data(), indices.data(), validRow,
	462	+ gShape0, gShape1, gShape2, gStride0, gStride1,
	463	+ gStride2, gStride3);
	464	+ } else {
	465	+ constexpr uint32_t kC0 = C0_SIZE_BYTE / sizeof(T);
	466	+ const uint32_t nLogicalCols = static_cast<uint32_t>(gShape0 * gShape1) * kC0;
	467	+ const uint32_t tableSize = static_cast<uint32_t>(gShape2 * FRACTAL_NZ_ROW) * nLogicalCols;
	468	+ MScatterElemNzImpl<Atomic, Oob, T, TIdx, SrcTile, IdxTile>(
	469	+ tablePtr, src.data(), indices.data(), validRow, validCol, tableSize, gShape0, gShape1, gStride0,
	470	+ gStride1, gStride2, gStride3, gStride4, nLogicalCols);
	471	+ }
	472	+ } else {
	473	+ if constexpr (Mode == Coalesce::Row) {
	474	+ const uint32_t tableRows =
	475	+ static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) *
	476	+ table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3));
	477	+ const uint32_t tableRowStride = static_cast<uint32_t>(table.GetStride(GlobalTensorDim::DIM_3));
	478	+ MScatterRowImpl<Atomic, Oob, T, TIdx, SrcTile, IdxTile>(tablePtr, src.data(), indices.data(), validRow,
	479	+ validCol, tableRows, tableRowStride);
	480	+ } else {
	481	+ const uint32_t tableSize =
	482	+ static_cast<uint32_t>(table.GetShape(GlobalTensorDim::DIM_0) * table.GetShape(GlobalTensorDim::DIM_1) *
	483	+ table.GetShape(GlobalTensorDim::DIM_2) * table.GetShape(GlobalTensorDim::DIM_3) *
	484	+ table.GetShape(GlobalTensorDim::DIM_4));
	485	+ MScatterElemImpl<Atomic, Oob, T, TIdx, SrcTile, IdxTile>(tablePtr, src.data(), indices.data(), validRow,
	486	+ validCol, tableSize);
	487	+ }
	488	+ }
	489	+}
	490	+
	491	+} // namespace pto
	492	+
	493	+#endif // MSCATTER_HPP

Mtests/npu/a2a3/src/st/testcase/CMakeLists.txt+2-0

	@@ -210,6 +210,8 @@ tconcatidx
210	tconcat	210	tconcat
211	textract_vec	211	textract_vec
212	tinsert_vec	212	tinsert_vec
		213	+mgather
		214	+mscatter
213	tpushpop_cv	215	tpushpop_cv
214	tpushpop_vc	216	tpushpop_vc
215	tpushpop_cv_nosplit	217	tpushpop_cv_nosplit

Atests/npu/a2a3/src/st/testcase/mgather/CMakeLists.txt+9-0

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

Atests/npu/a2a3/src/st/testcase/mgather/MGATHER.md+481-0

@@ -0,0 +1,481 @@
	1	+# MGATHER (A2/A3 Vec-Core)
	2	+
	3	+## Tile Operation Diagram
	4	+
	5	+![MGATHER tile operation](../../../../../../../docs/figures/isa/MGATHER.svg)
	6	+
	7	+## Introduction
	8	+
	9	+`MGATHER` performs an indexed gather from a GM `GlobalTensor` into a UB destination tile through a UB index tile, running on the A2/A3 AIV vector core. It is dispatched as a sequential walk driven from the scalar pipe; there is no async-launch or cross-core orchestration — the kernel is a single AIV function call. The operating mode is selected explicitly through the `Coalesce` template parameter:
	10	+
	11	+- `Coalesce::Row` (default) — gather full rows from `table[idx[r], :]` into `dst[r, :]`. The index tile is 1-D (`[1, R]`, row-major). For each row `r` the scalar pipe reads `idx[r]` and issues one MTE2 burst. For ND tables that burst is `copy_gm_to_ubuf_align_b8/b16/b32` of `validCol * sizeof(T)` bytes from `table + idx[r] * tableRowStride`. For NZ tables it is a multi-burst MTE2 transfer (`nBurst = number of column fractals`) that walks every `C0`-wide fractal of the source row.
	12	+- `Coalesce::Elem` — element-wise gather from a linearized `table` into `dst[R, C]` (or `dst[1, N]`) through `idx[R, C]`. The index tile must have the same valid shape as the destination. For every `(r, c)` the scalar pipe reads the index, applies the OOB remap, and copies the element with a scalar `dstUb[r, c] = tableGm[idx[r, c]]`. The (1, 1) shape is a degenerate case of the same loop.
	13	+
	14	+Both modes accept either an ND GM table (`Layout::ND`) paired with an ND/RowMajor UB tile, or an NZ GM table (`Layout::NZ`) paired with an NZ/ColMajor-fractal UB tile (see "NZ Layout Support" below).
	15	+
	16	+Out-of-bounds handling is selected through the `GatherOOB` template parameter. `MGATHER` has no atomic or conflict policy: every destination slot has exactly one defined source index, so collisions cannot occur.
	17	+
	18	+### Why Elem mode uses scalar GM reads
	19	+
	20	+`copy_gm_to_ubuf_align_b8/b16/b32` requires the UB destination address to be 32-byte aligned, and the destination must be a whole number of 32-byte burst chunks. A per-element MTE2 burst of `lenBurst = sizeof(T)` to `dstPtr + r * RowStride + c` does not satisfy that rule whenever `(c * sizeof(T)) % 32 != 0`, which covers almost every elem-mode lane. On the simulator the runtime accepts the misaligned burst; on real A2/A3 hardware the transfer silently drops, leaving the destination lane at its initial value (typically zero). Row mode does not hit this problem because each row write starts at `r * RowStride`, and `RowStride * sizeof(T)` is always a multiple of 32 bytes.
	21	+
	22	+The Elem mode therefore uses scalar GM→UB copies, which have element-level addressing granularity and place no alignment requirement on the destination. Atomic semantics are not needed for gather (no destination is written from multiple sources) and OOB::Zero collapses into a direct scalar zero-write. Per-element MTE2 dispatch is the only A2/A3 mechanism that could give pure vec-core elem-mode throughput, and the hardware's destination-alignment rule rules it out for arbitrary column offsets. Scalar GM↔UB is the (1, 1) fallback already validated on hardware; we extend it to all elem shapes.
	23	+
	24	+## Math Interpretation
	25	+
	26	+### Row Coalesce (`Coalesce::Row`)
	27	+
	28	+Destination `dst[R, C]`, index `idx[1, R]`, table `table[TableRows, C]`:
	29	+
	30	+$$ \mathrm{dst}_{r, j} = \mathrm{table}_{\mathrm{idx}_{r},\; j} \quad\text{for } 0 \le r < R,\; 0 \le j < C $$
	31	+
	32	+The kernel issues one GM→UB DMA burst per row through `copy_gm_to_ubuf_align_b`, with burst length `validCol sizeof(T)` bytes (the valid width, not the padded `Tile::Cols`). UB destination addressing uses `Tile::RowStride`, so partial-valid tiles padded for 32-byte burst alignment are supported transparently.
	33	+
	34	+### Element Coalesce (`Coalesce::Elem`)
	35	+
	36	+Destination `dst[R, C]`, index `idx[R, C]` (same valid shape as `dst`), flat table of length `TableSize`:
	37	+
	38	+$$ \mathrm{dst}_{r, c} = \mathrm{table}[\mathrm{idx}_{r, c}] \quad\text{for } 0 \le r < R,\; 0 \le c < C $$
	39	+
	40	+`TableSize = Shape[0] * Shape[1] * Shape[2] * Shape[3] * Shape[4]` of the `GlobalTensor` (5-D, any combination of static and dynamic dims). For ND tables this matches the linear element count directly; for NZ tables the kernel maps each scalar `idx` through a row-major (`logicalRow = idx / nLogicalCols`, `logicalCol = idx % nLogicalCols`) → NZ block-stride translation.
	41	+
	42	+For each `(r, c)`:
	43	+
	44	+1. Scalar pipe reads `rawIdx = idxPtr[r * IdxTile::RowStride + c]`.
	45	+2. `safeIdx = mgather_remap<Oob>(rawIdx, tableSize, doRead)` applies the OOB policy.
	46	+3. If `doRead`, `dstPtr[r * DstTile::RowStride + c] = tablePtr[gmOff]` (ND: `gmOff = safeIdx`; NZ: `gmOff = MGatherNZGmOffset(...)` after splitting `safeIdx` into logical row/col).
	47	+4. If `!doRead` and `Oob == GatherOOB::Zero`, `dstPtr[...] = static_cast<T>(0)`.
	48	+
	49	+NZ Elem walks block-col-major (outer block-col, then row, then column-within-block) so consecutive writes to UB stay within the same 32 B fractal block.
	50	+
	51	+### Out-of-Bounds Behaviour
	52	+
	53	+```cpp
	54	+enum class GatherOOB : uint8_t {
	55	+ Undefined = 0,
	56	+ Clamp = 1,
	57	+ Wrap = 2,
	58	+ Zero = 3
	59	+};
	60	+```
	61	+
	62	+`capacity` is `TableRows` (Row mode) or `TableSize` (Elem mode):
	63	+
	64	+- `Undefined`: caller guarantees `idx < capacity`; no remap is applied.
	65	+- `Clamp`: `idx = min(idx, capacity - 1)` before access.
	66	+- `Wrap`: `idx = idx % capacity` before access.
	67	+- `Zero`: out-of-bounds destinations receive `static_cast<T>(0)`.
	68	+ - Row mode (ND and NZ). OOB rows are not DMA'd; the destination row is filled with `T(0)`. ND fills the row inline on the scalar pipe (per-row `validCol` writes); NZ pre-zeroes the whole tile once before the DMA loop so every fractal slot has a defined value regardless of OOB membership.
	69	+ - Elem mode (ND and NZ). OOB lanes write `T(0)` directly from the scalar loop. The `if/else` inside the lane handles both branches with a single store.
	70	+
	71	+All dtypes are supported under every `GatherOOB` value; no `static_assert` restricts the dtype set for `Elem + GatherOOB::Zero`.
	72	+
	73	+## Assembly Syntax
	74	+
	75	+PTO-AS form: see [docs/assembly/PTO-AS.md](/docs/assembly/PTO-AS.md).
	76	+
	77	+```text
	78	+mgather.row %dst, %table, %idx : (!pto.tile<RxCxT>, !pto.memref<...>, !pto.tile<1xRxi32>)
	79	+mgather.elem %dst, %table, %idx : (!pto.tile<RxCxT>, !pto.memref<...>, !pto.tile<RxCxi32>)
	80	+```
	81	+
	82	+OOB-aware variants append the mode suffix (`mgather.row.clamp`, `mgather.elem.zero`, etc.).
	83	+
	84	+## C++ Intrinsic
	85	+
	86	+Declared in `include/pto/common/pto_instr.hpp` and `include/pto/npu/a2a3/MGather.hpp`:
	87	+
	88	+```cpp
	89	+template <Coalesce CMode = Coalesce::Row,
	90	+ GatherOOB Oob = GatherOOB::Undefined,
	91	+ typename TileDst, typename GlobalTable, typename TileIdx,
	92	+ typename... WaitEvents>
	93	+PTO_INST RecordEvent MGATHER(TileDst& dst, GlobalTable& table, TileIdx& idx,
	94	+ WaitEvents&... events);
	95	+```
	96	+
	97	+The kernel iterates over `TileDst::ValidRow * TileDst::ValidCol` logical positions; physical UB strides come from each tile's `RowStride` (which equals padded `Cols` for `BLayout::RowMajor`).
	98	+
	99	+Parameters:
	100	+
	101	+- `dst` : UB destination tile (`TileType::Vec`); shape `[R, C]`. `BLayout::RowMajor` for ND tables, `BLayout::ColMajor` + `SLayout::RowMajor` + `SFractalSize=512` for NZ tables.
	102	+- `table` : Source GM `GlobalTensor` with `Layout::ND` (linear contiguous) or `Layout::NZ` (fractal `[B, BlockCols, BlockRows, 16, 32/sizeof(T)]`). The `GlobalTensor::DType` must be `__gm__ T` matching the destination element type.
	103	+- `idx` : UB index tile (`TileType::Vec`). For `Coalesce::Row`: 1-D `[1, R]` row-major. For `Coalesce::Elem`: same valid shape as `dst`, row-major.
	104	+- `CMode` : `Coalesce` — `Row` (default) or `Elem`. First template parameter, so the operating mode is always explicit at the call site.
	105	+- `Oob` : `GatherOOB` — out-of-bounds index handling.
	106	+
	107	+## Coalesce Mode
	108	+
	109	+```cpp
	110	+enum class Coalesce : uint8_t {
	111	+ Row = 0,
	112	+ Elem = 1
	113	+};
	114	+```
	115	+
	116	+## Constraints
	117	+
	118	+### Data Types
	119	+
	120	+`TileDst::DType` must be one of: `int8_t`, `uint8_t`, `int16_t`, `uint16_t`, `int32_t`, `uint32_t`, `half`, `bfloat16_t`, `float`. (No `float8_e4m3_t` / `float8_e5m2_t` / `hifloat8_t` on A2/A3 vec-core.)
	121	+
	122	+### Index Types
	123	+
	124	+`TileIdx::DType` must be `int32_t` or `uint32_t`.
	125	+
	126	+### Tile Constraints
	127	+
	128	+- `TileDst::Loc == TileType::Vec` (UB).
	129	+- `TileIdx::Loc == TileType::Vec` (UB). The index tile is always `BLayout::RowMajor + SLayout::NoneBox` (ND) regardless of the table layout.
	130	+- Source and table must share the same element type `T` (`GlobalTable::DType == __gm__ T`).
	131	+- The destination tile's bulk + sub layout must be paired with the table layout exactly:
	132	+ - `GlobalTable::layout == Layout::ND` ⇒ `TileDst` is `BLayout::RowMajor + SLayout::NoneBox`.
	133	+ - `GlobalTable::layout == Layout::NZ` ⇒ `TileDst` is `BLayout::ColMajor + SLayout::RowMajor + SFractalSize == TileConfig::fractalABSize` (= 512 B). In addition:
	134	+ - `GlobalTable::staticShape[3] == FRACTAL_NZ_ROW` (= 16),
	135	+ - `GlobalTable::staticShape[4] == C0_SIZE_BYTE / sizeof(T)` (= 32 B / element width),
	136	+ - `TileDst::Cols % (C0_SIZE_BYTE / sizeof(T)) == 0` (whole `C0` columns per fractal block-col),
	137	+ - `TileDst::Rows % FRACTAL_NZ_ROW == 0` (whole `16`-row fractal blocks).
	138	+- Padded `TileDst::Cols * sizeof(T)` must be 32-byte aligned in both layouts (the same DMA-burst rule that `TLOAD` / `TSTORE` enforce). `ValidCol` / `ValidRow` are not constrained by this rule — they only set the kernel's iteration bounds.
	139	+- For `Coalesce::Row`: `TileIdx::ValidRow == 1`, `TileIdx::ValidCol == TileDst::ValidRow` (a 1-D row of `R` indices).
	140	+- For `Coalesce::Elem`: `TileIdx::ValidRow == TileDst::ValidRow` and `TileIdx::ValidCol == TileDst::ValidCol`.
	141	+- Both row and elem modes require `TileDst::ValidRow >= 1` and `TileDst::ValidCol >= 1`.
	142	+
	143	+### Dynamic Runtime Shapes
	144	+
	145	+`MGATHER` supports both compile-time fixed shapes and runtime-dynamic shapes for the source `GlobalTensor` and the destination / index `Tile`s. Any dimension declared as `DYNAMIC` (`-1`) at template-instantiation time is resolved at runtime through the standard PTO accessors:
	146	+
	147	+- `Tile<…, RowMask, ColMask>` with `RowMask == -1` and/or `ColMask == -1` stores the runtime valid extents in the tile object; `MGATHER_IMPL` reads them through `dst.GetValidRow()` / `dst.GetValidCol()` and uses them to drive the loop bounds.
	148	+- `Shape<S0, S1, S2, S3, S4>` / `Stride<…>` with one or more `-1` entries are constructed with the runtime sizes; `MGATHER_IMPL` reads them through `table.GetShape(GlobalTensorDim::DIM_*)` and folds them into `tableRows` (Row mode) or `tableSize = ∏ shape[0..4]` (Elem mode).
	149	+
	150	+Static-asserts in `MGatherCheck` are gated on `if constexpr (DIM > 0)`, so they fire only for compile-time-known dimensions; mixed static/dynamic combinations check exactly the static dims and defer the dynamic ones to runtime arithmetic. Padded `Tile::Rows` / `Tile::Cols` are always compile-time (they govern the UB DMA-burst alignment); only the valid sub-region and the GM table extents may be dynamic.
	151	+
	152	+Example:
	153	+
	154	+```cpp
	155	+constexpr auto kPadCols = 16;
	156	+using DstTileT = Tile<TileType::Vec, float, 1, kPadCols, BLayout::RowMajor, -1, -1>;
	157	+using IdxTileT = Tile<TileType::Vec, int32_t, 1, kPadCols, BLayout::RowMajor, -1, -1>;
	158	+using TableShape = Shape<1, 1, 1, -1, -1>;
	159	+using TableStride = Stride<1, 1, 1, -1, -1>;
	160	+
	161	+int64_t validCols = 9, d3 = 3, d4 = 10, srcStride3 = 10;
	162	+TableShape tableShape(d3, d4);
	163	+TableStride tableStride(srcStride3, (int64_t)1);
	164	+GlobalTensor<float, TableShape, TableStride> tableGM(srcGm, tableShape, tableStride);
	165	+
	166	+DstTileT dstTile(1, validCols);
	167	+IdxTileT idxTile(1, validCols);
	168	+TASSIGN(dstTile, dstUbOffsetBytes);
	169	+TASSIGN(idxTile, idxUbOffsetBytes);
	170	+
	171	+MGATHER<Coalesce::Elem, GatherOOB::Undefined>(dstTile, tableGM, idxTile);
	172	+```
	173	+
	174	+At dispatch time `MGATHER_IMPL` resolves `validRows = 1`, `validCols = 9`, and `tableSize = 1·1·1·3·10 = 30`. The padded UB `Tile::Cols = 16` is purely a `TLOAD` burst-alignment artifact — the elem loop only walks the valid 9 elements.
	175	+
	176	+### Layout Support
	177	+
	178	+The kernel handles two paired layouts: ND-GM with ND-UB, and NZ-GM with NZ-UB. UB addressing is computed from the tile's `Rows` / `Cols` plus an optional fractal block-col stride; GM addressing is driven from the `GlobalTensor::GetStride(DIM_*)` accessors.
	179	+
	180	+\| Tile / Tensor \| Supported layouts \| Notes \|
	181	+\|---------------\|-------------------\|-------\|
	182	+\| `TileDst` (UB) — ND path \| `BLayout::RowMajor` + `SLayout::NoneBox` \| Row writes use `dstPtr + r * Tile::RowStride`. Elem writes use `dstPtr + r * Tile::RowStride + c`. \|
	183	+\| `TileDst` (UB) — NZ path \| `BLayout::ColMajor` + `SLayout::RowMajor` + `SFractalSize == 512` \| Block-col stride is `Tile::Rows * C0`; per-element offset is `(c / C0) * (Tile::Rows * C0) + r * C0 + (c % C0)`. \|
	184	+\| `TileIdx` (UB) — Row mode \| `[1, R]` `BLayout::RowMajor` + `SLayout::NoneBox` \| Linear `R`-element layout in UB; the kernel reads `idxPtr[row]` directly. Always ND, regardless of the table layout. \|
	185	+\| `TileIdx` (UB) — Elem mode \| `[R, C]` `BLayout::RowMajor` + `SLayout::NoneBox` \| Reads `idxPtr[r * Tile::RowStride + c]` per element. Always ND, regardless of the table layout. \|
	186	+\| `GlobalTable` (GM) — ND \| `Layout::ND` (linear contiguous addressing); 5-D `Shape<…, R, C>` \| Row mode addresses `table + idx[r] * tableRowStride`; Elem mode addresses `table + idx`. `tableRowStride = GetStride(DIM_3)` so non-trivial row strides (for example zero-padded ND tables) are honoured. \|
	187	+\| `GlobalTable` (GM) — NZ \| `Layout::NZ`; 5-D `Shape<B, BCols, BRows, 16, C0>` with `B == 1`, `staticShape[3] == 16`, `staticShape[4] == 32 / sizeof(T)` \| Row mode walks fractal block-cols through a multi-burst MTE2 (`nBurst = BCols`, `lenBurst = C0 * sizeof(T)`, `gmGap = stride1 - C0`); Elem mode resolves each `idx` into a NZ block-stride offset (`blockColCombined / blockColOuter1 / blockRow / rowInBlock / colInBlock`) and copies one element with a scalar load. \|
	188	+
	189	+### NZ Layout Support
	190	+
	191	+When `GlobalTable::layout == Layout::NZ` and `TileDst` is the matching `BLayout::ColMajor + SLayout::RowMajor + SFractalSize=512` tile, `MGATHER` runs the dedicated NZ paths (`MGatherRowNzImpl`, `MGatherElemNzImpl`).
	192	+
	193	+- Constants. `kC0 = C0_SIZE_BYTE / sizeof(T) = 32 / sizeof(T)`; `kFRow = FRACTAL_NZ_ROW = 16`. Each fractal block is `kFRow × kC0` elements (= 32 B × 16 = 512 B).
	194	+- Logical shape. Logical rows = `gShape2 * kFRow` (number of NZ row-blocks × 16). Logical cols = `gShape0 * gShape1 * kC0` (batch × col-blocks × C0). For Row mode `mgather_remap` clamps/wraps against the logical row count; for Elem mode it clamps/wraps against the total element count `(gShape2 * kFRow) * (gShape0 * gShape1 * kC0)`.
	195	+- Row mode. For each logical row `r`, the kernel maps `idx[r]` to `(srcBlockRow, srcRowInBlock) = (idx / kFRow, idx % kFRow)` and `(dstBlockRow, dstRowInBlock) = (r / kFRow, r % kFRow)`, then issues one multi-burst MTE2 transfer per outer batch (`nBurst = gShape1`, `lenBurst = kC0 * sizeof(T) = 32 B`, `gmGap = (gStride1 - kC0) * sizeof(T)`, `ubGap = TileDst::Rows - 1` blocks). Every column-fractal in the source row-of-blocks is gathered with a single instruction; the GM gap honours `gStride1` (the actual stride between block-cols, not the implicit `BlockRows * 16 * C0`). When `Oob == GatherOOB::Zero`, the kernel pre-fills the whole tile with `T(0)` before the DMA loop and simply skips DMAs for OOB rows.
	196	+- Elem mode. For each `(r, c)` the kernel maps `idx` to `(logicalRow, logicalCol) = (idx / nLogicalCols, idx % nLogicalCols)`, then to NZ physical offsets through `MGatherNZGmOffset` (which folds `gShape0/1` and `gStride0..4`, supporting both packed and stride-padded NZ tensors). The destination offset is `(c / kC0) * (TileDst::Rows * kC0) + r * kC0 + (c % kC0)`. The walk order is block-col → row → col-in-block so consecutive writes always target consecutive 32 B UB blocks; row-major iteration would alternate writes between the first and the second column-fractal block of each row, which complicates the scalar walk. Out-of-bounds lanes write `T(0)` inline when `Oob == GatherOOB::Zero`.
	197	+- Stride vs. valid-shape. `MGatherNzImpl` reads strides from the `GlobalTensor` runtime (`GetStride(DIM_)`), so packed NZ tensors (`gStride1 == gShape2 * gShape3 * gShape4`) and stride-padded NZ tensors (`gStride1 > gShape2 * gShape3 * gShape4`) both work without any caller-side adjustment — Row mode propagates `gStride1` into the multi-burst `gmGap`, and Elem mode threads every stride term into `MGatherNZGmOffset`.
	198	+
	199	+### Aligned vs Unaligned Tile Shapes
	200	+
	201	+The kernel does not care whether the tile's logical shape is "aligned" — it walks all `ValidRow * ValidCol` positions:
	202	+
	203	+- Row mode (ND): per-row DMA `lenBurst = validCol * sizeof(T)` (any byte length supported by `copy_gm_to_ubuf_align_b`); `Tile::RowStride sizeof(T)` is forced 32-byte aligned by the upstream `Tile` system, so subsequent rows always start on a 32-byte burst boundary.
	204	+- Row mode (NZ): one multi-burst transfer per logical row × outer-batch; `lenBurst = kC0 * sizeof(T) = 32 B` is fixed by the fractal layout, so per-row alignment is automatic. `validRow` does not have to be a multiple of `kFRow` — the kernel only walks the valid logical rows and writes their fractal-mapped UB slots, leaving fractal-padding rows untouched (caller-zeroed).
	205	+- Elem mode: one scalar GM→UB copy per element. The scalar pipe has element-level addressing granularity, so unaligned valid sub-regions inside an aligned padded tile work without further constraints. The padded `Tile::Cols * sizeof(T)` still has to be 32-byte aligned (enforced upstream so `TSTORE` of the destination tile works), but `ValidCol` can take any value `1 ≤ ValidCol ≤ Tile::Cols`.
	206	+
	207	+Callers handle "unaligned valid region" by:
	208	+
	209	+1. Padding the tile up to the nearest 32-byte alignment (for example valid `[3, 3]` int32 → tile `[3, 8]`), and
	210	+2. Either zero-initializing the padding (`TASSIGN`-then-clear) or only inspecting the valid region post-gather.
	211	+
	212	+### Minimum Tile Shape
	213	+
	214	+`MGatherCheck` accepts any `(ValidRow, ValidCol)` with `ValidRow, ValidCol >= 1` (including the degenerate `(1, 1)` for both Row and Elem modes).
	215	+
	216	+The actual lower bound on the padded `Tile<…, Rows, Cols, BLayout, ValidRow, ValidCol>` shape is enforced upstream by the `Tile` system because every `TLOAD` / `TSTORE` that brings data in/out of UB issues 32-byte GM↔UB DMA bursts. The contiguous-in-memory dim of the tile must therefore be a whole number of bursts:
	217	+
	218	+- `BLayout::RowMajor` ⇒ `Cols * sizeof(T) % 32 == 0` (one row = N×32 B).
	219	+
	220	+`ValidRow` / `ValidCol` are not constrained by this rule. So a logical `(1, 1)` int32 tile is expressed as `Tile<int32, 1, 8, RowMajor, 1, 1>` (one padded burst, one valid element); `(3, 3)` int32 as `Tile<int32, 3, 8, RowMajor, 3, 3>`; and so on. The smallest padded `Cols` per dtype for a row-major tile is:
	221	+
	222	+\| `T` \| Min `Cols` (`BLayout::RowMajor`) \|
	223	+\|-----\|---------------------------------\|
	224	+\| `int8` / `uint8` \| 32 \|
	225	+\| `int16` / `uint16` / `half` / `bfloat16` \| 16 \|
	226	+\| `int32` / `uint32` / `float` \| 8 \|
	227	+
	228	+The padded dimension is purely a TLOAD/TSTORE alignment artifact — `MGATHER` itself walks `ValidRow * ValidCol` positions and addresses through `Tile::RowStride`.
	229	+
	230	+### Mode Resolution
	231	+
	232	+Mode is explicit, not auto-detected. The static-asserts in `MGatherCheck` validate that the supplied tile shapes match the chosen `Coalesce` value:
	233	+
	234	+```text
	235	+Coalesce::Row : Idx.ValidRow == 1 && Idx.ValidCol == Dst.ValidRow
	236	+Coalesce::Elem : Idx.ValidRow == Dst.ValidRow && Idx.ValidCol == Dst.ValidCol
	237	+```
	238	+
	239	+## Pipe / Synchronisation Model
	240	+
	241	+The implementation centralises every pipe handshake the kernel needs. Callers do not need to insert any extra barriers beyond the standard `TLOAD` post-load `set_flag(PIPE_MTE2, PIPE_V)` / `wait_flag(PIPE_MTE2, PIPE_V)` pair that brings the index tile into a clean state on the vector pipe before `MGATHER`. `MGATHER` never uses `pipe_barrier(PIPE_ALL)` in the kernel — every wait is a specific producer→consumer pair, so unrelated pipes keep running in parallel.
	242	+
	243	+\| Phase \| Pipe transition \| What it guards \|
	244	+\|-------\|-----------------\|----------------\|
	245	+\| Pre-amble (Row, Elem ND, Elem NZ) \| `set_flag(PIPE_V, PIPE_S)` / `wait_flag(PIPE_V, PIPE_S)` and `set_flag(PIPE_MTE3, PIPE_S)` / `wait_flag(PIPE_MTE3, PIPE_S)`; Elem path also adds `set_flag(PIPE_MTE2, PIPE_S)` / `wait_flag(PIPE_MTE2, PIPE_S)` \| Make the index tile visible to scalar reads (V→S transitively waits for MTE2 through the caller's MTE2→V flag; the explicit MTE2→S in Elem mode is a defensive guard for callers that omit the V handshake). Also flush any pending vector / MTE3 writes that might overlap UB before the scalar loop starts. \|
	246	+\| Body (Row mode, ND) \| `copy_gm_to_ubuf_align_b` per row \| One DMA per row, `lenBurst = validCol sizeof(T)`; trip count = `validRow`. Issued from the scalar pipe, executed on PIPE_MTE2. \|
	247	+\| Body (Row mode, NZ) \| `copy_gm_to_ubuf_align_b` multi-burst per logical row × batch \| `nBurst = gShape1`, `lenBurst = C0 sizeof(T) = 32 B`, `gmGap = (gStride1 - C0) * sizeof(T)`, `ubGap = Tile::Rows - 1` blocks; trip count = `validRow * gShape0`. \|
	248	+\| Body (Elem mode, ND and NZ) \| Scalar `dstUb[r, c] = tableGm[gmOff]` per element \| Per-element scalar GM→UB copy; trip count = `validRow * validCol`. NZ walks block-col-major to keep each 32 B UB block written contiguously in time. OOB::Zero lanes write `T(0)` inline through the same scalar store. \|
	249	+\| Row mode post-amble \| `set_flag(PIPE_S, PIPE_MTE2)` / `wait_flag` then `set_flag(PIPE_MTE2, PIPE_V/MTE3)` / `wait_flag` and `set_flag(PIPE_S, PIPE_V/MTE3)` / `wait_flag` \| Drain the MTE2 DMAs before the next consumer touches the destination tile, and release the scalar pipe to V and MTE3 (any caller that issues `set_flag(PIPE_V, PIPE_MTE3)` after `MGATHER` therefore sees the gathered rows on both V and MTE3). \|
	250	+\| Elem mode post-amble \| `set_flag(PIPE_S, PIPE_V)` / `wait_flag`, `set_flag(PIPE_S, PIPE_MTE2)` / `wait_flag`, `set_flag(PIPE_S, PIPE_MTE3)` / `wait_flag` \| Make the scalar UB writes visible to V (for downstream vector ops), MTE2 (for follow-up gathers), and MTE3 (for `TSTORE`). The S→MTE3 flag is what bridges the gap between the scalar gather body and the caller's `set_flag(PIPE_V, PIPE_MTE3)` / `TSTORE` pair. \|
	251	+
	252	+## Examples
	253	+
	254	+### Row Coalesce — Embedding Lookup
	255	+
	256	+```cpp
	257	+#include <pto/pto-inst.hpp>
	258	+#include <pto/common/constants.hpp>
	259	+
	260	+using namespace pto;
	261	+
	262	+template <typename T, int R, int C, int TableRows>
	263	+__global__ AICORE void example_embedding_lookup(__gm__ T outPtr, __gm__ T tablePtr, __gm__ int32_t *idxPtr)
	264	+{
	265	+ using DstTile = Tile<TileType::Vec, T, R, C, BLayout::RowMajor, R, C>;
	266	+ using IdxTile = Tile<TileType::Vec, int32_t, 1, R, BLayout::RowMajor, 1, R>;
	267	+
	268	+ using TableShape = Shape<1, 1, 1, TableRows, C>;
	269	+ using TableStride = Stride<1, 1, 1, C, 1>;
	270	+ using TableTensor = GlobalTensor<T, TableShape, TableStride>;
	271	+
	272	+ using IdxShape = Shape<1, 1, 1, 1, R>;
	273	+ using IdxStride = Stride<1, 1, 1, R, 1>;
	274	+ using IdxTensor = GlobalTensor<int32_t, IdxShape, IdxStride>;
	275	+
	276	+ using OutShape = Shape<1, 1, 1, R, C>;
	277	+ using OutStride = Stride<1, 1, 1, C, 1>;
	278	+ using OutTensor = GlobalTensor<T, OutShape, OutStride>;
	279	+
	280	+ TableTensor tableGM(tablePtr);
	281	+ IdxTensor idxGM(idxPtr);
	282	+ OutTensor outGM(outPtr);
	283	+ DstTile dst; TASSIGN(dst, 0x0000);
	284	+ IdxTile idx; TASSIGN(idx, 0x1000);
	285	+
	286	+ TLOAD(idx, idxGM);
	287	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	288	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	289	+ MGATHER<Coalesce::Row, GatherOOB::Clamp>(dst, tableGM, idx);
	290	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	291	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	292	+ TSTORE(outGM, dst);
	293	+}
	294	+```
	295	+
	296	+### Element Coalesce — 2-D Random Access
	297	+
	298	+```cpp
	299	+#include <pto/pto-inst.hpp>
	300	+
	301	+using namespace pto;
	302	+
	303	+__global__ AICORE void example_elem_2d(__gm__ float outPtr, __gm__ float tablePtr, __gm__ int32_t *idxPtr)
	304	+{
	305	+ constexpr int R = 8, C = 32, TableSize = 256;
	306	+
	307	+ using DstTile = Tile<TileType::Vec, float, R, C, BLayout::RowMajor, R, C>;
	308	+ using IdxTile = Tile<TileType::Vec, int32_t, R, C, BLayout::RowMajor, R, C>;
	309	+
	310	+ using TableShape = Shape<1, 1, 1, 1, TableSize>;
	311	+ using TableStride = Stride<1, 1, 1, TableSize, 1>;
	312	+ using TableTensor = GlobalTensor<float, TableShape, TableStride>;
	313	+
	314	+ TableTensor tableGM(tablePtr);
	315	+ DstTile dst; TASSIGN(dst, 0x0000);
	316	+ IdxTile idx; TASSIGN(idx, 0x0800);
	317	+
	318	+ MGATHER<Coalesce::Elem, GatherOOB::Wrap>(dst, tableGM, idx);
	319	+}
	320	+```
	321	+
	322	+### Element Coalesce — `(1, 1)` Degenerate Case
	323	+
	324	+```cpp
	325	+#include <pto/pto-inst.hpp>
	326	+
	327	+using namespace pto;
	328	+
	329	+__global__ AICORE void example_scalar(__gm__ float outPtr, __gm__ float tablePtr, __gm__ int32_t *idxPtr)
	330	+{
	331	+ constexpr int TableSize = 32;
	332	+
	333	+ using DstTile = Tile<TileType::Vec, float, 1, 8, BLayout::RowMajor, 1, 1>;
	334	+ using IdxTile = Tile<TileType::Vec, int32_t, 1, 8, BLayout::RowMajor, 1, 1>;
	335	+
	336	+ using TableShape = Shape<1, 1, 1, 1, TableSize>;
	337	+ using TableStride = Stride<1, 1, 1, TableSize, 1>;
	338	+ using TableTensor = GlobalTensor<float, TableShape, TableStride>;
	339	+
	340	+ TableTensor tableGM(tablePtr);
	341	+ DstTile dst; TASSIGN(dst, 0x0000);
	342	+ IdxTile idx; TASSIGN(idx, 0x0080);
	343	+
	344	+ MGATHER<Coalesce::Elem>(dst, tableGM, idx);
	345	+}
	346	+```
	347	+
	348	+## Performance Considerations
	349	+
	350	+1. Row vs. Elem. Row coalesce achieves the best aggregate bandwidth — one wide DMA per logical row (ND) or one multi-burst DMA per logical row × batch (NZ). Elem coalesce issues one scalar GM read + UB write per active lane: there is no DMA-engine pipelining, and throughput is bound by the scalar pipe's GM access latency. Prefer Row whenever the indexing structure permits.
	351	+2. Sequential scalar loop (Elem). A2/A3 dispatches `MGATHER` as a single-thread sequential walk of the `validRow * validCol` lanes. Loop trip counts of `ValidCol ≤ 32 / sizeof(T)` rows are the sweet spot; large flat tiles are bound by scalar GM read latency. The block-col-major walk used for NZ keeps consecutive writes spatially-local in UB.
	352	+3. Why not per-element MTE2 in Elem mode. The MTE2 DMA `copy_gm_to_ubuf_align_b` intrinsics require a 32-byte aligned UB destination address, which a per-element burst at `dstPtr + r RowStride + c` cannot satisfy for arbitrary `c`. On hardware the misaligned burst drops silently, so a per-element MTE2 elem mode would zero out almost every lane. The scalar GM→UB path has element-level addressing granularity, matches the (1, 1) fallback that already passes on hardware, and supports every dtype uniformly.
	353	+4. DMA cost (Row).
	354	+ - ND: each row is one `copy_gm_to_ubuf_align_b` call with `nBurst = 1`, `lenBurst = validCol sizeof(T)`.
	355	+ - NZ: each (logical row, batch) pair is one `copy_gm_to_ubuf_align_b` call with `nBurst = gShape1` (column block-cols), `lenBurst = C0 sizeof(T) = 32 B`, `gmGap = (gStride1 - C0) * sizeof(T)`, `ubGap = Tile::Rows - 1` blocks. The kernel trip count is `validRow * gShape0`.
	356	+
	357	+ A2/A3 pipelines the MTE2 DMA bursts through the DMA engine, but back-pressure is bounded by `MAX_OUTSTANDING_MTE2`; for very large row counts the kernel still issues all DMAs unconditionally — there is no row chunking.
	358	+5. OOB policy cost.
	359	+ - `Undefined`: zero overhead — caller guarantees valid indices.
	360	+ - `Clamp` / `Wrap`: a single arithmetic remap per lane (`min` / `mod`).
	361	+ - `Zero`: Row mode skips DMAs for OOB rows and either writes `T(0)` per lane (ND) or pre-zeroes the whole tile (NZ); Elem mode writes `T(0)` inline through the same scalar store branch.
	362	+6. Single-pass dispatch. `MGATHER` is a regular AIV function call from the kernel (no async-launch or cross-core orchestration). The whole gather completes as a sequential scalar / MTE2 pipeline within one AIV invocation; concurrency comes from the DMA engine pipelining row DMAs behind the scalar issue loop, not from multiple worker threads.
	363	+
	364	+## Related Instructions
	365	+
	366	+- [`TLOAD`](/docs/isa/TLOAD.md): Contiguous block transfer from GM to Tile.
	367	+- [`MSCATTER`](../mscatter/MSCATTER.md): Indexed scatter from Tile to GM (inverse operation).
	368	+- [`TGATHER`](/docs/isa/TGATHER.md): Index-based gather within tiles (UB-to-UB on the same vec-core).
	369	+
	370	+## Test Cases
	371	+
	372	+The A2/A3 ST suite covers 63 cases distributed across data types, modes, OOB handling, alignment patterns, dynamic shapes, and the ND ↔ NZ layout pair (including a dedicated NZ + Elem + `OOB::Zero` case). Each case follows the standard A2/A3 ST pattern: `gen_data.py` produces `table.bin`, `indices.bin`, and `golden.bin`; `mgather_kernel.cpp` instantiates the kernel template and `<<<1, nullptr, stream>>>`-launches it; `main.cpp` (`MGATHERTest`) reads inputs, copies them to GM, runs the kernel, fetches the output, and compares against golden with `eps = 0.0f` (max-diff = 0).
	373	+
	374	+### Row Coalesce — `[1, R]` index form
	375	+
	376	+\| Case \| Data Type \| Dst Size \| TableRows \| OOB Mode \|
	377	+\|------\|-----------\|----------\|-----------\|----------\|
	378	+\| `case_row_float_8x32_64rows` \| float \| 8×32 \| 64 \| Undefined \|
	379	+\| `case_row_half_16x64_64rows` \| half \| 16×64 \| 64 \| Undefined \|
	380	+\| `case_row_bfloat16_16x16_64rows` \| bf16 \| 16×16 \| 64 \| Undefined \|
	381	+\| `case_row_int32_8x16_32rows` \| int32 \| 8×16 \| 32 \| Undefined \|
	382	+\| `case_row_uint32_8x16_32rows` \| uint32 \| 8×16 \| 32 \| Undefined \|
	383	+\| `case_row_int16_8x16_32rows` \| int16 \| 8×16 \| 32 \| Undefined \|
	384	+\| `case_row_uint16_8x16_32rows` \| uint16 \| 8×16 \| 32 \| Undefined \|
	385	+\| `case_row_int8_8x32_32rows` \| int8 \| 8×32 \| 32 \| Undefined \|
	386	+\| `case_row_uint8_8x32_32rows` \| uint8 \| 8×32 \| 32 \| Undefined \|
	387	+\| `case_row_float_clamp_8x32_8rows` \| float \| 8×32 \| 8 \| Clamp \|
	388	+\| `case_row_int32_wrap_8x16_8rows` \| int32 \| 8×16 \| 8 \| Wrap \|
	389	+\| `case_row_half_zero_8x32_8rows` \| half \| 8×32 \| 8 \| Zero \|
	390	+
	391	+### Row Coalesce — Unaligned / Odd Valid Rows / Padded
	392	+
	393	+\| Case \| Data Type \| Valid Dst \| Padded Dst \| OOB Mode \|
	394	+\|------\|-----------\|-----------\|------------\|----------\|
	395	+\| `case_row_int32_unaligned_3x8_8rows` \| int32 \| 3×8 \| 3×8 \| Undefined \|
	396	+\| `case_row_float_partial_4x16_in_8x16` \| float \| 4×16 \| 8×16 \| Undefined \|
	397	+\| `case_row_half_partial_5x32_in_8x32` \| half \| 5×32 \| 8×32 \| Undefined \|
	398	+\| `case_row_uint8_unaligned_3x32_32rows` \| uint8 \| 3×32 \| 3×32 \| Undefined \|
	399	+\| `case_row_int16_partial_3x16_in_4x16` \| int16 \| 3×16 \| 4×16 \| Clamp \|
	400	+
	401	+### Element Coalesce — 1-D destination `[1, N]`
	402	+
	403	+\| Case \| Data Type \| Valid N / TableSize \| OOB Mode \|
	404	+\|------\|-----------\|---------------------\|----------\|
	405	+\| `case_elem_float_64_128size` \| float \| 64 / 128 \| Undefined \|
	406	+\| `case_elem_half_64_128size` \| half \| 64 / 128 \| Undefined \|
	407	+\| `case_elem_bfloat16_64_128size` \| bf16 \| 64 / 128 \| Undefined \|
	408	+\| `case_elem_int32_32_64size` \| int32 \| 32 / 64 \| Undefined \|
	409	+\| `case_elem_uint32_32_64size` \| uint32 \| 32 / 64 \| Undefined \|
	410	+\| `case_elem_int16_32_64size` \| int16 \| 32 / 64 \| Undefined \|
	411	+\| `case_elem_uint16_32_64size` \| uint16 \| 32 / 64 \| Undefined \|
	412	+\| `case_elem_int8_64_128size` \| int8 \| 64 / 128 \| Undefined \|
	413	+\| `case_elem_uint8_64_128size` \| uint8 \| 64 / 128 \| Undefined \|
	414	+\| `case_elem_float_clamp_32_16size` \| float \| 32 / 16 \| Clamp \|
	415	+\| `case_elem_int32_wrap_32_16size` \| int32 \| 32 / 16 \| Wrap \|
	416	+\| `case_elem_half_zero_32_16size` \| half \| 32 / 16 \| Zero \|
	417	+
	418	+### Element Coalesce — 2-D destination `[R, C]`
	419	+
	420	+\| Case \| Data Type \| Dst Size \| TableSize \| OOB Mode \|
	421	+\|------\|-----------\|----------\|-----------\|----------\|
	422	+\| `case_elem2d_float_8x32_256size` \| float \| 8×32 \| 256 \| Undefined \|
	423	+\| `case_elem2d_int32_8x16_256size` \| int32 \| 8×16 \| 256 \| Undefined \|
	424	+\| `case_elem2d_half_4x32_256size` \| half \| 4×32 \| 256 \| Undefined \|
	425	+\| `case_elem2d_bfloat16_4x32_256size` \| bf16 \| 4×32 \| 256 \| Undefined \|
	426	+\| `case_elem2d_uint8_4x64_256size` \| uint8 \| 4×64 \| 256 \| Undefined \|
	427	+\| `case_elem2d_int8_4x64_256size` \| int8 \| 4×64 \| 256 \| Undefined \|
	428	+\| `case_elem2d_int16_4x32_256size` \| int16 \| 4×32 \| 256 \| Undefined \|
	429	+\| `case_elem2d_uint16_4x32_256size` \| uint16 \| 4×32 \| 256 \| Undefined \|
	430	+\| `case_elem2d_uint32_8x16_256size` \| uint32 \| 8×16 \| 256 \| Undefined \|
	431	+\| `case_elem2d_float_wrap_4x16_64size` \| float \| 4×16 \| 64 \| Wrap \|
	432	+\| `case_elem2d_int32_clamp_4x8_32size` \| int32 \| 4×8 \| 32 \| Clamp \|
	433	+\| `case_elem2d_half_zero_4x32_64size` \| half \| 4×32 \| 64 \| Zero \|
	434	+
	435	+### Element Coalesce — Unaligned / Padded / `(1, 1)`
	436	+
	437	+\| Case \| Data Type \| Valid Dst \| Padded Dst \| TableSize \| OOB Mode \|
	438	+\|------\|-----------\|-----------\|------------\|-----------\|----------\|
	439	+\| `case_elem2d_int32_unaligned_3x3_in_3x8_64size` \| int32 \| 3×3 \| 3×8 \| 64 \| Undefined \|
	440	+\| `case_elem2d_float_unaligned_5x5_in_5x8_64size` \| float \| 5×5 \| 5×8 \| 64 \| Undefined \|
	441	+\| `case_elem2d_half_unaligned_3x9_in_3x16_64size` \| half \| 3×9 \| 3×16 \| 64 \| Undefined \|
	442	+\| `case_elem2d_int8_unaligned_3x17_in_3x32_64size`\| int8 \| 3×17 \| 3×32 \| 64 \| Undefined \|
	443	+\| `case_elem_scalar_float_1x1_in_1x8_8size` \| float \| 1×1 \| 1×8 \| 8 \| Undefined \|
	444	+\| `case_elem_scalar_int32_1x1_in_1x8_8size` \| int32 \| 1×1 \| 1×8 \| 8 \| Undefined \|
	445	+\| `case_elem_scalar_half_1x1_in_1x16_16size` \| half \| 1×1 \| 1×16 \| 16 \| Undefined \|
	446	+
	447	+### Dynamic Runtime Shapes
	448	+
	449	+`Tile<…, -1, -1>` (runtime valid extents) paired with `GlobalTensor<…, Shape<1,1,1,-1,-1>, Stride<1,1,1,-1,-1>>` (runtime table shape / stride). The kernel resolves all extents at dispatch through `Tile::GetValidRow/Col()` and `GlobalTensor::GetShape(DIM_*)`; padded `Tile::Rows / Cols` remain compile-time so the UB layout and DMA bursts stay statically known.
	450	+
	451	+\| Case \| Mode \| Data Type \| Runtime Valid Dst \| Padded Dst \| Runtime Table \| OOB Mode \|
	452	+\|------\|------\|-----------\|-------------------\|------------\|----------------\|----------\|
	453	+\| `case_elem2d_dyn_float_4x8_64size` \| Elem \| float \| 4×8 \| 4×8 \| 1×64 \| Undefined \|
	454	+\| `case_elem2d_dyn_int32_3x3_in_3x8_64size` \| Elem \| int32 \| 3×3 \| 3×8 \| 1×64 \| Undefined \|
	455	+\| `case_row_dyn_int32_3x16_8rows` \| Row \| int32 \| 3×16 \| 3×16 \| 8 rows × 16 \| Undefined \|
	456	+\| `case_row_dyn_half_4x32_16rows` \| Row \| half \| 4×32 \| 4×32 \| 16 rows × 32 \| Undefined \|
	457	+
	458	+### NZ Layout — Row Coalesce
	459	+
	460	+GM table is `Layout::NZ` with shape `(1, BlockCols, BlockRows, 16, C0)`; UB destination is the matching `BLayout::ColMajor + SLayout::RowMajor + SFractalSize=512` fractal tile. Each row gather walks all `BlockCols` column-fractals in one multi-burst MTE2 transfer.
	461	+
	462	+\| Case \| Data Type \| Dst Size (logical) \| Block Layout (BR × BC × C0) \| OOB Mode \|
	463	+\|------\|-----------\|---------------------\|------------------------------\|----------\|
	464	+\| `case_row_nz_float_16x16_2blk` \| float \| 16×16 \| 2 × 2 × 8 \| Undefined \|
	465	+\| `case_row_nz_half_32x16_2blk` \| half \| 32×16 \| 2 × 1 × 16 \| Undefined \|
	466	+\| `case_row_nz_int32_16x16_2blk` \| int32 \| 16×16 \| 2 × 2 × 8 \| Undefined \|
	467	+\| `case_row_nz_int16_32x16_1blk` \| int16 \| 32×16 \| 2 × 1 × 16 \| Undefined \|
	468	+\| `case_row_nz_int8_16x32_1blk` \| int8 \| 16×32 \| 2 × 1 × 32 \| Undefined \|
	469	+\| `case_row_nz_float_clamp_16x8_1blk` \| float \| 16×8 \| 2 × 1 × 8 \| Clamp \|
	470	+\| `case_row_nz_half_zero_16x16_2blk` \| half \| 16×16 \| 2 × 1 × 16 \| Zero \|
	471	+
	472	+### NZ Layout — Element Coalesce
	473	+
	474	+GM table is `Layout::NZ`; UB destination is the matching NZ fractal tile. The kernel walks block-col-major, mapping each `idx` through the row-major (`logicalRow`, `logicalCol`) representation to a NZ block-stride GM offset.
	475	+
	476	+\| Case \| Data Type \| Dst Size (logical) \| Block Layout (BR × BC × C0) \| OOB Mode \|
	477	+\|------\|-----------\|---------------------\|------------------------------\|----------\|
	478	+\| `case_elem2d_nz_float_16x16_2blk` \| float \| 16×16 \| 2 × 2 × 8 \| Undefined \|
	479	+\| `case_elem2d_nz_half_16x16_1blk` \| half \| 16×16 \| 2 × 1 × 16 \| Undefined \|
	480	+\| `case_elem2d_nz_int32_16x8_1blk` \| int32 \| 16×8 \| 2 × 1 × 8 \| Undefined \|
	481	+\| `case_elem2d_nz_half_zero_16x16_1blk` \| half \| 16×16 \| 2 × 1 × 16 \| Zero \|

Atests/npu/a2a3/src/st/testcase/mgather/gen_data.py+298-0

@@ -0,0 +1,298 @@
	1	+#!/usr/bin/python3
	2	+# coding=utf-8
	3	+# --------------------------------------------------------------------------------
	4	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	5	+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	6	+# CANN Open Software License Agreement Version 2.0 (the "License").
	7	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	8	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	9	+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	10	+# See LICENSE in the root of the software repository for the full text of the License.
	11	+# --------------------------------------------------------------------------------
	12	+
	13	+import os
	14	+import numpy as np
	15	+
	16	+np.random.seed(42)
	17	+
	18	+
	19	+def make_table(dtype, count):
	20	+ if np.issubdtype(dtype, np.integer):
	21	+ info = np.iinfo(dtype)
	22	+ mod = min(info.max - info.min + 1, 251)
	23	+ else:
	24	+ mod = 251
	25	+ arr = (np.arange(1, count + 1) % mod) + 1
	26	+ return arr.astype(dtype)
	27	+
	28	+
	29	+def make_idx_random(rng, shape, max_val):
	30	+ return rng.integers(0, max_val, size=shape, dtype=np.int32)
	31	+
	32	+
	33	+def make_idx_with_oob(rng, shape, table_size, oob_count):
	34	+ flat = rng.integers(0, table_size, size=int(np.prod(shape)), dtype=np.int32)
	35	+ flat[:oob_count] = rng.integers(table_size, table_size * 2, size=oob_count, dtype=np.int32)
	36	+ rng.shuffle(flat)
	37	+ return flat.reshape(shape)
	38	+
	39	+
	40	+def golden_row(table, idx, dst_rows, dst_cols, oob):
	41	+ table_rows = table.shape[0]
	42	+ out = np.zeros((dst_rows, dst_cols), dtype=table.dtype)
	43	+ flat = idx.reshape(-1)
	44	+ for i in range(dst_rows):
	45	+ raw = int(flat[i])
	46	+ if oob == "undefined":
	47	+ safe = raw
	48	+ elif oob == "clamp":
	49	+ safe = min(max(raw, 0), table_rows - 1)
	50	+ elif oob == "wrap":
	51	+ safe = raw % table_rows
	52	+ elif oob == "zero":
	53	+ safe = raw
	54	+ else:
	55	+ raise ValueError(oob)
	56	+ if oob == "zero" and (raw < 0 or raw >= table_rows):
	57	+ out[i, :] = 0
	58	+ else:
	59	+ out[i, :] = table[safe, :]
	60	+ return out
	61	+
	62	+
	63	+def golden_elem(table_flat, idx, oob):
	64	+ out = np.zeros_like(idx, dtype=table_flat.dtype)
	65	+ table_size = table_flat.shape[0]
	66	+ flat_idx = idx.reshape(-1)
	67	+ flat_out = out.reshape(-1)
	68	+ for i in range(flat_idx.shape[0]):
	69	+ raw = int(flat_idx[i])
	70	+ if oob == "undefined":
	71	+ safe = raw
	72	+ flat_out[i] = table_flat[safe]
	73	+ elif oob == "clamp":
	74	+ safe = min(max(raw, 0), table_size - 1)
	75	+ flat_out[i] = table_flat[safe]
	76	+ elif oob == "wrap":
	77	+ safe = raw % table_size
	78	+ flat_out[i] = table_flat[safe]
	79	+ elif oob == "zero":
	80	+ if 0 <= raw < table_size:
	81	+ flat_out[i] = table_flat[raw]
	82	+ else:
	83	+ flat_out[i] = 0
	84	+ else:
	85	+ raise ValueError(oob)
	86	+ return flat_out.reshape(idx.shape)
	87	+
	88	+
	89	+def case_row(name, dtype, dst_rows, dst_cols, table_rows, oob="undefined", idx_kind="random"):
	90	+ rng = np.random.default_rng(hash(name) & 0xFFFFFFFF)
	91	+ table = make_table(dtype, table_rows * dst_cols).reshape(table_rows, dst_cols)
	92	+ if idx_kind == "random":
	93	+ idx = make_idx_random(rng, (dst_rows, 1), table_rows)
	94	+ elif idx_kind == "oob":
	95	+ idx = make_idx_with_oob(rng, (dst_rows, 1), table_rows, max(1, dst_rows // 2))
	96	+ else:
	97	+ raise ValueError(idx_kind)
	98	+ golden = golden_row(table, idx, dst_rows, dst_cols, oob)
	99	+ return table.reshape(-1), idx, golden
	100	+
	101	+
	102	+def case_elem(name, dtype, n, ts, oob="undefined", idx_kind="random"):
	103	+ rng = np.random.default_rng(hash(name) & 0xFFFFFFFF)
	104	+ table = make_table(dtype, ts)
	105	+ if idx_kind == "random":
	106	+ idx = make_idx_random(rng, (1, n), ts)
	107	+ elif idx_kind == "oob":
	108	+ idx = make_idx_with_oob(rng, (1, n), ts, max(1, n // 2))
	109	+ else:
	110	+ raise ValueError(idx_kind)
	111	+ golden = golden_elem(table, idx, oob)
	112	+ return table, idx, golden
	113	+
	114	+
	115	+def case_elem2d(name, dtype, r, c, ts, oob="undefined", idx_kind="random"):
	116	+ rng = np.random.default_rng(hash(name) & 0xFFFFFFFF)
	117	+ table = make_table(dtype, ts)
	118	+ if idx_kind == "random":
	119	+ idx = make_idx_random(rng, (r, c), ts)
	120	+ elif idx_kind == "oob":
	121	+ idx = make_idx_with_oob(rng, (r, c), ts, max(1, (r * c) // 2))
	122	+ else:
	123	+ raise ValueError(idx_kind)
	124	+ golden = golden_elem(table, idx, oob)
	125	+ return table, idx, golden
	126	+
	127	+
	128	+def nd_to_nz(arr_2d, c0):
	129	+ """Convert ND (R, C) array to NZ-flat layout: [1, C/c0, R/16, 16, c0]."""
	130	+ r, c = arr_2d.shape
	131	+ assert r % 16 == 0 and c % c0 == 0
	132	+ n_block_rows = r // 16
	133	+ n_block_cols = c // c0
	134	+ out = np.zeros((1, n_block_cols, n_block_rows, 16, c0), dtype=arr_2d.dtype)
	135	+ for bc in range(n_block_cols):
	136	+ for br in range(n_block_rows):
	137	+ out[0, bc, br] = arr_2d[br * 16 : (br + 1) * 16, bc * c0 : (bc + 1) * c0]
	138	+ return out.reshape(-1)
	139	+
	140	+
	141	+def case_row_nz(name, dtype, dst_rows, dst_cols, block_rows, block_cols, c0, oob="undefined", idx_kind="random"):
	142	+ rng = np.random.default_rng(hash(name) & 0xFFFFFFFF)
	143	+ table_rows = block_rows * 16
	144	+ table_cols = block_cols * c0
	145	+ assert dst_cols == table_cols
	146	+ table_nd = make_table(dtype, table_rows * table_cols).reshape(table_rows, table_cols)
	147	+ if idx_kind == "random":
	148	+ idx = make_idx_random(rng, (dst_rows, 1), table_rows)
	149	+ elif idx_kind == "oob":
	150	+ idx = make_idx_with_oob(rng, (dst_rows, 1), table_rows, max(1, dst_rows // 2))
	151	+ else:
	152	+ raise ValueError(idx_kind)
	153	+ golden_nd = golden_row(table_nd, idx, dst_rows, dst_cols, oob)
	154	+ table_nz = nd_to_nz(table_nd, c0)
	155	+ golden_nz = nd_to_nz(golden_nd, c0)
	156	+ return table_nz, idx, golden_nz
	157	+
	158	+
	159	+def case_elem2d_nz(name, dtype, dst_rows, dst_cols, block_rows, block_cols, c0, oob="undefined", idx_kind="random"):
	160	+ rng = np.random.default_rng(hash(name) & 0xFFFFFFFF)
	161	+ table_rows = block_rows * 16
	162	+ table_cols = block_cols * c0
	163	+ table_size = table_rows * table_cols
	164	+ table_nd = make_table(dtype, table_size).reshape(table_rows, table_cols)
	165	+ if idx_kind == "random":
	166	+ idx = make_idx_random(rng, (dst_rows, dst_cols), table_size)
	167	+ elif idx_kind == "oob":
	168	+ idx = make_idx_with_oob(rng, (dst_rows, dst_cols), table_size, max(1, (dst_rows * dst_cols) // 2))
	169	+ else:
	170	+ raise ValueError(idx_kind)
	171	+ golden_nd = golden_elem(table_nd.reshape(-1), idx, oob)
	172	+ table_nz = nd_to_nz(table_nd, c0)
	173	+ golden_nz = nd_to_nz(golden_nd.reshape(dst_rows, dst_cols), c0)
	174	+ return table_nz, idx, golden_nz
	175	+
	176	+
	177	+CASES = []
	178	+
	179	+
	180	+def add(name, fn):
	181	+ CASES.append((name, fn))
	182	+
	183	+
	184	+add("MGATHERTest.case_row_float_8x32_64rows", lambda n: case_row(n, np.float32, 8, 32, 64))
	185	+add("MGATHERTest.case_row_half_16x64_64rows", lambda n: case_row(n, np.float16, 16, 64, 64))
	186	+add("MGATHERTest.case_row_bfloat16_16x16_64rows", lambda n: case_row(n, np.uint16, 16, 16, 64))
	187	+add("MGATHERTest.case_row_int32_8x16_32rows", lambda n: case_row(n, np.int32, 8, 16, 32))
	188	+add("MGATHERTest.case_row_uint32_8x16_32rows", lambda n: case_row(n, np.uint32, 8, 16, 32))
	189	+add("MGATHERTest.case_row_int16_8x16_32rows", lambda n: case_row(n, np.int16, 8, 16, 32))
	190	+add("MGATHERTest.case_row_uint16_8x16_32rows", lambda n: case_row(n, np.uint16, 8, 16, 32))
	191	+add("MGATHERTest.case_row_int8_8x32_32rows", lambda n: case_row(n, np.int8, 8, 32, 32))
	192	+add("MGATHERTest.case_row_uint8_8x32_32rows", lambda n: case_row(n, np.uint8, 8, 32, 32))
	193	+add(
	194	+ "MGATHERTest.case_row_float_clamp_8x32_8rows",
	195	+ lambda n: case_row(n, np.float32, 8, 32, 8, oob="clamp", idx_kind="oob"),
	196	+)
	197	+add("MGATHERTest.case_row_int32_wrap_8x16_8rows", lambda n: case_row(n, np.int32, 8, 16, 8, oob="wrap", idx_kind="oob"))
	198	+add(
	199	+ "MGATHERTest.case_row_half_zero_8x32_8rows", lambda n: case_row(n, np.float16, 8, 32, 8, oob="zero", idx_kind="oob")
	200	+)
	201	+
	202	+add("MGATHERTest.case_row_int32_unaligned_3x8_8rows", lambda n: case_row(n, np.int32, 3, 8, 8))
	203	+add("MGATHERTest.case_row_float_partial_4x16_in_8x16", lambda n: case_row(n, np.float32, 4, 16, 8))
	204	+add("MGATHERTest.case_row_half_partial_5x32_in_8x32", lambda n: case_row(n, np.float16, 5, 32, 8))
	205	+add("MGATHERTest.case_row_uint8_unaligned_3x32_32rows", lambda n: case_row(n, np.uint8, 3, 32, 8))
	206	+add(
	207	+ "MGATHERTest.case_row_int16_partial_3x16_in_4x16",
	208	+ lambda n: case_row(n, np.int16, 3, 16, 8, oob="clamp", idx_kind="oob"),
	209	+)
	210	+
	211	+add("MGATHERTest.case_elem_float_64_128size", lambda n: case_elem(n, np.float32, 64, 128))
	212	+add("MGATHERTest.case_elem_half_64_128size", lambda n: case_elem(n, np.float16, 64, 128))
	213	+add("MGATHERTest.case_elem_bfloat16_64_128size", lambda n: case_elem(n, np.uint16, 64, 128))
	214	+add("MGATHERTest.case_elem_int32_32_64size", lambda n: case_elem(n, np.int32, 32, 64))
	215	+add("MGATHERTest.case_elem_uint32_32_64size", lambda n: case_elem(n, np.uint32, 32, 64))
	216	+add("MGATHERTest.case_elem_int16_32_64size", lambda n: case_elem(n, np.int16, 32, 64))
	217	+add("MGATHERTest.case_elem_uint16_32_64size", lambda n: case_elem(n, np.uint16, 32, 64))
	218	+add("MGATHERTest.case_elem_int8_64_128size", lambda n: case_elem(n, np.int8, 64, 128))
	219	+add("MGATHERTest.case_elem_uint8_64_128size", lambda n: case_elem(n, np.uint8, 64, 128))
	220	+add(
	221	+ "MGATHERTest.case_elem_float_clamp_32_16size",
	222	+ lambda n: case_elem(n, np.float32, 32, 16, oob="clamp", idx_kind="oob"),
	223	+)
	224	+add("MGATHERTest.case_elem_int32_wrap_32_16size", lambda n: case_elem(n, np.int32, 32, 16, oob="wrap", idx_kind="oob"))
	225	+add("MGATHERTest.case_elem_half_zero_32_16size", lambda n: case_elem(n, np.float16, 32, 16, oob="zero", idx_kind="oob"))
	226	+
	227	+add("MGATHERTest.case_elem2d_float_8x32_256size", lambda n: case_elem2d(n, np.float32, 8, 32, 256))
	228	+add("MGATHERTest.case_elem2d_int32_8x16_256size", lambda n: case_elem2d(n, np.int32, 8, 16, 256))
	229	+add("MGATHERTest.case_elem2d_half_4x32_256size", lambda n: case_elem2d(n, np.float16, 4, 32, 256))
	230	+add("MGATHERTest.case_elem2d_bfloat16_4x32_256size", lambda n: case_elem2d(n, np.uint16, 4, 32, 256))
	231	+add("MGATHERTest.case_elem2d_uint8_4x64_256size", lambda n: case_elem2d(n, np.uint8, 4, 64, 256))
	232	+add("MGATHERTest.case_elem2d_int8_4x64_256size", lambda n: case_elem2d(n, np.int8, 4, 64, 256))
	233	+add("MGATHERTest.case_elem2d_int16_4x32_256size", lambda n: case_elem2d(n, np.int16, 4, 32, 256))
	234	+add("MGATHERTest.case_elem2d_uint16_4x32_256size", lambda n: case_elem2d(n, np.uint16, 4, 32, 256))
	235	+add("MGATHERTest.case_elem2d_uint32_8x16_256size", lambda n: case_elem2d(n, np.uint32, 8, 16, 256))
	236	+add(
	237	+ "MGATHERTest.case_elem2d_float_wrap_4x16_64size",
	238	+ lambda n: case_elem2d(n, np.float32, 4, 16, 64, oob="wrap", idx_kind="oob"),
	239	+)
	240	+add(
	241	+ "MGATHERTest.case_elem2d_int32_clamp_4x8_32size",
	242	+ lambda n: case_elem2d(n, np.int32, 4, 8, 32, oob="clamp", idx_kind="oob"),
	243	+)
	244	+add(
	245	+ "MGATHERTest.case_elem2d_half_zero_4x32_64size",
	246	+ lambda n: case_elem2d(n, np.float16, 4, 32, 64, oob="zero", idx_kind="oob"),
	247	+)
	248	+
	249	+add("MGATHERTest.case_elem2d_int32_unaligned_3x3_in_3x8_64size", lambda n: case_elem2d(n, np.int32, 3, 3, 64))
	250	+add("MGATHERTest.case_elem2d_float_unaligned_5x5_in_5x8_64size", lambda n: case_elem2d(n, np.float32, 5, 5, 64))
	251	+add("MGATHERTest.case_elem2d_half_unaligned_3x9_in_3x16_64size", lambda n: case_elem2d(n, np.float16, 3, 9, 64))
	252	+add("MGATHERTest.case_elem2d_int8_unaligned_3x17_in_3x32_64size", lambda n: case_elem2d(n, np.int8, 3, 17, 64))
	253	+
	254	+add("MGATHERTest.case_elem_scalar_float_1x1_in_1x8_8size", lambda n: case_elem2d(n, np.float32, 1, 1, 8))
	255	+add("MGATHERTest.case_elem_scalar_int32_1x1_in_1x8_8size", lambda n: case_elem2d(n, np.int32, 1, 1, 8))
	256	+add("MGATHERTest.case_elem_scalar_half_1x1_in_1x16_16size", lambda n: case_elem2d(n, np.float16, 1, 1, 16))
	257	+
	258	+add("MGATHERTest.case_elem2d_dyn_float_4x8_64size", lambda n: case_elem2d(n, np.float32, 4, 8, 64))
	259	+add("MGATHERTest.case_elem2d_dyn_int32_3x3_in_3x8_64size", lambda n: case_elem2d(n, np.int32, 3, 3, 64))
	260	+add("MGATHERTest.case_row_dyn_int32_3x16_8rows", lambda n: case_row(n, np.int32, 3, 16, 8))
	261	+add("MGATHERTest.case_row_dyn_half_4x32_16rows", lambda n: case_row(n, np.float16, 4, 32, 16))
	262	+
	263	+add("MGATHERTest.case_row_nz_float_16x16_2blk", lambda n: case_row_nz(n, np.float32, 16, 16, 2, 2, 8))
	264	+add("MGATHERTest.case_row_nz_half_32x16_2blk", lambda n: case_row_nz(n, np.float16, 32, 16, 2, 1, 16))
	265	+add("MGATHERTest.case_row_nz_int32_16x16_2blk", lambda n: case_row_nz(n, np.int32, 16, 16, 2, 2, 8))
	266	+add("MGATHERTest.case_row_nz_int16_32x16_1blk", lambda n: case_row_nz(n, np.int16, 32, 16, 2, 1, 16))
	267	+add("MGATHERTest.case_row_nz_int8_16x32_1blk", lambda n: case_row_nz(n, np.int8, 16, 32, 2, 1, 32))
	268	+add(
	269	+ "MGATHERTest.case_row_nz_float_clamp_16x8_1blk",
	270	+ lambda n: case_row_nz(n, np.float32, 16, 8, 2, 1, 8, oob="clamp", idx_kind="oob"),
	271	+)
	272	+add(
	273	+ "MGATHERTest.case_row_nz_half_zero_16x16_2blk",
	274	+ lambda n: case_row_nz(n, np.float16, 16, 16, 2, 1, 16, oob="zero", idx_kind="oob"),
	275	+)
	276	+
	277	+add("MGATHERTest.case_elem2d_nz_float_16x16_2blk", lambda n: case_elem2d_nz(n, np.float32, 16, 16, 2, 2, 8))
	278	+add("MGATHERTest.case_elem2d_nz_half_16x16_1blk", lambda n: case_elem2d_nz(n, np.float16, 16, 16, 2, 1, 16))
	279	+add("MGATHERTest.case_elem2d_nz_int32_16x8_1blk", lambda n: case_elem2d_nz(n, np.int32, 16, 8, 2, 1, 8))
	280	+add(
	281	+ "MGATHERTest.case_elem2d_nz_half_zero_16x16_1blk",
	282	+ lambda n: case_elem2d_nz(n, np.float16, 16, 16, 2, 1, 16, oob="zero", idx_kind="oob"),
	283	+)
	284	+
	285	+
	286	+if __name__ == "__main__":
	287	+ for name, fn in CASES:
	288	+ if not os.path.exists(name):
	289	+ os.makedirs(name)
	290	+ original_dir = os.getcwd()
	291	+ os.chdir(name)
	292	+ table, idx, golden = fn(name)
	293	+ table.tofile("table.bin")
	294	+ idx.astype(np.int32).tofile("indices.bin")
	295	+ golden.tofile("golden.bin")
	296	+ os.chdir(original_dir)
	297	+ print(f"Generated {name}")
	298	+ print("All MGATHER A2/A3 test data generated successfully")

Atests/npu/a2a3/src/st/testcase/mgather/main.cpp+277-0

@@ -0,0 +1,277 @@
	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	+#include "test_common.h"
	12	+#include "acl/acl.h"
	13	+#include <gtest/gtest.h>
	14	+
	15	+using namespace std;
	16	+using namespace PtoTestCommon;
	17	+
	18	+class MGATHERTest : public testing::Test {
	19	+protected:
	20	+ void SetUp() override
	21	+ {}
	22	+ void TearDown() override
	23	+ {}
	24	+};
	25	+
	26	+static std::string GetGoldenDir()
	27	+{
	28	+ const testing::TestInfo *testInfo = testing::UnitTest::GetInstance()->current_test_info();
	29	+ return std::string("../") + testInfo->test_suite_name() + "." + testInfo->name();
	30	+}
	31	+
	32	+template <typename T, typename TIdx, typename Launcher>
	33	+void run_mgather_test(size_t tableCount, size_t idxCount, size_t outCount, Launcher launcher)
	34	+{
	35	+ size_t tableByteSize = tableCount * sizeof(T);
	36	+ size_t idxByteSize = idxCount * sizeof(TIdx);
	37	+ size_t outByteSize = outCount * sizeof(T);
	38	+
	39	+ aclInit(nullptr);
	40	+ aclrtSetDevice(0);
	41	+ aclrtStream stream;
	42	+ aclrtCreateStream(&stream);
	43	+
	44	+ T tableHost, outHost;
	45	+ TIdx *idxHost;
	46	+ T tableDevice, outDevice;
	47	+ TIdx *idxDevice;
	48	+
	49	+ aclrtMallocHost((void **)(&tableHost), tableByteSize);
	50	+ aclrtMallocHost((void **)(&idxHost), idxByteSize);
	51	+ aclrtMallocHost((void **)(&outHost), outByteSize);
	52	+
	53	+ aclrtMalloc((void **)&tableDevice, tableByteSize, ACL_MEM_MALLOC_HUGE_FIRST);
	54	+ aclrtMalloc((void **)&idxDevice, idxByteSize, ACL_MEM_MALLOC_HUGE_FIRST);
	55	+ aclrtMalloc((void **)&outDevice, outByteSize, ACL_MEM_MALLOC_HUGE_FIRST);
	56	+
	57	+ ReadFile(GetGoldenDir() + "/table.bin", tableByteSize, tableHost, tableByteSize);
	58	+ ReadFile(GetGoldenDir() + "/indices.bin", idxByteSize, idxHost, idxByteSize);
	59	+
	60	+ aclrtMemcpy(tableDevice, tableByteSize, tableHost, tableByteSize, ACL_MEMCPY_HOST_TO_DEVICE);
	61	+ aclrtMemcpy(idxDevice, idxByteSize, idxHost, idxByteSize, ACL_MEMCPY_HOST_TO_DEVICE);
	62	+
	63	+ aclrtMemset(outDevice, outByteSize, 0, outByteSize);
	64	+
	65	+ launcher(outDevice, tableDevice, idxDevice, stream);
	66	+
	67	+ aclrtSynchronizeStream(stream);
	68	+ aclrtMemcpy(outHost, outByteSize, outDevice, outByteSize, ACL_MEMCPY_DEVICE_TO_HOST);
	69	+
	70	+ WriteFile(GetGoldenDir() + "/output.bin", outHost, outByteSize);
	71	+
	72	+ aclrtFree(tableDevice);
	73	+ aclrtFree(idxDevice);
	74	+ aclrtFree(outDevice);
	75	+
	76	+ aclrtFreeHost(tableHost);
	77	+ aclrtFreeHost(idxHost);
	78	+ aclrtFreeHost(outHost);
	79	+ aclrtDestroyStream(stream);
	80	+ aclrtResetDevice(0);
	81	+ aclFinalize();
	82	+
	83	+ std::vector<T> golden(outCount);
	84	+ std::vector<T> devFinal(outCount);
	85	+ ReadFile(GetGoldenDir() + "/golden.bin", outByteSize, golden.data(), outByteSize);
	86	+ ReadFile(GetGoldenDir() + "/output.bin", outByteSize, devFinal.data(), outByteSize);
	87	+
	88	+ bool ret = ResultCmp<T>(golden, devFinal, 0.0f);
	89	+ EXPECT_TRUE(ret);
	90	+}
	91	+
	92	+#define DECLARE_LAUNCH(NAME, THOST, TIDX) void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream);
	93	+
	94	+DECLARE_LAUNCH(row_float_8x32_64rows, float, int32_t)
	95	+DECLARE_LAUNCH(row_half_16x64_64rows, aclFloat16, int32_t)
	96	+DECLARE_LAUNCH(row_bfloat16_16x16_64rows, uint16_t, int32_t)
	97	+DECLARE_LAUNCH(row_int32_8x16_32rows, int32_t, int32_t)
	98	+DECLARE_LAUNCH(row_uint32_8x16_32rows, uint32_t, int32_t)
	99	+DECLARE_LAUNCH(row_int16_8x16_32rows, int16_t, int32_t)
	100	+DECLARE_LAUNCH(row_uint16_8x16_32rows, uint16_t, int32_t)
	101	+DECLARE_LAUNCH(row_int8_8x32_32rows, int8_t, int32_t)
	102	+DECLARE_LAUNCH(row_uint8_8x32_32rows, uint8_t, int32_t)
	103	+DECLARE_LAUNCH(row_float_clamp_8x32_8rows, float, int32_t)
	104	+DECLARE_LAUNCH(row_int32_wrap_8x16_8rows, int32_t, int32_t)
	105	+DECLARE_LAUNCH(row_half_zero_8x32_8rows, aclFloat16, int32_t)
	106	+
	107	+DECLARE_LAUNCH(row_int32_unaligned_3x8_8rows, int32_t, int32_t)
	108	+DECLARE_LAUNCH(row_float_partial_4x16_in_8x16, float, int32_t)
	109	+DECLARE_LAUNCH(row_half_partial_5x32_in_8x32, aclFloat16, int32_t)
	110	+DECLARE_LAUNCH(row_uint8_unaligned_3x32_32rows, uint8_t, int32_t)
	111	+DECLARE_LAUNCH(row_int16_partial_3x16_in_4x16, int16_t, int32_t)
	112	+
	113	+DECLARE_LAUNCH(elem_float_64_128size, float, int32_t)
	114	+DECLARE_LAUNCH(elem_half_64_128size, aclFloat16, int32_t)
	115	+DECLARE_LAUNCH(elem_bfloat16_64_128size, uint16_t, int32_t)
	116	+DECLARE_LAUNCH(elem_int32_32_64size, int32_t, int32_t)
	117	+DECLARE_LAUNCH(elem_uint32_32_64size, uint32_t, int32_t)
	118	+DECLARE_LAUNCH(elem_int16_32_64size, int16_t, int32_t)
	119	+DECLARE_LAUNCH(elem_uint16_32_64size, uint16_t, int32_t)
	120	+DECLARE_LAUNCH(elem_int8_64_128size, int8_t, int32_t)
	121	+DECLARE_LAUNCH(elem_uint8_64_128size, uint8_t, int32_t)
	122	+DECLARE_LAUNCH(elem_float_clamp_32_16size, float, int32_t)
	123	+DECLARE_LAUNCH(elem_int32_wrap_32_16size, int32_t, int32_t)
	124	+DECLARE_LAUNCH(elem_half_zero_32_16size, aclFloat16, int32_t)
	125	+
	126	+DECLARE_LAUNCH(elem2d_float_8x32_256size, float, int32_t)
	127	+DECLARE_LAUNCH(elem2d_int32_8x16_256size, int32_t, int32_t)
	128	+DECLARE_LAUNCH(elem2d_half_4x32_256size, aclFloat16, int32_t)
	129	+DECLARE_LAUNCH(elem2d_bfloat16_4x32_256size, uint16_t, int32_t)
	130	+DECLARE_LAUNCH(elem2d_uint8_4x64_256size, uint8_t, int32_t)
	131	+DECLARE_LAUNCH(elem2d_int8_4x64_256size, int8_t, int32_t)
	132	+DECLARE_LAUNCH(elem2d_int16_4x32_256size, int16_t, int32_t)
	133	+DECLARE_LAUNCH(elem2d_uint16_4x32_256size, uint16_t, int32_t)
	134	+DECLARE_LAUNCH(elem2d_uint32_8x16_256size, uint32_t, int32_t)
	135	+DECLARE_LAUNCH(elem2d_float_wrap_4x16_64size, float, int32_t)
	136	+DECLARE_LAUNCH(elem2d_int32_clamp_4x8_32size, int32_t, int32_t)
	137	+DECLARE_LAUNCH(elem2d_half_zero_4x32_64size, aclFloat16, int32_t)
	138	+
	139	+DECLARE_LAUNCH(elem2d_int32_unaligned_3x3_in_3x8_64size, int32_t, int32_t)
	140	+DECLARE_LAUNCH(elem2d_float_unaligned_5x5_in_5x8_64size, float, int32_t)
	141	+DECLARE_LAUNCH(elem2d_half_unaligned_3x9_in_3x16_64size, aclFloat16, int32_t)
	142	+DECLARE_LAUNCH(elem2d_int8_unaligned_3x17_in_3x32_64size, int8_t, int32_t)
	143	+
	144	+DECLARE_LAUNCH(elem_scalar_float_1x1_in_1x8_8size, float, int32_t)
	145	+DECLARE_LAUNCH(elem_scalar_int32_1x1_in_1x8_8size, int32_t, int32_t)
	146	+DECLARE_LAUNCH(elem_scalar_half_1x1_in_1x16_16size, aclFloat16, int32_t)
	147	+
	148	+DECLARE_LAUNCH(elem2d_dyn_float_4x8_64size, float, int32_t)
	149	+DECLARE_LAUNCH(elem2d_dyn_int32_3x3_in_3x8_64size, int32_t, int32_t)
	150	+DECLARE_LAUNCH(row_dyn_int32_3x16_8rows, int32_t, int32_t)
	151	+DECLARE_LAUNCH(row_dyn_half_4x32_16rows, aclFloat16, int32_t)
	152	+
	153	+DECLARE_LAUNCH(row_nz_float_16x16_2blk, float, int32_t)
	154	+DECLARE_LAUNCH(row_nz_half_32x16_2blk, aclFloat16, int32_t)
	155	+DECLARE_LAUNCH(row_nz_int32_16x16_2blk, int32_t, int32_t)
	156	+DECLARE_LAUNCH(row_nz_int16_32x16_1blk, int16_t, int32_t)
	157	+DECLARE_LAUNCH(row_nz_int8_16x32_1blk, int8_t, int32_t)
	158	+DECLARE_LAUNCH(row_nz_float_clamp_16x8_1blk, float, int32_t)
	159	+DECLARE_LAUNCH(row_nz_half_zero_16x16_2blk, aclFloat16, int32_t)
	160	+DECLARE_LAUNCH(elem2d_nz_float_16x16_2blk, float, int32_t)
	161	+DECLARE_LAUNCH(elem2d_nz_half_16x16_1blk, aclFloat16, int32_t)
	162	+DECLARE_LAUNCH(elem2d_nz_int32_16x8_1blk, int32_t, int32_t)
	163	+DECLARE_LAUNCH(elem2d_nz_half_zero_16x16_1blk, aclFloat16, int32_t)
	164	+
	165	+#define ROW_TEST(NAME, THOST, TIDX, R, C, TR) \
	166	+ TEST_F(MGATHERTest, case_##NAME) \
	167	+ { \
	168	+ run_mgather_test<THOST, TIDX>((size_t)TR * C, (size_t)R, (size_t)R * C, Launch_##NAME); \
	169	+ }
	170	+
	171	+#define ELEM_TEST(NAME, THOST, TIDX, N, TS) \
	172	+ TEST_F(MGATHERTest, case_##NAME) \
	173	+ { \
	174	+ run_mgather_test<THOST, TIDX>((size_t)TS, (size_t)N, (size_t)N, Launch_##NAME); \
	175	+ }
	176	+
	177	+#define ELEM2D_TEST(NAME, THOST, TIDX, R, C, TS) \
	178	+ TEST_F(MGATHERTest, case_##NAME) \
	179	+ { \
	180	+ run_mgather_test<THOST, TIDX>((size_t)TS, (size_t)R * C, (size_t)R * C, Launch_##NAME); \
	181	+ }
	182	+
	183	+#define SCALAR_TEST(NAME, THOST, TIDX, TS) \
	184	+ TEST_F(MGATHERTest, case_##NAME) \
	185	+ { \
	186	+ run_mgather_test<THOST, TIDX>((size_t)TS, (size_t)1, (size_t)1, Launch_##NAME); \
	187	+ }
	188	+
	189	+ROW_TEST(row_float_8x32_64rows, float, int32_t, 8, 32, 64)
	190	+ROW_TEST(row_half_16x64_64rows, aclFloat16, int32_t, 16, 64, 64)
	191	+ROW_TEST(row_bfloat16_16x16_64rows, uint16_t, int32_t, 16, 16, 64)
	192	+ROW_TEST(row_int32_8x16_32rows, int32_t, int32_t, 8, 16, 32)
	193	+ROW_TEST(row_uint32_8x16_32rows, uint32_t, int32_t, 8, 16, 32)
	194	+ROW_TEST(row_int16_8x16_32rows, int16_t, int32_t, 8, 16, 32)
	195	+ROW_TEST(row_uint16_8x16_32rows, uint16_t, int32_t, 8, 16, 32)
	196	+ROW_TEST(row_int8_8x32_32rows, int8_t, int32_t, 8, 32, 32)
	197	+ROW_TEST(row_uint8_8x32_32rows, uint8_t, int32_t, 8, 32, 32)
	198	+ROW_TEST(row_float_clamp_8x32_8rows, float, int32_t, 8, 32, 8)
	199	+ROW_TEST(row_int32_wrap_8x16_8rows, int32_t, int32_t, 8, 16, 8)
	200	+ROW_TEST(row_half_zero_8x32_8rows, aclFloat16, int32_t, 8, 32, 8)
	201	+
	202	+ROW_TEST(row_int32_unaligned_3x8_8rows, int32_t, int32_t, 3, 8, 8)
	203	+ROW_TEST(row_float_partial_4x16_in_8x16, float, int32_t, 4, 16, 8)
	204	+ROW_TEST(row_half_partial_5x32_in_8x32, aclFloat16, int32_t, 5, 32, 8)
	205	+ROW_TEST(row_uint8_unaligned_3x32_32rows, uint8_t, int32_t, 3, 32, 8)
	206	+ROW_TEST(row_int16_partial_3x16_in_4x16, int16_t, int32_t, 3, 16, 8)
	207	+
	208	+ELEM_TEST(elem_float_64_128size, float, int32_t, 64, 128)
	209	+ELEM_TEST(elem_half_64_128size, aclFloat16, int32_t, 64, 128)
	210	+ELEM_TEST(elem_bfloat16_64_128size, uint16_t, int32_t, 64, 128)
	211	+ELEM_TEST(elem_int32_32_64size, int32_t, int32_t, 32, 64)
	212	+ELEM_TEST(elem_uint32_32_64size, uint32_t, int32_t, 32, 64)
	213	+ELEM_TEST(elem_int16_32_64size, int16_t, int32_t, 32, 64)
	214	+ELEM_TEST(elem_uint16_32_64size, uint16_t, int32_t, 32, 64)
	215	+ELEM_TEST(elem_int8_64_128size, int8_t, int32_t, 64, 128)
	216	+ELEM_TEST(elem_uint8_64_128size, uint8_t, int32_t, 64, 128)
	217	+ELEM_TEST(elem_float_clamp_32_16size, float, int32_t, 32, 16)
	218	+ELEM_TEST(elem_int32_wrap_32_16size, int32_t, int32_t, 32, 16)
	219	+ELEM_TEST(elem_half_zero_32_16size, aclFloat16, int32_t, 32, 16)
	220	+
	221	+ELEM2D_TEST(elem2d_float_8x32_256size, float, int32_t, 8, 32, 256)
	222	+ELEM2D_TEST(elem2d_int32_8x16_256size, int32_t, int32_t, 8, 16, 256)
	223	+ELEM2D_TEST(elem2d_half_4x32_256size, aclFloat16, int32_t, 4, 32, 256)
	224	+ELEM2D_TEST(elem2d_bfloat16_4x32_256size, uint16_t, int32_t, 4, 32, 256)
	225	+ELEM2D_TEST(elem2d_uint8_4x64_256size, uint8_t, int32_t, 4, 64, 256)
	226	+ELEM2D_TEST(elem2d_int8_4x64_256size, int8_t, int32_t, 4, 64, 256)
	227	+ELEM2D_TEST(elem2d_int16_4x32_256size, int16_t, int32_t, 4, 32, 256)
	228	+ELEM2D_TEST(elem2d_uint16_4x32_256size, uint16_t, int32_t, 4, 32, 256)
	229	+ELEM2D_TEST(elem2d_uint32_8x16_256size, uint32_t, int32_t, 8, 16, 256)
	230	+ELEM2D_TEST(elem2d_float_wrap_4x16_64size, float, int32_t, 4, 16, 64)
	231	+ELEM2D_TEST(elem2d_int32_clamp_4x8_32size, int32_t, int32_t, 4, 8, 32)
	232	+ELEM2D_TEST(elem2d_half_zero_4x32_64size, aclFloat16, int32_t, 4, 32, 64)
	233	+
	234	+ELEM2D_TEST(elem2d_int32_unaligned_3x3_in_3x8_64size, int32_t, int32_t, 3, 3, 64)
	235	+ELEM2D_TEST(elem2d_float_unaligned_5x5_in_5x8_64size, float, int32_t, 5, 5, 64)
	236	+ELEM2D_TEST(elem2d_half_unaligned_3x9_in_3x16_64size, aclFloat16, int32_t, 3, 9, 64)
	237	+ELEM2D_TEST(elem2d_int8_unaligned_3x17_in_3x32_64size, int8_t, int32_t, 3, 17, 64)
	238	+
	239	+SCALAR_TEST(elem_scalar_float_1x1_in_1x8_8size, float, int32_t, 8)
	240	+SCALAR_TEST(elem_scalar_int32_1x1_in_1x8_8size, int32_t, int32_t, 8)
	241	+SCALAR_TEST(elem_scalar_half_1x1_in_1x16_16size, aclFloat16, int32_t, 16)
	242	+
	243	+#define ELEM2D_DYN_TEST(NAME, THOST, TIDX, RVR, RVC, RTS) \
	244	+ TEST_F(MGATHERTest, case_##NAME) \
	245	+ { \
	246	+ run_mgather_test<THOST, TIDX>((size_t)RTS, (size_t)RVR * RVC, (size_t)RVR * RVC, Launch_##NAME); \
	247	+ }
	248	+
	249	+ELEM2D_DYN_TEST(elem2d_dyn_float_4x8_64size, float, int32_t, 4, 8, 64)
	250	+ELEM2D_DYN_TEST(elem2d_dyn_int32_3x3_in_3x8_64size, int32_t, int32_t, 3, 3, 64)
	251	+ROW_TEST(row_dyn_int32_3x16_8rows, int32_t, int32_t, 3, 16, 8)
	252	+ROW_TEST(row_dyn_half_4x32_16rows, aclFloat16, int32_t, 4, 32, 16)
	253	+
	254	+#define ROW_NZ_TEST(NAME, THOST, TIDX, R, C, BR, BC, C0) \
	255	+ TEST_F(MGATHERTest, case_##NAME) \
	256	+ { \
	257	+ run_mgather_test<THOST, TIDX>((size_t)BR * 16 * BC * C0, (size_t)R, (size_t)R * C, Launch_##NAME); \
	258	+ }
	259	+
	260	+#define ELEM2D_NZ_TEST(NAME, THOST, TIDX, R, C, BR, BC, C0) \
	261	+ TEST_F(MGATHERTest, case_##NAME) \
	262	+ { \
	263	+ run_mgather_test<THOST, TIDX>((size_t)BR * 16 * BC * C0, (size_t)R * C, (size_t)R * C, Launch_##NAME); \
	264	+ }
	265	+
	266	+ROW_NZ_TEST(row_nz_float_16x16_2blk, float, int32_t, 16, 16, 2, 2, 8)
	267	+ROW_NZ_TEST(row_nz_half_32x16_2blk, aclFloat16, int32_t, 32, 16, 2, 1, 16)
	268	+ROW_NZ_TEST(row_nz_int32_16x16_2blk, int32_t, int32_t, 16, 16, 2, 2, 8)
	269	+ROW_NZ_TEST(row_nz_int16_32x16_1blk, int16_t, int32_t, 32, 16, 2, 1, 16)
	270	+ROW_NZ_TEST(row_nz_int8_16x32_1blk, int8_t, int32_t, 16, 32, 2, 1, 32)
	271	+ROW_NZ_TEST(row_nz_float_clamp_16x8_1blk, float, int32_t, 16, 8, 2, 1, 8)
	272	+ROW_NZ_TEST(row_nz_half_zero_16x16_2blk, aclFloat16, int32_t, 16, 16, 2, 1, 16)
	273	+
	274	+ELEM2D_NZ_TEST(elem2d_nz_float_16x16_2blk, float, int32_t, 16, 16, 2, 2, 8)
	275	+ELEM2D_NZ_TEST(elem2d_nz_half_16x16_1blk, aclFloat16, int32_t, 16, 16, 2, 1, 16)
	276	+ELEM2D_NZ_TEST(elem2d_nz_int32_16x8_1blk, int32_t, int32_t, 16, 8, 2, 1, 8)
	277	+ELEM2D_NZ_TEST(elem2d_nz_half_zero_16x16_1blk, aclFloat16, int32_t, 16, 16, 2, 1, 16)

Atests/npu/a2a3/src/st/testcase/mgather/mgather_kernel.cpp+574-0

@@ -0,0 +1,574 @@
	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	+#include <pto/pto-inst.hpp>
	12	+#include <pto/common/pto_tile.hpp>
	13	+#include <pto/common/constants.hpp>
	14	+#include <pto/npu/a2a3/MGather.hpp>
	15	+#include "acl/acl.h"
	16	+
	17	+using namespace pto;
	18	+
	19	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kDstRows, uint32_t kDstCols, uint32_t kTableRows>
	20	+inline AICORE void runRow(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	21	+{
	22	+ using TableShape = pto::Shape<1, 1, 1, kTableRows, kDstCols>;
	23	+ using TableStride = pto::Stride<1, 1, 1, kDstCols, 1>;
	24	+ using IdxShape = pto::Shape<1, 1, 1, 1, kDstRows>;
	25	+ using IdxStride = pto::Stride<1, 1, 1, kDstRows, 1>;
	26	+ using OutShape = pto::Shape<1, 1, 1, kDstRows, kDstCols>;
	27	+ using OutStride = pto::Stride<1, 1, 1, kDstCols, 1>;
	28	+
	29	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table);
	30	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	31	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out);
	32	+
	33	+ using DstTile = Tile<TileType::Vec, T, kDstRows, kDstCols, BLayout::RowMajor, kDstRows, kDstCols>;
	34	+ using IdxTile = Tile<TileType::Vec, TIdx, 1, kDstRows, BLayout::RowMajor, 1, kDstRows>;
	35	+
	36	+ DstTile dstTile;
	37	+ IdxTile idxTile;
	38	+
	39	+ constexpr uint32_t idxBytes = ((1u * kDstRows * sizeof(TIdx) + 31u) / 32u) * 32u;
	40	+ TASSIGN(idxTile, 0x0);
	41	+ TASSIGN(dstTile, idxBytes);
	42	+
	43	+ TLOAD(idxTile, idxGlobal);
	44	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	45	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	46	+
	47	+ MGATHER<Coalesce::Row, Oob>(dstTile, tableGlobal, idxTile);
	48	+
	49	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	50	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	51	+
	52	+ TSTORE(outGlobal, dstTile);
	53	+}
	54	+
	55	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kValidRows, uint32_t kPadRows, uint32_t kPadCols,
	56	+ uint32_t kPadIdxCols, uint32_t kTableRows>
	57	+inline AICORE void runRowPadded(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	58	+{
	59	+ using TableShape = pto::Shape<1, 1, 1, kTableRows, kPadCols>;
	60	+ using TableStride = pto::Stride<1, 1, 1, kPadCols, 1>;
	61	+ using IdxShape = pto::Shape<1, 1, 1, 1, kValidRows>;
	62	+ using IdxStride = pto::Stride<1, 1, 1, kValidRows, 1>;
	63	+ using OutShape = pto::Shape<1, 1, 1, kValidRows, kPadCols>;
	64	+ using OutStride = pto::Stride<1, 1, 1, kPadCols, 1>;
	65	+
	66	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table);
	67	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	68	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out);
	69	+
	70	+ using DstTile = Tile<TileType::Vec, T, kPadRows, kPadCols, BLayout::RowMajor, kValidRows, kPadCols>;
	71	+ using IdxTile = Tile<TileType::Vec, TIdx, 1, kPadIdxCols, BLayout::RowMajor, 1, kValidRows>;
	72	+
	73	+ DstTile dstTile;
	74	+ IdxTile idxTile;
	75	+
	76	+ constexpr uint32_t idxBytes = ((1u * kPadIdxCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	77	+ TASSIGN(idxTile, 0x0);
	78	+ TASSIGN(dstTile, idxBytes);
	79	+
	80	+ TLOAD(idxTile, idxGlobal);
	81	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	82	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	83	+
	84	+ MGATHER<Coalesce::Row, Oob>(dstTile, tableGlobal, idxTile);
	85	+
	86	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	87	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	88	+
	89	+ TSTORE(outGlobal, dstTile);
	90	+}
	91	+
	92	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kDstCols, uint32_t kTableSize>
	93	+inline AICORE void runElem(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	94	+{
	95	+ using TableShape = pto::Shape<1, 1, 1, 1, kTableSize>;
	96	+ using TableStride = pto::Stride<1, 1, 1, kTableSize, 1>;
	97	+ using IdxShape = pto::Shape<1, 1, 1, 1, kDstCols>;
	98	+ using IdxStride = pto::Stride<1, 1, 1, kDstCols, 1>;
	99	+ using OutShape = pto::Shape<1, 1, 1, 1, kDstCols>;
	100	+ using OutStride = pto::Stride<1, 1, 1, kDstCols, 1>;
	101	+
	102	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table);
	103	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	104	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out);
	105	+
	106	+ using DstTile = Tile<TileType::Vec, T, 1, kDstCols, BLayout::RowMajor, 1, kDstCols>;
	107	+ using IdxTile = Tile<TileType::Vec, TIdx, 1, kDstCols, BLayout::RowMajor, 1, kDstCols>;
	108	+
	109	+ DstTile dstTile;
	110	+ IdxTile idxTile;
	111	+
	112	+ constexpr uint32_t idxBytes = ((1u * kDstCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	113	+ TASSIGN(idxTile, 0x0);
	114	+ TASSIGN(dstTile, idxBytes);
	115	+
	116	+ TLOAD(idxTile, idxGlobal);
	117	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	118	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	119	+
	120	+ MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxTile);
	121	+
	122	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	123	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	124	+
	125	+ TSTORE(outGlobal, dstTile);
	126	+}
	127	+
	128	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kDstRows, uint32_t kDstCols, uint32_t kTableSize>
	129	+inline AICORE void runElem2D(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	130	+{
	131	+ using TableShape = pto::Shape<1, 1, 1, 1, kTableSize>;
	132	+ using TableStride = pto::Stride<1, 1, 1, kTableSize, 1>;
	133	+ using IdxShape = pto::Shape<1, 1, 1, kDstRows, kDstCols>;
	134	+ using IdxStride = pto::Stride<1, 1, 1, kDstCols, 1>;
	135	+ using OutShape = pto::Shape<1, 1, 1, kDstRows, kDstCols>;
	136	+ using OutStride = pto::Stride<1, 1, 1, kDstCols, 1>;
	137	+
	138	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table);
	139	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	140	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out);
	141	+
	142	+ using DstTile = Tile<TileType::Vec, T, kDstRows, kDstCols, BLayout::RowMajor, kDstRows, kDstCols>;
	143	+ using IdxTile = Tile<TileType::Vec, TIdx, kDstRows, kDstCols, BLayout::RowMajor, kDstRows, kDstCols>;
	144	+
	145	+ DstTile dstTile;
	146	+ IdxTile idxTile;
	147	+
	148	+ constexpr uint32_t idxBytes = ((kDstRows * kDstCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	149	+ TASSIGN(idxTile, 0x0);
	150	+ TASSIGN(dstTile, idxBytes);
	151	+
	152	+ TLOAD(idxTile, idxGlobal);
	153	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	154	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	155	+
	156	+ MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxTile);
	157	+
	158	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	159	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	160	+
	161	+ TSTORE(outGlobal, dstTile);
	162	+}
	163	+
	164	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kValidRows, uint32_t kValidCols, uint32_t kPadRows,
	165	+ uint32_t kPadCols, uint32_t kTableSize>
	166	+inline AICORE void runElem2DPadded(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	167	+{
	168	+ using TableShape = pto::Shape<1, 1, 1, 1, kTableSize>;
	169	+ using TableStride = pto::Stride<1, 1, 1, kTableSize, 1>;
	170	+ using IdxShape = pto::Shape<1, 1, 1, kValidRows, kValidCols>;
	171	+ using IdxStride = pto::Stride<1, 1, 1, kValidCols, 1>;
	172	+ using OutShape = pto::Shape<1, 1, 1, kValidRows, kValidCols>;
	173	+ using OutStride = pto::Stride<1, 1, 1, kValidCols, 1>;
	174	+
	175	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table);
	176	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	177	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out);
	178	+
	179	+ using DstTile = Tile<TileType::Vec, T, kPadRows, kPadCols, BLayout::RowMajor, kValidRows, kValidCols>;
	180	+ using IdxTile = Tile<TileType::Vec, TIdx, kPadRows, kPadCols, BLayout::RowMajor, kValidRows, kValidCols>;
	181	+
	182	+ DstTile dstTile;
	183	+ IdxTile idxTile;
	184	+
	185	+ constexpr uint32_t idxBytes = ((kPadRows * kPadCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	186	+ TASSIGN(idxTile, 0x0);
	187	+ TASSIGN(dstTile, idxBytes);
	188	+
	189	+ TLOAD(idxTile, idxGlobal);
	190	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	191	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	192	+
	193	+ MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxTile);
	194	+
	195	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	196	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	197	+
	198	+ TSTORE(outGlobal, dstTile);
	199	+}
	200	+
	201	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kPadCols, uint32_t kTableSize>
	202	+inline AICORE void runElemScalar(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	203	+{
	204	+ using TableShape = pto::Shape<1, 1, 1, 1, kTableSize>;
	205	+ using TableStride = pto::Stride<1, 1, 1, kTableSize, 1>;
	206	+ using IdxShape = pto::Shape<1, 1, 1, 1, 1>;
	207	+ using IdxStride = pto::Stride<1, 1, 1, 1, 1>;
	208	+ using OutShape = pto::Shape<1, 1, 1, 1, 1>;
	209	+ using OutStride = pto::Stride<1, 1, 1, 1, 1>;
	210	+
	211	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table);
	212	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	213	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out);
	214	+
	215	+ using DstTile = Tile<TileType::Vec, T, 1, kPadCols, BLayout::RowMajor, 1, 1>;
	216	+ using IdxTile = Tile<TileType::Vec, TIdx, 1, kPadCols, BLayout::RowMajor, 1, 1>;
	217	+
	218	+ DstTile dstTile;
	219	+ IdxTile idxTile;
	220	+
	221	+ constexpr uint32_t idxBytes = ((1u * kPadCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	222	+ TASSIGN(idxTile, 0x0);
	223	+ TASSIGN(dstTile, idxBytes);
	224	+
	225	+ TLOAD(idxTile, idxGlobal);
	226	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	227	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	228	+
	229	+ MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxTile);
	230	+
	231	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	232	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	233	+
	234	+ TSTORE(outGlobal, dstTile);
	235	+}
	236	+
	237	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kPadRows, uint32_t kPadCols, int64_t kRtValidRows,
	238	+ int64_t kRtValidCols, int64_t kRtTableSize>
	239	+inline AICORE void runElem2DDyn(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	240	+{
	241	+ using TableShape = pto::Shape<1, 1, 1, -1, -1>;
	242	+ using TableStride = pto::Stride<1, 1, 1, -1, -1>;
	243	+ using IdxShape = pto::Shape<1, 1, 1, -1, -1>;
	244	+ using IdxStride = pto::Stride<1, 1, 1, -1, -1>;
	245	+ using OutShape = pto::Shape<1, 1, 1, -1, -1>;
	246	+ using OutStride = pto::Stride<1, 1, 1, -1, -1>;
	247	+
	248	+ TableShape tableShape((int64_t)1, kRtTableSize);
	249	+ TableStride tableStride(kRtTableSize, (int64_t)1);
	250	+ IdxShape idxShape(kRtValidRows, kRtValidCols);
	251	+ IdxStride idxStride(kRtValidCols, (int64_t)1);
	252	+ OutShape outShape(kRtValidRows, kRtValidCols);
	253	+ OutStride outStride(kRtValidCols, (int64_t)1);
	254	+
	255	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table, tableShape, tableStride);
	256	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices, idxShape, idxStride);
	257	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out, outShape, outStride);
	258	+
	259	+ using DstTile = Tile<TileType::Vec, T, kPadRows, kPadCols, BLayout::RowMajor, -1, -1>;
	260	+ using IdxTile = Tile<TileType::Vec, TIdx, kPadRows, kPadCols, BLayout::RowMajor, -1, -1>;
	261	+
	262	+ DstTile dstTile(static_cast<unsigned>(kRtValidRows), static_cast<unsigned>(kRtValidCols));
	263	+ IdxTile idxTile(static_cast<unsigned>(kRtValidRows), static_cast<unsigned>(kRtValidCols));
	264	+
	265	+ constexpr uint32_t idxBytes = ((kPadRows * kPadCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	266	+ TASSIGN(idxTile, 0x0);
	267	+ TASSIGN(dstTile, idxBytes);
	268	+
	269	+ TLOAD(idxTile, idxGlobal);
	270	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	271	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	272	+
	273	+ MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxTile);
	274	+
	275	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	276	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	277	+
	278	+ TSTORE(outGlobal, dstTile);
	279	+}
	280	+
	281	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kPadRows, uint32_t kPadCols, uint32_t kPadIdxCols,
	282	+ int64_t kRtValidRows, int64_t kRtValidCols, int64_t kRtTableRows>
	283	+inline AICORE void runRowDyn(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	284	+{
	285	+ using TableShape = pto::Shape<1, 1, 1, -1, -1>;
	286	+ using TableStride = pto::Stride<1, 1, 1, -1, -1>;
	287	+ using IdxShape = pto::Shape<1, 1, 1, 1, -1>;
	288	+ using IdxStride = pto::Stride<1, 1, 1, -1, -1>;
	289	+ using OutShape = pto::Shape<1, 1, 1, -1, -1>;
	290	+ using OutStride = pto::Stride<1, 1, 1, -1, -1>;
	291	+
	292	+ TableShape tableShape(kRtTableRows, kRtValidCols);
	293	+ TableStride tableStride(kRtValidCols, (int64_t)1);
	294	+ IdxShape idxShape(kRtValidRows);
	295	+ IdxStride idxStride(kRtValidRows, (int64_t)1);
	296	+ OutShape outShape(kRtValidRows, kRtValidCols);
	297	+ OutStride outStride(kRtValidCols, (int64_t)1);
	298	+
	299	+ GlobalTensor<T, TableShape, TableStride> tableGlobal(table, tableShape, tableStride);
	300	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices, idxShape, idxStride);
	301	+ GlobalTensor<T, OutShape, OutStride> outGlobal(out, outShape, outStride);
	302	+
	303	+ using DstTile = Tile<TileType::Vec, T, kPadRows, kPadCols, BLayout::RowMajor, -1, -1>;
	304	+ using IdxTile = Tile<TileType::Vec, TIdx, 1, kPadIdxCols, BLayout::RowMajor, 1, -1>;
	305	+
	306	+ DstTile dstTile(static_cast<unsigned>(kRtValidRows), static_cast<unsigned>(kRtValidCols));
	307	+ IdxTile idxTile(static_cast<unsigned>(kRtValidRows));
	308	+
	309	+ constexpr uint32_t idxBytes = ((1u * kPadIdxCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	310	+ TASSIGN(idxTile, 0x0);
	311	+ TASSIGN(dstTile, idxBytes);
	312	+
	313	+ TLOAD(idxTile, idxGlobal);
	314	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	315	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	316	+
	317	+ MGATHER<Coalesce::Row, Oob>(dstTile, tableGlobal, idxTile);
	318	+
	319	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	320	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	321	+
	322	+ TSTORE(outGlobal, dstTile);
	323	+}
	324	+
	325	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kDstRows, uint32_t kDstCols, uint32_t kBlockRows,
	326	+ uint32_t kBlockCols, uint32_t kC0>
	327	+inline AICORE void runRowNz(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	328	+{
	329	+ using TableShape = pto::Shape<1, kBlockCols, kBlockRows, 16, kC0>;
	330	+ using TableStride = pto::Stride<kBlockCols * kBlockRows * 16 * kC0, kBlockRows * 16 * kC0, 16 * kC0, kC0, 1>;
	331	+ using IdxShape = pto::Shape<1, 1, 1, 1, kDstRows>;
	332	+ using IdxStride = pto::Stride<1, 1, 1, kDstRows, 1>;
	333	+ using OutShape = pto::Shape<1, kBlockCols, kDstRows / 16, 16, kC0>;
	334	+ using OutStride = pto::Stride<kBlockCols (kDstRows / 16) 16 * kC0, (kDstRows / 16) * 16 * kC0, 16 * kC0, kC0, 1>;
	335	+
	336	+ GlobalTensor<T, TableShape, TableStride, Layout::NZ> tableGlobal(table);
	337	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	338	+ GlobalTensor<T, OutShape, OutStride, Layout::NZ> outGlobal(out);
	339	+
	340	+ using DstTile =
	341	+ Tile<TileType::Vec, T, kDstRows, kDstCols, BLayout::ColMajor, kDstRows, kDstCols, SLayout::RowMajor, 512>;
	342	+ using IdxTile = Tile<TileType::Vec, TIdx, 1, kDstRows, BLayout::RowMajor, 1, kDstRows>;
	343	+
	344	+ DstTile dstTile;
	345	+ IdxTile idxTile;
	346	+
	347	+ constexpr uint32_t idxBytes = ((1u * kDstRows * sizeof(TIdx) + 31u) / 32u) * 32u;
	348	+ TASSIGN(idxTile, 0x0);
	349	+ TASSIGN(dstTile, idxBytes);
	350	+
	351	+ TLOAD(idxTile, idxGlobal);
	352	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	353	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	354	+
	355	+ MGATHER<Coalesce::Row, Oob>(dstTile, tableGlobal, idxTile);
	356	+
	357	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	358	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	359	+
	360	+ TSTORE(outGlobal, dstTile);
	361	+}
	362	+
	363	+template <pto::GatherOOB Oob, typename T, typename TIdx, uint32_t kDstRows, uint32_t kDstCols, uint32_t kBlockRows,
	364	+ uint32_t kBlockCols, uint32_t kC0>
	365	+inline AICORE void runElem2DNz(__gm__ T out, __gm__ T table, __gm__ TIdx *indices)
	366	+{
	367	+ using TableShape = pto::Shape<1, kBlockCols, kBlockRows, 16, kC0>;
	368	+ using TableStride = pto::Stride<kBlockCols * kBlockRows * 16 * kC0, kBlockRows * 16 * kC0, 16 * kC0, kC0, 1>;
	369	+ using IdxShape = pto::Shape<1, 1, 1, kDstRows, kDstCols>;
	370	+ using IdxStride = pto::Stride<1, 1, 1, kDstCols, 1>;
	371	+ using OutShape = pto::Shape<1, kBlockCols, kDstRows / 16, 16, kC0>;
	372	+ using OutStride = pto::Stride<kBlockCols (kDstRows / 16) 16 * kC0, (kDstRows / 16) * 16 * kC0, 16 * kC0, kC0, 1>;
	373	+
	374	+ GlobalTensor<T, TableShape, TableStride, Layout::NZ> tableGlobal(table);
	375	+ GlobalTensor<TIdx, IdxShape, IdxStride> idxGlobal(indices);
	376	+ GlobalTensor<T, OutShape, OutStride, Layout::NZ> outGlobal(out);
	377	+
	378	+ using DstTile =
	379	+ Tile<TileType::Vec, T, kDstRows, kDstCols, BLayout::ColMajor, kDstRows, kDstCols, SLayout::RowMajor, 512>;
	380	+ using IdxTile = Tile<TileType::Vec, TIdx, kDstRows, kDstCols, BLayout::RowMajor, kDstRows, kDstCols>;
	381	+
	382	+ DstTile dstTile;
	383	+ IdxTile idxTile;
	384	+
	385	+ constexpr uint32_t idxBytes = ((kDstRows * kDstCols * sizeof(TIdx) + 31u) / 32u) * 32u;
	386	+ TASSIGN(idxTile, 0x0);
	387	+ TASSIGN(dstTile, idxBytes);
	388	+
	389	+ TLOAD(idxTile, idxGlobal);
	390	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	391	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
	392	+
	393	+ MGATHER<Coalesce::Elem, Oob>(dstTile, tableGlobal, idxTile);
	394	+
	395	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	396	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
	397	+
	398	+ TSTORE(outGlobal, dstTile);
	399	+}
	400	+
	401	+#define DEFINE_ROW(NAME, THOST, T, TIDX, R, C, TR, OOB) \
	402	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	403	+ { \
	404	+ runRow<pto::GatherOOB::OOB, T, TIDX, R, C, TR>(out, table, indices); \
	405	+ } \
	406	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	407	+ { \
	408	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	409	+ }
	410	+
	411	+#define DEFINE_ROW_PAD(NAME, THOST, T, TIDX, VR, PR, PC, PIC, TR, OOB) \
	412	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	413	+ { \
	414	+ runRowPadded<pto::GatherOOB::OOB, T, TIDX, VR, PR, PC, PIC, TR>(out, table, indices); \
	415	+ } \
	416	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	417	+ { \
	418	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	419	+ }
	420	+
	421	+#define DEFINE_ELEM(NAME, THOST, T, TIDX, N, TS, OOB) \
	422	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	423	+ { \
	424	+ runElem<pto::GatherOOB::OOB, T, TIDX, N, TS>(out, table, indices); \
	425	+ } \
	426	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	427	+ { \
	428	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	429	+ }
	430	+
	431	+#define DEFINE_ELEM2D(NAME, THOST, T, TIDX, R, C, TS, OOB) \
	432	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	433	+ { \
	434	+ runElem2D<pto::GatherOOB::OOB, T, TIDX, R, C, TS>(out, table, indices); \
	435	+ } \
	436	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	437	+ { \
	438	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	439	+ }
	440	+
	441	+#define DEFINE_ELEM2D_PAD(NAME, THOST, T, TIDX, VR, VC, PR, PC, TS, OOB) \
	442	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	443	+ { \
	444	+ runElem2DPadded<pto::GatherOOB::OOB, T, TIDX, VR, VC, PR, PC, TS>(out, table, indices); \
	445	+ } \
	446	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	447	+ { \
	448	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	449	+ }
	450	+
	451	+#define DEFINE_ELEM_SCALAR(NAME, THOST, T, TIDX, PC, TS, OOB) \
	452	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	453	+ { \
	454	+ runElemScalar<pto::GatherOOB::OOB, T, TIDX, PC, TS>(out, table, indices); \
	455	+ } \
	456	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	457	+ { \
	458	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	459	+ }
	460	+
	461	+#define DEFINE_ELEM2D_DYN(NAME, THOST, T, TIDX, PR, PC, RVR, RVC, RTS, OOB) \
	462	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	463	+ { \
	464	+ runElem2DDyn<pto::GatherOOB::OOB, T, TIDX, PR, PC, (int64_t)RVR, (int64_t)RVC, (int64_t)RTS>(out, table, \
	465	+ indices); \
	466	+ } \
	467	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	468	+ { \
	469	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	470	+ }
	471	+
	472	+#define DEFINE_ROW_DYN(NAME, THOST, T, TIDX, PR, PC, PIC, RVR, RVC, RTR, OOB) \
	473	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	474	+ { \
	475	+ runRowDyn<pto::GatherOOB::OOB, T, TIDX, PR, PC, PIC, (int64_t)RVR, (int64_t)RVC, (int64_t)RTR>(out, table, \
	476	+ indices); \
	477	+ } \
	478	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	479	+ { \
	480	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	481	+ }
	482	+
	483	+#define DEFINE_ROW_NZ(NAME, THOST, T, TIDX, R, C, BR, BC, C0, OOB) \
	484	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	485	+ { \
	486	+ runRowNz<pto::GatherOOB::OOB, T, TIDX, R, C, BR, BC, C0>(out, table, indices); \
	487	+ } \
	488	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	489	+ { \
	490	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	491	+ }
	492	+
	493	+#define DEFINE_ELEM2D_NZ(NAME, THOST, T, TIDX, R, C, BR, BC, C0, OOB) \
	494	+ extern "C" __global__ AICORE void runMGATHER_##NAME(__gm__ T out, __gm__ T table, __gm__ TIDX *indices) \
	495	+ { \
	496	+ runElem2DNz<pto::GatherOOB::OOB, T, TIDX, R, C, BR, BC, C0>(out, table, indices); \
	497	+ } \
	498	+ void Launch_##NAME(THOST out, THOST table, TIDX indices, void stream) \
	499	+ { \
	500	+ runMGATHER_##NAME<<<1, nullptr, stream>>>(reinterpret_cast<T >(out), reinterpret_cast<T >(table), indices); \
	501	+ }
	502	+
	503	+DEFINE_ROW(row_float_8x32_64rows, float, float, int32_t, 8, 32, 64, Undefined)
	504	+DEFINE_ROW(row_half_16x64_64rows, aclFloat16, half, int32_t, 16, 64, 64, Undefined)
	505	+DEFINE_ROW(row_bfloat16_16x16_64rows, uint16_t, bfloat16_t, int32_t, 16, 16, 64, Undefined)
	506	+DEFINE_ROW(row_int32_8x16_32rows, int32_t, int32_t, int32_t, 8, 16, 32, Undefined)
	507	+DEFINE_ROW(row_uint32_8x16_32rows, uint32_t, uint32_t, int32_t, 8, 16, 32, Undefined)
	508	+DEFINE_ROW(row_int16_8x16_32rows, int16_t, int16_t, int32_t, 8, 16, 32, Undefined)
	509	+DEFINE_ROW(row_uint16_8x16_32rows, uint16_t, uint16_t, int32_t, 8, 16, 32, Undefined)
	510	+DEFINE_ROW(row_int8_8x32_32rows, int8_t, int8_t, int32_t, 8, 32, 32, Undefined)
	511	+DEFINE_ROW(row_uint8_8x32_32rows, uint8_t, uint8_t, int32_t, 8, 32, 32, Undefined)
	512	+DEFINE_ROW(row_float_clamp_8x32_8rows, float, float, int32_t, 8, 32, 8, Clamp)
	513	+DEFINE_ROW(row_int32_wrap_8x16_8rows, int32_t, int32_t, int32_t, 8, 16, 8, Wrap)
	514	+DEFINE_ROW(row_half_zero_8x32_8rows, aclFloat16, half, int32_t, 8, 32, 8, Zero)
	515	+
	516	+DEFINE_ROW_PAD(row_int32_unaligned_3x8_8rows, int32_t, int32_t, int32_t, 3, 3, 8, 8, 8, Undefined)
	517	+DEFINE_ROW_PAD(row_float_partial_4x16_in_8x16, float, float, int32_t, 4, 8, 16, 8, 8, Undefined)
	518	+DEFINE_ROW_PAD(row_half_partial_5x32_in_8x32, aclFloat16, half, int32_t, 5, 8, 32, 8, 8, Undefined)
	519	+DEFINE_ROW_PAD(row_uint8_unaligned_3x32_32rows, uint8_t, uint8_t, int32_t, 3, 3, 32, 8, 8, Undefined)
	520	+DEFINE_ROW_PAD(row_int16_partial_3x16_in_4x16, int16_t, int16_t, int32_t, 3, 4, 16, 8, 8, Clamp)
	521	+
	522	+DEFINE_ELEM(elem_float_64_128size, float, float, int32_t, 64, 128, Undefined)
	523	+DEFINE_ELEM(elem_half_64_128size, aclFloat16, half, int32_t, 64, 128, Undefined)
	524	+DEFINE_ELEM(elem_bfloat16_64_128size, uint16_t, bfloat16_t, int32_t, 64, 128, Undefined)
	525	+DEFINE_ELEM(elem_int32_32_64size, int32_t, int32_t, int32_t, 32, 64, Undefined)
	526	+DEFINE_ELEM(elem_uint32_32_64size, uint32_t, uint32_t, int32_t, 32, 64, Undefined)
	527	+DEFINE_ELEM(elem_int16_32_64size, int16_t, int16_t, int32_t, 32, 64, Undefined)
	528	+DEFINE_ELEM(elem_uint16_32_64size, uint16_t, uint16_t, int32_t, 32, 64, Undefined)
	529	+DEFINE_ELEM(elem_int8_64_128size, int8_t, int8_t, int32_t, 64, 128, Undefined)
	530	+DEFINE_ELEM(elem_uint8_64_128size, uint8_t, uint8_t, int32_t, 64, 128, Undefined)
	531	+DEFINE_ELEM(elem_float_clamp_32_16size, float, float, int32_t, 32, 16, Clamp)
	532	+DEFINE_ELEM(elem_int32_wrap_32_16size, int32_t, int32_t, int32_t, 32, 16, Wrap)
	533	+DEFINE_ELEM(elem_half_zero_32_16size, aclFloat16, half, int32_t, 32, 16, Zero)
	534	+
	535	+DEFINE_ELEM2D(elem2d_float_8x32_256size, float, float, int32_t, 8, 32, 256, Undefined)
	536	+DEFINE_ELEM2D(elem2d_int32_8x16_256size, int32_t, int32_t, int32_t, 8, 16, 256, Undefined)
	537	+DEFINE_ELEM2D(elem2d_half_4x32_256size, aclFloat16, half, int32_t, 4, 32, 256, Undefined)
	538	+DEFINE_ELEM2D(elem2d_bfloat16_4x32_256size, uint16_t, bfloat16_t, int32_t, 4, 32, 256, Undefined)
	539	+DEFINE_ELEM2D(elem2d_uint8_4x64_256size, uint8_t, uint8_t, int32_t, 4, 64, 256, Undefined)
	540	+DEFINE_ELEM2D(elem2d_int8_4x64_256size, int8_t, int8_t, int32_t, 4, 64, 256, Undefined)
	541	+DEFINE_ELEM2D(elem2d_int16_4x32_256size, int16_t, int16_t, int32_t, 4, 32, 256, Undefined)
	542	+DEFINE_ELEM2D(elem2d_uint16_4x32_256size, uint16_t, uint16_t, int32_t, 4, 32, 256, Undefined)
	543	+DEFINE_ELEM2D(elem2d_uint32_8x16_256size, uint32_t, uint32_t, int32_t, 8, 16, 256, Undefined)
	544	+DEFINE_ELEM2D(elem2d_float_wrap_4x16_64size, float, float, int32_t, 4, 16, 64, Wrap)
	545	+DEFINE_ELEM2D(elem2d_int32_clamp_4x8_32size, int32_t, int32_t, int32_t, 4, 8, 32, Clamp)
	546	+DEFINE_ELEM2D(elem2d_half_zero_4x32_64size, aclFloat16, half, int32_t, 4, 32, 64, Zero)
	547	+
	548	+DEFINE_ELEM2D_PAD(elem2d_int32_unaligned_3x3_in_3x8_64size, int32_t, int32_t, int32_t, 3, 3, 3, 8, 64, Undefined)
	549	+DEFINE_ELEM2D_PAD(elem2d_float_unaligned_5x5_in_5x8_64size, float, float, int32_t, 5, 5, 5, 8, 64, Undefined)
	550	+DEFINE_ELEM2D_PAD(elem2d_half_unaligned_3x9_in_3x16_64size, aclFloat16, half, int32_t, 3, 9, 3, 16, 64, Undefined)
	551	+DEFINE_ELEM2D_PAD(elem2d_int8_unaligned_3x17_in_3x32_64size, int8_t, int8_t, int32_t, 3, 17, 3, 32, 64, Undefined)
	552	+
	553	+DEFINE_ELEM_SCALAR(elem_scalar_float_1x1_in_1x8_8size, float, float, int32_t, 8, 8, Undefined)
	554	+DEFINE_ELEM_SCALAR(elem_scalar_int32_1x1_in_1x8_8size, int32_t, int32_t, int32_t, 8, 8, Undefined)
	555	+DEFINE_ELEM_SCALAR(elem_scalar_half_1x1_in_1x16_16size, aclFloat16, half, int32_t, 16, 16, Undefined)
	556	+
	557	+DEFINE_ELEM2D_DYN(elem2d_dyn_float_4x8_64size, float, float, int32_t, 4, 8, 4, 8, 64, Undefined)
	558	+DEFINE_ELEM2D_DYN(elem2d_dyn_int32_3x3_in_3x8_64size, int32_t, int32_t, int32_t, 3, 8, 3, 3, 64, Undefined)
	559	+
	560	+DEFINE_ROW_DYN(row_dyn_int32_3x16_8rows, int32_t, int32_t, int32_t, 3, 16, 8, 3, 16, 8, Undefined)
	561	+DEFINE_ROW_DYN(row_dyn_half_4x32_16rows, aclFloat16, half, int32_t, 4, 32, 16, 4, 32, 16, Undefined)
	562	+
	563	+DEFINE_ROW_NZ(row_nz_float_16x16_2blk, float, float, int32_t, 16, 16, 2, 2, 8, Undefined)
	564	+DEFINE_ROW_NZ(row_nz_half_32x16_2blk, aclFloat16, half, int32_t, 32, 16, 2, 1, 16, Undefined)
	565	+DEFINE_ROW_NZ(row_nz_int32_16x16_2blk, int32_t, int32_t, int32_t, 16, 16, 2, 2, 8, Undefined)
	566	+DEFINE_ROW_NZ(row_nz_int16_32x16_1blk, int16_t, int16_t, int32_t, 32, 16, 2, 1, 16, Undefined)
	567	+DEFINE_ROW_NZ(row_nz_int8_16x32_1blk, int8_t, int8_t, int32_t, 16, 32, 2, 1, 32, Undefined)
	568	+DEFINE_ROW_NZ(row_nz_float_clamp_16x8_1blk, float, float, int32_t, 16, 8, 2, 1, 8, Clamp)
	569	+DEFINE_ROW_NZ(row_nz_half_zero_16x16_2blk, aclFloat16, half, int32_t, 16, 16, 2, 1, 16, Zero)
	570	+
	571	+DEFINE_ELEM2D_NZ(elem2d_nz_float_16x16_2blk, float, float, int32_t, 16, 16, 2, 2, 8, Undefined)
	572	+DEFINE_ELEM2D_NZ(elem2d_nz_half_16x16_1blk, aclFloat16, half, int32_t, 16, 16, 2, 1, 16, Undefined)
	573	+DEFINE_ELEM2D_NZ(elem2d_nz_int32_16x8_1blk, int32_t, int32_t, int32_t, 16, 8, 2, 1, 8, Undefined)
	574	+DEFINE_ELEM2D_NZ(elem2d_nz_half_zero_16x16_1blk, aclFloat16, half, int32_t, 16, 16, 2, 1, 16, Zero)

Atests/npu/a2a3/src/st/testcase/mscatter/CMakeLists.txt+9-0

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

Atests/npu/a2a3/src/st/testcase/mscatter/MSCATTER.md+509-0

Atests/npu/a2a3/src/st/testcase/mscatter/gen_data.py+354-0

Atests/npu/a2a3/src/st/testcase/mscatter/main.cpp+296-0

Atests/npu/a2a3/src/st/testcase/mscatter/mscatter_kernel.cpp+559-0

Mtests/run_st.sh+27-0