| 133 | } |
| 134 | |
| 135 | absl::optional<Shape> GetXfeedShape(bool is_infeed, |
| 136 | const HloModuleProto& module, |
| 137 | const Options& opts) { |
| 138 | std::vector<HloInstructionProto> xfeed_instrs; |
| 139 | for (const auto& comp : module.computations()) { |
| 140 | for (const auto& instruction : comp.instructions()) { |
| 141 | if (instruction.opcode() == HloOpcodeString(is_infeed |
| 142 | ? HloOpcode::kInfeed |
| 143 | : HloOpcode::kOutfeed)) { |
| 144 | xfeed_instrs.push_back(instruction); |
| 145 | } |
| 146 | } |
| 147 | } |
| 148 | |
| 149 | auto log_xfeed_instrs = [&] { |
| 150 | for (const auto& infeed : xfeed_instrs) { |
| 151 | LOG(ERROR) << " " << ShapeUtil::HumanString(Shape(infeed.shape())) << " " |
| 152 | << infeed.name(); |
| 153 | } |
| 154 | }; |
| 155 | |
| 156 | auto find_instruction_from_id_or_die = [&](int64 id) { |
| 157 | for (const auto& comp : module.computations()) { |
| 158 | for (const auto& instruction : comp.instructions()) { |
| 159 | if (instruction.id() == id) { |
| 160 | return instruction; |
| 161 | } |
| 162 | } |
| 163 | } |
| 164 | LOG(FATAL) << "No instruction with id " << id; |
| 165 | }; |
| 166 | |
| 167 | absl::optional<Shape> xfeed_shape; |
| 168 | string xfeed_name = is_infeed ? "infeed" : "outfeed"; |
| 169 | string fake_xfeed_shape = |
| 170 | is_infeed ? opts.fake_infeed_shape : opts.fake_outfeed_shape; |
| 171 | bool generate_fake_xfeed = |
| 172 | is_infeed ? opts.generate_fake_infeed : opts.generate_fake_outfeed; |
| 173 | if (!fake_xfeed_shape.empty()) { |
| 174 | xfeed_shape = std::move(ParseShape(fake_xfeed_shape)).ValueOrDie(); |
| 175 | } else if (generate_fake_xfeed) { |
| 176 | CHECK_LT(xfeed_instrs.size(), 2) |
| 177 | << "--generate_fake_" << xfeed_name |
| 178 | << " only works if the model has 0 or 1 " << xfeed_name << " ops."; |
| 179 | if (xfeed_instrs.empty()) { |
| 180 | LOG(INFO) << "Not generating fake " << xfeed_name |
| 181 | << " shape; model has no " << xfeed_name << "s."; |
| 182 | } else if (xfeed_instrs.size() == 1) { |
| 183 | // kInfeed instructions should have a shape (buffer, token). kOutfeed |
| 184 | // instructions should have operand 0 of shape `buffer`. We want to xfeed |
| 185 | // just `buffer`. |
| 186 | xfeed_shape = is_infeed |
| 187 | ? Shape(xfeed_instrs.front().shape()).tuple_shapes(0) |
| 188 | : Shape(find_instruction_from_id_or_die( |
| 189 | xfeed_instrs.front().operand_ids(0)) |
| 190 | .shape()); |
| 191 | LOG(INFO) << "Generating fake " << xfeed_name << " with inferred shape: " |
| 192 | << ShapeUtil::HumanString(*xfeed_shape); |
no test coverage detected