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 #include "tensorflow/compiler/xla/service/hlo_domain_verifier.h"
17
18 #include <set>
19
20 #include "tensorflow/compiler/xla/service/hlo_computation.h"
21 #include "tensorflow/compiler/xla/service/hlo_domain_map.h"
22 #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h"
23 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
24 #include "tensorflow/compiler/xla/service/hlo_opcode.h"
25 #include "tensorflow/compiler/xla/types.h"
26
27 namespace xla {
28
29 class HloDomainVerifier::RunContext {
30 public:
RunContext(HloModule * module,HloDomainVerifier * verifier)31 RunContext(HloModule* module, HloDomainVerifier* verifier)
32 : module_(module), verifier_(verifier) {}
33
34 Status Run();
35
36 private:
37 // If the verifier caller passed an empty vector for kinds, we collect all the
38 // available domain types.
39 Status PopulateDomainKinds();
40
41 HloModule* module_;
42 HloDomainVerifier* verifier_;
43 };
44
PopulateDomainKinds()45 Status HloDomainVerifier::RunContext::PopulateDomainKinds() {
46 if (verifier_->kinds_.empty()) {
47 // The caller specified no domain kinds, collect all the ones available.
48 std::set<string> kinds;
49 for (HloComputation* computation : module_->computations()) {
50 for (HloInstruction* instruction : computation->instructions()) {
51 if (instruction->opcode() == HloOpcode::kDomain) {
52 TF_RET_CHECK(instruction->user_side_metadata().Kind() ==
53 instruction->operand_side_metadata().Kind())
54 << instruction->ToString();
55 kinds.insert(string(instruction->user_side_metadata().Kind()));
56 }
57 }
58 }
59 verifier_->kinds_.insert(verifier_->kinds_.end(), kinds.begin(),
60 kinds.end());
61 }
62 return Status::OK();
63 }
64
Run()65 Status HloDomainVerifier::RunContext::Run() {
66 VLOG(4) << "Running HLO Domain Verifier";
67 TF_RETURN_IF_ERROR(PopulateDomainKinds());
68 for (HloComputation* computation : module_->computations()) {
69 for (auto& kind : verifier_->kinds_) {
70 // First create the domain instruction sets. A domain instruction set is
71 // the set of instructions whose edges never cross a kDomain instruction.
72 TF_ASSIGN_OR_RETURN(std::unique_ptr<HloDomainMap> domain_map,
73 HloDomainMap::Create(computation, kind));
74 // Verify every domain populated within the map.
75 for (auto& domain : domain_map->GetDomains()) {
76 TF_RETURN_IF_ERROR(VerifyDomain(*domain).status());
77 }
78 }
79 }
80 return Status::OK();
81 }
82
Run(HloModule * module)83 StatusOr<bool> HloDomainVerifier::Run(HloModule* module) {
84 RunContext run_context(module, this);
85 TF_RETURN_IF_ERROR(run_context.Run());
86 return false;
87 }
88
VerifyDomain(const DomainMetadata::Domain & domain)89 StatusOr<const DomainMetadata*> HloDomainVerifier::VerifyDomain(
90 const DomainMetadata::Domain& domain) {
91 const DomainMetadata* ref_metadata = nullptr;
92 VLOG(4) << "Reach set:";
93 for (HloInstruction* instruction : domain.instructions) {
94 VLOG(4) << " " << instruction->name();
95 }
96 VLOG(4) << " Domains:";
97 for (HloInstruction* instruction : domain.enter_domains) {
98 const DomainMetadata& meta = instruction->user_side_metadata();
99 VLOG(4) << " User side: " << instruction->name();
100 VLOG(4) << " " << meta.ToString();
101 if (ref_metadata == nullptr) {
102 ref_metadata = &meta;
103 } else {
104 TF_RET_CHECK(meta.Matches(*ref_metadata))
105 << "Metadata mismatch at instruction " << instruction->name() << " : "
106 << meta.ToString() << " vs " << ref_metadata->ToString();
107 }
108 }
109 for (HloInstruction* instruction : domain.exit_domains) {
110 const DomainMetadata& meta = instruction->operand_side_metadata();
111 VLOG(4) << " Operand side: " << instruction->name();
112 VLOG(4) << " " << meta.ToString();
113 if (ref_metadata == nullptr) {
114 ref_metadata = &meta;
115 } else {
116 TF_RET_CHECK(meta.Matches(*ref_metadata))
117 << "Metadata mismatch at instruction " << instruction->name() << " : "
118 << meta.ToString() << " vs " << ref_metadata->ToString();
119 }
120 }
121 return ref_metadata;
122 }
123
124 } // namespace xla
125