Method dev_tensor_memset

src/core/impl/tensor.cpp:725–763

/* =================== misc =================== */

void mgb::dev_tensor_memset(const DeviceTensorND& tensor, int val) {
    auto&& env = CompNodeEnv::from_comp_node(tensor.comp_node());
    env.activate();
    size_t size = tensor.layout().span().dist_byte();
    switch (env.property().type) {
#if MGB_CUDA
        case CompNode::DeviceType::CUDA:
            MGB_CUDA_CHECK(cudaMemsetAsync(
                    tensor.raw_ptr(), val, size, env.cuda_env().stream));
            break;
#endif
#if MGB_ATLAS
        case CompNode::DeviceType::ATLAS:
#if MGB_USE_ATLAS_ASYNC_API
            MGB_ATLAS_CHECK(aclrtMemsetAsync(
                    tensor.raw_ptr(), -1, val, size, env.atlas_env().stream));
#else
            MGB_ATLAS_CHECK(aclrtMemset(tensor.raw_ptr(), -1, val, size));
#endif
            break;
#endif
#if MGB_CAMBRICON
        case CompNode::DeviceType::CAMBRICON:
            MGB_CNRT_CHECK(
                    cnrtMemsetAsync(tensor.raw_ptr(), val, size, env.cnrt_env().queue));
            break;
#endif
        case CompNode::DeviceType::CPU: {
            auto fill = [tensor, size, val]() {
                std::memset(tensor.as_megdnn().raw_ptr(), val, size);
            };
            env.cpu_env().dispatch(fill);
        } break;
        default:
            mgb_throw(
                    MegBrainError, "unhandled comp node in dev_tensor_memset: %s",
                    tensor.comp_node().to_string().c_str());
    }
}

namespace mgb {
template class TensorStorage<HostTensorStorageTrait>;