Home | History | Annotate | Download | only in service 
      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