Optimize where_index_op(prefix sum) #30601
Conversation
update Paddle to newest version
Merge newest Paddle code
merge newest Paddle code
|
Thanks for your contribution! |
thisjiang
changed the title
|
Sorry to inform you that cde9c19's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually. |
…f imperfect self-kernel, test=develop
… optimize-whereindexop-prefix
|
Sorry to inform you that 86f20e9's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually. |
… optimize-whereindexop-prefix
| __host__ __device__ bool operator()(const T &val) { | ||
| return static_cast<bool>(val); | ||
| } | ||
| }; |
| for (int64_t i = 0; i < numel; i++) { | ||
| if (static_cast<bool>(cond_data[i])) { | ||
| h_true_index.push_back(i); | ||
| struct CheckTrue { |
There was a problem hiding this comment.
感觉这个功能封装成一个struct的意义不大，一般将一些操作定义成Fuctor，是希望将Functor作为函数的参数，从而支持设置多种同类型不同的操作。
另外，结构名不建议使用动宾结构来命名。
| memory::Alloc(platform::CPUPlace(), (rank + 1) * sizeof(int64_t)); | ||
| int64_t *ptr_stride = reinterpret_cast<int64_t *>(d_tmp_mem->ptr()); | ||
| int64_t *ptr_true_num = ptr_stride + rank; | ||
| int64_t *h_stride = reinterpret_cast<int64_t *>(h_tmp_mem->ptr()); |
There was a problem hiding this comment.
已重命名。ptr_stride是个数组，其中每个元素存的是每个维度的步长值（也就是stride），现在已重命名为stride_array。
| auto d_tmp_mem = memory::Alloc(dev_ctx, (numel + rank) * sizeof(int64_t)); | ||
| auto h_tmp_mem = |
| cub::DeviceScan::InclusiveSum(nullptr, cub_tmp_size, d_true_num, d_true_num, | ||
| numel, dev_ctx.stream()); | ||
| auto cub_tmp = memory::Alloc(dev_ctx, cub_tmp_size * sizeof(int64_t)); | ||
| void *ptr_mem = cub_tmp->ptr(); |
There was a problem hiding this comment.
已修改为cub_data标识这是cub库函数用到的数据
…Paddle into optimize-whereindexop-prefix
| const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; | ||
|
|
||
| for (int64_t idx = tid; idx < numel; idx += gridDim.x * blockDim.x) { | ||
| true_num_array[idx] = static_cast<bool>(cond_data[idx]) ? 1 : 0; |
There was a problem hiding this comment.
Improve performance: true_num_array[idx] = static_cast<bool>(cond_data[idx]) ? 1 : 0; --> true_num_array[idx] = static_cast<int64_t>(static_cast<bool>(cond_data[idx]));
| for (int64_t i = 0; i < numel; i++) { | ||
| if (static_cast<bool>(cond_data[i])) { | ||
| h_true_index.push_back(i); | ||
| __global__ void KeGetTrueNum(const T *cond_data, const int64_t numel, |
There was a problem hiding this comment.
KeGetTrueNum: What is Ke ? If it means kernel, just use GetTrueNum.
| } | ||
|
|
||
| template <typename T> | ||
| __global__ void KeSetTrueIndex(int64_t *out_ptr, const T *cond_data, |
… optimize-whereindexop-prefix
PR types
Performance optimization
PR changes
OPs
Describe
优化
优化效果:
竞品对比（测试数据大小
[100, 100, 100]）：V100-SXM2-32GB机器循环100次tf.wheretorch.nonzero可读性性优化d7fd2da：
优化变量名并增加若干注释
优化345442eb：
优化方法：
cub::DeviceScan::InclusiveSum方法替代自写的不完美的prefix sum kernel优化点：
alloc操作移到kernel前KeScanPrefixSumkernel替换为cub库带的cub::DeviceScan::InclusiveSum，提高可靠性out->Resize一次，用InclusiveSum而非ExclusiveSum的一个原因也是为了可以直接得到true_numKeGetTrueNum和KeSetTrueIndex此时都做到了访存合并优化2cde9c19：
优化方法：
thrust::inclusive_scan优化点：
thrust::inclusive_scan的defalut stream导致的不稳定性尽管这里起了4个kernel，但由于每个kernel都起来多个block，因此计算速度比单一kernel但仅1个block要快得多。
优化1
优化方法：
优化点：
true_index计算过程分为三步，每步起一个kernelKeGetTrueNum统计cond_data是否为true，是则将true_num_array设为1，否则设为0thrust::inclusive_scan计算true_num_array的前缀和（高优化点），由此得到true_index对应的索引值KeSetTrueIndex将对应索引值写入out_ptr中待优化点：
thrust::inclusive_scan在defalut stream上的表现，自写prefix sumkernel代替defalut stream上的thrust::inclusive_scandev_ctx.Wait()步骤先前优化版本PR30556
优化方法：
融合为一个kernel