wehub-resource-sync
400 lines
17 KiB
Python
400 lines
17 KiB
Python
"""A compiler pass that rewrites IR for pipeline parallelism."""
|
|
|
|
from typing import Dict, List, Optional, Tuple # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import relax, tirx
|
|
from tvm.ir.module import IRModule
|
|
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="PipelineParallelRewrite")
|
|
class PipelineParallelRewrite:
|
|
"""A compiler pass that rewrites IR for pipeline parallelism."""
|
|
|
|
def transform_module(
|
|
self,
|
|
mod: IRModule,
|
|
_ctx: tvm.transform.PassContext,
|
|
) -> IRModule:
|
|
"""IRModule-level transformation"""
|
|
return _PipelineParallelRewriter(mod.clone()).transform()
|
|
|
|
|
|
@mutator
|
|
class _PipelineParallelRewriter(PyExprMutator):
|
|
def __init__(self, mod: IRModule):
|
|
super().__init__(mod)
|
|
self.mod = mod
|
|
self.old_packed_params_var: relax.Var
|
|
self.new_main_packed_params_var: relax.Var
|
|
self.new_stage_func_packed_params: relax.Var
|
|
self.undefined_shape_vars_remap: Dict[tirx.Var, tirx.Var] # noqa: UP006
|
|
self.undefined_param_shape_vars_remap: Dict[tirx.Var, tirx.Var] # noqa: UP006
|
|
|
|
def transform(self) -> IRModule:
|
|
"""Entry point of the transformation"""
|
|
for g_var, func in self.mod.functions_items():
|
|
if not isinstance(func, relax.Function) or "pipeline_parallel_stages" not in func.attrs:
|
|
continue
|
|
num_stages = int(func.attrs["pipeline_parallel_stages"])
|
|
if num_stages == 1:
|
|
continue
|
|
|
|
pipeline_stages, stage_send_vars, stage_receive_vars = _extract_pipeline_stages(func)
|
|
assert len(pipeline_stages) == num_stages, (
|
|
"Number of pipeline stages mismatches: "
|
|
f"expecting {num_stages} stages, but {len(pipeline_stages)} are found in the IR."
|
|
)
|
|
|
|
required_func_params = _analyze_required_func_params(pipeline_stages, func.params)
|
|
|
|
assert "num_input" in func.attrs
|
|
num_input = int(func.attrs["num_input"])
|
|
assert (
|
|
len(func.params) == num_input + 1
|
|
and isinstance(func.params[num_input], relax.Var)
|
|
and func.params[num_input].name_hint == "packed_params"
|
|
), 'Only the extra "packed_params" parameter is allowed'
|
|
self.old_packed_params_var = func.params[num_input]
|
|
self.new_main_packed_params_var = relax.Var("packed_params", relax.ObjectType())
|
|
for required_params in required_func_params:
|
|
for i, param in enumerate(required_params):
|
|
if param.same_as(self.old_packed_params_var):
|
|
required_params.pop(i)
|
|
break
|
|
func_output = func.body.body
|
|
assert isinstance(func_output, relax.Var)
|
|
|
|
stage_func_gvs = []
|
|
caller_args_list = []
|
|
for i in range(num_stages):
|
|
stage_func_gv, caller_args = self._create_stage_func(
|
|
g_var.name_hint + f"_stage{i}",
|
|
pipeline_stages[i],
|
|
required_func_params[i],
|
|
stage_receive_vars[i],
|
|
stage_send_vars[i],
|
|
func.attrs,
|
|
func_output=func_output if i == num_stages - 1 else None,
|
|
)
|
|
stage_func_gvs.append(stage_func_gv)
|
|
caller_args_list.append(caller_args)
|
|
|
|
# Create and update the entry function, which dispatches toz the stage functions
|
|
# according to the disco worker group id.
|
|
bb = relax.BlockBuilder()
|
|
params = [*list(func.params[:-1]), self.new_main_packed_params_var]
|
|
with bb.function(g_var.name_hint, params=params):
|
|
dispatch_func_args = []
|
|
for stage_func_gv, caller_args in zip(stage_func_gvs, caller_args_list):
|
|
dispatch_func_args.append([stage_func_gv, *caller_args])
|
|
output = bb.emit(
|
|
relax.op.call_builtin_with_ctx(
|
|
"mlc.multi_gpu.DispatchFunctionByGroup",
|
|
args=[dispatch_func_args],
|
|
ty_args=relax.ObjectType(),
|
|
)
|
|
)
|
|
dispatch_func_gv = bb.emit_func_output(output)
|
|
dispatch_func = bb.finalize()[dispatch_func_gv]
|
|
self.builder_.update_func(g_var, dispatch_func)
|
|
|
|
return self.builder_.finalize()
|
|
|
|
def _create_stage_func(
|
|
self,
|
|
func_name: str,
|
|
stage_bindings: List[relax.Binding], # noqa: UP006
|
|
required_func_params: List[relax.Var], # noqa: UP006
|
|
stage_receive_vars: List[relax.Var], # noqa: UP006
|
|
stage_send_vars: List[relax.Var], # noqa: UP006
|
|
func_attrs: tvm.ir.DictAttrs,
|
|
func_output: Optional[relax.Var],
|
|
) -> Tuple[tvm.ir.GlobalVar, List[relax.Expr]]: # noqa: UP006
|
|
self.undefined_shape_vars_remap = {}
|
|
self.undefined_param_shape_vars_remap = {}
|
|
|
|
# Prepare the func parameters (except the shape variables and packed params)
|
|
params, args = self._prepare_stage_func_params_and_args(required_func_params)
|
|
for new_param, old_param in zip(params, required_func_params):
|
|
self.set_var_remap(old_param, new_param)
|
|
# Create new packed params
|
|
self.new_stage_func_packed_params = relax.Var("packed_params", relax.ObjectType())
|
|
self.set_var_remap(self.old_packed_params_var, self.new_stage_func_packed_params)
|
|
|
|
new_func_outputs = []
|
|
with self.builder_.function(func_name, pure=False):
|
|
with self.builder_.dataflow():
|
|
# Emit the tensors received from last stage.
|
|
for receive_var in stage_receive_vars:
|
|
new_receive_var = self.builder_.emit(
|
|
relax.call_dps_packed(
|
|
"runtime.disco.recv_from_prev_group",
|
|
args=[],
|
|
out_ty=self._update_struct_info(receive_var.ty),
|
|
),
|
|
name_hint=receive_var.name_hint,
|
|
)
|
|
self.set_var_remap(receive_var, new_receive_var)
|
|
# Process the bindings in this stage.
|
|
for stage_binding in stage_bindings:
|
|
if stage_binding.var in stage_send_vars or stage_binding.var.same_as(
|
|
func_output
|
|
):
|
|
assert isinstance(stage_binding, relax.VarBinding)
|
|
new_var = self.builder_.emit_output(
|
|
self.visit_expr(stage_binding.value),
|
|
name_hint=stage_binding.var.name_hint,
|
|
)
|
|
self.set_var_remap(stage_binding.var, new_var)
|
|
new_func_outputs.append(new_var)
|
|
else:
|
|
self.visit_binding(stage_binding)
|
|
# Emit the calls to send tensors to the next stage.
|
|
for send_var in stage_send_vars:
|
|
new_send_var = self.get_var_remap(send_var)
|
|
self.builder_.emit(
|
|
relax.Call(
|
|
relax.ExternFunc("runtime.disco.send_to_next_group"),
|
|
args=[new_send_var],
|
|
ty_args=None,
|
|
)
|
|
)
|
|
# Create the param for the shape variables.
|
|
shape_var_params = []
|
|
shape_var_args = []
|
|
for (
|
|
shape_var_arg,
|
|
shape_var_param,
|
|
) in self.undefined_shape_vars_remap.items():
|
|
if shape_var_arg not in self.undefined_param_shape_vars_remap:
|
|
shape_var_params.append(shape_var_param)
|
|
shape_var_args.append(shape_var_arg)
|
|
params.append(relax.Var("s", relax.ShapeType(shape_var_params)))
|
|
args.append(relax.ShapeExpr(shape_var_args))
|
|
# Add the packed params.
|
|
params.append(self.new_stage_func_packed_params)
|
|
args.append(self.new_main_packed_params_var)
|
|
# Conclude the function.
|
|
if func_output is not None:
|
|
assert len(new_func_outputs) == 1
|
|
new_gv = self.builder_.emit_func_output(
|
|
(
|
|
new_func_outputs[0]
|
|
if len(new_func_outputs) == 1
|
|
and isinstance(new_func_outputs[0].ty, relax.TupleType)
|
|
else new_func_outputs
|
|
),
|
|
params=params,
|
|
)
|
|
|
|
new_func = (
|
|
self.builder_.get()[new_gv]
|
|
.with_attrs(func_attrs)
|
|
.with_attr("num_input", len(params) - 1)
|
|
.without_attr("global_symbol")
|
|
.without_attr("pipeline_parallel_stages")
|
|
)
|
|
self.builder_.update_func(new_gv, new_func)
|
|
return new_gv, args
|
|
|
|
def visit_var_binding_(self, binding: relax.VarBinding) -> None:
|
|
if not isinstance(binding.value, relax.TupleGetItem):
|
|
super().visit_var_binding_(binding)
|
|
return
|
|
|
|
tuple_value = self.visit_expr(binding.value.tuple_value)
|
|
if not tuple_value.same_as(self.new_stage_func_packed_params):
|
|
super().visit_var_binding_(binding)
|
|
return
|
|
|
|
assert isinstance(binding.var.ty, relax.TensorType)
|
|
cur_num_undefined_param_shape_vars = len(self.undefined_param_shape_vars_remap)
|
|
new_tensor_struct_info = self._update_struct_info(
|
|
binding.var.ty, self.undefined_param_shape_vars_remap
|
|
)
|
|
has_new_undefined_shape_var = (
|
|
len(self.undefined_param_shape_vars_remap) != cur_num_undefined_param_shape_vars
|
|
)
|
|
self.undefined_shape_vars_remap = {
|
|
**self.undefined_shape_vars_remap,
|
|
**self.undefined_param_shape_vars_remap,
|
|
}
|
|
ret_sinfo = (
|
|
new_tensor_struct_info if not has_new_undefined_shape_var else relax.ObjectType()
|
|
)
|
|
call = relax.call_pure_packed(
|
|
"vm.builtin.tuple_getitem",
|
|
self.new_stage_func_packed_params,
|
|
relax.prim_value(binding.value.index),
|
|
ty_args=ret_sinfo,
|
|
)
|
|
new_binding_var = self.builder_.emit(call, binding.var.name_hint)
|
|
if has_new_undefined_shape_var:
|
|
new_binding_var = self.builder_.match_cast(
|
|
new_binding_var, new_tensor_struct_info, binding.var.name_hint + "_cast"
|
|
)
|
|
self.set_var_remap(binding.var, new_binding_var)
|
|
|
|
def visit_call_(self, call: relax.Call) -> relax.Call:
|
|
call = super().visit_call_(call)
|
|
return relax.Call(
|
|
call.op,
|
|
call.args,
|
|
call.attrs,
|
|
ty_args=[self._update_struct_info(struct_info) for struct_info in call.ty_args],
|
|
)
|
|
|
|
def _prepare_stage_func_params_and_args(
|
|
self,
|
|
required_func_params: List[relax.Var], # noqa: UP006
|
|
) -> Tuple[List[relax.Var], List[relax.Expr]]: # noqa: UP006
|
|
params: List[relax.Var] = [] # noqa: UP006
|
|
args: List[relax.Expr] = [] # noqa: UP006
|
|
for required_param in required_func_params:
|
|
struct_info = self._update_struct_info(required_param.ty)
|
|
params.append(relax.Var(required_param.name_hint, struct_info))
|
|
args.append(required_param)
|
|
|
|
return params, args
|
|
|
|
def _update_struct_info(
|
|
self,
|
|
struct_info: relax.Type,
|
|
undefined_var_remap: Optional[Dict[tirx.Var, tirx.Var]] = None, # noqa: UP006
|
|
) -> relax.Type:
|
|
if undefined_var_remap is None:
|
|
undefined_var_remap = self.undefined_shape_vars_remap
|
|
if isinstance(struct_info, relax.TensorType):
|
|
return (
|
|
relax.TensorType(
|
|
self._update_shape(struct_info.shape.values, undefined_var_remap),
|
|
struct_info.dtype,
|
|
)
|
|
if struct_info.shape is not None and isinstance(struct_info.shape, relax.ShapeExpr)
|
|
else struct_info
|
|
)
|
|
if isinstance(struct_info, relax.ShapeType):
|
|
return (
|
|
relax.ShapeType(self._update_shape(struct_info.values, undefined_var_remap))
|
|
if struct_info.values is not None
|
|
else struct_info
|
|
)
|
|
if isinstance(struct_info, relax.ObjectType):
|
|
return relax.ObjectType()
|
|
if isinstance(struct_info, relax.TupleType):
|
|
return relax.TupleType(
|
|
[self._update_struct_info(field_sinfo) for field_sinfo in struct_info.fields]
|
|
)
|
|
return struct_info
|
|
|
|
def _copy_undefined_var(
|
|
self,
|
|
expr: tirx.Expr,
|
|
undefined_var_remap: Dict[tirx.Var, tirx.Var], # noqa: UP006
|
|
) -> None:
|
|
def _visit_expr(e: tirx.Expr) -> None:
|
|
if isinstance(e, tirx.Var) and e not in undefined_var_remap:
|
|
new_var = tirx.Var(e.name, e.ty)
|
|
undefined_var_remap[e] = new_var
|
|
|
|
tirx.stmt_functor.post_order_visit(expr, _visit_expr)
|
|
|
|
def _update_shape(
|
|
self,
|
|
shape: List[tirx.Expr], # noqa: UP006
|
|
undefined_var_remap: Dict[tirx.Var, tirx.Var], # noqa: UP006
|
|
) -> List[tirx.Expr]: # noqa: UP006
|
|
new_shape = []
|
|
for v in shape:
|
|
self._copy_undefined_var(v, undefined_var_remap)
|
|
new_shape.append(tirx.stmt_functor.substitute(v, undefined_var_remap))
|
|
return new_shape
|
|
|
|
|
|
def _extract_pipeline_stages(
|
|
func: relax.Function,
|
|
) -> Tuple[List[List[relax.Binding]], List[List[relax.Var]], List[List[relax.Var]]]: # noqa: UP006
|
|
pipeline_stages: List[List[relax.Binding]] = [] # noqa: UP006
|
|
stage_send_vars: List[List[relax.Var]] = [] # noqa: UP006
|
|
stage_receive_vars: List[List[relax.Var]] = [] # noqa: UP006
|
|
|
|
# Requiring that the function has only one body block which is a dataflow block
|
|
assert isinstance(func.body, relax.SeqExpr)
|
|
assert len(func.body.blocks) == 1
|
|
assert isinstance(func.body.blocks[0], relax.DataflowBlock)
|
|
bindings = func.body.blocks[0].bindings
|
|
|
|
boundary_var = None
|
|
current_stage_bindings: List[relax.Binding] = [] # noqa: UP006
|
|
current_stage_receive_vars: List[relax.Var] = [] # noqa: UP006
|
|
for binding in bindings:
|
|
if (
|
|
isinstance(binding, relax.VarBinding)
|
|
and isinstance(binding.value, relax.Call)
|
|
and binding.value.op == tvm.ir.Op.get("relax.call_pure_packed")
|
|
and binding.value.args[0].global_symbol == "mlc.pipeline_parallel_stage_boundary"
|
|
):
|
|
assert len(current_stage_bindings) > 0
|
|
pipeline_stages.append(current_stage_bindings)
|
|
assert all(receive_var is not None for receive_var in current_stage_receive_vars)
|
|
stage_receive_vars.append(current_stage_receive_vars)
|
|
args = binding.value.args[1:]
|
|
assert len(args) >= 1 and all(isinstance(arg, relax.Var) for arg in args)
|
|
stage_send_vars.append(list(args))
|
|
|
|
boundary_var = binding.var
|
|
current_stage_bindings = []
|
|
current_stage_receive_vars = [boundary_var] if len(args) == 1 else [None for _ in args]
|
|
elif (
|
|
isinstance(binding, relax.VarBinding)
|
|
and isinstance(binding.value, relax.TupleGetItem)
|
|
and binding.value.tuple_value.same_as(boundary_var)
|
|
):
|
|
current_stage_receive_vars[binding.value.index] = binding.var
|
|
else:
|
|
current_stage_bindings.append(binding)
|
|
|
|
assert len(current_stage_bindings) > 0
|
|
pipeline_stages.append(current_stage_bindings)
|
|
assert all(receive_var is not None for receive_var in current_stage_receive_vars)
|
|
stage_receive_vars.append(current_stage_receive_vars)
|
|
stage_send_vars.append([])
|
|
|
|
return pipeline_stages, stage_send_vars, stage_receive_vars
|
|
|
|
|
|
def _analyze_required_func_params(
|
|
pipeline_stages: List[List[relax.Binding]], # noqa: UP006
|
|
func_params: List[relax.Var], # noqa: UP006
|
|
) -> List[List[relax.Var]]: # noqa: UP006
|
|
analyzer = _RequiredFuncParamAnalyzer(func_params)
|
|
required_func_params: List[List[relax.Var]] = [] # noqa: UP006
|
|
for stage_bindings in pipeline_stages:
|
|
required_params: List[relax.Var] # noqa: UP006
|
|
required_params = analyzer.run(stage_bindings)
|
|
required_func_params.append(required_params)
|
|
return required_func_params
|
|
|
|
|
|
@visitor
|
|
class _RequiredFuncParamAnalyzer(PyExprVisitor):
|
|
"""The IR visitor which analyzes the required func parameters in each pipeline stage."""
|
|
|
|
def __init__(self, func_params: List[relax.Var]) -> None: # noqa: UP006
|
|
self.func_params = set(func_params)
|
|
self.required_params: List[relax.Var] # noqa: UP006
|
|
|
|
def run(self, stage_bindings: List[relax.Binding]) -> List[relax.Var]: # noqa: UP006
|
|
"""Entry point of the visitor."""
|
|
self.required_params = []
|
|
for binding in stage_bindings:
|
|
self.visit_binding(binding)
|
|
return self.required_params
|
|
|
|
def visit_var_(self, var: relax.Var) -> None:
|
|
if var in self.func_params:
|
|
if var not in self.required_params:
|
|
self.required_params.append(var)
|