| 56 | |
| 57 | template <typename T> |
| 58 | void CheckNonFiniteImpl::_exec( |
| 59 | _megdnn_in const TensorNDArray& srcs, _megdnn_tensor_out dst, |
| 60 | _megdnn_workspace workspace) { |
| 61 | check_exec(srcs, dst, workspace.size); |
| 62 | typedef CheckNonFiniteOp<T, dt_int32, dt_int32> Op; |
| 63 | auto stream = cuda_stream(this->handle()); |
| 64 | SmallVector<size_t> workspace_sizes{ |
| 65 | sizeof(T*) * m_size, |
| 66 | sizeof(size_t) * m_size, |
| 67 | }; |
| 68 | WorkspaceBundle workspace_cpu(nullptr, workspace_sizes), |
| 69 | workspace_gpu(nullptr, workspace_sizes); |
| 70 | auto total_workspace_size = workspace_cpu.total_size_in_bytes(); |
| 71 | void* workspace_cpu_raw = malloc(total_workspace_size); |
| 72 | megdnn_assert_internal(workspace_cpu_raw); |
| 73 | void* workspace_gpu_raw = workspace.raw_ptr; |
| 74 | workspace_cpu = WorkspaceBundle(workspace_cpu_raw, workspace_sizes); |
| 75 | workspace_gpu = WorkspaceBundle(workspace_gpu_raw, workspace_sizes); |
| 76 | |
| 77 | auto srcs_cpu = static_cast<T**>(workspace_cpu.get(0)); |
| 78 | auto srcs_gpu = static_cast<T**>(workspace_gpu.get(0)); |
| 79 | auto srcs_total_nr_elems_cpu = static_cast<size_t*>(workspace_cpu.get(1)); |
| 80 | auto srcs_total_nr_elems_gpu = static_cast<size_t*>(workspace_gpu.get(1)); |
| 81 | |
| 82 | // srcs |
| 83 | // cut the tensor to a fixed length of total_nr_elems_max |
| 84 | size_t i = 0; |
| 85 | for (const auto& src : srcs) { |
| 86 | size_t src_nr_elems = src.layout.total_nr_elems(); |
| 87 | size_t nr_elems = DIVUP(src_nr_elems, total_nr_elems_max); |
| 88 | for (size_t j = 0; j < nr_elems; ++j, ++i) { |
| 89 | srcs_cpu[i] = src.ptr<T>() + j * total_nr_elems_max; |
| 90 | if (j + 1 == nr_elems && src_nr_elems % total_nr_elems_max) { |
| 91 | srcs_total_nr_elems_cpu[i] = src_nr_elems % total_nr_elems_max; |
| 92 | } else { |
| 93 | srcs_total_nr_elems_cpu[i] = total_nr_elems_max; |
| 94 | } |
| 95 | } |
| 96 | } |
| 97 | for (size_t i = 0; i < workspace_cpu.nr_workspace(); ++i) { |
| 98 | cuda_check(cudaMemcpyAsync( |
| 99 | workspace_gpu.get(i), workspace_cpu.get(i), workspace_cpu.get_size(i), |
| 100 | cudaMemcpyHostToDevice, stream)); |
| 101 | } |
| 102 | cuda_check(cudaStreamAddCallback( |
| 103 | stream, callback_free, static_cast<void*>(workspace_cpu_raw), 0)); |
| 104 | |
| 105 | run_reduce<Op, false>( |
| 106 | static_cast<dt_int32*>( |
| 107 | (void*)((char*)workspace_gpu_raw + |
| 108 | workspace_gpu.total_size_in_bytes())), |
| 109 | 1, m_size * total_nr_elems_max, 1, stream, |
| 110 | Op(srcs_gpu, srcs_total_nr_elems_gpu, dst.ptr<dt_int32>(), |
| 111 | total_nr_elems_max, static_cast<T>(param().scale))); |
| 112 | } |
| 113 | |
| 114 | } // namespace cuda |
| 115 | } // namespace megdnn |
nothing calls this directly
no test coverage detected