From c9cd01cf305f04112138af4f6f8393d660d7c4f2 Mon Sep 17 00:00:00 2001 From: Laramie Leavitt Date: Tue, 5 Mar 2024 18:49:39 -0800 Subject: [PATCH] examples/extract_slice: translate input spec to 0. IWYU driver/downsample Cleanups while I was attempting to use a simple invocation of the extract_slice example to repro https://github.com/google/tensorstore/issues/140 PiperOrigin-RevId: 613040988 Change-Id: I53a6653bb9a777d5dcf9a3101ab64218dd34ad94 --- examples/BUILD | 3 + examples/extract_slice.cc | 19 ++++-- tensorstore/driver/downsample/BUILD | 58 ++++++++++++++++++- tensorstore/driver/downsample/downsample.cc | 43 ++++++++++++++ .../downsample/downsample_array_test.cc | 5 +- .../downsample/downsample_nditerable.cc | 14 +++-- .../driver/downsample/downsample_test.cc | 24 +++++++- .../driver/downsample/downsample_util.cc | 16 ++++- 8 files changed, 169 insertions(+), 13 deletions(-) diff --git a/examples/BUILD b/examples/BUILD index c2f8e0297..1a834a3cc 100644 --- a/examples/BUILD +++ b/examples/BUILD @@ -98,6 +98,7 @@ tensorstore_cc_binary( "//tensorstore:open", "//tensorstore:open_mode", "//tensorstore:spec", + "//tensorstore:strided_layout", "//tensorstore/index_space:dim_expression", "//tensorstore/index_space:index_transform", "//tensorstore/internal/image", @@ -106,6 +107,7 @@ tensorstore_cc_binary( "//tensorstore/internal/image:png", "//tensorstore/internal/image:webp", "//tensorstore/util:json_absl_flag", + "//tensorstore/util:result", "//tensorstore/util:span", "//tensorstore/util:status", "//tensorstore/util:str_cat", @@ -116,5 +118,6 @@ tensorstore_cc_binary( "@com_google_absl//absl/strings", "@com_google_riegeli//riegeli/bytes:fd_writer", "@com_google_riegeli//riegeli/bytes:std_io", + "@com_google_riegeli//riegeli/bytes:writer", ], ) diff --git a/examples/extract_slice.cc b/examples/extract_slice.cc index f96d3926c..3c0f45d1a 100644 --- a/examples/extract_slice.cc +++ b/examples/extract_slice.cc @@ -12,15 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Extracts a slice of a volumetric dataset, outputtting it as a 2d jpeg image. +// Extracts a slice of a volumetric dataset, outputting it as a 2d image. // // extract_slice --output_file=/tmp/foo.jpg --input_spec=... -#include -#include +#include + #include +#include #include #include +#include #include #include "absl/flags/flag.h" @@ -30,6 +32,7 @@ #include #include "riegeli/bytes/fd_writer.h" #include "riegeli/bytes/std_io.h" +#include "riegeli/bytes/writer.h" #include "tensorstore/array.h" #include "tensorstore/context.h" #include "tensorstore/data_type.h" @@ -46,8 +49,10 @@ #include "tensorstore/open.h" #include "tensorstore/open_mode.h" #include "tensorstore/spec.h" +#include "tensorstore/strided_layout.h" #include "tensorstore/tensorstore.h" #include "tensorstore/util/json_absl_flag.h" +#include "tensorstore/util/result.h" #include "tensorstore/util/span.h" #include "tensorstore/util/status.h" #include "tensorstore/util/str_cat.h" @@ -56,7 +61,6 @@ namespace { using ::tensorstore::Context; using ::tensorstore::Index; -using ::tensorstore::StrCat; using ::tensorstore::internal_image::AvifWriter; using ::tensorstore::internal_image::ImageInfo; using ::tensorstore::internal_image::ImageWriter; @@ -115,6 +119,10 @@ absl::Status Run(tensorstore::Spec input_spec, std::string output_filename) { std::cerr << std::endl << "Before: " << *transform << std::endl; + // DimRange(...).IndexSlice(0) transform below assumes that extra dimensions + // are 0-based; so make all dimensions 0-based here. + transform = transform | tensorstore::AllDims().TranslateTo(0); + // By convention, assume that the first dimension is Y, and the second is X, // and the third is C. The C++ api could use some help with labelling missing // dimensions, actually... @@ -126,6 +134,7 @@ absl::Status Run(tensorstore::Spec input_spec, std::string output_filename) { has_y = has_y || l == "y"; has_c = has_c || l == "c"; } + if (has_y) { transform = transform | tensorstore::Dims("y").MoveTo(0); } @@ -139,6 +148,8 @@ absl::Status Run(tensorstore::Spec input_spec, std::string output_filename) { transform = transform | tensorstore::DimRange(2, -1).IndexSlice(0); } + std::cerr << std::endl << "After: " << *transform << std::endl; + auto constrained_input = input | *transform; TENSORSTORE_RETURN_IF_ERROR(constrained_input); diff --git a/tensorstore/driver/downsample/BUILD b/tensorstore/driver/downsample/BUILD index 4f9c0cdcc..24afc5e08 100644 --- a/tensorstore/driver/downsample/BUILD +++ b/tensorstore/driver/downsample/BUILD @@ -30,23 +30,52 @@ tensorstore_cc_library( ":downsample_nditerable", ":downsample_util", ":grid_occupancy_map", + "//tensorstore:array", + "//tensorstore:array_storage_statistics", + "//tensorstore:box", + "//tensorstore:chunk_layout", + "//tensorstore:codec_spec", + "//tensorstore:context", + "//tensorstore:contiguous_layout", + "//tensorstore:data_type", "//tensorstore:downsample_method", + "//tensorstore:index", + "//tensorstore:json_serialization_options", + "//tensorstore:open_mode", + "//tensorstore:open_options", + "//tensorstore:rank", + "//tensorstore:resize_options", + "//tensorstore:schema", "//tensorstore:spec", + "//tensorstore:transaction", "//tensorstore/driver", + "//tensorstore/driver:chunk", "//tensorstore/index_space:dim_expression", + "//tensorstore/index_space:dimension_units", "//tensorstore/index_space:index_transform", + "//tensorstore/index_space:transformed_array", + "//tensorstore/internal:arena", "//tensorstore/internal:intrusive_ptr", + "//tensorstore/internal:lock_collection", "//tensorstore/internal:nditerable_transformed_array", "//tensorstore/internal/json_binding", + "//tensorstore/kvstore", "//tensorstore/serialization", "//tensorstore/util:executor", + "//tensorstore/util:future", + "//tensorstore/util:iterate", "//tensorstore/util:result", "//tensorstore/util:span", + "//tensorstore/util:status", + "//tensorstore/util:str_cat", + "//tensorstore/util/execution", "//tensorstore/util/execution:any_receiver", "//tensorstore/util/execution:sender_util", "//tensorstore/util/garbage_collection", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/status", + "@com_google_absl//absl/synchronization", ], alwayslink = True, ) @@ -82,11 +111,12 @@ tensorstore_cc_test( deps = [ ":downsample_array", "//tensorstore:array", + "//tensorstore:data_type", + "//tensorstore:downsample_method", "//tensorstore:index", "//tensorstore/index_space:dim_expression", "//tensorstore/index_space:transformed_array", "//tensorstore/util:span", - "//tensorstore/util:status_testutil", "@com_github_nlohmann_json//:nlohmann_json", "@com_google_googletest//:gtest_main", ], @@ -101,11 +131,11 @@ tensorstore_cc_library( "//tensorstore:data_type", "//tensorstore:downsample_method", "//tensorstore:index", + "//tensorstore:rank", "//tensorstore/internal:arena", "//tensorstore/internal:elementwise_function", "//tensorstore/internal:nditerable", "//tensorstore/internal:nditerable_buffer_management", - "//tensorstore/internal:nditerable_util", "//tensorstore/internal:unique_with_intrusive_allocator", "//tensorstore/util:extents", "//tensorstore/util:iterate", @@ -146,10 +176,19 @@ tensorstore_cc_test( srcs = ["downsample_test.cc"], deps = [ ":downsample", + "//tensorstore", + "//tensorstore:array", + "//tensorstore:chunk_layout", "//tensorstore:context", + "//tensorstore:data_type", "//tensorstore:downsample", + "//tensorstore:downsample_method", + "//tensorstore:index", "//tensorstore:open", + "//tensorstore:open_mode", + "//tensorstore:schema", "//tensorstore:spec", + "//tensorstore:static_cast", "//tensorstore/driver", "//tensorstore/driver:driver_testutil", "//tensorstore/driver/array", @@ -157,11 +196,18 @@ tensorstore_cc_test( "//tensorstore/driver/n5", "//tensorstore/driver/zarr", "//tensorstore/index_space:dim_expression", + "//tensorstore/index_space:index_transform", + "//tensorstore/index_space:transformed_array", "//tensorstore/internal:global_initializer", "//tensorstore/internal:json_gtest", "//tensorstore/kvstore/memory", "//tensorstore/util:status_testutil", + "//tensorstore/util:str_cat", + "//tensorstore/util:unit", + "//tensorstore/util/execution", "//tensorstore/util/execution:sender_util", + "@com_github_nlohmann_json//:nlohmann_json", + "@com_google_absl//absl/status", "@com_google_googletest//:gtest_main", ], ) @@ -176,12 +222,20 @@ tensorstore_cc_library( "//tensorstore:downsample_method", "//tensorstore:index", "//tensorstore:index_interval", + "//tensorstore:rank", "//tensorstore/index_space:index_transform", + "//tensorstore/index_space:output_index_method", + "//tensorstore/internal:integer_overflow", + "//tensorstore/util:byte_strided_pointer", + "//tensorstore/util:division", "//tensorstore/util:iterate", "//tensorstore/util:result", "//tensorstore/util:span", + "//tensorstore/util:status", + "//tensorstore/util:str_cat", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/status", "@com_google_absl//absl/strings", ], ) diff --git a/tensorstore/driver/downsample/downsample.cc b/tensorstore/driver/downsample/downsample.cc index 412d2e8a1..ae1d404da 100644 --- a/tensorstore/driver/downsample/downsample.cc +++ b/tensorstore/driver/downsample/downsample.cc @@ -14,10 +14,28 @@ #include "tensorstore/driver/downsample/downsample.h" +#include + +#include +#include #include +#include +#include #include "absl/base/thread_annotations.h" +#include "absl/container/inlined_vector.h" #include "absl/status/status.h" +#include "absl/synchronization/mutex.h" +#include "tensorstore/array.h" +#include "tensorstore/array_storage_statistics.h" +#include "tensorstore/box.h" +#include "tensorstore/chunk_layout.h" +#include "tensorstore/codec_spec.h" +#include "tensorstore/context.h" +#include "tensorstore/contiguous_layout.h" +#include "tensorstore/data_type.h" +#include "tensorstore/downsample_method.h" +#include "tensorstore/driver/chunk.h" #include "tensorstore/driver/downsample/downsample_array.h" #include "tensorstore/driver/downsample/downsample_method_json_binder.h" #include "tensorstore/driver/downsample/downsample_nditerable.h" @@ -25,20 +43,45 @@ #include "tensorstore/driver/downsample/grid_occupancy_map.h" #include "tensorstore/driver/driver.h" #include "tensorstore/driver/driver_handle.h" +#include "tensorstore/driver/driver_spec.h" #include "tensorstore/driver/read.h" #include "tensorstore/driver/registry.h" +#include "tensorstore/index.h" #include "tensorstore/index_space/dim_expression.h" +#include "tensorstore/index_space/dimension_units.h" +#include "tensorstore/index_space/index_domain.h" #include "tensorstore/index_space/index_domain_builder.h" +#include "tensorstore/index_space/index_transform.h" #include "tensorstore/index_space/index_transform_builder.h" +#include "tensorstore/index_space/transformed_array.h" +#include "tensorstore/internal/arena.h" #include "tensorstore/internal/intrusive_ptr.h" +#include "tensorstore/internal/json_binding/json_binding.h" #include "tensorstore/internal/json_binding/std_array.h" +#include "tensorstore/internal/lock_collection.h" #include "tensorstore/internal/nditerable_transformed_array.h" +#include "tensorstore/json_serialization_options.h" +#include "tensorstore/kvstore/kvstore.h" +#include "tensorstore/kvstore/spec.h" +#include "tensorstore/open_mode.h" +#include "tensorstore/open_options.h" +#include "tensorstore/rank.h" +#include "tensorstore/resize_options.h" +#include "tensorstore/schema.h" #include "tensorstore/serialization/std_vector.h" // IWYU pragma: keep #include "tensorstore/spec.h" +#include "tensorstore/transaction.h" #include "tensorstore/util/execution/any_receiver.h" +#include "tensorstore/util/execution/execution.h" #include "tensorstore/util/execution/sender_util.h" #include "tensorstore/util/executor.h" +#include "tensorstore/util/future.h" #include "tensorstore/util/garbage_collection/std_vector.h" // IWYU pragma: keep +#include "tensorstore/util/iterate.h" +#include "tensorstore/util/result.h" +#include "tensorstore/util/span.h" +#include "tensorstore/util/status.h" +#include "tensorstore/util/str_cat.h" namespace tensorstore { namespace internal_downsample { diff --git a/tensorstore/driver/downsample/downsample_array_test.cc b/tensorstore/driver/downsample/downsample_array_test.cc index 40805501e..cfb0f3a76 100644 --- a/tensorstore/driver/downsample/downsample_array_test.cc +++ b/tensorstore/driver/downsample/downsample_array_test.cc @@ -14,15 +14,18 @@ #include "tensorstore/driver/downsample/downsample_array.h" +#include + #include #include #include #include "tensorstore/array.h" +#include "tensorstore/data_type.h" +#include "tensorstore/downsample_method.h" #include "tensorstore/index.h" #include "tensorstore/index_space/dim_expression.h" #include "tensorstore/index_space/transformed_array.h" #include "tensorstore/util/span.h" -#include "tensorstore/util/status_testutil.h" namespace { diff --git a/tensorstore/driver/downsample/downsample_nditerable.cc b/tensorstore/driver/downsample/downsample_nditerable.cc index 204f8862d..859be0288 100644 --- a/tensorstore/driver/downsample/downsample_nditerable.cc +++ b/tensorstore/driver/downsample/downsample_nditerable.cc @@ -14,11 +14,18 @@ #include "tensorstore/driver/downsample/downsample_nditerable.h" +#include +#include + #include #include #include +#include +#include +#include #include #include +#include #include #include "absl/log/absl_log.h" @@ -33,8 +40,8 @@ #include "tensorstore/internal/elementwise_function.h" #include "tensorstore/internal/nditerable.h" #include "tensorstore/internal/nditerable_buffer_management.h" -#include "tensorstore/internal/nditerable_util.h" #include "tensorstore/internal/unique_with_intrusive_allocator.h" +#include "tensorstore/rank.h" #include "tensorstore/util/extents.h" #include "tensorstore/util/iterate.h" #include "tensorstore/util/span.h" @@ -49,7 +56,6 @@ namespace tensorstore { namespace internal_downsample { - namespace { using ::tensorstore::internal::ArenaAllocator; @@ -1005,7 +1011,7 @@ class DownsampledNDIterator : public NDIterator::Base { Index base_inclusive_min = initial_base_index * downsample_factor - downsample_dim_origin[i]; if (base_dim >= base_iteration_rank - 2) { - // This downsampling factor applies to one of the 2 inner-most + // This downsampling factor applies to one of the 2 innermost // dimensions. const DimensionIndex inner_dim_i = base_dim - (base_iteration_rank - 2); --num_outer_downsample_dims; @@ -1238,7 +1244,7 @@ class DownsampledNDIterable : public NDIterable::Base { DimensionIndex base_iteration_dimensions_[kMaxRank]; }; - std::ptrdiff_t GetWorkingMemoryBytesPerElement( + ptrdiff_t GetWorkingMemoryBytesPerElement( NDIterable::IterationLayoutView layout, IterationBufferKind buffer_kind) const override { NDIterable::IterationLayoutView base_layout; diff --git a/tensorstore/driver/downsample/downsample_test.cc b/tensorstore/driver/downsample/downsample_test.cc index faa3c3ae0..ae5ee0709 100644 --- a/tensorstore/driver/downsample/downsample_test.cc +++ b/tensorstore/driver/downsample/downsample_test.cc @@ -14,19 +14,41 @@ #include "tensorstore/downsample.h" +#include + +#include +#include +#include + #include #include +#include "absl/status/status.h" +#include +#include "tensorstore/array.h" +#include "tensorstore/chunk_layout.h" #include "tensorstore/context.h" +#include "tensorstore/data_type.h" +#include "tensorstore/downsample_method.h" #include "tensorstore/driver/array/array.h" -#include "tensorstore/driver/driver.h" #include "tensorstore/driver/driver_testutil.h" +#include "tensorstore/driver/read.h" +#include "tensorstore/index.h" #include "tensorstore/index_space/dim_expression.h" +#include "tensorstore/index_space/index_transform.h" +#include "tensorstore/index_space/transformed_array.h" #include "tensorstore/internal/global_initializer.h" #include "tensorstore/internal/json_gtest.h" #include "tensorstore/open.h" +#include "tensorstore/open_mode.h" +#include "tensorstore/schema.h" #include "tensorstore/spec.h" +#include "tensorstore/static_cast.h" +#include "tensorstore/tensorstore.h" +#include "tensorstore/util/execution/execution.h" #include "tensorstore/util/execution/sender_util.h" #include "tensorstore/util/status_testutil.h" +#include "tensorstore/util/str_cat.h" +#include "tensorstore/util/unit.h" namespace { diff --git a/tensorstore/driver/downsample/downsample_util.cc b/tensorstore/driver/downsample/downsample_util.cc index 01f9c12c7..8f68adaf5 100644 --- a/tensorstore/driver/downsample/downsample_util.cc +++ b/tensorstore/driver/downsample/downsample_util.cc @@ -15,21 +15,35 @@ #include "tensorstore/driver/downsample/downsample_util.h" #include -#include +#include +#include +#include +#include +#include #include "absl/base/optimization.h" #include "absl/container/inlined_vector.h" +#include "absl/status/status.h" #include "absl/strings/str_join.h" #include "tensorstore/array.h" #include "tensorstore/box.h" +#include "tensorstore/downsample_method.h" #include "tensorstore/index.h" #include "tensorstore/index_interval.h" +#include "tensorstore/index_space/index_domain.h" #include "tensorstore/index_space/index_transform.h" #include "tensorstore/index_space/internal/identity_transform.h" #include "tensorstore/index_space/internal/transform_rep.h" +#include "tensorstore/index_space/output_index_method.h" +#include "tensorstore/internal/integer_overflow.h" +#include "tensorstore/rank.h" +#include "tensorstore/util/byte_strided_pointer.h" +#include "tensorstore/util/division.h" #include "tensorstore/util/iterate.h" #include "tensorstore/util/result.h" #include "tensorstore/util/span.h" +#include "tensorstore/util/status.h" +#include "tensorstore/util/str_cat.h" namespace tensorstore { namespace internal_downsample {