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   // Constructs a ComputationLayout from a ProgramShape. The layouts of the
36   // parameters and results are set to the default layout. Layouts in the
37   // ProgramShape are ignored.
38   explicit ComputationLayout(const ProgramShape& program_shape);
39 
40   // Returns the layout of a particular parameter.
41   const ShapeLayout& parameter_layout(int64 param_no) const {
42     return parameter_layouts_[param_no];
43   }
44   ShapeLayout* mutable_parameter_layout(int64 param_no) {
45     return &parameter_layouts_[param_no];
46   }
47 
48   // Returns the number of parameters in the computation.
49   int parameter_count() const { return parameter_layouts_.size(); }
50 
51   // Returns the ShapeLayouts of the parameters of the computation.
52   const std::vector<ShapeLayout>& parameter_layouts() const {
53     return parameter_layouts_;
54   }
55 
56   // Returns the ShapeLayout of a result of the computation.
57   const ShapeLayout& result_layout() const { return result_layout_; }
58   ShapeLayout* mutable_result_layout() { return &result_layout_; }
59 
60   // Returns the shape of the particular parameter or result of the computation
61   // with layout.
62   const Shape& parameter_shape(int64 param_no) const {
63     return parameter_layouts_[param_no].shape();
64   }
65   const Shape& result_shape() const { return result_layout_.shape(); }
66 
67   // Sets layouts of all parameters and the result to the default layout.
68   void SetToDefaultLayout();
69 
70   // Returns true if all layouts (parameters and result) have been set.
71   bool LayoutIsSet() const;
72 
73   // Returns a string representation of this object.
74   string ToString() const;
75 
76  private:
77   std::vector<ShapeLayout> parameter_layouts_;
78   ShapeLayout result_layout_;
79 };
80 
81 }  // namespace xla
82 
83 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_LAYOUT_H_
84