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_SERVICE_COMPUTATION_LAYOUT_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_LAYOUT_H_ 18 19 #include <memory> 20 #include <string> 21 #include <vector> 22 23 #include "tensorflow/compiler/xla/shape_layout.h" 24 #include "tensorflow/compiler/xla/types.h" 25 #include "tensorflow/compiler/xla/xla_data.pb.h" 26 #include "tensorflow/core/platform/types.h" 27 28 namespace xla { 29 30 // Class which contains the layouts of the parameters and results of a 31 // computation. The layouts are stored as ShapeLayouts with immutable shapes and 32 // mutable layouts. 33 class ComputationLayout { 34 public: 35 // Creates a new ComputationLayout with the given result layout. ComputationLayout(ShapeLayout result_layout)36 explicit ComputationLayout(ShapeLayout result_layout) 37 : result_layout_(std::move(result_layout)) {} 38 39 // Constructs a ComputationLayout from a ProgramShape. The layouts of the 40 // parameters and results are set to the default layout. Layouts in the 41 // ProgramShape are ignored if ignore_layouts is true. 42 explicit ComputationLayout(const ProgramShape& program_shape, 43 bool ignore_layouts = true); 44 45 // Adds a new parameter layout to the computation layout. add_parameter_layout(ShapeLayout shape_layout)46 void add_parameter_layout(ShapeLayout shape_layout) { 47 parameter_layouts_.push_back(std::move(shape_layout)); 48 } 49 50 // Returns the layout of a particular parameter. parameter_layout(int64 param_no)51 const ShapeLayout& parameter_layout(int64 param_no) const { 52 return parameter_layouts_[param_no]; 53 } mutable_parameter_layout(int64 param_no)54 ShapeLayout* mutable_parameter_layout(int64 param_no) { 55 return ¶meter_layouts_[param_no]; 56 } 57 58 // Returns the number of parameters in the computation. parameter_count()59 int parameter_count() const { return parameter_layouts_.size(); } 60 61 // Returns the ShapeLayouts of the parameters of the computation. parameter_layouts()62 const std::vector<ShapeLayout>& parameter_layouts() const { 63 return parameter_layouts_; 64 } 65 66 // Returns the ShapeLayout of a result of the computation. result_layout()67 const ShapeLayout& result_layout() const { return result_layout_; } mutable_result_layout()68 ShapeLayout* mutable_result_layout() { return &result_layout_; } 69 70 // Returns the shape of the particular parameter or result of the computation 71 // with layout. parameter_shape(int64 param_no)72 const Shape& parameter_shape(int64 param_no) const { 73 return parameter_layouts_[param_no].shape(); 74 } result_shape()75 const Shape& result_shape() const { return result_layout_.shape(); } 76 77 // Sets layouts of all parameters and the result to the default layout. 78 void SetToDefaultLayout(); 79 80 // Returns true if all layouts (parameters and result) have been set. 81 bool LayoutIsSet() const; 82 83 // Returns a string representation of this object. 84 string ToString() const; 85 86 // Create a ProgramShape proto based on the parameter and result shapes held 87 // within this object. 88 ProgramShape ComputeProgramShape() const; 89 90 private: 91 std::vector<ShapeLayout> parameter_layouts_; 92 ShapeLayout result_layout_; 93 }; 94 95 } // namespace xla 96 97 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_LAYOUT_H_ 98