From b1981dd64b96716eafedfd37d134f3b4bca29f12 Mon Sep 17 00:00:00 2001 From: Sean Nijjar Date: Mon, 16 Dec 2024 17:33:20 +0000 Subject: [PATCH 1/2] add centralized CCL pybind file --- ttnn/cpp/pybind11/operations/__init__.hpp | 8 ++--- .../ccl/all_gather/all_gather_pybind.cpp | 5 ---- ttnn/cpp/ttnn/operations/ccl/ccl_pybind.cpp | 18 ++++++++++++ ttnn/cpp/ttnn/operations/ccl/ccl_pybind.hpp | 29 +++++++++++++++++++ 4 files changed, 49 insertions(+), 11 deletions(-) create mode 100644 ttnn/cpp/ttnn/operations/ccl/ccl_pybind.cpp create mode 100644 ttnn/cpp/ttnn/operations/ccl/ccl_pybind.hpp diff --git a/ttnn/cpp/pybind11/operations/__init__.hpp b/ttnn/cpp/pybind11/operations/__init__.hpp index b2bcb41671b..b4522d99b5b 100644 --- a/ttnn/cpp/pybind11/operations/__init__.hpp +++ b/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); diff --git a/ttnn/cpp/ttnn/operations/ccl/all_gather/all_gather_pybind.cpp b/ttnn/cpp/ttnn/operations/ccl/all_gather/all_gather_pybind.cpp index 210ac9d243a..b8404b749aa 100644 --- a/ttnn/cpp/ttnn/operations/ccl/all_gather/all_gather_pybind.cpp +++ b/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 void bind_all_gather(pybind11::module& module, const ccl_operation_t& operation, const char* doc) { - py::enum_(module, "Topology") - .value("Ring", ttnn::ccl::Topology::Ring) - .value("Linear", ttnn::ccl::Topology::Linear); - bind_registered_operation( module, operation, diff --git a/ttnn/cpp/ttnn/operations/ccl/ccl_pybind.cpp b/ttnn/cpp/ttnn/operations/ccl/ccl_pybind.cpp new file mode 100644 index 00000000000..00c9f4d36f0 --- /dev/null +++ b/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_(module, "Topology") + .value("Ring", ttnn::ccl::Topology::Ring) + .value("Linear", ttnn::ccl::Topology::Linear); +} + +} // namespace ttnn::operations::ccl diff --git a/ttnn/cpp/ttnn/operations/ccl/ccl_pybind.hpp b/ttnn/cpp/ttnn/operations/ccl/ccl_pybind.hpp new file mode 100644 index 00000000000..4c15186eda2 --- /dev/null +++ b/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 +#include + +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 From d0d12b7bef85bb9ac4ea8d193247240350edd6da Mon Sep 17 00:00:00 2001 From: Sean Nijjar Date: Mon, 16 Dec 2024 17:53:47 +0000 Subject: [PATCH 2/2] add centralized CCL experimental python module binding --- ttnn/cpp/ttnn/operations/ccl/CMakeLists.txt | 1 + .../operations/experimental/ccl/CMakeLists.txt | 1 + .../ccl/ccl_experimental_pybind.cpp | 18 ++++++++++++++++++ .../ccl/ccl_experimental_pybind.hpp | 16 ++++++++++++++++ .../experimental/experimental_pybind.cpp | 8 +++----- 5 files changed, 39 insertions(+), 5 deletions(-) create mode 100644 ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.cpp create mode 100644 ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.hpp diff --git a/ttnn/cpp/ttnn/operations/ccl/CMakeLists.txt b/ttnn/cpp/ttnn/operations/ccl/CMakeLists.txt index 148d928be91..1d68e569f8f 100644 --- a/ttnn/cpp/ttnn/operations/ccl/CMakeLists.txt +++ b/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 diff --git a/ttnn/cpp/ttnn/operations/experimental/ccl/CMakeLists.txt b/ttnn/cpp/ttnn/operations/experimental/ccl/CMakeLists.txt index 82767c44a09..1a6ce549e40 100644 --- a/ttnn/cpp/ttnn/operations/experimental/ccl/CMakeLists.txt +++ b/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 diff --git a/ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.cpp b/ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.cpp new file mode 100644 index 00000000000..07df56d1f98 --- /dev/null +++ b/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 diff --git a/ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.hpp b/ttnn/cpp/ttnn/operations/experimental/ccl/ccl_experimental_pybind.hpp new file mode 100644 index 00000000000..91c022ea0fd --- /dev/null +++ b/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 +#include + +namespace py = pybind11; + +namespace ttnn::operations::experimental::ccl { + +void py_module(py::module& module); + +} // namespace ttnn::operations::experimental::ccl diff --git a/ttnn/cpp/ttnn/operations/experimental/experimental_pybind.cpp b/ttnn/cpp/ttnn/operations/experimental/experimental_pybind.cpp index d6f9431947f..1facf7093d0 100644 --- a/ttnn/cpp/ttnn/operations/experimental/experimental_pybind.cpp +++ b/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