diff --git a/include/tvm/relax/distributed/transform.h b/include/tvm/relax/distributed/transform.h index 335d9a62e974..5cf15ee65e43 100644 --- a/include/tvm/relax/distributed/transform.h +++ b/include/tvm/relax/distributed/transform.h @@ -48,6 +48,13 @@ using DataflowBlock = tvm::relax::DataflowBlock; */ TVM_DLL Pass PropagateSharding(); +/*! + * \brief Legalize redistribute op to ccl op. + * + * \return The Pass. + */ +TVM_DLL Pass LegalizeRedistribute(); + } // 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 9d1e7800306a..517189392af7 100644 --- a/python/tvm/relax/distributed/transform/__init__.py +++ b/python/tvm/relax/distributed/transform/__init__.py @@ -16,4 +16,4 @@ # under the License. """Relax distributed-related transformations. """ -from .transform import PropagateSharding +from .transform import PropagateSharding, LegalizeRedistribute diff --git a/python/tvm/relax/distributed/transform/transform.py b/python/tvm/relax/distributed/transform/transform.py index 7c8bc49e0d64..b1f6bde04496 100644 --- a/python/tvm/relax/distributed/transform/transform.py +++ b/python/tvm/relax/distributed/transform/transform.py @@ -30,3 +30,16 @@ def PropagateSharding() -> tvm.ir.transform.Pass: The registered pass """ return _ffi_api.PropagateSharding() # 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 diff --git a/python/tvm/relax/op/distributed/__init__.py b/python/tvm/relax/op/distributed/__init__.py index 35226b879056..cd38b90bcef4 100644 --- a/python/tvm/relax/op/distributed/__init__.py +++ b/python/tvm/relax/op/distributed/__init__.py @@ -16,4 +16,4 @@ # under the License. """Operators serving for distributed Relax.""" -from .distributed import annotate_sharding, redistribute +from .distributed import annotate_sharding, redistribute, redistribute_replica_to_shard diff --git a/python/tvm/relax/op/distributed/distributed.py b/python/tvm/relax/op/distributed/distributed.py index bcc5ca2f49e5..7e0e6fa23d4c 100644 --- a/python/tvm/relax/op/distributed/distributed.py +++ b/python/tvm/relax/op/distributed/distributed.py @@ -59,3 +59,29 @@ def redistribute(input: Expr, device_mesh: DeviceMesh, placement: Placement) -> The tensor after redistribution. """ return _ffi_api.redistribute(input, device_mesh, placement) # type: ignore + + +def redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int) -> Expr: + """Slice tensor into several parts along one axis, + and each worker takes one part. + input.struct_info.shape[axis] % num_workers == 0 is required. + Each worker must have an identical copy of the input. + This is a specialized version of redistribute op. + + Parameters + ---------- + input : relax.Expr + The buffer to be sliced into equal parts. + + num_worker : int + The number of workers, i.e. the number of parts the given buffer should be sliced into. + + axis : int + The axis of the tensor to be sliced. + + Returns + ------- + result : relax.Expr + Sliced Tensor kept by each device. + """ + return _ffi_api.redistribute_replica_to_shard(input, num_workers, axis) diff --git a/python/tvm/relax/transform/legalize_ops/__init__.py b/python/tvm/relax/transform/legalize_ops/__init__.py index e222e7d4efc8..e3b3213a38b5 100644 --- a/python/tvm/relax/transform/legalize_ops/__init__.py +++ b/python/tvm/relax/transform/legalize_ops/__init__.py @@ -19,6 +19,7 @@ from . import ccl from . import create from . import datatype +from . import distributed from . import grad from . import image from . import index diff --git a/python/tvm/relax/transform/legalize_ops/distributed.py b/python/tvm/relax/transform/legalize_ops/distributed.py new file mode 100644 index 000000000000..d540628e0e2b --- /dev/null +++ b/python/tvm/relax/transform/legalize_ops/distributed.py @@ -0,0 +1,43 @@ +# 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 +"""Default legalization function for distir-related operators.""" +from tvm import tir, relax +from ...block_builder import BlockBuilder +from ...expr import Call, Expr +from ...op import call_pure_packed +from ...struct_info import ShapeStructInfo +from .common import register_legalize + + +@register_legalize("relax.dist.redistribute_replica_to_shard") +def _redistribute_replica_to_shard(_bb: BlockBuilder, call: Call) -> Expr: + num_workers = call.attrs.num_workers + axis = call.attrs.axis + worker_id_symbol = tir.Var("worker_id", "int64") + worker_id_var = _bb.emit( + call_pure_packed("runtime.disco.worker_id", sinfo_args=[ShapeStructInfo(None)]) + ) + _bb.match_cast(worker_id_var, ShapeStructInfo([worker_id_symbol])) + + split_axis_size = call.args[0].struct_info.shape[axis] + return relax.op.strided_slice( + call.args[0], + axes=[axis], + begin=[worker_id_symbol * split_axis_size // num_workers], + end=[(worker_id_symbol + 1) * split_axis_size // num_workers], + ) diff --git a/python/tvm/script/ir_builder/relax/distributed/ir.py b/python/tvm/script/ir_builder/relax/distributed/ir.py index 51ef84620cc1..28ff6ae7328a 100644 --- a/python/tvm/script/ir_builder/relax/distributed/ir.py +++ b/python/tvm/script/ir_builder/relax/distributed/ir.py @@ -14,7 +14,7 @@ # 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 +# pylint: disable=redefined-builtin, wrong-import-order, no-member, invalid-name, unused-import """IRBuilder for distributed Relax dialect""" from typing import Union, List, Tuple, Optional @@ -32,6 +32,7 @@ from tvm.relax.op.distributed import ( redistribute as _redistribute, annotate_sharding as _annotate_sharding, + redistribute_replica_to_shard, ) from tvm.relax.distributed import DeviceMesh, Placement from . import _ffi_api diff --git a/python/tvm/script/parser/relax/__init__.py b/python/tvm/script/parser/relax/__init__.py index 0b1ed168de1d..704189060b26 100644 --- a/python/tvm/script/parser/relax/__init__.py +++ b/python/tvm/script/parser/relax/__init__.py @@ -34,18 +34,15 @@ else: from .entry import function, macro -__all__ = ( - _relax.__all__ - + dist.__all__ - + [ - "Callable", - "Object", - "Prim", - "Shape", - "Tensor", - "Tuple", - "function", - "macro", - "match_cast", - ] -) +__all__ = _relax.__all__ + [ + "dist", + "Callable", + "Object", + "Prim", + "Shape", + "Tensor", + "Tuple", + "function", + "macro", + "match_cast", +] diff --git a/python/tvm/script/parser/relax/dist.py b/python/tvm/script/parser/relax/dist.py index f9c78f980fd0..c4e81a2ae208 100644 --- a/python/tvm/script/parser/relax/dist.py +++ b/python/tvm/script/parser/relax/dist.py @@ -25,7 +25,13 @@ from tvm.relax.distributed import DeviceMesh, Placement, DTensorStructInfo, device_mesh from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder.ir import IRModuleFrame -from tvm.script.ir_builder.relax.distributed import call_tir, const, annotate_sharding, redistribute +from tvm.script.ir_builder.relax.distributed import ( + call_tir, + const, + annotate_sharding, + redistribute, + redistribute_replica_to_shard, +) from .entry import StructInfoProxy, TensorProxy diff --git a/src/relax/distributed/transform/legalize_redistribute.cc b/src/relax/distributed/transform/legalize_redistribute.cc new file mode 100644 index 000000000000..5a67f0351cf7 --- /dev/null +++ b/src/relax/distributed/transform/legalize_redistribute.cc @@ -0,0 +1,123 @@ +/* + * 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/legalize_redistribute.cc + * \brief Pass for legalizing redistribute op to ccl op. + */ + +#include +#include +#include +#include +#include +#include + +#include "../../../tir/schedule/transform.h" +#include "../../op/ccl/ccl.h" +#include "../../op/distributed/distributed.h" + +namespace tvm { +namespace relax { +namespace distributed { + +class RedistributeLegalizer : public ExprMutator { + public: + static IRModule LegalizeRedistribute(IRModule mod) { + return RedistributeLegalizer(mod).Legalize(); + } + + private: + explicit RedistributeLegalizer(IRModule mod) : ExprMutator(mod) {} + + IRModule Legalize() { + auto mod = builder_->GetContextIRModule(); + for (const auto& [gv, base_func] : mod->functions) { + const auto* func_ = base_func.as(); + if (func_ == nullptr) { + continue; + } + Expr new_func_body = VisitExpr(func_->body); + auto new_func = make_object(*func_); + new_func->body = new_func_body; + builder_->UpdateFunction(gv, Function(new_func)); + } + return builder_->GetContextIRModule(); + } + using ExprMutator::VisitExpr_; + Expr VisitExpr_(const CallNode* op) final { + Call call = Downcast(ExprMutator::VisitExpr_(op)); + static Op redistribute_op = Op::Get("relax.dist.redistribute"); + if (call->op.same_as(redistribute_op)) { + const auto* attrs = call->attrs.as(); + ICHECK(attrs); + const auto* input_sinfo = call->args[0]->struct_info_.as(); + ICHECK(input_sinfo); + // As the first step, we only support redistribute in the same device mesh, + // and the device mesh must be 1d + // todo: extend the ccl ops so that it can support 2d device mesh, and different sharding + // dimension + ICHECK(StructuralEqual()(input_sinfo->device_mesh, attrs->device_mesh)); + ICHECK(input_sinfo->device_mesh->shape.size() == 1); + // only support "S[x]"-> "R" and "R" -> "S[x]" + PlacementSpec input_spec = input_sinfo->placement->dim_specs[0]; + PlacementSpec output_spec = attrs->placement->dim_specs[0]; + if (input_spec->kind == PlacementSpecKind::kReplica && + output_spec->kind == PlacementSpecKind::kReplica) { + // "R" -> "R" + return call->args[0]; + } else if (input_spec->kind == PlacementSpecKind::kSharding && + output_spec->kind == PlacementSpecKind::kSharding) { + // "S[x]" -> "S[y]" + if (input_spec->axis != output_spec->axis) { + LOG(FATAL) << "AlltoAll not implemented yet"; + } else { + return call->args[0]; + } + } else if (input_spec->kind == PlacementSpecKind::kSharding && + output_spec->kind == PlacementSpecKind::kReplica) { + // "S[x]" -> "R" + LOG(FATAL) << "Allgather not implemented yet"; + } else if (input_spec->kind == PlacementSpecKind::kReplica && + output_spec->kind == PlacementSpecKind::kSharding) { + // "R" -> "S[x]" + return redistribute_replica_to_shard(call->args[0], attrs->device_mesh->shape[0], + output_spec->axis); + } else { + LOG(FATAL) << "Unsupported redistribute op"; + } + } + return call; + } +}; + +namespace transform { + +Pass LegalizeRedistribute() { + runtime::TypedPackedFunc pass_func = + [=](IRModule m, PassContext pc) { return RedistributeLegalizer::LegalizeRedistribute(m); }; + return CreateModulePass(pass_func, 1, "LegalizeRedistribute", {}); +} +TVM_REGISTER_GLOBAL("relax.distributed.transform.LegalizeRedistribute") + .set_body_typed(LegalizeRedistribute); +} // namespace transform + +} // namespace distributed +} // namespace relax +} // namespace tvm diff --git a/src/relax/op/distributed/distributed.cc b/src/relax/op/distributed/distributed.cc index 6def4ef0198d..658b7cde34dc 100644 --- a/src/relax/op/distributed/distributed.cc +++ b/src/relax/op/distributed/distributed.cc @@ -24,6 +24,7 @@ #include "distributed.h" +#include #include #include @@ -85,5 +86,86 @@ TVM_REGISTER_OP("relax.dist.redistribute") .set_attr("dist.FInferStructInfo", InferDistStructInfoRedistribute) .set_attr("FPurity", Bool(true)); +StructInfo InferStructInfoRtoS(const Call& call, const BlockBuilder& ctx) { + TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx); + DataType output_dtype = input_sinfo->dtype; + + const auto* attrs = call->attrs.as(); + int num_workers = attrs->num_workers; + + arith::Analyzer* analyzer = ctx->GetAnalyzer(); + auto input_shape = input_sinfo->GetShape(); + CHECK(input_shape.defined()) + << "input tensor of redistribute_replica_to_shard should have defined shape."; + + if (analyzer->CanProve(floormod(input_shape.value()[attrs->axis], PrimExpr(num_workers))) != 0) { + ctx->ReportFatal(Diagnostic::Error(call) + << "redistribute_replica_to_shard expects the size of axis " << attrs->axis + << " of input tensor to be " + "divisible by the " + "num_workers. However, the axis " + << attrs->axis << " of input tensor is " << input_shape.value()[attrs->axis] + << " while num_workers is " << num_workers); + } + + Array output_shape = input_shape.value(); + output_shape.Set(attrs->axis, div(output_shape[attrs->axis], num_workers)); + if (input_sinfo->vdevice.defined()) { + return TensorStructInfo(ShapeExpr(output_shape), output_dtype, input_sinfo->vdevice.value()); + } + return TensorStructInfo(ShapeExpr(output_shape), output_dtype); +} + +StructInfo InferDistStructInfoRtoS(const Call& call, const BlockBuilder& ctx) { + using namespace distributed; + Array input_dtensor_sinfos = GetInputDTensorStructInfo(call, ctx); + ICHECK(input_dtensor_sinfos.size() == 1); + DTensorStructInfo input_dtensor_sinfo = input_dtensor_sinfos[0]; + TensorStructInfo tensor_sinfo = input_dtensor_sinfo->tensor_sinfo; + const auto* attrs = call->attrs.as(); + int num_workers = attrs->num_workers; + arith::Analyzer* analyzer = ctx->GetAnalyzer(); + auto input_shape = tensor_sinfo->GetShape(); + CHECK(input_shape.defined()) + << "input tensor of redistribute_replica_to_shard should have defined shape."; + + if (analyzer->CanProve(floormod(input_shape.value()[attrs->axis], PrimExpr(num_workers))) != 0) { + ctx->ReportFatal(Diagnostic::Error(call) + << "redistribute_replica_to_shard expects the size of axis " << attrs->axis + << " of input tensor to be " + "divisible by the " + "num_workers. However, the axis " + << attrs->axis << " of input tensor is " << input_shape.value()[attrs->axis] + << " while num_workers is " << num_workers); + } + + DeviceMesh device_mesh = input_dtensor_sinfo->device_mesh; + // FIXME: this is a hack where there's only 1d mesh + ICHECK(device_mesh->shape.size() == 1); + ICHECK(input_dtensor_sinfo->placement->dim_specs[0]->kind == PlacementSpecKind::kReplica); + return DTensorStructInfo(tensor_sinfo, device_mesh, + Placement::FromText("S[" + std::to_string(attrs->axis) + "]")); +} + +Expr redistribute_replica_to_shard(Expr input, int num_workers, int axis) { + ObjectPtr attrs = make_object(); + attrs->num_workers = std::move(num_workers); + attrs->axis = std::move(axis); + static const Op& op = Op::Get("relax.dist.redistribute_replica_to_shard"); + + return Call(op, {std::move(input)}, Attrs{attrs}, {}); +} + +TVM_REGISTER_GLOBAL("relax.op.dist.redistribute_replica_to_shard") + .set_body_typed(redistribute_replica_to_shard); + +TVM_REGISTER_OP("relax.dist.redistribute_replica_to_shard") + .set_num_inputs(1) + .add_argument("input", "Tensor", "The buffer to be sliced.") + .set_attrs_type() + .set_attr("FInferStructInfo", InferStructInfoRtoS) + .set_attr("dist.FInferStructInfo", InferDistStructInfoRtoS) + .set_attr("FPurity", Bool(true)); + } // namespace relax } // namespace tvm diff --git a/src/relax/op/distributed/distributed.h b/src/relax/op/distributed/distributed.h index 06f926edb1c4..67ccc0fc6b15 100644 --- a/src/relax/op/distributed/distributed.h +++ b/src/relax/op/distributed/distributed.h @@ -27,6 +27,7 @@ #include #include "../op_common.h" +#include "utils.h" namespace tvm { namespace relax { @@ -51,6 +52,17 @@ Expr annotate_sharding(Expr input, distributed::DeviceMesh device_mesh, Expr redistribute(Expr input, distributed::DeviceMesh device_mesh, distributed::Placement placement); +/*! + * \brief slice tensor into several parts along one axis, + and each worker takes one part. + Assumes input is already broadcasted. + This is a specialized version of redistribute op. + * \param input The input tensor. + * \param num_workers The number of workers. + * \param axis The tensor axis to slice. + * \return The result. + */ +Expr redistribute_replica_to_shard(Expr input, int num_workers, int axis); } // namespace relax } // namespace tvm diff --git a/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py b/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py new file mode 100644 index 000000000000..92e9e5a4d403 --- /dev/null +++ b/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py @@ -0,0 +1,69 @@ +# 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 +import tvm +from tvm import relax +import tvm.testing + + +def test_simple(): + @I.ir_module + class Before: + I.module_attrs({"device_num": 2}) + I.module_global_infos({"mesh": [R.device_mesh((2,), I.Range(0, 2))]}) + + @R.function + def foo( + x1: R.DTensor((128, 128), "float32", "mesh[0]", "R"), + x2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"), + ): + R.func_attr({"num_input": 1}) + # scatter + lv0 = R.dist.redistribute(x1, "mesh[0]", "S[1]") + # do nothing + lv1 = R.dist.redistribute(x2, "mesh[0]", "S[0]") + return (lv0, lv1) + + @I.ir_module + class Expected: + I.module_attrs({"device_num": 2}) + I.module_global_infos({"mesh": [R.device_mesh((2,), I.Range(0, 2))]}) + + @R.function + def foo( + x1: R.DTensor((128, 128), "float32", "mesh[0]", "R"), + x2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"), + ) -> R.Tuple( + R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"), + R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"), + ): + R.func_attr({"num_input": 1}) + lv0: R.DTensor( + (128, 128), "float32", "mesh[0]", "S[1]" + ) = R.dist.redistribute_replica_to_shard(x1, num_workers=2, axis=1) + lv1: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]") = x2 + return (lv0, lv1) + + after = relax.distributed.transform.LegalizeRedistribute()(Before) + tvm.ir.assert_structural_equal(after, Expected) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/relax/test_op_distributed.py b/tests/python/relax/test_op_distributed.py new file mode 100644 index 000000000000..dc5440c0c113 --- /dev/null +++ b/tests/python/relax/test_op_distributed.py @@ -0,0 +1,58 @@ +# 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. +import pytest +import tvm +from tvm._ffi.base import TVMError +import tvm.testing +from tvm import relax +from tvm.script.parser import relax as R + + +def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: relax.StructInfo): + ret = bb.normalize(call) + tvm.ir.assert_structural_equal(ret.struct_info, expected_sinfo) + + +def test_redistribute_R_to_S(): + bb = relax.BlockBuilder() + mesh = R.device_mesh((4,), list(range(4))) + x = relax.Var("x", R.DTensor((3, 4), "float32", device_mesh=mesh, placement="R")) + + _check_inference( + bb, + R.distributed.redistribute_replica_to_shard(x, num_workers=4, axis=1), + R.DTensor((3, 4), "float32", device_mesh=mesh, placement="S[1]"), + ) + + # wrong: indivisible + with pytest.raises(TVMError): + bb.normalize(R.distributed.redistribute_replica_to_shard(x, num_workers=4, axis=0)) + + y = relax.Var("y", R.Tensor((3, 4), "float32")) + _check_inference( + bb, + R.distributed.redistribute_replica_to_shard(y, num_workers=4, axis=1), + R.Tensor((3, 1), "float32"), + ) + + # wrong: indivisible + with pytest.raises(TVMError): + bb.normalize(R.distributed.redistribute_replica_to_shard(y, num_workers=4, axis=0)) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/relax/test_transform_legalize_ops_distributed.py b/tests/python/relax/test_transform_legalize_ops_distributed.py new file mode 100644 index 000000000000..c1f462e36caa --- /dev/null +++ b/tests/python/relax/test_transform_legalize_ops_distributed.py @@ -0,0 +1,64 @@ +# 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. + +import tvm +import tvm.testing +from tvm import relax +from tvm.relax.transform import LegalizeOps +from tvm.script import ir as I +from tvm.script import relax as R +from tvm.script import tir as T + + +def test_redistribute_replica_to_shard(): + # fmt: off + @tvm.script.ir_module + class Before: + @R.function + def main(x: R.Tensor((10, 10), "float32")) -> R.Tensor((10, 5), "float32"): + gv0 = R.dist.redistribute_replica_to_shard(x, num_workers=2, axis=1) + return gv0 + + @I.ir_module + class Expected: + @T.prim_func(private=True) + def strided_slice(A: T.Buffer((T.int64(10), T.int64(10)), "float32"), redistribute_replica_to_shard: T.Buffer((T.int64(10), T.int64(5)), "float32"), worker_id: T.int64): + T.func_attr({"tir.noalias": T.bool(True)}) + # with T.block("root"): + for i0, i1 in T.grid(T.int64(10), T.int64(5)): + with T.block("redistribute_replica_to_shard"): + v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) + T.reads(A[v_i0, worker_id * T.int64(5) + v_i1]) + T.writes(redistribute_replica_to_shard[v_i0, v_i1]) + redistribute_replica_to_shard[v_i0, v_i1] = A[v_i0, worker_id * T.int64(5) + v_i1] + + @R.function + def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 5), dtype="float32"): + worker_id = T.int64() + cls = Expected + gv: R.Shape(ndim=-1) = R.call_pure_packed("runtime.disco.worker_id", sinfo_args=(R.Shape(ndim=-1),)) + gv1: R.Shape([worker_id]) = R.match_cast(gv, R.Shape([worker_id])) + gv0 = R.call_tir(cls.strided_slice, (x,), out_sinfo=R.Tensor((10, 5), dtype="float32"), tir_vars=R.shape([worker_id])) + return gv0 + # fmt: on + + mod = LegalizeOps()(Before) + tvm.ir.assert_structural_equal(mod, Expected) + + +if __name__ == "__main__": + tvm.testing.main()