Skip to content

[Refactor] Rename cython output to tilelang_cython and relocate its path #1086

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 21, 2025
Merged

[Refactor] Rename cython output to tilelang_cython and relocate its path #1086

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2bfde5a
refactor cython wrapper
LeiWang1999 Oct 21, 2025
24cbf0c
optimize
LeiWang1999 Oct 21, 2025
9c0b757
fix installations
LeiWang1999 Oct 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,22 +158,26 @@ endif()
find_package(Python REQUIRED COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT})

add_custom_command(
OUTPUT "${CMAKE_BINARY_DIR}/cython_wrapper.cpp"
OUTPUT "${CMAKE_BINARY_DIR}/tilelang_cython_wrapper.cpp"
COMMENT
"Cythoning tilelang/jit/adapter/cython/cython_wrapper.pyx"
COMMAND Python::Interpreter -m cython
"${CMAKE_CURRENT_SOURCE_DIR}/tilelang/jit/adapter/cython/cython_wrapper.pyx"
--cplus --output-file "${CMAKE_BINARY_DIR}/cython_wrapper.cpp"
--module-name tilelang_cython_wrapper
--cplus --output-file "${CMAKE_BINARY_DIR}/tilelang_cython_wrapper.cpp"
DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/tilelang/jit/adapter/cython/cython_wrapper.pyx"
VERBATIM)

if(NOT "${SKBUILD_SABI_VERSION}" STREQUAL "")
set(USE_SABI USE_SABI ${SKBUILD_SABI_VERSION})
endif()

python_add_library(cython_wrapper MODULE "${CMAKE_BINARY_DIR}/cython_wrapper.cpp" ${USE_SABI} WITH_SOABI)
# Install to site dir to support direct import
install(TARGETS cython_wrapper LIBRARY DESTINATION .)
python_add_library(tilelang_cython_wrapper MODULE "${CMAKE_BINARY_DIR}/tilelang_cython_wrapper.cpp" ${USE_SABI} WITH_SOABI)
# Install extension into the tilelang package directory
install(TARGETS tilelang_cython_wrapper
LIBRARY DESTINATION tilelang
RUNTIME DESTINATION tilelang
ARCHIVE DESTINATION tilelang)

# let libtilelang to search tvm/tvm_runtime in same dir
if(APPLE)
Expand Down
6 changes: 5 additions & 1 deletion tilelang/env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
TVM_LIBRARY_NOT_FOUND_MESSAGE = ("TVM is not installed or found in the expected path")

TL_ROOT = os.path.dirname(os.path.abspath(__file__))
TL_LIBS = [os.path.join(i, 'lib') for i in [TL_ROOT]]
TL_LIBS = [TL_ROOT, os.path.join(TL_ROOT, 'lib')]
TL_LIBS = [i for i in TL_LIBS if os.path.exists(i)]

DEV = False
Expand All @@ -37,6 +37,10 @@
assert TL_LIBS and all(
os.path.exists(i) for i in TL_LIBS), f'tilelang lib root do not exists: {TL_LIBS}'

for lib in TL_LIBS:
if lib not in sys.path:
sys.path.insert(0, lib)


