pytorch
diff --git a/‎test/distributed/test_c10d_functional_native.py
Copy file name to clipboard
+18Lines changed: 18 additions & 0 deletions b/‎test/distributed/test_c10d_functional_native.py
Copy file name to clipboard
+18Lines changed: 18 additions & 0 deletions
diff --git a/‎test/distributed/test_c10d_nccl.py
Copy file name to clipboard
+42Lines changed: 42 additions & 0 deletions b/‎test/distributed/test_c10d_nccl.py
Copy file name to clipboard
+42Lines changed: 42 additions & 0 deletions
diff --git a/‎test/distributed/test_inductor_collectives.py
Copy file name to clipboard
+71-1Lines changed: 71 additions & 1 deletion b/‎test/distributed/test_inductor_collectives.py
Copy file name to clipboard
+71-1Lines changed: 71 additions & 1 deletion
diff --git a/‎torch/csrc/distributed/c10d/Functional.cpp
Copy file name to clipboard
+4-95Lines changed: 4 additions & 95 deletions b/‎torch/csrc/distributed/c10d/Functional.cpp
Copy file name to clipboard
+4-95Lines changed: 4 additions & 95 deletions
diff --git a/‎torch/csrc/distributed/c10d/Functional.hpp
Copy file name to clipboard
+1-9Lines changed: 1 addition & 9 deletions b/‎torch/csrc/distributed/c10d/Functional.hpp
Copy file name to clipboard
+1-9Lines changed: 1 addition & 9 deletions
@@ -405,18 +405,36 @@ def test_broadcast(self) -> None:
         assert output.eq(expect).all()
         assert output.completed
 
+    @skip_if_lt_x_gpu(2)
+    def test_wait_tensor(self) -> None:
+        self._init_process_group()
+
+        input = torch.full((10, 10), float(self.rank), device=self.device)
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
+        output = torch.ops._c10d_functional.all_reduce(
+            input,
+            "avg",
+            "default",
+        )
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 1)
+        torch.ops._c10d_functional.wait_tensor(output)
+        # `wait_tensor(output)` will pop the work from the work registry immediately
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
+
     @skip_if_lt_x_gpu(2)
     def test_unwaited(self) -> None:
         # Verify that the process can terminate gracefully
         # even with unwaited tensors
         self._init_process_group()
 
         input = torch.full((10, 10), float(self.rank), device=self.device)
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
         output = torch.ops._c10d_functional.all_reduce(
             input,
             "avg",
             "default",
         )
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 1)
 
     @skip_if_lt_x_gpu(2)
     def test_py_work(self) -> None:
 
@@ -3088,6 +3088,48 @@ def test_nccl_barrier_device_ids_function_argument(self):
         with self.assertRaisesRegex(TypeError, "Invalid function argument"):
             c10d.barrier(device_ids=self.rank)
 
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    def test_unwaited(self) -> None:
+        # Verify that the process can terminate gracefully
+        # even with unwaited tensors
+        store = c10d.FileStore(self.file_name, self.world_size)
+        c10d.init_process_group(
+            backend="nccl", rank=self.rank, world_size=self.world_size, store=store
+        )
+
+        input = torch.full((10240, 10240), float(self.rank), device=f"cuda:{self.rank}")
+        dist.all_reduce(input, op=dist.ReduceOp.SUM, async_op=True)
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 1)
+        # Running another collective on the same tensor should still work
+        dist.all_reduce(input, op=dist.ReduceOp.SUM, async_op=True)
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 2)
+
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    def test_wait_tensor(self) -> None:
+        # Verify that c10d_functional.wait_tensor() can be invoked on
+        # output tensor of non-functional collective
+        store = c10d.FileStore(self.file_name, self.world_size)
+        c10d.init_process_group(
+            backend="nccl", rank=self.rank, world_size=self.world_size, store=store
+        )
+
+        input1 = torch.full((10, 10), float(self.rank), device=f"cuda:{self.rank}")
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
+        dist.all_reduce(input1, op=dist.ReduceOp.SUM, async_op=True)
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 1)
+        torch.ops.c10d_functional.wait_tensor(input1)
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
+
+        input2 = torch.full((10, 10), float(self.rank), device=f"cuda:{self.rank}")
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
+        work = dist.all_reduce(input2, op=dist.ReduceOp.SUM, async_op=True)
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 1)
+        work.wait()
+        self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
+        self.assertEqual(input1, input2)
+
     @requires_nccl()
     @skip_if_lt_x_gpu(2)
     @with_dist_debug_levels(levels=["DETAIL"])
 
