已合并
GridSample2D算子新增bf16数据类型支持 #715
qiuyi创建于 4月14日
GridSample2D算子新增bf16数据类型支持 #715
已合并
qiuyi创建于 4月14日
8 个文件变更+293-95
@@ -50,8 +50,6 @@ static const int64_t SUPPORT_CHANNEL_310P = 32;
50 50 
51// 根据API定义,需要列出所能支持的所有dtype51// 根据API定义,需要列出所能支持的所有dtype
52static const std::initializer_list<op::DataType> DTYPE_SUPPORT_LIST = {52static const std::initializer_list<op::DataType> DTYPE_SUPPORT_LIST = {
53- op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_DOUBLE};
54-static const std::initializer_list<op::DataType> DTYPE_SUPPORT_LIST_REGBASE = {
55 op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16, op::DataType::DT_DOUBLE};53 op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16, op::DataType::DT_DOUBLE};
R
Rrenruhai4月16日

这两个dtype_support_list完全一样了,没必要分开了吧

likedislike
56 54 
57static bool CheckNotNull(const aclTensor *input, const aclTensor *grid, const aclTensor *out)55static bool CheckNotNull(const aclTensor *input, const aclTensor *grid, const aclTensor *out)
@@ -82,13 +80,16 @@ static bool CheckDtypeValid(const aclTensor *input, const aclTensor *grid, const
82 // 检查input、grid、out的数据类型是否一致80 // 检查input、grid、out的数据类型是否一致
83 OP_CHECK_DTYPE_NOT_MATCH(grid, input->GetDataType(), return false);81 OP_CHECK_DTYPE_NOT_MATCH(grid, input->GetDataType(), return false);
84 OP_CHECK_DTYPE_NOT_MATCH(out, input->GetDataType(), return false);82 OP_CHECK_DTYPE_NOT_MATCH(out, input->GetDataType(), return false);
83+ auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
85 84 
86 // 检查input的数据类型是否在gridsampler2d算子的支持列表内85 // 检查input的数据类型是否在gridsampler2d算子的支持列表内
87- if (IsRegBase()) {86+ if (curArch == NpuArch::DAV_2002 && input->GetDataType() == op::DataType::DT_BF16) {
88- OP_CHECK_DTYPE_NOT_SUPPORT(input, DTYPE_SUPPORT_LIST_REGBASE, return false);87+ OP_LOGD("input dtype does not support bf16 on this chip.");
88+ return false;
89 } else {89 } else {
90 OP_CHECK_DTYPE_NOT_SUPPORT(input, DTYPE_SUPPORT_LIST, return false);90 OP_CHECK_DTYPE_NOT_SUPPORT(input, DTYPE_SUPPORT_LIST, return false);
91 }91 }
92+
92 return true;93 return true;
93}94}
94 95 
@@ -228,8 +229,9 @@ static bool CheckAiCoreSuppport(const aclTensor *input, int64_t interpolationMod
228 }229 }
229 230 
230 const auto &inputShape = input->GetViewShape();231 const auto &inputShape = input->GetViewShape();
231- if (input->GetDataType() != op::DataType::DT_FLOAT && input->GetDataType() != op::DataType::DT_FLOAT16) {232+ if (input->GetDataType() != op::DataType::DT_FLOAT && input->GetDataType() != op::DataType::DT_FLOAT16 &&
232- OP_LOGD("Only support float16 or float32 on AICore, but got data type is %s",233+ input->GetDataType() != op::DataType::DT_BF16) {
234+ OP_LOGD("Only support float16, bfloat16 or float32 on AICore, but got data type is %s",
233 op::ToString(input->GetDataType()).GetString());235 op::ToString(input->GetDataType()).GetString());
234 return false;236 return false;
235 }237 }
@@ -84,13 +84,24 @@ ge::graphStatus GridSampleTiling::GetShapeAttrsInfo()
84 auto compileInfo = reinterpret_cast<const GridSampleCompileInfo *>(context_->GetCompileInfo());84 auto compileInfo = reinterpret_cast<const GridSampleCompileInfo *>(context_->GetCompileInfo());
85 OP_CHECK_NULL_WITH_CONTEXT(context_, compileInfo);85 OP_CHECK_NULL_WITH_CONTEXT(context_, compileInfo);
86 regBase = compileInfo->regBase;86 regBase = compileInfo->regBase;
87+ auto ascendc_platform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo());
88+ platform_ascendc::SocVersion gridSampleSocVersion = ascendc_platform.GetSocVersion();
89+ bool is310P = gridSampleSocVersion == platform_ascendc::SocVersion::ASCEND310P;
87 90 
88- OP_CHECK_IF((!regBase && dimension == 0 && xDtype != ge::DT_FLOAT && xDtype != ge::DT_FLOAT16),91+ OP_CHECK_IF((is310P && dimension == 0 && xDtype != ge::DT_FLOAT && xDtype != ge::DT_FLOAT16),
89 OP_LOGE(context_->GetNodeName(), "x datatype only support FLOAT32 or FLOAT16"),92 OP_LOGE(context_->GetNodeName(), "x datatype only support FLOAT32 or FLOAT16"),
90 return ge::GRAPH_FAILED);93 return ge::GRAPH_FAILED);
91- OP_CHECK_IF((!regBase && dimension == 0 && gridDtype != ge::DT_FLOAT && gridDtype != ge::DT_FLOAT16),94+ OP_CHECK_IF((is310P && dimension == 0 && gridDtype != ge::DT_FLOAT && gridDtype != ge::DT_FLOAT16),
92 OP_LOGE(context_->GetNodeName(), "grid datatype only support FLOAT32 or FLOAT16"),95 OP_LOGE(context_->GetNodeName(), "grid datatype only support FLOAT32 or FLOAT16"),
93 return ge::GRAPH_FAILED);96 return ge::GRAPH_FAILED);
97+ OP_CHECK_IF((!regBase && dimension == 0 && xDtype != ge::DT_FLOAT && xDtype != ge::DT_FLOAT16 &&
98+ xDtype != ge::DT_BF16),
99+ OP_LOGE(context_->GetNodeName(), "x datatype only support FLOAT32, FLOAT16, BFLOAT16"),
R
Rrenruhai4月16日

aclnn里判断了310P不支持BF16呀,这里校验是不是丢了

likedislike
100+ return ge::GRAPH_FAILED);
101+ OP_CHECK_IF((!regBase && dimension == 0 && gridDtype != ge::DT_FLOAT && gridDtype != ge::DT_FLOAT16 &&
102+ gridDtype != ge::DT_BF16),
103+ OP_LOGE(context_->GetNodeName(), "grid datatype only support FLOAT32, FLOAT16, BFLOAT16"),
104+ return ge::GRAPH_FAILED);
94 OP_CHECK_IF((regBase && dimension == 0 && xDtype != ge::DT_FLOAT && xDtype != ge::DT_FLOAT16 && xDtype != ge::DT_BF16),105 OP_CHECK_IF((regBase && dimension == 0 && xDtype != ge::DT_FLOAT && xDtype != ge::DT_FLOAT16 && xDtype != ge::DT_BF16),
95 OP_LOGE(context_->GetNodeName(), "x datatype only support FLOAT32, FLOAT16, BFLOAT16"),106 OP_LOGE(context_->GetNodeName(), "x datatype only support FLOAT32, FLOAT16, BFLOAT16"),
96 return ge::GRAPH_FAILED);107 return ge::GRAPH_FAILED);
@@ -169,9 +180,7 @@ ge::graphStatus GridSampleTiling::GetShapeAttrsInfo()
169 (inC * inH * inW <= X_MAX_HWC_FACTOR)) {180 (inC * inH * inW <= X_MAX_HWC_FACTOR)) {
170 tempType = FULL_LOAD_TYPE;181 tempType = FULL_LOAD_TYPE;
171 hwFactor = TILING_HW_FACTOR;182 hwFactor = TILING_HW_FACTOR;
172- auto ascendc_platform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo());183+ if ((outH * outW < BLOCK_NUM) && (inN < coreNumVar * BLOCK_NUM) && is310P) {
173- platform_ascendc::SocVersion gridSampleSocVersion = ascendc_platform.GetSocVersion();
174- if ((outH * outW < BLOCK_NUM) && (inN < coreNumVar * BLOCK_NUM) && (gridSampleSocVersion == platform_ascendc::SocVersion::ASCEND310P)) {
175 context_->SetScheduleMode(SCHEDULE_MODE);184 context_->SetScheduleMode(SCHEDULE_MODE);
176 }185 }
177 OP_LOGD(context_->GetNodeName(), "Get in FullLoad Template.");186 OP_LOGD(context_->GetNodeName(), "Get in FullLoad Template.");
@@ -94,6 +94,11 @@ extern "C" __global__ __aicore__ void grid_sample(GM_ADDR x, GM_ADDR grid, GM_AD
94 GridSample::GridSampler2DNearest<half> op;94 GridSample::GridSampler2DNearest<half> op;
95 op.Init(x, grid, y, userWS, &tilingData, pipe);95 op.Init(x, grid, y, userWS, &tilingData, pipe);
96 op.Process();96 op.Process();
97+ } else if (TILING_KEY_IS(1000231) || TILING_KEY_IS(1001231)) {
98+ // 2D nearest bf16 normal
99+ GridSample::GridSampler2DNearest<bfloat16_t> op;
100+ op.Init(x, grid, y, userWS, &tilingData, pipe);
101+ op.Process();
97 } else if (TILING_KEY_IS(1000222) || TILING_KEY_IS(1001222)) {102 } else if (TILING_KEY_IS(1000222) || TILING_KEY_IS(1001222)) {
98 // 2D Bicubic fp32 normal103 // 2D Bicubic fp32 normal
99 GridSample::GridSamplerBicubic2D<float> op;104 GridSample::GridSamplerBicubic2D<float> op;
@@ -104,6 +109,11 @@ extern "C" __global__ __aicore__ void grid_sample(GM_ADDR x, GM_ADDR grid, GM_AD
104 GridSample::GridSamplerBicubic2D<half> op;109 GridSample::GridSamplerBicubic2D<half> op;
105 op.Init(x, grid, y, userWS, &tilingData, pipe);110 op.Init(x, grid, y, userWS, &tilingData, pipe);
106 op.Process();111 op.Process();
112+ } else if (TILING_KEY_IS(1000232) || TILING_KEY_IS(1001232)) {
113+ // 2D Bicubic bf16 normal
114+ GridSample::GridSamplerBicubic2D<bfloat16_t> op;
115+ op.Init(x, grid, y, userWS, &tilingData, pipe);
116+ op.Process();
107 } else if (TILING_KEY_IS(1001220)) {117 } else if (TILING_KEY_IS(1001220)) {
108 // 2D Bilinear fp32 slide window118 // 2D Bilinear fp32 slide window
109 GridSample::GridSampler2DSlideWindow<float> op;119 GridSample::GridSampler2DSlideWindow<float> op;
@@ -114,6 +124,11 @@ extern "C" __global__ __aicore__ void grid_sample(GM_ADDR x, GM_ADDR grid, GM_AD
114 GridSample::GridSampler2DFP16SlideWindow<half> op;124 GridSample::GridSampler2DFP16SlideWindow<half> op;
115 op.Init(x, grid, y, userWS, &tilingData, pipe);125 op.Init(x, grid, y, userWS, &tilingData, pipe);
116 op.Process();126 op.Process();
127+ } else if (TILING_KEY_IS(1000230) || TILING_KEY_IS(1001230)) {
128+ // 2D Bilinear bf16 sliceWindow
129+ GridSample::GridSampler2DFP16SlideWindow<bfloat16_t> op;
130+ op.Init(x, grid, y, userWS, &tilingData, pipe);
131+ op.Process();
117 } else if (TILING_KEY_IS(2000220) || TILING_KEY_IS(2001220)) {132 } else if (TILING_KEY_IS(2000220) || TILING_KEY_IS(2001220)) {
118 // 2D Bilinear fp32 fullLoad general133 // 2D Bilinear fp32 fullLoad general
119 GridSample::GridSampler2DFullLoad<float, 0> op;134 GridSample::GridSampler2DFullLoad<float, 0> op;
@@ -124,6 +139,11 @@ extern "C" __global__ __aicore__ void grid_sample(GM_ADDR x, GM_ADDR grid, GM_AD
124 GridSample::GridSampler2DFullLoad<half, 0> op;139 GridSample::GridSampler2DFullLoad<half, 0> op;
125 op.Init(x, grid, y, userWS, &tilingData, pipe);140 op.Init(x, grid, y, userWS, &tilingData, pipe);
126 op.Process();141 op.Process();
142+ } else if (TILING_KEY_IS(2000230) || TILING_KEY_IS(2001230)) {
143+ // 2D Bilinear bf16 fullLoad general
144+ GridSample::GridSampler2DFullLoad<bfloat16_t, 0> op;
145+ op.Init(x, grid, y, userWS, &tilingData, pipe);
146+ op.Process();
127 } else if (TILING_KEY_IS(2100220) || TILING_KEY_IS(2101220)) {147 } else if (TILING_KEY_IS(2100220) || TILING_KEY_IS(2101220)) {
128 // 2D Bilinear fp32 fullLoad C=1 and small input148 // 2D Bilinear fp32 fullLoad C=1 and small input
129 GridSample::GridSampler2DFullLoad<float, 1> op;149 GridSample::GridSampler2DFullLoad<float, 1> op;
@@ -134,6 +154,11 @@ extern "C" __global__ __aicore__ void grid_sample(GM_ADDR x, GM_ADDR grid, GM_AD
134 GridSample::GridSampler2DFullLoad<half, 1> op;154 GridSample::GridSampler2DFullLoad<half, 1> op;
135 op.Init(x, grid, y, userWS, &tilingData, pipe);155 op.Init(x, grid, y, userWS, &tilingData, pipe);
136 op.Process();156 op.Process();
157+ } else if (TILING_KEY_IS(2100230) || TILING_KEY_IS(2101230)) {
158+ // 2D Bilinear bf16 fullLoad C=1 and small input
159+ GridSample::GridSampler2DFullLoad<bfloat16_t, 1> op;
160+ op.Init(x, grid, y, userWS, &tilingData, pipe);
161+ op.Process();
137 } else if (TILING_KEY_IS(2200220) || TILING_KEY_IS(2201220)) {162 } else if (TILING_KEY_IS(2200220) || TILING_KEY_IS(2201220)) {
138 // 2D Bilinear fp32 fullLoad C=32 and large input163 // 2D Bilinear fp32 fullLoad C=32 and large input
139 GridSample::GridSampler2DFullLoad<float, 2> op;164 GridSample::GridSampler2DFullLoad<float, 2> op;
@@ -144,6 +169,11 @@ extern "C" __global__ __aicore__ void grid_sample(GM_ADDR x, GM_ADDR grid, GM_AD
144 GridSample::GridSampler2DFullLoad<half, 2> op;169 GridSample::GridSampler2DFullLoad<half, 2> op;
145 op.Init(x, grid, y, userWS, &tilingData, pipe);170 op.Init(x, grid, y, userWS, &tilingData, pipe);
146 op.Process();171 op.Process();
172+ } else if (TILING_KEY_IS(2200230) || TILING_KEY_IS(2201230)) {
173+ // 2D Bilinear bf16 fullLoad C=32 and large input
174+ GridSample::GridSampler2DFullLoad<bfloat16_t, 2> op;
175+ op.Init(x, grid, y, userWS, &tilingData, pipe);
176+ op.Process();
147 } else if (TILING_KEY_IS(1010320)) {177 } else if (TILING_KEY_IS(1010320)) {
148 // 3D Bilinear fp32 normal178 // 3D Bilinear fp32 normal
149 GridSample::GridSampler3D<float> op;179 GridSample::GridSampler3D<float> op;
@@ -77,6 +77,10 @@ private:
77 int32_t calCElems, int32_t loopOffset, LocalTensor<float> coeffTy, LocalTensor<float> interp1dUb);77 int32_t calCElems, int32_t loopOffset, LocalTensor<float> coeffTy, LocalTensor<float> interp1dUb);
78 __aicore__ inline void MTE3ForNCHW(int64_t gmYBaseOffset, int32_t calCElems, int64_t calHwNum, int32_t loopElems,78 __aicore__ inline void MTE3ForNCHW(int64_t gmYBaseOffset, int32_t calCElems, int64_t calHwNum, int32_t loopElems,
79 LocalTensor<float> interp1dUb, GlobalTensor<float> dstGm, int32_t interp1dIdx);79 LocalTensor<float> interp1dUb, GlobalTensor<float> dstGm, int32_t interp1dIdx);
80+ __aicore__ inline void CubicZeroWeight(LocalTensor<float> weightTx, LocalTensor<float> coeffTx,
81+ LocalTensor<uint8_t> weightMaskUb, LocalTensor<uint64_t> maskUbTmp, int32_t loopIdx);
82+ __aicore__ inline void interp1dCompute(int64_t outBaseOffset, int32_t calCElems, int32_t loopOffset,
83+ int32_t cIdx, int32_t loopElems, int32_t interp1dIdx, LocalTensor<float> coeffTy, LocalTensor<float> interp1dUb);
80 __aicore__ inline void CubicInterp1d(int32_t nIdx, int64_t outBaseOffset, int32_t loopIdx, int32_t loopOffset,84 __aicore__ inline void CubicInterp1d(int32_t nIdx, int64_t outBaseOffset, int32_t loopIdx, int32_t loopOffset,
81 int32_t loopElems, LocalTensor<int32_t> coordinatesUb, LocalTensor<float> coeffTx, LocalTensor<float> coeffTy,85 int32_t loopElems, LocalTensor<int32_t> coordinatesUb, LocalTensor<float> coeffTx, LocalTensor<float> coeffTy,
82 LocalTensor<uint8_t> weightMaskUb, int32_t cIdx, int32_t calCElems, LocalTensor<float> interp1dUb,86 LocalTensor<uint8_t> weightMaskUb, int32_t cIdx, int32_t calCElems, LocalTensor<float> interp1dUb,
@@ -743,8 +747,8 @@ __aicore__ inline void GridSamplerBicubic2D<T>::OutTransposeFp32(
743 LocalTensor<float> dstList[16];747 LocalTensor<float> dstList[16];
744 LocalTensor<float> srcList[16];748 LocalTensor<float> srcList[16];
745 749 
746- event_t eventVS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));750+ event_t eventS_V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
747- event_t eventSV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));751+ event_t eventV_S = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
748 752 
749 TransDataTo5HDParams transDataParams;753 TransDataTo5HDParams transDataParams;
750 transDataParams.dstHighHalf = false;754 transDataParams.dstHighHalf = false;
@@ -754,39 +758,39 @@ __aicore__ inline void GridSamplerBicubic2D<T>::OutTransposeFp32(
754 transDataParams.dstRepStride = 2;758 transDataParams.dstRepStride = 2;
755 transDataParams.srcRepStride = 16;759 transDataParams.srcRepStride = 16;
756 760 
757- for (int32_t i = 0; i < 16; i++) {761+ for (int32_t iValue = 0; iValue < 16; iValue++) {
758- srcList[i] = xLocal[i * 8];762+ srcList[iValue] = xLocal[iValue * 8];
759 }763 }
760 764 
761- for (int32_t i = 0; i < 8; i++) {765+ for (int32_t iValue = 0; iValue < 8; iValue++) {
762- dstList[i * 2] = outValueUb[i * TRANSE_REP_STRIDE];766+ dstList[iValue * 2] = outValueUb[iValue * TRANSE_REP_STRIDE];
763- dstList[i * 2 + 1] = outValueUb[i * TRANSE_REP_STRIDE + 8];767+ dstList[iValue * 2 + 1] = outValueUb[iValue * TRANSE_REP_STRIDE + 8];
764 }768 }
765 769 
766- SetFlag<HardEvent::S_V>(eventSV);770+ SetFlag<HardEvent::S_V>(eventS_V);
767- WaitFlag<HardEvent::S_V>(eventSV);771+ WaitFlag<HardEvent::S_V>(eventS_V);
768 TransDataTo5HD<float>(dstList, srcList, transDataParams);772 TransDataTo5HD<float>(dstList, srcList, transDataParams);
769- SetFlag<HardEvent::V_S>(eventVS);773+ SetFlag<HardEvent::V_S>(eventV_S);
770- WaitFlag<HardEvent::V_S>(eventVS);774+ WaitFlag<HardEvent::V_S>(eventV_S);
771 } else if (channelAlign <= 64) {775 } else if (channelAlign <= 64) {
772 transDataParams.repeatTimes = channelAlign / 8;776 transDataParams.repeatTimes = channelAlign / 8;
773 transDataParams.dstRepStride = TRANSE_REP_STRIDE;777 transDataParams.dstRepStride = TRANSE_REP_STRIDE;
774 transDataParams.srcRepStride = 1;778 transDataParams.srcRepStride = 1;
775- for (int32_t j = 0; j < 8; j++) {779+ for (int32_t jVal = 0; jVal < 8; jVal++) {
776 for (int32_t i = 0; i < 16; i++) {780 for (int32_t i = 0; i < 16; i++) {
777- srcList[i] = xLocal[i * channelAlign + j * 16 * channelAlign];781+ srcList[i] = xLocal[i * channelAlign + jVal * 16 * channelAlign];
778 }782 }
779 783 
780 for (int32_t i = 0; i < 8; i++) {784 for (int32_t i = 0; i < 8; i++) {
781- dstList[i * 2] = outValueUb[i * TRANSE_REP_STRIDE + j * 16];785+ dstList[i * 2] = outValueUb[i * TRANSE_REP_STRIDE + jVal * 16];
782- dstList[i * 2 + 1] = outValueUb[i * TRANSE_REP_STRIDE + 8 + j * 16];786+ dstList[i * 2 + 1] = outValueUb[i * TRANSE_REP_STRIDE + 8 + jVal * 16];
783 }787 }
784 788 
785- SetFlag<HardEvent::S_V>(eventSV);789+ SetFlag<HardEvent::S_V>(eventS_V);
786- WaitFlag<HardEvent::S_V>(eventSV);790+ WaitFlag<HardEvent::S_V>(eventS_V);
787 TransDataTo5HD<float>(dstList, srcList, transDataParams);791 TransDataTo5HD<float>(dstList, srcList, transDataParams);
788- SetFlag<HardEvent::V_S>(eventVS);792+ SetFlag<HardEvent::V_S>(eventV_S);
789- WaitFlag<HardEvent::V_S>(eventVS);793+ WaitFlag<HardEvent::V_S>(eventV_S);
790 }794 }
791 }795 }
792}796}
@@ -897,6 +901,45 @@ __aicore__ inline void GridSamplerBicubic2D<T>::MTE3ForNCHW(int64_t gmYBaseOffse
897 }901 }
898}902}
899 903 
904+template <typename T>
905+__aicore__ inline void GridSamplerBicubic2D<T>::CubicZeroWeight(LocalTensor<float> weightTx, LocalTensor<float> coeffTx,
906+ LocalTensor<uint8_t> weightMaskUb, LocalTensor<uint64_t> maskUbTmp, int32_t loopIdx)
907+{
908+ weightTx = coeffTmpBuf_.Get<float>(CAL_H_W_BLOCK);
909+ CoordinatesSelectScalar(coeffTx, weightTx, weightMaskUb, 0.0f);
910+ 
911+ auto weightMaskUbTmp = weightMaskUb.ReinterpretCast<uint64_t>();
912+ int32_t maskOffset = loopIdx * 2;
913+ event_t eventSV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
914+ SetFlag<HardEvent::S_V>(eventSV);
915+ WaitFlag<HardEvent::S_V>(eventSV);
916+ maskUbTmp.SetValue(0, weightMaskUbTmp.GetValue(maskOffset));
917+ maskUbTmp.SetValue(1, weightMaskUbTmp.GetValue(maskOffset + 1));
918+}
919+ 
920+template <typename T>
921+__aicore__ inline void GridSamplerBicubic2D<T>::interp1dCompute(int64_t outBaseOffset, int32_t calCElems, int32_t loopOffset,
922+ int32_t cIdx, int32_t loopElems, int32_t interp1dIdx, LocalTensor<float> coeffTy, LocalTensor<float> interp1dUb)
923+{
924+ ApplyCoeffTy(calCElems, loopOffset, coeffTy, interp1dUb);
925+ 
926+ event_t eventIdVToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
927+ SetFlag<HardEvent::V_MTE3>(eventIdVToMte3);
928+ WaitFlag<HardEvent::V_MTE3>(eventIdVToMte3);
929+ 
930+ if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
931+ int64_t gmYOffset = CAL_H_W_BLOCK * inputC_ * blockIDX + loopOffset + cIdx * CHANNEL_BLOCK * CAL_H_W_BLOCK;
932+ MTE3ForNCHW(gmYOffset, calCElems, CAL_H_W_BLOCK, loopElems, interp1dUb, gmWorkspace_, interp1dIdx);
933+ } else {
934+ int64_t gmYOffset = outBaseOffset + loopOffset + cIdx * CHANNEL_BLOCK * gridHW_;
935+ MTE3ForNCHW(gmYOffset, calCElems, gridHW_, loopElems, interp1dUb, gmY_, interp1dIdx);
936+ }
937+ 
938+ event_t eventMte3V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
939+ SetFlag<HardEvent::MTE3_V>(eventMte3V);
940+ WaitFlag<HardEvent::MTE3_V>(eventMte3V);
941+}
942+ 
900template <typename T>943template <typename T>
901__aicore__ inline void GridSamplerBicubic2D<T>::CubicInterp1d(int32_t nIdx, int64_t outBaseOffset, int32_t loopIdx,944__aicore__ inline void GridSamplerBicubic2D<T>::CubicInterp1d(int32_t nIdx, int64_t outBaseOffset, int32_t loopIdx,
902 int32_t loopOffset, int32_t loopElems, LocalTensor<int32_t> coordinatesUb, LocalTensor<float> coeffTx,945 int32_t loopOffset, int32_t loopElems, LocalTensor<int32_t> coordinatesUb, LocalTensor<float> coeffTx,
@@ -908,19 +951,13 @@ __aicore__ inline void GridSamplerBicubic2D<T>::CubicInterp1d(int32_t nIdx, int6
908 951 
909 LocalTensor<float> weightTx = coeffTx;952 LocalTensor<float> weightTx = coeffTx;
910 if (paddingMode_ == PADDING_MODE_ZEROS) {953 if (paddingMode_ == PADDING_MODE_ZEROS) {
911- weightTx = coeffTmpBuf_.Get<float>(CAL_H_W_BLOCK);954+ CubicZeroWeight(weightTx, coeffTx, weightMaskUb, maskUbTmp, loopIdx);
912- CoordinatesSelectScalar(coeffTx, weightTx, weightMaskUb, 0.0f);
913- 
914- auto weightMaskUbTmp = weightMaskUb.ReinterpretCast<uint64_t>();
915- int32_t maskOffset = loopIdx * 2;
916- event_t eventSV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
917- SetFlag<HardEvent::S_V>(eventSV);
918- WaitFlag<HardEvent::S_V>(eventSV);
919- maskUbTmp.SetValue(0, weightMaskUbTmp.GetValue(maskOffset));
920- maskUbTmp.SetValue(1, weightMaskUbTmp.GetValue(maskOffset + 1));
921 }955 }
922 956 
923 LocalTensor<T> xLocal = xBuf_.AllocTensor<T>();957 LocalTensor<T> xLocal = xBuf_.AllocTensor<T>();
958+ if constexpr (IsSameType<T, bfloat16_t>::value) { // T: fp16
959+ xLocal = yFp16Buf_.AllocTensor<T>();
960+ }
924 int32_t channelAlign = Ceil(calCElems, B32_ALIGN_FACTOR) * B32_ALIGN_FACTOR;961 int32_t channelAlign = Ceil(calCElems, B32_ALIGN_FACTOR) * B32_ALIGN_FACTOR;
925 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {962 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
926 channelAlign = Ceil(calCElems, B16_ALIGN_FACTOR) * B16_ALIGN_FACTOR;963 channelAlign = Ceil(calCElems, B16_ALIGN_FACTOR) * B16_ALIGN_FACTOR;
@@ -935,7 +972,12 @@ __aicore__ inline void GridSamplerBicubic2D<T>::CubicInterp1d(int32_t nIdx, int6
935 SetFlag<HardEvent::MTE2_V>(eventMte2V);972 SetFlag<HardEvent::MTE2_V>(eventMte2V);
936 WaitFlag<HardEvent::MTE2_V>(eventMte2V);973 WaitFlag<HardEvent::MTE2_V>(eventMte2V);
937 974 
938- if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) { // T: fp16975+ if constexpr (IsSameType<T, bfloat16_t>::value) { // T: bf16
976+ LocalTensor<float> xFp32Ub = xBuf_.Get<float>();
977+ Cast(xFp32Ub, xLocal, RoundMode::CAST_NONE, channelAlign * TRANSE_REP_STRIDE);
978+ PipeBarrier<PIPE_V>();
979+ OutTransposeFp32(channelAlign, xFp32Ub, outValueUb);
980+ } else if constexpr (IsSameType<T, half>::value) { // T: fp16
939 LocalTensor<T> yFp16Ub = yFp16Buf_.Get<T>();981 LocalTensor<T> yFp16Ub = yFp16Buf_.Get<T>();
940 OutTransposeFp16(channelAlign, xLocal, yFp16Ub);982 OutTransposeFp16(channelAlign, xLocal, yFp16Ub);
941 PipeBarrier<PIPE_V>();983 PipeBarrier<PIPE_V>();
@@ -948,11 +990,7 @@ __aicore__ inline void GridSamplerBicubic2D<T>::CubicInterp1d(int32_t nIdx, int6
948 if (paddingMode_ == PADDING_MODE_ZEROS) {990 if (paddingMode_ == PADDING_MODE_ZEROS) {
949 for (size_t i = 0; i < calCElems; i++) {991 for (size_t i = 0; i < calCElems; i++) {
950 int32_t ubOffset = i * TRANSE_REP_STRIDE;992 int32_t ubOffset = i * TRANSE_REP_STRIDE;
951- Select(outValueUb[ubOffset],993+ Select(outValueUb[ubOffset], maskUbTmp, outValueUb[ubOffset], 0.0f, SELMODE::VSEL_TENSOR_SCALAR_MODE,
952- maskUbTmp,
953- outValueUb[ubOffset],
954- 0.0f,
955- SELMODE::VSEL_TENSOR_SCALAR_MODE,
956 TRANSE_REP_STRIDE);994 TRANSE_REP_STRIDE);
957 }995 }
958 PipeBarrier<PIPE_V>();996 PipeBarrier<PIPE_V>();
@@ -961,23 +999,7 @@ __aicore__ inline void GridSamplerBicubic2D<T>::CubicInterp1d(int32_t nIdx, int6
961 ApplyCoeffTx(calCElems, loopOffset, weightTx, outValueUb, interp1dUb, interp1dIdx);999 ApplyCoeffTx(calCElems, loopOffset, weightTx, outValueUb, interp1dUb, interp1dIdx);
962 1000 
963 if (interp1dIdx % 4 == 3) {1001 if (interp1dIdx % 4 == 3) {
964- ApplyCoeffTy(calCElems, loopOffset, coeffTy, interp1dUb);1002+ interp1dCompute(outBaseOffset, calCElems, loopOffset, cIdx, loopElems, interp1dIdx, coeffTy, interp1dUb);
965- 
966- event_t eventIdVToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
967- SetFlag<HardEvent::V_MTE3>(eventIdVToMte3);
968- WaitFlag<HardEvent::V_MTE3>(eventIdVToMte3);
969- 
970- if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
971- int64_t gmYOffset = CAL_H_W_BLOCK * inputC_ * blockIDX + loopOffset + cIdx * CHANNEL_BLOCK * CAL_H_W_BLOCK;
972- MTE3ForNCHW(gmYOffset, calCElems, CAL_H_W_BLOCK, loopElems, interp1dUb, gmWorkspace_, interp1dIdx);
973- } else {
974- int64_t gmYOffset = outBaseOffset + loopOffset + cIdx * CHANNEL_BLOCK * gridHW_;
975- MTE3ForNCHW(gmYOffset, calCElems, gridHW_, loopElems, interp1dUb, gmY_, interp1dIdx);
976- }
977- 
978- event_t eventMte3V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
979- SetFlag<HardEvent::MTE3_V>(eventMte3V);
980- WaitFlag<HardEvent::MTE3_V>(eventMte3V);
981 }1003 }
982}1004}
983 1005 
@@ -61,6 +61,8 @@ private:
61 __aicore__ inline void MTE2ForNHWC(int32_t nIdx, int32_t cIdx, int32_t calCElems, int32_t channelAlign,61 __aicore__ inline void MTE2ForNHWC(int32_t nIdx, int32_t cIdx, int32_t calCElems, int32_t channelAlign,
62 int32_t loopOffset, int32_t loopElems, LocalTensor<int32_t> coorUb, LocalTensor<T> xLocal);62 int32_t loopOffset, int32_t loopElems, LocalTensor<int32_t> coorUb, LocalTensor<T> xLocal);
63 __aicore__ inline void OutTranspose(int32_t channelAlign, LocalTensor<T> xLocal, LocalTensor<T> outValueUb);63 __aicore__ inline void OutTranspose(int32_t channelAlign, LocalTensor<T> xLocal, LocalTensor<T> outValueUb);
64+ __aicore__ inline void OutTransposeBf16(int32_t channelAlign, LocalTensor<float> xLocal,
65+ LocalTensor<float> outValueUb);
64 __aicore__ inline void MTE3ForNCHW(int32_t nIdx, int32_t cIdx, int32_t calCElems, int32_t channelAlign,66 __aicore__ inline void MTE3ForNCHW(int32_t nIdx, int32_t cIdx, int32_t calCElems, int32_t channelAlign,
65 int32_t hwIdx, int32_t loopOffset, int32_t loopElems, int64_t outBaseOffset, LocalTensor<float> weightUb,67 int32_t hwIdx, int32_t loopOffset, int32_t loopElems, int64_t outBaseOffset, LocalTensor<float> weightUb,
66 LocalTensor<float> outValueUb, bool isAutomicAdd);68 LocalTensor<float> outValueUb, bool isAutomicAdd);
@@ -795,6 +797,61 @@ __aicore__ inline void GridSampler2DFP16SlideWindow<T>::OutTranspose(
795 }797 }
796}798}
797 799 
800+template <typename T>
801+__aicore__ inline void GridSampler2DFP16SlideWindow<T>::OutTransposeBf16(
802+ int32_t channelAlign, LocalTensor<float> xLocal, LocalTensor<float> outValueUb)
803+{
804+ LocalTensor<float> dstList[16];
805+ LocalTensor<float> srcList[16];
806+ 
807+ event_t event_vs = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
808+ event_t event_sv = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
809+ 
810+ TransDataTo5HDParams transDataParams;
811+ transDataParams.dstHighHalf = false;
812+ transDataParams.srcHighHalf = false;
813+ if (channelAlign == 8) {
814+ transDataParams.repeatTimes = 8;
815+ transDataParams.dstRepStride = 2;
816+ transDataParams.srcRepStride = 16;
817+ 
818+ for (int32_t iVal = 0; iVal < 16; iVal++) {
819+ srcList[iVal] = xLocal[iVal * 8];
820+ }
821+ 
822+ for (int32_t iVal = 0; iVal < 8; iVal++) {
823+ dstList[iVal * 2] = outValueUb[iVal * TRANSE_REP_STRIDE];
824+ dstList[iVal * 2 + 1] = outValueUb[iVal * TRANSE_REP_STRIDE + 8];
825+ }
826+ 
827+ SetFlag<HardEvent::S_V>(event_sv);
828+ WaitFlag<HardEvent::S_V>(event_sv);
829+ TransDataTo5HD<float>(dstList, srcList, transDataParams);
830+ SetFlag<HardEvent::V_S>(event_vs);
831+ WaitFlag<HardEvent::V_S>(event_vs);
832+ } else if (channelAlign <= 64) {
833+ transDataParams.repeatTimes = channelAlign / 8;
834+ transDataParams.dstRepStride = TRANSE_REP_STRIDE;
835+ transDataParams.srcRepStride = 1;
836+ for (int32_t j = 0; j < 8; j++) {
837+ for (int32_t iVal = 0; iVal < 16; iVal++) {
838+ srcList[iVal] = xLocal[iVal * channelAlign + j * 16 * channelAlign];
839+ }
840+ 
841+ for (int32_t iVal = 0; iVal < 8; iVal++) {
842+ dstList[iVal * 2] = outValueUb[iVal * TRANSE_REP_STRIDE + j * 16];
843+ dstList[iVal * 2 + 1] = outValueUb[iVal * TRANSE_REP_STRIDE + 8 + j * 16];
844+ }
845+ 
846+ SetFlag<HardEvent::S_V>(event_sv);
847+ WaitFlag<HardEvent::S_V>(event_sv);
848+ TransDataTo5HD<float>(dstList, srcList, transDataParams);
849+ SetFlag<HardEvent::V_S>(event_vs);
850+ WaitFlag<HardEvent::V_S>(event_vs);
851+ }
852+ }
853+}
854+ 
798template <typename T>855template <typename T>
799__aicore__ inline void GridSampler2DFP16SlideWindow<T>::MTE3ForNCHW(int32_t nIdx, int32_t cIdx, int32_t calCElems,856__aicore__ inline void GridSampler2DFP16SlideWindow<T>::MTE3ForNCHW(int32_t nIdx, int32_t cIdx, int32_t calCElems,
800 int32_t channelAlign, int32_t hwIdx, int32_t loopOffset, int32_t loopElems, int64_t outBaseOffset,857 int32_t channelAlign, int32_t hwIdx, int32_t loopOffset, int32_t loopElems, int64_t outBaseOffset,
@@ -887,7 +944,12 @@ __aicore__ inline void GridSampler2DFP16SlideWindow<T>::PointBilinear(int32_t nI
887 maskUbTmpLocal.SetValue(2, weightMaskUbTmp.GetValue(maskOffset));944 maskUbTmpLocal.SetValue(2, weightMaskUbTmp.GetValue(maskOffset));
888 maskUbTmpLocal.SetValue(3, weightMaskUbTmp.GetValue(maskOffset + 1));945 maskUbTmpLocal.SetValue(3, weightMaskUbTmp.GetValue(maskOffset + 1));
889 auto weightMaskUbTmpfp32 = maskUbTmpLocal.ReinterpretCast<float>();946 auto weightMaskUbTmpfp32 = maskUbTmpLocal.ReinterpretCast<float>();
890- LocalTensor<T> xLocal = xBuf_.AllocTensor<T>();947+ LocalTensor<T> xLocal;
948+ if constexpr (IsSameType<T, bfloat16_t>::value) { // T: bf16
949+ xLocal = outValueFp16Buf_.AllocTensor<T>();
950+ } else {
951+ xLocal = xBuf_.AllocTensor<T>();
952+ }
891 // channel先按64大小循环953 // channel先按64大小循环
892 for (int32_t cIdx = 0; cIdx < channelLoop_; cIdx++) {954 for (int32_t cIdx = 0; cIdx < channelLoop_; cIdx++) {
893 int32_t calCElems = perLoopChannel_;955 int32_t calCElems = perLoopChannel_;
@@ -904,10 +966,17 @@ __aicore__ inline void GridSampler2DFP16SlideWindow<T>::PointBilinear(int32_t nI
904 SetFlag<HardEvent::MTE2_V>(eventMte2V);966 SetFlag<HardEvent::MTE2_V>(eventMte2V);
905 WaitFlag<HardEvent::MTE2_V>(eventMte2V);967 WaitFlag<HardEvent::MTE2_V>(eventMte2V);
906 968 
907- LocalTensor<T> outValueFp16Ub = outValueFp16Buf_.Get<T>();969+ if constexpr (IsSameType<T, bfloat16_t>::value) { // T: bf16
908- OutTranspose(channelAlign, xLocal, outValueFp16Ub);970+ LocalTensor<float> xLocalFp32Ub = xBuf_.Get<float>();
909- PipeBarrier<PIPE_V>();971+ Cast(xLocalFp32Ub, xLocal, RoundMode::CAST_NONE, channelAlign * TRANSE_REP_STRIDE);
910- Cast(outValueUb, outValueFp16Ub, RoundMode::CAST_NONE, calCElems * TRANSE_REP_STRIDE);972+ PipeBarrier<PIPE_V>();
973+ OutTransposeBf16(channelAlign, xLocalFp32Ub, outValueUb);
974+ } else { // T: fp16
975+ LocalTensor<T> outValueFp16Ub = outValueFp16Buf_.Get<T>();
976+ OutTranspose(channelAlign, xLocal, outValueFp16Ub);
977+ PipeBarrier<PIPE_V>();
978+ Cast(outValueUb, outValueFp16Ub, RoundMode::CAST_NONE, calCElems * TRANSE_REP_STRIDE);
979+ }
911 PipeBarrier<PIPE_V>();980 PipeBarrier<PIPE_V>();
912 if (calCElems >= 16) {981 if (calCElems >= 16) {
913 BinaryRepeatParams repParams{1, 1, 0, 8, 8, 0};982 BinaryRepeatParams repParams{1, 1, 0, 8, 8, 0};
@@ -997,16 +1066,22 @@ __aicore__ inline void GridSampler2DFP16SlideWindow<T>::PointBilinearXInLocal(in
997 }1066 }
998 }1067 }
999 1068 
1000- for (size_t i = 0; i < inputC_; i++) {1069+ if constexpr (IsSameType<T, bfloat16_t>::value) {
1001- auto ubOffset = i * CAL_H_W_BLOCK;1070+ Cast(outValueUb, outValueFp16Ub, RoundMode::CAST_NONE, inputC_ * CAL_H_W_BLOCK);
1002- if constexpr (IsSameType<T, bfloat16_t>::value) {1071+ PipeBarrier<PIPE_V>();
1003- Select(outValueFp16Ub[ubOffset],1072+ for (size_t i = 0; i < inputC_; i++) {
1073+ auto ubOffset = i * CAL_H_W_BLOCK;
1074+ Select(outValueUb[ubOffset],
1004 weightMaskUbTmp,1075 weightMaskUbTmp,
1005- outValueFp16Ub[ubOffset],1076+ outValueUb[ubOffset],
1006- ToBfloat16(0.0),1077+ (float)0.0,
1007 SELMODE::VSEL_TENSOR_SCALAR_MODE,1078 SELMODE::VSEL_TENSOR_SCALAR_MODE,
1008 CAL_H_W_BLOCK);1079 CAL_H_W_BLOCK);
1009- } else {1080+ }
1081+ PipeBarrier<PIPE_V>();
1082+ } else {
1083+ for (size_t i = 0; i < inputC_; i++) {
1084+ auto ubOffset = i * CAL_H_W_BLOCK;
1010 Select(outValueFp16Ub[ubOffset],1085 Select(outValueFp16Ub[ubOffset],
1011 weightMaskUbTmp,1086 weightMaskUbTmp,
1012 outValueFp16Ub[ubOffset],1087 outValueFp16Ub[ubOffset],
@@ -1014,11 +1089,10 @@ __aicore__ inline void GridSampler2DFP16SlideWindow<T>::PointBilinearXInLocal(in
1014 SELMODE::VSEL_TENSOR_SCALAR_MODE,1089 SELMODE::VSEL_TENSOR_SCALAR_MODE,
1015 CAL_H_W_BLOCK);1090 CAL_H_W_BLOCK);
1016 }1091 }
1092+ PipeBarrier<PIPE_V>();
1093+ Cast(outValueUb, outValueFp16Ub, RoundMode::CAST_NONE, inputC_ * CAL_H_W_BLOCK);
1094+ PipeBarrier<PIPE_V>();
1017 }1095 }
1018- PipeBarrier<PIPE_V>();
1019- 
1020- Cast(outValueUb, outValueFp16Ub, RoundMode::CAST_NONE, inputC_ * CAL_H_W_BLOCK);
1021- PipeBarrier<PIPE_V>();
1022 1096 
1023 int32_t trans_loop = Ceil(calHWElems, B32_MASK);1097 int32_t trans_loop = Ceil(calHWElems, B32_MASK);
1024 // 权重处理1098 // 权重处理
@@ -66,6 +66,8 @@ private:
66 __aicore__ inline void MTE3ForC32(GlobalTensor<float> gm_, int32_t calCElems, int32_t loopElems,66 __aicore__ inline void MTE3ForC32(GlobalTensor<float> gm_, int32_t calCElems, int32_t loopElems,
67 LocalTensor<float> weightUb, LocalTensor<float> outValueUb, bool isAutomicAdd);67 LocalTensor<float> weightUb, LocalTensor<float> outValueUb, bool isAutomicAdd);
68 __aicore__ inline void OutTranspose(int32_t channelAlign, LocalTensor<T> xLocal, LocalTensor<T> outValueUb);68 __aicore__ inline void OutTranspose(int32_t channelAlign, LocalTensor<T> xLocal, LocalTensor<T> outValueUb);
69+ __aicore__ inline void OutTransposeBf16(int32_t channelAlign, LocalTensor<half> xLocal,
70+ LocalTensor<half> outValueUb);
69 __aicore__ inline void PointBilinearForHalf(int32_t calHWElems, LocalTensor<int32_t> coordinatesUb, LocalTensor<float> weightUb, 71 __aicore__ inline void PointBilinearForHalf(int32_t calHWElems, LocalTensor<int32_t> coordinatesUb, LocalTensor<float> weightUb,
70 LocalTensor<uint8_t> weightMaskUb, LocalTensor<float> outValueUb, bool isAutomicAdd);72 LocalTensor<uint8_t> weightMaskUb, LocalTensor<float> outValueUb, bool isAutomicAdd);
71 __aicore__ inline void PointBilinear(int32_t nIdx, int32_t hwIdx, int32_t calHWElems,73 __aicore__ inline void PointBilinear(int32_t nIdx, int32_t hwIdx, int32_t calHWElems,
@@ -853,6 +855,41 @@ __aicore__ inline void GridSampler2DFullLoad<T, templateCNum>::OutTranspose(
853 }855 }
854}856}
855 857 
858+template <typename T, int templateCNum>
859+__aicore__ inline void GridSampler2DFullLoad<T, templateCNum>::OutTransposeBf16(
860+ int32_t channelAlign, LocalTensor<half> xLocal, LocalTensor<half> outValueUb)
861+{
862+ const int64_t TRANSE_REP_STRIDE = 512;
863+ LocalTensor<half> dstList[16];
864+ LocalTensor<half> srcList[16];
865+ 
866+ event_t eventVS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
867+ event_t eventSV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
868+ 
869+ TransDataTo5HDParams transDataParams;
870+ transDataParams.dstHighHalf = false;
871+ transDataParams.srcHighHalf = false;
872+ 
873+ if (channelAlign == 32 / sizeof(half)) {
874+ transDataParams.repeatTimes = 8 * 4;
875+ transDataParams.dstRepStride = sizeof(half) / 2;
876+ transDataParams.srcRepStride = 16;
877+ 
878+ for (int32_t i = 0; i < 16; i++) {
879+ srcList[i] = xLocal[i * 32 / sizeof(half)];
880+ }
881+ for (int32_t i = 0; i < 16; i++) {
882+ dstList[i] = outValueUb[i * TRANSE_REP_STRIDE];
883+ }
884+ 
885+ SetFlag<HardEvent::S_V>(eventSV);
886+ WaitFlag<HardEvent::S_V>(eventSV);
887+ TransDataTo5HD<half>(dstList, srcList, transDataParams);
888+ SetFlag<HardEvent::V_S>(eventVS);
889+ WaitFlag<HardEvent::V_S>(eventVS);
890+ }
891+}
892+ 
856template <typename T, int templateCNum>893template <typename T, int templateCNum>
857__aicore__ inline void GridSampler2DFullLoad<T, templateCNum>::PointBilinearForHalf(int32_t calHWElems, LocalTensor<int32_t> coordinatesUb, 894__aicore__ inline void GridSampler2DFullLoad<T, templateCNum>::PointBilinearForHalf(int32_t calHWElems, LocalTensor<int32_t> coordinatesUb,
858 LocalTensor<float> weightUb, LocalTensor<uint8_t> weightMaskUb, LocalTensor<float> outValueUb, bool isAutomicAdd)895 LocalTensor<float> weightUb, LocalTensor<uint8_t> weightMaskUb, LocalTensor<float> outValueUb, bool isAutomicAdd)
@@ -876,20 +913,25 @@ __aicore__ inline void GridSampler2DFullLoad<T, templateCNum>::PointBilinearForH
876 uint32_t srcBaseAddr = cIdx * perLoopChannel_ * sizeof(T) + (uint32_t)c_idx * sizeof(T);913 uint32_t srcBaseAddr = cIdx * perLoopChannel_ * sizeof(T) + (uint32_t)c_idx * sizeof(T);
877 Gather(outValueFP16Local[c_idx * calHWBlock], xLocal, coorUb, srcBaseAddr, calHWBlock);914 Gather(outValueFP16Local[c_idx * calHWBlock], xLocal, coorUb, srcBaseAddr, calHWBlock);
878 }915 }
879- 
880 PipeBarrier<PIPE_V>();916 PipeBarrier<PIPE_V>();
881- for (size_t i = 0; i < calCElems; i++) {917+ 
882- ubOffset = i * calHWBlock;918+ if constexpr (IsSameType<T, bfloat16_t>::value) {
883- if constexpr (IsSameType<T, bfloat16_t>::value) {919+ Cast(outValueUb, outValueFP16Local, RoundMode::CAST_NONE, calCElems * calHWBlock);
884- Select(outValueFP16Local[ubOffset], weightMaskUb, outValueFP16Local[ubOffset], ToBfloat16(0.0), SELMODE::VSEL_TENSOR_SCALAR_MODE, calHWBlock);920+ PipeBarrier<PIPE_V>();
885- } else {921+ for (size_t i = 0; i < calCElems; i++) {
922+ ubOffset = i * calHWBlock;
923+ Select(outValueUb[ubOffset], weightMaskUb, outValueUb[ubOffset], (float)0.0,
924+ SELMODE::VSEL_TENSOR_SCALAR_MODE, calHWBlock);
925+ }
926+ } else {
927+ for (size_t i = 0; i < calCElems; i++) {
928+ ubOffset = i * calHWBlock;
886 Select(outValueFP16Local[ubOffset], weightMaskUb, outValueFP16Local[ubOffset], half(0.0), SELMODE::VSEL_TENSOR_SCALAR_MODE, calHWBlock);929 Select(outValueFP16Local[ubOffset], weightMaskUb, outValueFP16Local[ubOffset], half(0.0), SELMODE::VSEL_TENSOR_SCALAR_MODE, calHWBlock);
887 }930 }
931+ PipeBarrier<PIPE_V>();
932+ Cast(outValueUb, outValueFP16Local, RoundMode::CAST_NONE, calCElems * calHWBlock);
888 }933 }
889 934 
890- PipeBarrier<PIPE_V>();
891- Cast(outValueUb, outValueFP16Local, RoundMode::CAST_NONE, calCElems * calHWBlock);
892- 
893 PipeBarrier<PIPE_V>();935 PipeBarrier<PIPE_V>();
894 MTE3ForNCHWToWorkSpace(cIdx, calCElems, loop_elems, weightUb, outValueUb, isAutomicAdd);936 MTE3ForNCHWToWorkSpace(cIdx, calCElems, loop_elems, weightUb, outValueUb, isAutomicAdd);
895 event_t eventMte3V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));937 event_t eventMte3V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
@@ -987,15 +1029,23 @@ __aicore__ inline void GridSampler2DFullLoad<T, templateCNum>::PointBilinearC32F
987 params);1029 params);
988 1030 
989 PipeBarrier<PIPE_V>();1031 PipeBarrier<PIPE_V>();
990- OutTranspose(32 / sizeof(T), tmpBufTotal[2 * 1024 / sizeof(uint16_t)], outValueFP16Local);1032+ LocalTensor<half> outValueFP16LocalHalf;
1033+ if constexpr (IsSameType<T, bfloat16_t>::value) { // T: bf16
1034+ outValueFP16LocalHalf = outValueFP16Local.template ReinterpretCast<half>();
1035+ LocalTensor<half> tmpBufHalf =
1036+ tmpBufTotal[2 * 1024 / sizeof(uint16_t)].template ReinterpretCast<half>();
1037+ OutTransposeBf16(32 / sizeof(T), tmpBufHalf, outValueFP16LocalHalf);
1038+ } else {
1039+ OutTranspose(32 / sizeof(T), tmpBufTotal[2 * 1024 / sizeof(uint16_t)], outValueFP16Local);
1040+ }
991 PipeBarrier<PIPE_V>();1041 PipeBarrier<PIPE_V>();
992 for (size_t i = 0; i < calCElems; i++) {1042 for (size_t i = 0; i < calCElems; i++) {
993 ubOffset = i * C32_H_W_BLOCK;1043 ubOffset = i * C32_H_W_BLOCK;
994 if constexpr (IsSameType<T, bfloat16_t>::value) {1044 if constexpr (IsSameType<T, bfloat16_t>::value) {
995- Select(outValueFP16Local[ubOffset],1045+ Select(outValueFP16LocalHalf[ubOffset],
996 weightMaskUb[HWLoop * C32_H_W_BLOCK / 8],1046 weightMaskUb[HWLoop * C32_H_W_BLOCK / 8],
997- outValueFP16Local[ubOffset],1047+ outValueFP16LocalHalf[ubOffset],
998- ToBfloat16(0.0),1048+ half(0.0),
999 SELMODE::VSEL_TENSOR_SCALAR_MODE,1049 SELMODE::VSEL_TENSOR_SCALAR_MODE,
1000 C32_H_W_BLOCK);1050 C32_H_W_BLOCK);
1001 } else {1051 } else {
@@ -839,7 +839,12 @@ __aicore__ inline void GridSampler2DNearest<T>::PointNearest(int32_t nIdx, int32
839 maskUbTmp.SetValue(0, weightMaskUbTmp.GetValue(maskOffset));839 maskUbTmp.SetValue(0, weightMaskUbTmp.GetValue(maskOffset));
840 maskUbTmp.SetValue(1, weightMaskUbTmp.GetValue(maskOffset + 1));840 maskUbTmp.SetValue(1, weightMaskUbTmp.GetValue(maskOffset + 1));
841 841 
842- LocalTensor<T> xLocal = xBuf_.AllocTensor<T>();842+ LocalTensor<T> xLocal;
843+ if constexpr (IsSameType<T, bfloat16_t>::value) { // T: bf16
844+ xLocal = yFp16Buf_.AllocTensor<T>();
845+ } else {
846+ xLocal = xBuf_.AllocTensor<T>();
847+ }
843 for (int32_t cIdx = 0; cIdx < channelLoop_; cIdx++) {848 for (int32_t cIdx = 0; cIdx < channelLoop_; cIdx++) {
844 int32_t calCElems = perLoopChannel_;849 int32_t calCElems = perLoopChannel_;
845 if (cIdx == channelLoop_ - 1) {850 if (cIdx == channelLoop_ - 1) {
@@ -858,7 +863,12 @@ __aicore__ inline void GridSampler2DNearest<T>::PointNearest(int32_t nIdx, int32
858 SetFlag<HardEvent::MTE2_V>(eventMte2V);863 SetFlag<HardEvent::MTE2_V>(eventMte2V);
859 WaitFlag<HardEvent::MTE2_V>(eventMte2V);864 WaitFlag<HardEvent::MTE2_V>(eventMte2V);
860 865 
861- if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) { // T: fp16866+ if constexpr (IsSameType<T, bfloat16_t>::value) { // T: bf16
867+ LocalTensor<float> xFp32Ub = xBuf_.Get<float>();
868+ Cast(xFp32Ub, xLocal, RoundMode::CAST_NONE, channelAlign * TRANSE_REP_STRIDE);
869+ PipeBarrier<PIPE_V>();
870+ OutTransposeFp32(channelAlign, xFp32Ub, outValueUb);
871+ } else if constexpr (IsSameType<T, half>::value) { // T: fp16
862 LocalTensor<T> yFp16Ub = yFp16Buf_.Get<T>();872 LocalTensor<T> yFp16Ub = yFp16Buf_.Get<T>();
863 OutTransposeFp16(channelAlign, xLocal, yFp16Ub);873 OutTransposeFp16(channelAlign, xLocal, yFp16Ub);
864 PipeBarrier<PIPE_V>();874 PipeBarrier<PIPE_V>();
@@ -910,6 +920,7 @@ __aicore__ inline void GridSampler2DNearest<T>::PointNearest(int32_t nIdx, int32
910 }920 }
911 }921 }
912}922}
923+ 
913template <typename T>924template <typename T>
914__aicore__ inline void GridSampler2DNearest<T>::CopyOutFp16(int32_t nIdx, int32_t hwIdx, int32_t calHWElems)925__aicore__ inline void GridSampler2DNearest<T>::CopyOutFp16(int32_t nIdx, int32_t hwIdx, int32_t calHWElems)
915{926{
@@ -132,7 +132,7 @@ TEST_F(l2_grid_sampler2d_test, case_5)
132 132 
133 uint64_t workspaceSize = 0;133 uint64_t workspaceSize = 0;
134 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);134 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
135- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);135+ EXPECT_EQ(aclRet, ACL_SUCCESS);
136}136}
137 137 
138// float16138// float16