xref: /aosp_15_r20/external/tensorflow/tensorflow/core/framework/tensor_slice.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)

/* Copyright 2015 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/framework/tensor_slice.h"

#include <limits>
#include <vector>

#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/logging.h"

namespace tensorflow {

TensorSlice::TensorSlice(int dim) { SetFullSlice(dim); }

TensorSlice::TensorSlice(const TensorSliceProto& proto) {
  starts_.reserve(proto.extent_size());
  lengths_.reserve(proto.extent_size());
  for (const auto& e : proto.extent()) {
    starts_.push_back(e.start());
    lengths_.push_back(GetExtentLength(e));
  }
}

TensorSlice::TensorSlice(
    std::initializer_list<std::pair<int64_t, int64_t>> extents) {
  starts_.reserve(extents.size());
  lengths_.reserve(extents.size());
  for (const auto& e : extents) {
    starts_.push_back(e.first);
    lengths_.push_back(e.second);
  }
}

Status TensorSlice::BuildTensorSlice(const TensorSliceProto& proto,
                                     TensorSlice* output) {
  output->Clear();
  output->starts_.reserve(proto.extent_size());
  output->lengths_.reserve(proto.extent_size());
  for (const auto& e : proto.extent()) {
    int64_t l = GetExtentLength(e);
    if (e.start() != 0 || l != kFullExtent) {
      if (e.start() < 0 || l <= 0) {
        return errors::InvalidArgument(
            "Expected non-negative start and positive length but got start = ",
            e.start(), ", length = ", l, ": extent = ", e.ShortDebugString());
      }
      // Calculating the extent end must not cause signed integer overflow.
      if (static_cast<uint64_t>(e.start()) + static_cast<uint64_t>(e.length()) >
          std::numeric_limits<int64_t>::max()) {
        return errors::InvalidArgument(
            "Extent end exceeds the maximum possible size: extent = ",
            e.ShortDebugString());
      }
    }
    output->starts_.push_back(e.start());
    output->lengths_.push_back(l);
  }

  return OkStatus();
}

Status TensorSlice::Parse(const string& str, TensorSlice* slice) {
  std::vector<string> items = str_util::Split(str, ':', str_util::SkipEmpty());
  slice->starts_.reserve(items.size());
  slice->lengths_.reserve(items.size());
  for (const string& x : items) {
    int64_t s, l;
    if (x == "-") {
      // "everything"
      s = 0;
      l = kFullExtent;
    } else {
      std::vector<string> sl = str_util::Split(x, ',', str_util::SkipEmpty());
      if (sl.size() != 2 || !strings::safe_strto64(sl[0], &s) ||
          !strings::safe_strto64(sl[1], &l)) {
        return errors::InvalidArgument(
            "Expected a pair of numbers or '-' "
            "but got '",
            x, "': string = ", str);
      }
      if (s < 0 || l <= 0) {
        return errors::InvalidArgument(
            "Expected non-negative start and "
            "positive length but got start = ",
            s, ", length = ", l, ": string = ", str);
      }
    }
    slice->starts_.push_back(s);
    slice->lengths_.push_back(l);
  }

  return OkStatus();
}

void TensorSlice::Clear() {
  starts_.clear();
  lengths_.clear();
}

bool TensorSlice::IsFull() const {
  for (int d = 0; d < dims(); ++d) {
    if (!IsFullAt(d)) return false;
  }
  return true;
}

void TensorSlice::SetFullSlice(int dim) {
  Clear();
  starts_.reserve(dim);
  lengths_.reserve(dim);
  for (int d = 0; d < dim; ++d) {
    starts_.push_back(0);
    lengths_.push_back(kFullExtent);
  }
}

void TensorSlice::Extend(int dim) {
  int old_dim = dims();
  DCHECK_LE(old_dim, dim);
  starts_.resize(dim);
  lengths_.resize(dim);
  for (int d = old_dim; d < dim; ++d) {
    starts_[d] = 0;
    lengths_[d] = kFullExtent;
  }
}

void TensorSlice::AsProto(TensorSliceProto* proto) const {
  for (int d = 0; d < dims(); ++d) {
    TensorSliceProto::Extent* e = proto->add_extent();
    // We only need to record the explicit slice for non-full slices
    if (!IsFullAt(d)) {
      e->set_start(starts_[d]);
      e->set_length(lengths_[d]);
    }
  }
}

string TensorSlice::DebugString() const {
  string buffer;
  bool first = true;
  for (int d = 0; d < dims(); ++d) {
    if (!first) {
      buffer.append(":");
    }
    if (IsFullAt(d)) {
      buffer.append("-");
    } else {
      strings::StrAppend(&buffer, starts_[d], ",", lengths_[d]);
    }
    first = false;
  }
  return buffer;
}

bool TensorSlice::Intersect(const TensorSlice& other,
                            TensorSlice* result) const {
  // First, if two slices have different ranks, they obviously don't overlap
  // -- in fact they are not compatible.
  if (dims() != other.dims()) {
    return false;
  }

  // Setting the result to the right dimension
  if (result) {
    result->SetFullSlice(dims());
  }
  // The two slices overlap if they overlap in all dimensions.
  for (int d = 0; d < dims(); ++d) {
    if (IsFullAt(d)) {
      if (result) {
        result->set_start(d, other.start(d));
        result->set_length(d, other.length(d));
      }
    } else if (other.IsFullAt(d)) {
      if (result) {
        result->set_start(d, start(d));
        result->set_length(d, length(d));
      }
    } else {
      // If we have an intersection here, it should have a start that is the
      // max of the two starts and an end that is the min of the two ends.
      int64_t s = std::max(start(d), other.start(d));
      int64_t l = std::min(end(d), other.end(d)) - s;
      if (l > 0) {
        // We have a real intersection
        if (result) {
          result->set_start(d, s);
          result->set_length(d, l);
        }
      } else {
        // We don't have an intersection for this dimension -- thus we don't
        // have any intersection at all.
        if (result) {
          result->Clear();
        }
        return false;
      }
    }
  }
  // If we are here, we know there is overlap in every dimension.
  return true;
}

bool TensorSlice::operator==(const TensorSlice& other) const {
  return dims() == other.dims() && starts_ == other.starts_ &&
         lengths_ == other.lengths_;
}

void TensorSlice::ComputeRelative(const TensorSlice& sub,
                                  TensorSlice* relative) const {
  DCHECK_EQ(dims(), sub.dims());
  relative->SetFullSlice(dims());
  for (int d = 0; d < dims(); ++d) {
    if (IsFullAt(d)) {
      relative->set_start(d, sub.start(d));
      relative->set_length(d, sub.length(d));
    } else {
      // Otherwise the relative start is the difference between the start of
      // sub and the start of base
      relative->set_start(d, sub.start(d) - start(d));
      relative->set_length(d, sub.length(d));
    }
  }
}

void TensorSlice::UpdateToCover(const TensorSlice& other) {
  DCHECK_EQ(dims(), other.dims());
  for (int d = 0; d < dims(); ++d) {
    if (!IsFullAt(d)) {
      if (other.IsFullAt(d)) {
        starts_[d] = 0;
        lengths_[d] = kFullExtent;
      } else {
        const auto new_end = std::max(end(d), other.end(d));
        set_start(d, std::min(start(d), other.start(d)));
        set_length(d, new_end - start(d));
      }
    }
  }
}

// static
bool TensorSlice::HasExtentLength(const TensorSliceProto::Extent& extent) {
  return extent.has_length_case() == TensorSliceProto::Extent::kLength;
}

// static
int64_t TensorSlice::GetExtentLength(const TensorSliceProto::Extent& extent) {
  if (!HasExtentLength(extent)) return -1;
  return extent.length();
}

Status TensorSlice::SliceTensorShape(const TensorShape& shape,
                                     TensorShape* result_shape) const {
  result_shape->Clear();
  // Mismatching ranks: we can't apply the slice at all.
  if (shape.dims() != dims()) {
    return errors::Internal("Mismatching ranks: shape = ", shape.DebugString(),
                            ", slice = ", DebugString());
  }
  for (int d = 0; d < dims(); ++d) {
    if (IsFullAt(d)) {
      result_shape->AddDim(shape.dim_size(d));
    } else {
      // Check if the extent applies to the dimension
      if (end(d) <= shape.dim_size(d)) {
        // Yes: the end is within the range of the dim -- we adjust the result
        // shape so that its size along this dimension is the length of the
        // slice.
        result_shape->AddDim(length(d));
      } else {
        // The extent doesn't apply to the dimension
        result_shape->Clear();
        return errors::Internal("Extent in dimension ", d,
                                " out of bounds: shape = ", shape.DebugString(),
                                ", slice = ", DebugString());
      }
    }
  }
  // If we are here, we have successfully applied the shape.
  return OkStatus();
}

const int64_t TensorSlice::kFullExtent = -1;

}  // namespace tensorflow