/**
 * \file dnn/src/rocm/argsort/backward.cpp.hip
 * 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 "hcc_detail/hcc_defs_prologue.h"

#include "src/rocm/utils.h.hip"
#include "./argsort.h.hip"
#include "./backward.h.hip"

using namespace megdnn;
using namespace rocm;
using namespace argsort;

namespace {

template <typename T>
__global__ void backward_kernel(uint32_t dst_w, uint32_t src_w,
                                uint32_t src_size, T* dst, const T* src_data,
                                const int* src_idx) {
    uint32_t idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < src_size) {
        uint32_t r = idx / src_w;
        dst[r * dst_w + src_idx[idx]] = src_data[idx];
    }
}

}  // namespace

template <typename T>
void argsort::backward_proxy(uint32_t dst_h, uint32_t dst_w, uint32_t src_w,
                             T* dst, const T* src_data, const int* src_idx,
                             hipStream_t stream) {
    if (dst_w != src_w) {
        hipMemsetAsync(dst, 0, dst_h * dst_w * sizeof(T), stream);
    }

    uint32_t src_size = dst_h * src_w;
    backward_kernel<<<DIVUP(src_size, 512), 512, 0, stream>>>(
            dst_w, src_w, src_size, dst, src_data, src_idx);
    after_kernel_launch();
}

namespace megdnn {
namespace rocm {
namespace argsort {

#define INST(T)                                                             \
    template void backward_proxy(uint32_t dst_h, uint32_t dst_w,            \
                                 uint32_t src_w, T* dst, const T* src_data, \
                                 const int* src_idx, hipStream_t stream);
ARGSORT_FOREACH_CTYPE(INST)
#undef INST

}  // namespace argsort
}  // namespace rocm
}  // namespace megdnn

// vim: syntax=cpp.doxygen