MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / EmitXfeedTransfer

Method EmitXfeedTransfer

tensorflow/compiler/xla/service/cpu/ir_emitter.cc:447–535  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

445}
446
447Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
448 llvm::Value* program_buffer_address) {
449 int64 length = ByteSizeOf(shape);
450 if (length <= 0 || length > std::numeric_limits<int32>::max()) {
451 return InvalidArgument(
452 "xfeed (infeed or outfeed) buffer length %d is outside the valid "
453 "size range",
454 length);
455 }
456 int32 length_32 = static_cast<int32>(length);
457
458 int32 shape_length;
459 TF_ASSIGN_OR_RETURN(
460 llvm::Value * shape_ptr,
461 llvm_ir::EncodeSelfDescribingShapeConstant(shape, &shape_length, &b_));
462
463 llvm::Type* int32_type = b_.getInt32Ty();
464 llvm::Type* i8_ptr_type = llvm::Type::getInt8PtrTy(module_->getContext());
465 llvm::FunctionType* acquire_type = llvm::FunctionType::get(
466 i8_ptr_type,
467 {/*run_options*/ i8_ptr_type, /*buffer_length*/ int32_type,
468 /*shape_ptr*/ i8_ptr_type, /*shape_length*/ int32_type},
469 /*isVarArg=*/false);
470
471 llvm::Function* acquire_func;
472 if (kind == XfeedKind::kInfeed) {
473 acquire_func = llvm::dyn_cast<llvm::Function>(
474 module_
475 ->getOrInsertFunction(
476 runtime::kAcquireInfeedBufferForDequeueSymbolName, acquire_type)
477 .getCallee());
478 } else {
479 acquire_func = llvm::dyn_cast<llvm::Function>(
480 module_
481 ->getOrInsertFunction(
482 runtime::kAcquireOutfeedBufferForPopulationSymbolName,
483 acquire_type)
484 .getCallee());
485 }
486 acquire_func->setCallingConv(llvm::CallingConv::C);
487
488 llvm::FunctionType* release_type = llvm::FunctionType::get(
489 b_.getVoidTy(),
490 {/*run_options*/ i8_ptr_type, /*buffer_length*/ int32_type,
491 /*buffer_ptr*/ i8_ptr_type, /*shape_ptr*/ i8_ptr_type,
492 /*shape_length*/ int32_type},
493 /*isVarArg=*/false);
494
495 llvm::Function* release_func;
496 if (kind == XfeedKind::kInfeed) {
497 release_func = llvm::dyn_cast<llvm::Function>(
498 module_
499 ->getOrInsertFunction(
500 runtime::kReleaseInfeedBufferAfterDequeueSymbolName,
501 release_type)
502 .getCallee());
503 } else {
504 release_func = llvm::dyn_cast<llvm::Function>(

Callers

nothing calls this directly

Calls 8

InvalidArgumentFunction · 0.85
getContextMethod · 0.80
ByteSizeOfFunction · 0.50
maxFunction · 0.50
TF_ASSIGN_OR_RETURNFunction · 0.50
getFunction · 0.50
CallFunction · 0.50

Tested by

no test coverage detected