Refactor get_property. Add and functions to DataNode.

Signed-off-by: Michał Zientkiewicz <mzient@gmail.com>

dali/operators/util/get_property.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@ DALI_SCHEMA(GetProperty)
 
 The type of the output will depend on the ``key`` of the requested property.)code")
     .NumInput(1)
+    .InputDevice(0, InputDevice::Metadata)
     .NumOutput(1)
     .AddArg("key",
             R"code(Specifies, which property is requested.
@@ -38,6 +39,110 @@ The following properties are supported:
 )code",
             DALI_STRING);
 
+template <typename Backend, typename SampleShapeFunc, typename CopySampleFunc>
+void GetPerSample(TensorList<CPUBackend> &out, const TensorList<Backend> &in,
+                  SampleShapeFunc &&sample_shape, CopySampleFunc &&copy_sample) {
+  int N = in.num_samples();
+  TensorListShape<> tls;
+  for (int i = 0; i < N; i++) {
+    auto shape = sample_shape(in, i);
+    if (i == 0)
+      tls.resize(N, shape.sample_dim());
+    tls.set_tensor_shape(i, shape);
+  }
+  out.Resize(tls, DALI_UINT8);
+  for (int i = 0; i < N; i++) {
+    copy_sample(out, in, i);
+  }
+}
+
+template <typename Backend>
+void SourceInfoToTL(TensorList<CPUBackend> &out, const TensorList<Backend> &in) {
+  GetPerSample(out, in,
+    [](auto &in, int idx) {
+      auto &info = in.GetMeta(idx).GetSourceInfo();
+      return TensorShape<1>(info.length());
+    },
+    [](auto &out, auto &in, int idx) {
+      auto &info = in.GetMeta(idx).GetSourceInfo();
+      std::memcpy(out.raw_mutable_tensor(idx), info.c_str(), info.length());
+    });
+}
+
+template <typename Backend>
+void SourceInfoToTL(TensorList<GPUBackend> &out, const TensorList<Backend> &in) {
+  TensorList<CPUBackend> tmp;
+  tmp.set_pinned(true);
+  SourceInfoToTL(tmp, in);
+  tmp.set_order(out.order());
+  out.Copy(tmp);
+}
+
+template <typename OutputBackend>
+void SourceInfoToTL(TensorList<OutputBackend> &out, const Workspace &ws) {
+  ws.Output<OutputBackend>(0).set_order(ws.output_order());
+  if (ws.InputIsType<CPUBackend>(0))
+    return SourceInfoToTL(out, ws.Input<CPUBackend>(0));
+  else if (ws.InputIsType<GPUBackend>(0))
+    return SourceInfoToTL(out, ws.Input<GPUBackend>(0));
+  else
+    DALI_FAIL("Internal error - input 0 is neither CPU nor GPU.");
+}
+
+template <typename Backend>
+void RepeatTensor(TensorList<Backend> &tl, const Tensor<Backend> &t, int N) {
+  tl.Reset();
+  tl.set_device_id(t.device_id());
+  tl.SetSize(N);
+  tl.set_sample_dim(t.ndim());
+  tl.set_type(t.type());
+  tl.SetLayout(t.GetLayout());
+  for (int i = 0; i < N; i++)
+    tl.SetSample(i, t);
+}
+
+template <typename Backend>
+void RepeatFirstSample(TensorList<Backend> &tl, int N) {
+  Tensor<Backend> t;
+  TensorShape<> shape = tl[0].shape();
+  t.ShareData(unsafe_sample_owner(tl, 0), shape.num_elements(), tl.is_pinned(),
+              shape, tl.type(), tl.device_id(), tl.order());
+  t.SetMeta(tl.GetMeta(0));
+  RepeatTensor(tl, t, N);
+}
+
+void LayoutToTL(TensorList<CPUBackend> &out, const Workspace &ws) {
+  TensorLayout l = ws.GetInputLayout(0);
+  out.Resize(uniform_list_shape(1, { l.size() }), DALI_UINT8);
+  memcpy(out.raw_mutable_tensor(0), l.data(), l.size());
+  RepeatFirstSample(out, ws.GetInputBatchSize(0));
+}
+
+void LayoutToTL(TensorList<GPUBackend> &out, const Workspace &ws) {
+  TensorLayout l = ws.GetInputLayout(0);
+  Tensor<CPUBackend> tmp_cpu;
+  Tensor<GPUBackend> tmp_gpu;
+  tmp_cpu.Resize(TensorShape<1>(l.size()), DALI_UINT8);
+  memcpy(tmp_cpu.raw_mutable_data(), l.data(), l.size());
+  tmp_cpu.set_order(ws.output_order());
+  tmp_gpu.set_order(ws.output_order());
+  tmp_gpu.Copy(tmp_cpu);
+
+  RepeatTensor(out, tmp_gpu, ws.GetInputBatchSize(0));
+}
+
+template <typename Backend>
+auto GetProperty<Backend>::GetPropertyReader(std::string_view key) -> PropertyReader {
+  if (key == "source_info") {
+    return static_cast<PropertyReaderFunc &>(SourceInfoToTL<Backend>);
+  } else if (key == "layout") {
+    return static_cast<PropertyReaderFunc &>(LayoutToTL);
+  } else {
+    DALI_FAIL(make_string("Unsupported property key: ", key));
+  }
+}
+
+
 DALI_REGISTER_OPERATOR(GetProperty, GetProperty<CPUBackend>, CPU)
 DALI_REGISTER_OPERATOR(GetProperty, GetProperty<GPUBackend>, GPU)
 

