Add out-of-bounds-policy (including pad support) to Slice/Crop

dali/kernels/slice/slice_gpu.cuh

@@ -106,25 +106,33 @@ __device__ void SliceFunc(OutputType *__restrict__ out, const InputType *__restr
   if (Dims > 1 && out_strides[Dims - 1] == in_strides[Dims - 1] && anchor[Dims - 1] == 0 &&
       channel_dim != Dims - 1) {
     const int NextDims = Dims > 1 ? Dims - 1 : 1;
-    SliceFunc<NextDims, OutputType, InputType, false>(out, in, out_strides, in_strides, anchor,
-                                                      in_shape, fill_values, channel_dim, offset,
-                                                      block_end);
+    SliceFunc<NextDims, OutputType, InputType, false>(
+        out, in, out_strides, in_strides, anchor, in_shape, fill_values, channel_dim, offset,
+        block_end);
     return;
   }
 
+  constexpr int LastDim = Dims - 1;
+  int64_t inner_in_anchor = anchor[LastDim];
+  int64_t inner_in_extent = in_shape[LastDim];
+  if (!AllDims) {  // if we fused dimensions, adjust inner dimension's anchor and extent
+    inner_in_anchor = anchor[LastDim] * in_strides[LastDim];
+    inner_in_extent = Dims > 1 ? in_strides[LastDim - 1] : in_shape[LastDim] * in_strides[LastDim];
+  }
+
   for (; offset < block_end; offset += blockDim.x) {
     uint64_t idx = offset;
     uint64_t out_idx = idx;
 
     // If no dimensions were skipped (AllDims=true) we can avoid division in the last dimension,
-    // because know the stride is 1
+    // because know the strides are 1 (or we treat them as 1 if we fused dimensions)
     int i_c = 0;
     int i_d;
     bool out_of_bounds = false;
     uint64_t in_idx = 0;
 
     #pragma unroll
-    for (int d = 0; d < Dims - AllDims; d++) {
+    for (int d = 0; d < Dims - 1; d++) {
       i_d = div_mod(idx, idx, out_strides[d]);
       if (d == channel_dim)
         i_c = i_d;
@@ -133,16 +141,13 @@ __device__ void SliceFunc(OutputType *__restrict__ out, const InputType *__restr
         in_idx += i_d * in_strides[d];
     }
 
-    // Here we handle the last dimension (if we didn't in the main loop)
-    if (AllDims) {
-      constexpr int d = Dims - 1;
-      i_d = idx;  // We know that out_strides[d] is 1
-      if (d == channel_dim)
-        i_c = i_d;
-      out_of_bounds |= is_out_of_bounds(anchor[d] + i_d, in_shape[d]);
-      if (!out_of_bounds)
-        in_idx += i_d;  // We know that in_strides[d] is 1
-    }
+    constexpr int d = LastDim;
+    i_d = idx;  // out_strides[d] is 1
+    if (AllDims && d == channel_dim)
+      i_c = i_d;
+    out_of_bounds |= is_out_of_bounds(inner_in_anchor + i_d, inner_in_extent);
+    if (!out_of_bounds)
+      in_idx += i_d;  // in_strides[d] is 1
 
     // Fill values are reused a lot, so let's make sure they are cached (by using __ldg())
     out[out_idx] = out_of_bounds ? __ldg(&fill_values[i_c]) : clamp<OutputType>(in[in_idx]);

dali/kernels/slice/slice_kernel_test.h

@@ -242,6 +242,22 @@ struct ArgsGen_CompletelyOutOfBounds{
   }
 };
 
+template <typename OutputType, int Dims = 3>
+struct ArgsGen_SingleValuePad {
+  SliceArgs<OutputType, Dims> Get(const TensorShape<Dims>& input_shape) {
+    SliceArgs<OutputType, 3> args;
+    args.anchor[0] = -input_shape[0] / 2;
+    args.anchor[1] = -input_shape[1] / 2;
+    args.anchor[2] = 0;
+    args.shape[0] = 2 * input_shape[0];
+    args.shape[1] = 2 * input_shape[1];
+    args.shape[2] = input_shape[2];
+    args.fill_values = {128};
+    args.channel_dim = -1;
+    return args;
+  }
+};
+
 template <typename OutputType, int Dims = 3>
 struct ArgsGen_MultiChannelPad {
   SliceArgs<OutputType, Dims> Get(const TensorShape<Dims>& input_shape) {
@@ -332,6 +348,7 @@ using SLICE_TEST_TYPES = ::testing::Types<
     SliceTestArgs<int, int, 1, 1, 22, ArgsGen_RightSideOutOfBounds<int, 1>>,
     SliceTestArgs<int, int, 2, 1, 22, ArgsGen_RightSideOutOfBounds<int, 2>>,
     SliceTestArgs<int, int, 2, 1, 22, ArgsGen_CompletelyOutOfBounds<int, 2>>,
+    SliceTestArgs<int, int, 3, 1, 20, ArgsGen_SingleValuePad<int, 3>, 20, 20, 3>,
     SliceTestArgs<int, int, 3, 1, 20, ArgsGen_MultiChannelPad<int, 3>, 20, 20, 3>,
     SliceTestArgs<int, int, 3, 1, 20, ArgsGen_MultiChannelPad_ChFirst<int, 3>, 3, 20, 20>,
     SliceTestArgs<int, int, 3, 1, 20, ArgsGen_PadAlsoChDim<int, 3>, 20, 20, 3>,

dali/operators/generic/slice/out_of_bounds_attr.cc

@@ -0,0 +1,35 @@
+// Copyright (c) 2020, NVIDIA CORPORATION. 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 <vector>
+#include "dali/pipeline/operator/common.h"
+#include "dali/pipeline/operator/operator.h"
+
+namespace dali {
+
+DALI_SCHEMA(OutOfBoundsAttr)
+    .DocStr(R"code(Out-of-bounds slicing attributes placeholder)code")
+    .AddOptionalArg("out_of_bounds_policy",
+        R"code(Determines the policy when slicing out of bounds of the input.
+Supported values are:
+
+- "error" (default) : Attempting to slice outside of the bounds of the image will produce an error.
+- "pad": The input will be padded as needed with zeros or any other value specified with ``fill_values`` argument.
+- "trim_to_shape": The slice window will be cut to the bounds of the input.))code", "error")
+    .AddOptionalArg("fill_values",
+        R"code(Determines padding values, only relevant if ``out_of_bounds_policy`` is set to "pad".
+If a scalar is provided, it will be used for all the channels. If multiple values are given, there should be as many values as
+channels (extent of dimension 'C' in the layout) in the output slice.)code", std::vector<float>{0.f});
+
+}  // namespace dali

dali/operators/generic/slice/out_of_bounds_policy.h

@@ -0,0 +1,102 @@
+// Copyright (c) 2020, NVIDIA CORPORATION. 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_SLICE_OUT_OF_BOUNDS_POLICY_H_
+#define DALI_OPERATORS_GENERIC_SLICE_OUT_OF_BOUNDS_POLICY_H_
+
+#include <string>
+#include "dali/core/math_util.h"
+#include "dali/core/tensor_shape.h"
+#include "dali/core/tensor_shape_print.h"
+#include "dali/pipeline/operator/common.h"
+
+namespace dali {
+
+template <bool inclusive_end>
+DALI_HOST_DEV DALI_FORCEINLINE bool is_out_of_bounds(int64_t idx, int64_t data_extent) {
+  if (inclusive_end)  // check idx is within [0, data_extent]
+    return static_cast<uint64_t>(idx) > static_cast<uint64_t>(data_extent);
+  else                // check idx is within [0, data_extent)
+    return static_cast<uint64_t>(idx) >= static_cast<uint64_t>(data_extent);
+}
+
+/**
+ * @brief Determines what to do if slice parameters point to outside of the input bounds
+ */
+enum class OutOfBoundsPolicy {
+  Error,        // sampling out of bounds will throw an error
+  TrimToShape,  // Slice shape will be trimmed to fit the input bounds (potentially empty output)
+  Pad,  // Slicing out of bounds will result in padding with zeroes or any other provided value(s)
+};
+
+inline OutOfBoundsPolicy GetOutOfBoundsPolicy(const OpSpec &spec) {
+  bool has_out_of_bounds_policy = spec.HasArgument("out_of_bounds_policy");
+  OutOfBoundsPolicy policy = OutOfBoundsPolicy::Error;
+  if (has_out_of_bounds_policy) {
+    auto policy_str = spec.GetArgument<std::string>("out_of_bounds_policy");
+    if (policy_str == "pad") {
+      policy = OutOfBoundsPolicy::Pad;
+    } else if (policy_str == "trim_to_shape") {
+      policy = OutOfBoundsPolicy::TrimToShape;
+    } else if (policy_str == "error") {
+      policy = OutOfBoundsPolicy::Error;
+    } else {
+      DALI_FAIL(
+          make_string("Not supported out_of_bounds_policy: ", policy_str,
+                      ". Supported values are \"pad\", \"trim_to_shape\", \"error\" (default)"));
+    }
+  }
+  return policy;
+}
+
+template <int Dims>
+void ApplySliceBoundsPolicy(OutOfBoundsPolicy policy, const TensorShape<Dims> &input_shape,
+                            TensorShape<Dims> &slice_anchor, TensorShape<Dims> &slice_shape) {
+  DALI_ENFORCE(
+      input_shape.size() == slice_anchor.size() && input_shape.size() == slice_shape.size(),
+      "Slice arguments should have the same number of dimensions as the input");
+  switch (policy) {
+    case OutOfBoundsPolicy::Pad:
+      // nothing to do
+      break;
+
+    case OutOfBoundsPolicy::TrimToShape:
+      for (int d = 0; d < input_shape.size(); d++) {
+        auto slice_start = clamp<int64_t>(slice_anchor[d], 0, input_shape[d]);
+        auto slice_end   = clamp<int64_t>(slice_anchor[d] + slice_shape[d], 0, input_shape[d]);
+        assert(slice_end >= slice_start);
+        slice_anchor[d] = slice_start;
+        slice_shape[d] = slice_end - slice_start;
+      }
+      break;
+
+    case OutOfBoundsPolicy::Error:
+    default:
+      for (int d = 0; d < input_shape.size(); d++) {
+        // start within [0, extent), and end within [0, extent]
+        if (is_out_of_bounds<false>(slice_anchor[d], input_shape[d]) ||
+            is_out_of_bounds<true>(slice_anchor[d] + slice_shape[d], input_shape[d])) {
+          DALI_FAIL(make_string(
+              "Slice can't be place out of bounds with current policy. Got: input_shape={",
+              input_shape, "}, slice_anchor={", slice_anchor, "}, slice_shape={", slice_shape,
+              "}"));
+        }
+      }
+    break;
+  }
+}
+
+}  // namespace dali
+
+#endif  // DALI_OPERATORS_GENERIC_SLICE_OUT_OF_BOUNDS_POLICY_H_

dali/operators/generic/slice/slice.cc

@@ -13,26 +13,27 @@
 // limitations under the License.
 
 #include "dali/operators/generic/slice/slice.h"
+#include "dali/kernels/slice/slice_cpu.h"
 
 namespace dali {
 
 DALI_SCHEMA(Slice)
     .DocStr(
-        R"code(Extract a subtensor or `slice` with a given shape and anchor.
-Inputs must be supplied as 3 separate tensors in a specific order: `data`, `anchor` and `shape`.
-Both `anchor` and `shape` coordinates must be within the interval
+        R"code(Extract a subtensor or ``slice`` with a given shape and anchor.
+Inputs must be supplied as 3 separate tensors in a specific order: ``data``, ``anchor`` and ``shape``.
+Both ``anchor`` and ``shape`` coordinates must be within the interval
 [0.0, 1.0] for normalized coordinates, or within the image shape for absolute
-coordinates. Both `anchor` and `shape` inputs will provide as many dimensions as specified
-with arguments `axis_names` or `axes`. By default `Slice` operator uses normalized
-coordinates and `WH` order for the slice arguments.)code")
+coordinates. Both ``anchor`` and ``shape`` inputs will provide as many dimensions as specified
+with arguments ``axis_names`` or ``axes``. By default ``Slice`` operator uses normalized
+coordinates and ``WH`` order for the slice arguments.)code")
     .NumInput(3)
     .NumOutput(1)
     .InputDox(0, "data", "TensorList", "Batch containing input data")
-    .InputDox(1, "anchor", "1D TensorList of floats",
+    .InputDox(1, "anchor", "1D TensorList of float",
                  R"code(Input containing either normalized or absolute coordinates
 (depending on the value of `normalized_anchor`) for the starting point of the
 slice (x0, x1, x2, ...).)code")
-    .InputDox(2, "shape", "1D TensorList of floats",
+    .InputDox(2, "shape", "1D TensorList of float",
                  R"code(Input containing either normalized or absolute coordinates
 (depending on the value of `normalized_shape`) for the dimensions of the slice
 (s0, s1, s2, ...).)code")
@@ -42,27 +43,8 @@ slice (x0, x1, x2, ...).)code")
       R"code(The color space of input and output image)code",
       DALI_RGB, false)
     .AddParent("SliceBase")
-    .AddParent("SliceAttr");
-
-template <>
-void Slice<CPUBackend>::DataDependentSetup(SampleWorkspace &ws) {
-  slice_attr_.ProcessArguments(ws);
-  const auto &images = ws.Input<CPUBackend>(kImagesInId);
-  auto data_idx = ws.data_idx();
-  auto crop_window_generator = slice_attr_.GetCropWindowGenerator(data_idx);
-  DALI_ENFORCE(crop_window_generator);
-  auto layout = InputLayout(ws, 0);
-  if (layout.empty())
-    layout = GetDefaultLayout(images.shape().size());
-  CropWindow win = crop_window_generator(images.shape(), layout);
-  slice_shapes_[data_idx] = std::vector<int64_t>(win.shape.begin(), win.shape.end());
-  slice_anchors_[data_idx] = std::vector<int64_t>(win.anchor.begin(), win.anchor.end());
-}
-
-template <>
-void Slice<CPUBackend>::RunImpl(SampleWorkspace &ws) {
-  SliceBase<CPUBackend>::RunImpl(ws);
-}
+    .AddParent("SliceAttr")
+    .AddParent("OutOfBoundsAttr");
 
 DALI_REGISTER_OPERATOR(Slice, Slice<CPUBackend>, CPU);
 

dali/operators/generic/slice/slice.cu

@@ -12,28 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <vector>
 #include "dali/operators/generic/slice/slice.h"
+#include "dali/kernels/slice/slice_gpu.cuh"
 
 namespace dali {
 
-template <>
-void Slice<GPUBackend>::DataDependentSetup(DeviceWorkspace &ws) {
-  slice_attr_.ProcessArguments(ws);
-  const auto &images = ws.Input<GPUBackend>(kImagesInId);
-  for (int data_idx = 0; data_idx < batch_size_; data_idx++) {
-    const auto img_shape = images.tensor_shape(data_idx);
-    auto crop_window_generator = slice_attr_.GetCropWindowGenerator(data_idx);
-    DALI_ENFORCE(crop_window_generator);
-    auto layout = InputLayout(ws, 0);
-    if (layout.empty())
-      layout = GetDefaultLayout(img_shape.size());
-    CropWindow win = crop_window_generator(img_shape, layout);
-    slice_shapes_[data_idx] = std::vector<int64_t>(win.shape.begin(), win.shape.end());
-    slice_anchors_[data_idx] = std::vector<int64_t>(win.anchor.begin(), win.anchor.end());
-  }
-}
-
 DALI_REGISTER_OPERATOR(Slice, Slice<GPUBackend>, GPU);
 
 }  // namespace dali

dali/operators/generic/slice/slice.h

@@ -34,37 +34,15 @@ class Slice : public SliceBase<Backend> {
     , slice_attr_(spec) {}
 
  protected:
-  using SliceBase<Backend>::input_type_;
-  using SliceBase<Backend>::output_type_;
-  using SliceBase<Backend>::slice_anchors_;
-  using SliceBase<Backend>::slice_shapes_;
-
-  void RunImpl(Workspace<Backend> &ws) override {
-    SliceBase<Backend>::RunImpl(ws);
+  void ProcessCroppingAttrs(const workspace_t<Backend> &ws) override {
+    slice_attr_.ProcessArguments(ws);
   }
 
-  void SetupSharedSampleParams(Workspace<Backend> &ws) override {
-    DALI_ENFORCE(ws.NumInput() == 3,
-      "Expected 3 inputs. Received: " + std::to_string(ws.NumInput()));
-    SliceBase<Backend>::SetupSharedSampleParams(ws);
+  const CropWindowGenerator& GetCropWindowGenerator(std::size_t data_idx) const override {
+    return slice_attr_.GetCropWindowGenerator(data_idx);
   }
 
-  void DataDependentSetup(Workspace<Backend> &ws) override;
-
  private:
-  inline TensorLayout GetDefaultLayout(int ndims) {
-    switch (ndims) {
-      case 2:
-        return "HW";
-      case 3:
-        return "HWC";
-      case 4:
-        return "DHWC";
-      default:
-        return "";
-    }
-  }
-
   SliceAttr slice_attr_;
 
   static const int kImagesInId = 0;