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,148 @@
+# 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.
+"""Rewrite all-reduce operation to customized all-reduce impl with IPC memory.
+The pass is written in Python for experiment, fast development.
+"""
+
+import tvm
+from tvm import relax
+from tvm.ir.module import IRModule
+from tvm.relax.expr import Expr, Var
+from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
+
+
+@tvm.transform.module_pass(opt_level=0, name="IPCAllReduceRewrite")
+class IPCAllReduceRewrite:
+    """Rewrite all-reduce operation to customized all-reduce impl with IPC memory."""
+
+    def __init__(self, allreduce_strategy: int) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        allreduce_strategy : int
+            The all-reduce strategy. Only "1" and "2" are supported.
+            "1" stands for one-shot, and "2" stands for two-shot.
+        """
+        self.allreduce_strategy = allreduce_strategy
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        fcustom_allreduce = tvm.get_global_func(
+            "runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True
+        )
+        if fcustom_allreduce is None:
+            # Customized allreduce is not available.
+            return mod
+
+        binding_replacement_map = _Visitor(self.allreduce_strategy).visit(mod)
+        return _Rewriter(mod, binding_replacement_map).transform()
+
+
+@visitor
+class _Visitor(PyExprVisitor):  # pylint: disable=abstract-method
+    def __init__(self, allreduce_strategy: int) -> None:
+        self.allreduce_strategy = allreduce_strategy
+        self.alloc_map: dict[Var, relax.Call] = {}
+        self.binding_replacement_map: dict[relax.Expr, relax.Expr] = {}
+        self.builtin_alloc_tensor_op = tvm.ir.Op.get("relax.builtin.alloc_tensor")
+        self.reshape_op = tvm.ir.Op.get("relax.reshape")
+
+    def visit(self, mod: IRModule) -> dict[relax.Expr, relax.Expr]:
+        """Entry point"""
+        for _, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                self.alloc_map.clear()
+                self.visit_expr(func)
+        return self.binding_replacement_map
+
+    def visit_var_binding_(self, binding: relax.VarBinding):
+        super().visit_var_binding_(binding)
+        if (
+            isinstance(binding.value, relax.Call)
+            and binding.value.op == self.builtin_alloc_tensor_op
+        ):
+            self.alloc_map[binding.var] = binding.value
+        elif isinstance(binding.value, relax.Var) and binding.value in self.alloc_map:
+            self.alloc_map[binding.var] = self.alloc_map[binding.value]
+        elif (
+            isinstance(binding.value, relax.Call)
+            and binding.value.op == self.reshape_op
+            and binding.value.args[0] in self.alloc_map
+        ):
+            self.alloc_map[binding.var] = self.alloc_map[binding.value.args[0]]
+
+    def visit_call_(self, call: relax.Call) -> None:  # pylint: disable=arguments-renamed
+        if (
+            not isinstance(call.op, relax.ExternFunc)
+            or call.op.global_symbol != "runtime.disco.allreduce"
+            or call.args[1].values[0] != 0
+        ):
+            # Return if the call is not a summation all-reduce.
+            return
+
+        assert len(call.args) == 4
+        allreduce_input, _strategy, _ingroup, allreduce_output = call.args
+        alloc_tensor = self.alloc_map.get(allreduce_input, None)
+        if alloc_tensor is None or alloc_tensor.args[3].value != "global":
+            # Return if the allocation of all-reduce input is not recorded,
+            # or the scope of the allocation is not global.
+            return
+
+        # Set the scope of the alloc_tensor to IPC memory.
+        alloc_tensor = self.alloc_map[allreduce_input]
+        self.binding_replacement_map[alloc_tensor] = relax.op.builtin.alloc_tensor(
+            alloc_tensor.args[0],
+            alloc_tensor.args[1],
+            alloc_tensor.args[2],
+            relax.StringImm("ipc_memory"),
+        )
+
+        self.binding_replacement_map[call] = relax.Call(
+            relax.ExternFunc("runtime.disco.cuda_ipc.custom_allreduce"),
+            # The "cuda_ipc.custom_allreduce" implementation does not
+            # yet support num_groups>1, and therefore does not use the
+            # `in_group` argument.
+            [allreduce_input, relax.prim_value(self.allreduce_strategy), allreduce_output],
+        )
+
+
+@mutator
+class _Rewriter(PyExprMutator):
+    """Rewrite the IRModule according to the binding replacement provided by the visitor."""
+
+    def __init__(
+        self, mod: IRModule, binding_replacement_map: dict[relax.Expr, relax.Expr]
+    ) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.binding_replacement_map = binding_replacement_map
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, relax.Function):
+                updated_func = self.visit_expr(func)
+                self.builder_.update_func(g_var, updated_func)
+        return self.builder_.get()
+
+    def visit_call_(self, call: relax.Call) -> Expr:  # pylint: disable=arguments-renamed
+        return (
+            super().visit_call_(self.binding_replacement_map[call])
+            if call in self.binding_replacement_map
+            else super().visit_call_(call)
+        )