Method allreduce

python/tvm/runtime/disco/session.py:484–513

Perform an allreduce operation on an array. Parameters ---------- array : DRef The array to be reduced. op : str = "sum" The reduce operation to be performed. Available options are: - "sum" - "prod" - "min"

(
        self,
        src: DRef,
        dst: DRef,
        op: str = "sum",  # pylint: disable=invalid-name
        in_group: bool = True,
    )

        func(from_array, in_group, to_array)

    def allreduce(
        self,
        src: DRef,
        dst: DRef,
        op: str = "sum",  # pylint: disable=invalid-name
        in_group: bool = True,
    ) -> DRef:
        """Perform an allreduce operation on an array.

        Parameters
        ----------
        array : DRef
            The array to be reduced.

        op : str = "sum"
            The reduce operation to be performed. Available options are:
            - "sum"
            - "prod"
            - "min"
            - "max"
            - "avg"

        in_group : bool
            Whether the reduce operation performs globally or in group as default.
        """
        if op not in REDUCE_OPS:
            raise ValueError(f"Unsupported reduce op: {op}. Available ops are: {REDUCE_OPS.keys()}")
        op = Shape([REDUCE_OPS[op]])
        func = self._get_cached_method("runtime.disco.allreduce")
        func(src, op, in_group, dst)

    def allgather(
        self,