Fix slice to selection range conversion (#29)

Summary:
Pull Request resolved: #29

Previous implementation was converting `Slice(offset, sizes=[size, 1, 1, ...], strides=[stride, ...])` into a selection with `Range(offset, offset + size, stride)`, which is wrong. It should be `Range(offset, offset + size * stride, stride)`.

Reviewed By: shayne-fletcher, andrewjcg

Differential Revision: D75102147

fbshipit-source-id: 979c3712e65a49cd4e04de71921b82baead721bc

Author: Sam Lurye

controller/src/lib.rs

@@ -310,7 +310,11 @@ fn slice_to_selection(slice: Slice) -> Selection {
         ([], []) => dsl::range(slice.offset()..=slice.offset(), dsl::true_()),
         // Special case trivial range `Selection`.
         ([size, rsizes @ ..], [stride, ..]) if rsizes.iter().all(|s| *s == 1) => dsl::range(
-            Range(slice.offset(), Some(slice.offset() + *size), *stride),
+            Range(
+                slice.offset(),
+                Some(slice.offset() + *size * *stride),
+                *stride,
+            ),
             dsl::true_(),
         ),
         // Fallback to more heavy-weight translation for everything else.

python/tests/test_controller.py

@@ -11,7 +11,6 @@
 import sys
 import traceback
 from contextlib import contextmanager
-from enum import Enum
 from typing import Generator
 
 import monarch
@@ -324,8 +323,6 @@ def test_movement(self, backend_type):
             _ = b.to_mesh(sm1)
 
     def test_broadcast_one(self, backend_type):
-        if backend_type == BackendType.RS:
-            pytest.skip("deadlocks on rust")
         with self.local_device_mesh(2, 2, backend_type) as device_mesh:
             for dim in ("host", "gpu"):
                 subset = device_mesh(**{dim: 1})
@@ -341,9 +338,6 @@ def test_broadcast_one(self, backend_type):
                     assert torch.allclose(a.expand(2, -1), b, rtol=0, atol=0)
 
     def test_broadcast_two(self, backend_type):
-        if backend_type == BackendType.RS:
-            pytest.skip("deadlocks on rust")
-
         with self.local_device_mesh(2, 2, backend_type) as device_mesh:
             subset = device_mesh(host=1, gpu=1)
             with subset.activate():