已合并
fix: 修复 Gather 回绕及 BrcInline/Cast/Concat 截断(#256 #258 #259 #260) #1862
m0_56712298创建于 8月28日
fix: 修复 Gather 回绕及 BrcInline/Cast/Concat 截断(#256 #258 #259 #260) #1862
已合并
共 9 个文件变更+197-10
| @@ -28,7 +28,8 @@ inline __aicore__ void BinaryBrcInlineApiWithTwoVectorizedAxis( | |||
| 28 | int64_t elem_in_one_block = 32 / dtype_size; | 28 | int64_t elem_in_one_block = 32 / dtype_size; |
| 29 | int64_t cut_quotient = element / elem_in_one_repeat; | 29 | int64_t cut_quotient = element / elem_in_one_repeat; |
| 30 | int64_t cut_reminder = element - cut_quotient * elem_in_one_repeat; | 30 | int64_t cut_reminder = element - cut_quotient * elem_in_one_repeat; |
| 31 | - if (cut_quotient >= block) { | 31 | + int64_t repeat_stride_in_blocks = first_axis_v_stride / elem_in_one_block; |
| 32 | + if (cut_quotient >= block || repeat_stride_in_blocks > 255) { | ||
| 32 | // 将block层外抛作为for循环，原有的element整体使用counter模式 | 33 | // 将block层外抛作为for循环，原有的element整体使用counter模式 |
| 33 | for (int64_t outer_for = 0; outer_for < block; outer_for++) { | 34 | for (int64_t outer_for = 0; outer_for < block; outer_for++) { |
| 34 | FUNC1(dstLocal[outer_for * first_axis_v_stride], src0Local[outer_for * first_axis_v_stride * is_input0_block_brc], | 35 | FUNC1(dstLocal[outer_for * first_axis_v_stride], src0Local[outer_for * first_axis_v_stride * is_input0_block_brc], |
| @@ -40,7 +41,7 @@ inline __aicore__ void BinaryBrcInlineApiWithTwoVectorizedAxis( | |||
| 40 | constexpr uint8_t dst_block_stride = 1; | 41 | constexpr uint8_t dst_block_stride = 1; |
| 41 | constexpr uint8_t src0_block_stride = 1; | 42 | constexpr uint8_t src0_block_stride = 1; |
| 42 | constexpr uint8_t src1_block_stride = 1; | 43 | constexpr uint8_t src1_block_stride = 1; |
| 43 | - uint8_t dst_repeat_stride = first_axis_v_stride / elem_in_one_block; | 44 | + uint8_t dst_repeat_stride = static_cast<uint8_t>(repeat_stride_in_blocks); |
| 44 | uint8_t src0_repeat_stride = dst_repeat_stride * is_input0_block_brc; | 45 | uint8_t src0_repeat_stride = dst_repeat_stride * is_input0_block_brc; |
| 45 | uint8_t src1_repeat_stride = dst_repeat_stride * is_input1_block_brc; | 46 | uint8_t src1_repeat_stride = dst_repeat_stride * is_input1_block_brc; |
| 46 | uint32_t calcSize = 0; | 47 | uint32_t calcSize = 0; |
| @@ -131,8 +131,17 @@ inline __aicore__ void CastExtendWithMaskMode(const AscendC::LocalTensor<OutT> & | |||
| 131 | uint32_t repeat_reminder = repeat_times - repeat_throw_for_extent * MAX_REPEAT_TIMES; | 131 | uint32_t repeat_reminder = repeat_times - repeat_throw_for_extent * MAX_REPEAT_TIMES; |
| 132 | uint16_t dst_block_stride = 1; | 132 | uint16_t dst_block_stride = 1; |
| 133 | uint16_t src_block_stride = 1; | 133 | uint16_t src_block_stride = 1; |
| 134 | - uint8_t dst_repeat_stride = output_last_dim_stride * sizeof(OutT) / ONE_BLK_SIZE; | 134 | + uint32_t dst_repeat_stride_in_blocks = output_last_dim_stride / (ONE_BLK_SIZE / sizeof(OutT)); |
| 135 | - uint8_t src_repeat_stride = input_last_dim_stride * sizeof(InT) / ONE_BLK_SIZE; | 135 | + uint32_t src_repeat_stride_in_blocks = input_last_dim_stride / (ONE_BLK_SIZE / sizeof(InT)); |
| 136 | + if (dst_repeat_stride_in_blocks > MAX_REPEAT_TIMES || src_repeat_stride_in_blocks > MAX_REPEAT_TIMES) { | ||
| 137 | + for (uint32_t i = 0; i < repeat_times; ++i) { | ||
| 138 | + AscendC::Cast(dst[i * output_last_dim_stride], src[i * input_last_dim_stride], GetRoundMode<InT, OutT>(), | ||
| 139 | + last_dim); | ||
| 140 | + } | ||
| 141 | + return; | ||
| 142 | + } | ||
| 143 | + uint8_t dst_repeat_stride = static_cast<uint8_t>(dst_repeat_stride_in_blocks); | ||
| 144 | + uint8_t src_repeat_stride = static_cast<uint8_t>(src_repeat_stride_in_blocks); | ||
| 136 | AscendC::SetMaskNorm(); | 145 | AscendC::SetMaskNorm(); |
| 137 | if constexpr (sizeof(InT) > sizeof(OutT)) { | 146 | if constexpr (sizeof(InT) > sizeof(OutT)) { |
| 138 | AscendC::SetVectorMask<InT, MaskMode::NORMAL>(last_dim); | 147 | AscendC::SetVectorMask<InT, MaskMode::NORMAL>(last_dim); |
| @@ -18,6 +18,20 @@ using namespace AscendC; | |||
| 18 | 18 | ||
| 19 | enum class BinaryOp { Add, Sub, Div, Mul }; | 19 | enum class BinaryOp { Add, Sub, Div, Mul }; |
| 20 | 20 | ||
| 21 | +constexpr int64_t kLargeFirstAxisStride = 2048; | ||
| 22 | +constexpr float kCounterPathMarker = 1.0F; | ||
| 23 | +constexpr float kRepeatPathMarker = -1.0F; | ||
| 24 | + | ||
| 25 | +inline __aicore__ void MarkCounterPath(const LocalTensor<float> &dst, const LocalTensor<float> &, | ||
| 26 | + const LocalTensor<float> &, const int32_t &) { | ||
| 27 | + dst.SetValue(0, kCounterPathMarker); | ||
| 28 | +} | ||
| 29 | + | ||
| 30 | +inline __aicore__ void MarkRepeatPath(const LocalTensor<float> &dst, const LocalTensor<float> &, | ||
| 31 | + const LocalTensor<float> &, uint64_t, const uint8_t, const BinaryRepeatParams &) { | ||
| 32 | + dst.SetValue(0, kRepeatPathMarker); | ||
| 33 | +} | ||
| 34 | + | ||
| 21 | template <class T> | 35 | template <class T> |
| 22 | void UbToGm(T *gm, LocalTensor<T> &local, uint64_t size) { | 36 | void UbToGm(T *gm, LocalTensor<T> &local, uint64_t size) { |
| 23 | for (int i = 0; i < size; i++) { | 37 | for (int i = 0; i < size; i++) { |
| @@ -148,6 +162,35 @@ void TestBrcInlineTwoDimApi(std::vector<int64_t> &input1_shape, std::vector<int6 | |||
| 148 | EXPECT_EQ(diff_count, 0); | 162 | EXPECT_EQ(diff_count, 0); |
| 149 | } | 163 | } |
| 150 | 164 | ||
| 165 | +TEST(TestApiBroInline, TestLargeRepeatStrideFallsBackToCounterPath) { | ||
| 166 | + constexpr int64_t kOuterSize = 2; | ||
| 167 | + constexpr int64_t kInnerSize = 64; | ||
| 168 | + constexpr int64_t kBufferSize = kOuterSize * kLargeFirstAxisStride; | ||
| 169 | + constexpr int64_t kOutputSize = 2; | ||
| 170 | + auto *output = static_cast<float *>(AscendC::GmAlloc(kOutputSize * sizeof(float))); | ||
| 171 | + | ||
| 172 | + auto kernel = [](float *output) { | ||
| 173 | + TPipe pipe; | ||
| 174 | + TBuf<TPosition::VECCALC> output_buf; | ||
| 175 | + pipe.InitBuffer(output_buf, kBufferSize * sizeof(float)); | ||
| 176 | + LocalTensor<float> local_output = output_buf.Get<float>(); | ||
| 177 | + local_output.SetValue(0, 0.0F); | ||
| 178 | + local_output.SetValue(kLargeFirstAxisStride, 0.0F); | ||
| 179 | + BinaryBrcInlineApiWithTwoVectorizedAxis<float>(local_output, local_output, local_output, kOuterSize, kInnerSize, 0, | ||
| 180 | + 0, kLargeFirstAxisStride, sizeof(float), &MarkCounterPath, | ||
| 181 | + &MarkRepeatPath); | ||
| 182 | + output[0] = local_output.GetValue(0); | ||
| 183 | + output[1] = local_output.GetValue(kLargeFirstAxisStride); | ||
| 184 | + }; | ||
| 185 | + | ||
| 186 | + AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 187 | + ICPU_RUN_KF(kernel, 1, output); | ||
| 188 | + | ||
| 189 | + EXPECT_EQ(output[0], kCounterPathMarker); | ||
| 190 | + EXPECT_EQ(output[1], kCounterPathMarker); | ||
| 191 | + AscendC::GmFree(output); | ||
| 192 | +} | ||
| 193 | + | ||
| 151 | TEST(TestApiBroInline, Test_Add_Float_Two_Axis) { | 194 | TEST(TestApiBroInline, Test_Add_Float_Two_Axis) { |
| 152 | using AddCounterPtr = | 195 | using AddCounterPtr = |
| 153 | void (*)(const LocalTensor<float> &, const LocalTensor<float> &, const LocalTensor<float> &, const int32_t &); | 196 | void (*)(const LocalTensor<float> &, const LocalTensor<float> &, const LocalTensor<float> &, const int32_t &); |
| @@ -30,6 +30,55 @@ constexpr int gen_index_five = 5; | |||
| 30 | constexpr int gen_index_div = 1000; | 30 | constexpr int gen_index_div = 1000; |
| 31 | constexpr float gen_float_suffix = 0.12; | 31 | constexpr float gen_float_suffix = 0.12; |
| 32 | 32 | ||
| 33 | +TEST(TestApiCast, LargeDstRepeatStrideFallsBackToLoop) { | ||
| 34 | + constexpr uint32_t kFirstDim = 2; | ||
| 35 | + constexpr uint32_t kLastDim = 8; | ||
| 36 | + constexpr uint32_t kInputLastDimStride = 16; | ||
| 37 | + constexpr uint32_t kOutputLastDimStride = 2048; | ||
| 38 | + constexpr uint32_t kInputSize = kFirstDim * kInputLastDimStride; | ||
| 39 | + constexpr uint32_t kOutputSize = kFirstDim * kOutputLastDimStride; | ||
| 40 | + constexpr float kSentinel = -1.0F; | ||
| 41 | + auto *x = static_cast<half *>(AscendC::GmAlloc(sizeof(half) * kInputSize)); | ||
| 42 | + auto *y = static_cast<float *>(AscendC::GmAlloc(sizeof(float) * kOutputSize)); | ||
| 43 | + for (uint32_t i = 0; i < kInputSize; ++i) { | ||
| 44 | + x[i] = 0; | ||
| 45 | + } | ||
| 46 | + for (uint32_t i = 0; i < kLastDim; ++i) { | ||
| 47 | + x[i] = 1; | ||
| 48 | + x[kInputLastDimStride + i] = 2; | ||
| 49 | + } | ||
| 50 | + for (uint32_t i = 0; i < kOutputSize; ++i) { | ||
| 51 | + y[i] = kSentinel; | ||
| 52 | + } | ||
| 53 | + | ||
| 54 | + auto kernel = [](half *x, float *y) { | ||
| 55 | + TPipe tpipe; | ||
| 56 | + TBuf<TPosition::VECCALC> xbuf, ybuf, tmp; | ||
| 57 | + tpipe.InitBuffer(xbuf, sizeof(half) * kInputSize); | ||
| 58 | + tpipe.InitBuffer(ybuf, sizeof(float) * kOutputSize); | ||
| 59 | + tpipe.InitBuffer(tmp, 8192); | ||
| 60 | + auto l_x = xbuf.Get<half>(); | ||
| 61 | + auto l_y = ybuf.Get<float>(); | ||
| 62 | + auto l_tmp = tmp.Get<uint8_t>(); | ||
| 63 | + GmToUb(l_x, x, kInputSize); | ||
| 64 | + GmToUb(l_y, y, kOutputSize); | ||
| 65 | + CastExtend(l_y, l_x, l_tmp, kFirstDim, kLastDim, kInputLastDimStride, kOutputLastDimStride, sizeof(float)); | ||
| 66 | + UbToGm(y, l_y, kOutputSize); | ||
| 67 | + }; | ||
| 68 | + | ||
| 69 | + AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 70 | + ICPU_RUN_KF(kernel, 1, x, y); | ||
| 71 | + | ||
| 72 | + for (uint32_t i = 0; i < kLastDim; ++i) { | ||
| 73 | + EXPECT_EQ(y[i], 1.0F); | ||
| 74 | + EXPECT_EQ(y[kOutputLastDimStride + i], 2.0F); | ||
| 75 | + } | ||
| 76 | + EXPECT_EQ(y[kLastDim], kSentinel); | ||
| 77 | + EXPECT_EQ(y[kOutputLastDimStride + kLastDim], kSentinel); | ||
| 78 | + AscendC::GmFree(x); | ||
| 79 | + AscendC::GmFree(y); | ||
| 80 | +} | ||
| 81 | + | ||
| 33 | template <typename InT, typename OutT> | 82 | template <typename InT, typename OutT> |
| 34 | void CastExtendCalc(InT *x, OutT *y, int size) { | 83 | void CastExtendCalc(InT *x, OutT *y, int size) { |
| 35 | TPipe tpipe; | 84 | TPipe tpipe; |
| @@ -6,6 +6,7 @@ add_library(test_regbase_ascendc OBJECT | |||
| 6 | ./test_cast.cpp | 6 | ./test_cast.cpp |
| 7 | ./test_compare.cpp | 7 | ./test_compare.cpp |
| 8 | ./test_concat.cpp | 8 | ./test_concat.cpp |
| 9 | + ./test_gather.cpp | ||
| 9 | ./test_logical_and.cpp | 10 | ./test_logical_and.cpp |
| 10 | ./test_logical_not.cpp | 11 | ./test_logical_not.cpp |
| 11 | ./test_logical_or.cpp | 12 | ./test_logical_or.cpp |
| @@ -135,3 +135,26 @@ TEST_F(RegbaseApiConcatTest, ConcatSuccess) { | |||
| 135 | ConcatTest<uint8_t, 2>(16, {31, 33}); | 135 | ConcatTest<uint8_t, 2>(16, {31, 33}); |
| 136 | ConcatTest<uint8_t, 2>(16, {255, 257}); | 136 | ConcatTest<uint8_t, 2>(16, {255, 257}); |
| 137 | } | 137 | } |
| 138 | + | ||
| 139 | +TEST_F(RegbaseApiConcatTest, GenMaskRegPreservesHighBits) { | ||
| 140 | + constexpr uint32_t kMaskWordCount = 8; | ||
| 141 | + constexpr uint32_t kExpectedMask = 0x000fffffU; | ||
| 142 | + auto output = static_cast<uint32_t *>(AscendC::GmAlloc(kMaskWordCount * sizeof(uint32_t))); | ||
| 143 | + auto kernel = [output] { | ||
| 144 | + TPipe tpipe; | ||
| 145 | + TBuf<TPosition::VECCALC> index_buf; | ||
| 146 | + tpipe.InitBuffer(index_buf, kMaskWordCount * sizeof(uint32_t)); | ||
| 147 | + auto index = index_buf.Get<uint32_t>(); | ||
| 148 | + | ||
| 149 | + concat::GenMaskReg<uint16_t>(16, 16, 10, index.GetPhyAddr()); | ||
| 150 | + UbToGm(output, index, kMaskWordCount); | ||
| 151 | + }; | ||
| 152 | + | ||
| 153 | + AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 154 | + ICPU_RUN_KF(kernel, 1); | ||
| 155 | + | ||
| 156 | + for (uint32_t i = 0; i < kMaskWordCount; ++i) { | ||
| 157 | + EXPECT_EQ(output[i], kExpectedMask); | ||
| 158 | + } | ||
| 159 | + AscendC::GmFree(output); | ||
| 160 | +} | ||
| @@ -0,0 +1,59 @@ | |||
| 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 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +using namespace AscendC; | ||
| 17 | + | ||
| 18 | +namespace { | ||
| 19 | +constexpr uint32_t kBufferSize = 32; | ||
| 20 | +constexpr uint32_t kDstOffset = 1; | ||
| 21 | +constexpr uint64_t kCopyLength = 1; | ||
| 22 | +constexpr uint8_t kSentinel = 0xA5; | ||
| 23 | +constexpr uint8_t kSourceValue = 0x3C; | ||
| 24 | +} // namespace | ||
| 25 | + | ||
| 26 | +TEST(TestApiGather, DataCopySimdSimtSmallLengthWithinPadding) { | ||
| 27 | + auto *src = static_cast<uint8_t *>(AscendC::GmAlloc(kBufferSize)); | ||
| 28 | + auto *dst = static_cast<uint8_t *>(AscendC::GmAlloc(kBufferSize)); | ||
| 29 | + for (uint32_t i = 0; i < kBufferSize; ++i) { | ||
| 30 | + src[i] = kSourceValue + i; | ||
| 31 | + dst[i] = kSentinel; | ||
| 32 | + } | ||
| 33 | + | ||
| 34 | + auto kernel = [](uint8_t *src, uint8_t *dst) { | ||
| 35 | + TPipe pipe; | ||
| 36 | + TBuf<TPosition::VECCALC> dst_buf; | ||
| 37 | + pipe.InitBuffer(dst_buf, kBufferSize); | ||
| 38 | + LocalTensor<uint8_t> local_dst = dst_buf.Get<uint8_t>(); | ||
| 39 | + GmToUb(local_dst, dst, kBufferSize); | ||
| 40 | + | ||
| 41 | + GlobalTensor<uint8_t> global_src; | ||
| 42 | + global_src.SetGlobalBuffer((__gm__ uint8_t *)src); | ||
| 43 | + DataCopySimdSimt(local_dst, global_src, kDstOffset, 0, kCopyLength, false); | ||
| 44 | + UbToGm(dst, local_dst, kBufferSize); | ||
| 45 | + }; | ||
| 46 | + | ||
| 47 | + AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 48 | + ICPU_RUN_KF(kernel, 1, src, dst); | ||
| 49 | + | ||
| 50 | + EXPECT_EQ(dst[kDstOffset], kSourceValue); | ||
| 51 | + for (uint32_t i = 0; i < kBufferSize; ++i) { | ||
| 52 | + if (i != kDstOffset) { | ||
| 53 | + EXPECT_EQ(dst[i], kSentinel) << "unexpected write at index " << i; | ||
| 54 | + } | ||
| 55 | + } | ||
| 56 | + | ||
| 57 | + AscendC::GmFree(src); | ||
| 58 | + AscendC::GmFree(dst); | ||
| 59 | +} | ||
| @@ -253,7 +253,7 @@ __aicore__ inline AscendC::MicroAPI::MaskReg GenMaskReg(uint32_t block_size, uin | |||
| 253 | // 不对齐, 必然有tail | 253 | // 不对齐, 必然有tail |
| 254 | const uint16_t all_num = (gather_mask_repeat_stride / kDataBlockSize) - 1; | 254 | const uint16_t all_num = (gather_mask_repeat_stride / kDataBlockSize) - 1; |
| 255 | const uint16_t tail_num = gather_mask_dim_size % (32 / sizeof(U)); | 255 | const uint16_t tail_num = gather_mask_dim_size % (32 / sizeof(U)); |
| 256 | - const uint16_t tail_mask = (1U << (tail_num * sizeof(U))) - 1U; | 256 | + const uint32_t tail_mask = (1U << (tail_num * sizeof(U))) - 1U; |
| 257 | if (all_num == 0) { | 257 | if (all_num == 0) { |
| 258 | for (uint16_t i = 0; i < loop_times; ++i) { | 258 | for (uint16_t i = 0; i < loop_times; ++i) { |
| 259 | index_addr[i] = tail_mask; | 259 | index_addr[i] = tail_mask; |
| @@ -474,17 +474,19 @@ inline __aicore__ void DataCopySimdSimt(AscendC::LocalTensor<T1> &dst, const Asc | |||
| 474 | if ((dst_p * sizeof(T1)) % 32 != 0) { | 474 | if ((dst_p * sizeof(T1)) % 32 != 0) { |
| 475 | current_addr = dst_p * sizeof(T1); | 475 | current_addr = dst_p * sizeof(T1); |
| 476 | padding = (32 - (current_addr % 32)) % 32; | 476 | padding = (32 - (current_addr % 32)) % 32; |
| 477 | + uint32_t padding_elements = padding / sizeof(T1); | ||
| 478 | + uint32_t padding_copy_length = length <= padding_elements ? static_cast<uint32_t>(length) : padding_elements; | ||
| 477 | int32_t event_id_mte2_to_v = static_cast<int32_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::MTE2_V)); | 479 | int32_t event_id_mte2_to_v = static_cast<int32_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::MTE2_V)); |
| 478 | AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(event_id_mte2_to_v); | 480 | AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(event_id_mte2_to_v); |
| 479 | AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(event_id_mte2_to_v); | 481 | AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(event_id_mte2_to_v); |
| 480 | - AscendC::Simt::VF_CALL<GatherSimtContinuous<T1>>(AscendC::Simt::Dim3(128), dst_p1, x1_gm, padding / sizeof(T1), | 482 | + AscendC::Simt::VF_CALL<GatherSimtContinuous<T1>>(AscendC::Simt::Dim3(128), dst_p1, x1_gm, padding_copy_length, |
| 481 | dst_p, src_p, is_out); | 483 | dst_p, src_p, is_out); |
| 482 | - dst_p += padding / sizeof(T1); | 484 | + if (length <= padding_elements) { |
| 483 | - src_p += padding / sizeof(T1); | ||
| 484 | - length -= padding / sizeof(T1); | ||
| 485 | - if (length <= 0) { | ||
| 486 | return; | 485 | return; |
| 487 | } | 486 | } |
| 487 | + dst_p += padding_elements; | ||
| 488 | + src_p += padding_elements; | ||
| 489 | + length -= padding_elements; | ||
| 488 | if (unlikely(is_out)) { | 490 | if (unlikely(is_out)) { |
| 489 | T1 value{}; | 491 | T1 value{}; |
| 490 | AscendC::Duplicate(dst[dst_p], value, length); | 492 | AscendC::Duplicate(dst[dst_p], value, length); |