Function CheckDLPackDeviceBeforeCapsule

dali/python/backend_impl.cc:1226–1269

template <typename Backend>
void CheckDLPackDeviceBeforeCapsule(py::object obj, size_t idx,
                                    int &expected_device_id,
                                    AccessOrder &copy_order) {
  if (!py::hasattr(obj, "__dlpack__")) {
    throw py::type_error(make_string(
        "Object at position ", idx, " does not support the DLPack protocol."));
  }
  if (!py::hasattr(obj, "__dlpack_device__")) {
    return;
  }

  py::tuple dev_info = obj.attr("__dlpack_device__")();
  int dev_type = dev_info[0].cast<int>();
  if constexpr (std::is_same_v<Backend, CPUBackend>) {
    (void)expected_device_id;
    (void)copy_order;
    if (dev_type != kDLCPU && dev_type != kDLCUDAHost) {
      throw py::value_error(make_string(
          "All tensors must reside in CPU memory. "
          "Tensor at position ", idx, " has DLPack device type ", dev_type, "."));
    }
  } else {
    if (dev_type != kDLCUDA) {
      throw py::value_error(make_string(
          "All tensors must reside in GPU memory. "
          "Tensor at position ", idx, " has DLPack device type ", dev_type, "."));
    }
    int dev_id = dev_info[1].cast<int>();
    if (expected_device_id == -1) {
      expected_device_id = dev_id;
      // When no explicit stream was given, look up the UserStream for this device
      // so every __dlpack__() call below uses the same concrete stream handle.
      if (copy_order == AccessOrder::host()) {
        copy_order = AccessOrder(UserStream::Get()->GetStream(
            static_cast<size_t>(expected_device_id)));
      }
    } else if (dev_id != expected_device_id) {
      throw py::value_error(make_string(
          "All tensors must reside on the same GPU device. "
          "Tensor at position ", idx, " is on GPU ", dev_id,
          " but expected GPU ", expected_device_id, "."));
    }
  }
}

template <typename Backend>
py::capsule GetDLPackCapsule(py::object obj, py::object stream_handle) {