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