@@ -1,4 +1,5 @@
 # Owner(s): ["module: dynamo"]
+import datetime
 import functools
 import unittest
 from unittest.mock import patch
@@ -14,7 +15,7 @@
 from torch._dynamo.testing import CompileCounter
 from torch._dynamo.utils import same
 from torch._inductor.compile_fx import compile_fx as inductor_compile_fx
-from torch._inductor.utils import run_and_get_triton_code
+from torch._inductor.utils import run_and_get_code, run_and_get_triton_code
 from torch.distributed.distributed_c10d import GroupMember
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.testing._internal.common_distributed import (
@@ -28,6 +29,7 @@
     instantiate_parametrized_tests,
     parametrize,
     requires_cuda,
+    skipIfRocm,
 )
 from torch.testing._internal.inductor_utils import HAS_GPU
 
@@ -245,6 +247,74 @@ def compile(func, example_inputs):
             )
             self.assertTrue(same(eager_out, inductor_out, tol=0.001))
 
+    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
+    @skip_if_lt_x_gpu(2)
+    @skipIfRocm
+    def test_eager_async_allreduce_inductor_wait(self):
+        import torch.distributed as dist
+
+        def all_reduce_non_functional_eager(x):
+            y = x * x
+            work = dist.all_reduce(y, op=dist.ReduceOp.SUM, async_op=True)
+            assert isinstance(work, torch.distributed.Work)
+            return work, y
+
+        def all_reduce_wait(work, y):  # potentially compiled
+            if torch.compiler.is_dynamo_compiling():
+                torch.ops.c10d_functional.wait_tensor(y)
+            else:
+                work.wait(datetime.timedelta(seconds=10))
+            # Under compile, if `wait_tensor(y)` above is correctly executed,
+            # `y`'s data is in its final form and the output of this function will match eager;
+            # otherwise, `y * y` will run in parallel with `all_reduce(y)` and the output of this function
+            # will not match eager.
+            return y * y
+
+        with _dynamo_dist_per_rank_init(self.rank, self.world_size):
+            x = torch.ones(12800, 12800, device="cuda") + self.rank
+            self.assertEqual(torch._C._distributed_c10d._get_work_registry_size(), 0)
+
+            # NOTE: We run for 10 iterations each, to ensure that the GPU execution is way behind CPU
+            # and that `y * y` on CPU side will be issued before `all_reduce(y)` on GPU side is done,
+            # thus guaranteeing that in the bad case `y * y` on GPU side will run in parallel with `all_reduce(y)`
+            # thus will produce the wrong result that fails the unit test.
+
+            # Test: pure-eager
+            all_reduce_wait_eager = all_reduce_wait
+            for _ in range(10):
+                work, y = all_reduce_non_functional_eager(x)
+                self.assertEqual(
+                    torch._C._distributed_c10d._get_work_registry_size(), 1
+                )
+                out_ref = all_reduce_wait_eager(work, y)
+                # `work.wait()` will pop the work from the work registry immediately
+                self.assertEqual(
+                    torch._C._distributed_c10d._get_work_registry_size(), 0
+                )
+
+            # Test: issue comm in eager -> wait for comm in compile
+            all_reduce_wait_compiled = torch.compile(
+                all_reduce_wait,
+                backend="inductor",
+                fullgraph=True,
+            )
+            for _ in range(10):
+                work, y = all_reduce_non_functional_eager(x)
+                self.assertEqual(
+                    torch._C._distributed_c10d._get_work_registry_size(), 1
+                )
+                out_compiled, triton_codes = run_and_get_code(
+                    all_reduce_wait_compiled, work, y
+                )
+                # `wait_tensor(y)` will pop the work from the work registry immediately
+                self.assertEqual(
+                    torch._C._distributed_c10d._get_work_registry_size(), 0
+                )
+                FileCheck().check(
+                    "torch.ops._c10d_functional.wait_tensor.default("
+                ).run(triton_codes[0])
+            self.assertEqual(out_ref, out_compiled)
+
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     @skip_if_lt_x_gpu(2)
     @patch.object(torch._inductor.config, "allow_buffer_reuse", True)
 
