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_SHARDING_METADATA_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_SHARDING_METADATA_H_ 18 19 #include <string> 20 #include <utility> 21 22 #include "absl/container/flat_hash_map.h" 23 #include "absl/types/span.h" 24 #include "tensorflow/compiler/xla/service/hlo_domain_metadata.h" 25 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 26 #include "tensorflow/compiler/xla/service/hlo_sharding.h" 27 #include "tensorflow/core/lib/core/status.h" 28 29 namespace xla { 30 31 // A DomainMetadata implementation that internally wraps a sharding attribute. 32 class ShardingMetadata : public DomainMetadata { 33 public: ShardingMetadata(std::shared_ptr<const HloSharding> sharding)34 explicit ShardingMetadata(std::shared_ptr<const HloSharding> sharding) 35 : sharding_(std::move(sharding)) {} 36 37 std::unique_ptr<DomainMetadata> Clone() const override; 38 Kind()39 absl::string_view Kind() const override { return KindName(); } 40 41 bool Matches(const DomainMetadata& other) const override; 42 43 template <typename H> AbslHashValue(H h,const ShardingMetadata & sharding_metadata)44 friend H AbslHashValue(H h, const ShardingMetadata& sharding_metadata) { 45 const bool has_sharding = sharding_metadata.sharding_ != nullptr; 46 if (has_sharding) { 47 h = H::combine(std::move(h), *sharding_metadata.sharding_); 48 } 49 return H::combine(std::move(h), has_sharding); 50 } 51 Hash()52 size_t Hash() const override { return absl::HashOf(*this); } 53 54 std::string ToString() const override; 55 sharding()56 const HloSharding* sharding() const { return sharding_.get(); } 57 KindName()58 static absl::string_view KindName() { return "sharding"; } 59 60 static StatusOr<const ShardingMetadata*> ToShardingMetadata( 61 const DomainMetadata* metadata); 62 63 // Apply the specified domain metadata onto the specified domain. If no 64 // metadata is specified then apply sharding heuristics and normalize the 65 // instructions whose sharding deviates from the one which is inferred as to 66 // be the original one. Policy wise, HLO passes are allowed to create new 67 // unassigned instructions, but if they do create assigned ones, they have to 68 // conform to the ones around. 69 static Status NormalizeShardingDomain(const DomainMetadata::Domain& domain, 70 const DomainMetadata* metadata); 71 72 private: 73 std::shared_ptr<const HloSharding> sharding_; 74 }; 75 76 // If the sharding between root and instruction changes then returns a 77 // ShardingMetadata based kDomain instruction what can be used to separate 78 // operand and instruction. 79 // Returns nullptr if there is no need for a domain separation. 80 class ShardingDomainCreator { 81 public: 82 HloInstruction* operator()(HloInstruction* instruction, HloInstruction* root, 83 HloInstruction* operand); 84 85 private: 86 // Map from instruction and user sharding to domain users to CSE identical 87 // domains. 88 struct DomainCseMapKey { 89 const HloInstruction* instruction; 90 std::shared_ptr<const HloSharding> sharding; 91 92 bool operator==(const DomainCseMapKey& other) const; 93 94 template <typename H> AbslHashValueDomainCseMapKey95 friend H AbslHashValue(H h, const DomainCseMapKey& key) { 96 h = H::combine(std::move(h), key.instruction); 97 const bool has_sharding = key.sharding != nullptr; 98 if (has_sharding) { 99 h = H::combine(std::move(h), *key.sharding); 100 } 101 return H::combine(std::move(h), has_sharding); 102 } 103 }; 104 absl::flat_hash_map<DomainCseMapKey, HloInstruction*> domain_cse_map_; 105 }; 106 107 } // namespace xla 108 109 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_SHARDING_METADATA_H_ 110