#9609: Reorganize libs into ttnn (#9870)

* #9322: Move tt_eager to under ttnn as experimental

Added symlink tt_eager --> ttnn/cpp/ttnn/experimental
Removed tt_lib.so
Removed tt_eager libraries from cmake
Restructured ttnn cmake to build a single lib from ttnn + tt_eager sources
ttnn.so now contains both tt_lib and ttnn bindings
Repackaged the wheel

.github/workflows/build-artifact.yaml

@@ -36,7 +36,7 @@ jobs:
           cmake --build build --target tests
           cmake --build build --target install
       - name: 'Tar files'
-        run: tar -cvf ttm_${{ matrix.arch }}.tar build/hw build/lib tt_eager/tt_lib/*.so ttnn/ttnn/*.so build/programming_examples build/test build/tools runtime
+        run: tar -cvf ttm_${{ matrix.arch }}.tar build/hw build/lib ttnn/ttnn/*.so build/programming_examples build/test build/tools runtime
       - name: 'Upload Artifact'
         uses: actions/upload-artifact@v4
         with:

CMakeLists.txt

@@ -200,11 +200,6 @@ install(TARGETS tt_metal
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
     COMPONENT tt_build_artifacts
 )
-install(TARGETS tt_eager
-    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
-    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
-    COMPONENT tt_build_artifacts
-)
 install(TARGETS ttnn_lib
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
@@ -215,20 +210,11 @@ install(TARGETS ttnn
     COMPONENT tt_build_artifacts
 )
 if(WITH_PYTHON_BINDINGS)
-    install(TARGETS tt_lib
-        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
-        COMPONENT tt_build_artifacts
-    )
-
     # Install .so into src files for pybinds implementation
     install(FILES ${PROJECT_BINARY_DIR}/lib/_ttnn.so
     DESTINATION ${PROJECT_SOURCE_DIR}/ttnn/ttnn
     COMPONENT tt_pybinds
     )
-    install(FILES ${PROJECT_BINARY_DIR}/lib/_C.so
-    DESTINATION ${PROJECT_SOURCE_DIR}/tt_eager/tt_lib
-    COMPONENT tt_pybinds
-    )
 endif()
 
 # Temporary workaround for Issue #8767

create_venv.sh

@@ -34,7 +34,6 @@ python3 -m pip install -r $(pwd)/tt_metal/python_env/requirements-dev.txt
 
 echo "Installing tt-metal"
 pip install -e .
-pip install -e ttnn
 
 echo "Generating git hooks"
 pre-commit install

models/demos/resnet/tt/metalResnetBlock50.py

@@ -23,11 +23,11 @@
     resnet50_optimized_conv,
     resnet50_1x1_conv_s2_as_downsample_and_matmul,
 )
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.tt_py_composite_conv import (
+from ttnn.operations.conv.tt_py_composite_conv import (
     TTPyCompositeConv,
     SlidingWindowOpParamsWithParallelConfig,
 )
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.tt_py_max_pool import TTPyMaxPool
+from ttnn.operations.conv.tt_py_max_pool import TTPyMaxPool
 
 from models.utility_functions import (
     _nearest_32,

models/demos/t3000/falcon40b/scripts/get_falcon40b_perf_numbers.sh

@@ -12,7 +12,7 @@ dtype="BFLOAT8_B" #BFLOAT16
 for seq_len in ${seq_lens[@]}; do
     echo "Running seq_len: $seq_len"
     output_folder="generated/profiler/reports/"
-    WH_ARCH_YAML=wormhole_b0_80_arch_eth_dispatch.yaml python3 tt_eager/tracy.py -r -m "pytest models/demos/t3000/falcon40b/tests/test_perf_falcon.py::test_perf_bare_metal[${dtype}-DRAM-falcon_40b-layers_1-prefill_seq${seq_len}-8chips]"
+    WH_ARCH_YAML=wormhole_b0_80_arch_eth_dispatch.yaml python3 -m tracy -r -m "pytest models/demos/t3000/falcon40b/tests/test_perf_falcon.py::test_perf_bare_metal[${dtype}-DRAM-falcon_40b-layers_1-prefill_seq${seq_len}-8chips]"
     # get latest folder in output folder
     latest_created_folder=$(ls -td $output_folder/* | head -n 1)
     # find csv file that starts with "ops_perf_results"

models/utility_functions.py

@@ -14,7 +14,7 @@
 import pytest
 
 from tt_lib.fused_ops.conv import conv as TtConv
-from tt_lib.device import Arch
+from ttnn.device import Arch
 
 
 ### Math operations ###

setup.py

@@ -38,7 +38,7 @@ def get_arch_name():
     return attempt_get_env_var("ARCH_NAME")
 
 
-def get_metal_eager_local_version_scheme(metal_build_config, version):
+def get_metal_local_version_scheme(metal_build_config, version):
     from setuptools_scm.version import ScmVersion, guess_next_version
 
     arch_name = metal_build_config.arch_name
@@ -49,7 +49,7 @@ def get_metal_eager_local_version_scheme(metal_build_config, version):
         return ""
 
 
-def get_metal_eager_main_version_scheme(metal_build_config, version):
+def get_metal_main_version_scheme(metal_build_config, version):
     from setuptools_scm.version import ScmVersion, guess_next_version
 
     is_release_version = version.distance is None or version.distance == 0
@@ -71,8 +71,8 @@ def get_metal_eager_main_version_scheme(metal_build_config, version):
 
 def get_version(metal_build_config):
     return {
-        "version_scheme": partial(get_metal_eager_main_version_scheme, metal_build_config),
-        "local_scheme": partial(get_metal_eager_local_version_scheme, metal_build_config),
+        "version_scheme": partial(get_metal_main_version_scheme, metal_build_config),
+        "local_scheme": partial(get_metal_local_version_scheme, metal_build_config),
     }
 
 
@@ -97,12 +97,12 @@ def get_build_env():
             **os.environ.copy(),
             "TT_METAL_HOME": Path(__file__).parent,
             "CXX": "clang++-17",
-            # Currently, the ttnn (ttnn/_ttnn.so) and tt_lib (tt_lib/_C.so)
+            # Currently, the ttnn (ttnn/_ttnn.so)
             # both link to the tt_metal runtime. The specific thing in
             # ttnn linking to tt_metal is likely the implementation of
             # ttnn.manage_device.
             # However, because of the singleton design of
-            # tt_cluster in tt_metal, both tt_lib and ttnn will have a
+            # tt_cluster in tt_metal, ttnn will have a
             # copy of the cluster object, causing a hang during
             # device operations in ttnn, such as calling
             # output = ttnn.to_torch(output).
@@ -115,9 +115,7 @@ def get_build_env():
     # This should only run when building the wheel. Should not be running for any dev flow
     # Taking advantage of the fact devs run editable pip install -> "pip install -e ."
     def run(self) -> None:
-        assert (
-            len(self.extensions) == 2
-        ), f"Detected {len(self.extensions)} extensions, but should be only 2: tt_lib_csrc and ttnn"
+        assert len(self.extensions) == 1, f"Detected {len(self.extensions)} extensions, but should be only 2: ttnn"
 
         if self.is_editable_install_():
             assert (
@@ -145,11 +143,11 @@ def run(self) -> None:
         subprocess.check_call(["ls", "-hal", "build/lib"], cwd=source_dir, env=build_env)
         subprocess.check_call(["ls", "-hal", "runtime"], cwd=source_dir, env=build_env)
 
-        tt_build_dir = self.build_lib + "/tt_lib/build"
+        tt_build_dir = self.build_lib + "/ttnn/build"
         os.makedirs(tt_build_dir, exist_ok=True)
         self.copy_tree(source_dir / "build/lib", tt_build_dir + "/lib")
         self.copy_tree(source_dir / "runtime", self.build_lib + "/runtime")
-        arch_name_file = self.build_lib + "/tt_lib/.ARCH_NAME"
+        arch_name_file = self.build_lib + "/ttnn/.ARCH_NAME"
         subprocess.check_call(f"echo {metal_build_config.arch_name} > {arch_name_file}", shell=True)
 
         # Move built SOs into appropriate locations
@@ -174,18 +172,16 @@ def is_editable_install_(self):
 # Include tt_metal_C for kernels and src/ and tools
 # And any kernels inside `tt_eager/tt_dnn. We must keep all ops kernels inside
 # tt_dnn
-packages = ["tt_lib", "tt_metal", "tt_lib.models", "tt_eager.tt_dnn", "ttnn", "ttnn.cpp"]
+packages = ["tt_lib", "tt_metal", "tt_lib.models", "tt_eager.tt_dnn", "ttnn", "ttnn.cpp", "tracy"]
 
 # Empty sources in order to force extension executions
-eager_lib_C = Extension("tt_lib._C", sources=[])
 ttnn_lib_C = Extension("ttnn._ttnn", sources=[])
 
-ext_modules = [eager_lib_C, ttnn_lib_C]
+ext_modules = [ttnn_lib_C]
 
 BuildConstants = namedtuple("BuildConstants", ["so_src_location"])
 
 build_constants_lookup = {
-    eager_lib_C: BuildConstants(so_src_location="lib/_C.so"),
     ttnn_lib_C: BuildConstants(so_src_location="lib/_ttnn.so"),
 }
 
@@ -194,12 +190,12 @@ def is_editable_install_(self):
     use_scm_version=get_version(metal_build_config),
     packages=packages,
     package_dir={
-        "": "tt_eager",
-        "tt_metal": "tt_metal",
-        "tt_lib.models": "models",
-        "tt_eager.tt_dnn": "tt_eager/tt_dnn",
-        "ttnn": "ttnn/ttnn",
-        "ttnn.cpp": "ttnn/cpp",
+        "": "ttnn",  # only this is relevant in case of editable install mode
+        "tracy": "ttnn/tracy",
+        "tt_metal": "tt_metal",  # kernels depend on headers here
+        "tt_eager.tt_dnn": "tt_eager/tt_dnn",  # we need to pickup kernels
+        "ttnn.cpp": "ttnn/cpp",  # we need to pickup kernels
+        "tt_lib.models": "models",  # make sure ttnn does not depend on model and remove!!!
     },
     include_package_data=True,
     long_description_content_type="text/markdown",

tests/CMakeLists.txt

@@ -9,7 +9,4 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/tt_eager)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ttnn/unit_tests/gtests)
 
 set(TESTS_DEPENDS_LIST metal_tests eager_tests unit_tests_ttnn ttnn watcher_dump)
-if(WITH_PYTHON_BINDINGS)
-    list(APPEND TESTS_DEPENDS_LIST tt_lib)
-endif()
 add_custom_target(tests DEPENDS ${TESTS_DEPENDS_LIST})

tests/tt_eager/CMakeLists.txt

@@ -1,6 +1,6 @@
 
 add_library(test_eager_common_libs INTERFACE)
-target_link_libraries(test_eager_common_libs INTERFACE tt_eager test_common_libs)
+target_link_libraries(test_eager_common_libs INTERFACE ttnn test_common_libs)
 
 set(TT_EAGER_TESTS_OPS
     ops/ccl/test_all_gather_utils

tests/tt_eager/python_api_testing/sweep_tests/comparison_funcs.py

@@ -1 +1 @@
-../../../../tt_eager/tt_lib/_internal/comparison_funcs.py
+../../../../ttnn/tt_lib/_internal/comparison_funcs.py

tests/tt_eager/python_api_testing/unit_testing/misc/test_configs_for_untilize_with_halo_and_conv.py

@@ -7,7 +7,7 @@
 import numpy
 from loguru import logger
 
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.untilize_with_halo_config_generation_and_validation import (
+from ttnn.operations.conv.untilize_with_halo_config_generation_and_validation import (
     trace_conv_to_generate_data_top_left_indices_and_pad_metadata,
     validate_input_padded_tensor_and_data_top_left_indices_and_pad_metadata,
     decompose_conv_into_shards_and_generate_tensor_metadata,
@@ -17,12 +17,12 @@
     generate_untilize_with_halo_kernel_configs,
     validate_untilize_with_halo_kernel_configs,
 )
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.sliding_window_op_config_generation_and_validation import (
+from ttnn.operations.conv.sliding_window_op_config_generation_and_validation import (
     generate_sliding_window_op_sharded_input_top_left_indices,
     validate_conv_sharded_input_top_left_indices,
     validate_max_pool_sharded_input_top_left_indices,
 )
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.tt_py_untilize_with_halo import TTPyUntilizeWithHalo
+from ttnn.operations.conv.tt_py_untilize_with_halo import TTPyUntilizeWithHalo
 from tests.tt_eager.python_api_testing.sweep_tests.comparison_funcs import comp_equal, comp_allclose_and_pcc, comp_pcc
 from tt_lib.utils import _nearest_y
 import tt_lib as ttl

tests/tt_eager/python_api_testing/unit_testing/misc/test_optimized_conv_v2.py

@@ -18,7 +18,7 @@
     format_tensor,
 )
 
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.tt_py_composite_conv import (
+from ttnn.operations.conv.tt_py_composite_conv import (
     TTPyCompositeConv,
     SlidingWindowOpParams,
 )

tests/tt_eager/python_api_testing/unit_testing/misc/test_resnet50_untilize_with_halo_and_conv_v2.py

@@ -19,7 +19,7 @@
 )
 from models.utility_functions import skip_for_grayskull
 
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.tt_py_composite_conv import (
+from ttnn.operations.conv.tt_py_composite_conv import (
     TTPyCompositeConv,
     SlidingWindowOpParamsWithParallelConfig,
 )

tests/tt_eager/python_api_testing/unit_testing/misc/test_untilize_with_halo_and_max_pool_v2.py

@@ -9,7 +9,7 @@
 
 import torch
 
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.tt_py_max_pool import (
+from ttnn.operations.conv.tt_py_max_pool import (
     TTPyMaxPool,
     SlidingWindowOpParamsWithParallelConfig,
 )

tests/tt_eager/python_api_testing/unit_testing/misc/test_untilize_with_halo_v2.py

@@ -5,7 +5,7 @@
 import pytest
 import torch
 import numpy
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.untilize_with_halo_config_generation_and_validation import (
+from ttnn.operations.conv.untilize_with_halo_config_generation_and_validation import (
     trace_conv_to_generate_data_top_left_indices_and_pad_metadata,
     validate_input_padded_tensor_and_data_top_left_indices_and_pad_metadata,
     decompose_conv_into_shards_and_generate_tensor_metadata,
@@ -14,12 +14,12 @@
     validate_tensor_metadata,
     validate_untilize_with_halo_kernel_configs,
 )
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.sliding_window_op_config_generation_and_validation import (
+from ttnn.operations.conv.sliding_window_op_config_generation_and_validation import (
     generate_sliding_window_op_sharded_input_top_left_indices,
     validate_conv_sharded_input_top_left_indices,
     validate_max_pool_sharded_input_top_left_indices,
 )
-from tt_eager.tt_dnn.op_library.sliding_window_op_infra.tt_py_untilize_with_halo import (
+from ttnn.operations.conv.tt_py_untilize_with_halo import (
     TTPyUntilizeWithHalo,
     SlidingWindowOpParamsWithParallelConfig,
 )

tests/ttnn/unit_tests/gtests/CMakeLists.txt

@@ -10,7 +10,7 @@ set(TTNN_UNIT_TESTS_SRC
 
 add_executable(unit_tests_ttnn ${TTNN_UNIT_TESTS_SRC})
 
-target_link_libraries(unit_tests_ttnn PUBLIC test_common_libs ttnn_lib tt_metal tt_eager)
+target_link_libraries(unit_tests_ttnn PUBLIC test_common_libs ttnn tt_metal)
 target_include_directories(unit_tests_ttnn PRIVATE
     ${UMD_HOME}
     ${PROJECT_SOURCE_DIR}

tests/ttnn/unit_tests/test_memory_config.py

@@ -42,15 +42,15 @@ def test_serialize_memory_config_with_shard_spec(tmp_path):
 
 def test_serialize_memory_config_with_shard_spec_over_core_range_set(tmp_path):
     memory_config_path = str(tmp_path / "memory_config.bin")
-    shard_spec_32_cores_grid = ttnn._tt_lib.tensor.CoreRangeSet(
+    shard_spec_32_cores_grid = ttnn.CoreRangeSet(
         {
-            ttnn._tt_lib.tensor.CoreRange(
-                ttnn._tt_lib.tensor.CoreCoord(0, 0),
-                ttnn._tt_lib.tensor.CoreCoord(4, 4),
+            ttnn.CoreRange(
+                ttnn.CoreCoord(0, 0),
+                ttnn.CoreCoord(4, 4),
             ),
-            ttnn._tt_lib.tensor.CoreRange(
-                ttnn._tt_lib.tensor.CoreCoord(5, 5),
-                ttnn._tt_lib.tensor.CoreCoord(7, 7),
+            ttnn.CoreRange(
+                ttnn.CoreCoord(5, 5),
+                ttnn.CoreCoord(7, 7),
             ),
         }
     )
@@ -61,7 +61,7 @@ def test_serialize_memory_config_with_shard_spec_over_core_range_set(tmp_path):
     memory_config = ttnn.MemoryConfig(
         ttnn.TensorMemoryLayout.HEIGHT_SHARDED,
         ttnn.BufferType.L1,
-        ttnn._tt_lib.tensor.ShardSpec(
+        ttnn.ShardSpec(
             shard_spec_32_cores_grid,
             [
                 num_attention_heads * sequence_lengtb // 32,

tt_eager/CMakeLists.txt

@@ -3,27 +3,6 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/tensor)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/queue)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/tt_dnn)
 
-set(TT_EAGER_OBJECTS
-    $<TARGET_OBJECTS:tt_dnn>
-    $<TARGET_OBJECTS:tensor>
-    $<TARGET_OBJECTS:queue>
-)
-
-# Defaults to a shared lib, static lib for production python wheel
-add_library(tt_eager ${TT_EAGER_OBJECTS})
-target_link_libraries(tt_eager PUBLIC compiler_flags linker_flags tt_metal ttnn_lib)     # linker_flags = -rdynamic if tracy enabled
-target_include_directories(tt_eager PUBLIC
-    ${UMD_HOME}
-    ${PROJECT_SOURCE_DIR}
-    ${PROJECT_SOURCE_DIR}/tt_metal
-    ${PROJECT_SOURCE_DIR}/tt_eager
-    ${PROJECT_SOURCE_DIR}/ttnn/cpp/
-)
-
-set_target_properties(tt_eager PROPERTIES
-    INSTALL_RPATH "${PROJECT_BINARY_DIR}/lib"
-)
-
 if(WITH_PYTHON_BINDINGS)
     add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/tt_lib)
 endif()

tt_eager/queue/CMakeLists.txt

@@ -1,15 +1,6 @@
 
 set(QUEUE_SRCS
     ${CMAKE_CURRENT_SOURCE_DIR}/queue.cpp
-)
 
-add_library(queue OBJECT ${QUEUE_SRCS})
-target_link_libraries(queue PUBLIC metal_header_directories compiler_flags)
-target_include_directories(queue PUBLIC
-    ${UMD_HOME}
-    ${PROJECT_SOURCE_DIR}
-    ${PROJECT_SOURCE_DIR}/tt_metal
-    ${PROJECT_SOURCE_DIR}/tt_eager
-    ${CMAKE_CURRENT_SOURCE_DIR}
-    ${PROJECT_SOURCE_DIR}/tt_metal/third_party/fmt
+    CACHE INTERNAL "Queue sources to reuse in ttnn build"
 )

tt_eager/tensor/CMakeLists.txt

@@ -5,16 +5,6 @@ set(TENSOR_SRCS
     ${CMAKE_CURRENT_SOURCE_DIR}/types.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/tensor_utils.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/serialization.cpp
-)
-
-add_library(tensor OBJECT ${TENSOR_SRCS})
 
-target_link_libraries(tensor PUBLIC metal_header_directories compiler_flags)
-target_include_directories(tensor PUBLIC
-    ${UMD_HOME}
-    ${PROJECT_SOURCE_DIR}
-    ${PROJECT_SOURCE_DIR}/tt_metal
-    ${PROJECT_SOURCE_DIR}/tt_eager
-    ${CMAKE_CURRENT_SOURCE_DIR}
-    ${PROJECT_SOURCE_DIR}/tt_metal/third_party/fmt
+    CACHE INTERNAL "Tensor sources to reuse in ttnn build"
 )