[WIP] Histogram CPU

Input shape part.

Signed-off-by: Piotr Rak <prak@nvidia.com>

dali/operators/generic/reduce/axes_helper.h

@@ -0,0 +1,60 @@
+// Copyright (c) 2020-2021, 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.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+
+#ifndef DALI_OPERATORS_GENERIC_REDUCE_AXIS_HELPER_H__
+#define DALI_OPERATORS_GENERIC_REDUCE_AXIS_HELPER_H__
+
+#include <vector>
+
+#include "dali/pipeline/operator/operator.h"
+
+namespace dali {
+namespace detail {
+
+class AxesHelper {
+ public:
+  explicit inline AxesHelper(const OpSpec &spec) {
+    has_axes_arg_ = spec.TryGetRepeatedArgument(axes_, "axes");
+    has_axis_names_arg_ = spec.TryGetArgument(axis_names_, "axis_names");
+    has_empty_axes_arg_ =
+        (has_axes_arg_ && axes_.empty()) || (has_axis_names_arg_ && axis_names_.empty());
+
+    DALI_ENFORCE(!has_axes_arg_ || !has_axis_names_arg_,
+                 "Arguments `axes` and `axis_names` are mutually exclusive");
+  }
+
+  void PrepareAxes(const TensorLayout &layout, int sample_dim) {
+    if (has_axis_names_arg_) {
+      axes_ = GetDimIndices(layout, axis_names_).to_vector();
+      return;
+    }
+
+    if (!has_axes_arg_) {
+      axes_.resize(sample_dim);
+      std::iota(axes_.begin(), axes_.end(), 0);
+    }
+  }
+
+  bool has_axes_arg_;
+  bool has_axis_names_arg_;
+  bool has_empty_axes_arg_;
+  std::vector<int> axes_;
+  TensorLayout axis_names_;
+};
+
+}  // namespace detail
+}  // namespace dali
+
+#endif  // DALI_OPERATORS_GENERIC_REDUCE_AXIS_HELPER_H__

dali/operators/generic/reduce/histogram.cc

@@ -0,0 +1,17 @@
+// Copyright (c) 2021, 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.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "dali/operators/generic/reduce/histogram.h"
+
+using namespace dali::hist_detail;

dali/operators/generic/reduce/histogram.h

