#8835: add example template of a ttnn operation

tenstorrent · Jul 2, 2024 · 0023da1 · 0023da1
1 parent 1cb7515
commit 0023da1
Show file tree

Hide file tree

Showing 13 changed files with 43 additions and 36 deletions.
diff --git a/docs/source/ttnn/ttnn/adding_new_ttnn_operation.rst b/docs/source/ttnn/ttnn/adding_new_ttnn_operation.rst
@@ -40,37 +40,37 @@ In order to add a new device operation, follow the directory structure shown bel
 .. note::
  Add as many program factories as needed
 
-A concrete example of a device operation can be found in `ttnn/cpp/ttnn/operations/example/example/device`
+A concrete example of a device operation can be found in `ttnn/cpp/ttnn/operations/examples/example/device`
 
-`ttnn/cpp/ttnn/operations/example/example/device/example_device_operation.hpp`:
+`ttnn/cpp/ttnn/operations/examples/example/device/example_device_operation.hpp`:
 
-.. literalinclude::  example_operation/example/device/example_device_operation.hpp
+.. literalinclude::  examples/example/device/example_device_operation.hpp
 
-`ttnn/cpp/ttnn/operations/example/example/device/example_device_operation.cpp`:
+`ttnn/cpp/ttnn/operations/examples/example/device/example_device_operation.cpp`:
 
-.. literalinclude::  example_operation/example/device/example_device_operation.cpp
+.. literalinclude::  examples/example/device/example_device_operation.cpp
 
-`ttnn/cpp/ttnn/operations/example/example/device/single_core_program_factory.cpp`:
+`ttnn/cpp/ttnn/operations/examples/example/device/single_core_program_factory.cpp`:
 
-.. literalinclude::  example_operation/example/device/single_core_program_factory.cpp
+.. literalinclude::  examples/example/device/single_core_program_factory.cpp
 
-`ttnn/cpp/ttnn/operations/example/example/device/multi_core_program_factory.cpp`:
+`ttnn/cpp/ttnn/operations/examples/example/device/multi_core_program_factory.cpp`:
 
-.. literalinclude::  example_operation/example/device/multi_core_program_factory.cpp
+.. literalinclude::  examples/example/device/multi_core_program_factory.cpp
 
 
 Step 2: Implement the operation in C++
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-In order to add a new operation, follow the directory structure shown below:
+In order to add a new operation, add the following file:
 
 `ttnn/cpp/ttnn/operations/<category>/<operation_name>/<operation_name>.hpp`
 
 A concrete example:
 
-`ttnn/cpp/ttnn/operations/example/example/example.hpp`:
+`ttnn/cpp/ttnn/operations/examples/example/example.hpp`:
 
-.. literalinclude::  example_operation/example/example.hpp
+.. literalinclude::  examples/example/example.hpp
 
 
 Python Implementation
@@ -79,15 +79,22 @@ Python Implementation
 Step 1: Add Python binding
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-In order to add a new operation, follow the directory structure shown below:
+In order to add a python binding for the operation, follow the directory structure shown below:
 
 `ttnn/python/ttnn/operations/<category>/<operation_name>/<operation_name>_pybind.hpp`
+`ttnn/python/ttnn/operations/<category>/<category>_pybind.hpp`
 
 A concrete example:
 
-`ttnn/python/ttnn/operations/example/example/example_pybind.hpp`:
+`ttnn/python/ttnn/operations/examples/example/example_pybind.hpp`:
+
+.. literalinclude::  examples/example/example_pybind.hpp
+
+`ttnn/python/ttnn/operations/examples/examples_pybind.hpp`:
+
+.. literalinclude::  examples/example/example_pybind.hpp
 
-.. literalinclude::  example_operation/example/example_pybind.hpp
+Finally, call the module defined in `examples/example/example_pybind.hpp` wherever you want it to be added.
 
 
 

