Method exec

dnn/src/cambricon/argmxx/opr_impl.cpp:76–113

}

void ArgmaxForwardImpl::exec(
        _megdnn_tensor_in src, _megdnn_tensor_out dst, _megdnn_workspace workspace) {
    check_exec(src.layout, dst.layout, workspace.size);
    auto _handle = cnnl_handle(this->handle());
    ArgmxxCnnlDescs descs(src.layout, dst.layout, param().axis);
    auto ws_bundle = make_bundle(src.layout, dst.layout);
    ws_bundle.set(workspace.raw_ptr);

    void* target_src = src.raw_ptr();
    if (!src.layout.is_contiguous()) {
        TensorLayout dst;
        dst = src.layout;
        dst.init_contiguous_stride();
        CnnlTensorDescriptor dst_desc;
        dst_desc.set(dst);
        cnnl_check(cnnlCopy(
                _handle, descs.src_desc.desc(), src.raw_ptr(), dst_desc.desc(),
                ws_bundle.get(2)));
        target_src = ws_bundle.get(2);
    }

    switch (src.layout.dtype.enumv()) {
#define cb(t)                                                                  \
    case DTypeTrait<t>::enumv:                                                 \
        do {                                                                   \
            using ct = DTypeTrait<t>::ctype;                                   \
            dispatch_argmax_with_ctype<ct>(                                    \
                    _handle, static_cast<ct*>(target_src), dst.ptr<int32_t>(), \
                    ws_bundle, descs);                                         \
            return;                                                            \
        } while (0);
        cb(::megdnn::dtype::Float32);
        DNN_INC_FLOAT16(cb(::megdnn::dtype::Float16));
        MEGDNN_FOREACH_COMPUTING_DTYPE_INT(cb);
        default:
            megdnn_throw("unsupported dtype in cambricon ArgmaxImpl");
    }
}

}  // namespace cambricon
}  // namespace megdnn