chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,106 @@
+# 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,missing-docstring, invalid-name, unused-import, redefined-outer-name
+# ruff: noqa: F401
+
+from typing import Any, Optional, Union
+
+from tvm.ir import Range
+from tvm.relax import TensorType
+from tvm.relax.distributed import DeviceMesh, DTensorType, Placement, device_mesh
+from tvm.relax.script.builder.distributed import (
+    annotate_sharding,
+    call_tir,
+    call_tir_local_view,
+    const,
+    redistribute,
+    redistribute_replica_to_shard,
+)
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder.ir import IRModuleFrame
+from tvm.tirx import Expr
+
+from .entry import TensorProxy, TypeProxy
+
+############################### R.DTensor ###############################
+
+
+class DTensorProxy(TypeProxy):
+    tensor_ty_proxy: TensorProxy
+    device_mesh: DeviceMesh
+    placement: Placement
+
+    def __init__(
+        self,
+        tensor_ty_proxy: TensorProxy,
+        device_mesh: DeviceMesh,
+        placement: Placement,
+    ) -> None:
+        self.device_mesh = device_mesh
+        self.placement = placement
+        self.tensor_ty_proxy = tensor_ty_proxy
+        super().__init__()
+
+    def get_symbolic_vars(self) -> set[str]:
+        return self.tensor_ty_proxy.get_symbolic_vars()
+
+    def as_ty(self, dict_globals: dict[str, Any] | None = None) -> DTensorType:
+        return DTensorType(
+            self.tensor_ty_proxy.as_ty(dict_globals),
+            self.device_mesh,
+            self.placement,
+        )
+
+
+def DTensor(
+    shape: list[Expr | str] | None = None,
+    dtype: str | None = None,
+    device_mesh: DeviceMesh | str = DeviceMesh([], Range(0, 1)),
+    placement: Placement | str = "",
+    *,
+    ndim: int = -1,
+) -> DTensorProxy:
+    # scalar tensor case
+    if shape is not None and len(shape) == 0:
+        shape = []
+    if isinstance(shape, str) and dtype is None:
+        dtype = shape
+        shape = None
+
+    if shape is not None and not isinstance(shape, tuple | list):
+        raise ValueError(f"shape must be a list or tuple, but got: {shape}")
+    if isinstance(device_mesh, str):
+        if not IRBuilder.is_in_scope():
+            return (
+                DTensorProxy(
+                    TensorProxy(shape, dtype, None, ndim), DeviceMesh([], Range(0, 1)), ""
+                ),
+            )
+        name, index = device_mesh.split("[")
+        index = int(index[:-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)
+    if isinstance(placement, str):
+        placement = Placement.from_text(placement)
+    return DTensorProxy(TensorProxy(shape, dtype, None, ndim), device_mesh, placement)
+
+
+__all__ = ["DTensor", "device_mesh"]