| 345 | explicit SplitVOpGPU(OpKernelConstruction* c) : Base(c) {} |
| 346 | |
| 347 | void Compute(OpKernelContext* context) override { |
| 348 | bool done = false; |
| 349 | std::vector<Tlen> split_sizes_vec; |
| 350 | Base::ComputeEasyCases(context, &done, &split_sizes_vec); |
| 351 | if (!context->status().ok() || done) { |
| 352 | return; |
| 353 | } |
| 354 | const int32 num_split = Base::num_outputs(); |
| 355 | const Tensor& input = context->input(0); |
| 356 | const TensorShape& input_shape = input.shape(); |
| 357 | const int32 split_dim_orig = context->input(2).flat<int32>()(0); |
| 358 | const int32 split_dim = |
| 359 | split_dim_orig < 0 ? split_dim_orig + input.dims() : split_dim_orig; |
| 360 | OP_REQUIRES( |
| 361 | context, |
| 362 | FastBoundsCheck(input.NumElements(), std::numeric_limits<int32>::max()), |
| 363 | errors::InvalidArgument("Split on GPU requires input size " |
| 364 | "< max int32")); |
| 365 | |
| 366 | int32 prefix_dim_size; |
| 367 | int32 split_dim_size; |
| 368 | int32 suffix_dim_size; |
| 369 | std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) = |
| 370 | Base::template SetDims<int32>(input_shape, split_dim); |
| 371 | |
| 372 | // use the same approach as concat (see documentation there) |
| 373 | // reshape to 2D |
| 374 | |
| 375 | if (num_split > 16) { |
| 376 | GpuDeviceArrayOnHost<T*> ptrs(context, num_split); |
| 377 | OP_REQUIRES_OK(context, ptrs.Init()); |
| 378 | |
| 379 | GpuDeviceArrayOnHost<Tlen> offsets(context, num_split + 1); |
| 380 | OP_REQUIRES_OK(context, offsets.Init()); |
| 381 | |
| 382 | Tlen offset = 0; |
| 383 | int entry = split_sizes_vec[0]; |
| 384 | bool fixed_size = |
| 385 | std::all_of(split_sizes_vec.begin(), split_sizes_vec.end(), |
| 386 | [&entry](int n) { return n == entry; }); |
| 387 | |
| 388 | for (int i = 0; i < num_split; ++i) { |
| 389 | TensorShape output_shape(input_shape); |
| 390 | output_shape.set_dim(split_dim, split_sizes_vec[i]); |
| 391 | Tensor* result = nullptr; |
| 392 | OP_REQUIRES_OK(context, |
| 393 | context->allocate_output(i, output_shape, &result)); |
| 394 | ptrs.Set(i, result->flat<T>().data()); |
| 395 | offsets.Set(i, offset); |
| 396 | offset += split_sizes_vec[i] * suffix_dim_size; |
| 397 | } |
| 398 | offsets.Set(num_split, offset); |
| 399 | OP_REQUIRES_OK(context, ptrs.Finalize()); |
| 400 | OP_REQUIRES_OK(context, offsets.Finalize()); |
| 401 | |
| 402 | if (input.NumElements() > 0) { |
| 403 | SplitVOpGPULaunch<T, Tlen>().Run( |
| 404 | context->eigen_device<GPUDevice>(), fixed_size, |
nothing calls this directly
no test coverage detected