Commit 9d3d6da2 authored by Yong Tang's avatar Yong Tang Committed by Rasmus Munk Larsen
Browse files

Add GPU support for Bucketize op (#13922)



* Add GPU support for Bucketize op

This fix tries to add GPU support for `Bucketize` op. Before
this PR only CPU implementation is available. This PR add GPU
implementation with a CUDA kernel.

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Add GPU kernel for Bucketize op.

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Update test cases to invoke GPU implementation of Bucketize

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Move Eigen header to the top

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Address review feedback

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Use binary search for finding the upper bound with GPU

similiar to std::upper_bound

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Use CudaDeviceArrayOnHost instead of pinned memory

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>
parent a620cd39
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment