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_SHAPE_LAYOUT_H_ 17 #define TENSORFLOW_COMPILER_XLA_SHAPE_LAYOUT_H_ 18 19 #include <string> 20 21 #include "tensorflow/compiler/xla/shape_util.h" 22 #include "tensorflow/compiler/xla/types.h" 23 #include "tensorflow/compiler/xla/xla_data.pb.h" 24 #include "tensorflow/core/lib/core/status.h" 25 26 namespace xla { 27 28 // A ShapeLayout object encapsulates the layout of a particular shape (including 29 // tuples). This differs from the Layout proto which describes the layout of a 30 // single array. ShapeLayout contains a Layout proto for each array in the shape 31 // (a tuple can have more than one array). For array shapes, this object 32 // trivially holds a single Layout. Logically, ShapeLayout holds a nonmutable 33 // shape with mutable layouts. 34 class ShapeLayout { 35 public: 36 // Constructs a ShapeLayout of the given shape. Layouts are copied from the 37 // shape parameter. 38 explicit ShapeLayout(const Shape& shape) : shape_(shape) {} 39 40 // Assigns the layouts in this ShapeLayout to the Layout fields of the given 41 // shape. 'to_shape' and the shape of the ShapeLayout object must be 42 // compatible. 43 tensorflow::Status AssignLayoutToShape(Shape* to_shape) const; 44 45 // Returns true if the Layouts in this ShapeLayout match the layouts in the 46 // given shape. Returns false otherwise. If the given shape is not compatible 47 // with the ShapeLayout's shape, then false is returned. 48 bool MatchesLayoutInShape(const Shape& shape) const; 49 50 // Copies the layout from the given shape into this ShapeLayout. 'other_shape' 51 // must be compatible with the ShapeLayout's shape, and 'other_shape' must 52 // have a layout (LayoutUtil::HasLayout). 53 tensorflow::Status CopyLayoutFromShape(const Shape& other_shape); 54 55 // Clears (Layout::Clear) all the Layouts stored in this object. 56 void Clear(); 57 58 // Sets all Layouts stored in this object to the default layout. 59 void SetToDefaultLayout(); 60 61 // Returns the shape (with layouts). 62 const Shape& shape() const { return shape_; } 63 64 // Checks that a layout is set for the shape, and returns a reference to the 65 // layout directly on the shape. Shape must not be a tuple. 66 const Layout& layout() const; 67 68 // Returns true if all layouts have been set for this ShapeLayout object. That 69 // is, every array has a layout. 70 bool LayoutIsSet() const; 71 72 // Resets the layout on the shape to the provided layout. Shape must not be a 73 // tuple. 74 void ResetLayout(const Layout& layout); 75 76 // Returns a string representation of this object. 77 string ToString() const { return ShapeUtil::HumanStringWithLayout(shape_); } 78 79 // Tests for equality of both shape and layout (ShapeUtil::Equal). 80 bool operator==(const ShapeLayout& other) const; 81 bool operator!=(const ShapeLayout& other) const; 82 83 private: 84 Shape shape_; 85 }; 86 87 } // namespace xla 88 89 #endif // TENSORFLOW_COMPILER_XLA_SHAPE_LAYOUT_H_ 90
