Create standalone CCL and CCL experimental python module bindings

ttnn/cpp/pybind11/operations/__init__.hpp

@@ -11,9 +11,7 @@
 #include "pybind11/operations/core.hpp"
 #include "pybind11/operations/creation.hpp"
 #include "ttnn/operations/bernoulli/bernoulli_pybind.hpp"
-#include "ttnn/operations/ccl/all_gather/all_gather_pybind.hpp"
-#include "ttnn/operations/ccl/reduce_scatter/reduce_scatter_pybind.hpp"
-#include "ttnn/operations/ccl/barrier/barrier_pybind.hpp"
+#include "ttnn/cpp/ttnn/operations/ccl/ccl_pybind.hpp"
 #include "ttnn/operations/conv/conv_pybind.hpp"
 #include "ttnn/operations/data_movement/data_movement_pybind.hpp"
 #include "ttnn/operations/eltwise/binary/binary_pybind.hpp"
@@ -93,9 +91,7 @@ void py_module(py::module& module) {
     complex_unary_backward::py_module(m_complex_unary_backward);
 
     auto m_ccl = module.def_submodule("ccl", "collective communication operations");
-    ccl::py_bind_all_gather(m_ccl);
-    ccl::py_bind_reduce_scatter(m_ccl);
-    ccl::py_bind_barrier(m_ccl);
+    ccl::py_module(m_ccl);
 
     auto m_creation = module.def_submodule("creation", "creation operations");
     creation::py_module(m_creation);

ttnn/cpp/ttnn/operations/ccl/CMakeLists.txt

@@ -7,6 +7,7 @@ set(CCL_TTNN_SRCS
     ${CMAKE_CURRENT_SOURCE_DIR}/common/types/ccl_types_args_emitters.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/common/uops/ccl_command.cpp
     # CCL Ops
+    ${CMAKE_CURRENT_SOURCE_DIR}/ccl_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/all_gather/all_gather.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/all_gather/all_gather_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/all_gather/device/all_gather_op.cpp

ttnn/cpp/ttnn/operations/ccl/all_gather/all_gather_pybind.cpp

@@ -9,7 +9,6 @@
 
 #include "ttnn/cpp/pybind11/decorators.hpp"
 #include "ttnn/operations/ccl/all_gather/all_gather.hpp"
-#include "ttnn/operations/ccl/ccl_host_datastructures.hpp"
 #include "ttnn/distributed/types.hpp"
 
 namespace ttnn::operations::ccl {
@@ -18,10 +17,6 @@ namespace detail {
 
 template <typename ccl_operation_t>
 void bind_all_gather(pybind11::module& module, const ccl_operation_t& operation, const char* doc) {
-    py::enum_<ttnn::ccl::Topology>(module, "Topology")
-        .value("Ring", ttnn::ccl::Topology::Ring)
-        .value("Linear", ttnn::ccl::Topology::Linear);
-
     bind_registered_operation(
         module,
         operation,

ttnn/cpp/ttnn/operations/ccl/ccl_pybind.cpp

@@ -0,0 +1,18 @@
+// SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "ttnn/cpp/ttnn/operations/ccl/ccl_pybind.hpp"
+
+#include "ttnn/operations/ccl/ccl_host_datastructures.hpp"
+#include "ttnn/cpp/ttnn/operations/ccl/erisc_datamover_builder.hpp"
+
+namespace ttnn::operations::ccl {
+
+void py_bind_common(pybind11::module& module) {
+    py::enum_<ttnn::ccl::Topology>(module, "Topology")
+        .value("Ring", ttnn::ccl::Topology::Ring)
+        .value("Linear", ttnn::ccl::Topology::Linear);
+}
+
+}  // namespace ttnn::operations::ccl

ttnn/cpp/ttnn/operations/ccl/ccl_pybind.hpp

@@ -0,0 +1,29 @@
+// SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "ttnn/operations/ccl/all_gather/all_gather_pybind.hpp"
+#include "ttnn/operations/ccl/reduce_scatter/reduce_scatter_pybind.hpp"
+#include "ttnn/operations/ccl/barrier/barrier_pybind.hpp"
+
+// #include "ttnn/cpp/pybind11/decorators.hpp"
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+namespace ttnn::operations::ccl {
+
+void py_bind_common(pybind11::module& module);
+
+void py_module(py::module& module) {
+    ccl::py_bind_common(module);
+    ccl::py_bind_all_gather(module);
+    ccl::py_bind_reduce_scatter(module);
+    ccl::py_bind_barrier(module);
+}
+
+}  // namespace ttnn::operations::ccl

ttnn/cpp/ttnn/operations/experimental/ccl/CMakeLists.txt

@@ -1,5 +1,6 @@
 set(CCL_EXPERIMENTAL_TTNN_SRCS
     #Experimental Ops
+    ${CMAKE_CURRENT_SOURCE_DIR}/ccl_experimental_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/all_gather_matmul/all_gather_matmul.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/all_gather_matmul/all_gather_matmul_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/all_gather_matmul/device/all_gather_matmul_op.cpp

ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.cpp

@@ -0,0 +1,18 @@
+// SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.hpp"
+#include "ttnn/operations/experimental/ccl/all_gather_matmul/all_gather_matmul_pybind.hpp"
+#include "ttnn/operations/experimental/ccl/all_reduce/all_reduce_pybind.hpp"
+
+namespace py = pybind11;
+
+namespace ttnn::operations::experimental::ccl {
+
+void py_module(py::module& module) {
+    ccl::py_bind_all_gather_matmul(module);
+    ccl::py_bind_all_reduce(module);
+}
+
+}  // namespace ttnn::operations::experimental::ccl

ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.hpp

@@ -0,0 +1,16 @@
+// SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+namespace ttnn::operations::experimental::ccl {
+
+void py_module(py::module& module);
+
+}  // namespace ttnn::operations::experimental::ccl

ttnn/cpp/ttnn/operations/experimental/experimental_pybind.cpp

@@ -33,8 +33,7 @@
 #include "ttnn/cpp/ttnn/operations/experimental/copy/typecast/typecast_pybind.hpp"
 #include "ttnn/cpp/ttnn/operations/experimental/matmul/attn_matmul/attn_matmul_pybind.hpp"
 #include "ttnn/cpp/ttnn/operations/experimental/matmul/group_attn_matmul/group_attn_matmul_pybind.hpp"
-#include "ttnn/operations/experimental/ccl/all_gather_matmul/all_gather_matmul_pybind.hpp"
-#include "ttnn/operations/experimental/ccl/all_reduce/all_reduce_pybind.hpp"
+#include "ttnn/operations/experimental/ccl/ccl_experimental_pybind.hpp"
 #include "ttnn/operations/experimental/plusone/plusone_pybind.hpp"
 namespace ttnn::operations::experimental {
 
@@ -77,9 +76,8 @@ void py_module(py::module& module) {
     plusone::detail::bind_experimental_plusone_operation(module);
 
     // CCL ops
-    auto m_experimental_ccl = module.def_submodule("ccl", "experiemental collective communication operations");
-    ccl::py_bind_all_gather_matmul(m_experimental_ccl);
-    ccl::py_bind_all_reduce(m_experimental_ccl);
+    auto m_experimental_ccl = module.def_submodule("ccl", "experimental collective communication operations");
+    ccl::py_module(m_experimental_ccl);
 }
 
 }  // namespace ttnn::operations::experimental