diff --git a/src/relax/distributed/global_info.cc b/src/relax/distributed/global_info.cc index b2b691a3dccc..58f4ab7e6a4b 100644 --- a/src/relax/distributed/global_info.cc +++ b/src/relax/distributed/global_info.cc @@ -20,19 +20,38 @@ #include #include +#include + namespace tvm { namespace relax { namespace distributed { TVM_FFI_STATIC_INIT_BLOCK() { DeviceMeshNode::RegisterReflection(); } -DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array device_ids) { - int prod = 1; - for (int i = 0; i < static_cast(shape.size()); i++) { - prod *= shape[i]; +namespace { + +int64_t MeshSize(const ffi::Shape& shape) { + bool empty = false; + for (int64_t dim : shape) { + TVM_FFI_CHECK_GE(dim, 0, ValueError) << "Device mesh dimensions must be non-negative"; + empty |= dim == 0; } + if (empty) return 0; + int64_t size = 1; + for (int64_t dim : shape) { + TVM_FFI_CHECK_LE(size, std::numeric_limits::max() / dim, ValueError) + << "Device mesh shape product exceeds int64"; + size *= dim; + } + return size; +} + +} // namespace + +DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array device_ids) { + int64_t size = MeshSize(shape); ffi::ObjectPtr n = ffi::make_object(); - TVM_FFI_ICHECK_EQ(prod, static_cast(device_ids.size())) + TVM_FFI_ICHECK_EQ(static_cast(size), device_ids.size()) << "The number of device ids must match the product of the shape"; n->shape = std::move(shape); n->device_ids = std::move(device_ids); @@ -42,17 +61,22 @@ DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array device_ids) { DeviceMesh::DeviceMesh(ffi::Shape shape, Range device_range) { ffi::ObjectPtr n = ffi::make_object(); ffi::Array device_ids; - int range_start = device_range->min.as()->value; - int range_extent = device_range->extent.as()->value; - for (int i = range_start; i < range_start + range_extent; i++) { - device_ids.push_back(i); + const auto* start = device_range->min.as(); + const auto* extent = device_range->extent.as(); + TVM_FFI_CHECK(start && extent, ValueError) << "Device mesh range must be constant"; + int64_t range_start = start->value; + int64_t range_extent = extent->value; + TVM_FFI_CHECK_GE(range_extent, 0, ValueError) << "Device mesh range extent must be non-negative"; + TVM_FFI_ICHECK_EQ(MeshSize(shape), range_extent) + << "The number of device ids must match the product of the shape"; + if (range_extent > 0) { + TVM_FFI_CHECK_LE(range_start, std::numeric_limits::max() - (range_extent - 1), + ValueError) + << "Device mesh range exceeds int64"; } - int prod = 1; - for (int i = 0; i < static_cast(shape.size()); i++) { - prod *= shape[i]; + for (int64_t i = 0; i < range_extent; ++i) { + device_ids.push_back(range_start + i); } - TVM_FFI_ICHECK_EQ(prod, static_cast(device_ids.size())) - << "The number of device ids must match the product of the shape"; n->device_ids = std::move(device_ids); n->shape = std::move(shape); n->device_range = std::move(device_range); diff --git a/tests/python/relax/distributed/test_distributed_dtensor_type.py b/tests/python/relax/distributed/test_distributed_dtensor_type.py index 1e67ea16c0ef..3cb669371cf3 100644 --- a/tests/python/relax/distributed/test_distributed_dtensor_type.py +++ b/tests/python/relax/distributed/test_distributed_dtensor_type.py @@ -41,6 +41,41 @@ def _check_json_roundtrip(x): return xret +@pytest.mark.parametrize("use_range", [False, True]) +@pytest.mark.parametrize( + "shape,count", + [((2, 2), 3), ((2**31, 2), 0), ((2**32,), 0), ((2**62, 4), 0), ((-2, -2), 4), ((0, -1), 0)], +) +def test_device_mesh_invalid_cardinality(shape, count, use_range): + ids = Range(0, count) if use_range else list(range(count)) + with pytest.raises((ValueError, tvm.error.InternalError)): + rx.distributed.DeviceMesh(shape, ids) + + +@pytest.mark.parametrize("use_range", [False, True]) +@pytest.mark.parametrize("shape,count", [((2, 2), 4), ((), 1), ((0, 2), 0), ((2**62, 4, 0), 0)]) +def test_device_mesh_cardinality(shape, count, use_range): + ids = Range(0, count) if use_range else list(range(count)) + mesh = rx.distributed.DeviceMesh(shape, ids) + assert tuple(mesh.shape) == shape + assert list(mesh.device_ids) == list(range(count)) + + +def test_device_mesh_range_large_start(): + mesh = rx.distributed.DeviceMesh((2,), Range(2**32, 2**32 + 2)) + assert list(mesh.device_ids) == [2**32, 2**32 + 1] + + +@pytest.mark.parametrize( + "start,extent,shape", + [(0, -1, (0,)), (2**63 - 1, 2, (2,)), (tirx.Var("start", "int64"), 1, (1,))], +) +def test_device_mesh_invalid_range(start, extent, shape): + device_range = Range.from_min_extent(start, extent) + with pytest.raises(ValueError): + rx.distributed.DeviceMesh(shape, device_range) + + def test_dtensor_type(): n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")