@@ -6,80 +6,10 @@
 #include <torch/csrc/distributed/c10d/Functional.hpp>
 #include <torch/csrc/distributed/c10d/GroupRegistry.hpp>
 #include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
-#include <torch/csrc/distributed/c10d/RankLocal.hpp>
 #include <utility>
 
 namespace {
 
-class WorkRegistry {
- public:
-  void register_work(
-      const at::Tensor& tensor,
-      const c10::intrusive_ptr<c10d::Work>& work) {
-    auto storage = tensor.storage().getWeakStorageImpl();
-    std::unique_lock lock(lock_);
-    auto [it, inserted] = registry_.try_emplace(std::move(storage), work);
-    TORCH_CHECK(
-        inserted || it->second != work,
-        "The tensor storage is already associated with another work.");
-  }
-
-  c10::intrusive_ptr<c10d::Work> pop_work(const at::Tensor& tensor) {
-    const auto storage = tensor.storage().getWeakStorageImpl();
-    std::unique_lock lock(lock_);
-    auto it = registry_.find(storage);
-    if (it == registry_.end()) {
-      return nullptr;
-    }
-    auto work = it->second;
-    registry_.erase(it);
-    return work;
-  }
-
-  ~WorkRegistry() {
-    // If there are still unwaited work objects, their corresponding process
-    // groups should have already been destroyed at this stage. Any attempts to
-    // wait for these work objects or to destroy them will only result in
-    // confusing errors. Therefore, we simply issue a warning and intentionally
-    // allow the unwaited work objects to leak.
-    if (!registry_.empty()) {
-      TORCH_WARN(
-          "At the time of process termination, there are still ",
-          registry_.size(),
-          " unwaited c10d_functional collective calls. "
-          "Please review your program to ensure c10d_functional.wait_tensor() "
-          "is invoked on all tensors returned from c10d_functional collective "
-          "ops before they are used.");
-    }
-    for (auto& it : registry_) {
-      it.second.release();
-    }
-  }
-
- private:
-  std::unordered_map<
-      c10::weak_intrusive_ptr<c10::StorageImpl>,
-      c10::intrusive_ptr<c10d::Work>>
-      registry_;
-  std::mutex lock_;
-};
-
-static WorkRegistry process_registry;
-
-} // namespace
-
-namespace c10d {
-
-void register_work(
-    const at::Tensor& tensor,
-    const c10::intrusive_ptr<c10d::Work>& work) {
-  RankLocal<WorkRegistry>::get().register_work(tensor, work);
-}
-
-} // namespace c10d
-
-namespace {
-
 const std::unordered_map<std::string, c10d::ReduceOp> str_to_reduce_op = {
     {"sum", c10d::ReduceOp(c10d::ReduceOp::RedOpType::SUM)},
     {"avg", c10d::ReduceOp(c10d::ReduceOp::RedOpType::AVG)},
@@ -112,7 +42,6 @@ at::Tensor& all_reduce_(
   std::vector<at::Tensor> inputs{input};
   auto group = c10d::resolve_process_group(group_name);
   auto work = group->allreduce(inputs, opts);
-  c10d::register_work(input, work);
   return input;
 }
 
@@ -135,9 +64,6 @@ std::vector<at::Tensor> all_reduce_coalesced_(
 
   auto group = c10d::resolve_process_group(group_name);
   auto work = group->allreduce_coalesced(inputs, opts);
-  for (const auto& tensor : inputs) {
-    c10d::register_work(tensor, work);
-  }
   return inputs;
 }
 
@@ -178,9 +104,6 @@ std::vector<at::Tensor> all_gather_into_tensor_coalesced(
 
   auto group = c10d::resolve_process_group(group_name);
   auto work = group->allgather_into_tensor_coalesced(outputs, inputs);
-  for (const auto& tensor : outputs) {
-    c10d::register_work(tensor, work);
-  }
   return outputs;
 }
 
@@ -202,7 +125,6 @@ at::Tensor& all_gather_into_tensor_out(
 
   auto group = c10d::resolve_process_group(group_name);
   auto work = group->_allgather_base(output, input, opts);
-  c10d::register_work(output, work);
   return output;
 }
 
@@ -238,9 +160,6 @@ std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
 
   auto group = c10d::resolve_process_group(group_name);
   auto work = group->reduce_scatter_tensor_coalesced(outputs, inputs, opts);
-  for (const auto& tensor : outputs) {
-    c10d::register_work(tensor, work);
-  }
   return outputs;
 }
 
@@ -272,7 +191,6 @@ at::Tensor all_to_all_single(
       const_cast<at::Tensor&>(input),
       output_split_sizes,
       input_split_sizes);
-  c10d::register_work(output, work);
   return output;
 }
 
