| 34 | |
| 35 | template <typename T, typename IntType> |
| 36 | void ConcatGPUCall( |
| 37 | OpKernelContext* c, |
| 38 | const std::vector<std::unique_ptr<typename TTypes<T, 2>::ConstMatrix>>& |
| 39 | inputs_flat, |
| 40 | typename TTypes<T, 2>::Tensor* output_flat) { |
| 41 | GpuDeviceArrayOnHost<const T*> input_ptrs(c, inputs_flat.size()); |
| 42 | OP_REQUIRES_OK(c, input_ptrs.Init()); |
| 43 | for (int i = 0; i < inputs_flat.size(); ++i) { |
| 44 | input_ptrs.Set(i, inputs_flat[i]->data()); |
| 45 | } |
| 46 | OP_REQUIRES_OK(c, input_ptrs.Finalize()); |
| 47 | |
| 48 | GpuDeviceArrayOnHost<IntType> output_scan(c, inputs_flat.size() + 1); |
| 49 | OP_REQUIRES_OK(c, output_scan.Init()); |
| 50 | IntType scan = 0; |
| 51 | output_scan.Set(0, scan); |
| 52 | bool one_size_input = true; |
| 53 | for (int i = 0; i < inputs_flat.size(); ++i) { |
| 54 | if (one_size_input && i < inputs_flat.size() - 1 && |
| 55 | inputs_flat[i]->dimension(1) != inputs_flat[i + 1]->dimension(1)) { |
| 56 | one_size_input = false; |
| 57 | } |
| 58 | scan += inputs_flat[i]->dimension(1); |
| 59 | output_scan.Set(i + 1, scan); |
| 60 | } |
| 61 | if (!one_size_input) OP_REQUIRES_OK(c, output_scan.Finalize()); |
| 62 | |
| 63 | ConcatGPUImpl<T, IntType>(c->eigen_gpu_device(), input_ptrs.data(), |
| 64 | output_scan.data(), one_size_input, |
| 65 | inputs_flat[0]->dimension(1), output_flat); |
| 66 | } |
| 67 | |
| 68 | } // end namespace |
| 69 | |