use shard_tensor, rm print

csrc/ir/nodes.cpp

@@ -4519,7 +4519,7 @@ std::vector<int64_t> computeStrides(
 
   std::optional<std::vector<int64_t>> out_order = ir_utils::computePermutation(
       TensorDomain::noReductions(logical_domain),
-      TensorDomain::noReductions(allocation_domain));
+      TensorDomain::noDevices(TensorDomain::noReductions(allocation_domain)));
   NVF_CHECK(
       out_order.has_value(),
       "Valid permute from logical to allocation domain was not found.");

tests/python/test_multidevice.py

@@ -162,7 +162,7 @@ def multidevice_schedule(self) -> None:
 
 
 @pytest.mark.mpi
-def test_matmul_loop_split(mpi_test):
+def test_matmul_loop_split(multidevice_test):
     class Model(FusionDefinition):
         def __init__(self, num_devices, batch, sequence, hidden):
             super().__init__()
@@ -174,7 +174,7 @@ def __init__(self, num_devices, batch, sequence, hidden):
         def definition(self):
             d, b, s, e = self._num_devices, self._batch, self._sequence, self._hidden
             self.inp = self.define_tensor([b, s, e])
-            self.weight = self.define_tensor([e, d * e], contiguity=[True, True])
+            self.weight = self.define_tensor([e, d * e])
             self.out = self.ops.matmul(self.inp, self.weight)
             self.add_output(self.out)
 
@@ -193,26 +193,23 @@ def multidevice_schedule(self):
             self.sched.parallelize(self.out, -3, nvfuser.ParallelType.mesh_x)
             self.sched.set_allocation_as_loop(self.out)
 
-    d = mpi_test.size
+    d = multidevice_test.size
     mesh = nvfuser.DeviceMesh(range(d))
-    rank = mpi_test.rank
+    rank = multidevice_test.rank
 
-    torch.cuda.set_device(mpi_test.local_rank)
+    torch.cuda.set_device(multidevice_test.local_rank)
 
     b, s, e = 2, 1024, 768
     inp_tensor = torch.randn(b, s, e, device="cuda")
     unsharded_weight_tensor = torch.randn(e, d * e)
-    sharded_weight_tensor = mpi_test.shard_tensor(unsharded_weight_tensor, -1, mesh)
+    sharded_weight_tensor = multidevice_test.shard_tensor(unsharded_weight_tensor, -1, mesh)
 
     fd = Model(d, b, s, e)
     out_tensors = fd.execute([inp_tensor, sharded_weight_tensor])
-    print(f"Output tensor: {out_tensors[0].shape}")
 
     # [b, s, d*e]
-    unsharded_out_tensor = torch.matmul(inp_tensor, unsharded_weight_tensor.cuda())
-    expected_out_tensor = unsharded_out_tensor.view([b, s, d, e]).permute(2, 0, 1, 3)[
-        rank : rank + 1
-    ]
+    unsharded_out_tensor = torch.matmul(inp_tensor.cpu(), unsharded_weight_tensor)
+    expected_out_tensor = multidevice_test.shard_tensor(unsharded_out_tensor, -1, mesh)
     # rtol is the same as the default for fp32. atol is slightly increased.
     torch.testing.assert_close(
         out_tensors[0], expected_out_tensor.squeeze(0), rtol=1.3e-6, atol=1e-3