/* Copyright 2022 The JAX Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>
#include <string_view>
#include <tuple>
#include <utility>
#include <vector>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/ascii.h"
#include "nanobind/nanobind.h"
#include "nanobind/stl/optional.h"  // IWYU pragma: keep
#include "nanobind/stl/pair.h"  // IWYU pragma: keep
#include "nanobind/stl/string.h"  // IWYU pragma: keep
#include "nanobind/stl/string_view.h"  // IWYU pragma: keep
#include "nanobind/stl/tuple.h"  // IWYU pragma: keep
#include "nanobind/stl/vector.h"  // IWYU pragma: keep
#include "jaxlib/absl_status_casters.h"
#include "jaxlib/gpu/gpu_kernel_helpers.h"
#include "jaxlib/gpu/triton.pb.h"
#include "jaxlib/gpu/triton_kernels.h"
#include "jaxlib/gpu/triton_utils.h"
#include "jaxlib/gpu/vendor.h"
#include "jaxlib/kernel_nanobind_helpers.h"

#define GPU_RETURN_IF_ERROR(expr) JAX_RETURN_IF_ERROR(JAX_AS_STATUS(expr))

namespace nb = nanobind;

namespace jax::JAX_GPU_NAMESPACE {

nb::dict Registrations() {
  nb::dict dict;
  nb::dict gpu_dict;
  gpu_dict["instantiate"] =
      EncapsulateFfiHandler(kTritonKernelCallFfiInstantiate);
  gpu_dict["execute"] = EncapsulateFfiHandler(kTritonKernelCallFfi);
  dict["triton_kernel_call_ffi"] = gpu_dict;
  return dict;
}

NB_MODULE(_triton, m) {
  nb::class_<Kernel>(m, "TritonKernel")
      .def(nb::init<std::string, uint32_t, uint32_t, uint32_t, std::string,
                    std::string, int, std::optional<uint32_t>,
                    std::optional<uint32_t>>(),
           nb::arg("kernel_name"), nb::arg("num_warps"), nb::arg("num_ctas"),
           nb::arg("shared_mem_bytes"), nb::arg("ptx"), nb::arg("ttir"),
           nb::arg("compute_capability"),
           nb::arg("global_scratch_size") = nb::none(),
           nb::arg("global_scratch_align") = nb::none());

  nb::class_<KernelCall::Parameter>(m, "TritonParameter");

  m.def("create_array_parameter",
        [](size_t bytes_to_zero, size_t ptr_divisibility) {
          return KernelCall::Parameter{
              KernelCall::Parameter::Array{bytes_to_zero, ptr_divisibility}};
        });

  m.def(
      "create_tma_descriptor_parameter",
      [](uint32_t elem_type, uint32_t swizzle, std::vector<uint64_t> shape,
         std::vector<uint64_t> strides, std::vector<uint32_t> block_shape,
         uint32_t oob_fill) {
        return KernelCall::Parameter{KernelCall::Parameter::TmaDescriptor{
            elem_type, swizzle, std::move(shape), std::move(strides),
            std::move(block_shape), oob_fill}};
      },
      nb::arg("elem_type"), nb::arg("swizzle"), nb::arg("shape"),
      nb::arg("strides"), nb::arg("block_shape"), nb::arg("oob_fill"));

  m.def("create_scalar_parameter",
        ValueOrThrowWrapper([](bool value, std::string_view dtype)
                                -> absl::StatusOr<KernelCall::Parameter> {
          if ((dtype == "i1") || (dtype == "B")) {
            return KernelCall::Parameter{value};
          } else {
            return absl::InvalidArgumentError(std::string("unknown dtype: ") +
                                              dtype.data());
          }
        }));

  m.def("create_scalar_parameter",
        ValueOrThrowWrapper([](nb::int_ value, std::string_view dtype)
                                -> absl::StatusOr<KernelCall::Parameter> {
          if (dtype == "i32") {
            return KernelCall::Parameter{static_cast<int32_t>(value)};
          } else if (dtype == "u32") {
            return KernelCall::Parameter{static_cast<uint32_t>(value)};
          } else if (dtype == "i64") {
            return KernelCall::Parameter{static_cast<int64_t>(value)};
          } else if (dtype == "u64") {
            return KernelCall::Parameter{static_cast<uint64_t>(value)};
          } else {
            return absl::InvalidArgumentError(std::string("unknown dtype: ") +
                                              dtype.data());
          }
        }));

  m.def("create_scalar_parameter",
        ValueOrThrowWrapper([](double value, std::string_view dtype)
                                -> absl::StatusOr<KernelCall::Parameter> {
          if (dtype == "fp32") {
            return KernelCall::Parameter{static_cast<float>(value)};
          } else if (dtype == "fp64") {
            return KernelCall::Parameter{static_cast<double>(value)};
          } else {
            return absl::InvalidArgumentError(std::string("unknown dtype: ") +
                                              dtype.data());
          }
        }));

  nb::class_<KernelCall>(m, "TritonKernelCall")
      .def(nb::init<Kernel, uint32_t, uint32_t, uint32_t,
                    std::vector<KernelCall::Parameter>>())
      .def("to_proto", [](const KernelCall& kernel_call, std::string name,
                          nb::bytes metadata) {
        jax_triton::TritonAnyKernelCall proto;
        *proto.mutable_kernel_call() = kernel_call.ToProto();
        proto.set_name(std::move(name));
        proto.set_metadata(metadata.c_str(), metadata.size());
        std::string s = proto.SerializeAsString();
        return nb::bytes(s.c_str(), s.size());
      });

  nb::class_<AutotunedKernelCall>(m, "TritonAutotunedKernelCall")
      .def("__init__",
           [](AutotunedKernelCall* call, std::string name,
              std::vector<std::pair<KernelCall, std::string>>
                  calls_and_descriptions,
              std::vector<std::tuple<size_t, size_t, size_t>>
                  input_output_aliases) {
             std::vector<AutotunedKernelCall::Config> configs;
             configs.reserve(calls_and_descriptions.size());
             for (auto& [kernel_call, desc] : calls_and_descriptions) {
               configs.push_back({std::move(kernel_call), std::move(desc)});
             }
             new (call) AutotunedKernelCall(std::move(name), std::move(configs),
                                            std::move(input_output_aliases));
           })
      .def("to_proto", [](const AutotunedKernelCall& kernel_call,
                          std::string name, nb::bytes metadata) {
        jax_triton::TritonAnyKernelCall proto;
        *proto.mutable_autotuned_kernel_call() = kernel_call.ToProto();
        proto.set_name(std::move(name));
        proto.set_metadata(metadata.c_str(), metadata.size());
        std::string s = proto.SerializeAsString();
        return nb::bytes(s.c_str(), s.size());
      });

  m.def("get_custom_call",
        [] { return EncapsulateFunction(&TritonKernelCall); });

  m.def("get_compute_capability",
        ValueOrThrowWrapper([](int device) -> absl::StatusOr<int> {
          int major, minor;
          GPU_RETURN_IF_ERROR(gpuInit(device));
          GPU_RETURN_IF_ERROR(gpuDeviceGetAttribute(
              &major, GPU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
          GPU_RETURN_IF_ERROR(gpuDeviceGetAttribute(
              &minor, GPU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device));
          return major * 10 + minor;
        }));

  m.def("get_arch_details",
        ValueOrThrowWrapper([](int device) -> absl::StatusOr<std::string> {
#ifdef JAX_GPU_HIP
          hipDeviceProp_t prop;
          GPU_RETURN_IF_ERROR(hipGetDeviceProperties(&prop, device));
          return prop.gcnArchName;
#else
          return absl::UnimplementedError("Not a HIP GPU");
#endif
        }));

  m.def("get_serialized_metadata",
        ValueOrThrowWrapper([](nb::bytes opaque) -> absl::StatusOr<nb::bytes> {
          JAX_ASSIGN_OR_RETURN(
              std::string metadata,
              GetTritonKernelCallSerializedMetadata(
                  std::string_view(opaque.c_str(), opaque.size())));
          return nb::bytes(metadata.c_str(), metadata.size());
        }));

  m.def("registrations", &Registrations);
}

}  // namespace jax::JAX_GPU_NAMESPACE