/* Copyright 2020 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 "jaxlib/dlpack.h"

#include <Python.h>

#include <cstdint>
#include <functional>
#include <memory>
#include <numeric>
#include <optional>
#include <string_view>
#include <tuple>
#include <utility>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_join.h"
#include "absl/types/span.h"
#include "include/dlpack/dlpack.h"
#include "nanobind/nanobind.h"
#include "nanobind/ndarray.h"
#include "jaxlib/nb_class_ptr.h"
#include "jaxlib/py_array.h"
#include "jaxlib/py_client.h"
#include "jaxlib/py_user_context.h"
#include "jaxlib/python_ref_manager.h"
#include "jaxlib/util.h"
#include "xla/layout.h"
#include "xla/pjrt/exceptions.h"
#include "xla/pjrt/host_memory_spaces.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/pjrt_common.h"
#include "xla/pjrt/pjrt_compiler.h"
#include "xla/python/dlpack_types.h"
#include "xla/python/ifrt/array.h"
#include "xla/python/ifrt/device.h"
#include "xla/python/ifrt/rtti.h"
#include "xla/python/pjrt_ifrt/pjrt_array.h"
#include "xla/python/pjrt_ifrt/pjrt_client.h"
#include "xla/python/pjrt_ifrt/pjrt_device.h"
#include "xla/python/types.h"
#include "xla/python/version.h"
#include "xla/shape_util.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/util.h"
#include "xla/xla_data.pb.h"

namespace ifrt = xla::ifrt;
namespace nb = nanobind;

namespace jax {
namespace {

const char* const kDlTensorCapsuleName = "dltensor";
const char* const kDlTensorVersionedCapsuleName = "dltensor_versioned";

template <typename ManagedTensor>
struct DLPackTensor {
  ~DLPackTensor();

  // `buffer_reference` is populated if we have shared (read-only) access.
  nb::object buffer_reference;

  // `external_reference` is always populated.
  std::unique_ptr<xla::PjRtBuffer::ExternalReference> external_reference;

  std::vector<int64_t> shape;
  std::vector<int64_t> strides;
  ManagedTensor tensor;
};

template <typename ManagedTensor>
DLPackTensor<ManagedTensor>::~DLPackTensor() {
  // We must release the external reference first before deleting the array.
  external_reference.reset();
  if (buffer_reference) {
    GlobalPyRefManager()->AddGarbage(
        absl::MakeSpan(&buffer_reference, /*size=*/1));
  }
}

template <typename ManagedTensor>
void DLPackTensorDeleter(ManagedTensor* t) {
  if (t) {
    delete static_cast<DLPackTensor<ManagedTensor>*>(t->manager_ctx);
  }
}

absl::StatusOr<std::vector<int64_t>> StridesToLayout(
    absl::Span<int64_t const> dims, absl::Span<int64_t const> strides) {
  CHECK_EQ(dims.size(), strides.size());
  std::vector<int64_t> minor_to_major(dims.size());
  std::iota(minor_to_major.begin(), minor_to_major.end(), 0);
  absl::c_sort(minor_to_major, [&](int a, int b) {
    if (strides[a] < strides[b]) {
      return true;
    }
    if (strides[a] > strides[b]) {
      return false;
    }
    // If two dimensions have the same stride, prefer the major-to-minor
    // interpretation of the ordering, since that's what JAX wants.
    return b < a;
  });
  int64_t stride = 1;
  for (int64_t d : minor_to_major) {
    if (dims[d] > 1 && strides[d] != stride) {
      return xla::Unimplemented(
          "Only DLPack tensors with trivial (compact) striding are supported; "
          "i.e., tensors whose striding represents a transposition of the "
          "underlying buffer but not broadcasting. Dimensions were: [%s], "
          "strides were [%s].",
          absl::StrJoin(dims, ","), absl::StrJoin(strides, ","));
    }
    stride *= dims[d];
  }
  return minor_to_major;
}

absl::StatusOr<DLDeviceType> DLDeviceTypeForDevice(
    const xla::PjRtDevice& device) {
  if (device.client()->platform_id() == xla::CpuId()) {
    return kDLCPU;
  } else if (device.client()->platform_id() == xla::CudaId()) {
    return kDLCUDA;
  } else if (device.client()->platform_id() == xla::RocmId()) {
    return kDLROCM;
  } else if (device.client()->platform_id() == xla::OneapiId()) {
    return kDLOneAPI;
  }
  return xla::InvalidArgument("Device %s cannot be used as a DLPack device.",
                              device.DebugString());
}

