android-12.0.0_r34/s

/* Copyright 2020 The TensorFlow Authors. 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 "tensorflow/core/kernels/data/experimental/snapshot_util.h"

#include <queue>

#include "absl/memory/memory.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/kernels/data/name_utils.h"
#include "tensorflow/core/lib/io/buffered_inputstream.h"
#include "tensorflow/core/lib/io/random_inputstream.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/lib/io/snappy/snappy_inputbuffer.h"
#include "tensorflow/core/lib/io/snappy/snappy_outputbuffer.h"
#include "tensorflow/core/lib/io/zlib_compression_options.h"
#include "tensorflow/core/lib/io/zlib_inputstream.h"
#include "tensorflow/core/lib/io/zlib_outputbuffer.h"
#include "tensorflow/core/platform/coding.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/random.h"
#include "tensorflow/core/platform/stringprintf.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/protobuf/snapshot.pb.h"

namespace tensorflow {
namespace data {
namespace snapshot_util {

/* static */ constexpr const int64
    CustomReader::kSnappyReaderInputBufferSizeBytes;
/* static */ constexpr const int64
    CustomReader::kSnappyReaderOutputBufferSizeBytes;

std::string HashDirectory(const std::string& path, uint64 hash) {
  return io::JoinPath(
      path, strings::Printf("%llu", static_cast<unsigned long long>(hash)));
}

std::string RunDirectory(const std::string& hash_directory, uint64 run_id) {
  return RunDirectory(
      hash_directory,
      strings::Printf("%llu", static_cast<unsigned long long>(run_id)));
}

std::string RunDirectory(const std::string& hash_directory,
                         const std::string& run_id) {
  return io::JoinPath(hash_directory, run_id);
}

std::string ShardDirectory(const std::string& run_directory, int64 shard_id) {
  return io::JoinPath(
      run_directory,
      strings::Printf("%08llu%s", static_cast<unsigned long long>(shard_id),
                      kShardDirectorySuffix));
}
std::string GetCheckpointFileName(const std::string& shard_directory,
                                  uint64 checkpoint_id) {
  return io::JoinPath(
      shard_directory,
      strings::Printf("%08llu.snapshot",
                      static_cast<unsigned long long>(checkpoint_id)));
}

Status Writer::Create(Env* env, const std::string& filename,
                      const std::string& compression_type, int version,
                      const DataTypeVector& dtypes,
                      std::unique_ptr<Writer>* out_writer) {
  switch (version) {
    case 1:
      *out_writer =
          absl::make_unique<CustomWriter>(filename, compression_type, dtypes);
      break;
    case 2:
      *out_writer =
          absl::make_unique<TFRecordWriter>(filename, compression_type);
      break;
    default:
      return errors::InvalidArgument("Snapshot writer version: ", version,
                                     " is not supported.");
  }

  return (*out_writer)->Initialize(env);
}

TFRecordWriter::TFRecordWriter(const std::string& filename,
                               const std::string& compression_type)
    : filename_(filename), compression_type_(compression_type) {}

Status TFRecordWriter::Initialize(tensorflow::Env* env) {
  TF_RETURN_IF_ERROR(env->NewAppendableFile(filename_, &dest_));

  record_writer_ = absl::make_unique<io::RecordWriter>(
      dest_.get(), io::RecordWriterOptions::CreateRecordWriterOptions(
                       /*compression_type=*/compression_type_));
  return Status::OK();
}

Status TFRecordWriter::WriteTensors(const std::vector<Tensor>& tensors) {
  for (const auto& tensor : tensors) {
    TensorProto proto;
    tensor.AsProtoTensorContent(&proto);
#if defined(TF_CORD_SUPPORT)
    // Creating raw pointer here because std::move() in a releases in OSS TF
    // will result in a smart pointer being moved upon function creation, which
    // will result in proto_buffer == nullptr when WriteRecord happens.
    auto proto_buffer = new std::string();
    proto.SerializeToString(proto_buffer);
    absl::Cord proto_serialized = absl::MakeCordFromExternal(
        *proto_buffer,
        [proto_buffer](absl::string_view) { delete proto_buffer; });
    TF_RETURN_IF_ERROR(record_writer_->WriteRecord(proto_serialized));
#else   // TF_CORD_SUPPORT
    TF_RETURN_IF_ERROR(record_writer_->WriteRecord(proto.SerializeAsString()));
#endif  // TF_CORD_SUPPORT
  }
  return Status::OK();
}

Status TFRecordWriter::Sync() {
  TF_RETURN_IF_ERROR(record_writer_->Flush());
  return dest_->Flush();
}

