| 240 | : DynamicStitchOpImplBase<T>(c, "DynamicStitchOp") {} |
| 241 | |
| 242 | void Compute(OpKernelContext* c) override { |
| 243 | OpInputList indices_inputs; |
| 244 | OpInputList data_inputs; |
| 245 | int first_dim_size; |
| 246 | int data_elements_size; |
| 247 | Tensor* merged = nullptr; |
| 248 | this->CheckArgsAndAllocateResult(c, &indices_inputs, &data_inputs, |
| 249 | &first_dim_size, &data_elements_size, |
| 250 | &merged); |
| 251 | if (!c->status().ok()) { |
| 252 | // Avoid segmentation faults if merged cannot be allocated and an error is |
| 253 | // passed back in the context. |
| 254 | return; |
| 255 | } |
| 256 | |
| 257 | // TODO(jeff): Currently we leave uninitialized any portions of |
| 258 | // merged that aren't covered by an index in indices. What should we do? |
| 259 | if (first_dim_size > 0) { |
| 260 | // because the collision requirements, we have to deal with |
| 261 | // collision first before send data to gpu kernel. |
| 262 | // TODO(ekelsen): Instead of doing a serial scan on the CPU to pick the |
| 263 | // last of duplicated indices, it could instead be done of the GPU |
| 264 | // implicitly using atomics to make sure the last index is the final |
| 265 | // write. |
| 266 | const int slice_size = merged->flat_outer_dims<T>().dimension(1); |
| 267 | GpuDeviceArrayOnHost<int32> indices_flat(c, first_dim_size); |
| 268 | GpuDeviceArrayOnHost<const T*> data_flat(c, data_elements_size); |
| 269 | OP_REQUIRES_OK(c, indices_flat.Init()); |
| 270 | OP_REQUIRES_OK(c, data_flat.Init()); |
| 271 | // initialize the indices_flat (-1 represents missing indices) |
| 272 | for (int i = 0; i < first_dim_size; ++i) { |
| 273 | indices_flat.Set(i, -1); |
| 274 | } |
| 275 | |
| 276 | // data_flat index |
| 277 | int32 idx = 0; |
| 278 | // sum of indices_inputs[i].NumElements() for compute indicies_flat value. |
| 279 | int32 base_size = 0; |
| 280 | for (int i = 0; i < indices_inputs.size(); ++i) { |
| 281 | auto indices_vec = indices_inputs[i].flat<int32>(); |
| 282 | auto data_ptr_base = data_inputs[i].template flat<T>().data(); |
| 283 | for (int j = 0; j < indices_vec.size(); ++j) { |
| 284 | // indices_flat's indices represent the indices of output. |
| 285 | // indices_flat's values represent the indices of input_data where the |
| 286 | // data located. |
| 287 | indices_flat.Set(indices_vec(j), base_size + j); |
| 288 | data_flat.Set( |
| 289 | idx, const_cast<T*>(reinterpret_cast<const T*>(data_ptr_base) + |
| 290 | j * slice_size)); |
| 291 | ++idx; |
| 292 | } |
| 293 | base_size += indices_vec.size(); |
| 294 | } |
| 295 | OP_REQUIRES_OK(c, indices_flat.Finalize()); |
| 296 | OP_REQUIRES_OK(c, data_flat.Finalize()); |
| 297 | |
| 298 | auto output = merged->template flat<T>().data(); |
| 299 | DynamicStitchGPUImpl<T>(c->eigen_gpu_device(), slice_size, first_dim_size, |