#4003: Adding pyind11 to ttnn

module.mk

@@ -138,7 +138,9 @@ LIBS_TO_BUILD += \
 	tools \
 	tt_metal \
 	tracy \
-	tt_eager
+	tt_eager \
+	ttnn \
+	ttnn/setup_local_so
 
 # These must be in dependency order (enforces no circular deps)
 include $(UMD_HOME)/device/module.mk

scripts/github/issues.py

@@ -10,51 +10,57 @@
 # See: https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens#creating-a-fine-grained-personal-access-token
 
 
-def get_pending_reviewers(pull_numbers, github_token):
+def get_reviewers_and_pending_reviewers(pull_numbers, github_token):
     owner = "tenstorrent-metal"
     repo = "tt-metal"
     headers = {"Authorization": f"token {github_token}"}
-    pending_reviewers = set()
+    reviewers_info = {}
 
     for pull_number in pull_numbers:
         requested_reviewers_url = f"https://api.github.com/repos/{owner}/{repo}/pulls/{pull_number}/requested_reviewers"
         reviews_url = f"https://api.github.com/repos/{owner}/{repo}/pulls/{pull_number}/reviews"
 
         requested_reviewers_response = requests.get(requested_reviewers_url, headers=headers)
-        requested_reviewers = (
-            {reviewer["login"] for reviewer in requested_reviewers_response.json().get("users", [])}
-            if requested_reviewers_response.status_code == 200
-            else set()
-        )
+        if requested_reviewers_response.status_code == 200:
+            reviewers = {reviewer["login"] for reviewer in requested_reviewers_response.json().get("users", [])}
+        else:
+            print(
+                f"Unable to find Requested Reviews for PR {pull_number}: HTTP {requested_reviewers_response.status_code} -> {requested_reviewers_response.text}"
+            )
+            continue
 
         reviews_response = requests.get(reviews_url, headers=headers)
-        reviews = reviews_response.json() if reviews_response.status_code == 200 else []
-
-        if any(review["state"] == "APPROVED" for review in reviews):
+        if reviews_response.status_code == 200:
+            reviews = reviews_response.json()
+        else:
+            print(
+                f"Unable to find Reviews for PR {pull_number}: HTTP {reviews_response.status_code} -> {requested_reviewers_response.text}"
+            )
             continue
 
-        pending_reviewers.update(requested_reviewers - {review["user"]["login"] for review in reviews})
-
-    return pending_reviewers
+        pr_details_url = f"https://api.github.com/repos/{owner}/{repo}/pulls/{pull_number}"
+        pr_details_response = requests.get(pr_details_url, headers=headers)
+        if pr_details_response.status_code == 200:
+            pr_data = pr_details_response.json()
+            base_branch = pr_data.get("base", {}).get("ref", "")
+            state = pr_data.get("state")
+            mergeable = pr_data.get("mergeable")
+        else:
+            print(f"Unable to find details for PR {pull_number}: HTTP {pr_details_response.status_code}")
+            continue
 
+        pending_reviewers = reviewers - {
+            review["user"]["login"] for review in reviews if review["state"] in ["APPROVED", "CHANGES_REQUESTED"]
+        }
 
-def get_reviewers(pull_numbers, github_token):
-    owner = "tenstorrent-metal"
-    repo = "tt-metal"
-    reviewers = set()
-    headers = {"Authorization": f"token {github_token}"}
-
-    for pull_number in pull_numbers:
-        url = f"https://api.github.com/repos/{owner}/{repo}/pulls/{pull_number}/requested_reviewers"
-        response = requests.get(url, headers=headers)
-        if response.status_code == 200:
-            data = response.json()
-            for reviewer in data["users"]:
-                reviewers.add(reviewer["login"])
-        else:
-            print(f"Failed to get reviewers for PR {pull_number}: {response.status_code}")
+        reviewers_info[pull_number] = {
+            "reviewers": reviewers,
+            "pending_reviewers": pending_reviewers,
+            "state": state,
+            "mergeable": mergeable,
+        }
 
-    return list(reviewers)
+    return reviewers_info
 
 
 def get_my_pending_reviews(username, github_token):
@@ -106,11 +112,12 @@ def main():
         print("You forget to provide the --token")
 
     if issues != None:
