Star1735
1646
Code
Code
Issues469
Pull Requests559
Actions
Discussion
Wiki
Analysis
Star1735
1646
Merged
fix(radix_top_k): 修复 WS 变体 V→MTE3 同步缺失与 CreateVecIndex4TopK 越界，统一跨核前缀和累加为 int64 #6396
Nice tryAdd on 24 days ago
fix(radix_top_k): 修复 WS 变体 V→MTE3 同步缺失与 CreateVecIndex4TopK 越界，统一跨核前缀和累加为 int64 #6396
Merged
Nice tryAdd on 24 days ago
Merged fromNice_try/ops-math_workaddutintocann/ops-mathmaster
Total: 7 files changed+102-62
Mmath/radix_top_k/README.md+1-1
@@ -97,4 +97,4 @@
97 97 
98## 约束说明98## 约束说明
99 99 
100-无100+- 本算子为选取类实现，输出为选中元素在排序轴上的索引序，不保证有序。
Mmath/radix_top_k/op_host/radix_top_k_tiling.cpp+7-0
@@ -35,6 +35,11 @@ static constexpr uint64_t LARGE_TILE_LEN = 10496;
35static constexpr uint64_t MAX_TILE_NUM_IN_UB = 6144;35static constexpr uint64_t MAX_TILE_NUM_IN_UB = 6144;
36// sortLen 超过 int32 上限时切换 int64 模板（索引值域与跨核计数可达 sortLen）36// sortLen 超过 int32 上限时切换 int64 模板（索引值域与跨核计数可达 sortLen）
37static constexpr uint64_t INT32_SORT_LEN_MAX = static_cast<uint64_t>(std::numeric_limits<int32_t>::max());37static constexpr uint64_t INT32_SORT_LEN_MAX = static_cast<uint64_t>(std::numeric_limits<int32_t>::max());
38+// isInt64Shape_（sortLen > INT32_SORT_LEN_MAX）必须蕴含 isLargeShape_（sortLen > LARGE_SORT_LEN）：
39+// needWorkspace / GetWorkspaceSize 两处都写成 `isLargeShape_ || isInt64Shape_`（后者冗余）。这里用
40+// static_assert 把两个常量的相对大小钉死，避免以后调大 LARGE_SORT_LEN 时 i64 分支被漏掉。
41+static_assert(LARGE_SORT_LEN < INT32_SORT_LEN_MAX,
42+ "isInt64Shape_ must imply isLargeShape_: LARGE_SORT_LEN must stay below INT32_SORT_LEN_MAX");
38// checkUb 中 UB 占用按「每 tile 元素数 × 系数」估算：43// checkUb 中 UB 占用按「每 tile 元素数 × 系数」估算：
39// int32 索引：xBuf ping/pong(fp32) + outValue(fp32) + outIndex(int32) = 16 元素/tile；44// int32 索引：xBuf ping/pong(fp32) + outValue(fp32) + outIndex(int32) = 16 元素/tile；
40// int64 索引：outIndexBuf 翻倍（int64），系数 16 -> 20。45// int64 索引：outIndexBuf 翻倍（int64），系数 16 -> 20。
@@ -133,6 +138,7 @@ ge::graphStatus RadixTopKTiling::DoOpTiling()
133 blockDim_ = tilingData_.formerCoreNum + tilingData_.tailCoreNum;138 blockDim_ = tilingData_.formerCoreNum + tilingData_.tailCoreNum;
134 tilingData_.coreNum = blockDim_;139 tilingData_.coreNum = blockDim_;
135 // 仅 UB 变体可复用 indices 内存（WS/INT64 变体的 tileTopK/tileHist 与 CopyOutResult 跨 core 冲突）140 // 仅 UB 变体可复用 indices 内存（WS/INT64 变体的 tileTopK/tileHist 与 CopyOutResult 跨 core 冲突）
141+ // 注：isInt64Shape_ ⇒ isLargeShape_（见文件头 static_assert），这里的 `|| isInt64Shape_` 是防御性冗余
136 if (isLargeShape_ || isInt64Shape_) {142 if (isLargeShape_ || isInt64Shape_) {
137 tilingData_.needWorkspace = true;143 tilingData_.needWorkspace = true;
138 } else {144 } else {
@@ -168,6 +174,7 @@ ge::graphStatus RadixTopKTiling::GetWorkspaceSize()
168 workspaceSize_ = 0;174 workspaceSize_ = 0;
169 return ge::GRAPH_SUCCESS;175 return ge::GRAPH_SUCCESS;
170 }176 }
177+ // 注：同上，isInt64Shape_ 已被 isLargeShape_ 蕴含
171 if (isLargeShape_ || isInt64Shape_) {178 if (isLargeShape_ || isInt64Shape_) {
172 workspaceSize_ = (NUM_VALUE_2BIT * tilingData_.coreNum + tilingData_.coreNum + tilingData_.coreNum +179 workspaceSize_ = (NUM_VALUE_2BIT * tilingData_.coreNum + tilingData_.coreNum + tilingData_.coreNum +
173 WS_INT32_PER_TILE * tilingData_.totalTileNum) *180 WS_INT32_PER_TILE * tilingData_.totalTileNum) *
Mmath/radix_top_k/op_kernel/radix_top_k_common.h+24-5
@@ -55,6 +55,13 @@ using namespace AscendC;
55 * @param curTileKVal 当前 tile 的 TopK 元素数55 * @param curTileKVal 当前 tile 的 TopK 元素数
56 * @param vecIndexExpr 预生成索引序列表达式（UB/WS 变体不同）56 * @param vecIndexExpr 预生成索引序列表达式（UB/WS 变体不同）
57 * 前置条件：curBuf, coreTopKOffset, maxIndexLen, tileOffset, curTileLen 在作用域内57 * 前置条件：curBuf, coreTopKOffset, maxIndexLen, tileOffset, curTileLen 在作用域内
58+ * @note 提取相与计数相的 NaN 规范化口径**不对称是有意的**（勿在此处"补齐"）：
59+ * 计数相 CanonicalizeDataForLargest 对 bf16/half、largest 两个方向都做，目的是让 ±NaN 落到
60+ * 同一个 twiddle 键（largest=true -> 0xFFFF 计入配额；largest=false -> 0x0000 不计入）；
61+ * 提取相只需保证"掩码口径与计数一致"——bf16 分支把原始数据 Cast 成 fp32 后比较，
62+ * Compares(LE, ±NaN) 恒 false ⇒ 取非后恒 true，符号/payload 不影响掩码；largest=false 走
63+ * Compares(LT)，NaN 参与比较恒 false ⇒ 天然不入选；元素本身为 NaN 的情形由 IsNaNValue +
64+ * Compare(x,x,EQ) 取非得来的 isNaN 掩码覆盖（同样与符号/payload 无关）。
58 */65 */
59#define PROCESS_ONE_TILE_TOPK_OR_SKIP(curTileKVal, vecIndexExpr) \66#define PROCESS_ONE_TILE_TOPK_OR_SKIP(curTileKVal, vecIndexExpr) \
60 if (curTileKVal > 0) { \67 if (curTileKVal > 0) { \
@@ -139,7 +146,8 @@ protected:
139 146 
140 __aicore__ inline void CopyOutBoundaryBinCumSum(LocalTensor<int32_t>& resTensor);147 __aicore__ inline void CopyOutBoundaryBinCumSum(LocalTensor<int32_t>& resTensor);
141 148 
142- __aicore__ inline void ComputeCumSumPrev(LocalTensor<int32_t>& resTensor, int32_t& cumSumValuePrev);149+ /** @brief 前驱 core 边界 bin 累计元素数（int64 累加：总量可达 sortLen > 2^31；int32 形状下同样安全） */
150+ __aicore__ inline void ComputeCumSumPrev(LocalTensor<int32_t>& resTensor, int64_t& cumSumValuePrev);
143 151 
144 __aicore__ inline void CopyOutCoreTopK(LocalTensor<int32_t>& resTensor, int32_t totalTileTopKInCore);152 __aicore__ inline void CopyOutCoreTopK(LocalTensor<int32_t>& resTensor, int32_t totalTileTopKInCore);
145 153 
@@ -454,7 +462,11 @@ __aicore__ inline uint32_t RadixTopKBaseKernel<T, largest>::CreateVecIndex4TopK(
454}462}
455 463 
456/**464/**
457- * @brief 单个 tile 内按 cmpMode 提取匹配 boundaryNumber 的元素并输出465+ * @brief 判断边界值是否为 NaN（提取相用：边界取到 NaN 时必须换 isNaN 掩码）
466+ * 位模式判定（指数位全 1 且尾数非 0），对 ±NaN 与任意 payload 都成立；
467+ * 不用浮点自比较（v != v）是为了避免浮点相等/自比较类告警。
468+ * @param value 待判定的边界值（float 或 half）
469+ * @return true 表示 value 是 NaN
458 */470 */
459template <typename T, bool largest>471template <typename T, bool largest>
460template <typename CT>472template <typename CT>
@@ -466,8 +478,13 @@ __aicore__ inline bool RadixTopKBaseKernel<T, largest>::IsNaNValue(const CT& val
466 } else {478 } else {
467 valueFp32 = static_cast<float>(value); // half -> fp32，NaN 语义保持479 valueFp32 = static_cast<float>(value); // half -> fp32，NaN 语义保持
468 }480 }
469- uint32_t bits = *reinterpret_cast<uint32_t*>(&valueFp32);481+ // 读 float 的位模式
470- return (bits & 0x7FFFFFFFU) > 0x7F800000U; // 指数位全 1 且尾数非 0 ⇒ NaN482+ union Fp32Bits {
483+ float f;
484+ uint32_t u;
485+ } bits;
486+ bits.f = valueFp32;
487+ return (bits.u & 0x7FFFFFFFU) > 0x7F800000U; // 指数位全 1 且尾数非 0 ⇒ NaN
471}488}
472 489 
473/**490/**
@@ -706,7 +723,7 @@ __aicore__ inline void RadixTopKBaseKernel<T, largest>::CopyOutBoundaryBinCumSum
706 */723 */
707template <typename T, bool largest>724template <typename T, bool largest>
708__aicore__ inline void RadixTopKBaseKernel<T, largest>::ComputeCumSumPrev(LocalTensor<int32_t>& resTensor,725__aicore__ inline void RadixTopKBaseKernel<T, largest>::ComputeCumSumPrev(LocalTensor<int32_t>& resTensor,
709- int32_t& cumSumValuePrev)726+ int64_t& cumSumValuePrev)
710{727{
711 DataCopyPad(resTensor, boundaryBinCumSumGm_,728 DataCopyPad(resTensor, boundaryBinCumSumGm_,
712 DataCopyExtParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t) * coreNum_), 0, 0, 0},729 DataCopyExtParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t) * coreNum_), 0, 0, 0},
@@ -714,6 +731,8 @@ __aicore__ inline void RadixTopKBaseKernel<T, largest>::ComputeCumSumPrev(LocalT
714 MTE2ToSSync();731 MTE2ToSSync();
715 cumSumValuePrev = 0;732 cumSumValuePrev = 0;
716 for (int32_t coreId = 0; coreId < blockIdx_; coreId++) {733 for (int32_t coreId = 0; coreId < blockIdx_; coreId++) {
734+ // 每 core 的边界 bin 计数 < 2^31（per-core 元素数上限），int64 累加对 int32 形状同样安全，
735+ // 故 UB/WS 两个变体共用本函数（不再各留一份只差累加器类型的实现）
717 cumSumValuePrev += resTensor.GetValue(coreId);736 cumSumValuePrev += resTensor.GetValue(coreId);
718 }737 }
719}738}
Mmath/radix_top_k/op_kernel/radix_top_k_ub.h+4-3
@@ -322,9 +322,9 @@ __aicore__ inline void RadixTopKUb<T, largest>::HandleLastRoundBoundary(LocalTen
322 this->CopyOutBoundaryBinCumSum(resTensor);322 this->CopyOutBoundaryBinCumSum(resTensor);
323 SyncAll();323 SyncAll();
324 324 
325- int32_t cumSumValuePrev = 0;325+ int64_t cumSumValuePrev = 0; // 基类 ComputeCumSumPrev 的本体为 int64 累加（与 WS 变体共用）
326 this->ComputeCumSumPrev(resTensor, cumSumValuePrev);326 this->ComputeCumSumPrev(resTensor, cumSumValuePrev);
327- int32_t remainCoreBoundaryNum = this->remainK_ - cumSumValuePrev;327+ int64_t remainCoreBoundaryNum = static_cast<int64_t>(this->remainK_) - cumSumValuePrev;
328 if (remainCoreBoundaryNum > 0) {328 if (remainCoreBoundaryNum > 0) {
329 for (int32_t tileId = 0; tileId < this->tileNum_; tileId++) {329 for (int32_t tileId = 0; tileId < this->tileNum_; tileId++) {
330 int32_t curTileBoundaryNum = tileHist[this->boundaryBin * this->tileNumAlign_].GetValue(tileId);330 int32_t curTileBoundaryNum = tileHist[this->boundaryBin * this->tileNumAlign_].GetValue(tileId);
@@ -335,7 +335,8 @@ __aicore__ inline void RadixTopKUb<T, largest>::HandleLastRoundBoundary(LocalTen
335 remainCoreBoundaryNum -= curTileBoundaryNum;335 remainCoreBoundaryNum -= curTileBoundaryNum;
336 tileTopK.SetValue(tileId, tileTopK.GetValue(tileId) + curTileBoundaryNum);336 tileTopK.SetValue(tileId, tileTopK.GetValue(tileId) + curTileBoundaryNum);
337 } else {337 } else {
338- tileTopK.SetValue(tileId, tileTopK.GetValue(tileId) + remainCoreBoundaryNum);338+ // per-core 配额 <= 该 tile 的边界 bin 计数 < 2^31，回写 tileTopK 的 int32 通道安全
339+ tileTopK.SetValue(tileId, tileTopK.GetValue(tileId) + static_cast<int32_t>(remainCoreBoundaryNum));
339 remainCoreBoundaryNum = 0;340 remainCoreBoundaryNum = 0;
340 break;341 break;
341 }342 }
Mmath/radix_top_k/op_kernel/radix_top_k_ws.h+15-22
@@ -70,8 +70,6 @@ protected:
70 __aicore__ inline void AddTileHistToTileTopK(LocalTensor<int32_t>& globalHist);70 __aicore__ inline void AddTileHistToTileTopK(LocalTensor<int32_t>& globalHist);
71 __aicore__ inline void HandleLastRoundBoundary(LocalTensor<int32_t>& resTensor, LocalTensor<float>& tileHistFp32);71 __aicore__ inline void HandleLastRoundBoundary(LocalTensor<int32_t>& resTensor, LocalTensor<float>& tileHistFp32);
72 __aicore__ inline void WriteCoreTopKFromWs(LocalTensor<int32_t>& resTensor, LocalTensor<float>& tileHistFp32);72 __aicore__ inline void WriteCoreTopKFromWs(LocalTensor<int32_t>& resTensor, LocalTensor<float>& tileHistFp32);
73- /** @brief 前驱 core 边界 bin 累计数（int64 累加，总量可达 sortLen > 2^31；int32 形状下同样安全） */
74- __aicore__ inline void ComputeCumSumPrevI64(LocalTensor<int32_t>& resTensor, int64_t& cumSumValuePrev);
75 73 
76 GlobalTensor<int32_t> tileTopKGm_; /**< WS 中 tileTopK 段的 Global Tensor */74 GlobalTensor<int32_t> tileTopKGm_; /**< WS 中 tileTopK 段的 Global Tensor */
77 GlobalTensor<int32_t> tileHistGm_; /**< WS 中 tileHist 段的 Global Tensor */75 GlobalTensor<int32_t> tileHistGm_; /**< WS 中 tileHist 段的 Global Tensor */
@@ -332,7 +330,8 @@ __aicore__ inline void RadixTopKWs<T, largest>::HandleLastRoundBoundary(LocalTen
332 SyncAll();330 SyncAll();
333 331 
334 int64_t cumSumValuePrev = 0;332 int64_t cumSumValuePrev = 0;
335- ComputeCumSumPrevI64(resTensor, cumSumValuePrev);333+ // 复用基类的前缀和（出参已是 int64）：本变体与 UB 变体口径一致，避免两份只差累加器类型的实现
334+ this->ComputeCumSumPrev(resTensor, cumSumValuePrev);
336 int64_t remainCoreBoundaryNum = static_cast<int64_t>(this->remainK_) - cumSumValuePrev;335 int64_t remainCoreBoundaryNum = static_cast<int64_t>(this->remainK_) - cumSumValuePrev;
337 if (remainCoreBoundaryNum > 0) {336 if (remainCoreBoundaryNum > 0) {
338 LocalTensor<int32_t> tileTopK = tileHistAndTopKBuf_.Get<int32_t>();337 LocalTensor<int32_t> tileTopK = tileHistAndTopKBuf_.Get<int32_t>();
@@ -383,7 +382,11 @@ __aicore__ inline void RadixTopKWs<T, largest>::HandleLastRoundBoundary(LocalTen
383 Add(tileTopK, tileTopK, tileHistBoundary, dataNum);382 Add(tileTopK, tileTopK, tileHistBoundary, dataNum);
384 PipeBarrier<PIPE_V>();383 PipeBarrier<PIPE_V>();
385 remainCoreBoundaryNum -= segSum;384 remainCoreBoundaryNum -= segSum;
386- SToMTE3Sync();385+ // Add 是 V 写 tileTopK，紧接着 MTE3 要读同一块 UB：PipeBarrier<PIPE_V> 只保证 V 流水线
386+ // 内部顺序，SToMTE3Sync 只覆盖 S→MTE3，两者都管不到 V→MTE3，MTE3 可能读到 Add 之前的
387+ // 旧值（该段 k 配额错乱）。同文件 AddTileHist2TileTopKInWs 对同样的 Add→CopyOut 序列
388+ // 用的就是 VToMTE3Sync，这里必须保持一致。
389+ VToMTE3Sync();
387 CopyTileTopKUb2Ws(tileTopK, startTileId, dataNum);390 CopyTileTopKUb2Ws(tileTopK, startTileId, dataNum);
388 MTE3ToMTE2Sync();391 MTE3ToMTE2Sync();
389 if (remainCoreBoundaryNum == 0) {392 if (remainCoreBoundaryNum == 0) {
@@ -449,24 +452,6 @@ __aicore__ inline void RadixTopKWs<T, largest>::WriteCoreTopKFromWs(LocalTensor<
449 this->CopyOutCoreTopK(resTensor, static_cast<int32_t>(totalTileTopKInCore));452 this->CopyOutCoreTopK(resTensor, static_cast<int32_t>(totalTileTopKInCore));
450}453}
451 454 
452-/**
453- * @brief 计算排在当前 core 之前所有 core 的边界 bin 累计元素数（int64 累加，总量可达 sortLen）
454- */
455-template <typename T, bool largest>
456-__aicore__ inline void RadixTopKWs<T, largest>::ComputeCumSumPrevI64(LocalTensor<int32_t>& resTensor,
457- int64_t& cumSumValuePrev)
458-{
459- DataCopyPad(
460- resTensor, this->boundaryBinCumSumGm_,
461- DataCopyExtParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t) * this->coreNum_), 0, 0, 0},
462- DataCopyPadExtParams<int32_t>{true, 0, 0, 0});
463- MTE2ToSSync();
464- cumSumValuePrev = 0;
465- for (int32_t coreId = 0; coreId < this->blockIdx_; coreId++) {
466- cumSumValuePrev += resTensor.GetValue(coreId);
467- }
468-}
469- 
470/**455/**
471 * @brief 为 TileTopK 阶段创建预生成索引，WS模板从 tileHistAndTopKBuf_ 剩余空间分配456 * @brief 为 TileTopK 阶段创建预生成索引，WS模板从 tileHistAndTopKBuf_ 剩余空间分配
472 * @return 实际可用的 maxIndexLen457 * @return 实际可用的 maxIndexLen
@@ -474,6 +459,14 @@ __aicore__ inline void RadixTopKWs<T, largest>::ComputeCumSumPrevI64(LocalTensor
474template <typename T, bool largest>459template <typename T, bool largest>
475__aicore__ inline uint32_t RadixTopKWs<T, largest>::CreateVecIndex4TopK()460__aicore__ inline uint32_t RadixTopKWs<T, largest>::CreateVecIndex4TopK()
476{461{
462+ // tileNumAlign_ 是 UB 内 tileHistAndTopKBuf_ 已占用的元素数（= 对齐后的 tileNum_）。i64 变体的
463+ // 多段 regime 会命中 tileNum_ > MAX_TILE_NUM_IN_UB：此时下面的无符号相减会下溢成巨大值，再被
464+ // tileLen_ 兜住，最终 GetWithOffset(maxIndexLen, tileNumAlign_ * 4) 的偏移会写出 6144 元素 buffer。
465+ // 这里把「不变量：仅当 tileNumAlign_ < MAX_TILE_NUM_IN_UB 才调用」显式落到代码，返回 0 时下游
466+ // 走「现算索引」分支，行为与 maxIndexLen==0 的既有兜底一致。
467+ if (this->tileNumAlign_ >= MAX_TILE_NUM_IN_UB) {
468+ return 0;
469+ }
477 uint64_t int32Align = BLOCK_SIZE / sizeof(int32_t);470 uint64_t int32Align = BLOCK_SIZE / sizeof(int32_t);
478 uint32_t maxIndexLen = Ops::Base::FloorDiv(MAX_TILE_NUM_IN_UB - this->tileNumAlign_, int32Align) * int32Align;471 uint32_t maxIndexLen = Ops::Base::FloorDiv(MAX_TILE_NUM_IN_UB - this->tileNumAlign_, int32Align) * int32Align;
479 if (this->tileLen_ < maxIndexLen) {472 if (this->tileLen_ < maxIndexLen) {
Mmath/radix_top_k/tests/ut/op_host/test_radix_top_k_tiling.cpp+46-29
@@ -20,32 +20,23 @@ using namespace ge;
20 20 
21class RadixTopKTilingTest : public testing::Test {21class RadixTopKTilingTest : public testing::Test {
22protected:22protected:
23- static void SetUpTestCase()23+ static void SetUpTestCase() { std::cout << "RadixTopKTilingTest SetUp" << std::endl; }
24- {24+ static void TearDownTestCase() { std::cout << "RadixTopKTilingTest TearDown" << std::endl; }
25- std::cout << "RadixTopKTilingTest SetUp" << std::endl;
26- }
27- static void TearDownTestCase()
28- {
29- std::cout << "RadixTopKTilingTest TearDown" << std::endl;
30- }
31};25};
32 26 
33namespace {27namespace {
34 28 
35constexpr uint64_t LARGE_SORT_LEN = 1448600000;29constexpr uint64_t LARGE_SORT_LEN = 1448600000;
30+constexpr uint64_t INT32_SORT_LEN_MAX = 2147483647; // 与 op_host 的 INT32_SORT_LEN_MAX 对齐
36constexpr uint64_t CORE_NUM = 48;31constexpr uint64_t CORE_NUM = 48;
37constexpr uint64_t UB_SIZE = 196608;32constexpr uint64_t UB_SIZE = 196608;
38 33 
39static thread_local int32_t g_kValue = 0;34static thread_local int32_t g_kValue = 0;
40 35 
41-gert::TilingContextPara MakeRadixTopKTilingContext(36+gert::TilingContextPara MakeRadixTopKTilingContext(const gert::StorageShape& xShape, ge::DataType xDtype,
42- const gert::StorageShape &xShape,37+ int32_t kValue, bool sorted = true, int64_t dim = -1,
43- ge::DataType xDtype,38+ bool largest = true, int64_t indicesDtype = 3,
44- int32_t kValue,39+ ge::DataType outIndicesDtype = ge::DT_INT32)
45- bool sorted = true,
46- int64_t dim = -1,
47- bool largest = true,
48- int64_t indicesDtype = 3)
49{40{
50 g_kValue = kValue;41 g_kValue = kValue;
51 gert::StorageShape kShape({1}, {1});42 gert::StorageShape kShape({1}, {1});
@@ -64,14 +55,13 @@ gert::TilingContextPara MakeRadixTopKTilingContext(
64 },55 },
65 {56 {
66 {xShape, xDtype, ge::FORMAT_ND},57 {xShape, xDtype, ge::FORMAT_ND},
67- {xShape, ge::DT_INT32, ge::FORMAT_ND},58+ {xShape, outIndicesDtype, ge::FORMAT_ND},
68 },59 },
69 {60 {
70 gert::TilingContextPara::OpAttr("sorted", Ops::Math::AnyValue::CreateFrom<bool>(sorted)),61 gert::TilingContextPara::OpAttr("sorted", Ops::Math::AnyValue::CreateFrom<bool>(sorted)),
71 gert::TilingContextPara::OpAttr("dim", Ops::Math::AnyValue::CreateFrom<int64_t>(dim)),62 gert::TilingContextPara::OpAttr("dim", Ops::Math::AnyValue::CreateFrom<int64_t>(dim)),
72 gert::TilingContextPara::OpAttr("largest", Ops::Math::AnyValue::CreateFrom<bool>(largest)),63 gert::TilingContextPara::OpAttr("largest", Ops::Math::AnyValue::CreateFrom<bool>(largest)),
73- gert::TilingContextPara::OpAttr("indices_dtype",64+ gert::TilingContextPara::OpAttr("indices_dtype", Ops::Math::AnyValue::CreateFrom<int64_t>(indicesDtype)),
74- Ops::Math::AnyValue::CreateFrom<int64_t>(indicesDtype)),
75 },65 },
76 compileInfoPtr);66 compileInfoPtr);
77}67}
@@ -83,10 +73,14 @@ uint64_t ComputeExpectedTilingKey(bool sorted, bool largest, uint64_t sortLen)
83 key |= (static_cast<uint64_t>(largest) << 1);73 key |= (static_cast<uint64_t>(largest) << 1);
84 bool isLargeShape = (sortLen > LARGE_SORT_LEN);74 bool isLargeShape = (sortLen > LARGE_SORT_LEN);
85 key |= (static_cast<uint64_t>(isLargeShape) << 2);75 key |= (static_cast<uint64_t>(isLargeShape) << 2);
76+ // 第 4 位 isInt64Shape：与 op_host 的 GET_TPL_TILING_KEY(sorted, largest, isLargeShape,
77+ // isInt64Shape) 一一对应；此前只编码 3 位，int64 主路径（sortLen > 2^31）恒测不到
78+ bool isInt64Shape = (sortLen > INT32_SORT_LEN_MAX);
79+ key |= (static_cast<uint64_t>(isInt64Shape) << 3);
86 return key;80 return key;
87}81}
88 82 
89-void CheckTilingKey(const gert::TilingContextPara &tilingCtx, uint64_t expectKey)83+void CheckTilingKey(const gert::TilingContextPara& tilingCtx, uint64_t expectKey)
90{84{
91 TilingInfo tilingInfo;85 TilingInfo tilingInfo;
92 ASSERT_TRUE(ExecuteTiling(tilingCtx, tilingInfo));86 ASSERT_TRUE(ExecuteTiling(tilingCtx, tilingInfo));
@@ -94,23 +88,20 @@ void CheckTilingKey(const gert::TilingContextPara &tilingCtx, uint64_t expectKey
94 EXPECT_EQ(tilingInfo.workspaceSizes.size(), 1);88 EXPECT_EQ(tilingInfo.workspaceSizes.size(), 1);
95}89}
96 90 
97-void CheckTilingWithWs(const gert::TilingContextPara &tilingCtx,91+void CheckTilingWithWs(const gert::TilingContextPara& tilingCtx, uint64_t expectKey, bool expectWsNonzero)
98- uint64_t expectKey, bool expectWsNonzero)
99{92{
100 TilingInfo tilingInfo;93 TilingInfo tilingInfo;
101 ASSERT_TRUE(ExecuteTiling(tilingCtx, tilingInfo));94 ASSERT_TRUE(ExecuteTiling(tilingCtx, tilingInfo));
102 EXPECT_EQ(tilingInfo.tilingKey, expectKey);95 EXPECT_EQ(tilingInfo.tilingKey, expectKey);
103 ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1);96 ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1);
104 if (expectWsNonzero) {97 if (expectWsNonzero) {
105- EXPECT_GT(tilingInfo.workspaceSizes[0], 0u)98+ EXPECT_GT(tilingInfo.workspaceSizes[0], 0u) << "k 较小时应触发 workspace 分配 (needWorkspace=true)";
106- << "k 较小时应触发 workspace 分配 (needWorkspace=true)";
107 } else {99 } else {
108- EXPECT_EQ(tilingInfo.workspaceSizes[0], 0u)100+ EXPECT_EQ(tilingInfo.workspaceSizes[0], 0u) << "k 较大时应复用 indices 内存 (needWorkspace=false)";
109- << "k 较大时应复用 indices 内存 (needWorkspace=false)";
110 }101 }
111}102}
112 103 
113-} // namespace104+} // namespace
114 105 
115// ============================================================106// ============================================================
116// UB 变体 fp16 — 1D shapes107// UB 变体 fp16 — 1D shapes
@@ -239,8 +230,7 @@ TEST_F(RadixTopKTilingTest, ub_fp16_tiny_1x1)
239TEST_F(RadixTopKTilingTest, ws_fp16_large_sort_len)230TEST_F(RadixTopKTilingTest, ws_fp16_large_sort_len)
240{231{
241 uint64_t bigSortLen = LARGE_SORT_LEN + 1;232 uint64_t bigSortLen = LARGE_SORT_LEN + 1;
242- auto tilingCtx = MakeRadixTopKTilingContext(233+ auto tilingCtx = MakeRadixTopKTilingContext({{1, bigSortLen}, {1, bigSortLen}}, ge::DT_FLOAT16, 10);
243- {{1, bigSortLen}, {1, bigSortLen}}, ge::DT_FLOAT16, 10);
244 uint64_t expectKey = ComputeExpectedTilingKey(true, true, bigSortLen);234 uint64_t expectKey = ComputeExpectedTilingKey(true, true, bigSortLen);
245 TilingInfo tilingInfo;235 TilingInfo tilingInfo;
246 ASSERT_TRUE(ExecuteTiling(tilingCtx, tilingInfo));236 ASSERT_TRUE(ExecuteTiling(tilingCtx, tilingInfo));
@@ -274,3 +264,30 @@ TEST_F(RadixTopKTilingTest, k100_ub_needs_ws)
274 auto tilingCtx = MakeRadixTopKTilingContext({{65536}, {65536}}, ge::DT_FLOAT16, 100);264 auto tilingCtx = MakeRadixTopKTilingContext({{65536}, {65536}}, ge::DT_FLOAT16, 100);
275 CheckTilingWithWs(tilingCtx, ComputeExpectedTilingKey(true, true, 65536), true);265 CheckTilingWithWs(tilingCtx, ComputeExpectedTilingKey(true, true, 65536), true);
276}266}
267+ 
268+// ============================================================
269+// int64 主路径（sortLen > 2^31）：tiling key 第 4 位 + 入口防御校验
270+// ============================================================
271+ 
272+TEST_F(RadixTopKTilingTest, i64_shape_tiling_key_bit3)
273+{
274+ // sortLen = 2^31 > INT32_SORT_LEN_MAX ⇒ isInt64Shape=true（tilingKey bit3），并蕴含
275+ // isLargeShape（bit2）；indices 输出必须为 int64，且该变体必须申请 workspace
276+ constexpr uint64_t SORT_LEN = 2147483648ULL; // 2^31
277+ auto tilingCtx = MakeRadixTopKTilingContext(
278+ {{1, static_cast<int64_t>(SORT_LEN)}, {1, static_cast<int64_t>(SORT_LEN)}}, ge::DT_FLOAT16, 1001, true, -1,
279+ true, /*indicesDtype=*/9 /*DT_INT64*/,
280+ /*outIndicesDtype=*/ge::DT_INT64);
281+ CheckTilingWithWs(tilingCtx, ComputeExpectedTilingKey(true, true, SORT_LEN), true);
282+}
283+ 
284+TEST_F(RadixTopKTilingTest, i64_shape_reject_int32_indices)
285+{
286+ // sortLen > 2^31 但 indices 输出为 int32 ⇒ DoOpTiling/GetShapeAttrsInfo 必须返回 GRAPH_FAILED
287+ // （int32 索引表示不了 > 2^31 的下标，不能留到 kernel 再出错）
288+ constexpr uint64_t SORT_LEN = 2147483648ULL;
289+ auto tilingCtx = MakeRadixTopKTilingContext(
290+ {{1, static_cast<int64_t>(SORT_LEN)}, {1, static_cast<int64_t>(SORT_LEN)}}, ge::DT_FLOAT16, 1001);
291+ TilingInfo tilingInfo;
292+ EXPECT_FALSE(ExecuteTiling(tilingCtx, tilingInfo)) << "sortLen 超 int32 上限且 indices 非 int64 时必须失败";
293+}
Mmath/topk/op_api/topk.cpp+5-2
@@ -127,8 +127,11 @@ static bool IsRadixTopKSupport(const aclTensor* self, int64_t k, bool sorted, op
127 int64_t sortLen = inputShape.GetDim(dimNum - 1);127 int64_t sortLen = inputShape.GetDim(dimNum - 1);
128 bool socCheck = version == SocVersion::ASCEND910B || version == SocVersion::ASCEND910_93;128 bool socCheck = version == SocVersion::ASCEND910B || version == SocVersion::ASCEND910_93;
129 bool dtypeCheck = CheckType(self->GetDataType(), {op::DataType::DT_FLOAT16, op::DataType::DT_BF16});129 bool dtypeCheck = CheckType(self->GetDataType(), {op::DataType::DT_FLOAT16, op::DataType::DT_BF16});
130- // sortLen 超 int32 上限 → int64 大 shape 模板：只要求 indices 为 int64 且 k > 0，130+ // sortLen 超 int32 上限 → int64 大 shape 模板：
131- // 不套用下面的 k 比例门限，故此处直接判定131+ // * 只要求 indices 为 int64 且 k > 0 且 !sorted：该量级下替代路径是 TopKV2 的 workspace OOM（模型场景
132+ // 约 59.5GB），"k 大才有收益"的比例门限失去意义，故与 MIN_K 门限一并放掉；
133+ // RadixTopK 内核输出无序列，有序性靠调用方补 Sort + GatherElementsV2，
134+ // 而 int64 索引补不了（GatherElementsV2 只接受 int32 索引）
132 if (sortLen > MAX_INT32_INPUTSIZE) {135 if (sortLen > MAX_INT32_INPUTSIZE) {
133 return socCheck && dtypeCheck && (indicesDType == op::DataType::DT_INT64) && k > 0 && !sorted;136 return socCheck && dtypeCheck && (indicesDType == op::DataType::DT_INT64) && k > 0 && !sorted;
134 }137 }