Mark missed operator stateless, assure checkpointing tests coverage (#…

…5232)

* Track operators tested for checkpointing
* Mark missed operators as stateless
* Enable checkpointing in all stateless ops tests to verify they are properly marked
* Make list of unsupported ops
* Remove redundant stateless ops tests, add auto augment tests
* Make sure to run attr-ed test cases in CI

---------

Signed-off-by: Kamil Tokarski <[email protected]>

dali/operators/generic/resize/tensor_resize.h

@@ -23,6 +23,7 @@
 #include "dali/core/expand_dims.h"
 #include "dali/operators/image/resize/resize_base.h"
 #include "dali/operators/image/resize/tensor_resize_attr.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/operator.h"
 
@@ -31,7 +32,7 @@ namespace tensor_resize {
 
 
 template <typename Backend>
-class TensorResize : public Operator<Backend>
+class TensorResize : public StatelessOperator<Backend>
                    , protected ResizeBase<Backend> {
  public:
   explicit TensorResize(const OpSpec &spec);
@@ -166,7 +167,7 @@ class TensorResize : public Operator<Backend>
 
 template <typename Backend>
 TensorResize<Backend>::TensorResize(const OpSpec &spec)
-    : Operator<Backend>(spec), ResizeBase<Backend>(spec), resize_attr_(spec) {
+    : StatelessOperator<Backend>(spec), ResizeBase<Backend>(spec), resize_attr_(spec) {
   InitializeBackend();
 }
 

dali/operators/geometry/affine_transforms/combine_transforms.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2020-2023, 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.
@@ -20,10 +20,11 @@
 #include "dali/core/static_switch.h"
 #include "dali/kernels/kernel_manager.h"
 #include "dali/pipeline/data/types.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/op_spec.h"
-#include "dali/pipeline/workspace/workspace.h"
 #include "dali/pipeline/operator/operator.h"
 #include "dali/pipeline/operator/sequence_operator.h"
+#include "dali/pipeline/workspace/workspace.h"
 
 #define TRANSFORM_INPUT_TYPES (float)
 
@@ -49,10 +50,11 @@ Example: combining [T1, T2, T3] is equivalent to T3(T2(T1(...))) for default ord
   .AllowSequences()
   .AddParent("TransformAttr");
 
-class CombineTransformsCPU : public SequenceOperator<CPUBackend> {
+class CombineTransformsCPU : public SequenceOperator<CPUBackend, StatelessOperator> {
  public:
+  using Base = SequenceOperator<CPUBackend, StatelessOperator>;
   explicit CombineTransformsCPU(const OpSpec &spec) :
-      SequenceOperator<CPUBackend>(spec),
+      Base(spec),
       reverse_order_(spec.GetArgument<bool>("reverse_order")) {
   }
 

dali/operators/geometry/coord_transform.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2020-2023, 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.
@@ -21,6 +21,7 @@
 #include "dali/core/static_switch.h"
 #include "dali/kernels/kernel_manager.h"
 #include "dali/operators/geometry/mt_transform_attr.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/operator.h"
 #include "dali/pipeline/operator/sequence_operator.h"
 
@@ -30,9 +31,10 @@ namespace dali {
 #define COORD_TRANSFORM_DIMS (1, 2, 3, 4, 5, 6)
 
 template <typename Backend>
-class CoordTransform : public SequenceOperator<Backend, Operator, true>, private MTTransformAttr {
+class CoordTransform : public SequenceOperator<Backend, StatelessOperator, true>,
+                       private MTTransformAttr {
  public:
-  using Base = SequenceOperator<Backend, Operator, true>;
+  using Base = SequenceOperator<Backend, StatelessOperator, true>;
   explicit CoordTransform(const OpSpec &spec) : Base(spec), MTTransformAttr(spec) {
     dtype_ = spec_.template GetArgument<DALIDataType>("dtype");
   }

dali/operators/image/color/brightness_contrast.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2023, 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 @@
 #include "dali/core/static_switch.h"
 #include "dali/kernels/kernel_manager.h"
 #include "dali/pipeline/data/views.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/operator.h"
 #include "dali/pipeline/operator/sequence_operator.h"
@@ -54,17 +55,16 @@ const float kDefaultBrightnessShift = 0;
 const float kDefaultContrast = 1.f;
 
 template <typename Backend>
-class BrightnessContrastOp : public SequenceOperator<Backend> {
+class BrightnessContrastOp : public SequenceOperator<Backend, StatelessOperator> {
  public:
+  using Base = SequenceOperator<Backend, StatelessOperator>;
   ~BrightnessContrastOp() override = default;
 
   DISABLE_COPY_MOVE_ASSIGN(BrightnessContrastOp);
 
  protected:
   explicit BrightnessContrastOp(const OpSpec &spec)
-      : SequenceOperator<Backend>(spec),
-        output_type_(DALI_NO_TYPE),
-        input_type_(DALI_NO_TYPE) {
+      : Base(spec), output_type_(DALI_NO_TYPE), input_type_(DALI_NO_TYPE) {
     spec.TryGetArgument(output_type_arg_, "dtype");
   }
 
@@ -76,7 +76,7 @@ class BrightnessContrastOp : public SequenceOperator<Backend> {
   // the number of samples and parameters unnecessarily for FHWC when there are no
   // per-frame parameters provided.
   bool ShouldExpand(const Workspace &ws) override {
-    return SequenceOperator<Backend>::ShouldExpand(ws) && this->HasPerFrameArgInputs(ws);
+    return Base::ShouldExpand(ws) && this->HasPerFrameArgInputs(ws);
   }
 
   template <typename OutputType, typename InputType>
@@ -163,7 +163,7 @@ class BrightnessContrastCpu : public BrightnessContrastOp<CPUBackend> {
    * "overloaded virtual function `dali::Operator<dali::CPUBackend>::RunImpl` is only partially
    * overridden in class `dali::brightness_contrast::BrightnessContrast<dali::CPUBackend>`"
    */
-  using SequenceOperator<CPUBackend>::RunImpl;
+  using Base::RunImpl;
 
   ~BrightnessContrastCpu() override = default;
 

dali/operators/image/color/debayer.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2022-2023, 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.
@@ -20,6 +20,7 @@
 
 #include "dali/core/span.h"
 #include "dali/kernels/imgproc/color_manipulation/debayer/debayer.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/op_spec.h"
 #include "dali/pipeline/operator/operator.h"
@@ -99,10 +100,11 @@ class DebayerImplBase {
 
 
 template <typename Backend>
-class Debayer : public SequenceOperator<Backend> {
+class Debayer : public SequenceOperator<Backend, StatelessOperator> {
  public:
+  using Base = SequenceOperator<Backend, StatelessOperator>;
   explicit Debayer(const OpSpec &spec)
-      : SequenceOperator<Backend>(spec),
+      : Base(spec),
         alg_{debayer::parse_algorithm_name(spec.GetArgument<std::string>(debayer::kAlgArgName))} {
     if (!spec_.HasTensorArgument(debayer::kBluePosArgName)) {
       std::vector<int> blue_pos;

dali/operators/image/convolution/filter.h

@@ -22,6 +22,7 @@
 #include "dali/core/boundary.h"
 #include "dali/core/common.h"
 #include "dali/core/static_switch.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/operator.h"
 #include "dali/pipeline/operator/sequence_operator.h"
@@ -169,10 +170,11 @@ TensorListView<detail::storage_tag_map_t<Backend>, const In, 0> get_fill_values_
 }  // namespace filter
 
 template <typename Backend>
-class Filter : public SequenceOperator<Backend> {
+class Filter : public SequenceOperator<Backend, StatelessOperator> {
  public:
+  using Base = SequenceOperator<Backend, StatelessOperator>;
   inline explicit Filter(const OpSpec& spec)
-      : SequenceOperator<Backend>(spec),
+      : Base(spec),
         is_valid_mode_{filter::parse_is_valid_mode(spec.GetArgument<std::string>("mode"))} {
     spec.TryGetArgument(dtype_, "dtype");
   }
@@ -197,8 +199,7 @@ class Filter : public SequenceOperator<Backend> {
     // when there are no per-frame arguments, to reduce the number of instances of
     // per-sample data-structure when they are not needed.
     bool should_expand =
-        SequenceOperator<Backend>::ShouldExpand(ws) &&
-        (HasPerFramePositionalArgs(ws) || SequenceOperator<Backend>::HasPerFrameArgInputs(ws));
+        Base::ShouldExpand(ws) && (HasPerFramePositionalArgs(ws) || Base::HasPerFrameArgInputs(ws));
     if (should_expand && input_layout.size() && input_layout[0] == 'F') {
       assert(input_desc_.num_seq_dims >= 1);
       input_desc_.num_seq_dims--;

dali/operators/image/convolution/gaussian_blur.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2020-2023, 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.
@@ -153,7 +153,7 @@ bool GaussianBlur<CPUBackend>::ShouldExpand(const Workspace &ws) {
   const auto& input = ws.Input<CPUBackend>(0);
   auto layout = input.GetLayout();
   dim_desc_ = convolution_utils::ParseAndValidateDim(input.shape().sample_dim(), layout);
-  bool should_expand = SequenceOperator<CPUBackend>::ShouldExpand(ws);
+  bool should_expand = Base::ShouldExpand(ws);
   if (should_expand) {
     assert(dim_desc_.usable_axes_start > 0);
     dim_desc_.total_axes_count -= dim_desc_.usable_axes_start;

dali/operators/image/convolution/gaussian_blur.cu

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2020-2023, 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.
@@ -77,7 +77,7 @@ bool GaussianBlur<GPUBackend>::ShouldExpand(const Workspace &ws) {
   const auto& input = ws.Input<GPUBackend>(0);
   auto layout = input.GetLayout();
   dim_desc_ = convolution_utils::ParseAndValidateDim(input.shape().sample_dim(), layout);
-  bool should_expand = SequenceOperator<GPUBackend>::ShouldExpand(ws) && HasPerFrameArgInputs(ws);
+  bool should_expand = Base::ShouldExpand(ws) && HasPerFrameArgInputs(ws);
   if (should_expand) {
     assert(dim_desc_.usable_axes_start > 0);
     dim_desc_.total_axes_count -= dim_desc_.usable_axes_start;

dali/operators/image/convolution/gaussian_blur.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2020-2023, 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.
@@ -20,6 +20,7 @@
 
 #include "dali/operators/image/convolution/convolution_utils.h"
 #include "dali/operators/image/convolution/gaussian_blur_params.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/operator.h"
 #include "dali/pipeline/operator/sequence_operator.h"
@@ -38,10 +39,11 @@ namespace dali {
 #define GAUSSIAN_BLUR_SUPPORTED_AXES (1, 2, 3)
 
 template <typename Backend>
-class GaussianBlur : public SequenceOperator<Backend> {
+class GaussianBlur : public SequenceOperator<Backend, StatelessOperator> {
  public:
+  using Base = SequenceOperator<Backend, StatelessOperator>;
   inline explicit GaussianBlur(const OpSpec& spec)
-      : SequenceOperator<Backend>(spec) {
+      : Base(spec) {
     spec.TryGetArgument(dtype_, "dtype");
   }
 

dali/operators/image/convolution/laplacian.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2021-2023, 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.
@@ -191,7 +191,7 @@ bool Laplacian<CPUBackend>::ShouldExpand(const Workspace &ws) {
   const auto& input = ws.Input<CPUBackend>(0);
   auto layout = input.GetLayout();
   dim_desc_ = convolution_utils::ParseAndValidateDim(input.shape().sample_dim(), layout);
-  bool should_expand = SequenceOperator<CPUBackend>::ShouldExpand(ws);
+  bool should_expand = Base::ShouldExpand(ws);
   if (should_expand) {
     assert(dim_desc_.usable_axes_start > 0);
     dim_desc_.total_axes_count -= dim_desc_.usable_axes_start;

dali/operators/image/convolution/laplacian.cu

@@ -1,4 +1,4 @@
-// Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2022-2023, 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.
@@ -76,7 +76,7 @@ bool Laplacian<GPUBackend>::ShouldExpand(const Workspace &ws) {
   const auto& input = ws.Input<GPUBackend>(0);
   auto layout = input.GetLayout();
   dim_desc_ = convolution_utils::ParseAndValidateDim(input.shape().sample_dim(), layout);
-  bool should_expand = SequenceOperator<GPUBackend>::ShouldExpand(ws) && HasPerFrameArgInputs(ws);
+  bool should_expand = Base::ShouldExpand(ws) && HasPerFrameArgInputs(ws);
   if (should_expand) {
     assert(dim_desc_.usable_axes_start > 0);
     dim_desc_.total_axes_count -= dim_desc_.usable_axes_start;

dali/operators/image/convolution/laplacian.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2021-2023, 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.
@@ -20,6 +20,7 @@
 
 #include "dali/operators/image/convolution/convolution_utils.h"
 #include "dali/operators/image/convolution/laplacian_params.h"
+#include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/operator.h"
 #include "dali/pipeline/operator/sequence_operator.h"
@@ -39,10 +40,11 @@ namespace dali {
 
 
 template <typename Backend>
-class Laplacian : public SequenceOperator<Backend> {
+class Laplacian : public SequenceOperator<Backend, StatelessOperator> {
  public:
+  using Base = SequenceOperator<Backend, StatelessOperator>;
   inline explicit Laplacian(const OpSpec& spec)
-      : SequenceOperator<Backend>(spec) {
+      : Base(spec) {
     spec.TryGetArgument(dtype_, "dtype");
   }
 

dali/operators/image/remap/displacement_filter.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved.
+// Copyright (c) 2017-2023, 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.
@@ -35,6 +35,8 @@ struct HasParam <T, decltype((void) (typename T::Param()), 0)> : std::true_type
 
 class DisplacementIdentity {
  public:
+  // helper flag for checkpointing to select proper base class
+  static constexpr bool is_stateless = true;
   explicit DisplacementIdentity(const OpSpec& spec) {}
 
   DALI_HOST_DEV
@@ -47,10 +49,14 @@ class DisplacementIdentity {
   void Cleanup() {}
 };
 
+template <typename Backend, typename Displacement>
+using DisplacementBase =
+    std::conditional_t<Displacement::is_stateless, StatelessOperator<Backend>, Operator<Backend>>;
+
 template <typename Backend,
           class Displacement = DisplacementIdentity,
           bool per_channel_transform = false>
-class DisplacementFilter : public StatelessOperator<Backend> {};
+class DisplacementFilter : public DisplacementBase<Backend, Displacement> {};
 
 }  // namespace dali
 

dali/operators/image/remap/displacement_filter_impl_cpu.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-2023, 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.
@@ -64,10 +64,10 @@ void Warp(
 
 template <class Displacement, bool per_channel_transform>
 class DisplacementFilter<CPUBackend, Displacement, per_channel_transform>
-    : public Operator<CPUBackend> {
+    : public DisplacementBase<CPUBackend, Displacement> {
  public:
   explicit DisplacementFilter(const OpSpec &spec)
-      : Operator(spec),
+      : DisplacementBase<CPUBackend, Displacement>(spec),
         displace_(num_threads_, Displacement(spec)),
         interp_type_(spec.GetArgument<DALIInterpType>("interp_type")) {
     has_mask_ = spec.HasTensorArgument("mask");

dali/operators/image/remap/displacement_filter_impl_gpu.cuh

@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-2023, 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.
@@ -214,13 +214,12 @@ void DisplacementKernel_aligned32bit(
   }
 }
 
-template <class Displacement,
-          bool per_channel_transform>
-class DisplacementFilter<GPUBackend, Displacement,
-                         per_channel_transform> : public Operator<GPUBackend> {
+template <class Displacement, bool per_channel_transform>
+class DisplacementFilter<GPUBackend, Displacement, per_channel_transform>
+    : public DisplacementBase<GPUBackend, Displacement> {
  public:
   explicit DisplacementFilter(const OpSpec &spec) :
-      Operator(spec),
+      DisplacementBase<GPUBackend, Displacement>(spec),
       displace_(spec),
       interp_type_(spec.GetArgument<DALIInterpType>("interp_type")) {
     channel_block_setup_.SetBlockDim(ivec3{kAlignedBlockDim, 1, 1});