1 /* Copyright 2018 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_HLO_INPUT_OUTPUT_ALIAS_CONFIG_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INPUT_OUTPUT_ALIAS_CONFIG_H_ 18 19 #include <utility> 20 21 #include "absl/container/flat_hash_set.h" 22 #include "absl/types/optional.h" 23 #include "tensorflow/compiler/xla/service/hlo.pb.h" 24 #include "tensorflow/compiler/xla/shape_tree.h" 25 #include "tensorflow/compiler/xla/shape_util.h" 26 27 namespace xla { 28 29 class HloModule; 30 31 // This class specifies the alias map from output index to parameter number and 32 // parameter index in the entry computation. 33 class HloInputOutputAliasConfig { 34 public: 35 // The kind of aliases which can be set. A kUserAlias is one setup at 36 // compilation time by the user, and has to be respected. A kSystemAlias one 37 // might be setup by the compiler, if it decides it is convenient to do so. 38 enum AliasKind { 39 kNoAlias, 40 kUserAlias, 41 kSystemAlias, 42 }; 43 44 // Defines the alias information for a given output buffer. A given output 45 // buffer shape index can refer only to one parameter+index. 46 struct Alias { AliasAlias47 Alias(AliasKind kind, int64 parameter_number, ShapeIndex parameter_index) 48 : kind(kind), 49 parameter_number(parameter_number), 50 parameter_index(std::move(parameter_index)) {} 51 52 AliasKind kind; 53 int64 parameter_number; 54 ShapeIndex parameter_index; 55 }; 56 57 HloInputOutputAliasConfig() = default; 58 HloInputOutputAliasConfig(Shape output_shape)59 explicit HloInputOutputAliasConfig(Shape output_shape) 60 : alias_(output_shape) {} 61 62 virtual ~HloInputOutputAliasConfig() = default; 63 64 // Sets up alias config from `output_index` to `param_index` at 65 // `param_number`. 66 Status SetUpAlias(const ShapeIndex& output_index, int64 param_number, 67 const ShapeIndex& param_index, AliasKind kind); 68 69 // Returns the kind of alias for the given parameter number and parameter 70 // index. If no alias exists, AliasKind::kNoAlias is returned. 71 AliasKind ParameterAliasKind(int64 param_number, 72 const ShapeIndex& param_index) const; 73 74 // Returns true if the given parameter is aliased with one of the output 75 // buffers. ParameterHasAlias(int64 param_number,const ShapeIndex & param_index)76 bool ParameterHasAlias(int64 param_number, 77 const ShapeIndex& param_index) const { 78 return ParameterAliasKind(param_number, param_index) != AliasKind::kNoAlias; 79 } 80 81 // Checks whether the provided output index has already been aliased. 82 bool OutputHasAlias(const ShapeIndex& output_index) const; 83 84 // (De)Serializes an HloInputOutoutAliasConfig to/from an 85 // HloInputOutoutAliasProto. 86 HloInputOutputAliasProto ToProto() const; 87 88 static StatusOr<HloInputOutputAliasConfig> CreateFromProto( 89 const Shape& output_shape, const HloInputOutputAliasProto& proto); 90 91 // Returns the output index that the given parameter and parameter index is 92 // aliased with. A nullopt is returned if there is no output that is aliased 93 // with the parameter number and index. 94 absl::optional<ShapeIndex> GetAliasedOutput( 95 int64 param_number, const ShapeIndex& param_index) const; 96 97 // Returns the number of parameter and index of the parameter buffer that the 98 // given output buffer index is aliased with. A nullopt is returned if there 99 // is no parameter is aliased with the specific output. 100 absl::optional<Alias> GetAliasedParameter( 101 const ShapeIndex& output_index) const; 102 103 using AliasFn = 104 std::function<void(const ShapeIndex& output_index, const Alias&)>; 105 106 // Iterates through each aliased output and input. 107 void ForEachAlias(AliasFn fn) const; 108 109 using AliasFnWithStatus = 110 std::function<Status(const ShapeIndex& output_index, const Alias&)>; 111 112 // Verifies that the given config is valid for the given module. 113 // Specifically, the config's input and output should be in-bound and size of 114 // the aliased buffers should match. 115 Status Verify(const HloModule& module, 116 std::function<int64(const Shape&)> size_func_) const; 117 118 Status ForEachAliasWithStatus(AliasFnWithStatus fn) const; 119 120 string ToString() const; 121 122 private: 123 // A ShapeTree which indicates the list of buffers that's expected to be 124 // aliased. The key on this shape tree represents the output index. The value 125 // is an Alias data structure which defines the input parameter coordinates. 126 // If the value is nullopt, it means there is no parameter aliasing for this 127 // output. 128 ShapeTree<absl::optional<Alias>> alias_; 129 }; 130 131 std::ostream& operator<<(std::ostream& out, 132 const HloInputOutputAliasConfig& config); 133 134 } // namespace xla 135 136 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INPUT_OUTPUT_ALIAS_CONFIG_H_ 137