// coding=utf-8
//
// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation
// and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "disco.h"
#include "disco_cuda.cuh"

template <int BDIM_X, int ELXTH, typename REAL_T>
__device__ void disco_fwd_d(const int Hi, const int Wi, const int K, const int Ho, const int Wo, const int pscale,
                            const int64_t *__restrict__ roff, const int64_t *__restrict__ kers,
                            const int64_t *__restrict__ rows, const int64_t *__restrict__ cols,
                            const REAL_T *__restrict__ vals, const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
{

    const int tid = threadIdx.x;

    const int64_t bidx = blockIdx.x; // gloabl row
    const int64_t bidy = blockIdx.y; // bc

    int64_t soff = roff[bidx];
    int64_t eoff = roff[bidx + 1];

    const int64_t ker = kers[soff];
    const int64_t row = rows[soff];

    inp += bidy * Hi * Wi;
    out += bidy * K * Ho * Wo + ker * Ho * Wo + row * Wo;

    REAL_T __reg[ELXTH] = {0};

    // align to larger supported fp type
    extern __shared__ __align__(
        sizeof(double)) unsigned char __sh_ptr[]; // REAL_T __sh[2*Wi + ppscale*(BDIM_X*ELXTH - Wo)]
    REAL_T *__sh = reinterpret_cast<REAL_T *>(__sh_ptr);

    int col_prev = cols[soff];

    int h_prev = col_prev / Wi;
    int w_prev = col_prev % Wi;

    // copy current inp row in shmem
    for (int i = tid; i < Wi; i += BDIM_X) {
        const REAL_T v = inp[h_prev * Wi + i];
        __sh[i] = v;
        __sh[Wi + i] = v;
    }
    // locations __sh[2*Wi : ppscale*(BDIM_X*ELXTH-Wo)] are not used
    __syncthreads();

    // loops along the colums of CTA's row
    for (int64_t nz = soff; nz < eoff; nz++) {

        const int col = cols[nz];
        const REAL_T val = vals[nz];

        // if we are processing a nz with a col value
        // leading to a new row of inp then copy it
        // to shmem;
        // checks whether (h_prev < h) with:
        //  (col >= col_prev - (col_prev % Wi) + Wi)
        if (col >= col_prev - w_prev + Wi) {

            col_prev = col;
            h_prev = col / Wi;
            w_prev = col % Wi;

            __syncthreads();
            for (int i = tid; i < Wi; i += BDIM_X) {
                const REAL_T v = inp[h_prev * Wi + i];
                __sh[i] = v;
                __sh[Wi + i] = v;
            }
            __syncthreads();
        }

        const int w = w_prev + (col - col_prev);

#pragma unroll
        for (int i = 0; i < ELXTH; i++) {

            const int pp = i * BDIM_X + tid;

            // original lines:
            //
            //   if (pp >= Wo) break;
            //   const int wpp = (w + pscale*pp) % Wi;
            //
            // value of (w + pscale*pp) < (Wi + (Wi/Wo)*Wo) = 2*Wi
            // so we can allocate twice the amount of shmem,
            // replicate the current inp row and avoid the costly mod
            //
            // also, to avoid the conditional, sh can be extended to
            // cover the maximum location accessed during this loop
            //
            // REAL_T __sh[2*Wi + ppscale*NUM_REM]
            //
            //   Wi + (Wi/Wo)*BDIM_X*ELXTH = (since BDIM_X*ELXTH >= Wo) =
            // = Wi + (Wi/Wo)*(Wo + (BDIM_X*ELXTH - Wo)) =
            // = 2*Wi + ppscale*NUM_REM
            //
            // with NUM_REM = BDIM_X*ELXTH - Wo

            const int wpp = w + pscale * pp;

            __reg[i] += val * __sh[wpp];
        }
    }

#pragma unroll
    for (int i = 0; i < ELXTH; i++) {

        const int pp = i * BDIM_X + tid;
        if (pp >= Wo) break;

        out[pp] = __reg[i];
    }

    return;
}

template <int BDIM_X, int ELXTH, typename REAL_T>
__global__
__launch_bounds__(BDIM_X) void disco_fwd_blk_k(const int Hi, const int Wi, const int K, const int Ho, const int Wo,
                                               const int pscale, const int64_t *__restrict__ roff,
                                               const int64_t *__restrict__ kers, const int64_t *__restrict__ rows,
                                               const int64_t *__restrict__ cols, const REAL_T *__restrict__ vals,
                                               const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
{

    disco_fwd_d<BDIM_X, ELXTH>(Hi, Wi, K, Ho, Wo, pscale, roff, kers, rows, cols, vals, inp, out);

    return;
}

template <int NTH, int ELXTH, typename REAL_T>
static void launch_kernel(int BC, int Hi, int Wi, int K, int Ho, int Wo, int64_t nrows, int64_t *roff_d, int64_t *ker_d,
                          int64_t *row_d, int64_t *col_d, REAL_T *val_d, REAL_T *inp_d, REAL_T *out_d,
                          cudaStream_t stream)
{

    static_assert(sizeof(REAL_T) == 2 || sizeof(REAL_T) == 4 || sizeof(REAL_T) == 8);

    if constexpr (ELXTH <= ELXTH_MAX) {
        if (NTH * ELXTH >= Wo) {
            dim3 grid(nrows, BC);

            const int pscale = Wi / Wo;
            size_t shmem = sizeof(*out_d) * (Wi * 2 + pscale * (NTH * ELXTH - Wo));

            disco_fwd_blk_k<NTH, ELXTH><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d, row_d,
                                                                      col_d, val_d, inp_d, out_d);
        } else {
            launch_kernel<NTH, ELXTH + 1>(BC, Hi, Wi, K, Ho, Wo, nrows, roff_d, ker_d, row_d, col_d, val_d, inp_d,
                                          out_d, stream);
        }
    }
    return;
}

torch::Tensor disco_cuda_fwd(torch::Tensor inp, torch::Tensor roff_idx, torch::Tensor ker_idx, torch::Tensor row_idx,
                             torch::Tensor col_idx, torch::Tensor val, int64_t K, int64_t Ho, int64_t Wo)
{

    // some sanity checks
    CHECK_CUDA_INPUT_TENSOR(inp);
    CHECK_CUDA_INPUT_TENSOR(roff_idx);
    CHECK_CUDA_INPUT_TENSOR(ker_idx);
    CHECK_CUDA_INPUT_TENSOR(row_idx);
    CHECK_CUDA_INPUT_TENSOR(col_idx);
    CHECK_CUDA_INPUT_TENSOR(val);

    // extract some shapes
    int64_t B = inp.size(0);
    int64_t C = inp.size(1);
    int64_t BC = B * C;
    int64_t Hi = inp.size(2);
    int64_t Wi = inp.size(3);
    int64_t nrows = roff_idx.size(0) - 1;

    // allocate output
    int64_t out_dims[] = {B, C, K, Ho, Wo};
    auto options = torch::TensorOptions().device(inp.device()).dtype(inp.dtype());
    torch::Tensor out = torch::zeros(out_dims, options);

    // get stream
    auto stream = at::cuda::getCurrentCUDAStream().stream();

    // assert
    static_assert(0 == (ELXTH_MAX % 2));

    // pick the correct launch config
    if (Wo <= 64 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
                                       launch_kernel<64, 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 128 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
                                       launch_kernel<128, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 256 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
                                       launch_kernel<256, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 512 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
                                       launch_kernel<512, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 1024 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
                                       launch_kernel<1024, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else {
        fprintf(stderr, "%s:%d: error, unsupported Wo value (%ld), max supported is %d\n", __FILE__, __LINE__, Wo,
                1024 * ELXTH_MAX);
        exit(EXIT_FAILURE);
    }

    return out;
}