Status TFRecordWriter::Close() {
  if (record_writer_ != nullptr) {
    TF_RETURN_IF_ERROR(Sync());
    TF_RETURN_IF_ERROR(record_writer_->Close());
    TF_RETURN_IF_ERROR(dest_->Close());
    record_writer_ = nullptr;
    dest_ = nullptr;
  }
  return Status::OK();
}

TFRecordWriter::~TFRecordWriter() {
  Status s = Close();
  if (!s.ok()) {
    LOG(ERROR) << "Failed to close snapshot file " << filename_ << ": " << s;
  }
}

CustomWriter::CustomWriter(const std::string& filename,
                           const std::string& compression_type,
                           const DataTypeVector& dtypes)
    : filename_(filename),
      compression_type_(compression_type),
      dtypes_(dtypes) {}

Status CustomWriter::Initialize(tensorflow::Env* env) {
  TF_RETURN_IF_ERROR(env->NewAppendableFile(filename_, &dest_));
#if defined(IS_SLIM_BUILD)
  if (compression_type_ != io::compression::kNone) {
    LOG(ERROR) << "Compression is unsupported on mobile platforms. Turning "
               << "off compression.";
  }
#else   // IS_SLIM_BUILD
  if (compression_type_ == io::compression::kGzip) {
    zlib_underlying_dest_.swap(dest_);
    io::ZlibCompressionOptions zlib_options;
    zlib_options = io::ZlibCompressionOptions::GZIP();

    io::ZlibOutputBuffer* zlib_output_buffer = new io::ZlibOutputBuffer(
        zlib_underlying_dest_.get(), zlib_options.input_buffer_size,
        zlib_options.output_buffer_size, zlib_options);
    TF_CHECK_OK(zlib_output_buffer->Init());
    dest_.reset(zlib_output_buffer);
  }
#endif  // IS_SLIM_BUILD
  simple_tensor_mask_.reserve(dtypes_.size());
  for (const auto& dtype : dtypes_) {
    if (DataTypeCanUseMemcpy(dtype)) {
      simple_tensor_mask_.push_back(true);
      num_simple_++;
    } else {
      simple_tensor_mask_.push_back(false);
      num_complex_++;
    }
  }

  return Status::OK();
}

Status CustomWriter::WriteTensors(const std::vector<Tensor>& tensors) {
  if (compression_type_ != io::compression::kSnappy) {
    experimental::SnapshotRecord record;
    for (const auto& tensor : tensors) {
      TensorProto* t = record.add_tensor();
      tensor.AsProtoTensorContent(t);
    }
#if defined(TF_CORD_SUPPORT)
    auto record_buffer = new std::string();
    record.SerializeToString(record_buffer);
    absl::Cord record_serialized = absl::MakeCordFromExternal(
        *record_buffer,
        [record_buffer](absl::string_view) { delete record_buffer; });
    return WriteRecord(record_serialized);
#else   // TF_CORD_SUPPORT
    return WriteRecord(record.SerializeAsString());
#endif  // TF_CORD_SUPPORT
  }

  std::vector<const TensorBuffer*> tensor_buffers;
  tensor_buffers.reserve(num_simple_);
  std::vector<TensorProto> tensor_protos;
  tensor_protos.reserve(num_complex_);
  experimental::SnapshotTensorMetadata metadata;
  int64 total_size = 0;
  for (int i = 0, end = tensors.size(); i < end; ++i) {
    const Tensor& tensor = tensors[i];
    experimental::TensorMetadata* tensor_metadata =
        metadata.add_tensor_metadata();
    tensor.shape().AsProto(tensor_metadata->mutable_tensor_shape());
    int64 size = 0;
    if (simple_tensor_mask_[i]) {
      auto tensor_buffer = DMAHelper::buffer(&tensor);
      tensor_buffers.push_back(tensor_buffer);
      size = tensor_buffer->size();
    } else {
      TensorProto proto;
      tensor.AsProtoTensorContent(&proto);
      size = proto.ByteSizeLong();
      tensor_protos.push_back(std::move(proto));
    }
    tensor_metadata->set_tensor_size_bytes(size);
    total_size += size;
  }

  std::vector<char> uncompressed(total_size);
  char* position = uncompressed.data();
  int buffer_index = 0;
  int proto_index = 0;
  for (int i = 0, end = tensors.size(); i < end; ++i) {
    const auto& tensor_metadata = metadata.tensor_metadata(i);
    if (simple_tensor_mask_[i]) {
      memcpy(position, tensor_buffers[buffer_index]->data(),
             tensor_metadata.tensor_size_bytes());
      buffer_index++;
    } else {
      tensor_protos[proto_index].SerializeToArray(
          position, tensor_metadata.tensor_size_bytes());
      proto_index++;
    }
    position += tensor_metadata.tensor_size_bytes();
  }
  DCHECK_EQ(position, uncompressed.data() + total_size);

  string output;
  if (!port::Snappy_Compress(uncompressed.data(), total_size, &output)) {
    return errors::Internal("Failed to compress using snappy.");
  }

#if defined(TF_CORD_SUPPORT)
  auto metadata_buffer = new std::string();
  metadata.SerializeToString(metadata_buffer);
  absl::Cord metadata_serialized = absl::MakeCordFromExternal(
      *metadata_buffer,
      [metadata_buffer](absl::string_view) { delete metadata_buffer; });
#else
  std::string metadata_serialized = metadata.SerializeAsString();
#endif  // TF_CORD_SUPPORT
  TF_RETURN_IF_ERROR(WriteRecord(metadata_serialized));
  TF_RETURN_IF_ERROR(WriteRecord(output));
  return Status::OK();
}

