Method ForEachIndexInternal

tensorflow/compiler/xla/shape_util.h:764–819

  template <typename FnType>
  static Status ForEachIndexInternal(const Shape& shape,
                                     absl::Span<const int64> base,
                                     absl::Span<const int64> count,
                                     absl::Span<const int64> incr,
                                     const FnType& visitor_function,
                                     bool parallel = false) {
    if (ShapeUtil::IsZeroElementArray(shape)) {
      return Status::OK();
    }
    CHECK_EQ(shape.rank(), base.size());
    CHECK_EQ(incr.size(), base.size());
    CHECK_EQ(count.size(), base.size());
    const int64 rank = LayoutUtil::MinorToMajor(shape).size();
    // Allows handling R0 arrays, such that the visitor function will be called
    // once with the proper empty indexes.
    int64 n = -1;
    std::vector<int64> indexes(base.begin(), base.end());
    const int kNumThreads = tensorflow::port::MaxParallelism();
    absl::optional<tensorflow::thread::ThreadPool> pool;
    if (parallel) {
      pool.emplace(tensorflow::Env::Default(), "foreach", kNumThreads);
    }

    tensorflow::mutex mu;
    Status status;  // Guarded by mu

    while (n < rank) {
      if (pool != absl::nullopt) {
        pool->Schedule([indexes, &visitor_function, &mu, &status] {
          StatusOr<bool> result = visitor_function(indexes);
          if (!result.ok()) {
            tensorflow::mutex_lock lock(mu);
            status = status.ok() ? result.status() : status;
          }
        });
      } else {
        TF_ASSIGN_OR_RETURN(bool should_continue, visitor_function(indexes));
        if (!should_continue) {
          break;
        }
      }
      // Increments dimensions in minor to major order.
      for (n = 0; n < rank; ++n) {
        int64 dim = LayoutUtil::Minor(shape.layout(), n);
        indexes[dim] += incr[dim];
        if (indexes[dim] < base[dim] + count[dim]) {
          break;
        }
        indexes[dim] = base[dim];
      }
    }

    // Waits for the scheduled work to complete.
    pool.reset();
    return status;
  }

  TF_DISALLOW_COPY_AND_ASSIGN(ShapeUtil);