Add source info functions + setter APIs.

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

Author: mzient

dali/c_api_2/data_objects.cc

@@ -147,6 +147,16 @@ daliResult_t daliTensorResize(
   DALI_EPILOG();
 }
 
+daliResult_t daliTensorGetBufferPlacement(
+      daliTensor_h tensor,
+      daliBufferPlacement_t *out_placement) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor);
+  CHECK_OUTPUT(out_placement);
+  *out_placement = ptr->GetBufferPlacement();
+  DALI_EPILOG();
+}
+
 daliResult_t daliTensorSetLayout(
       daliTensor_h tensor,
       const char *layout) {
@@ -212,6 +222,24 @@ daliResult_t daliTensorGetShape(
   DALI_EPILOG();
 }
 
+daliResult_t daliTensorGetSourceInfo(
+      daliTensor_h tensor,
+      const char **out_source_info) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor);
+  CHECK_OUTPUT(out_source_info);
+  *out_source_info = ptr->GetSourceInfo();
+  DALI_EPILOG();
+}
+
+daliResult_t daliTensorSetSourceInfo(
+      daliTensor_h tensor,
+      const char *source_info) {
+  DALI_PROLOG();
+  ToPointer(tensor)->SetSourceInfo(source_info);
+  DALI_EPILOG();
+}
+
 //////////////////////////////////////////////////////////////////////////////
 // TensorList
 //////////////////////////////////////////////////////////////////////////////
@@ -294,6 +322,17 @@ daliResult_t daliTensorListResize(
   DALI_EPILOG();
 }
 