absl::StatusOr<DLDevice> DLDeviceForBuffer(const xla::PjRtBuffer& buffer) {
  DLDevice context;
  const xla::PjRtMemorySpace* memory_space = buffer.memory_space();
  if (memory_space != nullptr &&
      memory_space->kind() == xla::PinnedHostMemorySpace::kKind &&
      buffer.device()->client()->platform_id() == xla::CudaId()) {
    context.device_type = kDLCUDAHost;
  } else if (memory_space != nullptr &&
             memory_space->kind() == xla::PinnedHostMemorySpace::kKind &&
             buffer.device()->client()->platform_id() == xla::TpuId()) {
    context.device_type = kDLTPUHost;
  } else if (memory_space != nullptr &&
             memory_space->kind() == xla::PinnedHostMemorySpace::kKind &&
             buffer.device()->client()->platform_id() == xla::RocmId()) {
    context.device_type = kDLROCMHost;
  } else {
    ABSL_ASSIGN_OR_RETURN(context.device_type,
                          DLDeviceTypeForDevice(*buffer.device()));
  }
  context.device_id = buffer.device()->local_hardware_id().value();
  return context;
}

absl::Status VerifyDType(const DLTensor& dl_tensor) {
  if (dl_tensor.dtype.bits % 8 != 0) {
    return xla::InvalidArgument(
        "Unsupported DLPack tensor dtype: bits should be a multiple of 8, got "
        "%d",
        dl_tensor.dtype.bits);
  }

  if (dl_tensor.dtype.lanes != 1) {
    return xla::InvalidArgument(
        "Unsupported DLPack tensor dtype: lanes should be equal to 1, got %d",
        dl_tensor.dtype.lanes);
  }

  return absl::OkStatus();
}

absl::StatusOr<std::vector<int64_t>> GetByteStrides(const DLTensor& dl_tensor) {
  ABSL_RETURN_IF_ERROR(VerifyDType(dl_tensor));

  // Convert element strides from the number of elements to the number of bytes.
  std::vector<int64_t> strides;
  strides.reserve(dl_tensor.ndim);
  for (int i = 0; i < dl_tensor.ndim; ++i) {
    strides.push_back(dl_tensor.strides[i] * dl_tensor.dtype.bits / 8);
  }
  return strides;
}

// Makes a PjRtBuffer from a DLPack tensor. Returns a pair where the second
// element is true if a copy actually happened.
absl::StatusOr<std::pair<std::unique_ptr<xla::PjRtBuffer>, bool>>
MakePjrtBuffer(xla::PjRtDevice& device, const DLTensor& dl_tensor,
               std::function<void()> on_delete_callback,
               const xla::Shape& shape, xla::PrimitiveType element_type,
               absl::Span<int64_t const> dimensions,
               std::optional<bool> copy = std::nullopt,
               std::optional<std::intptr_t> stream = std::nullopt,
               std::optional<DLDeviceType> dl_device_type = std::nullopt) {
  void* data = static_cast<char*>(dl_tensor.data) + dl_tensor.byte_offset;

  // DLPack producers advertise the tensor device type via __dlpack_device__()
  // but the corresponding __dlpack__() call might return a capsule with a
  // different device type. E.g. pinned PyTorch tensors on CUDA advertise
  // kDLCUDAHost but store kDLCPU in the capsule. dl_device_type allows us to
  // explicitly override the capsule's device type and pick the correct
  // destination memory space, i.e. pinned_host in the above example.
  DLDeviceType effective_device_type =
      dl_device_type.value_or(dl_tensor.device.device_type);
  xla::PjRtMemorySpace* memory_space;
  if (effective_device_type == kDLCUDAHost ||
      effective_device_type == kDLROCMHost ||
      effective_device_type == kDLTPUHost) {
    ABSL_ASSIGN_OR_RETURN(memory_space, device.memory_space_by_kind(
                                            xla::PinnedHostMemorySpace::kKind));
  } else {
    ABSL_ASSIGN_OR_RETURN(memory_space, device.default_memory_space());
  }

  // On CPU, creating a view may fail because of unaligned data buffer
  // in which case we'll fallback to copy. On non-CPU, array-api copy
  // semantics is handled in dlpack._place_array function.
  bool fallback_to_copy =
      !copy.has_value() &&
      (effective_device_type == kDLCPU ||
       effective_device_type == kDLCUDAHost ||
       effective_device_type == kDLROCMHost);

  // Create a view.
  if (!copy.value_or(false)) {
    auto result = device.client()->CreateViewOfDeviceBuffer(
        data, shape, memory_space, on_delete_callback, stream);
    if (!(result.status().code() == absl::StatusCode::kInvalidArgument &&
          fallback_to_copy)) {
      ABSL_RETURN_IF_ERROR(result.status());
      return std::make_pair(*std::move(result), false);
    }
  }

  // Convert tensor strides (expressed in number of elements) to byte strides.
  std::optional<std::vector<int64_t>> byte_strides;
  if (dl_tensor.strides) {
    ABSL_ASSIGN_OR_RETURN(byte_strides, GetByteStrides(dl_tensor));
  }

  // Create a copy.
  ABSL_ASSIGN_OR_RETURN(
      auto buffer,
      device.client()->BufferFromHostBuffer(
          data, element_type, dimensions, byte_strides,
          xla::PjRtClient::HostBufferSemantics::kMutableZeroCopy,
          on_delete_callback, memory_space, /*device_layout=*/nullptr));
  return std::make_pair(std::move(buffer), true);
}