Status CustomWriter::Sync() { return dest_->Sync(); }

Status CustomWriter::Close() {
  if (dest_ != nullptr) {
    TF_RETURN_IF_ERROR(dest_->Close());
    dest_ = nullptr;
  }
  if (zlib_underlying_dest_ != nullptr) {
    TF_RETURN_IF_ERROR(zlib_underlying_dest_->Close());
    zlib_underlying_dest_ = nullptr;
  }
  return Status::OK();
}

CustomWriter::~CustomWriter() {
  Status s = Close();
  if (!s.ok()) {
    LOG(ERROR) << "Could not finish writing file: " << s;
  }
}

Status CustomWriter::WriteRecord(const StringPiece& data) {
  char header[kHeaderSize];
  core::EncodeFixed64(header, data.size());
  TF_RETURN_IF_ERROR(dest_->Append(StringPiece(header, sizeof(header))));
  return dest_->Append(data);
}

#if defined(TF_CORD_SUPPORT)
Status CustomWriter::WriteRecord(const absl::Cord& data) {
  char header[kHeaderSize];
  core::EncodeFixed64(header, data.size());
  TF_RETURN_IF_ERROR(dest_->Append(StringPiece(header, sizeof(header))));
  return dest_->Append(data);
}
#endif  // TF_CORD_SUPPORT

Status Reader::Create(Env* env, const std::string& filename,
                      const string& compression_type, int version,
                      const DataTypeVector& dtypes,
                      std::unique_ptr<Reader>* out_reader) {
  switch (version) {
    // CustomReader is able to read a legacy snapshot file format (v0) though
    // custom writer doesn't have the ability to write it any more since it is
    // strictly worse than V1.
    case 0:
    case 1:
      *out_reader = absl::make_unique<CustomReader>(filename, compression_type,
                                                    version, dtypes);
      break;
    case 2:
      *out_reader =
          absl::make_unique<TFRecordReader>(filename, compression_type, dtypes);
      break;
    default:
      return errors::InvalidArgument("Snapshot reader version: ", version,
                                     " is not supported.");
  }

  return (*out_reader)->Initialize(env);
}

Status Reader::SkipRecords(int64 num_records) {
  // TODO(frankchn): Optimize to not parse the entire Tensor and actually skip.
  for (int i = 0; i < num_records; ++i) {
    std::vector<Tensor> unused_tensors;
    TF_RETURN_IF_ERROR(ReadTensors(&unused_tensors));
  }
  return Status::OK();
}

class Reader::Dataset : public DatasetBase {
 public:
  explicit Dataset(const std::string& shard_dir, const std::string& compression,
                   const int64 version, const DataTypeVector& dtypes,
                   const std::vector<PartialTensorShape>& shapes,
                   const int64 start_index, DatasetContext::Params params)
      : DatasetBase(DatasetContext(std::move(params))),
        shard_dir_(shard_dir),
        compression_(compression),
        version_(version),
        dtypes_(dtypes),
        shapes_(shapes),
        start_index_(start_index) {}

  const DataTypeVector& output_dtypes() const override { return dtypes_; }

  const std::vector<PartialTensorShape>& output_shapes() const override {
    return shapes_;
  }

  std::string DebugString() const override {
    return "snapshot_util::Reader::Dataset";
  }

  Status InputDatasets(std::vector<const DatasetBase*>* inputs) const override {
    return Status::OK();
  }

  Status CheckExternalState() const override { return Status::OK(); }

