chore: import upstream snapshot with attribution

python/tvm/relax/distributed/__init__.py

@@ -0,0 +1,24 @@
+# isort: skip_file
+# 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.
+
+"""The infrastructure for distributed inference on Relax."""
+
+from .global_info import DeviceMesh, device_mesh
+from .type import Placement, DTensorType, PlacementSpec
+
+from . import transform

python/tvm/relax/distributed/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs for tvm.relax.distributed"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.distributed", __name__)

python/tvm/relax/distributed/global_info.py

@@ -0,0 +1,66 @@
+# 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=redefined-builtin, invalid-name
+"""Global Info Data structures for distributed tensor."""
+
+import tvm_ffi
+from tvm_ffi import Shape
+
+from tvm.ir import Range
+from tvm.ir.global_info import GlobalInfo
+
+from . import _ffi_api as ffi
+
+
+@tvm_ffi.register_object("relax.distributed.DeviceMesh")
+class DeviceMesh(GlobalInfo):
+    """Device mesh express a view of topology of devices,
+       represented by an n-d matrix of device ids.
+
+    Parameters
+    ----------
+    shape: Union[Shape, List[int], Tuple[int]]
+        Logical shape of device mesh
+    device_ids: Union[List[int], Range]
+        Represents the device id in the mesh
+    """
+
+    def __init__(self, shape: Shape | list[int] | tuple[int], device_ids: list[int] | Range):
+        if not isinstance(shape, Shape):
+            shape = Shape(shape)
+        device_range = None
+        if isinstance(device_ids, Range):
+            device_range = device_ids
+            device_ids = []
+        self.__init_handle_by_constructor__(ffi.DeviceMesh, shape, device_ids, device_range)  # type: ignore
+
+
+def device_mesh(shape: Shape, device_ids: list[int] | Range) -> DeviceMesh:
+    """Create a device mesh expression.
+    Parameters
+    ----------
+    shape : Shape
+        The shape of the device mesh.
+    device_ids: Union[List[int], Range]
+        Represents the device id in the mesh
+
+    Returns
+    -------
+    res : DeviceMesh
+        The device mesh.
+    """
+    return DeviceMesh(shape, device_ids)  # pylint: disable=no-member # type: ignore

python/tvm/relax/distributed/transform/__init__.py

@@ -0,0 +1,25 @@
+# isort: skip_file
+# 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.
+"""Relax distributed-related transformations."""
+
+from .transform import (
+    PropagateSharding,
+    LowerGlobalViewToLocalView,
+    LegalizeRedistribute,
+    LowerDistIR,
+)

python/tvm/relax/distributed/transform/_ffi_api.py

@@ -0,0 +1,20 @@
+# 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
+"""FFI APIs for tvm.relax.distributed.transform"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.distributed.transform", __name__)

python/tvm/relax/distributed/transform/transform.py

@@ -0,0 +1,68 @@
+# 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
+"""Relax distributed-related transformation passes."""
+
+import tvm.ir
+
+from . import _ffi_api
+
+
+def PropagateSharding() -> tvm.ir.transform.Pass:
+    """Propagate sharding information.
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.PropagateSharding()  # type: ignore
+
+
+def LowerGlobalViewToLocalView() -> tvm.ir.transform.Pass:
+    """Lower global view TIR to local view
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.LowerGlobalViewToLocalView()  # type: ignore
+
+
+def LegalizeRedistribute() -> tvm.ir.transform.Pass:
+    """Legalize redistribute op to ccl op.
+    S->R: R.ccl.allgather
+    R->S: R.dist.redistribute_replica_to_shard
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.LegalizeRedistribute()  # type: ignore
+
+
+def LowerDistIR() -> tvm.ir.transform.Pass:
+    """Lower DistIR to Relax
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.LowerDistIR()  # type: ignore

python/tvm/relax/distributed/type.py

@@ -0,0 +1,146 @@
+# 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=redefined-builtin, invalid-name
+"""Types for distributed tensor."""
+
+import enum
+
+import tvm_ffi
+
+from tvm.ir import Span
+from tvm.relax.type import TensorType, Type
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .global_info import DeviceMesh
+
+
+class PlacementSpecKind(enum.IntEnum):
+    kSharding = 0
+    kReplica = 1
+
+
+@tvm_ffi.register_object("relax.distributed.PlacementSpec")
+class PlacementSpec(Object):
+    """Describes how data is distributed in one dimension of the device mesh
+
+    Parameters
+    ----------
+    axis: int
+        If the kind is sharding, this value represents the tensor dimension to shard.
+        otherwise, axis is -1
+    kind: PlacementSpecKind
+        The kind of placement spec. Possible values: kSharding and kReplica.
+    """
+
+    axis: int
+    kind: PlacementSpecKind
+
+    def __init__(self, *args, **kwargs):
+        raise RuntimeError("PlacementSpec is not intended to be constructed directly, ")
+
+    @staticmethod
+    def sharding(axis: int) -> "PlacementSpec":
+        """Create a sharding placement spec
+
+        Parameters
+        ----------
+        axis: int
+            The tensor dimension to shard.
+
+        Returns
+        -------
+        placement_spec: PlacementSpec
+            The placement spec.
+        """
+        return _ffi_api.Sharding(axis)
+
+    @staticmethod
+    def replica() -> "PlacementSpec":
+        """Create a replica placement spec
+
+        Returns
+        -------
+        placement_spec: PlacementSpec
+            The placement spec.
+        """
+        return _ffi_api.Replica()
+
+
+@tvm_ffi.register_object("relax.distributed.Placement")
+class Placement(Object):
+    """Describes how data is distributed in each dimension of the device mesh
+
+    Parameters
+    ----------
+    dim_specs: List[PlacementSpec]
+        The placement spec for each dimension of the device mesh.
+    """
+
+    def __init__(self, dim_specs: list[PlacementSpec]):
+        self.__init_handle_by_constructor__(_ffi_api.Placement, dim_specs)  # type: ignore
+
+    @staticmethod
+    def from_text(text: str) -> "Placement":
+        """Create a placement from a text string.
+
+        Parameters
+        ----------
+        text: str
+            The text string.
+
+        Returns
+        -------
+        placement: Placement
+            The placement.
+        """
+        return _ffi_api.PlacementFromText(text)
+
+
+@tvm_ffi.register_object("relax.DTensorType")
+class DTensorType(Type):
+    """Type of a Distributed Tensor value.
+
+    Parameters
+    ----------
+    tensor_ty: TensorType
+        The tensor type carried by the distributed tensor.
+    device_mesh: DeviceMesh
+        The device mesh of the tensor.
+    placement: Placement
+        The placement of the tensor among the device mesh
+
+    """
+
+    tensor_ty: TensorType
+    device_mesh: DeviceMesh
+    placement: Placement
+
+    def __init__(
+        self,
+        tensor_ty: TensorType,
+        device_mesh: DeviceMesh,
+        placement: Placement,
+        span: Span = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.DTensorType,
+            tensor_ty,
+            device_mesh,
+            placement,
+            span,  # type: ignore
+        )