Add guards for USE_C10D_FOO in relevant c10d files (#59697)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59697
The c10d build process selectively adds files based on the `USE_C10D_FOO` flags (where `FOO` is one of `GLOO`, `NCCL` or `MPI`). Replicating this logic inside libtorch will be harder, since libtorch uses a simpler approach (i.e., it lists the files in `build_variables.bzl`). So instead we could always include all files, and "disable" each file as needed using `#ifdef`s. Note that this is not a new approach: we already do the same for all the files of the TensorPipe agent based on the flag `USE_TENSORPIPE`.
ghstack-source-id: 131169540
Test Plan: CI
Reviewed By: agolynski
Differential Revision: D28987577
fbshipit-source-id: 4c6195de4e9a58101dad9379537e8d055dfd38af
diff --git a/BUILD.bazel b/BUILD.bazel
index 2cdf157..8a116d6 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -1875,8 +1875,6 @@
"torch/lib/c10d/*.hpp",
],
exclude = [
- "torch/lib/c10d/ProcessGroupMPI.hpp",
- "torch/lib/c10d/ProcessGroupNCCL.hpp",
"torch/csrc/autograd/generated/VariableType.h",
"torch/csrc/autograd/generated/RegistrationDeclarations.h",
"torch/csrc/autograd/generated/variable_factories.h",
diff --git a/torch/lib/c10d/GlooDeviceFactory.cpp b/torch/lib/c10d/GlooDeviceFactory.cpp
index 01a6445..4166764 100644
--- a/torch/lib/c10d/GlooDeviceFactory.cpp
+++ b/torch/lib/c10d/GlooDeviceFactory.cpp
@@ -1,5 +1,7 @@
#include <c10d/GlooDeviceFactory.hpp>
+#ifdef USE_C10D_GLOO
+
#include <stdlib.h>
#include <c10/util/Exception.h>
@@ -162,3 +164,5 @@
}
} // namespace c10d
+
+#endif // USE_C10D_GLOO
diff --git a/torch/lib/c10d/GlooDeviceFactory.hpp b/torch/lib/c10d/GlooDeviceFactory.hpp
index 58821a9..7d03818 100644
--- a/torch/lib/c10d/GlooDeviceFactory.hpp
+++ b/torch/lib/c10d/GlooDeviceFactory.hpp
@@ -1,5 +1,7 @@
#pragma once
+#ifdef USE_C10D_GLOO
+
#include <string>
#include <c10/util/Registry.h>
@@ -26,3 +28,5 @@
const std::string& /* hostname */);
} // namespace c10d
+
+#endif // USE_C10D_GLOO
diff --git a/torch/lib/c10d/NCCLUtils.cpp b/torch/lib/c10d/NCCLUtils.cpp
index d0245fe..9e0566a 100644
--- a/torch/lib/c10d/NCCLUtils.cpp
+++ b/torch/lib/c10d/NCCLUtils.cpp
@@ -1,4 +1,7 @@
#include <c10d/NCCLUtils.hpp>
+
+#ifdef USE_C10D_NCCL
+
#include <mutex>
namespace c10d {
@@ -32,3 +35,5 @@
}
} // namespace c10d
+
+#endif // USE_C10D_NCCL
diff --git a/torch/lib/c10d/NCCLUtils.hpp b/torch/lib/c10d/NCCLUtils.hpp
index 8462930..0dec457 100644
--- a/torch/lib/c10d/NCCLUtils.hpp
+++ b/torch/lib/c10d/NCCLUtils.hpp
@@ -1,5 +1,7 @@
#pragma once
+#ifdef USE_C10D_NCCL
+
#include <stdio.h>
#include <stdlib.h>
@@ -200,3 +202,5 @@
};
} // namespace c10d
+
+#endif // USE_C10D_NCCL
diff --git a/torch/lib/c10d/ProcessGroupGloo.cpp b/torch/lib/c10d/ProcessGroupGloo.cpp
index 1b413c8..d423271 100644
--- a/torch/lib/c10d/ProcessGroupGloo.cpp
+++ b/torch/lib/c10d/ProcessGroupGloo.cpp
@@ -1,6 +1,7 @@
-#include <c10/util/irange.h>
#include <c10d/ProcessGroupGloo.hpp>
+#ifdef USE_C10D_GLOO
+
#include <c10d/GlooDeviceFactory.hpp>
#include <chrono>
#include <exception>
@@ -2816,3 +2817,5 @@
}
} // namespace c10d
+
+#endif // USE_C10D_GLOO
diff --git a/torch/lib/c10d/ProcessGroupGloo.hpp b/torch/lib/c10d/ProcessGroupGloo.hpp
index eeaf868..32e3799 100644
--- a/torch/lib/c10d/ProcessGroupGloo.hpp
+++ b/torch/lib/c10d/ProcessGroupGloo.hpp
@@ -1,5 +1,7 @@
#pragma once
+#ifdef USE_C10D_GLOO
+
#include <condition_variable>
#include <deque>
#include <mutex>
@@ -357,3 +359,5 @@
};
} // namespace c10d
+
+#endif // USE_C10D_GLOO
diff --git a/torch/lib/c10d/ProcessGroupMPI.cpp b/torch/lib/c10d/ProcessGroupMPI.cpp
index 1c56e68..0c47121 100644
--- a/torch/lib/c10d/ProcessGroupMPI.cpp
+++ b/torch/lib/c10d/ProcessGroupMPI.cpp
@@ -1,5 +1,7 @@
#include <c10d/ProcessGroupMPI.hpp>
+#ifdef USE_C10D_MPI
+
#include <limits>
#include <map>
@@ -920,3 +922,5 @@
}
} // namespace c10d
+
+#endif // USE_C10D_MPI
diff --git a/torch/lib/c10d/ProcessGroupMPI.hpp b/torch/lib/c10d/ProcessGroupMPI.hpp
index b63448e..9536331 100644
--- a/torch/lib/c10d/ProcessGroupMPI.hpp
+++ b/torch/lib/c10d/ProcessGroupMPI.hpp
@@ -1,5 +1,7 @@
#pragma once
+#ifdef USE_C10D_MPI
+
#include <condition_variable>
#include <deque>
#include <exception>
@@ -263,3 +265,5 @@
};
} // namespace c10d
+
+#endif // USE_C10D_MPI
diff --git a/torch/lib/c10d/ProcessGroupNCCL.cpp b/torch/lib/c10d/ProcessGroupNCCL.cpp
index 7e0320d..3f62cab 100644
--- a/torch/lib/c10d/ProcessGroupNCCL.cpp
+++ b/torch/lib/c10d/ProcessGroupNCCL.cpp
@@ -1,7 +1,7 @@
-#include <c10/util/irange.h>
-#include <c10/util/Optional.h>
#include <c10d/ProcessGroupNCCL.hpp>
+#ifdef USE_C10D_NCCL
+
#include <exception>
#include <map>
#include <stdexcept>
@@ -12,7 +12,9 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
+#include <c10/util/irange.h>
#include <c10/util/Logging.h>
+#include <c10/util/Optional.h>
#include <c10d/ParamCommsUtils.hpp>
#include <torch/csrc/cuda/nccl.h>
@@ -1925,3 +1927,5 @@
}
} // namespace c10d
+
+#endif // USE_C10D_NCCL
diff --git a/torch/lib/c10d/ProcessGroupNCCL.hpp b/torch/lib/c10d/ProcessGroupNCCL.hpp
index 5a467ff..bafb76d 100644
--- a/torch/lib/c10d/ProcessGroupNCCL.hpp
+++ b/torch/lib/c10d/ProcessGroupNCCL.hpp
@@ -1,5 +1,7 @@
#pragma once
+#ifdef USE_C10D_NCCL
+
#include <chrono>
#include <iostream>
#include <list>
@@ -560,3 +562,5 @@
};
} // namespace c10d
+
+#endif // USE_C10D_NCCL
diff --git a/torch/lib/c10d/ProcessGroupWrapper.cpp b/torch/lib/c10d/ProcessGroupWrapper.cpp
index 7d0da4d..c31e6f6 100644
--- a/torch/lib/c10d/ProcessGroupWrapper.cpp
+++ b/torch/lib/c10d/ProcessGroupWrapper.cpp
@@ -1,3 +1,7 @@
+#include <c10d/ProcessGroupWrapper.hpp>
+
+#ifdef USE_C10D_GLOO
+
#include <c10/core/ScalarType.h>
#include <c10/core/TensorOptions.h>
#include <c10/util/Exception.h>
@@ -6,7 +10,6 @@
#include <c10/util/irange.h>
#include <c10d/ProcessGroup.hpp>
#include <c10d/ProcessGroupGloo.hpp>
-#include <c10d/ProcessGroupWrapper.hpp>
#include <stdexcept>
namespace c10d {
@@ -321,3 +324,5 @@
}
} // namespace c10d
+
+#endif // USE_C10D_GLOO
diff --git a/torch/lib/c10d/ProcessGroupWrapper.hpp b/torch/lib/c10d/ProcessGroupWrapper.hpp
index ea80ea0..9ee4355 100644
--- a/torch/lib/c10d/ProcessGroupWrapper.hpp
+++ b/torch/lib/c10d/ProcessGroupWrapper.hpp
@@ -1,5 +1,7 @@
#pragma once
+#ifdef USE_C10D_GLOO
+
#include <c10d/ProcessGroup.hpp>
#include <c10d/ProcessGroupGloo.hpp>
#include <c10d/Types.hpp>
@@ -124,3 +126,5 @@
const std::vector<at::Tensor>& tensors) const;
};
} // namespace c10d
+
+#endif // USE_C10D_GLOO