 protected:
  Status AsGraphDefInternal(SerializationContext* ctx,
                            DatasetGraphDefBuilder* b,
                            Node** node) const override {
    // Not necessary perform any serialization as this dataset is only
    // constructed at runtime in C++ and will be reconstructed every time.
    return Status::OK();
  }

  std::unique_ptr<IteratorBase> MakeIteratorInternal(
      const string& prefix) const override {
    return absl::make_unique<Iterator>(Iterator::Params{
        this, name_utils::IteratorPrefix(node_name(), prefix)});
  }

 private:
  class Iterator : public DatasetIterator<Dataset> {
   public:
    explicit Iterator(const Params& params)
        : DatasetIterator<Dataset>(params), current_checkpoint_id_(0) {}

    Status Initialize(IteratorContext* ctx) override {
      TF_RETURN_IF_ERROR(Reader::Create(
          ctx->env(), GetCurrentFilename(), dataset()->compression_,
          dataset()->version_, dataset()->dtypes_, &reader_));
      bool end_of_sequence;
      for (int64 i = 0; i < dataset()->start_index_; ++i) {
        // TODO(frankchn): Optimize this to not parse every single element.
        std::vector<Tensor> unused;
        TF_RETURN_IF_ERROR(GetNextInternal(ctx, &unused, &end_of_sequence));
      }
      return Status::OK();
    }

   protected:
    Status GetNextInternal(IteratorContext* ctx,
                           std::vector<Tensor>* out_tensors,
                           bool* end_of_sequence) override {
      *end_of_sequence = false;
      Status s = reader_->ReadTensors(out_tensors);
      if (!errors::IsOutOfRange(s)) {
        return s;
      }
      Status status = AdvanceToNextFile(ctx->env());
      if (errors::IsNotFound(status)) {
        *end_of_sequence = true;
        return Status::OK();
      } else {
        return status;
      }
    }

    Status SaveInternal(SerializationContext* ctx,
                        IteratorStateWriter* writer) override {
      // Not necessary to save any state as this iterator will be reconstructed
      // from scratch when the parent snapshot dataset is restored from
      // checkpoint.
      return Status::OK();
    }

    Status RestoreInternal(IteratorContext* ctx,
                           IteratorStateReader* reader) override {
      // Not necessary to restore any state as this iterator will be
      // reconstructed from scratch when the parent snapshot dataset is restored
      // from checkpoint.
      return Status::OK();
    }

   private:
    std::string GetCurrentFilename() {
      return GetCheckpointFileName(dataset()->shard_dir_,
                                   current_checkpoint_id_);
    }

    Status AdvanceToNextFile(Env* env) {
      current_checkpoint_id_++;
      TF_RETURN_IF_ERROR(env->FileExists(GetCurrentFilename()));
      return Reader::Create(env, GetCurrentFilename(), dataset()->compression_,
                            dataset()->version_, dataset()->dtypes_, &reader_);
    }

    std::unique_ptr<Reader> reader_;

    // Stores the id current checkpoint file that we are in the process of
    // reading (e.g. if the file is currently 00000001.snapshot, then this will
    // be 1).
    uint64 current_checkpoint_id_;
  };

  const std::string shard_dir_;
  const std::string compression_;
  const int64 version_;
  const DataTypeVector dtypes_;
  const std::vector<PartialTensorShape> shapes_;
  const int64 start_index_;
};

class Reader::NestedDataset : public DatasetBase {
 public:
  explicit NestedDataset(std::vector<DatasetBase*> datasets,
                         DatasetContext::Params params)
      : DatasetBase(DatasetContext(std::move(params))), datasets_(datasets) {
    dtypes_.push_back(DT_VARIANT);
    gtl::InlinedVector<int64, 1> element_dim_sizes;
    element_dim_sizes.push_back(1);
    partial_shapes_.emplace_back(element_dim_sizes);
  }

  const DataTypeVector& output_dtypes() const override { return dtypes_; }

  const std::vector<PartialTensorShape>& output_shapes() const override {
    return partial_shapes_;
  }

  std::string DebugString() const override {
    return "snapshot_util::Reader::NestedDataset";
  }

  Status InputDatasets(std::vector<const DatasetBase*>* inputs) const override {
    inputs->clear();
    return Status::OK();
  }

  Status CheckExternalState() const override { return Status::OK(); }

 protected:
  Status AsGraphDefInternal(SerializationContext* ctx,
                            DatasetGraphDefBuilder* b,
                            Node** node) const override {
    // Not necessary perform any serialization as this dataset is only
    // constructed at runtime in C++ and will be reconstructed every time.
    return Status::OK();
  }

