已合并
SparseReshape输入和输出shape相同时直接输入到输出 #4607
kevin_huang1234创建于 13 天前
SparseReshape输入和输出shape相同时直接输入到输出 #4607
已合并
共 1 个文件变更+26-6
| @@ -33,7 +33,22 @@ using namespace AscendC; | |||
| 33 | 33 | ||
| 34 | static constexpr uint32_t THREAD_NUM = 512; | 34 | static constexpr uint32_t THREAD_NUM = 512; |
| 35 | 35 | ||
| 36 | -// VF kernel: Grid-Stride loop over non-zero elements | 36 | +// VF kernel (identity): direct copy indices, no flat index computation |
| 37 | +template <typename T> | ||
| 38 | +__simt_vf__ __aicore__ __launch_bounds__(THREAD_NUM) inline void OpSparseReshapeIdentitySimtKernel(int64_t nnz, | ||
| 39 | + int32_t rank, | ||
| 40 | + __gm__ T* indices, | ||
| 41 | + __gm__ T* yIndices) | ||
| 42 | +{ | ||
| 43 | + for (int64_t i = static_cast<int64_t>(blockIdx.x * blockDim.x + threadIdx.x); i < nnz; | ||
| 44 | + i += static_cast<int64_t>(blockDim.x * gridDim.x)) { | ||
| 45 | + for (int32_t j = 0; j < rank; j++) { | ||
| 46 | + yIndices[i * rank + j] = indices[i * rank + j]; | ||
| 47 | + } | ||
| 48 | + } | ||
| 49 | +} | ||
| 50 | + | ||
| 51 | +// VF kernel (non-identity): flat index computation + UintDiv decomposition | ||
| 37 | template <typename T> | 52 | template <typename T> |
| 38 | __simt_vf__ __aicore__ __launch_bounds__(THREAD_NUM) inline void OpSparseReshapeSimtKernel( | 53 | __simt_vf__ __aicore__ __launch_bounds__(THREAD_NUM) inline void OpSparseReshapeSimtKernel( |
| 39 | int64_t nnz, int32_t inputRank, int32_t outputRank, __ubuf__ int64_t* inputStrides, __ubuf__ int64_t* outputStrides, | 54 | int64_t nnz, int32_t inputRank, int32_t outputRank, __ubuf__ int64_t* inputStrides, __ubuf__ int64_t* outputStrides, |
| @@ -68,6 +83,13 @@ __aicore__ inline void WriteYShape(GM_ADDR y_shape, const SparseReshapeTilingDat | |||
| 68 | template <typename T> | 83 | template <typename T> |
| 69 | __aicore__ inline void PrepareAndLaunchVF(GM_ADDR indices, GM_ADDR y_indices, const SparseReshapeTilingData* td) | 84 | __aicore__ inline void PrepareAndLaunchVF(GM_ADDR indices, GM_ADDR y_indices, const SparseReshapeTilingData* td) |
| 70 | { | 85 | { |
| 86 | + __gm__ T* indicesGm = (__gm__ T*)indices; | ||
| 87 | + __gm__ T* yIndicesGm = (__gm__ T*)y_indices; | ||
| 88 | + if (td->isIdentityReshape) { | ||
| 89 | + asc_vf_call<OpSparseReshapeIdentitySimtKernel<T>>(dim3(THREAD_NUM), td->nnz, td->inputRank, indicesGm, | ||
| 90 | + yIndicesGm); | ||
| 91 | + return; | ||
| 92 | + } | ||
| 71 | LocalMemAllocator<AscendC::Hardware::UB> ubAlloc; | 93 | LocalMemAllocator<AscendC::Hardware::UB> ubAlloc; |
| 72 | LocalTensor<int64_t> inStridesUb = ubAlloc.Alloc<int64_t>(MAX_RANK); | 94 | LocalTensor<int64_t> inStridesUb = ubAlloc.Alloc<int64_t>(MAX_RANK); |
| 73 | LocalTensor<int64_t> outStridesUb = ubAlloc.Alloc<int64_t>(MAX_RANK); | 95 | LocalTensor<int64_t> outStridesUb = ubAlloc.Alloc<int64_t>(MAX_RANK); |
| @@ -85,8 +107,6 @@ __aicore__ inline void PrepareAndLaunchVF(GM_ADDR indices, GM_ADDR y_indices, co | |||
| 85 | divShiftUb.SetValue(d, shift); | 107 | divShiftUb.SetValue(d, shift); |
| 86 | } | 108 | } |
| 87 | DataSyncBarrier<MemDsbT::UB>(); | 109 | DataSyncBarrier<MemDsbT::UB>(); |
| 88 | - __gm__ T* indicesGm = (__gm__ T*)indices; | ||
| 89 | - __gm__ T* yIndicesGm = (__gm__ T*)y_indices; | ||
| 90 | asc_vf_call<OpSparseReshapeSimtKernel<T>>( | 110 | asc_vf_call<OpSparseReshapeSimtKernel<T>>( |
| 91 | dim3(THREAD_NUM), td->nnz, td->inputRank, td->outputRank, (__ubuf__ int64_t*)inStridesUb.GetPhyAddr(), | 111 | dim3(THREAD_NUM), td->nnz, td->inputRank, td->outputRank, (__ubuf__ int64_t*)inStridesUb.GetPhyAddr(), |
| 92 | (__ubuf__ int64_t*)outStridesUb.GetPhyAddr(), (__ubuf__ uint64_t*)divMagicUb.GetPhyAddr(), | 112 | (__ubuf__ int64_t*)outStridesUb.GetPhyAddr(), (__ubuf__ uint64_t*)divMagicUb.GetPhyAddr(), |
| @@ -94,9 +114,9 @@ __aicore__ inline void PrepareAndLaunchVF(GM_ADDR indices, GM_ADDR y_indices, co | |||
| 94 | } | 114 | } |
| 95 | 115 | ||
| 96 | // Main Process function | 116 | // Main Process function |
| 97 | -// Always use VF kernel path: it correctly handles both identity and non-identity | 117 | +// VF kernel handles both paths via grid-stride partitioning (no multi-core race): |
| 98 | -// reshape via grid-stride partitioning. The previous identity fast path (scalar | 118 | +// - isIdentityReshape==1: direct copy indices, skip flat index computation |
| 99 | -// copy) had a multi-core race condition where all cores copied all elements. | 119 | +// - isIdentityReshape==0: flat index + UintDiv decomposition |
| 100 | template <typename T> | 120 | template <typename T> |
| 101 | __aicore__ inline void Process(GM_ADDR indices, GM_ADDR shape, GM_ADDR new_shape, GM_ADDR y_indices, GM_ADDR y_shape, | 121 | __aicore__ inline void Process(GM_ADDR indices, GM_ADDR shape, GM_ADDR new_shape, GM_ADDR y_indices, GM_ADDR y_shape, |
| 102 | GM_ADDR workspace, const SparseReshapeTilingData* td) | 122 | GM_ADDR workspace, const SparseReshapeTilingData* td) |