[FIX][Relax] Validate DeviceMesh cardinality without overflow

src/relax/distributed/global_info.cc

@@ -20,19 +20,38 @@
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/distributed/global_info.h>
 
+#include <limits>
+
 namespace tvm {
 namespace relax {
 namespace distributed {
 
 TVM_FFI_STATIC_INIT_BLOCK() { DeviceMeshNode::RegisterReflection(); }
 
-DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array<int64_t> device_ids) {
-  int prod = 1;
-  for (int i = 0; i < static_cast<int>(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<int64_t>::max() / dim, ValueError)
+        << "Device mesh shape product exceeds int64";
+    size *= dim;
+  }
+  return size;
+}
+
+}  // namespace
+
+DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array<int64_t> device_ids) {
+  int64_t size = MeshSize(shape);
   ffi::ObjectPtr<DeviceMeshNode> n = ffi::make_object<DeviceMeshNode>();
-  TVM_FFI_ICHECK_EQ(prod, static_cast<int>(device_ids.size()))
+  TVM_FFI_ICHECK_EQ(static_cast<uint64_t>(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<int64_t> device_ids) {
 DeviceMesh::DeviceMesh(ffi::Shape shape, Range device_range) {
   ffi::ObjectPtr<DeviceMeshNode> n = ffi::make_object<DeviceMeshNode>();
   ffi::Array<int64_t> device_ids;
-  int range_start = device_range->min.as<IntImmNode>()->value;
-  int range_extent = device_range->extent.as<IntImmNode>()->value;
-  for (int i = range_start; i < range_start + range_extent; i++) {
-    device_ids.push_back(i);
+  const auto* start = device_range->min.as<IntImmNode>();
+  const auto* extent = device_range->extent.as<IntImmNode>();
+  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<int64_t>::max() - (range_extent - 1),
+                     ValueError)
+        << "Device mesh range exceeds int64";
   }
-  int prod = 1;
-  for (int i = 0; i < static_cast<int>(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<int>(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);

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")