This is an automated email from the ASF dual-hosted git repository.
yongwww pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new a6adaae5ef [Unity][DistIR] LowerDistIR (#16169)
a6adaae5ef is described below
commit a6adaae5ef30e5d4b31d6e549c6a9848e2a05e23
Author: Hongyi Jin <[email protected]>
AuthorDate: Wed Nov 29 11:19:41 2023 -0500
[Unity][DistIR] LowerDistIR (#16169)
* lower distir
* format
* format
* add whitespace
* fix lint
* fix warning
---
include/tvm/relax/distributed/transform.h | 6 +
python/tvm/relax/distributed/transform/__init__.py | 7 +-
.../tvm/relax/distributed/transform/transform.py | 11 +
src/relax/distributed/transform/lower_distir.cc | 271 ++++++++++++++
.../test_distributed_transform_lower_distir.py | 396 +++++++++++++++++++++
5 files changed, 690 insertions(+), 1 deletion(-)
diff --git a/include/tvm/relax/distributed/transform.h
b/include/tvm/relax/distributed/transform.h
index 2ac3492165..31727b181e 100644
--- a/include/tvm/relax/distributed/transform.h
+++ b/include/tvm/relax/distributed/transform.h
@@ -62,6 +62,12 @@ TVM_DLL Pass LowerGlobalViewToLocalView();
*/
TVM_DLL Pass LegalizeRedistribute();
+/*!
+ * \brief Lower DistIR to Relax
+ *
+ * \return The Pass.
+ */
+TVM_DLL Pass LowerDistIR();
} // namespace transform
} // namespace distributed
} // namespace relax
diff --git a/python/tvm/relax/distributed/transform/__init__.py
b/python/tvm/relax/distributed/transform/__init__.py
index 573b7d599d..a6158eaedb 100644
--- a/python/tvm/relax/distributed/transform/__init__.py
+++ b/python/tvm/relax/distributed/transform/__init__.py
@@ -16,4 +16,9 @@
# under the License.
"""Relax distributed-related transformations. """
-from .transform import PropagateSharding, LowerGlobalViewToLocalView,
LegalizeRedistribute
+from .transform import (
+ PropagateSharding,
+ LowerGlobalViewToLocalView,
+ LegalizeRedistribute,
+ LowerDistIR,
+)
diff --git a/python/tvm/relax/distributed/transform/transform.py
b/python/tvm/relax/distributed/transform/transform.py
index 36aa1fe9b6..e64366c79a 100644
--- a/python/tvm/relax/distributed/transform/transform.py
+++ b/python/tvm/relax/distributed/transform/transform.py
@@ -54,3 +54,14 @@ def LegalizeRedistribute() -> tvm.ir.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
diff --git a/src/relax/distributed/transform/lower_distir.cc
b/src/relax/distributed/transform/lower_distir.cc
new file mode 100644
index 0000000000..2bfd38b9af
--- /dev/null
+++ b/src/relax/distributed/transform/lower_distir.cc
@@ -0,0 +1,271 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/distributed/transform/lower_distir.cc
+ * \brief Pass for lowering DistIR into Relax
+ * This pass assumes all the TensorIR functions are in local view,
+ * so the pass only handles sharding relax tensor shape and
+ * inserting necessary broadcast and scatter for inputs.
+ */
+
+#include <tvm/relax/attrs/ccl.h>
+#include <tvm/relax/distributed/axis_group_graph.h>
+#include <tvm/relax/distributed/transform.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/tir/stmt_functor.h>
+
+#include "../../../tir/schedule/transform.h"
+#include "../../op/ccl/ccl.h"
+#include "../../op/tensor/manipulate.h"
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+namespace distributed {
+
+class DistIRSharder : public ExprMutator {
+ public:
+ static IRModule LowerDistIR(IRModule mod) { return
DistIRSharder(mod).Lower(); }
+
+ private:
+ explicit DistIRSharder(IRModule mod) : ExprMutator(mod) {}
+
+ IRModule Lower() {
+ auto mod = builder_->GetContextIRModule();
+ for (const auto& [gv, base_func] : mod->functions) {
+ const auto* func_ = base_func.as<FunctionNode>();
+ if (func_ == nullptr || !IsDistIRFunc(GetRef<Function>(func_))) {
+ continue;
+ }
+ Function func = RewriteFunction(GetRef<Function>(func_));
+ builder_->UpdateFunction(gv, func);
+ }
+ return builder_->GetContextIRModule();
+ }
+
+ ShapeExpr ShardShape(ShapeExpr orig_shape, DeviceMesh device_mesh, Placement
placement) {
+ ShapeTuple device_mesh_shape = device_mesh->shape;
+ Array<PrimExpr> new_tensor_shape_value = orig_shape->values;
+ for (int i = 0; i < static_cast<int>(device_mesh_shape.size()); i++) {
+ if (placement->dim_specs[i]->kind == PlacementSpecKind::kSharding) {
+ int shard_size = device_mesh_shape[i];
+ int axis = placement->dim_specs[i]->axis;
+ new_tensor_shape_value.Set(axis, floordiv(orig_shape->values[axis],
shard_size));
+ }
+ }
+ return ShapeExpr(new_tensor_shape_value);
+ }
+
+ TensorStructInfo ShardDTensorSinfo(DTensorStructInfo orig_sinfo) {
+ TensorStructInfo tensor_sinfo = orig_sinfo->tensor_sinfo;
+ ICHECK(tensor_sinfo->shape);
+ const auto* orig_shape = tensor_sinfo->shape.as<ShapeExprNode>();
+ auto new_tensor_sinfo =
make_object<TensorStructInfoNode>(*tensor_sinfo.get());
+ new_tensor_sinfo->shape =
+ ShardShape(GetRef<ShapeExpr>(orig_shape), orig_sinfo->device_mesh,
orig_sinfo->placement);
+ return TensorStructInfo(new_tensor_sinfo);
+ }
+
+ StructInfo ConvertSinfo(StructInfo orig_sinfo, bool shard_shape) {
+ if (const auto* dtensor_sinfo = orig_sinfo.as<DTensorStructInfoNode>()) {
+ if (shard_shape) {
+ return ShardDTensorSinfo(GetRef<DTensorStructInfo>(dtensor_sinfo));
+ } else {
+ return dtensor_sinfo->tensor_sinfo;
+ }
+ } else if (const auto* tuple_sinfo = orig_sinfo.as<TupleStructInfoNode>())
{
+ Array<StructInfo> new_fields;
+ for (const auto& field_sinfo : tuple_sinfo->fields) {
+ if (const auto* dtensor_sinfo =
field_sinfo.as<DTensorStructInfoNode>()) {
+ if (shard_shape) {
+
new_fields.push_back(ShardDTensorSinfo(GetRef<DTensorStructInfo>(dtensor_sinfo)));
+ } else {
+ new_fields.push_back(dtensor_sinfo->tensor_sinfo);
+ }
+ } else {
+ new_fields.push_back(field_sinfo);
+ }
+ }
+ return TupleStructInfo(new_fields);
+ } else {
+ return orig_sinfo;
+ }
+ }
+
+ Expr ShardInputParamTensorAndConstant(Expr input) {
+ ICHECK(input->struct_info_);
+ StructInfo old_sinfo = GetStructInfo(input);
+ StructInfo new_sinfo = ConvertSinfo(old_sinfo, false);
+ if (const auto* var = input.as<VarNode>()) {
+ Var new_param(var->name_hint(), new_sinfo);
+ return new_param;
+ } else if (const auto* constant = input.as<ConstantNode>()) {
+ for (const auto& spec :
Downcast<DTensorStructInfo>(old_sinfo)->placement->dim_specs) {
+ ICHECK(spec->kind == PlacementSpecKind::kReplica);
+ }
+ Constant new_constant(constant->data, new_sinfo);
+ return new_constant;
+ } else {
+ LOG(FATAL) << "Cannot shard tensor which is not Var or Constant: " <<
input;
+ throw;
+ }
+ }
+
+ void EmitBroadcastOrScatter(Expr old_expr, Expr new_expr, DTensorStructInfo
dtensor_sinfo) {
+ // FIXME: this is a hack that only works for 1d device mesh
+ ICHECK(dtensor_sinfo->device_mesh->shape.size() == 1);
+ PlacementSpec sharding_spec = dtensor_sinfo->placement->dim_specs[0];
+ if (sharding_spec->kind == PlacementSpecKind::kReplica) {
+ Var new_var = builder_->Emit(broadcast_from_worker0(new_expr));
+ if (const auto* var = old_expr.as<VarNode>()) {
+ var_remap_[var->vid] = new_var;
+ } else {
+ tuple_getitem_remap_[Downcast<TupleGetItem>(old_expr)] = new_var;
+ }
+ } else if (sharding_spec->kind == PlacementSpecKind::kSharding) {
+ Var scatter_var = builder_->Emit(scatter_from_worker0(
+ new_expr, dtensor_sinfo->device_mesh->shape[0],
sharding_spec->axis));
+ if (const auto* var = old_expr.as<VarNode>()) {
+ var_remap_[var->vid] = scatter_var;
+ } else {
+ tuple_getitem_remap_[Downcast<TupleGetItem>(old_expr)] = scatter_var;
+ }
+ } else {
+ LOG(FATAL) << "Unsupported placement spec";
+ }
+ }
+
+ void InputPreprocessing() {
+ for (int i = 0; i < static_cast<int>(func_->params.size()); i++) {
+ Var param = func_->params[i];
+ if (const auto* dtensor_sinfo =
GetStructInfoAs<DTensorStructInfoNode>(param)) {
+ EmitBroadcastOrScatter(param, new_params_[i],
GetRef<DTensorStructInfo>(dtensor_sinfo));
+ } else if (const auto* tuple_sinfo =
GetStructInfoAs<TupleStructInfoNode>(param)) {
+ for (int j = 0; j < static_cast<int>(tuple_sinfo->fields.size()); j++)
{
+ if (const auto* dtensor_sinfo =
tuple_sinfo->fields[j].as<DTensorStructInfoNode>()) {
+ EmitBroadcastOrScatter(TupleGetItem(param, j),
TupleGetItem(new_params_[i], j),
+ GetRef<DTensorStructInfo>(dtensor_sinfo));
+ }
+ }
+ }
+ }
+ }
+
+ Function RewriteFunction(Function func) {
+ Array<Var> new_params;
+ for (const Var& var : func->params) {
+ Var new_param = Downcast<Var>(ShardInputParamTensorAndConstant(var));
+ var_remap_[var->vid] = new_param;
+ new_params.push_back(new_param);
+ }
+ func_ = func;
+ new_params_ = new_params;
+ auto new_body = VisitWithNewScope(func->body, new_params);
+ Function new_func(new_params, new_body, NullOpt, func->is_pure,
func->attrs);
+ return new_func;
+ }
+
+ void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode*
val) {
+ if (tuple_getitem_remap_.count(GetRef<TupleGetItem>(val))) {
+ var_remap_[binding->var->vid] =
tuple_getitem_remap_[GetRef<TupleGetItem>(val)];
+ } else {
+ ExprMutator::VisitBinding_(binding, val);
+ }
+ }
+
+ BindingBlock VisitBindingBlock_(const BindingBlockNode* block) {
+ builder_->BeginBindingBlock();
+ InputPreprocessing();
+ for (Binding binding : block->bindings) {
+ this->VisitBinding(binding);
+ }
+ return builder_->EndBlock();
+ }
+
+ BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) {
+ builder_->BeginDataflowBlock();
+ InputPreprocessing();
+ for (auto binding : block->bindings) {
+ this->VisitBinding(binding);
+ }
+ return builder_->EndBlock();
+ }
+
+ Call HandleSpecialCaseinDTensorLowering(const CallNode* call, Var
binding_var) {
+ static Op reshape_op = Op::Get("relax.reshape");
+ static Op call_tir_op = Op::Get("relax.call_tir");
+ static Op call_tir_local_view_op =
Op::Get("relax.dist.call_tir_local_view");
+ if (call->op.same_as(reshape_op)) {
+ ICHECK(call->args[1].as<ShapeExprNode>());
+ const auto* out_sinfo =
GetStructInfoAs<DTensorStructInfoNode>(binding_var);
+ ICHECK(out_sinfo);
+ auto new_call_node = make_object<CallNode>(*call);
+ new_call_node->args.Set(1, ShardShape(Downcast<ShapeExpr>(call->args[1]),
+ out_sinfo->device_mesh,
out_sinfo->placement));
+ return Call(new_call_node);
+ } else if (call->op.same_as(call_tir_local_view_op)) {
+ auto new_call_node = make_object<CallNode>(*call);
+ new_call_node->op = call_tir_op;
+ new_call_node->sinfo_args = {ConvertSinfo(GetStructInfo(binding_var),
true)};
+ return Call(new_call_node);
+ } else if (call->op.same_as(call_tir_op)) {
+ LOG(FATAL) << "call_tir should be lowered to call_tir_local_view before
lowering to relax";
+ } else if (const auto* extern_func = call->op.as<ExternFuncNode>()) {
+ auto new_call_node = make_object<CallNode>(*call);
+ if (extern_func->global_symbol ==
"vm.builtin.distributed.attention_kv_cache_append") {
+ new_call_node->op = ExternFunc("vm.builtin.attention_kv_cache_append");
+ } else if (extern_func->global_symbol ==
"vm.builtin.distributed.attention_kv_cache_view") {
+ new_call_node->op = ExternFunc("vm.builtin.attention_kv_cache_view");
+ auto orig_shape = Downcast<ShapeExpr>(call->args[1]);
+ const auto* out_sinfo =
GetStructInfoAs<DTensorStructInfoNode>(binding_var);
+ ICHECK(out_sinfo);
+ ShapeExpr new_shape = ShardShape(orig_shape, out_sinfo->device_mesh,
out_sinfo->placement);
+ new_call_node->args.Set(1, new_shape);
+ new_call_node->sinfo_args = {TensorStructInfo(new_shape,
out_sinfo->tensor_sinfo->dtype)};
+ }
+ return Call(new_call_node);
+ }
+ return GetRef<Call>(call);
+ }
+
+ void VisitBinding_(const VarBindingNode* binding, const CallNode* val) {
+ Call new_call =
+ Downcast<Call>(this->VisitExpr(HandleSpecialCaseinDTensorLowering(val,
binding->var)));
+ ReEmitBinding(binding, builder_->Normalize(new_call));
+ }
+
+ Function func_;
+ Array<Var> new_params_;
+ std::unordered_map<TupleGetItem, Var, StructuralHash, StructuralEqual>
tuple_getitem_remap_;
+};
+
+namespace transform {
+
+Pass LowerDistIR() {
+ runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
+ [=](IRModule m, PassContext pc) { return DistIRSharder::LowerDistIR(m);
};
+ return CreateModulePass(pass_func, 1, "LowerDistIR", {});
+}
+TVM_REGISTER_GLOBAL("relax.distributed.transform.LowerDistIR").set_body_typed(LowerDistIR);
+} // namespace transform
+
+} // namespace distributed
+} // namespace relax
+} // namespace tvm
diff --git
a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
new file mode 100644
index 0000000000..3df65b3ea6
--- /dev/null
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
@@ -0,0 +1,396 @@
+# 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.
+
+# type: ignore
+from tvm.script.parser import ir as I
+from tvm.script.parser import relax as R
+from tvm.script.parser import tir as T
+import tvm
+from tvm import relax
+from tvm.ir import assert_structural_equal
+import tvm.testing
+
+
+def test_mlp():
+ @I.ir_module
+ class MLP:
+ I.module_attrs({"device_num": 10})
+ I.module_global_infos(
+ {"mesh": [R.device_mesh((2,), I.Range(0, 2)), R.device_mesh((1,),
I.Range(4, 5))]}
+ )
+
+ @T.prim_func(private=True)
+ def gelu1(
+ A: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ T_multiply: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ T_multiply_1 = T.alloc_buffer((T.int64(128), T.int64(64)))
+ compute = T.alloc_buffer((T.int64(128), T.int64(64)))
+ T_multiply_2 = T.alloc_buffer((T.int64(128), T.int64(64)))
+ T_add = T.alloc_buffer((T.int64(128), T.int64(64)))
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_multiply"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(A[v_ax0, v_ax1])
+ T.writes(T_multiply_1[v_ax0, v_ax1])
+ T_multiply_1[v_ax0, v_ax1] = A[v_ax0, v_ax1] *
T.float32(0.70710678118654757)
+ for i0, i1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("compute"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(T_multiply_1[v_i0, v_i1])
+ T.writes(compute[v_i0, v_i1])
+ compute[v_i0, v_i1] = T.erf(T_multiply_1[v_i0, v_i1])
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_multiply_1"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(compute[v_ax0, v_ax1])
+ T.writes(T_multiply_2[v_ax0, v_ax1])
+ T_multiply_2[v_ax0, v_ax1] = compute[v_ax0, v_ax1] *
T.float32(0.5)
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_add"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(T_multiply_2[v_ax0, v_ax1])
+ T.writes(T_add[v_ax0, v_ax1])
+ T_add[v_ax0, v_ax1] = T.float32(0.5) + T_multiply_2[v_ax0,
v_ax1]
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_multiply_2"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(A[v_ax0, v_ax1], T_add[v_ax0, v_ax1])
+ T.writes(T_multiply[v_ax0, v_ax1])
+ T_multiply[v_ax0, v_ax1] = A[v_ax0, v_ax1] * T_add[v_ax0,
v_ax1]
+
+ @T.prim_func(private=True)
+ def matmul1(
+ A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+ B: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ matmul_1: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1, k in T.grid(T.int64(128), T.int64(64), T.int64(128)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
+ T.reads(A[v_i0, v_k], B[v_k, v_i1])
+ T.writes(matmul_1[v_i0, v_i1])
+ with T.init():
+ matmul_1[v_i0, v_i1] = T.float32(0)
+ matmul_1[v_i0, v_i1] = matmul_1[v_i0, v_i1] + A[v_i0, v_k]
* B[v_k, v_i1]
+
+ @T.prim_func(private=True)
+ def matmul2(
+ A: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ B: T.Buffer((T.int64(64), T.int64(128)), "float32"),
+ matmul_1: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1, k in T.grid(T.int64(128), T.int64(128), T.int64(64)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
+ T.reads(A[v_i0, v_k], B[v_k, v_i1])
+ T.writes(matmul_1[v_i0, v_i1])
+ with T.init():
+ matmul_1[v_i0, v_i1] = T.float32(0)
+ matmul_1[v_i0, v_i1] = matmul_1[v_i0, v_i1] + A[v_i0, v_k]
* B[v_k, v_i1]
+
+ @R.function
+ def foo(
+ x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
+ weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
+ weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
+ ) -> R.DTensor((128, 128), "float32", "mesh[0]", "R"):
+ R.func_attr({"num_input": 1})
+ cls = MLP
+ lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.dist.call_tir_local_view(
+ cls.matmul1,
+ (x, weight1),
+ out_sinfo=R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
+ )
+ lv1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.dist.call_tir_local_view(
+ cls.gelu1, (lv0,), out_sinfo=R.DTensor((128, 128), "float32",
"mesh[0]", "S[1]")
+ )
+ lv2: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") = lv1
+ gv: R.DTensor((128, 128), "float32", "mesh[0]", "R") =
R.dist.call_tir_local_view(
+ cls.matmul2,
+ (lv2, weight2),
+ out_sinfo=R.DTensor((128, 128), "float32", "mesh[0]", "R"),
+ )
+ lv3: R.DTensor((128, 128), "float32", "mesh[0]", "R") =
R.ccl.allreduce(
+ gv, op_type="sum"
+ )
+ return lv3
+
+ @I.ir_module
+ class LoweredMLP:
+ I.module_attrs({"device_num": 10})
+ I.module_global_infos(
+ {"mesh": [R.device_mesh((2,), I.Range(0, 2)), R.device_mesh((1,),
I.Range(4, 5))]}
+ )
+
+ @R.function
+ def foo(
+ x: R.Tensor((128, 128), dtype="float32"),
+ weight1: R.Tensor((128, 128), dtype="float32"),
+ weight2: R.Tensor((128, 128), dtype="float32"),
+ ) -> R.Tensor((128, 128), dtype="float32"):
+ R.func_attr({"num_input": 1})
+ cls = LoweredMLP
+ gv: R.Tensor((128, 128), dtype="float32") =
R.ccl.broadcast_from_worker0(x)
+ gv1: R.Tensor((128, 64), dtype="float32") =
R.ccl.scatter_from_worker0(
+ weight1, num_workers=2, axis=1
+ )
+ gv2: R.Tensor((64, 128), dtype="float32") =
R.ccl.scatter_from_worker0(
+ weight2, num_workers=2, axis=0
+ )
+ lv0 = R.call_tir(
+ MLP.get_global_var("matmul1"),
+ (gv, gv1),
+ out_sinfo=R.Tensor((128, 64), dtype="float32"),
+ )
+ lv1 = R.call_tir(
+ MLP.get_global_var("gelu1"), (lv0,), out_sinfo=R.Tensor((128,
64), dtype="float32")
+ )
+ lv2: R.Tensor((128, 64), dtype="float32") = lv1
+ gv_1 = R.call_tir(
+ MLP.get_global_var("matmul2"),
+ (lv2, gv2),
+ out_sinfo=R.Tensor((128, 128), dtype="float32"),
+ )
+ lv3: R.Tensor((128, 128), dtype="float32") = R.ccl.allreduce(gv_1,
op_type="sum")
+ return lv3
+
+ for gv, func in MLP.functions_items():
+ if gv.name_hint != "foo":
+ LoweredMLP[gv] = func
+
+ mod = MLP
+ mod = relax.distributed.transform.LowerDistIR()(mod)
+ tvm.ir.assert_structural_equal(mod, LoweredMLP)
+
+
+def test_mlp_with_tuple():
+ @I.ir_module
+ class MLPWithTuple:
+ I.module_attrs({"device_num": 10})
+ I.module_global_infos(
+ {"mesh": [R.device_mesh((2,), I.Range(0, 2)), R.device_mesh((1,),
I.Range(4, 5))]}
+ )
+
+ @T.prim_func(private=True)
+ def gelu1(
+ A: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ T_multiply: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ T_multiply_1 = T.alloc_buffer((T.int64(128), T.int64(64)))
+ compute = T.alloc_buffer((T.int64(128), T.int64(64)))
+ T_multiply_2 = T.alloc_buffer((T.int64(128), T.int64(64)))
+ T_add = T.alloc_buffer((T.int64(128), T.int64(64)))
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_multiply"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(A[v_ax0, v_ax1])
+ T.writes(T_multiply_1[v_ax0, v_ax1])
+ T_multiply_1[v_ax0, v_ax1] = A[v_ax0, v_ax1] *
T.float32(0.70710678118654757)
+ for i0, i1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("compute"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(T_multiply_1[v_i0, v_i1])
+ T.writes(compute[v_i0, v_i1])
+ compute[v_i0, v_i1] = T.erf(T_multiply_1[v_i0, v_i1])
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_multiply_1"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(compute[v_ax0, v_ax1])
+ T.writes(T_multiply_2[v_ax0, v_ax1])
+ T_multiply_2[v_ax0, v_ax1] = compute[v_ax0, v_ax1] *
T.float32(0.5)
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_add"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(T_multiply_2[v_ax0, v_ax1])
+ T.writes(T_add[v_ax0, v_ax1])
+ T_add[v_ax0, v_ax1] = T.float32(0.5) + T_multiply_2[v_ax0,
v_ax1]
+ for ax0, ax1 in T.grid(T.int64(128), T.int64(64)):
+ with T.block("T_multiply_2"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(A[v_ax0, v_ax1], T_add[v_ax0, v_ax1])
+ T.writes(T_multiply[v_ax0, v_ax1])
+ T_multiply[v_ax0, v_ax1] = A[v_ax0, v_ax1] * T_add[v_ax0,
v_ax1]
+
+ @T.prim_func(private=True)
+ def matmul11(
+ A: T.Buffer((T.int64(64), T.int64(64)), "float32"),
+ B: T.Buffer((T.int64(64), T.int64(128)), "float32"),
+ matmul: T.Buffer((T.int64(64), T.int64(128)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1, k in T.grid(T.int64(64), T.int64(128), T.int64(64)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
+ T.reads(A[v_i0, v_k], B[v_k, v_i1])
+ T.writes(matmul[v_i0, v_i1])
+ with T.init():
+ matmul[v_i0, v_i1] = T.float32(0)
+ matmul[v_i0, v_i1] = matmul[v_i0, v_i1] + A[v_i0, v_k] *
B[v_k, v_i1]
+
+ @T.prim_func(private=True)
+ def matmul2(
+ A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+ B: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ matmul: T.Buffer((T.int64(128), T.int64(64)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1, k in T.grid(T.int64(128), T.int64(64), T.int64(128)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
+ T.reads(A[v_i0, v_k], B[v_k, v_i1])
+ T.writes(matmul[v_i0, v_i1])
+ with T.init():
+ matmul[v_i0, v_i1] = T.float32(0)
+ matmul[v_i0, v_i1] = matmul[v_i0, v_i1] + A[v_i0, v_k] *
B[v_k, v_i1]
+
+ @T.prim_func(private=True)
+ def split11(
+ A: T.Buffer((128, 64), "float32"),
+ T_split: T.Buffer((64, 64), "float32"),
+ T_split_1: T.Buffer((64, 64), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for ax1, ax2 in T.grid(64, 64):
+ with T.block("T_split"):
+ v_ax1, v_ax2 = T.axis.remap("SS", [ax1, ax2])
+ T.reads(A[v_ax1, v_ax2])
+ T.writes(T_split[v_ax1, v_ax2])
+ T_split[v_ax1, v_ax2] = A[v_ax1, v_ax2]
+ for ax1, ax2 in T.grid(64, 64):
+ with T.block("T_split_1"):
+ v_ax1, v_ax2 = T.axis.remap("SS", [ax1, ax2])
+ T.reads(A[v_ax1 + 64, v_ax2])
+ T.writes(T_split_1[v_ax1, v_ax2])
+ T_split_1[v_ax1, v_ax2] = A[v_ax1 + 64, v_ax2]
+
+ @R.function
+ def foo(
+ x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
+ weight_packed: R.Tuple(
+ R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
+ R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
+ ),
+ ) -> R.DTensor((64, 128), "float32", "mesh[0]", "R"):
+ cls = MLPWithTuple
+ weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
weight_packed[0]
+ lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.dist.call_tir_local_view(
+ cls.matmul2,
+ (x, weight1),
+ out_sinfo=R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
+ )
+ lv1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.dist.call_tir_local_view(
+ cls.gelu1, (lv0,), out_sinfo=R.DTensor((128, 128), "float32",
"mesh[0]", "S[1]")
+ )
+ gv: R.Tuple(
+ R.DTensor((64, 128), "float32", "mesh[0]", "S[1]"),
+ R.DTensor((64, 128), "float32", "mesh[0]", "S[1]"),
+ ) = R.dist.call_tir_local_view(
+ cls.split11,
+ (lv1,),
+ out_sinfo=[
+ R.DTensor((64, 128), "float32", "mesh[0]", "S[1]"),
+ R.DTensor((64, 128), "float32", "mesh[0]", "S[1]"),
+ ],
+ )
+ lv2: R.DTensor((64, 128), "float32", "mesh[0]", "S[1]") = gv[0]
+ lv3: R.DTensor((64, 128), "float32", "mesh[0]", "S[1]") = lv2
+ weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]") =
weight_packed[1]
+ gv_1: R.DTensor((64, 128), "float32", "mesh[0]", "R") =
R.dist.call_tir_local_view(
+ cls.matmul11,
+ (lv3, weight2),
+ out_sinfo=R.DTensor((64, 128), "float32", "mesh[0]", "R"),
+ )
+ lv4: R.DTensor((64, 128), "float32", "mesh[0]", "R") =
R.ccl.allreduce(
+ gv_1, op_type="sum"
+ )
+ return lv4
+
+ @I.ir_module
+ class LoweredMLPWithTuple:
+ I.module_attrs({"device_num": 10})
+ I.module_global_infos(
+ {"mesh": [R.device_mesh((2,), I.Range(0, 2)), R.device_mesh((1,),
I.Range(4, 5))]}
+ )
+
+ @R.function
+ def foo(
+ x: R.Tensor((128, 128), dtype="float32"),
+ weight_packed: R.Tuple(
+ R.Tensor((128, 128), dtype="float32"), R.Tensor((128, 128),
dtype="float32")
+ ),
+ ) -> R.Tensor((64, 128), dtype="float32"):
+ cls = LoweredMLPWithTuple
+ gv: R.Tensor((128, 128), dtype="float32") =
R.ccl.broadcast_from_worker0(x)
+ gv1: R.Tensor((128, 128), dtype="float32") = weight_packed[0]
+ gv2: R.Tensor((128, 64), dtype="float32") =
R.ccl.scatter_from_worker0(
+ gv1, num_workers=2, axis=1
+ )
+ gv3: R.Tensor((128, 128), dtype="float32") = weight_packed[1]
+ gv4: R.Tensor((64, 128), dtype="float32") =
R.ccl.scatter_from_worker0(
+ gv3, num_workers=2, axis=0
+ )
+ lv0 = R.call_tir(
+ MLPWithTuple.get_global_var("matmul2"),
+ (gv, gv2),
+ out_sinfo=R.Tensor((128, 64), dtype="float32"),
+ )
+ lv1 = R.call_tir(
+ MLPWithTuple.get_global_var("gelu1"),
+ (lv0,),
+ out_sinfo=R.Tensor((128, 64), dtype="float32"),
+ )
+ gv_1 = R.call_tir(
+ MLPWithTuple.get_global_var("split11"),
+ (lv1,),
+ out_sinfo=[
+ R.Tensor((64, 64), dtype="float32"),
+ R.Tensor((64, 64), dtype="float32"),
+ ],
+ )
+ lv2: R.Tensor((64, 64), dtype="float32") = gv_1[0]
+ lv3: R.Tensor((64, 64), dtype="float32") = lv2
+ gv_1_1 = R.call_tir(
+ MLPWithTuple.get_global_var("matmul11"),
+ (lv3, gv4),
+ out_sinfo=R.Tensor((64, 128), dtype="float32"),
+ )
+ lv4: R.Tensor((64, 128), dtype="float32") =
R.ccl.allreduce(gv_1_1, op_type="sum")
+ return lv4
+
+ for gv, func in MLPWithTuple.functions_items():
+ if gv.name_hint != "foo":
+ LoweredMLPWithTuple[gv] = func
+
+ mod = MLPWithTuple
+ mod = relax.distributed.transform.LowerDistIR()(mod)
+ tvm.ir.assert_structural_equal(mod, LoweredMLPWithTuple)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