dali/operators/util/get_property.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -17,8 +17,8 @@
 
 #include <memory>
 #include <string>
+#include <string_view>
 #include <vector>
-#include "dali/operators/util/property.h"
 #include "dali/pipeline/data/type_traits.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/checkpointing/stateless_operator.h"
@@ -32,41 +32,29 @@ class GetProperty : public StatelessOperator<Backend> {
   explicit GetProperty(const OpSpec &spec)
       : StatelessOperator<Backend>(spec),
         property_key_(spec.template GetArgument<std::string>("key")),
-        property_(PropertyFactory()) {}
-
-  ~GetProperty() override = default;
-  DISABLE_COPY_MOVE_ASSIGN(GetProperty);
+        property_reader_(GetPropertyReader(property_key_)) {}
 
  protected:
   bool CanInferOutputs() const override {
-    return true;
+    return false;  // we may broadcast a common value to all samples
   }
 
   bool SetupImpl(std::vector<OutputDesc> &output_desc, const Workspace &ws) override {
-    const auto &input = ws.Input<Backend>(0);
-    output_desc.resize(1);
-    output_desc[0].shape = property_->GetShape(input);
-    output_desc[0].type = property_->GetType(input);
-    return true;
+    return false;
   }
 
   void RunImpl(Workspace &ws) override {
-    property_->FillOutput(ws);
+    property_reader_(ws.Output<Backend>(0), ws);
   }
 
  private:
-  std::unique_ptr<tensor_property::Property<Backend>> PropertyFactory() {
-    if (property_key_ == "source_info") {
-      return std::make_unique<tensor_property::SourceInfo<Backend>>();
-    } else if (property_key_ == "layout") {
-      return std::make_unique<tensor_property::Layout<Backend>>();
-    } else {
-      DALI_FAIL(make_string("Unknown property key: ", property_key_));
-    }
-  }
+  using PropertyReaderFunc = void(TensorList<Backend> &, const Workspace &);
+  using PropertyReader = std::function<PropertyReaderFunc>;
+
+  std::string property_key_;
+  PropertyReader property_reader_;
 
-  const std::string property_key_;
-  std::unique_ptr<tensor_property::Property<Backend>> property_;
+  static PropertyReader GetPropertyReader(std::string_view key);
 };
 
 }  // namespace dali

dali/python/nvidia/dali/data_node.py

@@ -276,6 +276,32 @@ def shape(self, *, dtype=None, device="cpu"):
             self._check_gpu2cpu()
         return fn.shapes(self, dtype=dtype, device=device)
 
+    def property(self, key, *, device="cpu"):
+        """Returns a metadata property associated with a DataNode
+
+        Parameters
+        ----------
+        key : str
+            The name of the metadata item. Currently supported:
+            "source_info"   - the file name or location in the dataset where the data originated
+                              (each sample is a 1D uint8 tensor)
+            "layout"        - the layout string
+                              (each sample is a 1D uint8 tensor)
+        device : str, optional
+            The device, where the value is returned; defaults to CPU.
+        """
+
+        from . import fn
+
+        if device == "cpu":
+            self._check_gpu2cpu()
+
+        return fn.get_property(self, key=key, device=device)
+
+    def source_info(self, *, device="cpu"):
+        """Returns the "source_info" property. Equivalent to self.meta("source_info")."""
+        return self.property("source_info", device=device)
+
     def _check_gpu2cpu(self):
         if self.device == "gpu" and self.source and self.source.pipeline:
             if not self.source.pipeline._exec_dynamic: