已合并
[a2a3 TTrans] : add NCDHW -> FRACTAL_Z_3D feature #977
Sathi Sarveswara Reddy创建于 5月22日
[a2a3 TTrans] : add NCDHW -> FRACTAL_Z_3D feature #977
已合并
Sathi Sarveswara Reddy创建于 5月22日
8 个文件变更+676-2
@@ -1,4 +1,4 @@
1-# TTRANS1+# TTRANS
2 2 
3 3 
4## Tile Operation Diagram4## Tile Operation Diagram
@@ -73,6 +73,7 @@ PTO_INST RecordEvent TTRANS(TileDataDst &dst, TileDataSrc &src, TileDataTmp &tmp
73 - Transpose of ConvTile for `TileType::Vec` is supported。 Element size must be `1``2` or `4` bytes. Supported element types are `uint32_t``int32_t``float``uint16_t``int16_t``half``bfloat16_t``uint8_t``int8_t`.73 - Transpose of ConvTile for `TileType::Vec` is supported。 Element size must be `1``2` or `4` bytes. Supported element types are `uint32_t``int32_t``float``uint16_t``int16_t``half``bfloat16_t``uint8_t``int8_t`.
74 - Format transformation from `NCHW` to `NC1HWC0` is supported, while `C1 == (C + C0 - 1)/C0`,HW matches alignment constraint,which means `H*W*sizeof(T)==0`. C0 means `c0_size`, which `C0 * sizeof(T) == 32`。C0 can also be 4.74 - Format transformation from `NCHW` to `NC1HWC0` is supported, while `C1 == (C + C0 - 1)/C0`,HW matches alignment constraint,which means `H*W*sizeof(T)==0`. C0 means `c0_size`, which `C0 * sizeof(T) == 32`。C0 can also be 4.
75 - Format transformation from `NC1HWC0` to `FRACTAL_Z` is supported, while `N1 == (N + N0 - 1)/N0`。N0 should be 16.75 - Format transformation from `NC1HWC0` to `FRACTAL_Z` is supported, while `N1 == (N + N0 - 1)/N0`。N0 should be 16.
76+ - Format transformation from `NCDHW` to `FRACTAL_Z_3D` is supported, with the destination shape `[D * C1 * H * W, N1, N0, C0]`, where `C1 == (C + C0 - 1)/C0` and `N1 == (N + N0 - 1)/N0`. `N0` is `16`. `C0` depends on element width: `64` for 4-bit data, `32` for 8-bit data, `16` for 16-bit data and `8` for 32-bit data. The temporary tile must be large enough to stage the intermediate `(D + 1) * N * C1 * C0 * H * W` elements together with an `H * W * alignedC0` sub-scratch (`alignedC0` is `dstC0` rounded up to `16` for 16/32-bit data and to `32` for 8-bit data).
76 77 
77## Examples78## Examples
78 79 
@@ -644,6 +644,292 @@ __tf__ PTO_INTERNAL void TTransConvGNC1HWC02GC1HWNC0(typename TileDataDst::TileD
644 }644 }
645}645}
646 646 
647+template <typename T>
648+PTO_INTERNAL void ConvNCDHW2DNCHWUnalign(__ubuf__ T *dst, __ubuf__ T *src, unsigned srcN, unsigned srcC, unsigned srcD,
649+ unsigned srcH, unsigned srcW, unsigned dstC0)
650+{
651+ unsigned dstC1 = (srcC + dstC0 - 1) / dstC0;
652+ unsigned paddedC = dstC1 * dstC0;
653+ unsigned hw = srcH * srcW;
654+ unsigned ncStride = srcD * hw;
655+ unsigned dPlane = srcN * paddedC * hw;
656+#ifndef __PTO_AUTO__
657+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
658+#else
659+ set_flag(PIPE_V, PIPE_S, EVENT_ID0);
660+ wait_flag(PIPE_V, PIPE_S, EVENT_ID0);
661+#endif
662+ for (unsigned d = 0; d < srcD; d++) {
663+ for (unsigned n = 0; n < srcN; n++) {
664+ for (unsigned c = 0; c < paddedC; c++) {
665+ __ubuf__ T *dstPtr = dst + d * dPlane + n * paddedC * hw + c * hw;
666+ if (c < srcC) {
667+ __ubuf__ T *srcPtr = src + n * srcC * ncStride + c * ncStride + d * hw;
668+ for (unsigned i = 0; i < hw; ++i) {
669+ dstPtr[i] = srcPtr[i];
670+ }
671+ } else {
672+ for (unsigned i = 0; i < hw; ++i) {
673+ dstPtr[i] = static_cast<T>(0);
674+ }
675+ }
676+ }
677+ }
678+ }
679+#ifndef __PTO_AUTO__
680+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
681+#else
682+ set_flag(PIPE_S, PIPE_V, EVENT_ID0);
683+ wait_flag(PIPE_S, PIPE_V, EVENT_ID0);
684+#endif
685+}
686+ 
687+template <typename T>
688+PTO_INTERNAL void ConvNCDHW2DNCHW(__ubuf__ T *dst, __ubuf__ T *src, unsigned srcN, unsigned srcC, unsigned srcD,
689+ unsigned srcH, unsigned srcW, unsigned dstC0)
690+{
691+ unsigned hw = srcH * srcW;
692+ if ((hw * sizeof(T)) % BLOCK_BYTE_SIZE != 0) {
693+ ConvNCDHW2DNCHWUnalign<T>(dst, src, srcN, srcC, srcD, srcH, srcW, dstC0);
694+ return;
695+ }
696+ unsigned dstC1 = (srcC + dstC0 - 1) / dstC0;
697+ unsigned paddedC = dstC1 * dstC0;
698+ unsigned padC = paddedC - srcC;
699+ unsigned ncStride = srcD * hw;
700+ unsigned dPlane = srcN * paddedC * hw;
701+ uint32_t lenBurst = (hw * sizeof(T)) / BLOCK_BYTE_SIZE;
702+ uint16_t srcGap = 0;
703+ uint16_t dstGap = (uint16_t)((dPlane - hw) * sizeof(T) / BLOCK_BYTE_SIZE);
704+ for (unsigned n = 0; n < srcN; n++) {
705+ for (unsigned c = 0; c < srcC; c++) {
706+ __ubuf__ T *srcPtr = src + n * srcC * ncStride + c * ncStride;
707+ __ubuf__ T *dstPtr = dst + n * paddedC * hw + c * hw;
708+ pto_copy_ubuf_to_ubuf(dstPtr, srcPtr, (uint16_t)srcD, (uint16_t)lenBurst, srcGap, dstGap);
709+ }
710+ }
711+ if (padC > 0) {
712+#ifndef __PTO_AUTO__
713+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
714+#else
715+ set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
716+ wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
717+#endif
718+ for (unsigned d = 0; d < srcD; d++) {
719+ for (unsigned n = 0; n < srcN; n++) {
720+ for (unsigned c = srcC; c < paddedC; c++) {
721+ __ubuf__ T *dstPtr = dst + d * dPlane + n * paddedC * hw + c * hw;
722+ for (unsigned i = 0; i < hw; ++i) {
723+ dstPtr[i] = static_cast<T>(0);
724+ }
725+ }
726+ }
727+ }
728+#ifndef __PTO_AUTO__
729+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
730+#else
731+ set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
732+ wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
733+#endif
734+ }
735+}
736+ 
737+template <typename TileDataDst, typename TileDataSrc, typename TileDataTmp, unsigned blockSizeElem>
738+__tf__ PTO_INTERNAL void TTransConvNCDHW2FractalZ3D(typename TileDataDst::TileDType __out__ dst,
739+ typename TileDataSrc::TileDType __in__ src,
740+ typename TileDataTmp::TileDType __in__ tmp, unsigned srcN,
741+ unsigned srcC, unsigned srcD, unsigned srcH, unsigned srcW,
742+ unsigned dstN0, unsigned dstC0)
743+{
744+ using Tdst = typename TileDataDst::DType;
745+ using Tsrc = typename TileDataSrc::DType;
746+ using Ttmp = typename TileDataTmp::DType;
747+ 
748+ __ubuf__ Tdst *dstPtrOrig = (__ubuf__ Tdst *)__cce_get_tile_ptr(dst);
749+ __ubuf__ Tsrc *srcPtrOrig = (__ubuf__ Tsrc *)__cce_get_tile_ptr(src);
750+ __ubuf__ Ttmp *tmpPtrOrig = (__ubuf__ Ttmp *)__cce_get_tile_ptr(tmp);
751+ 
752+ unsigned dstC1 = (srcC + dstC0 - 1) / dstC0;
753+ unsigned dstN1 = (srcN + dstN0 - 1) / dstN0;
754+ unsigned paddedC = dstC1 * dstC0;
755+ unsigned paddedN = dstN1 * dstN0;
756+ unsigned hw = srcH * srcW;
757+ unsigned ncplaneSize = srcN * paddedC * hw;
758+ unsigned c1hw = dstC1 * hw;
759+ unsigned dstSliceSize = c1hw * paddedN * dstC0;
760+ bool useScalarNCHW = ((dstC0 % blockSizeElem) != 0) || ((hw % blockSizeElem) != 0) || hw / blockSizeElem > 255;
761+ bool useScalarC1HW = ((dstC0 * sizeof(Tsrc)) % BLOCK_BYTE_SIZE) != 0;
762+ 
763+ ConvNCDHW2DNCHW<Tsrc>(tmpPtrOrig, srcPtrOrig, srcN, srcC, srcD, srcH, srcW, dstC0);
764+#ifndef __PTO_AUTO__
765+ PtoSetWaitFlag<PIPE_MTE3, PIPE_V>();
766+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
767+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
768+#else
769+ set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
770+ wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
771+ set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
772+ wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
773+ set_flag(PIPE_S, PIPE_V, EVENT_ID0);
774+ wait_flag(PIPE_S, PIPE_V, EVENT_ID0);
775+#endif
776+ 
777+ __ubuf__ Ttmp *stagePtr = tmpPtrOrig + srcD * ncplaneSize;
778+ __ubuf__ Ttmp *subTmpPtr = stagePtr + ncplaneSize;
779+ for (unsigned d = 0; d < srcD; d++) {
780+ __ubuf__ Tsrc *planePtr = tmpPtrOrig + d * ncplaneSize;
781+ __ubuf__ Tsrc *nc1hwc0Ptr;
782+ unsigned srcStride = hw;
783+ unsigned dstStride = dstC0;
784+ if (useScalarNCHW) {
785+ ConvNCHW2NC1HWC0Unalign<Tsrc, blockSizeElem>(stagePtr, planePtr, srcN, srcC, srcH, srcW, dstC0);
786+ nc1hwc0Ptr = stagePtr;
787+ } else {
788+ if (d > 0) {
789+#ifndef __PTO_AUTO__
790+ PtoSetWaitFlag<PIPE_MTE3, PIPE_V>();
791+#else
792+ set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
793+ wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
794+#endif
795+ }
796+ unsigned validCol = hw;
797+ unsigned validRow = dstC0;
798+ unsigned nStride = dstC1 * dstC0 * hw;
799+ unsigned cStride = dstC0 * hw;
800+ for (unsigned n = 0; n < srcN; n++) {
801+ for (unsigned c = 0; c < dstC1; c++) {
802+ __ubuf__ Tsrc *innerSrc = planePtr + n * nStride + c * cStride;
803+ __ubuf__ Tsrc *innerDst = planePtr + n * nStride + c * cStride;
804+ TTransRepeatXOperation<Tsrc, blockSizeElem>(innerDst, innerSrc, subTmpPtr, validRow, validCol,
805+ dstStride, srcStride);
806+ }
807+ }
808+ nc1hwc0Ptr = planePtr;
809+ }
810+ 
811+ __ubuf__ Tdst *dstSlicePtr = dstPtrOrig + d * dstSliceSize;
812+ if (!useScalarC1HW) {
813+ if (useScalarNCHW) {
814+#ifndef __PTO_AUTO__
815+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
816+#else
817+ set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
818+ wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
819+#endif
820+ } else {
821+#ifndef __PTO_AUTO__
822+ PtoSetWaitFlag<PIPE_V, PIPE_MTE3>();
823+#else
824+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
825+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
826+#endif
827+ }
828+ uint32_t burstNum = c1hw;
829+ uint32_t lenBurst = (dstC0 * sizeof(Tsrc) + BLOCK_BYTE_SIZE - 1) / BLOCK_BYTE_SIZE;
830+ uint32_t srcGap = 0;
831+ uint32_t dstGap = (paddedN * dstC0 * sizeof(Tsrc) + BLOCK_BYTE_SIZE - 1) / BLOCK_BYTE_SIZE - lenBurst;
832+ unsigned nStride = c1hw * dstC0;
833+ for (unsigned i = 0; i < srcN; i++) {
834+ __ubuf__ Tsrc *innerSrc = nc1hwc0Ptr + i * nStride;
835+ __ubuf__ Tdst *innerDst = dstSlicePtr + i * dstC0;
836+ pto_copy_ubuf_to_ubuf(innerDst, innerSrc, burstNum, lenBurst, srcGap, dstGap);
837+ }
838+ unsigned remain = paddedN - srcN;
839+ if (remain != 0U) {
840+#ifndef __PTO_AUTO__
841+ PtoSetWaitFlag<PIPE_MTE3, PIPE_S>();
842+#else
843+ set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
844+ wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
845+#endif
846+ unsigned dstStride2 = paddedN * dstC0;
847+ __ubuf__ Tdst *padBase = dstSlicePtr + srcN * dstC0;
848+ for (unsigned r = 0; r < remain; r++) {
849+ for (unsigned i = 0; i < c1hw; ++i) {
850+ for (unsigned j = 0; j < dstC0; ++j) {
851+ padBase[r * dstC0 + i * dstStride2 + j] = static_cast<Tdst>(0);
852+ }
853+ }
854+ }
855+#ifndef __PTO_AUTO__
856+ PtoSetWaitFlag<PIPE_S, PIPE_MTE3>();
857+#else
858+ set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
859+ wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
860+#endif
861+ }
862+ } else {
863+ unsigned validCol = dstC0;
864+ unsigned validRow = c1hw;
865+ unsigned srcStride2 = dstC0;
866+ unsigned dstStride2 = paddedN * dstC0;
867+ unsigned nStride = c1hw * dstC0;
868+ unsigned remain = paddedN - srcN;
869+ if (!useScalarNCHW) {
870+#ifndef __PTO_AUTO__
871+ PtoSetWaitFlag<PIPE_V, PIPE_S>();
872+#else
873+ set_flag(PIPE_V, PIPE_S, EVENT_ID0);
874+ wait_flag(PIPE_V, PIPE_S, EVENT_ID0);
875+#endif
876+ }
877+ for (unsigned i = 0; i < srcN; i++) {
878+ __ubuf__ Tsrc *innerSrc = nc1hwc0Ptr + i * nStride;
879+ __ubuf__ Tdst *innerDst = dstSlicePtr + i * dstC0;
880+ for (unsigned r = 0; r < validRow; ++r) {
881+ for (unsigned j = 0; j < validCol; ++j) {
882+ innerDst[r * dstStride2 + j] = innerSrc[r * srcStride2 + j];
883+ }
884+ }
885+ }
886+ for (unsigned r = 0; r < remain; r++) {
887+ __ubuf__ Tdst *innerDst = dstSlicePtr + (r + srcN) * dstC0;
888+ for (unsigned i = 0; i < validRow; ++i) {
889+ for (unsigned j = 0; j < validCol; ++j) {
890+ innerDst[i * dstStride2 + j] = static_cast<Tdst>(0);
891+ }
892+ }
893+ }
894+ if (!useScalarNCHW) {
895+#ifndef __PTO_AUTO__
896+ PtoSetWaitFlag<PIPE_S, PIPE_V>();
897+#else
898+ set_flag(PIPE_S, PIPE_V, EVENT_ID0);
899+ wait_flag(PIPE_S, PIPE_V, EVENT_ID0);
900+#endif
901+ }
902+ }
903+ }
904+}
905+ 
906+template <typename TileDataDst, typename TileDataSrc, typename TileDataTmp>
907+PTO_INTERNAL void CheckConv3DTile(TileDataDst &dst, TileDataSrc &src, TileDataTmp &tmp)
908+{
909+#ifdef _DEBUG
910+ using T = typename TileDataSrc::DType;
911+ constexpr const int UB_SIZE = 196608; // 192*1024 B
912+ if constexpr (TileDataSrc::layout == Layout::NCDHW && TileDataDst::layout == Layout::FRACTAL_Z_3D) {
913+ unsigned srcN = src.GetShape(GlobalTensorDim::DIM_0);
914+ unsigned srcC = src.GetShape(GlobalTensorDim::DIM_1);
915+ unsigned srcD = src.GetShape(GlobalTensorDim::DIM_2);
916+ unsigned srcH = src.GetShape(GlobalTensorDim::DIM_3);
917+ unsigned srcW = src.GetShape(GlobalTensorDim::DIM_4);
918+ unsigned dstDC1HW = dst.GetShape(GlobalTensorDim::DIM_0);
919+ unsigned dstN1 = dst.GetShape(GlobalTensorDim::DIM_1);
920+ unsigned dstN0 = dst.GetShape(GlobalTensorDim::DIM_2);
921+ unsigned dstC0 = dst.GetShape(GlobalTensorDim::DIM_3);
922+ unsigned dstC1 = (srcC + dstC0 - 1) / dstC0;
923+ unsigned srcSize = srcN * srcC * srcD * srcH * srcW;
924+ unsigned dstSize = dstDC1HW * dstN1 * dstN0 * dstC0;
925+ unsigned tmpSize = TileDataTmp::Rows * TileDataTmp::Cols;
926+ PTO_ASSERT(srcD * dstC1 * srcH * srcW == dstDC1HW && dstN1 == (srcN + dstN0 - 1) / dstN0,
927+ "expect same size for src and dst.");
928+ PTO_ASSERT((srcSize + dstSize + tmpSize) * sizeof(T) < UB_SIZE, "ERROR: memory usage exceeds UB limit!");
929+ }
930+#endif
931+}
932+ 
647template <typename TileDataDst, typename TileDataSrc, typename TileDataTmp>933template <typename TileDataDst, typename TileDataSrc, typename TileDataTmp>
648PTO_INTERNAL void CheckConvTile(TileDataDst &dst, TileDataSrc &src, TileDataTmp &tmp)934PTO_INTERNAL void CheckConvTile(TileDataDst &dst, TileDataSrc &src, TileDataTmp &tmp)
649{935{
@@ -780,6 +1066,17 @@ PTO_INTERNAL void TTransImplConvTile(TileDataDst &dst, TileDataSrc &src, TileDat
780 unsigned dstC0 = dst.GetShape(GlobalTensorDim::DIM_4);1066 unsigned dstC0 = dst.GetShape(GlobalTensorDim::DIM_4);
781 TTransConvNCHW2NC1HWC0<TileDataDst, TileDataSrc, TileDataTmp, blockSizeElem>(dst.data(), src.data(), tmp.data(),1067 TTransConvNCHW2NC1HWC0<TileDataDst, TileDataSrc, TileDataTmp, blockSizeElem>(dst.data(), src.data(), tmp.data(),
782 srcN, srcC, srcH, srcW, dstC0);1068 srcN, srcC, srcH, srcW, dstC0);
1069+ } else if constexpr (TileDataSrc::layout == Layout::NCDHW && TileDataDst::layout == Layout::FRACTAL_Z_3D) {
1070+ CheckConv3DTile<TileDataDst, TileDataSrc, TileDataTmp>(dst, src, tmp);
1071+ unsigned srcN = src.GetShape(GlobalTensorDim::DIM_0);
1072+ unsigned srcC = src.GetShape(GlobalTensorDim::DIM_1);
1073+ unsigned srcD = src.GetShape(GlobalTensorDim::DIM_2);
1074+ unsigned srcH = src.GetShape(GlobalTensorDim::DIM_3);
1075+ unsigned srcW = src.GetShape(GlobalTensorDim::DIM_4);
1076+ unsigned dstN0 = dst.GetShape(GlobalTensorDim::DIM_2);
1077+ unsigned dstC0 = dst.GetShape(GlobalTensorDim::DIM_3);
1078+ TTransConvNCDHW2FractalZ3D<TileDataDst, TileDataSrc, TileDataTmp, blockSizeElem>(
1079+ dst.data(), src.data(), tmp.data(), srcN, srcC, srcD, srcH, srcW, dstN0, dstC0);
783 }1080 }
784}1081}
785 1082 
@@ -112,6 +112,7 @@ tmov
112tload112tload
113ttrans113ttrans
114ttrans_conv114ttrans_conv
115+ttrans_3d
115tmrgsort116tmrgsort
116tfillpad117tfillpad
117tgather118tgather
@@ -0,0 +1,11 @@
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+pto_vec_st(ttrans_3d)
@@ -0,0 +1,93 @@
1+#!/user/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+ 
15+import numpy as np
16+ 
17+np.random.seed(19)
18+ 
19+ 
20+def ncdhw_to_fractal_z_3d(ncdhw_tensor: np.ndarray, c0: int, n0: int) -> np.ndarray:
21+ if ncdhw_tensor.ndim != 5:
22+ raise ValueError(f"The input must be a 5-dimensional NCDHW tensor, current dim : {ncdhw_tensor.ndim}")
23+ 
24+ n, c, d, h, w = ncdhw_tensor.shape
25+ c1 = (c + c0 - 1) // c0
26+ n1 = (n + n0 - 1) // n0
27+ pad_c = c1 * c0 - c
28+ pad_n = n1 * n0 - n
29+ 
30+ padded = ncdhw_tensor
31+ if pad_c > 0:
32+ padded = np.pad(padded, ((0, 0), (0, pad_c), (0, 0), (0, 0), (0, 0)), mode="constant", constant_values=0)
33+ if pad_n > 0:
34+ padded = np.pad(padded, ((0, pad_n), (0, 0), (0, 0), (0, 0), (0, 0)), mode="constant", constant_values=0)
35+ 
36+ reshaped = padded.reshape(n1 * n0, c1, c0, d, h, w)
37+ transposed = np.transpose(reshaped, axes=(3, 1, 4, 5, 0, 2))
38+ fractal = transposed.reshape(d * c1 * h * w, n1, n0, c0)
39+ return fractal
40+ 
41+ 
42+def gen_golden_data(g_info):
43+ data_type = g_info.data_type
44+ src_n = g_info.src_n
45+ src_c = g_info.src_c
46+ src_d = g_info.src_d
47+ src_h = g_info.src_h
48+ src_w = g_info.src_w
49+ dst_n0 = g_info.dst_n0
50+ dst_c0 = g_info.dst_c0
51+ 
52+ input_arr = np.random.randint(1, 5, size=(src_n, src_c, src_d, src_h, src_w)).astype(data_type)
53+ output_arr = ncdhw_to_fractal_z_3d(input_arr, dst_c0, dst_n0)
54+ 
55+ input_arr.tofile("./input.bin")
56+ output_arr.tofile("./golden.bin")
57+ 
58+ 
59+class TTRANS3DParams:
60+ def __init__(self, case_name, data_type, src_n, src_c, src_d, src_h, src_w, dst_n0, dst_c0):
61+ self.case_name = case_name
62+ self.data_type = data_type
63+ self.src_n = src_n
64+ self.src_c = src_c
65+ self.src_d = src_d
66+ self.src_h = src_h
67+ self.src_w = src_w
68+ self.dst_n0 = dst_n0
69+ self.dst_c0 = dst_c0
70+ 
71+ 
72+if __name__ == "__main__":
73+ case_params_list = [
74+ TTRANS3DParams("TTRANS3DTest.case1_float32_2_4_2_2_2", np.float32, 2, 4, 2, 2, 2, 16, 8),
75+ TTRANS3DParams("TTRANS3DTest.case2_float32_4_5_2_2_4", np.float32, 4, 5, 2, 2, 4, 16, 8),
76+ TTRANS3DParams("TTRANS3DTest.case3_int32_17_3_3_2_2", np.int32, 17, 3, 3, 2, 2, 16, 8),
77+ TTRANS3DParams("TTRANS3DTest.case4_half_5_6_2_2_4", np.float16, 5, 6, 2, 2, 4, 16, 16),
78+ TTRANS3DParams("TTRANS3DTest.case5_half_19_14_2_4_2", np.float16, 19, 14, 2, 4, 2, 16, 16),
79+ TTRANS3DParams("TTRANS3DTest.case6_int16_8_13_2_3_4", np.int16, 8, 13, 2, 3, 4, 16, 16),
80+ TTRANS3DParams("TTRANS3DTest.case7_uint16_4_8_2_2_3", np.uint16, 4, 8, 2, 2, 3, 16, 16),
81+ TTRANS3DParams("TTRANS3DTest.case8_int8_7_28_2_3_4", np.int8, 7, 28, 2, 3, 4, 16, 32),
82+ TTRANS3DParams("TTRANS3DTest.case9_int8_16_26_3_3_4", np.int8, 16, 26, 3, 3, 4, 16, 32),
83+ TTRANS3DParams("TTRANS3DTest.case10_uint8_9_18_2_2_4", np.uint8, 9, 18, 2, 2, 4, 16, 32),
84+ ]
85+ 
86+ for case_params in case_params_list:
87+ case_name = case_params.case_name
88+ if not os.path.exists(case_name):
89+ os.makedirs(case_name)
90+ original_dir = os.getcwd()
91+ os.chdir(case_name)
92+ gen_golden_data(case_params)
93+ os.chdir(original_dir)
@@ -0,0 +1,135 @@
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+template <typename T, int dstDC1HW, int dstN1, int dstN0, int dstC0, int srcN, int srcC, int srcD, int srcH, int srcW>
19+void LaunchTTRANS3D(T *out, T *src, void *stream);
20+ 
21+class TTRANS3DTest : public testing::Test {
22+protected:
23+ void SetUp() override
24+ {}
25+ void TearDown() override
26+ {}
27+};
28+ 
29+std::string GetGoldenDir()
30+{
31+ const testing::TestInfo *testInfo = testing::UnitTest::GetInstance()->current_test_info();
32+ const std::string caseName = testInfo->name();
33+ std::string suiteName = testInfo->test_suite_name();
34+ std::string fullPath = "../" + suiteName + "." + caseName;
35+ return fullPath;
36+}
37+ 
38+template <typename T, int dstDC1HW, int dstN1, int dstN0, int dstC0, int srcN, int srcC, int srcD, int srcH, int srcW>
39+void test_ttrans_3d()
40+{
41+ size_t srcFileSize = static_cast<size_t>(srcN) * srcC * srcD * srcH * srcW * sizeof(T);
42+ size_t dstFileSize = static_cast<size_t>(dstDC1HW) * dstN1 * dstN0 * dstC0 * sizeof(T);
43+ 
44+ aclInit(nullptr);
45+ aclrtSetDevice(0);
46+ aclrtStream stream;
47+ aclrtCreateStream(&stream);
48+ 
49+ T *dstHost, *srcHost;
50+ T *dstDevice, *srcDevice;
51+ 
52+ aclrtMallocHost((void **)(&dstHost), dstFileSize);
53+ aclrtMallocHost((void **)(&srcHost), srcFileSize);
54+ 
55+ aclrtMalloc((void **)&dstDevice, dstFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
56+ aclrtMalloc((void **)&srcDevice, srcFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
57+ 
58+ ReadFile(GetGoldenDir() + "/input.bin", srcFileSize, srcHost, srcFileSize);
59+ 
60+ aclrtMemcpy(srcDevice, srcFileSize, srcHost, srcFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
61+ LaunchTTRANS3D<T, dstDC1HW, dstN1, dstN0, dstC0, srcN, srcC, srcD, srcH, srcW>(dstDevice, srcDevice, stream);
62+ 
63+ aclrtSynchronizeStream(stream);
64+ aclrtMemcpy(dstHost, dstFileSize, dstDevice, dstFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
65+ 
66+ WriteFile(GetGoldenDir() + "/output.bin", dstHost, dstFileSize);
67+ 
68+ aclrtFree(dstDevice);
69+ aclrtFree(srcDevice);
70+ 
71+ aclrtFreeHost(dstHost);
72+ aclrtFreeHost(srcHost);
73+ aclrtDestroyStream(stream);
74+ aclrtResetDevice(0);
75+ aclFinalize();
76+ 
77+ std::vector<T> golden(dstFileSize / sizeof(T));
78+ std::vector<T> result(dstFileSize / sizeof(T));
79+ ReadFile(GetGoldenDir() + "/golden.bin", dstFileSize, golden.data(), dstFileSize);
80+ ReadFile(GetGoldenDir() + "/output.bin", dstFileSize, result.data(), dstFileSize);
81+ 
82+ bool ret = ResultCmp(golden, result, 0.001f);
83+ 
84+ EXPECT_TRUE(ret);
85+}
86+ 
87+TEST_F(TTRANS3DTest, case1_float32_2_4_2_2_2)
88+{
89+ test_ttrans_3d<float, 8, 1, 16, 8, 2, 4, 2, 2, 2>();
90+}
91+ 
92+TEST_F(TTRANS3DTest, case2_float32_4_5_2_2_4)
93+{
94+ test_ttrans_3d<float, 16, 1, 16, 8, 4, 5, 2, 2, 4>();
95+}
96+ 
97+TEST_F(TTRANS3DTest, case3_int32_17_3_3_2_2)
98+{
99+ test_ttrans_3d<int32_t, 12, 2, 16, 8, 17, 3, 3, 2, 2>();
100+}
101+ 
102+TEST_F(TTRANS3DTest, case4_half_5_6_2_2_4)
103+{
104+ test_ttrans_3d<aclFloat16, 16, 1, 16, 16, 5, 6, 2, 2, 4>();
105+}
106+ 
107+TEST_F(TTRANS3DTest, case5_half_19_14_2_4_2)
108+{
109+ test_ttrans_3d<aclFloat16, 16, 2, 16, 16, 19, 14, 2, 4, 2>();
110+}
111+ 
112+TEST_F(TTRANS3DTest, case6_int16_8_13_2_3_4)
113+{
114+ test_ttrans_3d<int16_t, 24, 1, 16, 16, 8, 13, 2, 3, 4>();
115+}
116+ 
117+TEST_F(TTRANS3DTest, case7_uint16_4_8_2_2_3)
118+{
119+ test_ttrans_3d<uint16_t, 12, 1, 16, 16, 4, 8, 2, 2, 3>();
120+}
121+ 
122+TEST_F(TTRANS3DTest, case8_int8_7_28_2_3_4)
123+{
124+ test_ttrans_3d<int8_t, 24, 1, 16, 32, 7, 28, 2, 3, 4>();
125+}
126+ 
127+TEST_F(TTRANS3DTest, case9_int8_16_26_3_3_4)
128+{
129+ test_ttrans_3d<int8_t, 36, 1, 16, 32, 16, 26, 3, 3, 4>();
130+}
131+ 
132+TEST_F(TTRANS3DTest, case10_uint8_9_18_2_2_4)
133+{
134+ test_ttrans_3d<uint8_t, 16, 1, 16, 32, 9, 18, 2, 2, 4>();
135+}
@@ -0,0 +1,131 @@
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/constants.hpp>
13+#include <pto/common/debug.h>
14+#include "acl/acl.h"
15+ 
16+using namespace pto;
17+ 
18+template <typename T, int dstDC1HW, int dstN1, int dstN0, int dstC0, int srcN, int srcC, int srcD, int srcH, int srcW>
19+__global__ AICORE void runTTRANS3D(__gm__ T __out__ *out, __gm__ T __in__ *src)
20+{
21+ static_assert(dstN1 == (srcN + dstN0 - 1) / dstN0);
22+ constexpr int dstC1 = (srcC + dstC0 - 1) / dstC0;
23+ static_assert(dstDC1HW == srcD * dstC1 * srcH * srcW);
24+ 
25+ constexpr int srcElemNum = srcN * srcC * srcD * srcH * srcW;
26+ constexpr int dstElemNum = dstDC1HW * dstN1 * dstN0 * dstC0;
27+ constexpr int paddedC = dstC1 * dstC0;
28+ constexpr int ncplaneElem = srcN * paddedC * srcH * srcW;
29+ constexpr int tmpElemNum = (srcD + 1) * ncplaneElem;
30+ constexpr unsigned yTileSizeElem = (sizeof(T) == 1) ? 32 : 16;
31+ constexpr int subTmpW = (dstC0 + yTileSizeElem - 1) / yTileSizeElem * yTileSizeElem;
32+ constexpr int subTmpElemNum = srcH * srcW * subTmpW;
33+ constexpr int tmpTotalElem = tmpElemNum + subTmpElemNum;
34+ constexpr int elemPerBlock = 32 / sizeof(T);
35+ constexpr int srcAlignedElem = (srcElemNum + elemPerBlock - 1) / elemPerBlock * elemPerBlock;
36+ constexpr int tmpAlignedElem = (tmpTotalElem + elemPerBlock - 1) / elemPerBlock * elemPerBlock;
37+ 
38+ constexpr int srcBuffer = srcAlignedElem * sizeof(T);
39+ constexpr int dstBuffer = dstElemNum * sizeof(T);
40+ constexpr int tmpBuffer = tmpAlignedElem * sizeof(T);
41+ 
42+ using ShapeDim5Src = Shape<1, 1, 1, 1, srcElemNum>;
43+ using StrideDim5Src = pto::Stride<srcElemNum, srcElemNum, srcElemNum, srcElemNum, 1>;
44+ using GlobalDataInSrc = GlobalTensor<T, ShapeDim5Src, StrideDim5Src>;
45+ 
46+ using ShapeDim5Dst = Shape<1, 1, 1, 1, dstElemNum>;
47+ using StrideDim5Dst = pto::Stride<dstElemNum, dstElemNum, dstElemNum, dstElemNum, 1>;
48+ using GlobalDataInDst = GlobalTensor<T, ShapeDim5Dst, StrideDim5Dst>;
49+ 
50+ using SrcFlatTileData = Tile<TileType::Vec, T, 1, srcAlignedElem, BLayout::RowMajor, 1, srcElemNum>;
51+ SrcFlatTileData src0Tile;
52+ TASSIGN(src0Tile, 0x0);
53+ using SrcConvTileData =
54+ ConvTile<TileType::Vec, T, srcAlignedElem, Layout::NCDHW, ConvTileShape<srcN, srcC, srcD, srcH, srcW>>;
55+ SrcConvTileData srcTile;
56+ static_assert(srcTile.totalDimCount == 5);
57+ TASSIGN(srcTile, 0x0);
58+#ifdef __PTO_AUTO__
59+ TRESHAPE(src0Tile, srcTile);
60+#endif
61+ 
62+ using DstFlatTileData = Tile<TileType::Vec, T, 1, dstElemNum, BLayout::RowMajor, 1, dstElemNum>;
63+ DstFlatTileData dst0Tile;
64+ TASSIGN(dst0Tile, 0x0 + srcBuffer);
65+ using DstConvTileData =
66+ ConvTile<TileType::Vec, T, dstElemNum, Layout::FRACTAL_Z_3D, ConvTileShape<dstDC1HW, dstN1, dstN0, dstC0>>;
67+ DstConvTileData dstTile;
68+ static_assert(dstTile.totalDimCount == 4);
69+ TASSIGN(dstTile, 0x0 + srcBuffer);
70+#ifdef __PTO_AUTO__
71+ TRESHAPE(dst0Tile, dstTile);
72+#endif
73+ using ZeroTileData =
74+ Tile<TileType::Vec, int32_t, 1, dstElemNum * sizeof(T) / 4, BLayout::RowMajor, 1, dstElemNum * sizeof(T) / 4>;
75+ ZeroTileData dst1Tile;
76+ TASSIGN(dst1Tile, 0x0 + srcBuffer);
77+#ifdef __PTO_AUTO__
78+ TRESHAPE(dst1Tile, dstTile);
79+#endif
80+ 
81+ using TmpTileData = Tile<TileType::Vec, T, 1, tmpAlignedElem, BLayout::RowMajor, 1, tmpAlignedElem>;
82+ TmpTileData tmpTile;
83+ TASSIGN(tmpTile, 0x0 + srcBuffer + dstBuffer);
84+ 
85+ GlobalDataInSrc srcGlobal(src);
86+ GlobalDataInDst dstGlobal(out);
87+ TLOAD(src0Tile, srcGlobal);
88+#ifndef __PTO_AUTO__
89+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
90+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
91+#endif
92+ TSUB(dst1Tile, dst1Tile, dst1Tile);
93+#ifndef __PTO_AUTO__
94+ set_flag(PIPE_V, PIPE_S, EVENT_ID0);
95+ wait_flag(PIPE_V, PIPE_S, EVENT_ID0);
96+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
97+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
98+#endif
99+ TTRANS(dstTile, srcTile, tmpTile);
100+#ifndef __PTO_AUTO__
101+ set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
102+ wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0);
103+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
104+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
105+#endif
106+ TSTORE(dstGlobal, dst0Tile);
107+}
108+ 
109+template <typename T, int dstDC1HW, int dstN1, int dstN0, int dstC0, int srcN, int srcC, int srcD, int srcH, int srcW>
110+void LaunchTTRANS3D(T *out, T *src, void *stream)
111+{
112+ if constexpr (std::is_same_v<T, aclFloat16>) {
113+ runTTRANS3D<half, dstDC1HW, dstN1, dstN0, dstC0, srcN, srcC, srcD, srcH, srcW>
114+ <<<1, nullptr, stream>>>((half *)(out), (half *)(src));
115+ } else {
116+ runTTRANS3D<T, dstDC1HW, dstN1, dstN0, dstC0, srcN, srcC, srcD, srcH, srcW><<<1, nullptr, stream>>>(out, src);
117+ }
118+}
119+ 
120+// NCDHW -> FRACTAL_Z_3D
121+template void LaunchTTRANS3D<float, 8, 1, 16, 8, 2, 4, 2, 2, 2>(float *out, float *src, void *stream);
122+template void LaunchTTRANS3D<float, 16, 1, 16, 8, 4, 5, 2, 2, 4>(float *out, float *src, void *stream);
123+template void LaunchTTRANS3D<int32_t, 12, 2, 16, 8, 17, 3, 3, 2, 2>(int32_t *out, int32_t *src, void *stream);
124+template void LaunchTTRANS3D<aclFloat16, 16, 1, 16, 16, 5, 6, 2, 2, 4>(aclFloat16 *out, aclFloat16 *src, void *stream);
125+template void LaunchTTRANS3D<aclFloat16, 16, 2, 16, 16, 19, 14, 2, 4, 2>(aclFloat16 *out, aclFloat16 *src,
126+ void *stream);
127+template void LaunchTTRANS3D<int16_t, 24, 1, 16, 16, 8, 13, 2, 3, 4>(int16_t *out, int16_t *src, void *stream);
128+template void LaunchTTRANS3D<uint16_t, 12, 1, 16, 16, 4, 8, 2, 2, 3>(uint16_t *out, uint16_t *src, void *stream);
129+template void LaunchTTRANS3D<int8_t, 24, 1, 16, 32, 7, 28, 2, 3, 4>(int8_t *out, int8_t *src, void *stream);
130+template void LaunchTTRANS3D<int8_t, 36, 1, 16, 32, 16, 26, 3, 3, 4>(int8_t *out, int8_t *src, void *stream);
131+template void LaunchTTRANS3D<uint8_t, 16, 1, 16, 32, 9, 18, 2, 2, 4>(uint8_t *out, uint8_t *src, void *stream);
@@ -258,6 +258,10 @@ if [ "$ENABLE_A3" = "true" ]; then # A2A3
258 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather -g MGATHERTest.case_elem2d_int16_4x32_256size258 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather -g MGATHERTest.case_elem2d_int16_4x32_256size
259 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather -g MGATHERTest.case_row_dyn_int32_3x16_8rows259 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather -g MGATHERTest.case_row_dyn_int32_3x16_8rows
260 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather -g MGATHERTest.case_elem2d_nz_int32_16x8_1blk260 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather -g MGATHERTest.case_elem2d_nz_int32_16x8_1blk
261+ python3 tests/script/run_st.py $ARGS -w -v a3 -t ttrans_3d -g TTRANS3DTest.case3_int32_17_3_3_2_2
262+ python3 tests/script/run_st.py $ARGS -w -v a3 -t ttrans_3d -g TTRANS3DTest.case1_float32_2_4_2_2_2
263+ python3 tests/script/run_st.py $ARGS -w -v a3 -t ttrans_3d -g TTRANS3DTest.case7_uint16_4_8_2_2_3
264+ python3 tests/script/run_st.py $ARGS -w -v a3 -t ttrans_3d -g TTRANS3DTest.case10_uint8_9_18_2_2_4
261 265 
262 if [ "$IS_AUTO_MODE" = "false" ]; then266 if [ "$IS_AUTO_MODE" = "false" ]; then
263 # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode;267 # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode;
@@ -363,6 +367,7 @@ if [ "$ENABLE_A3" = "true" ]; then # A2A3
363 python3 tests/script/run_st.py $ARGS -w -v a3 -t tinsert_vec367 python3 tests/script/run_st.py $ARGS -w -v a3 -t tinsert_vec
364 python3 tests/script/run_st.py $ARGS -w -v a3 -t mscatter368 python3 tests/script/run_st.py $ARGS -w -v a3 -t mscatter
365 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather369 python3 tests/script/run_st.py $ARGS -w -v a3 -t mgather
370+ python3 tests/script/run_st.py $ARGS -w -v a3 -t ttrans_3d
366 if [ "$IS_AUTO_MODE" = "false" ]; then371 if [ "$IS_AUTO_MODE" = "false" ]; then
367 # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode;372 # this testcase has to directly call CCE intrinsics now, which won't compile for auto mode;
368 # besides, auto-sync doesn't work with CCE intrisics373 # besides, auto-sync doesn't work with CCE intrisics
@@ -727,4 +732,4 @@ if [ "$ENABLE_COMM" = "true" ]; then
727 python3 tests/script/run_st.py $ARGS -v a3 -t comm/tput_async732 python3 tests/script/run_st.py $ARGS -v a3 -t comm/tput_async
728 python3 tests/script/run_st.py $ARGS -v a3 -t comm/tget_async733 python3 tests/script/run_st.py $ARGS -v a3 -t comm/tget_async
729 fi734 fi
730-fi735+fi