@@ -0,0 +1,248 @@
+// Copyright (c) 2021, 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.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DALI_OPERATORS_GENERIC_REDUCE_HISTOGRAM_H_
+#define DALI_OPERATORS_GENERIC_REDUCE_HISTOGRAM_H_
+
+#include "dali/kernels/kernel_manager.h"
+#include "dali/operators/generic/reduce/axes_helper.h"
+#include "dali/pipeline/operator/operator.h"
+
+namespace dali {
+namespace hist_detail {
+
+inline bool is_simple_reduction(const std::vector<int> &reduction_axes, int ndims,
+                                int hist_dim = 1) {
+  int dim = ndims - hist_dim;
+
+  // Starting from inner dimension, look if we should reduce this axis
+  // If we can do so until we can collapse next dimention.
+  auto raxis = reduction_axes.rbegin();
+
+  for (; raxis != reduction_axes.rend(); ++raxis) {
+    // TODO: Handle multidimentional histograms correctly
+    if (dim != *raxis) {
+      break;
+    }
+    --dim;
+  }
+
+  // If there's no reduction axes left, we won't need transpose any axis.
+  if (raxis == reduction_axes.rend())
+    return true;
+
+  return false;
+}
+
+// Collapses all inner dimensions
+std::vector<int> GetTransposeAxes(int ndims, const std::vector<int> &reduction_axes,
+                                  int hist_dim = 1) {
+  int dim = ndims - hist_dim;
+
+  // std::vector<int> dimensions_to_collapse;
+
+  // Starting from inner dimension, look if we should reduce this axis
+  // We can collapse those dimension into one, that can be used to calculate histogram.
+  auto raxis = reduction_axes.rbegin();
+
+  for (; raxis != reduction_axes.rend(); ++raxis, --dim) {
+    if (dim != *raxis)
+      break;
+    // dimensions_to_collapse.insert(dimensions_to_collapse.begin(), dim);
+  }
+
+  std::vector<int> axes_to_transpose;
+
+  // All axes that couldn't be collapsed, need to be transposed so we can collapse those together.
+  int naxes_to_transpose = std::distance(raxis, reduction_axes.rend());
+
+  axes_to_transpose.resize(naxes_to_transpose);
+
+  auto raxes_start = reduction_axes.begin() + naxes_to_transpose;
+  std::copy(reduction_axes.begin(), raxes_start, axes_to_transpose.begin());
+
+  return axes_to_transpose;
+}
+
+class TransposeForReduction {
+
+};
+
+}  // namespace hist_detail
+}  // namespace dali
+
+
+
+namespace dali {
+
+template <typename Backend>
+class Histogram : public Operator<Backend>, detail::AxesHelper {
+ public:
+  explicit inline Histogram(const OpSpec &spec)
+      : Operator<Backend>(spec), detail::AxesHelper(spec) {}
+
+  bool CanInferOutputs() const override {
+    return true;
+  }
+
+  ~Histogram() override = default;
+
+  void PrepareInputShapesForReduction(const TensorListShape<> &input_shapes, int hist_dim = 1) {
+    std::vector<TensorShape<>> reduction_input_shapes;
+    reduction_input_shapes.reserve(input_shapes.num_samples());
+
+    const int ndims = input_shapes.sample_dim();
+
+    // TODO: support higher dimentionality histograms (hist_dim)
+    std::array<std::pair<int, int>, 1> reduced_inner_dim(
+        {std::make_pair(ndims - axes_.size() - 1, ndims - 1)});
+
+    for (int i = 0; i < input_shapes.num_samples(); ++i) {
+      // Collapse inner dimensions
+      auto shape = collapse_dims(input_shapes.tensor_shape(i), reduced_inner_dim);
+      reduction_input_shapes.push_back(std::move(shape));
+    }
+
+    reduced_input_shapes_ = TensorListShape<>(reduction_input_shapes);
+  }
+
+  // For all subsequent dimensions go through the list and find axes not being
+  // reduced
+  SmallVector<int, 6> GetNonReductionAxes(const int ndims) const {
+    SmallVector<int, 6> non_reduction_axes;
+
+    int rax_search_start = 0;
+
+    for (int dim = 0; dim < ndims; ++dim) {
+      bool found = false;
+      int rax = rax_search_start;
+
+      while (rax < axes_.size() && axes_[rax] <= dim) {
+        // Since axes_ are sorted, we know we can skip ahead searching
+        // of next dimension to at least to next reduction axis
+        rax_search_start = rax + 1;
+
+        if (axes_[rax] == dim) {
+          found = true;
+          break;
+        }
+        ++rax;
+      }
+
+      if (!found) {
+        non_reduction_axes.push_back(dim);
+      }
+    }
+
+    return non_reduction_axes;
+  }
+
+  void PrepareInputShapesForTranspose(const TensorListShape<> &input_shapes, int hist_dim = 1) {
+    const int ndims = input_shapes.sample_dim();
+
+    auto shape_span
+    auto non_rediction_axes = GetNonReductionAxes(ndims);
+
+    // Axes inner of that can be collapsed directly and don't need to have an order changed
+    SmallVector<int, 6> inner_reduction_axes;
+
+    // Go through reduction axes in inner dimension order
+    // We stop when we find first axis that doesn't need reduction;
+    for (int rax = axes_.size() - 1, dim = ndims; rax >= 0; --rax, --dim) {
+      if (axes_[rax] != dim) {
+        break;
+      }
+
+      inner_reduction_axes.insert(inner_reduction_axes.begin(), dim);
+    }
+
+    auto axes_to_transpose = hist_detail::GetTransposeAxes(ndims, axes_);
+
+    SmallVector<int, 6> axes_order;
+    axes_order.reserve(ndims);
+    for (int axis : non_rediction_axes) {
+      axes_order.push_back(axis);
+    }
+
+    for (int axis : axes_to_transpose) {
+      axes_order.push_back(axis);
+    }
+
+    for (int axis : axes_to_transpose) {
+      axes_order.push_back(axis);
+    }
+
+    TensorListShape<> shapes;
+    permute_dims(shapes, input_shapes, axes_order);
+
+    PrepareInputShapesForReduction(shapes, hist_dim);
+
+    transpose_axes_order_ = std::move(axes_order);
+  }
+
+  void PreparingOutputShapes(int hist_ndim = 1) {
+    auto  GetNonReductionAxes(reduced_input_shapes_.sample_dim());
+
+    for (int i = 0; i < reduced_input_shapes_.num_samples(); ++i) {
+
+    }
+  }
+
+  bool SetupImpl(std::vector<OutputDesc> &output_desc, const workspace_t<Backend> &ws) override {
+    output_desc.resize(1);
+
+    auto &input = ws.template Input<Backend>(0);
+    const size_t ndims = input.shape().sample_dim();
+
+    PrepareAxes(input.GetLayout(), ndims);
+
+    // Empty reduction axes were specified, histogram calculation becomes identity operation
+    is_noop_ = (has_axes_arg_ || has_axis_names_arg_) && axes_.empty();
+    if (is_noop_) {
+      output_desc[0].type = input.type();
+      output_desc[0].shape = input.shape();
+    } else {
+      if (!axes_.empty()) {
+        // Ensure reduction axes are sorted in ascending order,
+        // so we can check dimentions that can be collapsed
+        std::sort(axes_.begin(), axes_.end());
+      }
+
+      if (hist_detail::is_simple_reduction(axes_, ndims)) {
+        PrepareInputShapesForReduction(input.shape());
+      } else {
+        PrepareInputShapesForTranspose(input.shape());
+        needs_transpose_ = true;
+      }
+    }
+    return true;
+  }
+
+  void RunImpl(workspace_t<Backend> &ws) override {}
+
+ private:
+  USE_OPERATOR_MEMBERS();
+  bool keep_dims_;
+  bool needs_transpose_ = false;
+  bool is_noop_ = false;
+  TensorListShape<> reduced_input_shapes_;
+  SmallVector<int, 6> transpose_axes_order_;
+  kernels::KernelManager kmgr_;
+  TensorView<StorageGPU, const int, 1> params_num_bins_gpu_;
+};
+
+}  // namespace dali
+
+
+#endif  // DALI_OPERATORS_GENERIC_REDUCE_HISTOGRAM_H_

dali/operators/generic/reduce/reduce.cc

@@ -99,6 +99,17 @@ DALI_SCHEMA(reductions__Max)
   .NumOutput(1)
   .AddParent("ReduceBase");
 
+DALI_SCHEMA(reductions__Histogram)
+  .DocStr("Calculates histogram.")
+  .NumInput(1, 3)
+  .NumOutput(1)
+  .AddArg("num_bins", "Number of bins", DALI_INT_VEC, true)
+  .AddParent("ReduceBase");
+
+
+using HistogramCPU = Histogram<CPUBackend>;
+DALI_REGISTER_OPERATOR(reductions__Histogram, HistogramCPU, CPU)
+
 using MeanCPU = MeanOp<kernels::MeanCPU, CPUBackend>;
 DALI_REGISTER_OPERATOR(reductions__Mean, MeanCPU, CPU);