Method exec_internal

dnn/src/cuda/split/opr_impl.cpp:26–99

template <typename T>
void SplitForwardImpl::exec_internal(
        _megdnn_tensor_in src, const TensorNDArray& dsts, _megdnn_workspace workspace) {
    // Please refer to ConcatForwardImpl. Implementations are similar.
    auto dsts_layout = apply_vector<TensorLayout>(m_get_layout, dsts);
    auto dsts_shape = apply_vector<TensorShape>(m_get_shape, dsts_layout);
    check_exec(src.layout, dsts_layout, workspace.size);
    size_t A, B, C;
    auto stream = cuda_stream(this->handle());

    // Pre-calculate B to determine cpu-side workspace size.
    B = src.layout.shape[param().axis];

    // workspace_cpu will be freed by cuda callback.
    SmallVector<size_t> workspace_sizes{
            sizeof(const T*) * dsts.size(),
            sizeof(size_t) * dsts.size(),
            sizeof(size_t) * B,
            sizeof(size_t) * B,
    };

    WorkspaceBundle workspace_cpu(nullptr, workspace_sizes),
            workspace_gpu(nullptr, workspace_sizes);

    auto total_workspace_size = workspace_cpu.total_size_in_bytes();
    void* workspace_cpu_raw = malloc(total_workspace_size);
    megdnn_assert_internal(workspace_cpu_raw);
    void* workspace_gpu_raw = static_cast<void*>(workspace.raw_ptr);
    workspace_cpu = WorkspaceBundle(workspace_cpu_raw, workspace_sizes);
    workspace_gpu = WorkspaceBundle(workspace_gpu_raw, workspace_sizes);

    auto dsts_cpu = static_cast<T**>(workspace_cpu.get(0));
    auto dsts_gpu = static_cast<T**>(workspace_gpu.get(0));
    for (size_t i = 0; i < dsts.size(); ++i) {
        dsts_cpu[i] = dsts[i].ptr<T>();
    }

    auto Bv_cpu = static_cast<size_t*>(workspace_cpu.get(1));
    auto Bv_gpu = static_cast<size_t*>(workspace_gpu.get(1));
    get_ABC(dsts_shape, A, Bv_cpu, C);

    auto table_outer_cpu = static_cast<size_t*>(workspace_cpu.get(2));
    auto table_outer_gpu = static_cast<size_t*>(workspace_gpu.get(2));
    auto table_inner_cpu = static_cast<size_t*>(workspace_cpu.get(3));
    auto table_inner_gpu = static_cast<size_t*>(workspace_gpu.get(3));
    {
        size_t outer_idx = 0, inner_idx = 0;

        for (size_t i = 0; i < B; ++i) {
            table_outer_cpu[i] = outer_idx;
            table_inner_cpu[i] = inner_idx;
            ++inner_idx;
            if (inner_idx == Bv_cpu[outer_idx]) {
                ++outer_idx;
                inner_idx = 0;
            }
        }
    }
    for (size_t i = 0; i < workspace_cpu.nr_workspace(); ++i) {