| 644 | } |
| 645 | |
| 646 | void RepeatDataToFrom(bool broadcast_flag, const vector<size_t> &repeats, |
| 647 | int axis, Tensor *dst, const Tensor &src, |
| 648 | const size_t num) { |
| 649 | if (repeats.size() == 1) { |
| 650 | broadcast_flag = true; |
| 651 | } else if (repeats.size() > 1) { |
| 652 | if (axis == Noaxis) { |
| 653 | LOG(FATAL) << "When repeats parameter is sequence, axis cannot be None"; |
| 654 | } |
| 655 | } |
| 656 | for (size_t i = 0; i < repeats.size(); i++) { |
| 657 | CHECK_GE(repeats[i], 0); |
| 658 | } |
| 659 | auto width = SizeOf(src.data_type()); |
| 660 | CHECK_EQ(width, SizeOf(dst->data_type())); |
| 661 | // size_t nBytes = num * width; |
| 662 | int chunk = width; |
| 663 | int axis_shape = 1; |
| 664 | int shape_outer = 1; |
| 665 | if (axis == Noaxis) { |
| 666 | axis_shape = 1; |
| 667 | shape_outer = Product(src.shape()); |
| 668 | } else { |
| 669 | for (int i = 0; i < axis; i++) { |
| 670 | shape_outer *= src.shape()[i]; |
| 671 | } |
| 672 | axis_shape = src.shape()[axis]; |
| 673 | for (int i = axis + 1; i < static_cast<int>(src.nDim()); i++) { |
| 674 | chunk *= src.shape()[i]; |
| 675 | } |
| 676 | } |
| 677 | |
| 678 | Device *dev = nullptr; |
| 679 | CopyDirection direct; |
| 680 | std::shared_ptr<Device> src_dev = src.device(), dst_dev = dst->device(); |
| 681 | if (dst_dev->lang() != src_dev->lang()) { |
| 682 | // let the none cpp device conduct copy op |
| 683 | if (dst_dev->lang() == kCpp) { |
| 684 | dev = src_dev.get(); |
| 685 | direct = kDeviceToHost; |
| 686 | } else if (src_dev->lang() == kCpp) { |
| 687 | dev = dst_dev.get(); |
| 688 | direct = kHostToDevice; |
| 689 | } else { |
| 690 | LOG(FATAL) |
| 691 | << "Not support mem repeat copy between Cuda and OpenCL device"; |
| 692 | } |
| 693 | } else { |
| 694 | dev = src_dev.get(); |
| 695 | direct = src_dev->lang() == kCpp ? kHostToHost : kDeviceToDevice; |
| 696 | } |
| 697 | |
| 698 | int dst_offset = 0; |
| 699 | int src_offset = 0; |
| 700 | Tensor &dstRef = *dst; |
| 701 | for (int i = 0; i < shape_outer; i++) { |
| 702 | for (int j = 0; j < axis_shape; j++) { |
| 703 | int temp = broadcast_flag ? repeats[0] : repeats[j]; |
no test coverage detected