chore: import upstream snapshot with attribution

python/tvm/relax/script/builder/distributed/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+# pylint: disable=unused-import
+"""Package tvm.script.ir_builder.relax.distributed"""
+
+from .ir import *  # pylint: disable=wildcard-import,redefined-builtin

python/tvm/relax/script/builder/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.script.ir_builder.relax.distributed"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("script.ir_builder.relax.distributed", __name__)  # pylint: disable=protected-access

python/tvm/relax/script/builder/distributed/ir.py

@@ -0,0 +1,170 @@
+# 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, wrong-import-order, no-member, invalid-name, unused-import
+# ruff: noqa: F401
+
+"""IRBuilder for distributed Relax dialect"""
+
+from numbers import Number
+from typing import Optional, Union
+
+import numpy as _np  # type: ignore
+
+import tvm
+from tvm import base as _base
+from tvm.ir import Call
+from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
+from tvm.relax.expr import Constant, Expr, ExternFunc, ShapeExpr
+from tvm.relax.expr import Tuple as RxTuple
+from tvm.relax.op.distributed import (
+    annotate_sharding as _annotate_sharding,
+)
+from tvm.relax.op.distributed import (
+    call_tir_local_view,
+    redistribute_replica_to_shard,
+)
+from tvm.relax.op.distributed import (
+    redistribute as _redistribute,
+)
+from tvm.relax.script.builder.ir import py_str
+from tvm.relax.utils import convert_to_expr
+from tvm.runtime import _tensor
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder.ir import IRModuleFrame
+
+from . import _ffi_api
+
+
+def call_tir(
+    func: str | Expr,
+    args: Expr,
+    out_ty: DTensorType | list[DTensorType],
+    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
+) -> Call:
+    """Distributed version of call_tir
+
+    Parameters:
+    ----------
+    func : Union[str, Expr]
+        The destination-passing-style function, can be ExternFunc or PrimFunc.
+
+    args : Expr
+        The input arguments.
+
+    out_ty : Union[DTensorType, List[DTensorType]]
+        The type information of the call_tir output.
+        It should be a single or a list of DTensorType. Each one denotes the
+        type information of a returned distributed tensor.
+
+    tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
+        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
+
+    Returns
+    -------
+    ret: Call
+        A call node for the call_tir operator.
+    """
+    if isinstance(func, str):
+        func = ExternFunc(func)
+
+    if isinstance(args, tuple | list):
+        args = RxTuple([convert_to_expr(a) for a in args])
+    elif isinstance(args, Expr) and not isinstance(args, RxTuple):  # type: ignore
+        args = RxTuple((args,))
+
+    if not isinstance(out_ty, list):
+        out_ty = [out_ty]
+
+    if isinstance(tir_vars, list | tuple):
+        tir_vars = ShapeExpr(tir_vars)
+
+    return _ffi_api.call_tir_dist(func, args, out_ty, tir_vars)  # type: ignore
+
+
+def const(
+    value: bool | int | float | _np.ndarray | tvm.runtime.Tensor,
+    ty: DTensorType,
+) -> Constant:
+    """Create a constant value.
+
+    Parameters
+    ----------
+    value: Union[bool, int, float, numpy.ndarray, tvm.runtime.Tensor]
+        The constant value.
+
+    dtype: Optional[str]
+        The data type of the resulting constant.
+
+    Note
+    ----
+    When dtype is None, we use the following rule:
+
+    - int maps to "int32"
+    - float maps to "float32"
+    - bool maps to "bool"
+    - other using the same default rule as numpy.
+    """
+    ty = tvm.runtime.convert(ty)
+    if not isinstance(ty, DTensorType):
+        raise TypeError("ty needs to be an instance of DTensorType. ")
+    dtype = str(ty.tensor_ty.dtype)
+    if isinstance(value, Number | (bool | list)):
+        value = _np.array(value, dtype=dtype)
+
+    if isinstance(value, _np.ndarray | _np.generic):
+        if dtype is not None:
+            value = value.astype(dtype)
+        value = _tensor.tensor(value)
+
+    if not isinstance(value, _tensor.Tensor):
+        raise ValueError("value has to be scalar or Tensor")
+
+    return Constant(value, ty)
+
+
+def _lookup_device_mesh(device_mesh_str: py_str) -> DeviceMesh:
+    if not IRBuilder.is_in_scope():
+        raise ValueError("device_mesh cannot be found in global info")
+    name, index_str = device_mesh_str.split("[")
+    index = int(index_str[:-1])
+    frames = IRBuilder.current().frames
+    for f in frames:
+        if isinstance(f, IRModuleFrame):
+            device_mesh = f.global_infos[name][index]
+            break
+    assert isinstance(device_mesh, DeviceMesh)
+    return device_mesh
+
+
+def annotate_sharding(
+    value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement
+) -> Expr:
+    if isinstance(device_mesh, py_str):
+        device_mesh = _lookup_device_mesh(device_mesh)
+    if isinstance(placement, py_str):
+        placement = Placement.from_text(placement)
+    return _annotate_sharding(value, device_mesh, placement)
+
+
+def redistribute(
+    value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement
+) -> Expr:
+    if isinstance(device_mesh, py_str):
+        device_mesh = _lookup_device_mesh(device_mesh)
+    if isinstance(placement, py_str):
+        placement = Placement.from_text(placement)
+    return _redistribute(value, device_mesh, placement)