Support TPREFETCH and PTO_PREFETCH as offical ISA (support both smda CMO and aiv kernel, single commit)-pto-isa-AtomGit

已合并

Support TPREFETCH and PTO_PREFETCH as offical ISA (support both smda CMO and aiv kernel, single commit) #73

ChanKaLok创建于 1月13日从ChanKaLok/pto-isa： pr_pto_prefetch1 合入到 cann/pto-isa： master

Support TPREFETCH and PTO_PREFETCH as offical ISA (support both smda CMO and aiv kernel, single commit) #73

已合并

ChanKaLok创建于 1月13日

从ChanKaLok/pto-isa pr_pto_prefetch1合入到cann/pto-isa master

共 6 个文件变更+148-38

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

	@@ -90,6 +90,12 @@ PTO_INST RecordEvent TLOAD(TileData &dst, GlobalData &src, WaitEvents&... events
90	return {};	90	return {};
91	}	91	}
92		92
		93	template <typename TileData, typename GlobalData>
		94	PTO_INST RecordEvent TPREFETCH(TileData &dst, GlobalData &src) {
		95	MAP_INSTR_IMPL(TPREFETCH, dst, src);
		96	return {};
		97	}
		98
93	template <typename TileDataDst, typename TileDataSrc0, typename T, typename... WaitEvents>	99	template <typename TileDataDst, typename TileDataSrc0, typename T, typename... WaitEvents>
94	PTO_INST RecordEvent TCMPS(TileDataDst &dst, TileDataSrc0 &src0, T src1, CmpMode cmpMode, WaitEvents&... events) {	100	PTO_INST RecordEvent TCMPS(TileDataDst &dst, TileDataSrc0 &src0, T src1, CmpMode cmpMode, WaitEvents&... events) {
95	TSYNC(events...);	101	TSYNC(events...);

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

	@@ -71,6 +71,7 @@ See LICENSE in the root of the software repository for the full text of the Lice
71	#include "pto/npu/a2a3/TScatter.hpp"	71	#include "pto/npu/a2a3/TScatter.hpp"
72	#include "pto/npu/a2a3/TColExpand.hpp"	72	#include "pto/npu/a2a3/TColExpand.hpp"
73	#include "pto/npu/a2a3/TTri.hpp"	73	#include "pto/npu/a2a3/TTri.hpp"
		74	#include "pto/npu/a2a3/TPrefetch.hpp"
74	#endif	75	#endif
75		76
76	#ifdef REGISTER_BASE	77	#ifdef REGISTER_BASE

Ainclude/pto/npu/a2a3/TPrefetch.hpp+92-0

@@ -0,0 +1,92 @@
	1	/**
	2	Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	CANN Open Software License Agreement Version 2.0 (the "License").
	5	Please refer to the License for details. You may not use this file except in compliance with the License.
	6	THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	See LICENSE in the root of the software repository for the full text of the License.
	9	*/
	10
	11	#ifndef TPREFETCH_HPP
	12	#define TPREFETCH_HPP
	13
	14	#include "TLoad.hpp"
	15
	16	namespace pto {
	17	template <typename TileData, typename GlobalData>
	18	__tf__ AICORE void TPrefetchDoCopy(typename TileData::TileDType __out__ dstTile,
	19	typename GlobalData::DType __in__ *srcPtr, uint16_t rowChunk, uint32_t colChunk, int stride3)
	20	{
	21	const uint16_t nBurst = rowChunk;
	22	const uint32_t lenBurst = static_cast<uint32_t>(colChunk * sizeof(typename GlobalData::DType));
	23	const uint32_t gmGap =
	24	static_cast<uint32_t>((stride3 - static_cast<int>(colChunk)) * sizeof(typename GlobalData::DType));
	25	const uint32_t ubGap = 0;
	26	const uint32_t ubPad = 0;
	27	__ubuf__ typename TileData::DType *dstPtr =
	28	(__ubuf__ typename TileData::DType *)__cce_get_tile_ptr(dstTile);
	29	TLoadInstrGm2ub<TileData, GlobalData>(dstPtr, srcPtr, nBurst, lenBurst, gmGap, ubGap, ubPad);
	30	}
	31
	32	template <typename TileData, typename GlobalData>
	33	PTO_INTERNAL void TPrefetchCopySlice(TileData &dst, typename GlobalData::DType *basePtr, int s3, int s4, int st3,
	34	int st4, uint16_t tileRows, uint32_t maxColsPerChunk, bool fits)
	35	{
	36	if (fits) {
	37	TPrefetchDoCopy<TileData, GlobalData>(dst.data(), basePtr, static_cast<uint16_t>(s3),
	38	static_cast<uint32_t>(s4), st3);
	39	return;
	40	}
	41
	42	for (int r = 0; r < s3; r += tileRows) {
	43	const uint16_t rowChunk = static_cast<uint16_t>((s3 - r) < tileRows ? (s3 - r) : tileRows);
	44	typename GlobalData::DType rowPtr = basePtr + r st3;
	45	for (int c = 0; c < s4; c += static_cast<int>(maxColsPerChunk)) {
	46	const uint32_t colChunk = static_cast<uint32_t>((s4 - c) < static_cast<int>(maxColsPerChunk)
	47	? (s4 - c)
	48	: maxColsPerChunk);
	49	typename GlobalData::DType srcPtr = rowPtr + c st4;
	50	TPrefetchDoCopy<TileData, GlobalData>(dst.data(), srcPtr, rowChunk, colChunk, st3);
	51	}
	52	}
	53	}
	54
	55	// Prefetch GlobalTensor into a Vec tile without layout/type checks (dst is temporary)
	56	template <typename TileData, typename GlobalData>
	57	PTO_INTERNAL void TPREFETCH_IMPL(TileData &dst, GlobalData &src)
	58	{
	59	const uint16_t tileRows = TileData::Rows;
	60	const uint32_t tileCols = TileData::Cols;
	61	const std::size_t tileRowBytes = static_cast<std::size_t>(tileCols) * sizeof(typename TileData::DType);
	62	const std::size_t tileBytes = static_cast<std::size_t>(tileRows) * tileRowBytes;
	63
	64	const int s0 = src.GetShape(pto::GlobalTensorDim::DIM_0);
	65	const int s1 = src.GetShape(pto::GlobalTensorDim::DIM_1);
	66	const int s2 = src.GetShape(pto::GlobalTensorDim::DIM_2);
	67	const int s3 = src.GetShape(pto::GlobalTensorDim::DIM_3);
	68	const int s4 = src.GetShape(pto::GlobalTensorDim::DIM_4);
	69
	70	const int st0 = src.GetStride(pto::GlobalTensorDim::DIM_0);
	71	const int st1 = src.GetStride(pto::GlobalTensorDim::DIM_1);
	72	const int st2 = src.GetStride(pto::GlobalTensorDim::DIM_2);
	73	const int st3 = src.GetStride(pto::GlobalTensorDim::DIM_3);
	74	const int st4 = src.GetStride(pto::GlobalTensorDim::DIM_4);
	75
	76	const std::size_t sliceBytes = static_cast<std::size_t>(s3) * s4 * sizeof(typename GlobalData::DType);
	77	const bool fits = sliceBytes <= tileBytes;
	78	const uint32_t maxColsPerChunk = static_cast<uint32_t>(tileRowBytes / sizeof(typename GlobalData::DType));
	79
	80	for (int n0 = 0; n0 < s0; ++n0) {
	81	for (int n1 = 0; n1 < s1; ++n1) {
	82	for (int n2 = 0; n2 < s2; ++n2) {
	83	typename GlobalData::DType basePtr = src.data() + n0 st0 + n1 * st1 + n2 * st2;
	84	TPrefetchCopySlice<TileData, GlobalData>(dst, basePtr, s3, s4, st3, st4, tileRows, maxColsPerChunk,
	85	fits);
	86	}
	87	}
	88	}
	89	}
	90	} // namespace pto
	91
	92	#endif

Dinclude/pto/npu/a2a3/custom/TPrefetch.hpp+0-25

	@@ -1,25 +0,0 @@
1	/**
2	Copyright (c) 2025 Huawei Technologies Co., Ltd.
3	This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4	CANN Open Software License Agreement Version 2.0 (the "License").
5	Please refer to the License for details. You may not use this file except in compliance with the License.
6	THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7	INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8	See LICENSE in the root of the software repository for the full text of the License.
9	*/
10
11	#ifndef TPREFETCH_HPP
12	#define TPREFETCH_HPP
13
14	#include <pto/npu/a2a3/TLoad.hpp>
15
16	namespace pto {
17	// Simple prefetch wrapper to keep call sites explicit
18	template <typename TileData, typename GlobalData>
19	AICORE inline void TPREFETCH(TileData &dst, GlobalData &src)
20	{
21	TLOAD(dst, src);
22	}
23	} // namespace pto
24
25	#endif

Rinclude/pto/npu/a2a3/custom/Pto_prefetch.hpp→include/pto/npu/a2a3/kernels/Pto_prefetch.hpp+35-7

	@@ -11,22 +11,27 @@ See LICENSE in the root of the software repository for the full text of the Lice
11	#ifndef PTO_PREFETCH_HPP	11	#ifndef PTO_PREFETCH_HPP
12	#define PTO_PREFETCH_HPP	12	#define PTO_PREFETCH_HPP
13		13
		14	#include <acl/acl.h>
14	#include <pto/pto-inst.hpp>	15	#include <pto/pto-inst.hpp>
15	#include <pto/npu/a2a3/TLoad.hpp>	16	#include <pto/npu/a2a3/TLoad.hpp>
16		17
17	namespace pto {	18	namespace pto {
18		19
19	constexpr uint32_t kPtoPrefetchTileBytes = 64U * 1024U;	20	constexpr uint32_t kPtoPrefetchTileBytes = 64U * 1024U;
		21	constexpr uint32_t kPtoPrefetchDefaultBlocks = 20U;
20		22
21	// Generic prefetch kernel: split a 1D tensor across blocks (get_blockdim()) and	23	#define PTO_AIV_ATTR __attribute__((aiv))
22	// issue TLOAD for fixed 64KB tiles using a 1D layout.	24	#define PTO_PREFETCH_DEVICE_ENABLED 1
23	template <typename DType>	25
24	__global__ AICORE __attribute__((aiv)) void PTO_PREFETCH(__gm__ DType *tensor, uint64_t total_elems)	26	#if PTO_PREFETCH_DEVICE_ENABLED
		27	namespace detail {
		28	PTO_INTERNAL void PtoPrefetchKernelBody(__gm__ uint8_t *tensor, uint64_t total_elems)
25	{	29	{
26	constexpr uint32_t tile_bytes = kPtoPrefetchTileBytes;	30	constexpr uint32_t tile_bytes = kPtoPrefetchTileBytes;
27	constexpr uint32_t tile_elems = tile_bytes / sizeof(DType);	31	constexpr uint32_t tile_elems = tile_bytes / sizeof(uint8_t);
28	static_assert(tile_elems > 0, "tile_elems must be positive");	32	static_assert(tile_elems > 0, "tile_elems must be positive");
29		33
		34	using DType = uint8_t;
30	using PrefetchTile = Tile<TileType::Vec, DType, 1, tile_elems, BLayout::RowMajor, 1, DYNAMIC>;	35	using PrefetchTile = Tile<TileType::Vec, DType, 1, tile_elems, BLayout::RowMajor, 1, DYNAMIC>;
31	using PrefetchShape = Shape<1, 1, 1, 1, DYNAMIC>;	36	using PrefetchShape = Shape<1, 1, 1, 1, DYNAMIC>;
32	using PrefetchStride = Stride<1, 1, 1, DYNAMIC, 1>;	37	using PrefetchStride = Stride<1, 1, 1, DYNAMIC, 1>;
	@@ -37,7 +42,6 @@ __global__ AICORE __attribute__((aiv)) void PTO_PREFETCH(__gm__ DType *tensor, u
37		42
38	const uint32_t blk = get_block_idx();	43	const uint32_t blk = get_block_idx();
39		44
40	// Allocate whole tiles to each block; only the final block may have a partial tile at the end.
41	const uint64_t total_tiles = (total_elems + tile_elems - 1ULL) / static_cast<uint64_t>(tile_elems);	45	const uint64_t total_tiles = (total_elems + tile_elems - 1ULL) / static_cast<uint64_t>(tile_elems);
42	const uint64_t tiles_per_block = (total_tiles + block_dim - 1ULL) / static_cast<uint64_t>(block_dim);	46	const uint64_t tiles_per_block = (total_tiles + block_dim - 1ULL) / static_cast<uint64_t>(block_dim);
43		47
	@@ -59,7 +63,31 @@ __global__ AICORE __attribute__((aiv)) void PTO_PREFETCH(__gm__ DType *tensor, u
59	PrefetchShape dyn_shape(1, 1, 1, 1, static_cast<int>(cur_elems));	63	PrefetchShape dyn_shape(1, 1, 1, 1, static_cast<int>(cur_elems));
60	PrefetchStride dyn_stride(1, 1, 1, static_cast<int>(cur_elems), 1);	64	PrefetchStride dyn_stride(1, 1, 1, static_cast<int>(cur_elems), 1);
61	GlobalTensor<DType, PrefetchShape, PrefetchStride> g(tensor + offset, dyn_shape, dyn_stride);	65	GlobalTensor<DType, PrefetchShape, PrefetchStride> g(tensor + offset, dyn_shape, dyn_stride);
62	TLOAD(tile, g);	66	TPREFETCH(tile, g);
		67	}
		68	}
		69	} // namespace detail
		70
		71	// Generic prefetch kernel: split a 1D tensor across blocks (get_blockdim()) and issue TPREFETCH
		72	__global__ AICORE PTO_AIV_ATTR void PTO_PREFETCH_AIV(__gm__ uint8_t *tensor, uint64_t total_elems)
		73	{
		74	detail::PtoPrefetchKernelBody(tensor, total_elems);
		75	}
		76	#endif // PTO_PREFETCH_DEVICE_ENABLED
		77
		78	// Host wrapper to launch PTO_PREFETCH with bytes input and optional SDMA/AIV core selection.
		79	// Use the template parameters to pick SDMA or AIV and to set aiv_cores for finer control.
		80	template <bool UseSdma = true, int AivCores = -1>
		81	void PTO_PREFETCH(__gm__ void *tensor, uint64_t tensor_bytes, aclrtStream stream)
		82	{
		83	if (tensor_bytes == 0)
		84	return;
		85
		86	if constexpr (UseSdma) {
		87	aclrtCmoAsync((void *)(uint64_t)tensor, static_cast<size_t>(tensor_bytes), ACL_RT_CMO_TYPE_PREFETCH, stream);
		88	} else {
		89	static_assert(AivCores > 0, "AivCores must be > 0 when UseSdma is false");
		90	PTO_PREFETCH_AIV<<<AivCores, nullptr, stream>>>((__gm__ uint8_t *)tensor, tensor_bytes);
63	}	91	}
64	}	92	}
65		93

Mkernels/manual/a2a3/flash_atten/fa_performance_kernel.cpp+14-6

	@@ -12,7 +12,7 @@ See LICENSE in the root of the software repository for the full text of the Lice
12	#include <pto/pto-inst.hpp>	12	#include <pto/pto-inst.hpp>
13		13
14	#include "fa_performance_kernel.h"	14	#include "fa_performance_kernel.h"
15	#include <pto/npu/a2a3/custom/Pto_prefetch.hpp>	15	#include <pto/npu/a2a3/kernels/Pto_prefetch.hpp>
16	#include <pto/npu/a2a3/custom/TSyncCVID.hpp>	16	#include <pto/npu/a2a3/custom/TSyncCVID.hpp>
17	#include "pto_macro_matmul.hpp"	17	#include "pto_macro_matmul.hpp"
18	#include "pto_macro_fa_softmax.hpp"	18	#include "pto_macro_fa_softmax.hpp"
	@@ -34,7 +34,7 @@ enum FftsBufferFlag : uint32_t {
34	CV_BLOCK_END = 7, // CV comm slot block end (CV_COMM_CTRL reserved in TSyncCVID)	34	CV_BLOCK_END = 7, // CV comm slot block end (CV_COMM_CTRL reserved in TSyncCVID)
35	};	35	};
36	#endif	36	#endif
37		37
38	enum CoreEvtID : uint32_t {	38	enum CoreEvtID : uint32_t {
39	QK_EVENT_ID0,	39	QK_EVENT_ID0,
40	QK_EVENT_ID1,	40	QK_EVENT_ID1,
	@@ -904,11 +904,19 @@ void LaunchTFA(uint16_t ffts, aclFloat16 q, aclFloat16 k, aclFloat16 v, aclF
904	warmup_kernel<<<24, nullptr, stream>>>();	904	warmup_kernel<<<24, nullptr, stream>>>();
905		905
906	const uint64_t tensor_elems = static_cast<uint64_t>(S0) * static_cast<uint64_t>(HEAD_SIZE);	906	const uint64_t tensor_elems = static_cast<uint64_t>(S0) * static_cast<uint64_t>(HEAD_SIZE);
907	constexpr uint32_t prefetch_blocks = 20;	907	const uint64_t tensor_bytes = tensor_elems * sizeof(half);
		908	constexpr bool kPrefetchUseSdma = true; //simulation cannot use sdma
		909	constexpr int kPrefetchAivCores = 40; // only used when kPrefetchUseSdma is false
908		910
909	PTO_PREFETCH<<<prefetch_blocks, nullptr, stream>>>((__gm__ half *)q, tensor_elems);	911	if constexpr (kPrefetchUseSdma) {
910	PTO_PREFETCH<<<prefetch_blocks, nullptr, stream>>>((__gm__ half *)k, tensor_elems);	912	PTO_PREFETCH((__gm__ void *)q, tensor_bytes, stream);
911	PTO_PREFETCH<<<prefetch_blocks, nullptr, stream>>>((__gm__ half *)v, tensor_elems);	913	PTO_PREFETCH((__gm__ void *)k, tensor_bytes, stream);
		914	PTO_PREFETCH((__gm__ void *)v, tensor_bytes, stream);
		915	} else {
		916	PTO_PREFETCH<false, kPrefetchAivCores>((__gm__ void *)q, tensor_bytes, stream);
		917	PTO_PREFETCH<false, kPrefetchAivCores>((__gm__ void *)k, tensor_bytes, stream);
		918	PTO_PREFETCH<false, kPrefetchAivCores>((__gm__ void *)v, tensor_bytes, stream);
		919	}
912		920
913	runTFA<S0, HEAD_SIZE, S1, CUBE_S0, CUBE_S1, TILE_S1, QK_PRELOAD, CV_FIFO_SIZE, INTERMEDIATE_CHECK,	921	runTFA<S0, HEAD_SIZE, S1, CUBE_S0, CUBE_S1, TILE_S1, QK_PRELOAD, CV_FIFO_SIZE, INTERMEDIATE_CHECK,
914	CV_FIFO_CONS_SYNC_PERIOD><<<block_rows, nullptr, stream>>>((__gm__ uint64_t )ffts, (half )q, (half *)k,	922	CV_FIFO_CONS_SYNC_PERIOD><<<block_rows, nullptr, stream>>>((__gm__ uint64_t )ffts, (half )q, (half *)k,