/**
 * \file dnn/src/cambricon/checksum/opr_impl.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#include "src/cambricon/checksum/opr_impl.h"
#include "src/cambricon/checksum/checksum.mlu.h"

#include "src/cambricon/utils.h"

#include <algorithm>

using namespace megdnn;
using namespace cambricon;

namespace {
void bang_c_wrapper(
        uint32_t* dst, const uint32_t* src, int nr_elems, cnrtQueue_t queue,
        cnrtCoreVersion_t core_version) {
    cnrtKernelParamsBuffer_t params;
    cnrt_check(cnrtGetKernelParamsBuffer(&params));
    cnrt_check(cnrtKernelParamsBufferAddParam(params, &dst, sizeof(uint32_t*)));
    cnrt_check(cnrtKernelParamsBufferAddParam(params, &src, sizeof(uint32_t*)));
    cnrt_check(cnrtKernelParamsBufferAddParam(params, &nr_elems, sizeof(int)));
    if (core_version == CNRT_MLU270) {
        cnrtDim3_t dim;
        dim.x = 16;
        dim.y = 1;
        dim.z = 1;
        cnrtFunctionType_t c = CNRT_FUNC_TYPE_UNION4;
        cnrt_check(cnrtInvokeKernel_V2(
                (void*)&checksum_kernel_union4, dim, params, c, queue));
    } else if (core_version == CNRT_MLU220) {
        cnrtDim3_t dim;
        dim.x = 4;
        dim.y = 1;
        dim.z = 1;
        cnrtFunctionType_t c = CNRT_FUNC_TYPE_UNION1;
        cnrt_check(cnrtInvokeKernel_V2(
                (void*)&checksum_kernel_union1, dim, params, c, queue));
    }
    after_kernel_launch();
    cnrt_check(cnrtDestroyKernelParamsBuffer(params));
}
}  // namespace

size_t ChecksumForwardImpl::get_workspace_in_bytes(const TensorLayout& /* data */) {
    size_t ws_size = sizeof(ChecksumForward::Result::checksum);
    return ws_size;
}

ChecksumForward::Result ChecksumForwardImpl::exec(
        _megdnn_tensor_in data, _megdnn_workspace workspace) {
    Result result;
    memset(&result, 0, sizeof(result));
    check_exec(data.layout, workspace.size);
    auto queue = cnrt_queue(handle());

    auto ptr = static_cast<uint8_t*>(data.raw_ptr());
    size_t size_all = data.layout.shape[0], size_ints = size_all / sizeof(uint32_t);
    auto last_val_size = std::min<size_t>(size_all, 4);
    cnrt_check(cnrtMemcpyAsync(
            &result.last_val, ptr + size_all - last_val_size, last_val_size, queue,
            CNRT_MEM_TRANS_DIR_DEV2HOST));
    if (size_ints) {
        auto&& device_info = current_device_info();
        bang_c_wrapper(
                reinterpret_cast<uint32_t*>(workspace.raw_ptr),
                static_cast<uint32_t*>(data.raw_ptr()), size_ints, queue,
                device_info.core_version);
        cnrt_check(cnrtMemcpyAsync(
                &result.checksum, workspace.raw_ptr, sizeof(result.checksum), queue,
                CNRT_MEM_TRANS_DIR_DEV2HOST));
    }
    cnrt_check(cnrtSyncQueue(queue));
    return result;
}

// vim: syntax=cpp.doxygen