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

* refactor cython wrapper

* optimize

* fix installations

CMakeLists.txt

@@ -158,12 +158,13 @@ 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)
 
@@ -171,9 +172,12 @@ 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)

tilelang/env.py

@@ -20,7 +20,7 @@ TL_TEMPLATE_NOT_FOUND_MESSAGE = ("TileLang is not installed or found in the expe
 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
@@ -37,6 +37,10 @@ if not os.path.exists(THIRD_PARTY_ROOT):
 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.

tilelang/jit/adapter/cython/adapter.py

@@ -22,10 +22,8 @@ from tilelang.utils.tensor import map_torch_type
 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