@@ -284,7 +202,6 @@ at::Tensor& broadcast_(at::Tensor& input, int64_t src, std::string group_name) {
 
   auto group = c10d::resolve_process_group(group_name);
   auto work = group->broadcast(inputs, opts);
-  c10d::register_work(input, work);
   return input;
 }
 
@@ -296,14 +213,6 @@ at::Tensor broadcast(
   return broadcast_(output, src, std::move(group_name));
 }
 
-at::Tensor wait_tensor(const at::Tensor& tensor) {
-  auto work = c10d::RankLocal<WorkRegistry>::get().pop_work(tensor);
-  if (work != nullptr) {
-    work->wait();
-  }
-  return tensor;
-}
-
 } // namespace
 
 TORCH_LIBRARY(_c10d_functional, m) {
@@ -389,7 +298,7 @@ TORCH_LIBRARY(_c10d_functional, m) {
   m.def(
       "wait_tensor(Tensor tensor) -> Tensor",
       torch::dispatch(
-          c10::DispatchKey::CompositeExplicitAutograd, ::wait_tensor),
+          c10::DispatchKey::CompositeExplicitAutograd, c10d::wait_tensor),
       {at::Tag::pt2_compliant_tag});
 }
 
@@ -438,7 +347,7 @@ class AllToAllSingle : public torch::autograd::Function<AllToAllSingle> {
     // TODO: track active cuda stream in wait
     out = c10::Dispatcher::singleton()
               .findSchemaOrThrow("_c10d_functional::wait_tensor", "")
-              .typed<decltype(wait_tensor)>()
+              .typed<decltype(c10d::wait_tensor)>()
               .call(out);
 
     return {out, at::Tensor(), at::Tensor(), at::Tensor()};
@@ -493,7 +402,7 @@ class ReduceScatterTensor
     // TODO: track active cuda stream in wait
     out = c10::Dispatcher::singleton()
               .findSchemaOrThrow("_c10d_functional::wait_tensor", "")
-              .typed<decltype(wait_tensor)>()
+              .typed<decltype(c10d::wait_tensor)>()
               .call(out);
 
     return {
@@ -549,7 +458,7 @@ class AllGatherIntoTensor
     // TODO: track active cuda stream in wait
     out = c10::Dispatcher::singleton()
               .findSchemaOrThrow("_c10d_functional::wait_tensor", "")
-              .typed<decltype(wait_tensor)>()
+              .typed<decltype(c10d::wait_tensor)>()
               .call(out);
 
     return {
 
@@ -1,11 +1,3 @@
 #pragma once
 
-#include <torch/csrc/distributed/c10d/Work.hpp>
-
-namespace c10d {
-
-C10_EXPORT void register_work(
-    const at::Tensor& tensor,
-    const c10::intrusive_ptr<c10d::Work>& work);
-
-} // namespace c10d
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>