Skip to content

GitLab

Commit a4fdef4c authored by zhanghj2's avatar zhanghj2
Browse files

优化softmax计算

parent 3a477917
Show whitespace changes
Inline Side-by-side
Showing with 14 additions and 12 deletions
csrc/softmax.h
View file @ a4fdef4c
......@@ -55,12 +55,13 @@ __device__ __forceinline__ void warp_allreduce_(Tensor<Engine0, Layout0> &dst, T
// smem_reduce(row, col) = dst(0);
}
__syncthreads();
if (tidx < 16)
{
smem_reduce(row + 64) = op(op(smem_reduce(row * 4), smem_reduce(row * 4 + 1)), op(smem_reduce(row * 4 + 2), smem_reduce(row * 4 + 3)));
}
__syncthreads();
dst(0) = smem_reduce(row + 64);
// if (tidx < 16)
// {
// smem_reduce(row + 64) = op(op(smem_reduce(row * 4), smem_reduce(row * 4 + 1)), op(smem_reduce(row * 4 + 2), smem_reduce(row * 4 + 3)));
// }
// __syncthreads();
// dst(0) = smem_reduce(row + 64);
dst(0) = op(op(smem_reduce(row * 4), smem_reduce(row * 4 + 1)), op(smem_reduce(row * 4 + 2), smem_reduce(row * 4 + 3)));
}
template<typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
......@@ -75,12 +76,13 @@ __device__ __forceinline__ void warp_allreduce_tp1(Tensor<Engine0, Layout0> &dst
smem_reduce[row * 2 + (warp_id / 4)] = dst[0];
}
__syncthreads();
if (col == 0 && warp_id < 4) {
// printf("sum %d %d %d %.2f %.2f \n", row, tidx, warp_id, smem_reduce[row * 2], smem_reduce[row * 2 + warp_id / 4]);
smem_reduce[128 + row] = op(smem_reduce[row * 2], smem_reduce[row * 2 + 1]);
}
__syncthreads();
dst(0) = smem_reduce(128 + row);
// if (col == 0 && warp_id < 4) {
// // printf("sum %d %d %d %.2f %.2f \n", row, tidx, warp_id, smem_reduce[row * 2], smem_reduce[row * 2 + warp_id / 4]);
// smem_reduce[128 + row] = op(smem_reduce[row * 2], smem_reduce[row * 2 + 1]);
// }
// __syncthreads();
// dst(0) = smem_reduce(128 + row);
dst(0) = op(smem_reduce[row * 2], smem_reduce[row * 2 + 1]);
}
template<typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
__device__ __forceinline__ void warp_allreduce_tp4(Tensor<Engine0, Layout0> &dst, Tensor<Engine1, Layout1> &smem_reduce, Operator &op) {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment