C API changes required for TF plugin. (#5898)

* Tensor and TensorList copy-out functions
* Query for Tensor and TensorList DType (including 0-sample tensor lists)
* Query for Tensor and TensorList byte size
* Utility macro to set optional field in param structures.

---------

Signed-off-by: Michal Zientkiewicz <[email protected]>

Author: mzient

dali/c_api_2/data_objects.cc

@@ -222,6 +222,26 @@ daliResult_t daliTensorGetShape(
   DALI_EPILOG();
 }
 
+daliResult_t daliTensorGetByteSize(
+      daliTensor_h tensor,
+      size_t *out_size) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor);
+  CHECK_OUTPUT(out_size);
+  *out_size = ptr->GetByteSize();
+  DALI_EPILOG();
+}
+
+daliResult_t daliTensorGetDType(
+      daliTensor_h tensor,
+      daliDataType_t *out_dtype) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor);
+  CHECK_OUTPUT(out_dtype);
+  *out_dtype = ptr->GetDType();
+  DALI_EPILOG();
+}
+
 daliResult_t daliTensorGetSourceInfo(
       daliTensor_h tensor,
       const char **out_source_info) {
@@ -391,6 +411,26 @@ daliResult_t daliTensorListGetShape(
   DALI_EPILOG();
 }
 
+daliResult_t daliTensorListGetByteSize(
+      daliTensorList_h tensor_list,
+      size_t *out_size) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor_list);
+  CHECK_OUTPUT(out_size);
+  *out_size = ptr->GetByteSize();
+  DALI_EPILOG();
+}
+
+daliResult_t daliTensorListGetDType(
+      daliTensorList_h tensor_list,
+      daliDataType_t *out_dtype) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor_list);
+  CHECK_OUTPUT(out_dtype);
+  *out_dtype = ptr->GetDType();
+  DALI_EPILOG();
+}
+
 daliResult_t daliTensorListGetTensorDesc(
       daliTensorList_h tensor_list,
       daliTensorDesc_t *out_tensor,
@@ -432,3 +472,16 @@ daliResult_t daliTensorListViewAsTensor(
   *out_tensor = t.release();  // no throwing allowed after this line
   DALI_EPILOG();
 }
+
+daliResult_t daliTensorListCopyOut(
+      daliTensorList_h tensor_list,
+      void *dst_buffer,
+      daliBufferPlacement_t dst_buffer_placement,
+      const cudaStream_t *stream,
+      daliCopyFlags_t flags) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor_list);
+  CHECK_OUTPUT(dst_buffer)
+  ptr->CopyOut(dst_buffer, dst_buffer_placement, ToOptional(stream), flags);
+  DALI_EPILOG();
+}

dali/c_api_2/data_objects.h

@@ -23,6 +23,7 @@
 #include <vector>
 #include "dali/dali.h"
 #include "dali/pipeline/data/tensor_list.h"
+#include "dali/pipeline/data/copy_to_external.h"
 #include "dali/c_api_2/ref_counting.h"
 #include "dali/c_api_2/validation.h"
 
@@ -43,6 +44,15 @@ struct _DALITensor {
 namespace dali {
 namespace c_api {
 
+constexpr mm::memory_kind_id GetMemoryKind(const daliBufferPlacement_t &placement) {
+  if (placement.device_type == DALI_STORAGE_GPU) {
+    return mm::memory_kind_id::device;
+  } else {
+    assert(placement.device_type == DALI_STORAGE_CPU);
+    return placement.pinned ? mm::memory_kind_id::pinned : mm::memory_kind_id::host;
+  }
+}
+
 //////////////////////////////////////////////////////////////////////////////
 // Interfaces
 //////////////////////////////////////////////////////////////////////////////
@@ -88,10 +98,21 @@ class ITensor : public _DALITensor, public RefCountedObject {
 
   virtual const TensorShape<> &GetShape() const & = 0;
 
+  virtual size_t GetByteSize() const = 0;
+
+  virtual daliDataType_t GetDType() const = 0;
+
   virtual const char *GetSourceInfo() const & = 0;
 
   virtual void SetSourceInfo(const char *source_info) = 0;
 
+  virtual void CopyOut(
+      void *dst_buffer,
+      daliBufferPlacement_t dst_buffer_placement,
+      std::optional<cudaStream_t> stream,
+      daliCopyFlags_t flags) = 0;
+
+
   /** Retrieves the underlying DALI Tensor<Backend> pointer.
    *
    * Returns a shared pointer to the underlying DALI object. If the backend doesn't match,
@@ -155,12 +176,22 @@ class ITensorList : public _DALITensorList, public RefCountedObject {
 
   virtual const TensorListShape<> &GetShape() const & = 0;
 
+  virtual size_t GetByteSize() const = 0;
+
+  virtual daliDataType_t GetDType() const = 0;
+
   virtual RefCountedPtr<ITensor> ViewAsTensor() const = 0;
 
   virtual const char *GetSourceInfo(int sample) const & = 0;
 
   virtual void SetSourceInfo(int sample, const char *source_info) = 0;
 
+  virtual void CopyOut(
+    void *dst_buffer,
+    daliBufferPlacement_t dst_buffer_placement,
+    std::optional<cudaStream_t> stream,
+    daliCopyFlags_t flags) = 0;
+
   /** Retrieves the underlying DALI TensorList<Backend> pointer.
    *
    * Returns a shared pointer to the underlying DALI object. If the backend doesn't match,
@@ -334,6 +365,14 @@ class TensorWrapper : public ITensor {
     return t_->shape();
   }
 
+  size_t GetByteSize() const override {
+    return t_->nbytes();
+  }
+
+  daliDataType_t GetDType() const override {
+    return t_->type();
+  }
+
   const char *GetSourceInfo() const & override {
     const char *info = t_->GetMeta().GetSourceInfo().c_str();
     if (info && !*info)
@@ -345,6 +384,19 @@ class TensorWrapper : public ITensor {
     t_->SetSourceInfo(source_info ? source_info : "");
   }
 
+  void CopyOut(
+      void *dst_buffer,
+      daliBufferPlacement_t dst_buffer_placement,
+      std::optional<cudaStream_t> stream,
+      daliCopyFlags_t flags) override {
+    Validate(dst_buffer_placement);
+    AccessOrder order = stream ? *stream : t_->order();
+    mm::memory_kind_id mem_kind = GetMemoryKind(dst_buffer_placement);
+    CopyToExternal(dst_buffer, mem_kind, *t_, order, flags & DALI_COPY_USE_KERNEL);
+    if (flags & DALI_COPY_SYNC)
+        AccessOrder::host().wait(order);
+  }
+
   const auto &NativePtr() const & {
     return t_;
   }
@@ -646,6 +698,14 @@ class TensorListWrapper : public ITensorList {
     return tl_->shape();
   }
 
+  size_t GetByteSize() const override {
+    return tl_->nbytes();
+  }
+
+  daliDataType_t GetDType() const override {
+    return tl_->type();
+  }
+
   const char *GetSourceInfo(int sample) const & override {
     ValidateSampleIdx(sample);
     const char *info = tl_->GetMeta(sample).GetSourceInfo().c_str();
@@ -684,6 +744,19 @@ class TensorListWrapper : public ITensorList {
     return Wrap(std::move(t));
   }
 
+  void CopyOut(
+      void *dst_buffer,
+      daliBufferPlacement_t dst_buffer_placement,
+      std::optional<cudaStream_t> stream,
+      daliCopyFlags_t flags) override {
+    Validate(dst_buffer_placement);
+    AccessOrder order = stream ? *stream : tl_->order();
+    mm::memory_kind_id mem_kind = GetMemoryKind(dst_buffer_placement);
+    CopyToExternal(dst_buffer, mem_kind, *tl_, order, flags & DALI_COPY_USE_KERNEL);
+    if (flags & DALI_COPY_SYNC)
+        AccessOrder::host().wait(order);
+  }
+
   const auto &NativePtr() const & {
     return tl_;
   }