已合并
UnsortedSegmentMin支持空tensor #1438
jinpenghe创建于 2月3日
UnsortedSegmentMin支持空tensor #1438
已合并
共 2 个文件变更+6-5
| @@ -79,7 +79,7 @@ __aicore__ inline constexpr T GetDtypeMax() | |||
| 79 | } else if constexpr (IsSameType<T, int64_t>::value) { | 79 | } else if constexpr (IsSameType<T, int64_t>::value) { |
| 80 | dtypeMax = INT64_MAX; | 80 | dtypeMax = INT64_MAX; |
| 81 | } else if constexpr (IsSameType<T, uint32_t>::value) { | 81 | } else if constexpr (IsSameType<T, uint32_t>::value) { |
| 82 | - dtypeMax = INT32_MAX; | 82 | + dtypeMax = UINT32_MAX; |
| 83 | } else if constexpr (IsSameType<T, uint64_t>::value) { | 83 | } else if constexpr (IsSameType<T, uint64_t>::value) { |
| 84 | dtypeMax = UINT64_MAX; | 84 | dtypeMax = UINT64_MAX; |
| 85 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 85 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| @@ -76,13 +76,14 @@ extern "C" __global__ __aicore__ void unsorted_segment_min(GM_ADDR x, GM_ADDR se | |||
| 76 | op.Process(); | 76 | op.Process(); |
| 77 | } | 77 | } |
| 78 | } else if (TILING_KEY_IS(TEMPLATE_ADD_TILING_KEY)) { | 78 | } else if (TILING_KEY_IS(TEMPLATE_ADD_TILING_KEY)) { |
| 79 | + REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 4000", UnsortedSegmentOutFlTilingData); | ||
| 80 | + GET_TILING_DATA_WITH_STRUCT(UnsortedSegmentOutFlTilingData, tilingData, tiling); | ||
| 79 | if constexpr ( | 81 | if constexpr ( |
| 80 | std::is_same<uint32_t, DTYPE_X>::value || std::is_same<uint64_t, DTYPE_X>::value || | 82 | std::is_same<uint32_t, DTYPE_X>::value || std::is_same<uint64_t, DTYPE_X>::value || |
| 81 | - std::is_same<int64_t, DTYPE_X>::value) { | 83 | + std::is_same<int64_t, DTYPE_X>::value) { // 空tensor处理 |
| 82 | - return; | 84 | + UnsortedSegment::KernelUnsortedSegmentFL<DTYPE_X, DTYPE_SEGMENT_IDS, MODE_FLAGE> op(&pipe); |
| 85 | + op.Init(x, segment_ids, output, &tilingData); | ||
| 83 | } else { | 86 | } else { |
| 84 | - REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 4000", UnsortedSegmentOutFlTilingData); | ||
| 85 | - GET_TILING_DATA_WITH_STRUCT(UnsortedSegmentOutFlTilingData, tilingData, tiling); | ||
| 86 | UnsortedSegment::KernelUnsortedSegmentFL<DTYPE_X, DTYPE_SEGMENT_IDS, MODE_FLAGE> op(&pipe); | 87 | UnsortedSegment::KernelUnsortedSegmentFL<DTYPE_X, DTYPE_SEGMENT_IDS, MODE_FLAGE> op(&pipe); |
| 87 | op.Init(x, segment_ids, output, &tilingData); | 88 | op.Init(x, segment_ids, output, &tilingData); |
| 88 | op.Process(); | 89 | op.Process(); |