wehub-resource-sync
126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name
|
|
# ruff: noqa: RUF005
|
|
"""Default legalization function for ccl operators."""
|
|
|
|
from tvm import arith, tirx, topi
|
|
from tvm.ir import Call
|
|
|
|
from ...block_builder import BlockBuilder
|
|
from ...expr import Expr, ShapeExpr
|
|
from ...op import call_dps_packed
|
|
from ...type import ShapeType, TensorType
|
|
from .common import register_legalize
|
|
|
|
|
|
@register_legalize("relax.ccl.allreduce")
|
|
def _allreduce(_bb: BlockBuilder, call: Call) -> Expr:
|
|
op_type_str = call.attrs.op_type
|
|
op_type_map = {
|
|
"sum": 0,
|
|
"prod": 1,
|
|
"min": 2,
|
|
"max": 3,
|
|
"avg": 4,
|
|
}
|
|
if op_type_str not in op_type_map:
|
|
raise ValueError(
|
|
f"Unsupported reduction operation: {op_type_str}. "
|
|
f"Supported operations are {op_type_map.keys()}."
|
|
)
|
|
return call_dps_packed(
|
|
"runtime.disco.allreduce",
|
|
[call.args[0], ShapeExpr([op_type_map[op_type_str]]), call.attrs.in_group],
|
|
out_ty=call.args[0].ty,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.ccl.allgather")
|
|
def _allgather(_bb: BlockBuilder, call: Call) -> Expr:
|
|
output_shape = []
|
|
arg_ty = call.args[0].ty
|
|
assert isinstance(arg_ty, TensorType), "The input type of allgather should be TensorType."
|
|
assert isinstance(arg_ty.shape.ty, ShapeType)
|
|
arg_shape = arg_ty.shape.ty
|
|
for i, shape_value in enumerate(arg_shape.values):
|
|
if i == 0:
|
|
output_shape.append(shape_value * call.attrs.num_workers)
|
|
else:
|
|
output_shape.append(shape_value)
|
|
return call_dps_packed(
|
|
"runtime.disco.allgather",
|
|
[call.args[0], call.attrs.in_group],
|
|
out_ty=TensorType(
|
|
shape=output_shape,
|
|
dtype=arg_ty.dtype,
|
|
vdevice=arg_ty.vdevice,
|
|
),
|
|
)
|
|
|
|
|
|
@register_legalize("relax.ccl.broadcast_from_worker0")
|
|
def _broadcast_from_worker0(_bb: BlockBuilder, call: Call) -> Expr:
|
|
return call_dps_packed(
|
|
"runtime.disco.broadcast_from_worker0",
|
|
[call.args[0], False],
|
|
out_ty=call.args[0].ty,
|
|
)
|
|
|
|
|
|
# Since collective communication ops are performed on contiguous memory,
|
|
# we need to reshape and transpose the input tensor to make sharding dimension in the highest order
|
|
def _transpose_for_ccl(_bb: BlockBuilder, expr: Expr, axis: int, num_workers: int):
|
|
assert isinstance(expr.ty, TensorType), "The input type should be TensorType."
|
|
assert isinstance(expr.ty.shape.ty, ShapeType)
|
|
arg_shape = expr.ty.shape.ty
|
|
new_shape = []
|
|
for i, shape_value in enumerate(arg_shape.values):
|
|
if i == axis:
|
|
modulo = arith.Analyzer().simplify(shape_value % num_workers)
|
|
assert modulo == 0, (
|
|
f"scatter_from_worker0 expects the size of axis {axis} of input tensor "
|
|
"to be divisible by num_workers. However, the axis 0 of input tensor "
|
|
f"is {shape_value} while num_workers is {num_workers}"
|
|
)
|
|
new_shape.append(num_workers)
|
|
new_shape.append(tirx.div(shape_value, num_workers))
|
|
else:
|
|
new_shape.append(shape_value)
|
|
reshape_var = _bb.emit_te(topi.reshape, expr, new_shape)
|
|
if axis == 0:
|
|
return reshape_var
|
|
permute_order = [axis] + list(range(axis)) + list(range(axis + 1, len(new_shape)))
|
|
transpose_var = _bb.emit_te(topi.transpose, reshape_var, permute_order)
|
|
return transpose_var
|
|
|
|
|
|
@register_legalize("relax.ccl.scatter_from_worker0")
|
|
def _scatter_from_worker0(_bb: BlockBuilder, call: Call) -> Expr:
|
|
transpose_var = _transpose_for_ccl(_bb, call.args[0], call.attrs.axis, call.attrs.num_workers)
|
|
output_shape = transpose_var.ty.shape.ty.values
|
|
output_shape = output_shape[1:]
|
|
return call_dps_packed(
|
|
"runtime.disco.scatter_from_worker0",
|
|
[transpose_var, False],
|
|
out_ty=TensorType(
|
|
shape=output_shape,
|
|
dtype=call.args[0].ty.dtype,
|
|
vdevice=call.args[0].ty.vdevice,
|
|
),
|
|
)
|