[fix] Bugfixes for missing ttnn package, tt-mlir CAPI library resolution, and broken pytests (#208)

- Problem: TTNN package instability and configure-time availability
check - ttnn package in `tt-mlir/third_party/tt-metal` may disappear or
change, and `ttlang_check_ttnn_available()` checked at configure time
when ttnn only available after build.
Fix: Created `cmake/modules/CopyTTNNPythonPackage.cmake` to copy ttnn to
`build/python_packages/ttnn/` during build. Removed
`ttlang_check_ttnn_available()` from `cmake/modules/TTLangUtils.cmake`.
Tests use runtime checks (`REQUIRES: ttnn` directive,
`pytest.importorskip()`).

- Problem: CAPI library resolution issues - `python/CMakeLists.txt`
could find wrong `libTTMLIRPythonCAPI.so` from different tt-mlir install
at configure time, and Python extension could resolve wrong library at
runtime, causing MLIRContext registry mismatches.
Fix: Made configure-time search paths mutually exclusive (only search
`TTMLIR_BUILD_DIR` when defined, else search `TTMLIR_PATH`). Added
`BUILD_RPATH` and `INSTALL_RPATH` properties to constrain runtime
library search to configured tt-mlir directories.

- Problem: Missing pytest infrastructure and broken test files - no
pytest check target or configuration, lit tests incorrectly collected by
pytest, `test_elementwise_ops.py` imported from non-existent
`examples/utils` (deleted in #67), `test_block_allocation.py` used wrong
function name `new_split_work_to_cores` and unsafe ttnn import.
Fix: Added `check-ttlang-pytest` target to `test/CMakeLists.txt`.
Created `test/python/conftest.py` with feature detection, markers,
fixtures, and `collect_ignore` list. Recovered the deleted
`examples/utils.py` that was required in pytests -- now
`test/python/utils.py`. Fixed corresponding imports and function names
in test files, added `# REQUIRES: ttnn` directive, changed to
`pytest.importorskip("ttnn")`.
  
- Problem: `test/python/simple_add_multitile.py` fails on qb because of
ordering differences.
  Fix: Use `CHECK-DAG:` appropriately.

Author: brnorris03

.github/workflows/call-test-hardware.yml

@@ -129,11 +129,3 @@ jobs:
           name: hardware-test-reports-${{ inputs.runs-on }}
           path: build/python-lit-report.xml
           retention-days: 7
-
-      - name: Upload system descriptor artifacts
-        uses: actions/upload-artifact@v4
-        if: always()
-        with:
-          name: hardware-system-desc-${{ inputs.runs-on }}
-          path: ttrt-artifacts
-          retention-days: 7

CMakeLists.txt

@@ -73,6 +73,7 @@ if(MLIR_ENABLE_BINDINGS_PYTHON AND TTLANG_ENABLE_BINDINGS_PYTHON)
   message(STATUS "Enabling Python API")
   set(TTLANG_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
   add_subdirectory(python)
+  include(CopyTTNNPythonPackage)
 endif()
 
 add_subdirectory(test)

cmake/modules/CopyTTNNPythonPackage.cmake

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Copy tt-metal's ttnn Python package to build tree.
+# TTNN is required for compiling and running tt-lang kernels.
+
+# Find ttnn source directory from TT_METAL_HOME or tt-mlir's third_party.
+# _TTMLIR_CMAKE_HOME_DIRECTORY is loaded from the tt-mlir build cache when
+# using TTMLIR_BUILD_DIR (see ttlang_setup_ttmlir_build_tree in TTLangUtils.cmake).
+if(DEFINED TT_METAL_HOME AND EXISTS "${TT_METAL_HOME}/ttnn/ttnn")
+  set(_TTNN_SOURCE_DIR "${TT_METAL_HOME}/ttnn/ttnn")
+elseif(DEFINED _TTMLIR_CMAKE_HOME_DIRECTORY AND EXISTS "${_TTMLIR_CMAKE_HOME_DIRECTORY}/third_party/tt-metal/src/tt-metal/ttnn/ttnn")
+  set(_TTNN_SOURCE_DIR "${_TTMLIR_CMAKE_HOME_DIRECTORY}/third_party/tt-metal/src/tt-metal/ttnn/ttnn")
+endif()
+
+if(DEFINED _TTNN_SOURCE_DIR)
+  message(STATUS "Found ttnn Python package at: ${_TTNN_SOURCE_DIR}")
+
+  add_custom_command(
+    OUTPUT ${CMAKE_BINARY_DIR}/python_packages/ttnn/__init__.py
+    COMMAND ${CMAKE_COMMAND} -E copy_directory
+      ${_TTNN_SOURCE_DIR}
+      ${CMAKE_BINARY_DIR}/python_packages/ttnn
+    COMMENT "Copying tt-metal ttnn Python package to build tree"
+  )
+
+  add_custom_target(copy-ttnn-python-package ALL
+    DEPENDS ${CMAKE_BINARY_DIR}/python_packages/ttnn/__init__.py
+  )
+else()
+  message(STATUS "ttnn Python package not found - Python lit tests may be skipped")
+endif()

cmake/modules/TTLangUtils.cmake

@@ -200,25 +200,6 @@ macro(ttlang_setup_ttmlir_build_tree BUILD_DIR)
   endif()
 endmacro()
 
-# ttlang_check_ttnn_available(OUTPUT_VAR)
-# Checks if the TTNN Python package is available at configure time.
-# Sets the variable named by OUTPUT_VAR to TRUE if available, FALSE otherwise.
-function(ttlang_check_ttnn_available OUTPUT_VAR)
-  execute_process(
-    COMMAND ${Python3_EXECUTABLE} -c "import ttnn"
-    RESULT_VARIABLE _ttnn_import_result
-    OUTPUT_QUIET
-    ERROR_QUIET
-  )
-  if(_ttnn_import_result EQUAL 0)
-    set(${OUTPUT_VAR} TRUE PARENT_SCOPE)
-    message(STATUS "TTNN Python package available")
-  else()
-    set(${OUTPUT_VAR} FALSE PARENT_SCOPE)
-    message(STATUS "TTNN Python package not available")
-  endif()
-endfunction()
-
 # ttlang_check_device_available(OUTPUT_VAR)
 # Checks if a Tenstorrent device is available at configure time by looking for
 # /dev/tenstorrent* files. This is faster than calling ttnn.GetNumAvailableDevices().

python/CMakeLists.txt

@@ -142,10 +142,11 @@ set(_TTMLIR_PYTHON_CAPI_HINTS)
 if(DEFINED TTMLIR_BUILD_DIR)
   list(APPEND _TTMLIR_PYTHON_CAPI_HINTS
     "${TTMLIR_BUILD_DIR}/python_packages/ttmlir/_mlir_libs")
-endif()
-if(DEFINED TTMLIR_PATH)
-  list(APPEND _TTMLIR_PYTHON_CAPI_HINTS
-    "${TTMLIR_PATH}/python_packages/ttmlir/_mlir_libs")
+else()
+  if(DEFINED TTMLIR_PATH)
+    list(APPEND _TTMLIR_PYTHON_CAPI_HINTS
+      "${TTMLIR_PATH}/python_packages/ttmlir/_mlir_libs")
+  endif()
 endif()
 list(APPEND _TTMLIR_PYTHON_CAPI_HINTS
   "${TTMLIR_TOOLCHAIN_DIR}/python_packages/ttmlir/_mlir_libs"
@@ -190,6 +191,22 @@ target_link_libraries(TTLangPythonModules.extension._ttlang.dso PRIVATE
   ${TTMLIR_PYTHON_CAPI_LIB}
 )
 
+# Ensure we do not accidentally resolve libTTMLIRPythonCAPI.so from a different
+# tt-mlir install tree at runtime. Mixed CAPI builds in one Python process can
+# lead to MLIRContext registry mismatches.
+#
+# Constrain RUNPATH to the configured tt-mlir build tree.
+if(DEFINED TTMLIR_BUILD_DIR)
+  set(_TTMLIR_PYTHON_LIBDIR "${TTMLIR_BUILD_DIR}/python_packages/ttmlir/_mlir_libs")
+  set(_TTMLIR_LIBDIR "${TTMLIR_BUILD_DIR}/lib")
+  if(EXISTS "${_TTMLIR_PYTHON_LIBDIR}" AND EXISTS "${_TTMLIR_LIBDIR}")
+    set_target_properties(TTLangPythonModules.extension._ttlang.dso PROPERTIES
+      BUILD_RPATH "\$ORIGIN:${_TTMLIR_LIBDIR}:${_TTMLIR_PYTHON_LIBDIR}"
+      INSTALL_RPATH "\$ORIGIN:${_TTMLIR_LIBDIR}:${_TTMLIR_PYTHON_LIBDIR}"
+    )
+  endif()
+endif()
+
 # ###############################################################################
 # Install
 # ###############################################################################

python/ttlang/ttl_api.py

@@ -455,7 +455,7 @@ def _compile_ttnn_kernel(
     )
 
     if verbose:
-        print(f"\nCompiled kernel ready (will execute on {len(kernel_paths)} kernels)")
+        print(f"\nCompiled kernel ready (compiled {len(kernel_paths)} threads)")
         print("=" * 60)
 
     return compiled_kernel

test/CMakeLists.txt

@@ -7,8 +7,7 @@ configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
 )
 
-# Check for TTNN Python package and Tenstorrent device at configure time
-ttlang_check_ttnn_available(TTLANG_HAS_TTNN)
+# Check for Tenstorrent device at configure time
 ttlang_check_device_available(TTLANG_HAS_DEVICE)
 
 if(TTLANG_HAS_DEVICE)
@@ -53,21 +52,14 @@ add_lit_testsuite(check-ttlang-mlir "Running ttlang MLIR tests"
 )
 set_target_properties(check-ttlang-mlir PROPERTIES FOLDER "Tests")
 
-# Python lit tests require TTNN. Run sequentially to avoid TLB resource contention
+# Python lit tests. Run sequentially to avoid TLB resource contention
 # when multiple tests open Tenstorrent devices.
-if(TTLANG_HAS_TTNN)
-  add_lit_testsuite(check-ttlang-python-lit "Running ttlang Python lit tests"
-    ${CMAKE_CURRENT_BINARY_DIR}/python
-    -j1 ${TTLANG_LIT_COMMON_ARGS} --xunit-xml-output python-lit-report.xml
-    DEPENDS ${TTLANG_PYTHON_TEST_DEPENDS}
-  )
-  set_target_properties(check-ttlang-python-lit PROPERTIES FOLDER "Tests")
-else()
-  add_custom_target(check-ttlang-python-lit
-    COMMAND ${CMAKE_COMMAND} -E echo "TTNN not available - skipping Python lit tests"
-    COMMENT "Skipping ttlang Python lit tests (TTNN not available)"
-  )
-endif()
+add_lit_testsuite(check-ttlang-python-lit "Running ttlang Python lit tests"
+  ${CMAKE_CURRENT_BINARY_DIR}/python
+  -j1 ${TTLANG_LIT_COMMON_ARGS} --xunit-xml-output python-lit-report.xml
+  DEPENDS ${TTLANG_PYTHON_TEST_DEPENDS}
+)
+set_target_properties(check-ttlang-python-lit PROPERTIES FOLDER "Tests")
 
 # Python bindings tests only (no TT device required).
 add_lit_testsuite(check-ttlang-python-bindings "Running ttlang Python binding tests"
@@ -83,8 +75,20 @@ add_custom_target(check-ttlang
   COMMENT "Running ttlang lit tests that never require a TT device"
 )
 
+# Pytest tests (test_*.py files).
+add_custom_target(check-ttlang-pytest
+  COMMAND ${Python3_EXECUTABLE} -m pytest
+    ${CMAKE_CURRENT_SOURCE_DIR}/python
+    -v --tb=short
+    --junitxml=${CMAKE_CURRENT_BINARY_DIR}/pytest-report.xml
+  DEPENDS ${TTLANG_PYTHON_TEST_DEPENDS}
+  COMMENT "Running ttlang pytest tests"
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/python
+)
+set_target_properties(check-ttlang-pytest PROPERTIES FOLDER "Tests")
+
 # All lit tests (including Python tests that require a TT device).
 add_custom_target(check-ttlang-all
-  DEPENDS check-ttlang-mlir check-ttlang-python-bindings check-ttlang-python-lit
-  COMMENT "Running all ttlang lit tests (including Python tests that require a TT device)"
+  DEPENDS check-ttlang-mlir check-ttlang-python-bindings check-ttlang-python-lit check-ttlang-pytest
+  COMMENT "Running all ttlang tests (lit and pytest)"
 )

test/lit.cfg.py

@@ -52,6 +52,8 @@
     "Inputs",
     "lit.cfg.py",
     "sim",
+    "conftest.py",
+    "utils.py",
 ]
 
 # Exclude pytest-style tests (test_*.py) from lit collection.
@@ -137,10 +139,19 @@
     "TT_METAL_RUNTIME_ROOT",
     "TT_MLIR_HOME",
     "SYSTEM_DESC_PATH",
+    "TTLANG_COMPILE_ONLY",
 ]:
     if env_var in os.environ:
         config.environment[env_var] = os.environ[env_var]
 
 # Add system platform feature for UNSUPPORTED directives
 if platform.system() == "Darwin":
     config.available_features.add("system-darwin")
+
+# Add TTNN feature if available
+try:
+    import ttnn
+
+    config.available_features.add("ttnn")
+except ImportError:
+    pass

test/python/conftest.py

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pytest configuration and fixtures for tt-lang Python tests."""
+
+import glob
+import importlib.util
+import os
+
+import pytest
+
+# Lit tests that should not be collected by pytest (they have # RUN: directives)
+collect_ignore = [
+    "conftest.py",
+    "test_ttnn_interop_add.py",
+    "test_dram_interleaved_add.py",
+    "utils.py",
+]
+
+# =============================================================================
+# Feature detection
+# =============================================================================
+
+_ttnn_available = False
+if importlib.util.find_spec("ttnn") is not None:
+    _ttnn_available = True
+
+_hardware_available = bool(glob.glob("/dev/tenstorrent*"))
+
+# Set compile-only mode if no hardware
+if not _hardware_available:
+    os.environ["TTLANG_COMPILE_ONLY"] = "1"
+
+
+# =============================================================================
+# Pytest markers
+# =============================================================================
+
+
+def pytest_configure(config):
+    """Register custom markers."""
+    config.addinivalue_line(
+        "markers", "requires_ttnn: skip test if ttnn is not available"
+    )
+    config.addinivalue_line(
+        "markers", "requires_device: skip test if no TT device is available"
+    )
+
+
+def pytest_collection_modifyitems(config, items):
+    """Skip tests based on available features."""
+    skip_ttnn = pytest.mark.skip(reason="TTNN not available")
+    skip_device = pytest.mark.skip(reason="No Tenstorrent device available")
+
+    for item in items:
+        if "requires_ttnn" in item.keywords and not _ttnn_available:
+            item.add_marker(skip_ttnn)
+        if "requires_device" in item.keywords and not _hardware_available:
+            item.add_marker(skip_device)
+
+
+# =============================================================================
+# Fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def ttnn_device():
+    """Fixture that provides a TTNN device, skipping if unavailable."""
+    if not _ttnn_available:
+        pytest.skip("TTNN not available")
+    if not _hardware_available:
+        pytest.skip("No Tenstorrent device available")
+
+    import ttnn
+
+    device = ttnn.open_device(device_id=0)
+    yield device
+    ttnn.close_device(device)

test/python/invalid/invalid_3d_grid.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+# REQUIRES: ttnn
 # RUN: not %python %s 2>&1 | FileCheck %s
 
 """
@@ -14,21 +15,16 @@
 
 os.environ["TTLANG_COMPILE_ONLY"] = "1"
 
-from ttlang import ttl, make_circular_buffer_like
-from ttlang.ttl_api import Program
+import ttnn
+from ttlang import make_circular_buffer_like, ttl
 from ttlang.operators import copy
-
-try:
-    import ttnn
-except ImportError:
-    print("TTNN not available - exiting")
-    exit(0)
+from ttlang.ttl_api import Program
 
 
 # CHECK: ValueError: TTNN interop only supports single-core grid (1, 1), got (1, 1, 1)
-# CHECK-NEXT:   --> {{.*}}invalid_3d_grid.py:34:1
+# CHECK-NEXT:   --> {{.*}}invalid_3d_grid.py:[[LINE:[0-9]+]]:1
 # CHECK-NEXT:    |
-# CHECK-NEXT: 34 | @ttl.kernel(grid=(1, 1, 1))
+# CHECK-NEXT: [[LINE]] | @ttl.kernel(grid=(1, 1, 1))
 # CHECK-NEXT:    | ^
 # CHECK-NEXT:    |
 @ttl.kernel(grid=(1, 1, 1))