-        reviewers = get_reviewers(issues, token)
-        print("Reviewers:", reviewers)
+        reviewers_info = get_reviewers_and_pending_reviewers(issues, token)
 
-        pending_reviewers = get_pending_reviewers(issues, token)
-        print("Pending Reviewers:", pending_reviewers)
+        for pr, info in reviewers_info.items():
+            print(
+                f"PR {pr} - Reviewers: {info['reviewers']}, State: {info['state']}, Pending Reviewers: {info['pending_reviewers']}, Mergeable: {info['mergeable']}"
+            )
 
     elif username != None:
         pending_reviews = get_my_pending_reviews(username, token)

tests/ttnn/sweep_tests/README.md

@@ -10,12 +10,24 @@ python tests/ttnn/sweep_tests/run_all_tests.py
 python tests/ttnn/sweep_tests/print_report.py [--detailed]
 ```
 
+## Using Pytest to run sweeps all the sweeps for one operation file
+```
+pytest <full-path-to-tt-metal>/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_<operation>
+Example for matmul: pytest tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_matmul
+```
+
+## Using Pytest to run a single sweep test by the index
+```
+pytest <full-path-to-tt-metal>/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_<operation>[<operation>.py-<index-of-test-instance>]
+Example for matmul: pytest tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_matmul[matmul.py-0]
+```
+
 ## Debugging sweeps
 ```
 python tests/ttnn/sweep_tests/run_failed_and_crashed_tests.py [--exclude add,linear] [--stepwise]
 ```
 
-## Running a single test
+## Running a single test without pytest
 ```
 python tests/ttnn/sweep_tests/run_single_test.py --test-name add --index 0
 ```

tt_eager/tt_lib/__init__.py

@@ -2,4 +2,4 @@
 
 # SPDX-License-Identifier: Apache-2.0
 
-from ._C import tensor, device, dtx, profiler, program_cache, operations, ttnn
+from ._C import tensor, device, dtx, profiler, program_cache, operations

tt_eager/tt_lib/csrc/tt_lib_bindings.cpp

@@ -16,7 +16,6 @@
 #include "tt_metal/detail/reports/memory_reporter.hpp"
 #include "tt_metal/detail/tt_metal.hpp"
 #include "tt_metal/tools/profiler/op_profiler.hpp"
-#include "ttnn/module.hpp"
 #include "type_caster.hpp"
 
 namespace py = pybind11;
@@ -331,8 +330,6 @@ PYBIND11_MODULE(_C, m) {
     py::module_ m_operations = m.def_submodule("operations", "Submodule for operations");
     tt::operations::py_module(m_operations);
 
-    py::module_ m_ttnn = m.def_submodule("ttnn", "Submodule for ttnn");
-    ttnn::py_module(m_ttnn);
 
 #if defined(TRACY_ENABLE)
     py::function tracy_decorator = py::module::import("tt_eager.tt_lib_profiler_wrapper").attr("callable_decorator");

ttnn/README.md

@@ -55,11 +55,11 @@ followed the instructions for [installing and building the software](https://git
     * The primary reason was because we needed a way to create a consolidated report per operation in the form of a csv file.  The idea was that each operation would get its own python file where all the test combinations are handled by a single run method.  Each permuation of the input combinations would become the header for the resulting csv which is then uploaded and reported on.
 * How do I run sweep tests with pytest?
     * To run all of the sweep tests for a given python operation file:
-        * pytest <full-path-to-tt-metal>/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_<operation>
-        * Example for matmul: pytest /home/ubuntu/git/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_matmul
+        * `pytest <full-path-to-tt-metal>/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_<operation>`
+        * Example for matmul: `pytest /home/ubuntu/git/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_matmul`
     * To run just one sample combination for an operation:
-        * pytest <full-path-to-tt-metal>/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_<operation>[<operation>.py-<index-of-test-instance>]
-        * Example for matmul: pytest /home/ubuntu/git/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_matmul[matmul.py-0]
+        * `pytest <full-path-to-tt-metal>/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_<operation>[<operation>.py-<index-of-test-instance>]`
+        * Example for matmul: `pytest /home/ubuntu/git/tt-metal/tests/ttnn/sweep_tests/test_all_sweep_tests.py::test_matmul[matmul.py-0]`
 * What if my device hangs?
     * Be sure that you have a clean build with the latest code update.  Updating code without rebuilding is often the source of the issue.
     * If you have a clean build, you can reset the board on the command line using the tt-smi command and the device id with: `tt-smi -tr 0` where 0 represents the device id.

tt_eager/tt_lib/csrc/ttnn/module.hpp → ttnn/cpp/module.cpp

@@ -2,23 +2,26 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-#pragma once
+// #pragma once
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 
 #include "operations/module.hpp"
 #include "types.hpp"
 
 namespace py = pybind11;
 
-namespace ttnn {
 
-void py_module(py::module& m_ttnn) {
+PYBIND11_MODULE(_ttnn, m_ttnn) {
+
+    // m_ttnn.attr("__name__") = "_ttnn";
+    m_ttnn.doc() = "Python bindings for TTNN";
+
+
     auto m_types = m_ttnn.def_submodule("types", "ttnn Types");
-    types::py_module(m_types);
+    ttnn::types::py_module(m_types);
 
     auto m_operations = m_ttnn.def_submodule("operations", "ttnn Operations");
-    operations::py_module(m_operations);
+    ttnn::operations::py_module(m_operations);
 }
-
-}  // namespace ttnn

tt_eager/tt_lib/csrc/ttnn/operations/binary.hpp → ttnn/cpp/operations/binary.hpp

@@ -5,8 +5,11 @@
 #pragma once
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 
 #include "tt_dnn/op_library/eltwise_binary/eltwise_binary_op.hpp"
+#include "tt_eager/tt_dnn/op_library/bcast/bcast_op.hpp"
+#include "tt_eager/tensor/tensor_utils.hpp"
 
 namespace py = pybind11;
 

tt_eager/tt_lib/csrc/ttnn/operations/module.hpp → ttnn/cpp/operations/module.hpp

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 
 #include "binary.hpp"
 

tt_eager/tt_lib/csrc/ttnn/types.hpp → ttnn/cpp/types.hpp

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 
 #include "tensor/tensor.hpp"
 

ttnn/module.mk

@@ -5,3 +5,49 @@
 ttnn/dev_install: python_env/dev
 	echo "Installing editable dev version of ttnn package..."
 	bash -c "source $(PYTHON_ENV)/bin/activate && pip install -e ttnn"
+
+TTNN_LIB = $(LIBDIR)/libttnn.so
+TTNN_PYBIND_LIB = $(LIBDIR)/_ttnn.so
+
+TTNN_DEFINES =
+
+TTNN_INCLUDES = -Itt_eager/tt_dnn $(TT_EAGER_INCLUDES) $(TT_LIB_INCLUDES)
+TTNN_PYBIND11_INCLUDES = $(TTNN_INCLUDES) $(shell python3-config --includes) -Itt_metal/third_party/pybind11/include
+
+TTNN_LDFLAGS = -ltt_dnn -ldtx -ltensor -ltt_metal -lyaml-cpp $(LDFLAGS)
+TTNN_LDFLAGS += -L$(LIBDIR) -ldevice -lqueue -ltracy -ltt_lib_csrc
+TTNN_CFLAGS = $(CFLAGS) -Werror -Wno-int-to-pointer-cast -fno-var-tracking
+
+TTNN_SRCS =
+
+TTNN_PYBIND_SRCS = \
+    ttnn/cpp/module.cpp
+
+TTNN_OBJS = $(addprefix $(OBJDIR)/, $(TTNN_SRCS:.cpp=.o))
+TTNN_PYBIND_OBJS = $(addprefix $(OBJDIR)/, $(TTNN_PYBIND_SRCS:.cpp=.o))
+
+TTNN_DEPS = $(addprefix $(OBJDIR)/, $(TTNN_SRCS:.cpp=.d))
+-include $(TTNN_DEPS)
+
+-include $(TTNN_PYBIND_SRCS:.cpp=.d)
+
+
+ttnn: $(TTNN_LIB) $(TTNN_PYBIND_LIB)
+
+$(TTNN_LIB): $(TTNN_OBJS) $(TT_DNN_LIB) $(TENSOR_LIB) $(DTX_LIB) $(TT_METAL_LIB) $(TT_LIBS_TO_BUILD)
+	@mkdir -p $(LIBDIR)
+	$(CXX) $(TTNN_CFLAGS) $(CXXFLAGS) $(SHARED_LIB_FLAGS) -o $@ $(TTNN_OBJS) $(TTNN_LDFLAGS)
+
+$(TTNN_PYBIND_LIB): $(TTNN_PYBIND_OBJS)
+	@mkdir -p $(LIBDIR)
+	$(CXX) $(TTNN_CFLAGS) $(CXXFLAGS) $(SHARED_LIB_FLAGS) -o $@ $(TTNN_PYBIND_OBJS) $(TTNN_LDFLAGS)
+
+.PHONY: ttnn/setup_local_so
+ttnn/setup_local_so: $(TTNN_PYBIND_LIB)
+	rm -f ttnn/ttnn/_ttnn.so
+	cp $^ ttnn/ttnn/_ttnn.so
+
+
+$(OBJDIR)/ttnn/cpp/%.o: ttnn/cpp/%.cpp
+	@mkdir -p $(@D)
+	$(CXX) $(TTNN_CFLAGS) $(CXXFLAGS) $(STATIC_LIB_FLAGS) $(TTNN_INCLUDES) -c -o $@ $<

ttnn/ttnn/__init__.py

@@ -4,6 +4,10 @@
 
 import pathlib
 
+import tt_lib
+
+import ttnn._ttnn
+
 MODEL_CACHE_PATH = pathlib.Path().home() / ".cache" / "tenstorrent"
 
 from ttnn.types import (
@@ -114,6 +118,7 @@
     multiply,
 )
 
+
 from ttnn.operations.relational import (
     gtz,
     ltz,

ttnn/ttnn/decorators.py

@@ -5,6 +5,8 @@
 from contextlib import contextmanager
 from functools import wraps
 import inspect
+import inspect
+import importlib
 
 from loguru import logger
 
@@ -71,6 +73,15 @@ def convert_torch_output_to_be_like_ttnn_output(torch_output, output):
     return torch_output
 
 
+def is_import_time():
+    for frame_info in inspect.stack():
+        if frame_info.frame.f_globals.get("__name__") == "__main__":
+            return False
+        if "importlib" in frame_info.filename:
+            return True
+    return False
+
+
 def document_input_tensors(name, function, validate_input_tensors):
     signature = inspect.signature(validate_input_tensors)
     arguments = {arg_name: None for arg_name in signature.parameters}
@@ -89,13 +100,14 @@ def wrapper(*_, **kwargs):
 
         return wrapper
 
-    original_validate_input_tensor = ttnn.validate_input_tensor
-    ttnn.validate_input_tensor = document_validate_input_tensor(ttnn.validate_input_tensor)
-    try:
-        validate_input_tensors(**arguments)
-    except:
-        pass
-    ttnn.validate_input_tensor = original_validate_input_tensor
+    if not is_import_time():
+        original_validate_input_tensor = ttnn.validate_input_tensor
+        ttnn.validate_input_tensor = document_validate_input_tensor(ttnn.validate_input_tensor)
+        try:
+            validate_input_tensors(**arguments)
+        except:
+            pass
+        ttnn.validate_input_tensor = original_validate_input_tensor
 
     if not tensor_schemas:
         # Handle the case for the functions without input tensors

ttnn/ttnn/operations/binary.py

@@ -56,7 +56,10 @@ def binary_function(
         return output_tensor
 
     binary_function.__name__ = f"ttnn.{name}"
-    binary_function.__doc__ = doc + binary_function.__doc__
+    if binary_function.__doc__ is not None:
+        binary_function.__doc__ = doc + binary_function.__doc__
+    else:
+        binary_function = doc
 
     setattr(THIS_MODULE, name, binary_function)
 
@@ -170,7 +173,7 @@ def add(
     """
     input_tensor_a = input_tensor_a.value
     input_tensor_b = input_tensor_b.value if isinstance(input_tensor_b, ttnn.Tensor) else input_tensor_b
-    output = ttl.ttnn.operations.binary.add(input_tensor_a, input_tensor_b, memory_config=memory_config)
+    output = ttnn._ttnn.operations.binary.add(input_tensor_a, input_tensor_b, memory_config=memory_config)
     return ttnn.Tensor(output)