| 2474 | } |
| 2475 | |
| 2476 | Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { |
| 2477 | absl::Span<HloInstruction* const> operands(custom_call->operands()); |
| 2478 | llvm::Type* i8_ptr_type = b_.getInt8PtrTy(); |
| 2479 | llvm::AllocaInst* operands_alloca = |
| 2480 | llvm_ir::EmitAllocaAtFunctionEntryWithCount( |
| 2481 | i8_ptr_type, b_.getInt32(operands.size()), "cc_operands_alloca", &b_); |
| 2482 | for (size_t i = 0; i < operands.size(); ++i) { |
| 2483 | const HloInstruction* operand = operands[i]; |
| 2484 | llvm::Value* operand_as_i8ptr = |
| 2485 | PointerCast(GetEmittedValueFor(operand), i8_ptr_type); |
| 2486 | llvm::Value* slot_in_operands_alloca = |
| 2487 | InBoundsGEP(operands_alloca, {b_.getInt64(i)}); |
| 2488 | Store(operand_as_i8ptr, slot_in_operands_alloca); |
| 2489 | } |
| 2490 | if (emit_code_for_msan_) { |
| 2491 | // Mark the alloca as initialized for msan. The buffer gets read by the |
| 2492 | // custom callee, which might be msan-instrumented. |
| 2493 | // TODO(b/66051036): Run the msan instrumentation pass instead. |
| 2494 | const llvm::DataLayout& dl = module_->getDataLayout(); |
| 2495 | llvm::Type* intptr_type = b_.getIntPtrTy(dl); |
| 2496 | auto* msan_unpoison_ir_function = llvm::cast<llvm::Function>( |
| 2497 | module_ |
| 2498 | ->getOrInsertFunction( |
| 2499 | "__msan_unpoison", |
| 2500 | llvm::FunctionType::get( |
| 2501 | /*Result=*/b_.getVoidTy(), |
| 2502 | /*Params=*/{i8_ptr_type, intptr_type}, /*isVarArg=*/false)) |
| 2503 | .getCallee()); |
| 2504 | Call(msan_unpoison_ir_function, |
| 2505 | {PointerCast(operands_alloca, i8_ptr_type), |
| 2506 | llvm::ConstantInt::get( |
| 2507 | intptr_type, *operands_alloca->getAllocationSizeInBits(dl) / 8)}); |
| 2508 | } |
| 2509 | auto* custom_call_ir_function = llvm::dyn_cast<llvm::Function>( |
| 2510 | module_ |
| 2511 | ->getOrInsertFunction( |
| 2512 | custom_call->custom_call_target(), |
| 2513 | llvm::FunctionType::get( |
| 2514 | /*Result=*/b_.getVoidTy(), |
| 2515 | /*Params=*/{i8_ptr_type, operands_alloca->getType()}, |
| 2516 | /*isVarArg=*/false)) |
| 2517 | .getCallee()); |
| 2518 | |
| 2519 | TF_RETURN_IF_ERROR(EmitTargetAddressForOp(custom_call)); |
| 2520 | // Write the tuple table if the output is a tuple. |
| 2521 | if (custom_call->shape().IsTuple()) { |
| 2522 | std::vector<llvm::Value*> base_ptrs; |
| 2523 | for (int i = 0; i < ShapeUtil::TupleElementCount(custom_call->shape()); |
| 2524 | ++i) { |
| 2525 | const Shape& elem_shape = |
| 2526 | ShapeUtil::GetTupleElementShape(custom_call->shape(), i); |
| 2527 | TF_RET_CHECK(!elem_shape.IsTuple()) << "Nested tuples not implemented"; |
| 2528 | TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice slice, |
| 2529 | assignment_.GetUniqueSlice(custom_call, {i})); |
| 2530 | llvm::Value* addr = EmitBufferPointer(slice, elem_shape); |
| 2531 | base_ptrs.push_back(addr); |
| 2532 | } |
| 2533 | llvm_ir::EmitTuple(GetIrArrayFor(custom_call), base_ptrs, &b_); |
nothing calls this directly
no test coverage detected