Implementation:Microsoft Onnxruntime CPU MpiSend
| Knowledge Sources | |
|---|---|
| Domains | Training, CPU_Kernels |
| Last Updated | 2026-02-10 04:00 GMT |
Overview
Concrete tool for sending tensors via MPI during distributed training on CPU in the ONNX Runtime training framework.
Description
This file implements the Send kernel, which sends tensor data to a remote MPI process. The kernel aggregates all input tensors (starting from index 2, after the control signal and destination rank) into a single contiguous buffer with alignment padding, then sends the entire buffer in a single MPI_Send call. If tensor shapes cannot be inferred at compile time by the receiver, the kernel first sends shape information via SendShapeInfo. A boolean control signal input must be true for the kernel to execute, and a boolean control signal output is set to true upon completion. Same-rank communication is not allowed.
Usage
This kernel is used in distributed training pipeline parallelism and model parallelism scenarios where tensors need to be transferred to different processes via MPI. It is the counterpart of the Recv kernel.
Code Reference
Source Location
- Repository: Microsoft_Onnxruntime
- File: orttraining/orttraining/training_ops/cpu/communication/send.cc
- Lines: 1-126
Signature
void Send::SendData(OpKernelContext* ctx, const int dst,
const int num_tensors,
size_t aggregated_aligned_tensor_bytes,
std::vector<size_t> tensor_offsets_in_bytes,
std::vector<size_t> tensor_sizes_in_bytes) const;
Status Send::Compute(OpKernelContext* ctx) const;
Import
#include "orttraining/orttraining/training_ops/cpu/communication/send.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| input_signal | Tensor(bool) | Yes | Control signal (must be true) |
| remote_rank | Tensor(int64) | Yes | MPI rank of the destination process |
| tensors (variadic) | Tensor(V) | Yes | One or more tensors to send |
Outputs
| Name | Type | Description |
|---|---|---|
| output_signal | Tensor(bool) | Completion signal (set to true) |
Usage Examples
ONNX_OPERATOR_KERNEL_EX(
Send, kMSDomain, 1, kCpuExecutionProvider,
KernelDefBuilder()
.InputMemoryType(OrtMemTypeDefault, 0)
.InputMemoryType(OrtMemTypeDefault, 1)
.OutputMemoryType(OrtMemTypeDefault, 0)
.TypeConstraint("TBool", DataTypeImpl::GetTensorType<bool>())
.TypeConstraint("TInt64", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
Send);