#include <torch/library.h>
#include <torch/all.h>

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

#include "registration.h"
#include "torch_binding.h"
#include "../metakernel_cuda/probes_launch.h"

namespace {

void chk(bool ok, const char* msg) { TORCH_CHECK(ok, msg); }

cudaStream_t stream_of(const torch::Tensor& t) {
  const at::cuda::CUDAGuard guard(t.device());
  return at::cuda::getCurrentCUDAStream();
}

}  // namespace

void mk_triad(torch::Tensor a, torch::Tensor b, torch::Tensor c,
              int64_t width, double s) {
  chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32,
      "a must be contiguous CUDA f32");
  chk(b.sizes() == a.sizes() && c.sizes() == a.sizes(), "size mismatch");
  chk(width == 1 || width == 2 || width == 4, "width must be 1|2|4");
  chk(a.numel() % width == 0, "numel must divide width");
  mk_triad_launch(a.const_data_ptr<float>(), b.const_data_ptr<float>(),
                  c.data_ptr<float>(), a.numel() / width, (int)width,
                  (float)s, stream_of(a));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_triad_passes(torch::Tensor a, torch::Tensor b, torch::Tensor c,
                     double s, int64_t passes) {
  chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32,
      "a must be contiguous CUDA f32");
  chk(a.numel() % 4 == 0, "numel must divide 4");
  chk(passes >= 1 && passes <= 1000000, "passes in [1, 1e6]");
  mk_triad_passes_launch(a.const_data_ptr<float>(), b.const_data_ptr<float>(),
                         c.data_ptr<float>(), a.numel() / 4, 4, (float)s,
                         (int)passes, stream_of(a));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_read(torch::Tensor a, torch::Tensor sink) {
  chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32 &&
          a.numel() % 4 == 0,
      "a must be contiguous CUDA f32, numel % 4 == 0");
  mk_read_launch(a.const_data_ptr<float>(), a.numel() / 4,
                 sink.data_ptr<float>(), stream_of(a));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_write(torch::Tensor c, double v) {
  chk(c.is_cuda() && c.is_contiguous() && c.dtype() == torch::kFloat32 &&
          c.numel() % 4 == 0,
      "c must be contiguous CUDA f32, numel % 4 == 0");
  mk_write_launch(c.data_ptr<float>(), c.numel() / 4, (float)v,
                  stream_of(c));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_gather(torch::Tensor a, torch::Tensor idx, torch::Tensor c,
               int64_t threads, int64_t dyn_smem) {
  chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32 &&
          a.numel() % 4 == 0,
      "a must be contiguous CUDA f32, numel % 4 == 0");
  chk(idx.is_cuda() && idx.dtype() == torch::kInt32 &&
          idx.numel() * 4 == c.numel() && c.numel() % 4 == 0,
      "idx i32 [n4] with c f32 [n4*4]");
  chk(threads >= 32 && threads <= 1024 && dyn_smem >= 0 &&
          dyn_smem <= 48 * 1024,
      "bad launch shape");
  mk_gather_launch(a.const_data_ptr<float>(), idx.const_data_ptr<int>(),
                   c.data_ptr<float>(), idx.numel(), (int)threads,
                   (int)dyn_smem, stream_of(a));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_fma_f64(int64_t blocks, int64_t threads, int64_t iters,
                torch::Tensor sink) {
  chk(sink.is_cuda() && sink.dtype() == torch::kFloat64 &&
          sink.numel() >= 1,
      "sink must be CUDA f64 [>=1]");
  mk_fma_f64_launch((int)blocks, (int)threads, iters,
                    sink.data_ptr<double>(), stream_of(sink));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_fma_dep(int64_t blocks, int64_t threads, int64_t dyn_smem,
                int64_t iters, torch::Tensor sink) {
  chk(sink.is_cuda() && sink.dtype() == torch::kFloat32 && sink.numel() >= 1,
      "sink must be CUDA f32 [>=1]");
  chk(dyn_smem >= 0 && dyn_smem <= 48 * 1024, "dyn_smem in [0, 48K]");
  mk_fma_dep_launch((int)blocks, (int)threads, (int)dyn_smem, iters,
                    sink.data_ptr<float>(), stream_of(sink));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_chase_global(torch::Tensor ring, int64_t hops, torch::Tensor out) {
  chk(ring.is_cuda() && ring.is_contiguous() &&
          ring.dtype() == torch::kInt32,
      "ring must be contiguous CUDA i32");
  chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 2,
      "out must be CUDA i64 [>=2]");
  mk_chase_global_launch(ring.const_data_ptr<int>(), hops,
                         reinterpret_cast<long long*>(out.data_ptr<int64_t>()), stream_of(ring));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_chase_shared(torch::Tensor ring, int64_t hops, torch::Tensor out) {
  chk(ring.is_cuda() && ring.is_contiguous() &&
          ring.dtype() == torch::kInt32,
      "ring must be contiguous CUDA i32");
  chk(ring.numel() <= 12288, "shared ring must fit 48 KB (<= 12288 ints)");
  chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 2,
      "out must be CUDA i64 [>=2]");
  mk_chase_shared_launch(ring.const_data_ptr<int>(), (int)ring.numel(),
                         hops, reinterpret_cast<long long*>(out.data_ptr<int64_t>()), stream_of(ring));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_fma_f32(int64_t blocks, int64_t threads, int64_t iters,
                torch::Tensor sink) {
  chk(sink.is_cuda() && sink.dtype() == torch::kFloat32 && sink.numel() >= 1,
      "sink must be CUDA f32 [>=1]");
  chk(blocks >= 1 && blocks <= 1 << 20 && threads >= 32 && threads <= 1024,
      "bad launch shape");
  mk_fma_f32_launch((int)blocks, (int)threads, iters,
                    sink.data_ptr<float>(), stream_of(sink));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

int64_t mk_mma(int64_t kind, int64_t blocks, int64_t warps, int64_t iters,
               torch::Tensor sink) {
  chk(sink.is_cuda() && sink.dtype() == torch::kFloat32 && sink.numel() >= 1,
      "sink must be CUDA f32 [>=1]");
  chk(kind >= 0 && kind <= 4, "kind must be 0..4");
  chk(blocks >= 1 && warps >= 1 && warps <= 32, "bad launch shape");
  cudaDeviceProp prop;
  cudaGetDeviceProperties(&prop, sink.device().index());
  int r = mk_mma_launch((int)kind, (int)blocks, (int)warps, iters,
                        sink.data_ptr<float>(), prop.major, prop.minor,
                        stream_of(sink));
  if (r == 0) C10_CUDA_KERNEL_LAUNCH_CHECK();
  return r;
}

void mk_atomics(int64_t mode, int64_t blocks, int64_t threads, int64_t iters,
                torch::Tensor buf, int64_t slots) {
  chk(buf.is_cuda() && buf.dtype() == torch::kFloat32 && buf.numel() >= 2,
      "buf must be CUDA f32 [>=2]");
  chk(mode >= 0 && mode <= 3, "mode must be 0..3");
  chk(slots >= 1 && slots <= buf.numel(), "slots in [1, buf.numel()]");
  mk_atomics_launch((int)mode, (int)blocks, (int)threads, iters,
                    buf.data_ptr<float>(), slots, stream_of(buf));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_spin(int64_t blocks, int64_t threads, int64_t ticks,
             torch::Tensor out, torch::Tensor sink) {
  chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 1,
      "out must be CUDA i64 [>=1]");
  chk(sink.is_cuda() && sink.dtype() == torch::kFloat32, "sink f32");
  mk_spin_launch((int)blocks, (int)threads, ticks,
                 reinterpret_cast<long long*>(out.data_ptr<int64_t>()), sink.data_ptr<float>(),
                 stream_of(out));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void mk_empty(torch::Tensor device_anchor) {
  mk_empty_launch(stream_of(device_anchor));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

int64_t mk_barrier(int64_t blocks, int64_t threads, int64_t iters,
                   torch::Tensor out) {
  chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 1,
      "out must be CUDA i64 [>=1]");
  int r = mk_barrier_launch((int)blocks, (int)threads, (int)iters,
                            reinterpret_cast<long long*>(out.data_ptr<int64_t>()), stream_of(out));
  if (r == 0) C10_CUDA_KERNEL_LAUNCH_CHECK();
  return r;
}

int64_t mk_occupancy(int64_t threads, int64_t dyn_smem,
                     torch::Tensor device_anchor) {
  const at::cuda::CUDAGuard guard(device_anchor.device());
  int blocks_per_sm = -1;
  int r = mk_occupancy_triad((int)threads, (int)dyn_smem, &blocks_per_sm);
  return r == 0 ? blocks_per_sm : -1;
}

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("mk_triad(Tensor a, Tensor b, Tensor! c, int width, float s) -> ()");
  ops.impl("mk_triad", torch::kCUDA, &mk_triad);
  ops.def(
      "mk_triad_passes(Tensor a, Tensor b, Tensor! c, float s, int passes)"
      " -> ()");
  ops.impl("mk_triad_passes", torch::kCUDA, &mk_triad_passes);
  ops.def("mk_read(Tensor a, Tensor! sink) -> ()");
  ops.impl("mk_read", torch::kCUDA, &mk_read);
  ops.def("mk_write(Tensor! c, float v) -> ()");
  ops.impl("mk_write", torch::kCUDA, &mk_write);
  ops.def(
      "mk_gather(Tensor a, Tensor idx, Tensor! c, int threads,"
      " int dyn_smem) -> ()");
  ops.impl("mk_gather", torch::kCUDA, &mk_gather);
  ops.def(
      "mk_fma_f64(int blocks, int threads, int iters, Tensor! sink) -> ()");
  ops.impl("mk_fma_f64", torch::kCUDA, &mk_fma_f64);
  ops.def(
      "mk_fma_dep(int blocks, int threads, int dyn_smem, int iters,"
      " Tensor! sink) -> ()");
  ops.impl("mk_fma_dep", torch::kCUDA, &mk_fma_dep);
  ops.def("mk_chase_global(Tensor ring, int hops, Tensor! out) -> ()");
  ops.impl("mk_chase_global", torch::kCUDA, &mk_chase_global);
  ops.def("mk_chase_shared(Tensor ring, int hops, Tensor! out) -> ()");
  ops.impl("mk_chase_shared", torch::kCUDA, &mk_chase_shared);
  ops.def(
      "mk_fma_f32(int blocks, int threads, int iters, Tensor! sink) -> ()");
  ops.impl("mk_fma_f32", torch::kCUDA, &mk_fma_f32);
  ops.def(
      "mk_mma(int kind, int blocks, int warps, int iters, Tensor! sink)"
      " -> int");
  ops.impl("mk_mma", torch::kCUDA, &mk_mma);
  ops.def(
      "mk_atomics(int mode, int blocks, int threads, int iters, Tensor! buf,"
      " int slots) -> ()");
  ops.impl("mk_atomics", torch::kCUDA, &mk_atomics);
  ops.def(
      "mk_spin(int blocks, int threads, int ticks, Tensor! out,"
      " Tensor! sink) -> ()");
  ops.impl("mk_spin", torch::kCUDA, &mk_spin);
  ops.def("mk_empty(Tensor device_anchor) -> ()");
  ops.impl("mk_empty", torch::kCUDA, &mk_empty);
  ops.def(
      "mk_barrier(int blocks, int threads, int iters, Tensor! out) -> int");
  ops.impl("mk_barrier", torch::kCUDA, &mk_barrier);
  ops.def(
      "mk_occupancy(int threads, int dyn_smem, Tensor device_anchor) -> int");
  ops.impl("mk_occupancy", torch::kCUDA, &mk_occupancy);
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)