#include <torch/csrc/stable/accelerator.h>
#include <torch/csrc/stable/device.h>
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/macros.h>
#include <torch/csrc/stable/ops.h>
#include <torch/csrc/stable/tensor.h>

#include "compute_sparse_tile.h"
#include "pt_stable_utils.h"
#include "sparse24_metadata.h"
#include "sparse24_pack.h"

using namespace xformers::sp24;

namespace {
template <typename KT, typename Metadata, typename Algorithm>
__global__ void __launch_bounds__(32 /* num_threads */, 20)
    sparse24_sparsify_both_ways_kernel(
        typename KT::Params p,
        Metadata metadata,
        Algorithm algo) {
  KT::sparse24_sparsify_both_ways_kernel(p, metadata, algo);
}

template <typename Element, typename MetadataFormat, bool kIsMeta>
std::
    tuple<
        torch::stable::Tensor, // packed
        torch::stable::Tensor, // packed_meta_reordered
        torch::stable::Tensor, // packed_trans
        torch::stable::Tensor, // packed_trans_meta_reordered
        torch::stable::Tensor // threads_masks
        >
    sparse24_sparsify_both_ways_typed(
        const torch::stable::Tensor input,
        std::string algorithm) {
  using KT = KernelTypes<Element>;
  std::optional<torch::stable::accelerator::DeviceGuard> device_guard;
  if (!kIsMeta) {
    device_guard.emplace(input.device().index());
  }

  STD_TORCH_CHECK(input.dim() == 2, "Can only sparsify 2d tensors");
  STD_TORCH_CHECK(
      input.stride(1) == 1,
      "Can only sparsify contiguous tensors. Sparsify the transpose otherwise.");

  auto rows = input.size(0);
  auto cols = input.size(1);

  auto [compressed, packed, packed_meta_reordered] =
      MetadataFormat::create_compressed_representation(rows, cols, input, true);
  auto [compressed_trans, packed_trans, packed_trans_meta_reordered] =
      MetadataFormat::create_compressed_representation(cols, rows, input, true);
  STD_TORCH_CHECK(
      input.size(1) % 32 == 0, "Number of cols should be multiple of 32");

  typename KT::Params p;
  p.input_s0 = input.stride(0);
  p.input_dim0 = input.size(0);
  p.input_dim1 = input.size(1);
  p.packed_stride = packed.stride(0);
  p.packed_trans_stride = packed_trans.stride(0);

  MetadataFormat metadata = MetadataFormat(
      packed_meta_reordered, packed_trans_meta_reordered, rows, cols);
  torch::stable::Tensor threads_masks = torch::stable::new_empty(
      input,
      {p.getBlocksGrid().x * p.getThreadsGrid().x,
       p.getBlocksGrid().y * p.getThreadsGrid().y,
       sizeof(p.threads_masks[0])},
      torch::headeronly::ScalarType::Byte);
  if (!kIsMeta) {
    p.input = (Element const*)input.data_ptr();
    p.packed = (Element*)packed.data_ptr();
    p.packed_trans = (Element*)packed_trans.data_ptr();
    p.threads_masks = (uint64_t*)threads_masks.data_ptr();
  }

  bool kernel_launched = false;
  auto launchKernel = [&](auto algo, std::string const& algo_name) {
    if (algo_name == algorithm) {
      kernel_launched = true;
      if (kIsMeta) {
        return;
      }
      size_t smem_bytes = 0;
      sparse24_sparsify_both_ways_kernel<KT>
          <<<p.getBlocksGrid(),
             p.getThreadsGrid(),
             smem_bytes,
             xf_getCurrentCUDAStream()>>>(p, metadata, algo);
    }
  };
  named_algorithms(launchKernel);
  STD_TORCH_CHECK(kernel_launched, "Unknown algorithm \"", algorithm, "\"");
  STD_CUDA_KERNEL_LAUNCH_CHECK();
  return std::make_tuple(
      compressed,
      packed_meta_reordered,
      compressed_trans,
      packed_trans_meta_reordered,
      threads_masks);
}

template <bool kIsMeta = false>
std::
    tuple<
        torch::stable::Tensor, // packed
        torch::stable::Tensor, // packed_meta_reordered
        torch::stable::Tensor, // packed_trans
        torch::stable::Tensor, // packed_trans_meta_reordered
        torch::stable::Tensor // threads_masks
        >
    sparse24_sparsify_both_ways(
        const torch::stable::Tensor input,
        std::string algorithm,
        std::string backend) {
  auto runTyped = [&](auto type) {
    using ElementT = decltype(type);
    if (backend == "cusparselt") {
      return sparse24_sparsify_both_ways_typed<
          ElementT,
          MetadataCuSparseLtSm80,
          kIsMeta>(input, algorithm);
    } else {
      STD_TORCH_CHECK(
          backend == "cutlass",
          "backend argument only supports `cutlass` or `cusparselt`");
      return sparse24_sparsify_both_ways_typed<
          ElementT,
          MetadataCutlassSm80,
          kIsMeta>(input, algorithm);
    }
  };

  if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
    return runTyped(cutlass::half_t());
  } else {
    STD_TORCH_CHECK(
        input.scalar_type() == torch::headeronly::ScalarType::Half ||
        input.scalar_type() == torch::headeronly::ScalarType::BFloat16);
    return runTyped(cutlass::bfloat16_t());
  }
}
} // namespace

STABLE_TORCH_LIBRARY_IMPL(xformers, CUDA, m) {
  m.impl(
      "sparse24_sparsify_both_ways",
      TORCH_BOX(sparse24_sparsify_both_ways<false>));
}

STABLE_TORCH_LIBRARY_IMPL(xformers, Meta, m) {
  m.impl(
      "sparse24_sparsify_both_ways",
      TORCH_BOX(sparse24_sparsify_both_ways<true>));
}
