wehub-resource-sync
34 lines
926 B
Python
34 lines
926 B
Python
"""Operators for pipeline parallelism."""
|
|
|
|
from typing import List # noqa: UP035
|
|
|
|
from tvm import relax
|
|
from tvm.relax.frontend.nn import Tensor, op
|
|
|
|
|
|
def pipeline_stage_boundary(*tensors: Tensor) -> List[Tensor]: # noqa: UP006
|
|
"""Pipeline parallelism stage boundary mark operator in MLC.
|
|
|
|
Parameters
|
|
----------
|
|
tensors : Tensor
|
|
The tensors to be passed to the next stage.
|
|
|
|
Returns
|
|
-------
|
|
tensors : List[Tensor]
|
|
The list of input tensors passed to the next stage.
|
|
"""
|
|
return op.wrap_nested(
|
|
relax.call_pure_packed(
|
|
"mlc.pipeline_parallel_stage_boundary",
|
|
*[tensor._expr for tensor in tensors],
|
|
ty_args=(
|
|
tensors[0]._expr.ty
|
|
if len(tensors) == 1
|
|
else relax.TupleType([tensor._expr.ty for tensor in tensors])
|
|
),
|
|
),
|
|
name="pipeline_stage_boundary",
|
|
)
|