1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_ARRAY4D_H_ 17 #define TENSORFLOW_COMPILER_XLA_ARRAY4D_H_ 18 19 #include <algorithm> 20 #include <functional> 21 #include <initializer_list> 22 #include <iterator> 23 #include <memory> 24 #include <numeric> 25 #include <random> 26 #include <string> 27 #include <vector> 28 29 #include "absl/strings/str_cat.h" 30 #include "absl/types/span.h" 31 #include "tensorflow/compiler/xla/array.h" 32 #include "tensorflow/compiler/xla/array2d.h" 33 #include "tensorflow/compiler/xla/types.h" 34 #include "tensorflow/core/platform/logging.h" 35 #include "tensorflow/core/platform/macros.h" 36 #include "tensorflow/core/platform/types.h" 37 38 namespace xla { 39 40 // Simple 4D array structure, similar in form to Array2D, for use primarily in 41 // testing and describing to XLA APIs values in the 4D array structures used 42 // in convolutions. 43 // 44 // The data layout is, in order from major to minor: 45 // 46 // First dimension: plane, batch, n1 47 // Second dimension: depth, feature, z, n2 48 // Third dimension: height, y, n3 49 // Fourth dimension: width, x, n4 50 // 51 // These dimensions are referred to by various names, so that is why 52 // more than one name is given above. See operator() for the exact 53 // calculation of 1d indices from 4d indices. 54 template <typename T> 55 class Array4D : public Array<T> { 56 public: Array4D()57 Array4D() : Array<T>(std::vector<int64>{0, 0, 0, 0}) {} 58 59 // Creates a 4D array, uninitialized values. Array4D(int64_t planes,int64_t depth,int64_t height,int64_t width)60 Array4D(int64_t planes, int64_t depth, int64_t height, int64_t width) 61 : Array<T>(std::vector<int64>{planes, depth, height, width}) {} 62 63 // Creates a 4D array, initialized to value. Array4D(int64_t planes,int64_t depth,int64_t height,int64_t width,T value)64 Array4D(int64_t planes, int64_t depth, int64_t height, int64_t width, T value) 65 : Array<T>(std::vector<int64>{planes, depth, height, width}, value) {} 66 67 // Creates a 4D array, filled with values. 68 // 69 // We need to set a default type for Container so that code like 70 // Array4D(1, 1, 1, 1, {1}) will work. The template cannot infer the 71 // initializer_list type in that case without this default. 72 template <typename Container = std::initializer_list<T>> Array4D(int64_t planes,int64_t depth,int64_t height,int64_t width,const Container & values)73 Array4D(int64_t planes, int64_t depth, int64_t height, int64_t width, 74 const Container& values) 75 : Array4D(planes, depth, height, width) { 76 this->SetValues(values); 77 } 78 79 // Construct an Array4D with the given nested initializer list. Array4D(std::initializer_list<std::initializer_list<std::initializer_list<std::initializer_list<T>>>> values)80 Array4D(std::initializer_list<std::initializer_list< 81 std::initializer_list<std::initializer_list<T>>>> 82 values) 83 : Array<T>(values) {} 84 85 // Creates an array of a floating-point type (half, bfloat16, float, 86 // or double) from the given nested initializer list of float values. 87 template <typename T2, typename = typename std::enable_if< 88 (std::is_same<T, Eigen::half>::value || 89 std::is_same<T, bfloat16>::value || 90 std::is_same<T, float>::value || 91 std::is_same<T, double>::value) && 92 std::is_same<T2, float>::value>::type> Array4D(std::initializer_list<std::initializer_list<std::initializer_list<std::initializer_list<T2>>>> values)93 Array4D(std::initializer_list<std::initializer_list< 94 std::initializer_list<std::initializer_list<T2>>>> 95 values) 96 : Array<T>(values) {} 97 98 // Numerically-named aliases for the various dimensions. This matches the 99 // dimension names used in array3d. n4()100 int64 n4() const { return this->dim(3); } n3()101 int64 n3() const { return this->dim(2); } n2()102 int64 n2() const { return this->dim(1); } n1()103 int64 n1() const { return this->dim(0); } 104 width()105 int64 width() const { return this->dim(3); } height()106 int64 height() const { return this->dim(2); } depth()107 int64 depth() const { return this->dim(1); } planes()108 int64 planes() const { return this->dim(0); } 109 110 // Fills all of the {p,z} with the array provided, which specifies {y,x}. FillWithYX(const Array2D<T> & value)111 void FillWithYX(const Array2D<T>& value) { 112 CHECK_EQ(value.height(), height()); 113 CHECK_EQ(value.width(), width()); 114 for (int64_t plane = 0; plane < planes(); ++plane) { 115 for (int64_t depth = 0; depth < this->depth(); ++depth) { 116 for (int64_t height = 0; height < this->height(); ++height) { 117 for (int64_t width = 0; width < this->width(); ++width) { 118 (*this)(plane, depth, height, width) = value(height, width); 119 } 120 } 121 } 122 } 123 } 124 125 // Fills all of the {p,x} with the array provided, which specifies {z,y}. FillWithZY(const Array2D<T> & value)126 void FillWithZY(const Array2D<T>& value) { 127 CHECK_EQ(value.height(), depth()); 128 CHECK_EQ(value.width(), height()); 129 for (int64_t plane = 0; plane < planes(); ++plane) { 130 for (int64_t depth = 0; depth < this->depth(); ++depth) { 131 for (int64_t height = 0; height < this->height(); ++height) { 132 for (int64_t width = 0; width < this->width(); ++width) { 133 (*this)(plane, depth, height, width) = value(depth, height); 134 } 135 } 136 } 137 } 138 } 139 140 // Fills all of the {x,y} with the array provided, which specifies {p,z}. FillWithPZ(const Array2D<T> & value)141 void FillWithPZ(const Array2D<T>& value) { 142 CHECK_EQ(value.height(), planes()); 143 CHECK_EQ(value.width(), depth()); 144 for (int64_t height = 0; height < this->height(); ++height) { 145 for (int64_t width = 0; width < this->width(); ++width) { 146 for (int64_t plane = 0; plane < planes(); ++plane) { 147 for (int64_t depth = 0; depth < this->depth(); ++depth) { 148 (*this)(plane, depth, height, width) = value(plane, depth); 149 } 150 } 151 } 152 } 153 } 154 155 // Fills each of the minor-dim matrices with a number designating which minor 156 // dim matrix is enclosed by the shape. FillWithMinorDimNum()157 void FillWithMinorDimNum() { 158 LOG(INFO) << "width: " << this->width(); 159 LOG(INFO) << "height: " << this->height(); 160 LOG(INFO) << "depth: " << this->depth(); 161 LOG(INFO) << "planes: " << this->planes(); 162 for (int64_t height = 0; height < this->height(); ++height) { 163 for (int64_t width = 0; width < this->width(); ++width) { 164 for (int64_t plane = 0; plane < planes(); ++plane) { 165 for (int64_t depth = 0; depth < this->depth(); ++depth) { 166 float this_val = plane * this->depth() + depth; 167 (*this)(plane, depth, height, width) = this_val; 168 } 169 } 170 } 171 } 172 } 173 }; 174 175 } // namespace xla 176 177 #endif // TENSORFLOW_COMPILER_XLA_ARRAY4D_H_ 178