template <typename ManagedTensor>
absl::StatusOr<std::unique_ptr<DLPackTensor<ManagedTensor>>> PackDLPackTensor(
    nb::handle py_buffer, std::optional<std::intptr_t> stream) {
  xla::ifrt::ArrayRef ifrt_array =
      nb::cast<PyArray>(py_buffer).ifrt_array_ref();
  if (ifrt_array == nullptr) {
    return xla::Unimplemented(
        "BufferToDLPackManagedTensor called on deleted array.");
  }
  auto* arr =
      xla::ifrt::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(ifrt_array.get());
  if (arr == nullptr) {
    throw xla::XlaRuntimeError(
        "This operation is implemented for a PjRt-compatible backend only.");
  }
  xla::PjRtBuffer* pjrt_buffer = arr->pjrt_buffers().front().get();

  if (pjrt_buffer->IsTuple()) {
    return xla::Unimplemented(
        "BufferToDLPackManagedTensor is not implemented for tuple "
        "buffers.");
  }
  if (pjrt_buffer->has_dynamic_dimensions()) {
    return xla::Unimplemented("DynamicShape is not implemented in DLPack.");
  }

  auto pack = std::make_unique<DLPackTensor<ManagedTensor>>();
  DLTensor& dt = pack->tensor.dl_tensor;
  {
    // AcquireExternalReference may block; there are no API guarantees.
    GlobalPyRefManager()->CollectGarbage();
    nb::gil_scoped_release gil_release;
    ABSL_ASSIGN_OR_RETURN(pack->external_reference,
                          pjrt_buffer->AcquireExternalReference());
    if (stream) {
      ABSL_RETURN_IF_ERROR(
          pack->external_reference->WaitUntilBufferReadyOnStream(*stream));
    } else {
      xla::ifrt::Array* ifrt_array_ptr = ifrt_array.get();
      ABSL_RETURN_IF_ERROR(
          AwaitBuffersReady(absl::MakeConstSpan(&ifrt_array_ptr, 1)));
    }
  }
  pack->buffer_reference = nb::borrow<nb::object>(py_buffer);

  dt.data = pack->external_reference->OpaqueDeviceMemoryDataPointer();
  pack->tensor.manager_ctx = pack.get();
  pack->tensor.deleter = DLPackTensorDeleter<ManagedTensor>;
  ABSL_ASSIGN_OR_RETURN(dt.device, DLDeviceForBuffer(*pjrt_buffer));
  dt.ndim = pjrt_buffer->dimensions().size();
  ABSL_ASSIGN_OR_RETURN(dt.dtype,
                        PrimitiveTypeToDLDataType(pjrt_buffer->element_type()));

  pack->shape = std::vector<int64_t>(pjrt_buffer->dimensions().begin(),
                                     pjrt_buffer->dimensions().end());

  // TODO(b/327524065): use PjRtLayout directly instead of xla::Layout
  xla::Layout xla_layout = pjrt_buffer->layout()->xla_layout();
  pack->strides = StridesForShape(pjrt_buffer->element_type(),
                                  pjrt_buffer->dimensions(), xla_layout);

  dt.shape = reinterpret_cast<std::int64_t*>(pack->shape.data());
  dt.strides = reinterpret_cast<std::int64_t*>(pack->strides.data());
  dt.byte_offset = 0;
  return pack;
}

}  // namespace

