diff options
Diffstat (limited to 'tensorflow/compiler/xla/array4d.h')
| -rw-r--r-- | tensorflow/compiler/xla/array4d.h | 272 |
1 files changed, 272 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/array4d.h b/tensorflow/compiler/xla/array4d.h new file mode 100644 index 0000000000..db51a57cf2 --- /dev/null +++ b/tensorflow/compiler/xla/array4d.h @@ -0,0 +1,272 @@ +/* Copyright 2017 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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_ARRAY4D_H_ +#define TENSORFLOW_COMPILER_XLA_ARRAY4D_H_ + +#include <algorithm> +#include <functional> +#include <initializer_list> +#include <iterator> +#include <numeric> +#include <random> +#include <string> +#include <vector> + +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// Simple 4D array structure, similar in form to Array2D, for use primarily in +// testing and describing to XLA APIs values in the 4D array structures used +// in convolutions. +// +// The data layout is, in order from major to minor: +// +// First dimension: plane, batch, n1 +// Second dimension: depth, feature, z, n2 +// Third dimension: height, y, n3 +// Fourth dimension: width, x, n4 +// +// These dimensions are referred to by various names, so that is why +// more than one name is given above. See operator() for the exact +// calculation of 1d indices from 4d indices. +template <typename T> +class Array4D { + public: + // Creates a 4D array, unitialized values. + Array4D(int64 planes, int64 depth, int64 height, int64 width) + : planes_(planes), depth_(depth), height_(height), width_(width) { + values_.resize(planes * depth * height * width); + } + + // Creates a 4D array, initalized to value. + Array4D(int64 planes, int64 depth, int64 height, int64 width, T value) + : Array4D(planes, depth, height, width) { + Fill(value); + } + + // Creates a 4D array, filled with values. + // + // We need to set a default type for Container so that code like + // Array4D(1, 1, 1, 1, {1}) will work. The template cannot infer the + // initializer_list type in that case without this default. + template <typename Container = std::initializer_list<T>> + Array4D(int64 planes, int64 depth, int64 height, int64 width, + const Container& values) + : Array4D(planes, depth, height, width) { + SetValues(values); + } + + // Construct an Array4D with the given nested initializer list. + Array4D(std::initializer_list<std::initializer_list< + std::initializer_list<std::initializer_list<T>>>> + values) + : Array4D(values.size(), values.begin()->size(), + values.begin()->begin()->size(), + values.begin()->begin()->begin()->size()) { + int64 plane = 0; + for (const auto values_in_plane : values) { + DCHECK_EQ(values_in_plane.size(), depth_); + int64 depth = 0; + for (const auto values_in_depth : values_in_plane) { + DCHECK_EQ(values_in_depth.size(), height_); + int64 height = 0; + for (const auto values_in_height : values_in_depth) { + DCHECK_EQ(values_in_height.size(), width_); + int64 width = 0; + for (const auto element_value : values_in_height) { + (*this)(plane, depth, height, width) = element_value; + ++width; + } + ++height; + } + ++depth; + } + ++plane; + } + } + + T& operator()(int64 plane, int64 depth, int64 height, int64 width) { + CHECK_LT(plane, planes_); + CHECK_LT(depth, depth_); + CHECK_LT(height, height_); + CHECK_LT(width, width_); + return values_[plane * (depth_ * height_ * width_) + + depth * (height_ * width_) + height * (width_) + width]; + } + const T& operator()(int64 plane, int64 depth, int64 height, + int64 width) const { + return const_cast<Array4D*>(this)->operator()(plane, depth, height, width); + } + + int64 width() const { return width_; } + int64 height() const { return height_; } + int64 depth() const { return depth_; } + int64 planes() const { return planes_; } + + // Numerically-named aliases for the various dimensions. This matches the + // dimension names used in array3d. + int64 n4() const { return width_; } + int64 n3() const { return height_; } + int64 n2() const { return depth_; } + int64 n1() const { return planes_; } + int64 num_elements() const { return values_.size(); } + + // Sets all the values in the array to values. + template <typename Container = std::initializer_list<T>> + void SetValues(const Container& container) { + CHECK_EQ(std::distance(std::begin(container), std::end(container)), + num_elements()); + values_.assign(std::begin(container), std::end(container)); + } + + // Fills the array with the given value. + void Fill(const T& value) { + std::fill(values_.begin(), values_.end(), value); + } + + // Fills the array with iota. + void FillIota(const T& value) { + std::iota(values_.begin(), values_.end(), value); + } + + // Fills the array with random variable with a deviation of value and a mean + // of mean. + void FillRandom(const T& value, const double mean = 0.0, + const int seed = 12345) { + std::mt19937 g(seed); + std::normal_distribution<double> distribution(mean, + static_cast<double>(value)); + for (auto& v : values_) { + v = static_cast<T>(distribution(g)); + } + } + + // Fills values with the sequence i*multiplier for i=0,1,... + void FillWithMultiples(float multiplier) { + for (int64 i = 0; i < num_elements(); ++i) { + values_[i] = i * multiplier; + } + } + + // Invokes a callback with the (indices, value_ptr) for each cell in the 4D + // array. + void Each(std::function<void(tensorflow::gtl::ArraySlice<int64>, T*)> f) { + for (int64 plane = 0; plane < planes(); ++plane) { + for (int64 depth = 0; depth < this->depth(); ++depth) { + for (int64 height = 0; height < this->height(); ++height) { + for (int64 width = 0; width < this->width(); ++width) { + auto& value = (*this)(plane, depth, height, width); + f({plane, depth, height, width}, &value); + } + } + } + } + } + + // Fills all of the {p,z} with the array provided, which specifies {y,x}. + void FillWithYX(const Array2D<T>& value) { + CHECK_EQ(value.height(), height()); + CHECK_EQ(value.width(), width()); + for (int64 plane = 0; plane < planes(); ++plane) { + for (int64 depth = 0; depth < this->depth(); ++depth) { + for (int64 height = 0; height < this->height(); ++height) { + for (int64 width = 0; width < this->width(); ++width) { + (*this)(plane, depth, height, width) = value(height, width); + } + } + } + } + } + + // Fills all of the {x,y} with the array provided, which specifies {p,z}. + void FillWithPZ(const Array2D<T>& value) { + CHECK_EQ(value.height(), planes()); + CHECK_EQ(value.width(), depth()); + for (int64 height = 0; height < this->height(); ++height) { + for (int64 width = 0; width < this->width(); ++width) { + for (int64 plane = 0; plane < planes(); ++plane) { + for (int64 depth = 0; depth < this->depth(); ++depth) { + (*this)(plane, depth, height, width) = value(plane, depth); + } + } + } + } + } + + // Fills each of the minor-dim matrices with a number designating which minor + // dim matrix is enclosed by the shape. + void FillWithMinorDimNum() { + LOG(INFO) << "width: " << this->width(); + LOG(INFO) << "height: " << this->height(); + LOG(INFO) << "depth: " << this->depth(); + LOG(INFO) << "planes: " << this->planes(); + for (int64 height = 0; height < this->height(); ++height) { + for (int64 width = 0; width < this->width(); ++width) { + for (int64 plane = 0; plane < planes(); ++plane) { + for (int64 depth = 0; depth < this->depth(); ++depth) { + float this_val = plane * this->depth() + depth; + (*this)(plane, depth, height, width) = this_val; + } + } + } + } + } + + // Returns a string representation of the 4D array suitable for debugging. + string ToString() const { + std::vector<string> pieces = { + tensorflow::strings::Printf("p=%lld,z=%lld,y=%lld,x=%lld {\n", planes(), + depth(), height(), width())}; + for (int64 plane = 0; plane < planes_; ++plane) { + pieces.push_back(" {\n"); + for (int64 depth = 0; depth < depth_; ++depth) { + pieces.push_back(" {\n"); + for (int64 height = 0; height < height_; ++height) { + pieces.push_back(" {"); + for (int64 width = 0; width < width_; ++width) { + pieces.push_back(tensorflow::strings::StrCat( + (*this)(plane, depth, height, width), ", ")); + } + pieces.push_back("},\n"); + } + pieces.push_back(" },\n"); + } + pieces.push_back(" },\n"); + } + pieces.push_back("}"); + return tensorflow::str_util::Join(pieces, ""); + } + + private: + int64 planes_; + int64 depth_; + int64 height_; + int64 width_; + std::vector<T> values_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_ARRAY4D_H_ |