Method exec

dnn/src/cambricon/indexing_multi_axis_vec/opr_impl.cpp:394–478

TODO: check length of indices and delete is_indices_defined The indexed axes are contiguous just like [..., indexed_axes, ...]

// TODO: check length of indices and delete is_indices_defined
// The indexed axes are contiguous just like [..., indexed_axes, ...]
void IndexingMultiAxisVecImpl::exec(
        _megdnn_tensor_in src, const IndexDesc& index, _megdnn_tensor_out dst,
        _megdnn_workspace workspace) {
    auto cambricon_handle = concrete_handle(handle());

    WorkspaceBundle wk_bundle = get_workspace_bundle(src.layout, index, dst.layout);
    size_t wk_size = wk_bundle.total_size_in_bytes();
    void* wk_ptr = cambricon_handle->alloc(wk_size);
    wk_bundle.set(wk_ptr);

    auto src_indices_and_defined = prepare_src_and_index(src, index, wk_bundle);
    auto prepared_src = std::get<0>(src_indices_and_defined);
    auto prepared_indices = std::get<1>(src_indices_and_defined);
    auto is_indices_defined = std::get<2>(src_indices_and_defined);

    CnnlTensorDescriptor cnnl_src_desc, cnnl_dst_desc, cnnl_output_dim_desc,
            cnnl_expand_output_dims_desc;
    cnnl_src_desc.set(&prepared_src);
    if (dst.layout.is_contiguous()) {
        cnnl_dst_desc.set(&dst);
    } else {
        TensorLayout dst_contig_layout(dst.layout);
        dst_contig_layout.init_contiguous_stride();
        cnnl_dst_desc.set(dst_contig_layout);
    }

    size_t len = 8;
    CnnlTensorDescriptor cnnl_full_index_desc[len];
    for (size_t i = 0; i < len; ++i) {
        if (!is_indices_defined[i]) {
        } else {
            cnnl_full_index_desc[i].set(prepared_indices[i].layout);
        }
    }
    cnnlTensorDescriptor_t cnnl_full_index_desc_t[len];
    for (size_t i = 0; i < 8; ++i) {
        cnnl_full_index_desc_t[i] = cnnl_full_index_desc[i].desc();
    }
    void* indices_ptr[len];
    for (size_t i = 0; i < len; ++i) {
        if (!is_indices_defined[i]) {
            indices_ptr[i] = nullptr;
        } else {
            indices_ptr[i] = prepared_indices[i].raw_ptr();
        }
    }

    void* output_dim = wk_bundle.get(2);
    void* output_dims = wk_bundle.get(3);
    std::vector<size_t> output_dim_shape(1, 1);
    cnnl_output_dim_desc.set(1, output_dim_shape, CNNL_DTYPE_INT32);
    std::vector<size_t> expand_output_dims_shape(1, 8);
    cnnl_expand_output_dims_desc.set(1, expand_output_dims_shape, CNNL_DTYPE_INT64);

    size_t advanced_index_wk_size = 0;
    cnnl_check(cnnlGetAdvancedIndexWorkspaceSize(
            cambricon_handle->cnnl_handle(), cnnl_src_desc.desc(),
            cnnl_full_index_desc_t, &advanced_index_wk_size));