Function TensorListFromListOfDLPackObjects

dali/python/backend_impl.cc:1282–1420

template <typename Backend>
std::shared_ptr<TensorList<Backend>> TensorListFromListOfDLPackObjects(
      py::list &list_of_objects,
      const std::optional<std::string> &layout,
      py::object stream,
      bool contiguous) {
  constexpr uint32_t rangeColor = TensorListDLPackRangeColor<Backend>();
  DomainTimeRange range(TensorListDLPackRangeName<Backend>(), rangeColor);

  if (list_of_objects.empty()) {
    auto ptr = std::make_shared<TensorList<Backend>>();
    if (layout.has_value()) {
      ptr->set_sample_dim(layout->length());
      ptr->SetLayout(*layout);
    }
    return ptr;
  }

  AccessOrder copy_order = AccessOrder::host();
  if constexpr (std::is_same_v<Backend, GPUBackend>) {
    if (!stream.is_none())
      copy_order = AccessOrderFromPythonStreamObj(stream);
  }

  // Preflight pass: validate device compatibility before consuming any capsule.
  // __dlpack__() is single-use, so mismatches must be caught before the handshake starts.
  int expected_device_id = -1;
  for (size_t i = 0; i < list_of_objects.size(); ++i) {
    py::object obj = list_of_objects[i];
    CheckDLPackDeviceBeforeCapsule<Backend>(obj, i, expected_device_id, copy_order);
  }

  if constexpr (std::is_same_v<Backend, GPUBackend>) {
    // If no __dlpack_device__ was found and no explicit stream was provided, fall back to the
    // current CUDA device so every __dlpack__() call receives a concrete consumer stream.
    if (copy_order == AccessOrder::host()) {
      int current_dev = -1;
      CUDA_CALL(cudaGetDevice(&current_dev));
      copy_order = AccessOrder(UserStream::Get()->GetStream(static_cast<size_t>(current_dev)));
    }
  }

  // Derive the DLPack consumer-stream handle from copy_order so the producer always
  // synchronizes with the exact same stream that DALI will use for the copy.
  py::object stream_handle = py::none();
  if constexpr (std::is_same_v<Backend, GPUBackend>) {
    if (copy_order.is_device()) {
      stream_handle = py::int_(reinterpret_cast<int64_t>(copy_order.stream()));
    }
  }

  std::optional<TensorList<Backend>> non_contiguous_tmp;
  std::shared_ptr<TensorList<Backend>> non_contiguous_out;

  if (contiguous)
    non_contiguous_tmp = TensorList<Backend>(list_of_objects.size());
  else
    non_contiguous_out = std::make_shared<TensorList<Backend>>(list_of_objects.size());