| 2394 | } |
| 2395 | |
| 2396 | Status IrEmitter::HandleFusion(HloInstruction* fusion) { |
| 2397 | auto* root = fusion->fused_expression_root(); |
| 2398 | if (llvm_ir::CanEmitFusedDynamicUpdateSliceInPlace(fusion, assignment_)) { |
| 2399 | VLOG(3) << "HandleFusion FusedDynamicUpdateSliceInPlace"; |
| 2400 | CpuElementalIrEmitter elemental_emitter(hlo_module_config_, this, module_); |
| 2401 | TF_RETURN_IF_ERROR(EmitTargetAddressForOp(fusion)); |
| 2402 | // Delegate to common implementation of fused in-place dynamic-update-slice. |
| 2403 | return llvm_ir::EmitFusedDynamicUpdateSliceInPlace( |
| 2404 | fusion, GetGeneratorForOperandIrArrays(fusion), GetIrArrayFor(fusion), |
| 2405 | &elemental_emitter, &b_); |
| 2406 | } else if (fusion->IsLoopFusion()) { |
| 2407 | VLOG(3) << "HandleFusion kLoop"; |
| 2408 | CpuElementalIrEmitter elemental_emitter(hlo_module_config_, this, module_); |
| 2409 | auto operands = GetIrArraysForOperandsOf(fusion); |
| 2410 | FusedIrEmitter fused_emitter(GetGeneratorForOperandIrArrays(fusion), |
| 2411 | &elemental_emitter); |
| 2412 | TF_RETURN_IF_ERROR(fusion->fused_expression_root()->Accept(&fused_emitter)); |
| 2413 | |
| 2414 | return EmitTargetElementLoop(fusion, fused_emitter.GetRootGenerator()); |
| 2415 | } else if (fusion->IsOutputFusion()) { |
| 2416 | VLOG(3) << "HandleFusion kOutput"; |
| 2417 | int64 dot_op_index = root->operand(0)->opcode() == HloOpcode::kDot ? 0 : 1; |
| 2418 | const HloInstruction* dot = root->operand(dot_op_index); |
| 2419 | CHECK_EQ(dot->opcode(), HloOpcode::kDot) |
| 2420 | << dot->ToString() << " " |
| 2421 | << fusion->fused_instructions_computation()->ToString(); |
| 2422 | |
| 2423 | int64 dot_lhs_param_number = dot->operand(0)->parameter_number(); |
| 2424 | int64 dot_rhs_param_number = dot->operand(1)->parameter_number(); |
| 2425 | int64 addend_param_number = |
| 2426 | root->operand(1 - dot_op_index)->parameter_number(); |
| 2427 | |
| 2428 | Shape target_shape = fusion->shape(); |
| 2429 | TF_RETURN_IF_ERROR(EmitTargetAddressForOp(fusion)); |
| 2430 | llvm_ir::IrArray target_array = GetIrArrayFor(fusion); |
| 2431 | |
| 2432 | llvm_ir::IrArray lhs_array( |
| 2433 | GetIrArrayFor(fusion->operand(dot_lhs_param_number))); |
| 2434 | llvm_ir::IrArray rhs_array( |
| 2435 | GetIrArrayFor(fusion->operand(dot_rhs_param_number))); |
| 2436 | llvm_ir::IrArray addend_array( |
| 2437 | GetIrArrayFor(fusion->operand(addend_param_number))); |
| 2438 | |
| 2439 | TF_RETURN_IF_ERROR( |
| 2440 | EmitDotOperation(*dot, target_array, lhs_array, rhs_array, |
| 2441 | &addend_array, GetExecutableRunOptionsArgument(), &b_, |
| 2442 | hlo_module_config_, target_machine_features_)); |
| 2443 | return Status::OK(); |
| 2444 | } else { |
| 2445 | return Unimplemented("Fusion kind not implemented on CPU"); |
| 2446 | } |
| 2447 | } |
| 2448 | |
| 2449 | Status IrEmitter::HandleCall(HloInstruction* call) { |
| 2450 | HloComputation* computation = call->to_apply(); |
nothing calls this directly
no test coverage detected