Merged
fix(math): preserve equal infinities in low precision logaddexp #6219
puppyAdd on Oct 11
fix(math): preserve equal infinities in low precision logaddexp #6219
Merged
Total: 5 files changed+270-3
| @@ -2,7 +2,7 @@ | |||
| 2 | # Copyright (c) 2026 Huawei Technologies Co., Ltd. | 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 | 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"). | 4 | # CANN Open Software License Agreement Version 2.0 (the "License"). |
| 5 | - | 5 | + |
| 6 | # Please refer to the License for details. You may not use this file except in compliance with the License. | 6 | # Please refer to the License for details. You may not use this file except in compliance with the License. |
| 7 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | 7 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, |
| 8 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | 8 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. |
| @@ -19,4 +19,8 @@ foreach(SUB_DIR ${CURRENT_DIRS}) | |||
| 19 | if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") | 19 | if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") |
| 20 | add_subdirectory(${SUB_DIR}) | 20 | add_subdirectory(${SUB_DIR}) |
| 21 | endif() | 21 | endif() |
| 22 | -endforeach() | 22 | +endforeach() |
| 23 | + | ||
| 24 | +if(ENABLE_TEST AND (UT_TEST_ALL OR OP_KERNEL_UT)) | ||
| 25 | + add_subdirectory(tests/ut/op_kernel) | ||
| 26 | +endif() | ||
| @@ -115,7 +115,9 @@ __aicore__ inline void LogAddExp<T_IN, T_COMPUTE>::Init(GM_ADDR x, GM_ADDR y, GM | |||
| 115 | // Buffer size must be 32-byte aligned for DataCopy/DataCopyPad | 115 | // Buffer size must be 32-byte aligned for DataCopy/DataCopyPad |
| 116 | constexpr int64_t ALIGN_BYTES = 32; | 116 | constexpr int64_t ALIGN_BYTES = 32; |
| 117 | constexpr int64_t ALIGN_ELEMS_IN = ALIGN_BYTES / sizeof(T_IN); | 117 | constexpr int64_t ALIGN_ELEMS_IN = ALIGN_BYTES / sizeof(T_IN); |
| 118 | - constexpr int64_t ALIGN_ELEMS_COMPUTE = ALIGN_BYTES / sizeof(T_COMPUTE); | 118 | + // The promoted path compares fp32 inputs in 256-byte repeats. |
| 119 | + constexpr int64_t ALIGN_BYTES_COMPUTE = std::is_same_v<T_IN, T_COMPUTE> ? ALIGN_BYTES : 256; | ||
| 120 | + constexpr int64_t ALIGN_ELEMS_COMPUTE = ALIGN_BYTES_COMPUTE / sizeof(T_COMPUTE); | ||
| 119 | int64_t alignedUbIn = ((ubLength_ + ALIGN_ELEMS_IN - 1) / ALIGN_ELEMS_IN) * ALIGN_ELEMS_IN; | 121 | int64_t alignedUbIn = ((ubLength_ + ALIGN_ELEMS_IN - 1) / ALIGN_ELEMS_IN) * ALIGN_ELEMS_IN; |
| 120 | int64_t alignedUbCompute = ((ubLength_ + ALIGN_ELEMS_COMPUTE - 1) / ALIGN_ELEMS_COMPUTE) * ALIGN_ELEMS_COMPUTE; | 122 | int64_t alignedUbCompute = ((ubLength_ + ALIGN_ELEMS_COMPUTE - 1) / ALIGN_ELEMS_COMPUTE) * ALIGN_ELEMS_COMPUTE; |
| 121 | 123 | ||
| @@ -320,10 +322,30 @@ __aicore__ inline void LogAddExp<T_IN, T_COMPUTE>::Compute(int64_t currentNum) | |||
| 320 | // fp16/bf16 path: cast to fp32, compute, cast back | 322 | // fp16/bf16 path: cast to fp32, compute, cast back |
| 321 | AscendC::LocalTensor<T_COMPUTE> tmp3 = tmpBuf3.Get<T_COMPUTE>(); | 323 | AscendC::LocalTensor<T_COMPUTE> tmp3 = tmpBuf3.Get<T_COMPUTE>(); |
| 322 | 324 | ||
| 325 | + constexpr int64_t COMPARE_ALIGN_ELEMS = 256 / sizeof(T_COMPUTE); | ||
| 326 | + int64_t compareNum = ((currentNum + COMPARE_ALIGN_ELEMS - 1) / COMPARE_ALIGN_ELEMS) * COMPARE_ALIGN_ELEMS; | ||
| 327 | + if (compareNum > currentNum) { | ||
| 328 | + // Compare reads a complete final repeat, including the padding. | ||
| 329 | + int64_t tailOffset = compareNum - COMPARE_ALIGN_ELEMS; | ||
| 330 | + AscendC::Duplicate(tmp1[tailOffset], static_cast<T_COMPUTE>(0.0f), COMPARE_ALIGN_ELEMS); | ||
| 331 | + AscendC::Duplicate(tmp3[tailOffset], static_cast<T_COMPUTE>(0.0f), COMPARE_ALIGN_ELEMS); | ||
| 332 | + if constexpr (std::is_same_v<T_IN, bfloat16_t>) { | ||
| 333 | + AscendC::Duplicate(tmp2[tailOffset], static_cast<T_COMPUTE>(0.0f), COMPARE_ALIGN_ELEMS); | ||
| 334 | + } | ||
| 335 | + } | ||
| 336 | + | ||
| 323 | AscendC::Cast(tmp1, xLocal, AscendC::RoundMode::CAST_NONE, currentNum); | 337 | AscendC::Cast(tmp1, xLocal, AscendC::RoundMode::CAST_NONE, currentNum); |
| 324 | AscendC::Cast(tmp3, yLocal, AscendC::RoundMode::CAST_NONE, currentNum); | 338 | AscendC::Cast(tmp3, yLocal, AscendC::RoundMode::CAST_NONE, currentNum); |
| 325 | AscendC::Max(tmp2, tmp1, tmp3, currentNum); | 339 | AscendC::Max(tmp2, tmp1, tmp3, currentNum); |
| 340 | + // Reuse the output buffer as a bitmap until the final cast. Equal | ||
| 341 | + // infinities need a zero difference; NaN inputs must remain NaN. | ||
| 342 | + AscendC::LocalTensor<uint8_t> unequalMask = outLocal.template ReinterpretCast<uint8_t>(); | ||
| 343 | + AscendC::Compare(unequalMask, tmp1, tmp3, AscendC::CMPMODE::EQ, compareNum); | ||
| 344 | + AscendC::Not(unequalMask.template ReinterpretCast<uint16_t>(), unequalMask.template ReinterpretCast<uint16_t>(), | ||
| 345 | + (compareNum + 15) / 16); | ||
| 326 | AscendC::Sub(tmp1, tmp1, tmp3, currentNum); | 346 | AscendC::Sub(tmp1, tmp1, tmp3, currentNum); |
| 347 | + AscendC::Select(tmp1, unequalMask, tmp1, static_cast<T_COMPUTE>(0.0f), | ||
| 348 | + AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, currentNum); | ||
| 327 | AscendC::Abs(tmp1, tmp1, currentNum); | 349 | AscendC::Abs(tmp1, tmp1, currentNum); |
| 328 | AscendC::Muls(tmp1, tmp1, static_cast<T_COMPUTE>(-1.0f), currentNum); | 350 | AscendC::Muls(tmp1, tmp1, static_cast<T_COMPUTE>(-1.0f), currentNum); |
| 329 | AscendC::Exp(tmp1, tmp1, currentNum); | 351 | AscendC::Exp(tmp1, tmp1, currentNum); |
| @@ -331,6 +353,33 @@ __aicore__ inline void LogAddExp<T_IN, T_COMPUTE>::Compute(int64_t currentNum) | |||
| 331 | AscendC::Ln(tmp1, tmp1, currentNum); | 353 | AscendC::Ln(tmp1, tmp1, currentNum); |
| 332 | AscendC::Add(tmp2, tmp2, tmp1, currentNum); | 354 | AscendC::Add(tmp2, tmp2, tmp1, currentNum); |
| 333 | AscendC::Cast(outLocal, tmp2, AscendC::RoundMode::CAST_ROUND, currentNum); | 355 | AscendC::Cast(outLocal, tmp2, AscendC::RoundMode::CAST_ROUND, currentNum); |
| 356 | + if constexpr (std::is_same_v<T_IN, bfloat16_t>) { | ||
| 357 | + // Keep nonfinite results independent of narrowing conversion behavior. | ||
| 358 | + // Select moves the BF16 encodings through half views without converting them. | ||
| 359 | + AscendC::LocalTensor<uint8_t> specialMask = tmp1.template ReinterpretCast<uint8_t>(); | ||
| 360 | + AscendC::LocalTensor<uint16_t> specialBits = tmp3.template ReinterpretCast<uint16_t>(); | ||
| 361 | + AscendC::LocalTensor<half> specialValues = tmp3.template ReinterpretCast<half>(); | ||
| 362 | + AscendC::LocalTensor<half> outBits = outLocal.template ReinterpretCast<half>(); | ||
| 363 | + int64_t maskBytes = (((compareNum + 7) / 8 + 31) / 32) * 32; | ||
| 364 | + AscendC::Duplicate(tmp1.template ReinterpretCast<uint16_t>(), static_cast<uint16_t>(0), maskBytes / 2); | ||
| 365 | + | ||
| 366 | + constexpr T_COMPUTE POS_INF = __builtin_huge_valf(); | ||
| 367 | + constexpr T_COMPUTE NEG_INF = -__builtin_huge_valf(); | ||
| 368 | + AscendC::CompareScalar(specialMask, tmp2, POS_INF, AscendC::CMPMODE::EQ, compareNum); | ||
| 369 | + AscendC::Duplicate(specialBits, static_cast<uint16_t>(0x7F80), currentNum); | ||
| 370 | + AscendC::Select(outBits, specialMask, specialValues, outBits, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, | ||
| 371 | + currentNum); | ||
| 372 | + AscendC::CompareScalar(specialMask, tmp2, NEG_INF, AscendC::CMPMODE::EQ, compareNum); | ||
| 373 | + AscendC::Duplicate(specialBits, static_cast<uint16_t>(0xFF80), currentNum); | ||
| 374 | + AscendC::Select(outBits, specialMask, specialValues, outBits, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, | ||
| 375 | + currentNum); | ||
| 376 | + AscendC::Compare(specialMask, tmp2, tmp2, AscendC::CMPMODE::EQ, compareNum); | ||
| 377 | + AscendC::Not(specialMask.template ReinterpretCast<uint16_t>(), | ||
| 378 | + specialMask.template ReinterpretCast<uint16_t>(), (compareNum + 15) / 16); | ||
| 379 | + AscendC::Duplicate(specialBits, static_cast<uint16_t>(0x7FC0), currentNum); | ||
| 380 | + AscendC::Select(outBits, specialMask, specialValues, outBits, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, | ||
| 381 | + currentNum); | ||
| 382 | + } | ||
| 334 | } | 383 | } |
| 335 | 384 | ||
| 336 | outputQueue.EnQue<T_IN>(outLocal); | 385 | outputQueue.EnQue<T_IN>(outLocal); |
| @@ -0,0 +1,51 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 2 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 3 | +# CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 4 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 5 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 6 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 7 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 8 | + | ||
| 9 | +function(add_log_add_exp_kernel_ut_case) | ||
| 10 | + if(NOT ASCEND_COMPUTE_UNIT STREQUAL "ascend910b") | ||
| 11 | + return() | ||
| 12 | + endif() | ||
| 13 | + list(FIND ASCEND_OP_NAME log_add_exp op_index) | ||
| 14 | + if(NOT "${ASCEND_OP_NAME}" STREQUAL "" AND op_index EQUAL -1) | ||
| 15 | + return() | ||
| 16 | + endif() | ||
| 17 | + | ||
| 18 | + set(soc_version Ascend910B1) | ||
| 19 | + set(case_target log_add_exp_${soc_version}_cases_obj) | ||
| 20 | + set(module_target ${OP_KERNEL_MODULE_NAME}_${soc_version}_cases_obj) | ||
| 21 | + add_library(${case_target} OBJECT | ||
| 22 | + ${CMAKE_CURRENT_SOURCE_DIR}/test_log_add_exp.cpp | ||
| 23 | + ${CMAKE_CURRENT_SOURCE_DIR}/log_add_exp_kernel_inst.cpp) | ||
| 24 | + target_compile_options(${case_target} PRIVATE -g -DUT_SOC_VERSION="${soc_version}") | ||
| 25 | + target_include_directories(${case_target} PRIVATE | ||
| 26 | + ${GTEST_INCLUDE} | ||
| 27 | + ${OP_TILING_INCLUDE} | ||
| 28 | + ${ASCEND_DIR}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/ | ||
| 29 | + ${ASCEND_DIR}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/impl/basic_api/ | ||
| 30 | + ${ASCEND_DIR}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/include/ | ||
| 31 | + ${ASCEND_DIR}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/include/basic_api/) | ||
| 32 | + target_link_libraries(${case_target} PRIVATE | ||
| 33 | + $<BUILD_INTERFACE:intf_llt_pub_asan_cxx17> tikicpulib::${soc_version} GTest::gtest) | ||
| 34 | + | ||
| 35 | + if(NOT TARGET ${module_target}) | ||
| 36 | + add_library(${module_target} OBJECT $<TARGET_OBJECTS:${case_target}>) | ||
| 37 | + else() | ||
| 38 | + target_sources(${module_target} PRIVATE $<TARGET_OBJECTS:${case_target}>) | ||
| 39 | + endif() | ||
| 40 | + target_link_libraries(${module_target} PRIVATE | ||
| 41 | + $<BUILD_INTERFACE:intf_llt_pub_asan_cxx17> $<TARGET_OBJECTS:${case_target}>) | ||
| 42 | + list(FIND fastOpTestSocVersions "${soc_version}" soc_index) | ||
| 43 | + if(soc_index EQUAL -1) | ||
| 44 | + set(fastOpTestSocVersions ${fastOpTestSocVersions} ${soc_version} | ||
| 45 | + CACHE STRING "fastOp Test SocVersions" FORCE) | ||
| 46 | + endif() | ||
| 47 | +endfunction() | ||
| 48 | + | ||
| 49 | +if(UT_TEST_ALL OR OP_KERNEL_UT) | ||
| 50 | + add_log_add_exp_kernel_ut_case() | ||
| 51 | +endif() | ||
| @@ -0,0 +1,31 @@ | |||
| 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 <cstdint> | ||
| 12 | +#include "../../../op_kernel/log_add_exp.h" | ||
| 13 | + | ||
| 14 | +// Keep kernel headers separate from gtest's host headers. | ||
| 15 | +template <typename T> | ||
| 16 | +__aicore__ inline void RunLogAddExp(GM_ADDR x, GM_ADDR y, GM_ADDR out, GM_ADDR tiling) | ||
| 17 | +{ | ||
| 18 | + NsLogAddExp::LogAddExp<T, float> op; | ||
| 19 | + op.Init(x, y, out, nullptr, reinterpret_cast<LogAddExpTilingData*>(tiling)); | ||
| 20 | + op.Process(); | ||
| 21 | +} | ||
| 22 | + | ||
| 23 | +extern "C" __global__ __aicore__ void log_add_exp_fp16_ut(GM_ADDR x, GM_ADDR y, GM_ADDR out, GM_ADDR tiling) | ||
| 24 | +{ | ||
| 25 | + RunLogAddExp<half>(x, y, out, tiling); | ||
| 26 | +} | ||
| 27 | + | ||
| 28 | +extern "C" __global__ __aicore__ void log_add_exp_bf16_ut(GM_ADDR x, GM_ADDR y, GM_ADDR out, GM_ADDR tiling) | ||
| 29 | +{ | ||
| 30 | + RunLogAddExp<bfloat16_t>(x, y, out, tiling); | ||
| 31 | +} | ||
| @@ -0,0 +1,132 @@ | |||
| 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 <array> | ||
| 12 | +#include <cmath> | ||
| 13 | +#include <cstdint> | ||
| 14 | +#include <cstring> | ||
| 15 | +#include <limits> | ||
| 16 | +#include <memory> | ||
| 17 | +#include "gtest/gtest.h" | ||
| 18 | +#include "tikicpulib.h" | ||
| 19 | +#include "../../../op_kernel/log_add_exp_tiling_data.h" | ||
| 20 | + | ||
| 21 | +extern "C" __global__ __aicore__ void log_add_exp_fp16_ut(GM_ADDR x, GM_ADDR y, GM_ADDR out, GM_ADDR tiling); | ||
| 22 | +extern "C" __global__ __aicore__ void log_add_exp_bf16_ut(GM_ADDR x, GM_ADDR y, GM_ADDR out, GM_ADDR tiling); | ||
| 23 | + | ||
| 24 | +namespace { | ||
| 25 | +struct GmDeleter { | ||
| 26 | + void operator()(uint8_t* ptr) const { AscendC::GmFree(ptr); } | ||
| 27 | +}; | ||
| 28 | +using GmBuffer = std::unique_ptr<uint8_t, GmDeleter>; | ||
| 29 | + | ||
| 30 | +uint16_t Encode(float value, bool bf16) | ||
| 31 | +{ | ||
| 32 | + if (bf16) { | ||
| 33 | + uint32_t bits; | ||
| 34 | + std::memcpy(&bits, &value, sizeof(bits)); | ||
| 35 | + return static_cast<uint16_t>(bits >> 16); | ||
| 36 | + } | ||
| 37 | + half value16 = static_cast<half>(value); | ||
| 38 | + uint16_t bits; | ||
| 39 | + std::memcpy(&bits, &value16, sizeof(bits)); | ||
| 40 | + return bits; | ||
| 41 | +} | ||
| 42 | + | ||
| 43 | +float Decode(uint16_t bits, bool bf16) | ||
| 44 | +{ | ||
| 45 | + if (bf16) { | ||
| 46 | + uint32_t bits32 = static_cast<uint32_t>(bits) << 16; | ||
| 47 | + float value; | ||
| 48 | + std::memcpy(&value, &bits32, sizeof(value)); | ||
| 49 | + return value; | ||
| 50 | + } | ||
| 51 | + half value16; | ||
| 52 | + std::memcpy(&value16, &bits, sizeof(bits)); | ||
| 53 | + return static_cast<float>(value16); | ||
| 54 | +} | ||
| 55 | + | ||
| 56 | +float Golden(float x, float y) | ||
| 57 | +{ | ||
| 58 | + if (std::isnan(x) || std::isnan(y)) { | ||
| 59 | + return std::numeric_limits<float>::quiet_NaN(); | ||
| 60 | + } | ||
| 61 | + if (x == y && std::isinf(x)) { | ||
| 62 | + return x; | ||
| 63 | + } | ||
| 64 | + return std::fmax(x, y) + std::log1p(std::exp(-std::fabs(x - y))); | ||
| 65 | +} | ||
| 66 | +} // namespace | ||
| 67 | + | ||
| 68 | +class LogAddExpSpecialValuesTest : public testing::TestWithParam<bool> { | ||
| 69 | +protected: | ||
| 70 | + void Run(int64_t count, int64_t ubFactor) | ||
| 71 | + { | ||
| 72 | + bool bf16 = GetParam(); | ||
| 73 | + float inf = std::numeric_limits<float>::infinity(); | ||
| 74 | + float nan = std::numeric_limits<float>::quiet_NaN(); | ||
| 75 | + const std::array<std::array<float, 2>, 24> pairs = { | ||
| 76 | + {{inf, inf}, {-inf, -inf}, {inf, -inf}, {-inf, inf}, {inf, 1.0f}, {1.0f, inf}, | ||
| 77 | + {-inf, 1.0f}, {1.0f, -inf}, {nan, inf}, {inf, nan}, {nan, -inf}, {-inf, nan}, | ||
| 78 | + {nan, nan}, {nan, 1.0f}, {1.0f, nan}, {nan, 0.0f}, {0.0f, 0.0f}, {1.0f, 2.0f}, | ||
| 79 | + {-2.0f, 3.0f}, {3.0f, -2.0f}, {-1.0f, -1.0f}, {1.0f, 1.0f}, {-100.0f, -90.0f}, {100.0f, 90.0f}}}; | ||
| 80 | + size_t bytes = ((count * sizeof(uint16_t) + 31) / 32) * 32; | ||
| 81 | + GmBuffer x(static_cast<uint8_t*>(AscendC::GmAlloc(bytes))); | ||
| 82 | + GmBuffer y(static_cast<uint8_t*>(AscendC::GmAlloc(bytes))); | ||
| 83 | + GmBuffer out(static_cast<uint8_t*>(AscendC::GmAlloc(bytes))); | ||
| 84 | + GmBuffer tiling(static_cast<uint8_t*>(AscendC::GmAlloc(sizeof(LogAddExpTilingData)))); | ||
| 85 | + ASSERT_NE(x, nullptr); | ||
| 86 | + ASSERT_NE(y, nullptr); | ||
| 87 | + ASSERT_NE(out, nullptr); | ||
| 88 | + ASSERT_NE(tiling, nullptr); | ||
| 89 | + auto* x16 = reinterpret_cast<uint16_t*>(x.get()); | ||
| 90 | + auto* y16 = reinterpret_cast<uint16_t*>(y.get()); | ||
| 91 | + for (int64_t i = 0; i < count; ++i) { | ||
| 92 | + x16[i] = Encode(pairs[i % pairs.size()][0], bf16); | ||
| 93 | + y16[i] = Encode(pairs[i % pairs.size()][1], bf16); | ||
| 94 | + } | ||
| 95 | + std::memset(out.get(), 0x5A, bytes); | ||
| 96 | + LogAddExpTilingData tilingData{}; | ||
| 97 | + tilingData.totalLength = count; | ||
| 98 | + tilingData.blockFactor = count; | ||
| 99 | + tilingData.ubFactor = ubFactor; | ||
| 100 | + tilingData.dimNum = 1; | ||
| 101 | + std::memcpy(tiling.get(), &tilingData, sizeof(tilingData)); | ||
| 102 | + AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 103 | + if (bf16) { | ||
| 104 | + ICPU_RUN_KF(log_add_exp_bf16_ut, 1, x.get(), y.get(), out.get(), tiling.get()); | ||
| 105 | + } else { | ||
| 106 | + ICPU_RUN_KF(log_add_exp_fp16_ut, 1, x.get(), y.get(), out.get(), tiling.get()); | ||
| 107 | + } | ||
| 108 | + auto* out16 = reinterpret_cast<uint16_t*>(out.get()); | ||
| 109 | + for (int64_t i = 0; i < count; ++i) { | ||
| 110 | + SCOPED_TRACE(i); | ||
| 111 | + float expected = Golden(Decode(x16[i], bf16), Decode(y16[i], bf16)); | ||
| 112 | + float actual = Decode(out16[i], bf16); | ||
| 113 | + if (std::isnan(expected)) { | ||
| 114 | + EXPECT_TRUE(std::isnan(actual)); | ||
| 115 | + } else if (std::isinf(expected)) { | ||
| 116 | + EXPECT_EQ(actual, expected); | ||
| 117 | + } else { | ||
| 118 | + ASSERT_TRUE(std::isfinite(actual)); | ||
| 119 | + EXPECT_NEAR(actual, expected, (bf16 ? 0.02f : 0.005f) * (1.0f + std::fabs(expected))); | ||
| 120 | + } | ||
| 121 | + } | ||
| 122 | + } | ||
| 123 | +}; | ||
| 124 | + | ||
| 125 | +TEST_P(LogAddExpSpecialValuesTest, ShortTile16) { Run(48, 16); } | ||
| 126 | +TEST_P(LogAddExpSpecialValuesTest, ShortTile32AndTail) { Run(65, 32); } | ||
| 127 | +TEST_P(LogAddExpSpecialValuesTest, ShortTile48AndTail) { Run(65, 48); } | ||
| 128 | +TEST_P(LogAddExpSpecialValuesTest, CompareBoundaryAndTail) { Run(65, 80); } | ||
| 129 | +TEST_P(LogAddExpSpecialValuesTest, MultipleTilesAndTail) { Run(257, 80); } | ||
| 130 | +TEST_P(LogAddExpSpecialValuesTest, DefaultUbTileAndTail) { Run(4913, 4912); } | ||
| 131 | + | ||
| 132 | +INSTANTIATE_TEST_SUITE_P(LowPrecision, LogAddExpSpecialValuesTest, testing::Values(false, true)); | ||