+daliResult_t daliTensorListGetBufferPlacement(
+      daliTensorList_h tensor_list,
+      daliBufferPlacement_t *out_placement) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor_list);
+  CHECK_OUTPUT(out_placement);
+  *out_placement = ptr->GetBufferPlacement();
+  DALI_EPILOG();
+}
+
+
 daliResult_t daliTensorListSetLayout(
       daliTensorList_h tensor_list,
       const char *layout) {
@@ -363,6 +402,26 @@ daliResult_t daliTensorListGetTensorDesc(
   DALI_EPILOG();
 }
 
+daliResult_t daliTensorListGetSourceInfo(
+      daliTensorList_h tensor_list,
+      const char **out_source_info,
+      int sample_idx) {
+  DALI_PROLOG();
+  auto *ptr = ToPointer(tensor_list);
+  CHECK_OUTPUT(out_source_info);
+  *out_source_info = ptr->GetSourceInfo(sample_idx);
+  DALI_EPILOG();
+}
+
+daliResult_t daliTensorListSetSourceInfo(
+      daliTensorList_h tensor_list,
+      int sample_idx,
+      const char *source_info) {
+  DALI_PROLOG();
+  ToPointer(tensor_list)->SetSourceInfo(sample_idx, source_info);
+  DALI_EPILOG();
+}
+
 daliResult_t daliTensorListViewAsTensor(
       daliTensorList_h tensor_list,
       daliTensor_h *out_tensor) {

dali/c_api_2/data_objects.h

@@ -73,6 +73,10 @@ class ITensor : public _DALITensor, public RefCountedObject {
 
   virtual const TensorShape<> &GetShape() const & = 0;
 
+  virtual const char *GetSourceInfo() const & = 0;
+
+  virtual void SetSourceInfo(const char *source_info) = 0;
+
   template <typename Backend>
   const std::shared_ptr<Tensor<Backend>> &Unwrap() const &;
 
@@ -129,6 +133,10 @@ class ITensorList : public _DALITensorList, public RefCountedObject {
 
   virtual RefCountedPtr<ITensor> ViewAsTensor() const = 0;
 
+  virtual const char *GetSourceInfo(int sample) const & = 0;
+
+  virtual void SetSourceInfo(int sample, const char *source_info) = 0;
+
   template <typename Backend>
   const std::shared_ptr<TensorList<Backend>> &Unwrap() const &;
 
@@ -297,6 +305,17 @@ class TensorWrapper : public ITensor {
     return t_->shape();
   }
 
+  const char *GetSourceInfo() const & override {
+    const char *info = t_->GetMeta().GetSourceInfo().c_str();
+    if (info && !*info)
+      return nullptr;
+    return info;
+  }
+
+  void SetSourceInfo(const char *source_info) override {
+    t_->SetSourceInfo(source_info ? source_info : "");
+  }
+
   const auto &NativePtr() const & {
     return t_;
   }
@@ -583,11 +602,9 @@ class TensorListWrapper : public ITensorList {
   }
 
   daliTensorDesc_t GetTensorDesc(int sample) const override {
-    auto &shape = tl_->shape();
-    if (sample < 0 || sample >= shape.num_samples())
-      throw std::out_of_range(make_string("The sample index ", sample, " is out of range. "
-        "Valid indices are [0..", shape.num_samples() - 1, "]."));
+    ValidateSampleIdx(sample);
     daliTensorDesc_t desc{};
+    auto &shape = tl_->shape();
     desc.ndim = shape.sample_dim();
     desc.data = tl_->raw_mutable_tensor(sample);
     desc.dtype = tl_->type();
@@ -600,6 +617,19 @@ class TensorListWrapper : public ITensorList {
     return tl_->shape();
   }
 
+  const char *GetSourceInfo(int sample) const & override {
+    ValidateSampleIdx(sample);
+    const char *info = tl_->GetMeta(sample).GetSourceInfo().c_str();
+    if (info && !*info)
+      return nullptr;  // return empty string as NULL
+    return info;
+  }
+
+  void SetSourceInfo(int sample, const char *source_info) override {
+    ValidateSampleIdx(sample);
+    tl_->SetSourceInfo(sample, source_info ? source_info : "");
+  }
+
   RefCountedPtr<ITensor> ViewAsTensor() const override {
     if (!tl_->IsContiguous())
       throw std::runtime_error(
@@ -629,6 +659,12 @@ class TensorListWrapper : public ITensorList {
     return tl_;
   }
 
+  inline void ValidateSampleIdx(int idx) const {
+    if (idx < 0 || idx >= tl_->num_samples())
+    throw std::out_of_range(make_string("The sample index ", idx, " is out of range. "
+      "Valid indices are [0..", tl_->num_samples() - 1, "]."));
+  }
+
  private:
   std::shared_ptr<TensorList<Backend>> tl_;
 };

dali/c_api_2/data_objects_test.cc

@@ -16,6 +16,7 @@
 #include <gtest/gtest.h>
 #include "dali/c_api_2/managed_handle.h"
 #include "dali/core/span.h"
+#include "dali/core/device_guard.h"
 
 TEST(CAPI2_TensorListTest, NullHandle) {
   daliTensorList_h h = nullptr;
@@ -65,6 +66,7 @@ inline auto CreateTensorList(daliBufferPlacement_t placement) {
 void TestTensorListResize(daliStorageDevice_t storage_device) {
   daliBufferPlacement_t placement{};
   placement.device_type = storage_device;
+  placement.pinned = true;
   int64_t shapes[] = {
     480, 640, 3,
     600, 800, 4,
@@ -74,6 +76,13 @@ void TestTensorListResize(daliStorageDevice_t storage_device) {
   daliDataType_t dtype = DALI_UINT32;
 
   auto tl = CreateTensorList(placement);
+
+  daliBufferPlacement_t test_placement{};
+  EXPECT_EQ(daliTensorListGetBufferPlacement(tl, &test_placement), DALI_SUCCESS);
+  EXPECT_EQ(test_placement.device_type, placement.device_type);
+  EXPECT_EQ(test_placement.device_id, placement.device_id);
+  EXPECT_EQ(test_placement.pinned, placement.pinned);
+
   EXPECT_EQ(daliTensorListResize(tl, 4, 3, nullptr, dtype, nullptr), DALI_ERROR_INVALID_ARGUMENT);
   EXPECT_EQ(daliTensorListResize(tl, -1, 3, shapes, dtype, nullptr), DALI_ERROR_INVALID_ARGUMENT);
   EXPECT_EQ(daliTensorListResize(tl, 4, -1, shapes, dtype, nullptr), DALI_ERROR_INVALID_ARGUMENT);
@@ -289,8 +298,14 @@ TEST(CAPI2_TensorListTest, AttachSamples) {
 
 
 TEST(CAPI2_TensorListTest, ViewAsTensor) {
+  int num_dev = 0;
+  CUDA_CALL(cudaGetDeviceCount(&num_dev));
+  // use the last device
+  dali::DeviceGuard dg(num_dev - 1);
+
   daliBufferPlacement_t placement{};
   placement.device_type = DALI_STORAGE_CPU;
+  placement.pinned = true;
   using element_t = int;
   daliDataType_t dtype = dali::type2id<element_t>::value;
   dali::TensorListShape<> lshape = dali::uniform_list_shape(4, { 480, 640, 3 });
@@ -327,6 +342,12 @@ TEST(CAPI2_TensorListTest, ViewAsTensor) {
   ASSERT_NE(ht, nullptr);
   dali::c_api::TensorHandle t(ht);
 
+  daliBufferPlacement_t tensor_placement{};
+  EXPECT_EQ(daliTensorGetBufferPlacement(ht, &tensor_placement), DALI_SUCCESS);
+  EXPECT_EQ(tensor_placement.device_type, placement.device_type);
+  EXPECT_EQ(tensor_placement.device_id, placement.device_id);
+  EXPECT_EQ(tensor_placement.pinned, placement.pinned);
+
   daliTensorDesc_t desc{};
   EXPECT_EQ(daliTensorGetDesc(t, &desc), DALI_SUCCESS) << daliGetLastErrorMessage();
   EXPECT_EQ(desc.data, data.get());
@@ -507,3 +528,38 @@ TEST(CAPI2_TensorTest, ResizeCPU) {
 TEST(CAPI2_TensorTest, ResizeGPU) {
   TestTensorResize(DALI_STORAGE_GPU);
 }
+
+TEST(CAPI2_TensorTest, SourceInfo) {
+  auto t = CreateTensor({});
+  const char *out_src_info = "junk";
+  EXPECT_EQ(daliTensorGetSourceInfo(t, &out_src_info), DALI_SUCCESS);
+  EXPECT_EQ(out_src_info, nullptr);
+
+  EXPECT_EQ(daliTensorSetSourceInfo(t, "source_info"), DALI_SUCCESS);
+  EXPECT_EQ(daliTensorGetSourceInfo(t, &out_src_info), DALI_SUCCESS);
+  EXPECT_STREQ(out_src_info, "source_info");
+}
+
+TEST(CAPI2_TensorListTest, SourceInfo) {
+  auto t = CreateTensorList({});
+  ASSERT_EQ(daliTensorListResize(t, 5, 0, nullptr, DALI_UINT8, nullptr), DALI_SUCCESS);
+
+  const char *out_src_info = "junk";
+  EXPECT_EQ(daliTensorListGetSourceInfo(t, &out_src_info, 0), DALI_SUCCESS);
+  EXPECT_EQ(out_src_info, nullptr);
+
+  EXPECT_EQ(daliTensorListSetSourceInfo(t, 0, "quick"), DALI_SUCCESS);
+  EXPECT_EQ(daliTensorListSetSourceInfo(t, 2, "brown"), DALI_SUCCESS);
+  EXPECT_EQ(daliTensorListSetSourceInfo(t, 4, "fox"), DALI_SUCCESS);
+
+  EXPECT_EQ(daliTensorListGetSourceInfo(t, &out_src_info, 0), DALI_SUCCESS);
+  EXPECT_STREQ(out_src_info, "quick");
+  EXPECT_EQ(daliTensorListGetSourceInfo(t, &out_src_info, 1), DALI_SUCCESS);
+  EXPECT_EQ(out_src_info, nullptr);
+  EXPECT_EQ(daliTensorListGetSourceInfo(t, &out_src_info, 2), DALI_SUCCESS);
+  EXPECT_STREQ(out_src_info, "brown");
+  EXPECT_EQ(daliTensorListGetSourceInfo(t, &out_src_info, 3), DALI_SUCCESS);
+  EXPECT_EQ(out_src_info, nullptr);
+  EXPECT_EQ(daliTensorListGetSourceInfo(t, &out_src_info, 4), DALI_SUCCESS);
+  EXPECT_STREQ(out_src_info, "fox");
+}

include/dali/dali.h

@@ -874,6 +874,22 @@ DALI_API daliResult_t daliTensorListGetSourceInfo(
   const char **out_source_info,
   int sample_idx);
 
+/** Sets the "source info" metadata of a tensor in a list.
+ *
+ * A tensor can be associated with a "source info" string, which typically is the file name,
+ * but can also contain an index in a container, key, etc.
+ *
+ * @param tensor_list [in]  The tensor list
+ * @param sample_idx  [in]  The index of the sample, whose source info will is being set.
+ * @param source_info [in]  A source info string (e.g. filename) to associate with the tensor.
+ *                          Passing NULL is equivalent to passing an empty string.
+ */
+DALI_API daliResult_t daliTensorListSetSourceInfo(
+  daliTensorList_h tensor_list,
+  int sample_idx,
+  const char *source_info);
+
+
 /** Gets the tensor descriptor of the specified sample.
  *
  * @param tensor_list [in]  The tensor list
@@ -1107,6 +1123,19 @@ DALI_API daliResult_t daliTensorGetSourceInfo(
   daliTensor_h tensor,
   const char **out_source_info);
 
+/** Sets the "source info" metadata of a tensor.
+ *
+ * A tensor can be associated with a "source info" string, which typically is the file name,
+ * but can also contain an index in a container, key, etc.
+ *
+ * @param tensor      [in]  The tensor
+ * @param source_info [in]  A source info string (e.g. filename) to associate with the tensor.
+ *                          Passing NULL is equivalent to passing an empty string.
+ */
+DALI_API daliResult_t daliTensorSetSourceInfo(
+  daliTensor_h tensor,
+  const char *source_info);
+
 /** Gets the descriptor of the data in the tensor.
  *
  * @param tensor      [in]  The tensor