absl::StatusOr<nb::capsule> BufferToDLPackManagedTensor(
    nb::handle py_buffer, std::optional<std::intptr_t> stream,
    std::optional<std::tuple<int64_t, int64_t>> max_version, bool copied) {
  if (max_version && std::get<0>(*max_version) >= 1) {
    ABSL_ASSIGN_OR_RETURN(
        auto pack,
        PackDLPackTensor<DLManagedTensorVersioned>(py_buffer, stream));
    pack->tensor.version =
        DLPackVersion{DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION};
    pack->tensor.flags = DLPACK_FLAG_BITMASK_READ_ONLY;
    if (copied) {
      pack->tensor.flags |= DLPACK_FLAG_BITMASK_IS_COPIED;
    }
    nb::capsule capsule = nb::steal<nb::capsule>(PyCapsule_New(
        &pack.release()->tensor, kDlTensorVersionedCapsuleName,
        [](PyObject* obj) noexcept {
          PyObject* exc = PyErr_GetRaisedException();
          DLManagedTensorVersioned* dlmt =
              static_cast<DLManagedTensorVersioned*>(
                  PyCapsule_GetPointer(obj, kDlTensorVersionedCapsuleName));
          if (dlmt) {
            DLPackTensorDeleter(dlmt);
          }
          PyErr_SetRaisedException(exc);
        }));
    if (!capsule.ptr()) {
      throw nb::python_error();
    }
    return capsule;
  }

  ABSL_ASSIGN_OR_RETURN(auto pack,
                        PackDLPackTensor<DLManagedTensor>(py_buffer, stream));
  // We cannot use nanobind's capsule object constructor because we need to
  // detect if the capsule name has been changed in the deleter, but nanobind
  // hides the underlying Python object from the deleter.
  nb::capsule capsule = nb::steal<nb::capsule>(
      PyCapsule_New(&pack.release()->tensor, kDlTensorCapsuleName,
                    [](PyObject* obj) noexcept {
                      PyObject* exc = PyErr_GetRaisedException();
                      DLManagedTensor* dlmt = static_cast<DLManagedTensor*>(
                          PyCapsule_GetPointer(obj, kDlTensorCapsuleName));
                      if (dlmt) {
                        DLPackTensorDeleter(dlmt);
                      }
                      // PyCapsule_GetPointer may have raised. Restore the
                      // previous exception if there was one.
                      PyErr_SetRaisedException(exc);
                    }));
  if (!capsule.ptr()) {
    throw nb::python_error();
  }
  return capsule;
}