diff --git a/docs/source/ttnn/ttnn/example_operation b/docs/source/ttnn/ttnn/example_operation
diff --git a/docs/source/ttnn/ttnn/examples b/docs/source/ttnn/ttnn/examples
@@ -0,0 +1 @@
+../../../../ttnn/cpp/ttnn/operations/examples/
diff --git a/ttnn/CMakeLists.txt b/ttnn/CMakeLists.txt
@@ -11,9 +11,9 @@ set(TTNN_SRCS
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/eltwise/binary/device/broadcast_width_multi_core_program_factory.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/eltwise/binary/device/element_wise_multi_core_program_factory.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/eltwise/unary/device/unary_op.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/example/example/device/example_device_operation.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/example/example/device/single_core_program_factory.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/example/example/device/multi_core_program_factory.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/examples/example/device/example_device_operation.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/examples/example/device/single_core_program_factory.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/examples/example/device/multi_core_program_factory.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/reduction/topk/device/topk_op.cpp
 )
 

diff --git a/ttnn/cpp/pybind11/operations/__init__.hpp b/ttnn/cpp/pybind11/operations/__init__.hpp
@@ -24,7 +24,7 @@
 #include "ttnn/operations/eltwise/binary/binary_pybind.hpp"
 #include "ttnn/operations/eltwise/binary_backward/binary_backward_pybind.hpp"
 #include "ttnn/operations/eltwise/unary/unary_pybind.hpp"
-#include "ttnn/operations/example/example_pybind.hpp"
+#include "ttnn/operations/examples/examples_pybind.hpp"
 #include "ttnn/operations/reduction/reduction_pybind.hpp"
 
 namespace py = pybind11;
