/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_LAYOUT_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_LAYOUT_H_

#include <memory>
#include <string>
#include <vector>

#include "tensorflow/compiler/xla/shape_layout.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/platform/types.h"

namespace xla {

// Class which contains the layouts of the parameters and results of a
// computation. The layouts are stored as ShapeLayouts with immutable shapes and
// mutable layouts.
class ComputationLayout {
 public:
  // Constructs a ComputationLayout from a ProgramShape. The layouts of the
  // parameters and results are set to the default layout. Layouts in the
  // ProgramShape are ignored.
  explicit ComputationLayout(const ProgramShape& program_shape);

  // Returns the layout of a particular parameter.
  const ShapeLayout& parameter_layout(int64 param_no) const {
    return parameter_layouts_[param_no];
  }
  ShapeLayout* mutable_parameter_layout(int64 param_no) {
    return &parameter_layouts_[param_no];
  }

  // Returns the number of parameters in the computation.
  int parameter_count() const { return parameter_layouts_.size(); }

  // Returns the ShapeLayouts of the parameters of the computation.
  const std::vector<ShapeLayout>& parameter_layouts() const {
    return parameter_layouts_;
  }

  // Returns the ShapeLayout of a result of the computation.
  const ShapeLayout& result_layout() const { return result_layout_; }
  ShapeLayout* mutable_result_layout() { return &result_layout_; }

  // Returns the shape of the particular parameter or result of the computation
  // with layout.
  const Shape& parameter_shape(int64 param_no) const {
    return parameter_layouts_[param_no].shape();
  }
  const Shape& result_shape() const { return result_layout_.shape(); }

  // Sets layouts of all parameters and the result to the default layout.
  void SetToDefaultLayout();

  // Returns true if all layouts (parameters and result) have been set.
  bool LayoutIsSet() const;

  // Returns a string representation of this object.
  string ToString() const;

 private:
  std::vector<ShapeLayout> parameter_layouts_;
  ShapeLayout result_layout_;
};

}  // namespace xla

#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_LAYOUT_H_