absl::StatusOr<nb::object> DLPackManagedTensorToBuffer(
    const nb::capsule& tensor, ifrt::Device* ifrt_device,
    nb_class_ptr<PyClient> client, std::optional<std::intptr_t> stream,
    std::optional<bool> copy, std::optional<DLDeviceType> dl_device_type) {
  ifrt::PjRtDevice* device =
      xla::ifrt::dyn_cast_or_null<ifrt::PjRtDevice>(ifrt_device);
  if (device == nullptr) {
    throw xla::XlaRuntimeError(
        "DLPack is supported for PjRt-compatible backends only.");
  }
  if (!device->IsAddressable()) {
    throw xla::XlaRuntimeError(
        "DLPack is only supported for devices addressable by the current "
        "process.");
  }
  std::string_view capsule_name(tensor.name());
  const DLTensor* dl_tensor = nullptr;
  std::function<void()> on_delete_callback;
  const char* used_capsule_name = nullptr;
  if (capsule_name == kDlTensorVersionedCapsuleName) {
    DLManagedTensorVersioned* dlmt =
        static_cast<DLManagedTensorVersioned*>(tensor.data());
    if (dlmt->version.major > DLPACK_MAJOR_VERSION) {
      return xla::InvalidArgument(
          "Unsupported DLPack version: major version %u > %u",
          dlmt->version.major, DLPACK_MAJOR_VERSION);
    }
    dl_tensor = &dlmt->dl_tensor;
    if (dlmt->deleter) {
      on_delete_callback = [dlmt]() { dlmt->deleter(dlmt); };
    }
    used_capsule_name = "used_dltensor_versioned";
  } else if (capsule_name == kDlTensorCapsuleName) {
    DLManagedTensor* dlmt = static_cast<DLManagedTensor*>(tensor.data());
    dl_tensor = &dlmt->dl_tensor;
    if (dlmt->deleter) {
      on_delete_callback = [dlmt]() { dlmt->deleter(dlmt); };
    }
    used_capsule_name = "used_dltensor";
  } else {
    return xla::InvalidArgument(
        "DLPack tensor must be a capsule with name \"dltensor\" or "
        "\"dltensor_versioned\", got \"%s\". "
        "Note that a DLPack tensor may be consumed at most once.",
        capsule_name);
  }
  if (dl_tensor->ndim < 0) {
    return xla::InvalidArgument(
        "Number of dimensions in DLManagedTensor must be nonnegative, got %d",
        dl_tensor->ndim);
  }
  absl::Span<int64_t const> dimensions(
      reinterpret_cast<int64_t*>(dl_tensor->shape), dl_tensor->ndim);
  ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType element_type,
                        xla::DLDataTypeToPrimitiveType(dl_tensor->dtype));

  bool has_custom_layout = dl_tensor->strides != nullptr;
  std::vector<int64_t> minor_to_major;
  if (dl_tensor->strides &&
      absl::c_find(dimensions, 0) == dimensions.end()) {
    absl::Span<int64_t const> strides(
        reinterpret_cast<int64_t*>(dl_tensor->strides),
        dl_tensor->ndim);
    ABSL_ASSIGN_OR_RETURN(minor_to_major, StridesToLayout(dimensions, strides));
  } else {
    minor_to_major.resize(dl_tensor->ndim);
    std::iota(minor_to_major.rbegin(), minor_to_major.rend(), 0);
  }
  xla::Shape shape = xla::ShapeUtil::MakeShapeWithDenseLayout(
      element_type, dimensions, minor_to_major);

  ABSL_ASSIGN_OR_RETURN(
      auto pjrt_buffer_and_copied,
      MakePjrtBuffer(*device->pjrt_device(), *dl_tensor,
                     std::move(on_delete_callback), shape, element_type,
                     dimensions, copy, stream, dl_device_type));
  if (pjrt_buffer_and_copied.second) {
    // A PjRtBuffer uses a default layout if it has been created using copy.
    has_custom_layout = false;
  }

  // We have taken ownership of the array inside the capsule; make sure the
  // capsule it cannot be used again.
  PyCapsule_SetName(tensor.ptr(), used_capsule_name);
  PyCapsule_SetDestructor(tensor.ptr(), nullptr);

  auto* ifrt_client = xla::ifrt::dyn_cast_or_null<ifrt::PjRtCompatibleClient>(
      client->ifrt_client());
  if (ifrt_client == nullptr) {
    throw xla::XlaRuntimeError(
        "This operation is implemented for a PjRt-compatible backend only.");
  }
  PyUserContextScope user_context_scope;
  ABSL_ASSIGN_OR_RETURN(
      auto ifrt_array,
      ifrt_client->CreatePjRtArray(std::move(pjrt_buffer_and_copied.first),
                                   has_custom_layout));
  return PyArray::MakeFromSingleDeviceArray(std::move(client),
                                            std::move(ifrt_array), false, true);
}

absl::StatusOr<nanobind::dlpack::dtype> PrimitiveTypeToNbDLDataType(
    xla::PrimitiveType type) {
  ABSL_ASSIGN_OR_RETURN(DLDataType dl_type, PrimitiveTypeToDLDataType(type));

  nanobind::dlpack::dtype nb_type;
  nb_type.lanes = dl_type.lanes;
  nb_type.bits = dl_type.bits;
  nb_type.code = dl_type.code;

  return nb_type;
}

}  // namespace jax