| 4450 | } |
| 4451 | |
| 4452 | Status AlgebraicSimplifierVisitor::HandleMap(HloInstruction* map) { |
| 4453 | auto* map_computation = map->to_apply(); |
| 4454 | auto* map_root = map_computation->root_instruction(); |
| 4455 | if (map_root->opcode() == HloOpcode::kParameter) { |
| 4456 | ReplaceInstructionIfSameShape( |
| 4457 | map, map->mutable_operand(map_root->parameter_number())); |
| 4458 | return Status::OK(); |
| 4459 | } |
| 4460 | if (map_root->opcode() == HloOpcode::kConstant) { |
| 4461 | if (!ShapeUtil::IsScalar(map_root->shape())) { |
| 4462 | return Status::OK(); |
| 4463 | } |
| 4464 | auto clone = map_root->CloneWithNewOperands(map_root->shape(), {}); |
| 4465 | if (ShapeUtil::IsScalar(map->shape())) { |
| 4466 | return ReplaceWithNewInstruction(map, std::move(clone)); |
| 4467 | } |
| 4468 | return ReplaceWithNewInstruction( |
| 4469 | map, |
| 4470 | HloInstruction::CreateBroadcast( |
| 4471 | map->shape(), computation_->AddInstruction(std::move(clone)), {})); |
| 4472 | } |
| 4473 | // Inline the map if the map computation only contains an elementwise |
| 4474 | // operation that can accept arbitrary shapes. |
| 4475 | if (map_root->opcode() == HloOpcode::kFusion || !map_root->IsElementwise()) { |
| 4476 | return Status::OK(); |
| 4477 | } |
| 4478 | std::vector<HloInstruction*> new_operands; |
| 4479 | for (auto* root_operand : map_root->operands()) { |
| 4480 | if (root_operand->opcode() != HloOpcode::kParameter) { |
| 4481 | return Status::OK(); |
| 4482 | } |
| 4483 | new_operands.push_back( |
| 4484 | map->mutable_operand(root_operand->parameter_number())); |
| 4485 | } |
| 4486 | auto clone = map_root->CloneWithNewOperands(map->shape(), new_operands); |
| 4487 | return ReplaceWithNewInstruction(map, std::move(clone)); |
| 4488 | } |
| 4489 | |
| 4490 | StatusOr<bool> AlgebraicSimplifier::Run(HloModule* module) { |
| 4491 | XLA_VLOG_LINES(2, |
nothing calls this directly
no test coverage detected