Method exec

dnn/src/rocm/argmxx/opr_impl.cpp:51–68

}

void ArgmaxForwardImpl::exec(
        _megdnn_tensor_in src, _megdnn_tensor_out dst, _megdnn_workspace workspace) {
    check_exec(src.layout, dst.layout, workspace.size);
    size_t A, B, C;
    reduce::get_ABC(src.layout, A, B, C, param().axis);
    auto stream = hip_stream(handle());
#define cb(dt)                                                                     \
    if (src.layout.dtype.enumv() == DTypeTrait<dt>::enumv) {                       \
        using ctype = typename DTypeTrait<dt>::ctype;                              \
        exec_impl<ctype, true>(                                                    \
                src.ptr<ctype>(), dst.ptr<dt_int32>(), workspace.raw_ptr, A, B, C, \
                stream);                                                           \
        return;                                                                    \
    }
    MEGDNN_FOREACH_COMPUTING_DTYPE_FLOAT(cb)
#undef cb
    megdnn_throw(ssprintf("Unsupported DType: %s", src.layout.dtype.name()));
}

size_t ArgminForwardImpl::get_workspace_in_bytes(
        const TensorLayout& src, const TensorLayout& dst) {