@@ -35,7 +35,7 @@ namespace operations {
 
 void py_module(py::module& module) {
     auto m_example = module.def_submodule("example", "example operation");
-    example::py_module(m_example);
+    examples::py_module(m_example);
 
     auto m_unary = module.def_submodule("unary", "unary operations");
     unary::py_module(m_unary);

diff --git a/ttnn/cpp/ttnn/device_operation.hpp b/ttnn/cpp/ttnn/device_operation.hpp
@@ -80,7 +80,7 @@ concept DeviceOperationConcept = requires {
 };
 
 template <typename device_operation_t>
-concept DeviceOperationConceptWithCustomProgramCacheConcept = DeviceOperationConcept<device_operation_t> and requires {
+concept DeviceOperationWithCustomProgramCacheConcept = DeviceOperationConcept<device_operation_t> and requires {
     [](auto&& program_factory,
        const typename device_operation_t::operation_attributes_t& operation_attributes,
        const typename device_operation_t::tensor_args_t& tensor_args) {
@@ -101,7 +101,7 @@ template <typename device_operation_t>
 inline auto compute_program_hash(
     const typename device_operation_t::operation_attributes_t& operation_attributes,
     const typename device_operation_t::tensor_args_t& tensor_args) {
-    if constexpr (DeviceOperationConceptWithCustomProgramCacheConcept<device_operation_t>) {
+    if constexpr (DeviceOperationWithCustomProgramCacheConcept<device_operation_t>) {
         ZoneScopedN("Compute custom program hash");
         return device_operation_t::compute_program_hash(operation_attributes, tensor_args);
     } else {

diff --git a/...ample/device/example_device_operation.cpp → ...ample/device/example_device_operation.cpp b/...ample/device/example_device_operation.cpp → ...ample/device/example_device_operation.cpp
@@ -4,7 +4,7 @@
 
 #include "example_device_operation.hpp"
 
-namespace ttnn::operations::example {
+namespace ttnn::operations::examples {
 
 ExampleDeviceOperation::program_factory_t ExampleDeviceOperation::select_program_factory(
     const operation_attributes_t& operation_attributes, const tensor_args_t& tensor_args) {

diff --git a/...ample/device/example_device_operation.hpp → ...ample/device/example_device_operation.hpp b/...ample/device/example_device_operation.hpp → ...ample/device/example_device_operation.hpp
@@ -12,7 +12,7 @@
 #include "ttnn/device_operation.hpp"
 #include "ttnn/types.hpp"
 
-namespace ttnn::operations::example {
+namespace ttnn::operations::examples {
 
 struct ExampleDeviceOperation {
     struct operation_attributes_t {

diff --git a/...ple/device/multi_core_program_factory.cpp → ...ple/device/multi_core_program_factory.cpp b/...ple/device/multi_core_program_factory.cpp → ...ple/device/multi_core_program_factory.cpp
@@ -4,7 +4,7 @@
 
 #include "example_device_operation.hpp"
 
-namespace ttnn::operations::example {
+namespace ttnn::operations::examples {
 ExampleDeviceOperation::MultiCore::cached_program_t ExampleDeviceOperation::MultiCore::create(
     const operation_attributes_t& operation_attributes,
     const tensor_args_t& tensor_args,

diff --git a/...le/device/single_core_program_factory.cpp → ...le/device/single_core_program_factory.cpp b/...le/device/single_core_program_factory.cpp → ...le/device/single_core_program_factory.cpp
@@ -4,7 +4,7 @@
 
 #include "example_device_operation.hpp"
 
-namespace ttnn::operations::example {
+namespace ttnn::operations::examples {
 ExampleDeviceOperation::SingleCore::cached_program_t ExampleDeviceOperation::SingleCore::create(
     const operation_attributes_t& operation_attributes,
     const tensor_args_t& tensor_args,
@@ -27,4 +27,4 @@ void ExampleDeviceOperation::SingleCore::override_runtime_arguments(
         cached_program.shared_variables.some_variable_from_create_to_use_in_override_runtime_arguments;
 }
 
-}  // namespace ttnn::operations::example
+}  // namespace ttnn::operations::examples
diff --git a/...nn/operations/example/example/example.hpp → ...n/operations/examples/example/example.hpp b/...nn/operations/example/example/example.hpp → ...n/operations/examples/example/example.hpp
@@ -7,7 +7,7 @@
 
 #include "device/example_device_operation.hpp"
 
-namespace ttnn::operations::example {
+namespace ttnn::operations::examples {
 
 struct ExampleOperation {
 
@@ -30,6 +30,6 @@ struct ExampleOperation {
 
 namespace ttnn {
 
-constexpr auto example = ttnn::register_operation<operations::example::ExampleOperation>("ttnn::example");
+constexpr auto example = ttnn::register_operation<operations::examples::ExampleOperation>("ttnn::example");
 
 }  // namespace ttnn
diff --git a/...ations/example/example/example_pybind.hpp → ...tions/examples/example/example_pybind.hpp b/...ations/example/example/example_pybind.hpp → ...tions/examples/example/example_pybind.hpp
@@ -8,12 +8,12 @@
 #include <pybind11/stl.h>
 
 #include "ttnn/cpp/pybind11/decorators.hpp"
-#include "ttnn/operations/example/example/example.hpp"
+#include "ttnn/operations/examples/example/example.hpp"
 #include "ttnn/types.hpp"
 
 namespace py = pybind11;
 
-namespace ttnn::operations::example {
+namespace ttnn::operations::examples {
 
 void bind_example_operation(py::module& module) {
     bind_registered_operation(
@@ -28,4 +28,4 @@ void bind_example_operation(py::module& module) {
             py::arg("queue_id") = 0});
 }
 
-}  // namespace ttnn::operations::example
+}  // namespace ttnn::operations::examples
diff --git a/...tnn/operations/example/example_pybind.hpp → ...n/operations/examples/examples_pybind.hpp b/...tnn/operations/example/example_pybind.hpp → ...n/operations/examples/examples_pybind.hpp
@@ -6,10 +6,10 @@
 
 #include <pybind11/pybind11.h>
 
-#include "ttnn/operations/example/example/example_pybind.hpp"
+#include "ttnn/operations/examples/example/example_pybind.hpp"
 
-namespace ttnn::operations::example {
+namespace ttnn::operations::examples {
 
 void py_module(py::module& module) { bind_example_operation(module); }
 
-}  // namespace ttnn::operations::example
+}  // namespace ttnn::operations::examples