  std::unique_ptr<IteratorBase> MakeIteratorInternal(
      const string& prefix) const override {
    return absl::make_unique<Iterator>(Iterator::Params{
        this, name_utils::IteratorPrefix(node_name(), prefix)});
  }

 private:
  std::vector<DatasetBase*> datasets_;
  DataTypeVector dtypes_;
  std::vector<PartialTensorShape> partial_shapes_;

  class Iterator : public DatasetIterator<NestedDataset> {
   public:
    explicit Iterator(const Params& params)
        : DatasetIterator<NestedDataset>(params), index_(0) {}

   protected:
    Status GetNextInternal(IteratorContext* ctx,
                           std::vector<Tensor>* out_tensors,
                           bool* end_of_sequence) override {
      const int64 num_datasets = dataset()->datasets_.size();
      *end_of_sequence = num_datasets == index_;
      if (!*end_of_sequence) {
        Tensor tensor(DT_VARIANT, TensorShape({}));

        TF_RETURN_IF_ERROR(
            StoreDatasetInVariantTensor(dataset()->datasets_[index_], &tensor));
        out_tensors->clear();
        out_tensors->push_back(std::move(tensor));

        index_++;
      }
      return Status::OK();
    }

    Status SaveInternal(SerializationContext* ctx,
                        IteratorStateWriter* writer) override {
      // Not necessary to save any state as this iterator will be reconstructed
      // from scratch when the parent snapshot dataset is restored from
      // checkpoint.
      return Status::OK();
    }

    Status RestoreInternal(IteratorContext* ctx,
                           IteratorStateReader* reader) override {
      // Not necessary to restore any state as this iterator will be
      // reconstructed from scratch when the parent snapshot dataset is restored
      // from checkpoint.
      return Status::OK();
    }

   private:
    int64 index_;
  };
};

Status Reader::MakeNestedDataset(Env* env,
                                 const std::vector<std::string>& shard_dirs,
                                 const string& compression_type, int version,
                                 const DataTypeVector& dtypes,
                                 const std::vector<PartialTensorShape>& shapes,
                                 const int64 start_index,
                                 DatasetBase** output) {
  std::vector<DatasetBase*> datasets;

  datasets.reserve(shard_dirs.size());
  for (const auto& shard_dir : shard_dirs) {
    // TODO(frankchn): The reading pattern could be controlled in a non-round
    // robin fashion, so we cannot assume a round-robin manner when restoring.
    int64 dataset_start_index = start_index / shard_dirs.size();
    if (start_index % shard_dirs.size() > datasets.size()) {
      dataset_start_index++;
    }

    datasets.push_back(
        new Dataset(shard_dir, compression_type, version, dtypes, shapes,
                    dataset_start_index,
                    DatasetContext::Params({"snapshot_util::Reader::Dataset",
                                            "snapshot_util_reader_Dataset"})));
  }

  // Rotate the vector such that the first dataset contains the next element
  // to be produced, but not if there are no shards at all (then we just
  // construct an empty dataset).
  if (!shard_dirs.empty()) {
    std::rotate(datasets.begin(),
                datasets.begin() + (start_index % shard_dirs.size()),
                datasets.end());
  }

  *output = new NestedDataset(
      datasets, DatasetContext::Params({"snapshot_util::Reader::NestedDataset",
                                        "snapshot_util_reader_NestedDataset"}));
  return Status::OK();
}

TFRecordReader::TFRecordReader(const std::string& filename,
                               const string& compression_type,
                               const DataTypeVector& dtypes)
    : filename_(filename),
      offset_(0),
      compression_type_(compression_type),
      dtypes_(dtypes) {}

Status TFRecordReader::Initialize(Env* env) {
  TF_RETURN_IF_ERROR(env->NewRandomAccessFile(filename_, &file_));

  record_reader_ = absl::make_unique<io::RecordReader>(
      file_.get(), io::RecordReaderOptions::CreateRecordReaderOptions(
                       /*compression_type=*/compression_type_));
  return Status::OK();
}

Status TFRecordReader::ReadTensors(std::vector<Tensor>* read_tensors) {
  read_tensors->reserve(dtypes_.size());
  for (int i = 0; i < dtypes_.size(); ++i) {
    tstring record;
    TF_RETURN_IF_ERROR(record_reader_->ReadRecord(&offset_, &record));

    TensorProto proto;
    proto.ParseFromArray(record.data(), record.size());

    Tensor tensor;
    if (!tensor.FromProto(proto)) {
      return errors::DataLoss("Unable to parse tensor from stored proto.");
    }

    read_tensors->push_back(std::move(tensor));
  }
  return Status::OK();
}