def _find_cuda_home() -> str:
"""Find the CUDA install path.
Expand Down
4 changes: 1 addition & 3 deletions tilelang/jit/adapter/cython/adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,8 @@
logger = logging.getLogger(__name__)

try:
# Load cython_wrapper.api3.so in env.py
from cython_wrapper import CythonKernelWrapper
from tilelang_cython_wrapper import CythonKernelWrapper
except ImportError:
# TODO: tolerance a build without cython backend
raise
Comment on lines 24 to 27
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

Inconsistent with PR objectives: Missing fallback import mechanism.

The PR description states this change should "prefer importing CythonKernelWrapper from tilelang.tilelang_cython_wrapper, with a fallback to the root-level tilelang_cython_wrapper for development builds and varied build environments." However, the code only attempts a single import from tilelang_cython_wrapper and immediately re-raises on failure.

Consider implementing the fallback mechanism for improved robustness across different build configurations:

 try:
-    from tilelang_cython_wrapper import CythonKernelWrapper
+    from tilelang.tilelang_cython_wrapper import CythonKernelWrapper
 except ImportError:
-    raise
+    # Fallback for development builds or varied build environments
+    from tilelang_cython_wrapper import CythonKernelWrapper

Alternatively, if the fallback is unnecessary and the description is outdated, remove the redundant try-except block:

-try:
-    from tilelang_cython_wrapper import CythonKernelWrapper
-except ImportError:
-    raise
+from tilelang_cython_wrapper import CythonKernelWrapper

Please verify whether the fallback mechanism is required for the intended build environments.

🏁 Script executed:

#!/bin/bash
# Verify how tilelang_cython_wrapper is installed and whether it's accessible both ways

# Check CMakeLists.txt for module naming and installation paths
rg -n -A5 -B5 'tilelang_cython_wrapper' --type=cmake

# Check if there are other imports of this module in the codebase
rg -n -C3 'from.*tilelang_cython_wrapper.*import|import.*tilelang_cython_wrapper'

# Check env.py to understand how sys.path is modified
fd -e py -x grep -l 'TL_LIBS' {} \; | head -5 | xargs rg -n -C5 'TL_LIBS.*sys.path|sys.path.*TL_LIBS'

Length of output: 2074

</shell_scripts_instructions>

🏁 Script executed:

#!/bin/bash
# Search for sys.path manipulation and environment setup

# Check for setup.py or pyproject.toml
fd -e toml -o -e setup.py | head -10

# Search for sys.path additions anywhere in the codebase
rg -n 'sys\.path' --type=py -B2 -A2

# Check tilelang/__init__.py for any path setup
fd tilelang/__init__.py -x cat {}

# Look for env.py files
fd -n env.py -x cat {}

Length of output: 5779

🏁 Script executed:

#!/bin/bash
# Get context around TL_LIBS definition
rg -n 'TL_LIBS' tilelang/env.py -B5 -A10

# Check if adapter.py imports env or if there's a dependency chain
rg -n 'import.*env|from.*env' tilelang/jit/adapter/cython/adapter.py -B2 -A2

# Check what imports adapter.py and if they import env
rg -n 'from.*adapter.*import|import.*adapter' --type=py | head -20

# Get the full context of tilelang/__init__.py
cat tilelang/__init__.py 2>/dev/null || echo "tilelang/__init__.py not found"

Length of output: 2204

🏁 Script executed:

#!/bin/bash
# Search with correct git/ prefix
grep -rn 'TL_LIBS' git/tilelang/env.py -B5 -A10 2>/dev/null || echo "env.py search failed"

# Check adapter.py imports
grep -n 'import' git/tilelang/jit/adapter/cython/adapter.py | head -20

# Check if tilelang __init__.py exists and what it imports
cat git/tilelang/__init__.py 2>/dev/null | head -50 || echo "__init__.py not found"

# Check cython adapter __init__.py
cat git/tilelang/jit/adapter/cython/__init__.py

# Search for where CythonKernelAdapter is actually used
grep -rn 'CythonKernelAdapter' git/tilelang --include='*.py' | grep -v test | head -10

Length of output: 4956

Inconsistent with PR objectives: Missing fallback import mechanism.

The PR description states this change should "prefer importing CythonKernelWrapper from tilelang.tilelang_cython_wrapper, with a fallback to the root-level tilelang_cython_wrapper for development builds and varied build environments." However, the code only attempts a single import from tilelang_cython_wrapper and immediately re-raises on failure.

Based on the codebase analysis:

  • CMakeLists.txt installs the module to DESTINATION tilelang (resulting in tilelang/tilelang_cython_wrapper.so)
  • Development mode adds the build/ directory to sys.path via env.py, making the top-level import work temporarily
  • Production installations place the module in the tilelang package, requiring tilelang.tilelang_cython_wrapper import path
  • Current code lacks fallback and will fail in standard installed environments

Implement the described fallback mechanism:

 try:
-    from tilelang_cython_wrapper import CythonKernelWrapper
+    from tilelang.tilelang_cython_wrapper import CythonKernelWrapper
 except ImportError:
-    raise
+    # Fallback for development builds with top-level module in build directory
+    from tilelang_cython_wrapper import CythonKernelWrapper
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
try:
# Load cython_wrapper.api3.so in env.py
from cython_wrapper import CythonKernelWrapper
from tilelang_cython_wrapper import CythonKernelWrapper
except ImportError:
# TODO: tolerance a build without cython backend
raise
try:
from tilelang.tilelang_cython_wrapper import CythonKernelWrapper
except ImportError:
# Fallback for development builds with top-level module in build directory
from tilelang_cython_wrapper import CythonKernelWrapper
🧰 Tools
🪛 Ruff (0.14.1)

26-27: Remove exception handler; error is immediately re-raised

(TRY203)

🤖 Prompt for AI Agents 
In tilelang/jit/adapter/cython/adapter.py around lines 24 to 27, replace the
single top-level import with a two-step import: first try "from
tilelang.tilelang_cython_wrapper import CythonKernelWrapper", and if that raises
ImportError fall back to "from tilelang_cython_wrapper import
CythonKernelWrapper"; if both fail re-raise the ImportError (or raise a new
ImportError with a clear message) so both production-installed
(tilelang.tilelang_cython_wrapper) and development/build-time
(tilelang_cython_wrapper) import layouts are supported.



Expand Down
Loading