C API 2.0: External source info (#5872)

Add an ability to obtain information about pipeline inputs including:
- name
- device
- dtype
- ndim
- layout

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

Author: mzient

dali/c_api_2/pipeline.cc

@@ -88,6 +88,74 @@ void PipelineWrapper::Prefetch() {
   pipeline_->Prefetch();
 }
 
+int PipelineWrapper::GetInputCount() const {
+  return pipeline_->GetInputOperators().size();
+}
+
+daliPipelineIODesc_t PipelineWrapper::GetInputDesc(int idx) const & {
+  auto &inputs = pipeline_->GetInputOperators();
+  size_t n = inputs.size();
+  if (idx < 0 || static_cast<size_t>(idx) >= n)
+    throw std::out_of_range(make_string(
+        "The input index ", idx, " is out of range. The valid range is [0..", n-1, "]."));
+
+  if (input_names_.size() != n) {
+    input_names_.clear();
+    input_names_.reserve(n);
+    for (auto it = inputs.begin(); it != inputs.end(); it++) {
+      input_names_.push_back(it->first);
+    }
+  }
+  return GetInputDesc(input_names_[idx]);
+}
+
+namespace {
+
+template <typename Backend>
+void FillPipelineDesc(daliPipelineIODesc_t &desc, const InputOperator<Backend> &inp) {
+  int ndim = inp.in_ndim();
+  if (ndim >= 0) {
+    desc.ndim_present = true;
+    desc.ndim = ndim;
+  }
+  auto dtype = inp.in_dtype();
+  if (dtype != DALI_NO_TYPE) {
+    desc.dtype_present = true;
+    desc.dtype = dtype;
+  }
+  auto &layout = inp.in_layout();
+  if (layout.size())
+    desc.layout = layout.c_str();
+  else
+    desc.layout = nullptr;
+}
+
+}  // namespace
+
+daliPipelineIODesc_t PipelineWrapper::GetInputDesc(std::string_view name) const & {
+  auto &inputs = pipeline_->GetInputOperators();
+  auto it = inputs.find(name);
+  if (it == inputs.end())
+    throw invalid_key(make_string("The input with the name \"", name, "\" was not found."));
+
+  daliPipelineIODesc_t desc{};
+  desc.name = it->first.c_str();
+  desc.device = it->second->op_type == OpType::GPU ? DALI_STORAGE_GPU : DALI_STORAGE_CPU;
+  auto *op = pipeline_->GetOperator(name);
+  if (auto *inp = dynamic_cast<InputOperator<CPUBackend> *>(op))
+    FillPipelineDesc(desc, *inp);
+  else if (auto *inp = dynamic_cast<InputOperator<GPUBackend> *>(op))
+    FillPipelineDesc(desc, *inp);
+  else if (auto *inp = dynamic_cast<InputOperator<MixedBackend> *>(op))
+    FillPipelineDesc(desc, *inp);
+  else
+    throw std::logic_error(make_string(
+      "Internal error - the operator \"", name, "\" was found in the input operators map, but "
+      "it's not an instance of InputOperator<Backend>."));
+  return desc;
+}
+
+
 int PipelineWrapper::GetFeedCount(std::string_view input_name) {
   return pipeline_->InputFeedCount(input_name);
 }
@@ -122,21 +190,22 @@ int PipelineWrapper::GetOutputCount() const {
   return pipeline_->output_descs().size();
 }
 
-daliPipelineOutputDesc_t PipelineWrapper::GetOutputDesc(int idx) const {
+daliPipelineIODesc_t PipelineWrapper::GetOutputDesc(int idx) const & {
   auto &outputs = pipeline_->output_descs();
   int nout = outputs.size();
   if (idx < 0 || idx >= nout)
     throw std::out_of_range(make_string(
         "The output index ", idx, " is out of range. "
         "Valid range is [0..", nout-1, "]."));
   auto &out = outputs[idx];
-  daliPipelineOutputDesc_t desc{};
+  daliPipelineIODesc_t desc{};
   desc.device         = static_cast<daliStorageDevice_t>(out.device);
   desc.dtype          = out.dtype;
   desc.dtype_present  = out.dtype != DALI_NO_TYPE;
   desc.name           = out.name.c_str();
   desc.ndim           = out.ndim;
   desc.ndim_present   = out.ndim >= 0;
+  desc.layout         = out.layout.c_str();
   return desc;
 }
 
@@ -256,6 +325,37 @@ daliResult_t daliPipelineFeedInput(
   DALI_EPILOG();
 }
 
+daliResult_t daliPipelineGetInputCount(daliPipeline_h pipeline, int *out_input_count) {
+  DALI_PROLOG();
+  auto p = ToPointer(pipeline);
+  CHECK_OUTPUT(out_input_count);
+  *out_input_count = p->GetInputCount();
+  DALI_EPILOG();
+}
+
+DALI_API daliResult_t daliPipelineGetInputDescByIdx(
+      daliPipeline_h pipeline,
+      daliPipelineIODesc_t *out_input_desc,
+      int index) {
+  DALI_PROLOG();
+  auto p = ToPointer(pipeline);
+  CHECK_OUTPUT(out_input_desc);
+  *out_input_desc = p->GetInputDesc(index);
+  DALI_EPILOG();
+}
+
+DALI_API daliResult_t daliPipelineGetInputDesc(
+      daliPipeline_h pipeline,
+      daliPipelineIODesc_t *out_input_desc,
+      const char *name) {
+  DALI_PROLOG();
+  auto p = ToPointer(pipeline);
+  CHECK_OUTPUT(out_input_desc);
+  NOT_NULL(name);
+  *out_input_desc = p->GetInputDesc(name);
+  DALI_EPILOG();
+}
+
 daliResult_t daliPipelineGetOutputCount(daliPipeline_h pipeline, int *out_count) {
   DALI_PROLOG();
   auto pipe = ToPointer(pipeline);
@@ -266,7 +366,7 @@ daliResult_t daliPipelineGetOutputCount(daliPipeline_h pipeline, int *out_count)
 
 daliResult_t daliPipelineGetOutputDesc(
       daliPipeline_h pipeline,
-      daliPipelineOutputDesc_t *out_desc,
+      daliPipelineIODesc_t *out_desc,
       int index) {
   DALI_PROLOG();
   auto pipe = ToPointer(pipeline);

dali/c_api_2/pipeline.h

@@ -60,7 +60,13 @@ class PipelineWrapper : public _DALIPipeline {
 
   int GetOutputCount() const;
 
-  daliPipelineOutputDesc_t GetOutputDesc(int idx) const;
+  daliPipelineIODesc_t GetOutputDesc(int idx) const &;
+
+  int GetInputCount() const;
+
+  daliPipelineIODesc_t GetInputDesc(int idx) const &;
+
+  daliPipelineIODesc_t GetInputDesc(std::string_view name) const &;
 
   /** Retrieves the underlying DALI Pipeline object */
   dali::Pipeline *Unwrap() const & {
@@ -77,6 +83,7 @@ class PipelineWrapper : public _DALIPipeline {
         AccessOrder order);
 
   std::unique_ptr<Pipeline> pipeline_;
+  mutable std::vector<std::string_view> input_names_;
 };
 
 PipelineWrapper *ToPointer(daliPipeline_h handle);

dali/c_api_2/pipeline_test.cc

@@ -159,13 +159,16 @@ std::string GetPipelineWithExternalSource(
       StorageDevice dev_type,
       int max_batch_size,
       int num_threads,
-      int device_id) {
+      int device_id,
+      bool extended_input_desc = false) {
   Pipeline p(max_batch_size, num_threads, device_id);
   OpSpec src = OpSpec("ExternalSource")
     .AddOutput("out", dev_type)
     .AddArg("device", dev_type == StorageDevice::CPU ? "cpu" : "gpu")
     .AddArg("name", "ext")
     .AddArg("batch_size", max_batch_size);
+  if (extended_input_desc)
+    src.AddArg("ndim", 3).AddArg("layout", "HWC").AddArg("dtype", DALI_UINT8);
   p.AddOperator(src, "ext");
   p.SetOutputDescs({ {"out", to_string(dev_type)} });
   return p.SerializeToProtobuf();
@@ -264,7 +267,7 @@ void TestPipelineRun(PipelineType ptype) {
   int count;
   CHECK_DALI(daliPipelineGetOutputCount(h, &count));
   ASSERT_EQ(count, 2);
-  daliPipelineOutputDesc_t desc{};
+  daliPipelineIODesc_t desc{};
   EXPECT_EQ(daliPipelineGetOutputDesc(h, &desc, -1), DALI_ERROR_OUT_OF_RANGE);
   EXPECT_EQ(daliPipelineGetOutputDesc(h, &desc, count), DALI_ERROR_OUT_OF_RANGE);
   daliClearLastError();
@@ -502,7 +505,47 @@ TEST(CAPI2_PipelineTest, FeedInputGPUAsync) {
   TestFeedInput<GPUBackend>({});
 }
 
+TEST(CAPI2_PipelineTest, InputDescSimple) {
+  auto proto = GetPipelineWithExternalSource(dali::StorageDevice::GPU, 4, 4, 0, false);
+  daliPipelineParams_t params{};
+  params.exec_type_present = true;
+  params.exec_type = DALI_EXEC_DYNAMIC;
 
+  auto h = Deserialize(proto, params);
+  CHECK_DALI(daliPipelineBuild(h));
+  int count = 0;
+  CHECK_DALI(daliPipelineGetInputCount(h, &count));
+  ASSERT_EQ(count, 1);
+  daliPipelineIODesc_t desc{};
+  CHECK_DALI(daliPipelineGetInputDescByIdx(h, &desc, 0));
+  EXPECT_EQ(desc.device, DALI_STORAGE_GPU);
+  EXPECT_STREQ(desc.name, "ext");
+  EXPECT_FALSE(desc.ndim_present);
+  EXPECT_FALSE(desc.dtype_present);
+  EXPECT_EQ(desc.layout, nullptr);
+}
+
+TEST(CAPI2_PipelineTest, InputDescExtended) {
+  auto proto = GetPipelineWithExternalSource(dali::StorageDevice::CPU, 4, 4, 0, true);
+  daliPipelineParams_t params{};
+  params.exec_type_present = true;
+  params.exec_type = DALI_EXEC_DYNAMIC;
+
+  auto h = Deserialize(proto, params);
+  CHECK_DALI(daliPipelineBuild(h));
+  int count = 0;
+  CHECK_DALI(daliPipelineGetInputCount(h, &count));
+  ASSERT_EQ(count, 1);
+  daliPipelineIODesc_t desc{};
+  CHECK_DALI(daliPipelineGetInputDescByIdx(h, &desc, 0));
+  EXPECT_EQ(desc.device, DALI_STORAGE_CPU);
+  EXPECT_STREQ(desc.name, "ext");
+  EXPECT_TRUE(desc.ndim_present);
+  EXPECT_EQ(desc.ndim, 3);
+  EXPECT_TRUE(desc.dtype_present);
+  EXPECT_EQ(desc.dtype, DALI_UINT8);
+  EXPECT_STREQ(desc.layout, "HWC");
+}
 
 }  // namespace test
 }  // namespace c_api

dali/pipeline/pipeline.h

@@ -263,6 +263,15 @@ class DLL_PUBLIC Pipeline {
    */
   DLL_PUBLIC const graph::OpNode *GetInputOperatorNode(std::string_view name);
 
+  /**
+   * @brief Get input operatos as a name-to-node mapping.
+   *
+   */
+  DLL_PUBLIC const auto &GetInputOperators() const & {
+    DALI_ENFORCE(built_);
+    return input_operators_;
+  }
+
   /** @{ */
   /**
    * @brief Performs some checks on the user-constructed pipeline, setups data

include/dali/dali.h

@@ -336,7 +336,7 @@ typedef struct _DALIPipelineParams {
 } daliPipelineParams_t;
 
 /** Describes an output of a DALI Pipeline */
-typedef struct _DALIPipelineOutputDesc {
+typedef struct _DALIPipelineIODesc {
   const char *name;
   daliStorageDevice_t device;
   struct {
@@ -345,7 +345,8 @@ typedef struct _DALIPipelineOutputDesc {
   };
   daliDataType_t dtype;
   int ndim;
-} daliPipelineOutputDesc_t;
+  const char *layout;
+} daliPipelineIODesc_t;
 
 /** Creates an empty pipeline. */
 DALI_API daliResult_t daliPipelineCreate(
@@ -461,6 +462,54 @@ DALI_API daliResult_t daliPipelineFeedInput(
   daliFeedInputFlags_t options,
   const cudaStream_t *stream);
 
+/** Gets the number of pipeline inputs.
+ *
+ * NOTE: The pipeline must be built before calling this function.
+ *
+ * @param pipeline        [in]  The pipeline
+ * @param out_input_count [out] A pointer to the location where the number of pipeline inputs is
+ *                              stored.
+ *
+ * @retval DALI_SUCCESS
+ * @retval DALI_ERROR_INVALID_OPERATION   the pipeline wasn't built before the call
+ */
+DALI_API daliResult_t daliPipelineGetInputCount(daliPipeline_h pipeline, int *out_input_count);
+
+/** Gets a descriptor of a pipeline input specified by index.
+ *
+ * NOTE: The pipeline must be built before calling this function.
+ *
+ * @param pipeline        [in]  The pipeline
+ * @param out_input_desc  [out] A pointer to the location where the descriptor is written.
+ * @param index           [in]  The 0-based index of the input. See `daliPipelineGetInputCount`.
+ *
+ * @retval DALI_SUCCESS
+ * @retval DALI_ERROR_INVALID_OPERATION   the pipeline wasn't built before the call
+ * @retval DALI_ERROR_OUT_OF_RANGE        the index is not a valid 0-based index of the an input
+ */
+DALI_API daliResult_t daliPipelineGetInputDescByIdx(
+  daliPipeline_h pipeline,
+  daliPipelineIODesc_t *out_input_desc,
+  int index);
+
+/** Gets a descriptor of a pipeline input specified by its name.
+ *
+ * NOTE: The pipeline must be built before calling this function.
+ *
+ * @param pipeline        [in]  The pipeline
+ * @param out_input_desc  [out] A pointer to the location where the descriptor is written.
+ * @param name            [in]  The name of the input whose descriptor to obtain.]
+ *
+ * @retval DALI_SUCCESS
+ * @retval DALI_ERROR_INVALID_OPERATION   the pipeline wasn't built before the call
+ * @retval DALI_ERROR_INVALID_KEY         if `input_name` is not a valid name of an input of the
+ *                                        pipeline
+ */
+DALI_API daliResult_t daliPipelineGetInputDesc(
+  daliPipeline_h pipeline,
+  daliPipelineIODesc_t *out_input_desc,
+  const char *name);
+
 /** Gets the number of pipeline outputs.
  *
  * @param pipeline  [in]  The pipeline
@@ -481,7 +530,7 @@ DALI_API daliResult_t daliPipelineGetOutputCount(daliPipeline_h pipeline, int *o
  */
 DALI_API daliResult_t daliPipelineGetOutputDesc(
   daliPipeline_h pipeline,
-  daliPipelineOutputDesc_t *out_desc,
+  daliPipelineIODesc_t *out_desc,
   int index);