CustomReader::CustomReader(const std::string& filename,
                           const string& compression_type, const int version,
                           const DataTypeVector& dtypes)
    : filename_(filename),
      compression_type_(compression_type),
      version_(version),
      dtypes_(dtypes) {}

Status CustomReader::Initialize(Env* env) {
  TF_RETURN_IF_ERROR(env->NewRandomAccessFile(filename_, &file_));
  input_stream_ = std::make_unique<io::RandomAccessInputStream>(file_.get());

#if defined(IS_SLIM_BUILD)
  if (compression_type_ != io::compression::kNone) {
    LOG(ERROR) << "Compression is unsupported on mobile platforms. Turning "
               << "off compression.";
  }
#else   // IS_SLIM_BUILD
  if (compression_type_ == io::compression::kGzip) {
    io::ZlibCompressionOptions zlib_options;
    zlib_options = io::ZlibCompressionOptions::GZIP();

    input_stream_ = absl::make_unique<io::ZlibInputStream>(
        input_stream_.release(), zlib_options.input_buffer_size,
        zlib_options.output_buffer_size, zlib_options, true);
  } else if (compression_type_ == io::compression::kSnappy) {
    if (version_ == 0) {
      input_stream_ = absl::make_unique<io::SnappyInputBuffer>(
          file_.get(), /*input_buffer_bytes=*/kSnappyReaderInputBufferSizeBytes,
          /*output_buffer_bytes=*/kSnappyReaderOutputBufferSizeBytes);
    } else {
      input_stream_ =
          absl::make_unique<io::BufferedInputStream>(file_.get(), 64 << 20);
    }
  }
#endif  // IS_SLIM_BUILD
  simple_tensor_mask_.reserve(dtypes_.size());
  for (const auto& dtype : dtypes_) {
    if (DataTypeCanUseMemcpy(dtype)) {
      simple_tensor_mask_.push_back(true);
      num_simple_++;
    } else {
      simple_tensor_mask_.push_back(false);
      num_complex_++;
    }
  }

  return Status::OK();
}

Status CustomReader::ReadTensors(std::vector<Tensor>* read_tensors) {
  profiler::TraceMe activity(
      [&]() { return absl::StrCat(kClassName, kSeparator, "ReadTensors"); },
      profiler::TraceMeLevel::kInfo);
  if (version_ == 0 || compression_type_ != io::compression::kSnappy) {
    return ReadTensorsV0(read_tensors);
  }
  if (version_ != 1) {
    return errors::InvalidArgument("Version: ", version_, " is not supported.");
  }
  if (compression_type_ != io::compression::kSnappy) {
    return errors::InvalidArgument("Compression ", compression_type_,
                                   " is not supported.");
  }

  experimental::SnapshotTensorMetadata metadata;
  tstring metadata_str;
  TF_RETURN_IF_ERROR(ReadRecord(&metadata_str));
  if (!metadata.ParseFromArray(metadata_str.data(), metadata_str.size())) {
    return errors::DataLoss("Could not parse SnapshotTensorMetadata");
  }
  read_tensors->reserve(metadata.tensor_metadata_size());

  std::vector<Tensor> simple_tensors;
  simple_tensors.reserve(num_simple_);
  std::vector<std::pair<std::unique_ptr<char[]>, size_t>> tensor_proto_strs;
  tensor_proto_strs.reserve(num_complex_);
  TF_RETURN_IF_ERROR(
      SnappyUncompress(&metadata, &simple_tensors, &tensor_proto_strs));

  int simple_index = 0;
  int complex_index = 0;
  for (int i = 0, end = simple_tensor_mask_.size(); i < end; ++i) {
    if (simple_tensor_mask_[i]) {
      read_tensors->push_back(std::move(simple_tensors[simple_index]));
      simple_index++;
    } else {
      auto tensor_proto_str = std::move(tensor_proto_strs[complex_index].first);
      size_t tensor_proto_size = tensor_proto_strs[complex_index].second;
      TensorProto tp;
      if (!tp.ParseFromArray(tensor_proto_str.get(), tensor_proto_size)) {
        return errors::Internal("Could not parse TensorProto");
      }
      Tensor t;
      if (!t.FromProto(tp)) {
        return errors::Internal("Could not parse Tensor");
      }
      read_tensors->push_back(std::move(t));
      complex_index++;
    }
  }
  return Status::OK();
}

Status CustomReader::ReadTensorsV0(std::vector<Tensor>* read_tensors) {
  experimental::SnapshotRecord record;
#if defined(PLATFORM_GOOGLE)
  absl::Cord c;
  TF_RETURN_IF_ERROR(ReadRecord(&c));
  record.ParseFromCord(c);
#else   // PLATFORM_GOOGLE
  tstring record_bytes;
  TF_RETURN_IF_ERROR(ReadRecord(&record_bytes));
  record.ParseFromArray(record_bytes.data(), record_bytes.size());
#endif  // PLATFORM_GOOGLE
  read_tensors->reserve(record.tensor_size());
  for (int i = 0; i < record.tensor_size(); ++i) {
    read_tensors->emplace_back();
    if (!read_tensors->back().FromProto(record.tensor(i))) {
      return errors::DataLoss("Unable to parse tensor from proto.");
    }
  }
  return Status::OK();
}

Status CustomReader::SnappyUncompress(
    const experimental::SnapshotTensorMetadata* metadata,
    std::vector<Tensor>* simple_tensors,
    std::vector<std::pair<std::unique_ptr<char[]>, size_t>>*
        tensor_proto_strs) {
  tstring compressed;
  TF_RETURN_IF_ERROR(ReadRecord(&compressed));
  size_t size;
  if (!port::Snappy_GetUncompressedLength(compressed.data(), compressed.size(),
                                          &size)) {
    return errors::Internal("Could not get snappy uncompressed length");
  }

  int num_tensors = metadata->tensor_metadata_size();
  std::vector<struct iovec> iov(num_tensors);
  int index = 0;
  int64 total_size = 0;
  for (int i = 0, end = simple_tensor_mask_.size(); i < end; ++i) {
    const auto& tensor_metadata = metadata->tensor_metadata(i);
    if (simple_tensor_mask_[i]) {
      TensorShape shape(tensor_metadata.tensor_shape());
      Tensor simple_tensor(dtypes_[i], shape);
      TensorBuffer* buffer = DMAHelper::buffer(&simple_tensor);
      iov[index].iov_base = buffer->data();
      iov[index].iov_len = buffer->size();
      simple_tensors->push_back(std::move(simple_tensor));
    } else {
      auto tensor_proto_str =
          absl::make_unique<char[]>(tensor_metadata.tensor_size_bytes());
      iov[index].iov_base = tensor_proto_str.get();
      iov[index].iov_len = tensor_metadata.tensor_size_bytes();
      tensor_proto_strs->push_back(std::make_pair(
          std::move(tensor_proto_str), tensor_metadata.tensor_size_bytes()));
    }
    total_size += iov[index].iov_len;
    index++;
  }
  const int64 size_int = size;
  if (size_int != total_size) {
    return errors::Internal("Uncompressed size mismatch. Snappy expects ", size,
                            " whereas the tensor metadata suggests ",
                            total_size);
  }
  if (!port::Snappy_UncompressToIOVec(compressed.data(), compressed.size(),
                                      iov.data(), num_tensors)) {
    return errors::Internal("Failed to perform snappy decompression.");
  }
  return Status::OK();
}

Status CustomReader::ReadRecord(tstring* record) {
  tstring header;
  TF_RETURN_IF_ERROR(input_stream_->ReadNBytes(kHeaderSize, &header));
  uint64 length = core::DecodeFixed64(header.data());
  return input_stream_->ReadNBytes(length, record);
}

#if defined(TF_CORD_SUPPORT)
Status CustomReader::ReadRecord(absl::Cord* record) {
  tstring header;
  TF_RETURN_IF_ERROR(input_stream_->ReadNBytes(kHeaderSize, &header));
  uint64 length = core::DecodeFixed64(header.data());
  if (compression_type_ == io::compression::kNone) {
    return input_stream_->ReadNBytes(length, record);
  } else {
    auto tmp_str = new tstring();
    TF_RETURN_IF_ERROR(input_stream_->ReadNBytes(length, tmp_str));
    absl::string_view tmp_str_view(*tmp_str);
    record->Append(absl::MakeCordFromExternal(
        tmp_str_view, [tmp_str](absl::string_view) { delete tmp_str; }));
    return Status::OK();
  }
}
#endif  // TF_CORD_SUPPORT

Status WriteMetadataFile(Env* env, const string& dir,
                         const experimental::SnapshotMetadataRecord* metadata) {
  string metadata_filename = io::JoinPath(dir, kMetadataFilename);
  TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(dir));
  std::string tmp_filename =
      absl::StrCat(metadata_filename, "-tmp-", random::New64());
  TF_RETURN_IF_ERROR(WriteBinaryProto(env, tmp_filename, *metadata));
  return env->RenameFile(tmp_filename, metadata_filename);
}

