Method alloc_cnnl_workspace

dnn/src/cambricon/elemwise/opr_impl.cpp:871–889

}  // anonymous namespace

WorkspaceBundle ElemwiseForwardImpl::alloc_cnnl_workspace(
        const TensorNDArray& src, const TensorND& dst) {
    TensorLayoutArray src_layouts(src.size());
    std::transform(
            src.begin(), src.end(), src_layouts.begin(),
            [](const TensorND& tensor) { return tensor.layout; });
    auto mode = m_param.mode;
    auto handle = concrete_handle(this->handle());
    int nr_operands = src.size();
    if (nr_operands == 1) {  // unary mode
        return get_unary_ws(handle, src_layouts, dst.layout, mode);
    } else if (nr_operands == 2) {  // binary mode
        return get_binary_ws(handle, src_layouts, dst.layout, mode);
    } else if (nr_operands == 3) {
        return get_ternary_ws(handle, src_layouts, dst.layout, mode);
    } else {
        megdnn_throw("elemwise unsupport mode");
    }
}

void ElemwiseForwardImpl::free_cnnl_workspace(const WorkspaceBundle& wk_bundle) {
    concrete_handle(handle())->free(wk_bundle.ptr());