Use forward-declaration of OpCheckpoint in operator.h - most ops are …

…stateless. Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
NVIDIA · Jun 11, 2024 · 5cbffaf · 5cbffaf
1 parent 32087da
commit 5cbffaf
Show file tree

Hide file tree

Showing 12 changed files with 17 additions and 7 deletions.
diff --git a/dali/operators/decoder/host/fused/host_decoder_random_crop.cc b/dali/operators/decoder/host/fused/host_decoder_random_crop.cc
@@ -21,6 +21,7 @@
 #include "dali/operators/decoder/host/fused/host_decoder_random_crop.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/checkpointing/snapshot_serializer.h"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 
 namespace dali {
 

diff --git a/dali/operators/decoder/nvjpeg/fused/nvjpeg_decoder_random_crop.cc b/dali/operators/decoder/nvjpeg/fused/nvjpeg_decoder_random_crop.cc
@@ -17,6 +17,7 @@
 #include "dali/operators/decoder/nvjpeg/fused/nvjpeg_decoder_random_crop.h"
 #include "dali/pipeline/operator/checkpointing/snapshot_serializer.h"
 #include "dali/util/random_crop_generator.h"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 
 namespace dali {
 

diff --git a/dali/operators/image/crop/random_crop_generator.h b/dali/operators/image/crop/random_crop_generator.h
@@ -28,6 +28,7 @@
 #include "dali/pipeline/operator/common.h"
 #include "dali/operators/image/crop/random_crop_attr.h"
 #include "dali/pipeline/operator/checkpointing/snapshot_serializer.h"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 
 namespace dali {
 

diff --git a/dali/operators/image/resize/random_resized_crop.h b/dali/operators/image/resize/random_resized_crop.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-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.
@@ -29,6 +29,7 @@
 #include "dali/operators/image/crop/random_crop_attr.h"
 #include "dali/kernels/imgproc/resample/params.h"
 #include "dali/pipeline/operator/checkpointing/snapshot_serializer.h"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 
 namespace dali {
 

diff --git a/dali/operators/imgcodec/roi_image_decoder.h b/dali/operators/imgcodec/roi_image_decoder.h
@@ -21,6 +21,7 @@
 #include "dali/operators/image/crop/random_crop_attr.h"
 #include "dali/operators/imgcodec/imgcodec.h"
 #include "dali/pipeline/operator/checkpointing/snapshot_serializer.h"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 #include "dali/pipeline/operator/common.h"
 #include "dali/pipeline/operator/operator.h"
 #include "dali/util/crop_window.h"

diff --git a/dali/operators/random/rng_base_cpu_test.cc b/dali/operators/random/rng_base_cpu_test.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2023-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.
@@ -61,7 +61,7 @@ class RNGCheckpointingTest : public ::testing::Test {
 
     // save and restore the pipeline
     auto node = original_pipe.GetOperatorNode("rng_op");
-    OpCheckpoint cpt(node->spec);
+    OpCheckpoint cpt("rng_op");
     node->op->SaveState(cpt, {});
     restored_pipe.GetOperatorNode("rng_op")->op->RestoreState(cpt);
 

diff --git a/dali/operators/random/rng_checkpointing_utils.h b/dali/operators/random/rng_checkpointing_utils.h
@@ -23,6 +23,7 @@
 #include "dali/pipeline/operator/checkpointing/snapshot_serializer.h"
 #include "dali/pipeline/util/batch_rng.h"
 #include "dali/operators/util/randomizer.cuh"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 
 namespace dali {
 namespace rng {

diff --git a/dali/operators/reader/reader_op.h b/dali/operators/reader/reader_op.h
@@ -29,6 +29,7 @@
 #include "dali/operators/reader/loader/loader.h"
 #include "dali/operators/reader/parser/parser.h"
 #include "dali/pipeline/operator/checkpointing/snapshot_serializer.h"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 #include "dali/pipeline/operator/name_utils.h"
 #include "dali/pipeline/operator/operator.h"
 

diff --git a/dali/pipeline/executor/lowered_graph.h b/dali/pipeline/executor/lowered_graph.h
@@ -23,6 +23,7 @@
 #include <string>
 #include <memory>
 #include <fstream>
+#include <optional>
 #include <set>
 
 #include "dali/core/common.h"

diff --git a/dali/pipeline/operator/operator.h b/dali/pipeline/operator/operator.h
@@ -34,14 +34,15 @@
 #include "dali/pipeline/operator/op_schema.h"
 #include "dali/pipeline/operator/op_spec.h"
 #include "dali/pipeline/operator/operator_factory.h"
-#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 #include "dali/pipeline/util/batch_utils.h"
 #include "dali/pipeline/workspace/workspace.h"
 #include "dali/pipeline/workspace/sample_workspace.h"
 #include "dali/pipeline/util/thread_pool.h"
 
 namespace dali {
 
+class OpCheckpoint;
+
 struct DLL_PUBLIC ReaderMeta {
   Index epoch_size = -1;          // raw epoch size
   Index epoch_size_padded = -1;   // epoch size with the padding at the end

diff --git a/dali/test/dali_test_checkpointing.h b/dali/test/dali_test_checkpointing.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2023-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.
@@ -68,7 +68,7 @@ class PipelineWrapper {
     clone.Build();
 
     for (const auto &[spec, id] : ops_) {
-      OpCheckpoint cpt(spec);
+      OpCheckpoint cpt(id);
       // TODO(mstaniewski): provide a stream, so operators with state kept
       // in device memory can be tested.
       GetOperator(id)->SaveState(cpt, {});

diff --git a/dali/test/operators/dummy_op.h b/dali/test/operators/dummy_op.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-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.
@@ -19,6 +19,7 @@
 #include <vector>
 
 #include "dali/pipeline/operator/operator.h"
+#include "dali/pipeline/operator/checkpointing/op_checkpoint.h"
 
 namespace dali {