Status ReadMetadataFile(Env* env, const string& dir,
                        experimental::SnapshotMetadataRecord* metadata,
                        bool* file_exists) {
  string metadata_filename = io::JoinPath(dir, kMetadataFilename);
  Status s = env->FileExists(metadata_filename);
  *file_exists = s.ok();

  if (*file_exists) {
    return ReadBinaryProto(env, metadata_filename, metadata);
  } else {
    return Status::OK();
  }
}

Status DumpDatasetGraph(Env* env, const std::string& path, uint64 hash,
                        const GraphDef* graph) {
  std::string hash_hex =
      strings::StrCat(strings::Hex(hash, strings::kZeroPad16));
  std::string graph_file =
      io::JoinPath(path, absl::StrCat(hash_hex, "-graph.pbtxt"));

  LOG(INFO) << "Graph hash is " << hash_hex << ", writing to " << graph_file;
  TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(path));
  return WriteTextProto(env, graph_file, *graph);
}

Status DetermineOpState(const std::string& mode_string, bool file_exists,
                        const experimental::SnapshotMetadataRecord* metadata,
                        const uint64 pending_snapshot_expiry_seconds,
                        Mode* mode) {
  if (mode_string == kModeRead) {
    // In read mode, we should expect a metadata file is written.
    if (!file_exists) {
      return errors::NotFound("Metadata file does not exist.");
    }
    LOG(INFO) << "Overriding mode to reader.";
    *mode = READER;
    return Status::OK();
  }

  if (mode_string == kModeWrite) {
    LOG(INFO) << "Overriding mode to writer.";
    *mode = WRITER;
    return Status::OK();
  }

  if (mode_string == kModePassthrough) {
    LOG(INFO) << "Overriding mode to passthrough.";
    *mode = PASSTHROUGH;
    return Status::OK();
  }

  if (!file_exists) {
    *mode = WRITER;
    return Status::OK();
  }

  if (metadata->finalized()) {
    // File found, snapshot has been finalized.
    *mode = READER;
    return Status::OK();
  }

  int64 expiration_timer = static_cast<int64>(EnvTime::NowMicros()) -
                           pending_snapshot_expiry_seconds * 1000000;

  if (metadata->creation_timestamp() >= expiration_timer) {
    // Someone else is already writing and time has not expired.
    *mode = PASSTHROUGH;
    return Status::OK();
  } else {
    // Time has expired, we write regardless.
    *mode = WRITER;
    return Status::OK();
  }
}

AsyncWriter::AsyncWriter(Env* env, int64 file_index,
                         const std::string& shard_directory,
                         uint64 checkpoint_id, const std::string& compression,
                         int64 version, const DataTypeVector& output_types,
                         std::function<void(Status)> done) {
  thread_ = absl::WrapUnique(env->StartThread(
      ThreadOptions(), absl::StrCat("writer_thread_", file_index),
      [this, env, shard_directory, checkpoint_id, compression, version,
       &output_types, done = std::move(done)] {
        done(WriterThread(env, shard_directory, checkpoint_id, compression,
                          version, output_types));
      }));
}

void AsyncWriter::Write(const std::vector<Tensor>& tensors) {
  mutex_lock l(mu_);
  ElementOrEOF element;
  element.value = tensors;
  deque_.push_back(std::move(element));
}

void AsyncWriter::SignalEOF() {
  mutex_lock l(mu_);
  ElementOrEOF be;
  be.end_of_sequence = true;
  deque_.push_back(std::move(be));
}

void AsyncWriter::Consume(ElementOrEOF* be) {
  mutex_lock l(mu_);
  mu_.Await(tensorflow::Condition(this, &AsyncWriter::ElementAvailable));
  *be = deque_.front();
  deque_.pop_front();
}

bool AsyncWriter::ElementAvailable() { return !deque_.empty(); }

Status AsyncWriter::WriterThread(Env* env, const std::string& shard_directory,
                                 uint64 checkpoint_id,
                                 const std::string& compression, int64 version,
                                 DataTypeVector output_types) {
  std::unique_ptr<snapshot_util::Writer> writer;
  TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(shard_directory));

  TF_RETURN_IF_ERROR(snapshot_util::Writer::Create(
      env, GetCheckpointFileName(shard_directory, checkpoint_id), compression,
      version, std::move(output_types), &writer));

  while (true) {
    ElementOrEOF be;
    Consume(&be);

    if (be.end_of_sequence) {
      TF_RETURN_IF_ERROR(writer->Close());
      break;
    }

    TF_RETURN_IF_ERROR(writer->WriteTensors(be.value));
  }
  return Status::OK();
}

}  // namespace snapshot_util
}  // namespace data
}  // namespace tensorflow