Merged
fix(math): preserve equal infinities in low precision logaddexp #6219
fix(math): preserve equal infinities in low precision logaddexp #6219
Merged
puppyAdd on Oct 11
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 of3 # 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/DataCopyPad115 // 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 back322 // 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));