Repository navigation
[Feature] Add full Windows support and fix related cross-platform issues - #2093
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis pull request adds comprehensive Windows support to TileLang's build system and runtime environment. Changes include CMake helpers for Python interpreter selection and Windows import library generation, a cross-platform dynamic library loading abstraction replacing POSIX-only dlopen calls, CUDA toolkit discovery enhancements supporting pip-installed CUDA, MSVC compiler integration for Windows builds, and Python-level runtime DLL path management. The changes span build configuration, CUDA stub infrastructure, packaging, and runtime initialization to enable TileLang to build and run on Windows with CUDA support. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/jit/adapter/wrapper.py (1)
60-86:⚠️ Potential issue | 🟡 MinorCPU backend
call()is not exported on Windows — inconsistent with CUDA/HIP paths.
PREDEF_HOST_FUNC(used by CUDA and HIP via inheritance) now wrapscallwithTL_EXPORT, butTLCPUSourceWrapper.CALL_PREFIX(lines 751–758) still emits plainextern "C" int32_t call(...). When the CPU backend lib is built on Windows viatilelang.contrib.msvc.create_shared, the MSVC/EXPORTauto-discovery only matches__tvm_ffi_*symbols (_TVM_FFI_EXPORT_REinmsvc.py), socallwill be missing from the DLL exports and the ctypes adapter lookup (self.lib.call) will fail with an "unresolved symbol" error.🐛 Proposed fix for CPU wrapper
CALL_PREFIX = textwrap.dedent(""" `#ifdef` __cplusplus - extern "C" + extern "C" TL_EXPORT `#endif` int32_t call({}) {{ return {}; }} """)Since
CALL_PREFIXis appended afterINIT_FUNC = PREDEF_INIT_FUNC.format(""), theTL_EXPORTmacro is already defined at the pointcallis emitted.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/jit/adapter/wrapper.py` around lines 60 - 86, The CPU backend's exported call symbol is not marked with TL_EXPORT, causing missing export on Windows; change the CPU wrapper emission so the call signature uses the same export macro as CUDA/HIP by replacing the plain 'extern "C" int32_t call(...)' emitted by TLCPUSourceWrapper.CALL_PREFIX with the TL_EXPORT-enabled prefix (as defined by PREDEF_HOST_FUNC / PREDEF_INIT_FUNC where TL_EXPORT is already defined), ensuring the CPU shared lib exposes call and the ctypes adapter can resolve self.lib.call.
🧹 Nitpick comments (20)
src/runtime/logging.cc (2)
108-124: Drop the unreachablereturn settings;afterLOG(FATAL).
LogFatalImplis declared[[noreturn]](line 62), soLOG(FATAL)never returns. The threereturn settings;statements on lines 110, 115, and 123 are dead code and may also trigger "unreachable code" warnings under MSVC/W4(which this PR now exercises). The very firstLOG(FATAL)at lines 100-102 already omits the trailing return for consistency — apply the same treatment here.♻️ Proposed fix
if (!std::getline(spec_stream, level, ',')) { LOG(FATAL) << "TVM_LOG_DEBUG ill-formed at position " << tell_pos(level) << ": expecting \"=<level>\" after \"" << name << "\""; - return settings; } if (level.empty()) { LOG(FATAL) << "TVM_LOG_DEBUG ill-formed at position " << tell_pos(level) << ": empty level after \"" << name << "\""; - return settings; } char *end_of_level = nullptr; int level_val = static_cast<int>(strtol(level.c_str(), &end_of_level, 10)); if (end_of_level != level.c_str() + level.size()) { LOG(FATAL) << "TVM_LOG_DEBUG ill-formed at position " << tell_pos(level) << ": invalid level: \"" << level << "\""; - return settings; }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/runtime/logging.cc` around lines 108 - 124, Remove the unreachable "return settings;" statements that follow the LOG(FATAL) calls in the level-parsing block (the checks for ill-formed token, empty level, and invalid numeric level) because LOG(FATAL) routes to LogFatalImpl which is [[noreturn]]; keep behavior consistent with the earlier LOG(FATAL) that already omits a trailing return and leave only the LOG(FATAL) calls in place.
3-4: Include<sstream>explicitly.
ParseSpecusesstd::istringstream(line 85) andVLogContext::str()usesstd::stringstream(line 150), but<sstream>is not directly included. The file currently relies on transitive includes through<tvm/runtime/logging.h>or implementation details, which is not guaranteed across different standard library implementations or versions. Add the explicit include to ensure this translation unit is self-sufficient.♻️ Proposed fix
`#include` <algorithm> `#include` <cstdlib> `#include` <ctime> `#include` <iomanip> `#include` <iostream> +#include <sstream> `#include` <string>🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/runtime/logging.cc` around lines 3 - 4, The translation unit uses std::istringstream and std::stringstream in ParseSpec and VLogContext::str() but doesn't include <sstream>, relying on transitive includes; add an explicit `#include` <sstream> at the top of the file so ParseSpec and VLogContext::str() have the required header and the compilation is self-contained.src/transform/layout_reducer.cc (2)
174-175: Nit: stale doxygen reference.The docstring still mentions
analyzer_->const_int_bound, but the implementation now routes throughGetThreadConstIntBound(...). Worth updating so the contract matches the code.📝 Proposed doc tweak
- * - Reads and relies on `thread_var_`, `analyzer_->const_int_bound`, and - * `var_to_buffer_`. + * - Reads and relies on `thread_var_`, `GetThreadConstIntBound(thread_var_, + * *analyzer_)`, and `var_to_buffer_`.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/transform/layout_reducer.cc` around lines 174 - 175, Update the function docstring to remove the stale reference to analyzer_->const_int_bound and instead document that it reads and relies on thread_var_, var_to_buffer_, and the GetThreadConstIntBound(...) helper; specifically replace mentions of analyzer_->const_int_bound with GetThreadConstIntBound(...) and briefly state the expected contract/return (e.g., that it provides per-thread constant integer bounds) so the doc matches the implementation.
391-403: Optional: short-circuit traversal oncehas_reducer_infois set.
ReducerInfoDetectorkeeps visiting sibling blocks after the flag flips totrue. Since the detector only needs existence, early-returning fromVisitStmtavoids unnecessary traversal of large function bodies.♻️ Proposed micro-optimization
class ReducerInfoDetector : public StmtVisitor { public: bool has_reducer_info = false; private: + void VisitStmt(const Stmt &stmt) final { + if (has_reducer_info) return; + StmtVisitor::VisitStmt(stmt); + } + void VisitStmt_(const BlockNode *op) final { if (op->annotations.count(attr::kReducerInfo)) { has_reducer_info = true; return; } StmtVisitor::VisitStmt_(op); } };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/transform/layout_reducer.cc` around lines 391 - 403, ReducerInfoDetector currently continues traversing even after has_reducer_info is true; modify its VisitStmt_ override(s) (at least VisitStmt_(const BlockNode *op)) to early-return immediately when has_reducer_info is true before examining annotations or recursing via StmtVisitor::VisitStmt_, and apply the same guard in any other VisitStmt_/visitor overrides in ReducerInfoDetector to prevent further recursion once the flag is set.cmake/find_pip_cuda.py (2)
186-210: Optional: emit diagnostics to stderr on failure.Both
sys.exit(1)paths (nocu_dir, nonvcc) exit silently. The caller incmake/FindPipCUDAToolkit.cmakeonly sees a non-zero exit code, which loses the root cause (no candidatecu*dir found vs.cudir found but missingnvcc). A one-lineprint(..., file=sys.stderr)would make auto-detect failures debuggable without changing the JSON contract.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/find_pip_cuda.py` around lines 186 - 210, The two silent failure exits in main() should emit a short diagnostic to stderr before exiting: when _find_cu_dir() returns None print a one-line message (e.g. "no CUDA directory candidate found") to sys.stderr, and when _find_nvcc(cu_dir) returns None print a one-line message (e.g. "nvcc not found under <cu_dir>") to sys.stderr; update the main() function to call print(..., file=sys.stderr) with these messages immediately prior to the existing sys.exit(1) calls so the JSON output contract is unchanged but failures are debuggable via stderr.
87-120: Optional: simplify the generator-in-closure pattern.
addis defined as a generator (hasyield) that emits at most one item, consumed viafor candidate in add(...). It works, but yielding insidecontextlib.suppress(OSError)obscures intent and makes control flow harder to follow. A non-generator helper is easier to read.♻️ Proposed refactor
def _candidate_nvidia_dirs(): seen: set[pathlib.Path] = set() - def add(path: pathlib.Path): - with contextlib.suppress(OSError): - resolved = path.resolve() - if resolved not in seen: - seen.add(resolved) - yield resolved + def add(path: pathlib.Path) -> pathlib.Path | None: + try: + resolved = path.resolve() + except OSError: + return None + if resolved in seen: + return None + seen.add(resolved) + return resolved try: import nvidia except ImportError: pass else: - for candidate in add(pathlib.Path(nvidia.__path__[0])): + candidate = add(pathlib.Path(nvidia.__path__[0])) + if candidate: yield candidate for entry in sys.path: if not entry: continue - for candidate in add(pathlib.Path(entry) / "nvidia"): + candidate = add(pathlib.Path(entry) / "nvidia") + if candidate: yield candidate # ... and similarly for site/user-site blocks🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/find_pip_cuda.py` around lines 87 - 120, The helper add inside _candidate_nvidia_dirs is implemented as a single-item generator which obscures control flow; change add(path: pathlib.Path) from a generator to a regular helper that resolves the path inside contextlib.suppress(OSError) and returns the resolved pathlib.Path or None, then replace usages `for candidate in add(...)` with a simple check like `candidate = add(...)` followed by `if candidate: yield candidate`; keep all callers (the nvidia.__path__ branch, sys.path loop, site.getsitepackages() and getusersitepackages() branches) using the returned value and yielding when non-None to preserve behavior.tilelang/jit/adapter/libgen.py (1)
152-158: Optional: specifyencoding/errorsfor robust capture on Windows.
text=Truewithoutencoding/errorsuseslocale.getpreferredencoding()witherrors='strict'. On a Windows console configured for a non-UTF-8 code page (e.g., cp936/cp1252), any non-ASCII byte in NVCC output (such as a localized error string or a path with non-ASCII chars) will raiseUnicodeDecodeErrorinsidesubprocess.run, which the outerexcept Exceptionthen reports as "Compile kernel failed because of UnicodeDecodeError..." — hiding the real compiler error.♻️ Proposed hardening
ret = subprocess.run( command, timeout=timeout, env=compiler_env, capture_output=True, text=True, + encoding="utf-8", + errors="replace", )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/jit/adapter/libgen.py` around lines 152 - 158, The subprocess.run invocation that assigns ret (the call using command, timeout, env=compiler_env, capture_output=True, text=True) can raise UnicodeDecodeError on Windows; update that call to explicitly pass encoding and errors (e.g., encoding='utf-8' and errors='replace' or errors='surrogateescape') so decoder falls back instead of throwing, and keep using ret.stdout/ret.stderr for logging; modify the subprocess.run call in libgen.py where ret is set to include these parameters.cmake/pypi-z3/FindZ3.cmake (2)
36-40: Deadif(WIN32)/elsebranches — both call the sameadd_library.Both sides of the conditional create
z3::libz3identically, so the wholeif(WIN32)/else/endifis a no-op wrapper. Simplify to a single line.♻️ Proposed cleanup
-if(WIN32) - add_library(z3::libz3 SHARED IMPORTED GLOBAL) -else() - add_library(z3::libz3 SHARED IMPORTED GLOBAL) -endif() +add_library(z3::libz3 SHARED IMPORTED GLOBAL)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/pypi-z3/FindZ3.cmake` around lines 36 - 40, The if(WIN32)/else block is redundant because both branches call the same add_library(z3::libz3 SHARED IMPORTED GLOBAL); remove the conditional and replace it with a single call to add_library with the identifier z3::libz3 (keep the SHARED IMPORTED GLOBAL properties intact) so the unnecessary if/else wrapper around add_library is eliminated.
19-30:Z3_IMPORT_LIBRARYinitial path is misleading;_z3_import_diris never actually populated.The initial
set(Z3_IMPORT_LIBRARY "${_z3_import_dir}/libz3.lib")on line 22 points to${CMAKE_BINARY_DIR}/z3-import/libz3.lib, buttilelang_generate_windows_import_library(defined inCMakeLists.txt) ignores that location and writes the generated.libto${CMAKE_BINARY_DIR}/windows-import-libs/libz3/<dllname>/libz3.lib, then overwritesZ3_IMPORT_LIBRARYvia itsOUTPUT_VAR. Net effects:
- The
file(MAKE_DIRECTORY "${_z3_import_dir}")call creates an always-empty dir that's never used.- The
if(NOT EXISTS "${Z3_IMPORT_LIBRARY}")guard is effectively always true on reconfigures (the file-not-found guard insidetilelang_generate_windows_import_libraryitself handles re-run caching).Either drop the pre-seeding entirely and let the helper own the path, or ask the helper to write to the expected location. Minor cleanup, not a correctness bug.
♻️ Proposed simplification
set(Z3_DLL "${Z3_LIBRARY}") if(Z3_DLL) - set(_z3_import_dir "${CMAKE_BINARY_DIR}/z3-import") - file(MAKE_DIRECTORY "${_z3_import_dir}") - set(Z3_IMPORT_LIBRARY "${_z3_import_dir}/libz3.lib") - if(NOT EXISTS "${Z3_IMPORT_LIBRARY}") - if(COMMAND tilelang_generate_windows_import_library) - tilelang_generate_windows_import_library("${Z3_DLL}" Z3_IMPORT_LIBRARY "libz3") - else() - message(FATAL_ERROR "tilelang_generate_windows_import_library is unavailable for Windows Z3 import-lib generation") - endif() - endif() + if(COMMAND tilelang_generate_windows_import_library) + tilelang_generate_windows_import_library("${Z3_DLL}" Z3_IMPORT_LIBRARY "libz3") + else() + message(FATAL_ERROR "tilelang_generate_windows_import_library is unavailable for Windows Z3 import-lib generation") + endif() endif()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/pypi-z3/FindZ3.cmake` around lines 19 - 30, The Z3 import-lib pre-seeding is misleading because _z3_import_dir and Z3_IMPORT_LIBRARY are set but never used by tilelang_generate_windows_import_library; remove the unused pre-seed and directory creation or defer to the helper: delete the file(MAKE_DIRECTORY "${_z3_import_dir}") and the set(Z3_IMPORT_LIBRARY "${_z3_import_dir}/libz3.lib") / if(NOT EXISTS ...) block and instead call tilelang_generate_windows_import_library when Z3_DLL is set so that the helper sets Z3_IMPORT_LIBRARY via its OUTPUT_VAR (or alternatively update tilelang_generate_windows_import_library to accept an explicit output path and write to _z3_import_dir), referencing the symbols Z3_DLL, _z3_import_dir, Z3_IMPORT_LIBRARY, and tilelang_generate_windows_import_library to locate code to change.src/target/stubs/nvrtc.cc (1)
32-45: Consider addingnvrtc64_140_0.dll/ future CUDA 14 to the Windows candidate list.The list tops out at
nvrtc64_130_0.dll(CUDA 13.x). Once CUDA 14 ships, this stub silently fails to locate the installed library on machines that only have CUDA 14 until the list is updated. Low priority, just a maintenance reminder — you could prepend"nvrtc64_140_0.dll"now, or add a comment pointing maintainers here.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/stubs/nvrtc.cc` around lines 32 - 45, The Windows nvRTC candidate array kLibNvrtcPaths in nvrtc.cc stops at CUDA 13 names so systems with CUDA 14 may fail to load; update kLibNvrtcPaths to include "nvrtc64_140_0.dll" (preferably prepended to the Windows list) or add a clear TODO comment above kLibNvrtcPaths noting to add future nvrtc64_XXX_0.dll entries for new CUDA major versions so the loader finds newer libraries without silent failures.scripts/windows/uv-sync.cmd (1)
8-8: Use%~1for the terminator check to handle quoted arguments robustly.If a passthrough argument is itself a quoted string (e.g.,
--config-setting "cmake.args=-DFOO=bar"),%1keeps the surrounding quotes, and the"%1"==""form compares""<quoted>"" == ""— which is still false, so this specific check works. However, using%~1(which strips surrounding quotes) is the conventional and more resilient pattern and matches your comparison style on line 9.🔧 Proposed fix
-if "%1"=="" goto args_done +if "%~1"=="" goto args_done🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/windows/uv-sync.cmd` at line 8, Replace the bare %1 check with the quote-stripped form so the terminator comparison handles quoted passthrough arguments correctly: change the conditional that compares "%1" to "" to use %~1 (the stripped-quote form) in the if test that jumps to the args_done label (the existing if ... goto args_done check), ensuring the check uses %~1 for robustness with quoted arguments.tilelang/__init__.py (3)
125-126: Silentexcept Exception: passswallows failures fromfind_libtvm_ffi.Static analysis (Ruff S110/BLE001) flagged this. Since
tvm_ffi.libinfo.find_libtvm_ffi()failing means we won't prepend that runtime path on Windows — and the subsequentimport tvmmay then fail with an opaqueDLL load failed, it's worth at least alogger.debugso the cause is discoverable.🔧 Proposed logging
try: from tvm_ffi import libinfo as tvm_ffi_libinfo runtime_paths.append(os.path.dirname(tvm_ffi_libinfo.find_libtvm_ffi())) - except Exception: - pass + except Exception: # noqa: BLE001 + logger.debug("Failed to resolve tvm_ffi runtime path", exc_info=True)Note:
loggeris defined at line 52 which runs before this helper is invoked, so it's safe to reference.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/__init__.py` around lines 125 - 126, The except block swallowing all exceptions around tvm_ffi.libinfo.find_libtvm_ffi() should log the exception so failures are discoverable; update the except Exception: pass in tilelang/__init__.py to catch Exception as e and call logger.debug (or logger.exception) with a clear message referencing find_libtvm_ffi and the exception, so if the runtime path isn't prepended and import tvm fails the root cause is visible (logger is available at module scope).
177-191: Duplicates the CUDA path list inlibinfo.get_dll_directories().The tuple of
cuda_homesubpaths (cuda_home,bin,bin/x86_64,lib/x64,nvvm/bin) is the same one built bytilelang/libinfo.py::get_dll_directories()(which you already call below in_load_tile_lang_lib). Two copies mean any future change to the discovery order has to be updated in both places.Since
get_dll_directories()can be called beforeimport tvm(it only depends ontilelang.envandtilelang.libinfo, which are already imported), you could reuse it here:♻️ Proposed DRY refactor
if sys.platform.startswith("win32") and sys.version_info >= (3, 8): _prepend_windows_runtime_paths() + from . import libinfo as _libinfo # noqa: E402 _dll_handles = [] - cuda_home = getattr(env, "CUDA_HOME", "") - if cuda_home: - for path in ( - cuda_home, - os.path.join(cuda_home, "bin"), - os.path.join(cuda_home, "bin", "x86_64"), - os.path.join(cuda_home, "lib", "x64"), - os.path.join(cuda_home, "nvvm", "bin"), - ): - if os.path.isdir(path): - _dll_handles.append(os.add_dll_directory(path)) + for path in _libinfo.get_dll_directories(): + _dll_handles.append(os.add_dll_directory(path))Note:
get_dll_directories()also addsTL_LIBS, which is desirable here (TVM DLLs live undertilelang/3rdparty/tvm). If you want to keep this pre-tvm-import step scoped to CUDA only, factor the CUDA-only list into a separate helper inlibinfo.py.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/__init__.py` around lines 177 - 191, The CUDA DLL directory list is duplicated here; instead of rebuilding the same tuple, call tilelang.libinfo.get_dll_directories() (or add a small helper in tilelang.libinfo like get_cuda_dll_directories() if you want CUDA-only behavior) from the __init__.py block that runs on Windows (the code using _prepend_windows_runtime_paths, _dll_handles, cuda_home); use the returned iterable to iterate and call os.add_dll_directory for each existing path, and remove the hard-coded tuple so discovery logic is centralized and not duplicated (also keep the existing use of _load_tile_lang_lib which already relies on libinfo.get_dll_directories()).
54-57: Mutatingos.RTLD_LAZYat import time affects the entire process.This assignment to
os.RTLD_LAZY = 0on Windows is visible to every other module in the interpreter. It's safe becauseosdoesn't define it on Windows, but any third-party code that checkshasattr(os, 'RTLD_LAZY')to branch on POSIX-vs-Windows will now see it as truthy. An alternative would be to readgetattr(os, "RTLD_LAZY", 0)inside_lazy_load_lib(which already guards withhasattrat line 153) rather than mutating theosmodule.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/__init__.py` around lines 54 - 57, The file currently sets os.RTLD_LAZY = 0 on Windows which mutates the global os module; instead remove that import-time assignment and update the loader to use a local fallback flag (e.g., read rtld_flag = getattr(os, "RTLD_LAZY", 0)) when building the CDLL flags inside the _lazy_load_lib function so you no longer inject RTLD_LAZY into os; locate references to os.RTLD_LAZY and replace them with a local getattr fallback in _lazy_load_lib (or compute a module-level constant like _RTLD_LAZY = getattr(os, "RTLD_LAZY", 0) without assigning to os) to preserve current behavior without mutating the os module.tilelang/contrib/cc.py (1)
54-60: Duplicated Windows MSVC discovery acrossget_ccandget_cplus_compiler.The Windows branch in both functions is byte-identical. Consider extracting a small helper (e.g.,
_find_msvc_cl()) to keep them in sync and simplify future changes.♻️ Proposed helper extraction
+def _find_msvc_cl(): + from tilelang.contrib.msvc import get_msvc_subprocess_env, get_env_path + compiler_env = get_msvc_subprocess_env() + return shutil.which("cl.exe", path=get_env_path(compiler_env or {})) + + def get_cc(): ... if _is_windows_like(): - from tilelang.contrib.msvc import get_msvc_subprocess_env, get_env_path - - compiler_env = get_msvc_subprocess_env() - cl = shutil.which("cl.exe", path=get_env_path(compiler_env or {})) - return cl + return _find_msvc_cl()Also applies to: 85-91
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/cc.py` around lines 54 - 60, The Windows-specific MSVC discovery logic duplicated in get_cc and get_cplus_compiler should be extracted into a helper (e.g., _find_msvc_cl) that performs the same steps: import get_msvc_subprocess_env and get_env_path from tilelang.contrib.msvc, call get_msvc_subprocess_env(), compute the env path via get_env_path(compiler_env or {}), locate "cl.exe" with shutil.which(..., path=...), and return the discovered path; then replace the duplicated blocks in get_cc and get_cplus_compiler with a call to _find_msvc_cl() so behavior remains identical and both functions stay in sync.tilelang/libinfo.py (1)
49-56: Minor: simplify the redundantfor…elseand clarify the error message.The outer-loop
elsefires only when noreturnoccurred, which is exactly the same as falling through. The message also callsTL_LIBSthe "List of candidates", but semantically those are search directories (candidates are the fulllib_root/lib_namecombinations).♻️ Suggested simplification
- for lib_root in TL_LIBS: - for lib_name in lib_names: - lib_dll_path = os.path.join(lib_root, lib_name) - if os.path.exists(lib_dll_path) and os.path.isfile(lib_dll_path): - return lib_dll_path - else: - message = f"Cannot find libraries: {', '.join(lib_names)}\n" + "List of candidates:\n" + "\n".join(TL_LIBS) - raise RuntimeError(message) + for lib_root in TL_LIBS: + for lib_name in lib_names: + lib_dll_path = os.path.join(lib_root, lib_name) + if os.path.isfile(lib_dll_path): + return lib_dll_path + raise RuntimeError( + f"Cannot find libraries: {', '.join(lib_names)}\n" + f"Searched in: {', '.join(TL_LIBS)}" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/libinfo.py` around lines 49 - 56, The nested for…else is redundant: iterate TL_LIBS and lib_names, return lib_dll_path on first existing file, and if none found raise a RuntimeError after both loops; remove the for...else construct and place the error raise after the loops. Update the error message to list the full candidate paths (each os.path.join(lib_root, lib_name)) instead of calling TL_LIBS "List of candidates", referencing TL_LIBS, lib_names and lib_dll_path so the message shows the exact searched file paths.tilelang/contrib/msvc.py (2)
174-179: Minor: mark the MD5 hash as non-security.Ruff S324 will flag this on every CI run. The hash is used purely to disambiguate patched filenames, so pass
usedforsecurity=False(Python 3.9+) to document intent and silence the warning.- prefix = hashlib.md5(os.path.dirname(os.path.abspath(path)).encode()).hexdigest()[:8] + prefix = hashlib.md5( + os.path.dirname(os.path.abspath(path)).encode(), + usedforsecurity=False, + ).hexdigest()[:8]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` around lines 174 - 179, The MD5 call that generates prefix (variable prefix via hashlib.md5(...).hexdigest()) should explicitly mark the hash as non-security by passing usedforsecurity=False to hashlib.md5 to silence Ruff S324 and document intent; update the hashlib.md5 invocation in this function/block that computes prefix so it becomes hashlib.md5(<data>, usedforsecurity=False) (Python 3.9+).
275-287:subprocess.Popenwith no timeout can hang the build.If
cl.exe(or the linker spawned by it) deadlocks — which does happen on Windows with certain antivirus/filesystem interactions — the build will hang indefinitely with no way to recover short of killing the parent. Considerproc.communicate(timeout=…)with a generous timeout (e.g., 30 min) or usesubprocess.run(..., timeout=...)and surface a clearer error.♻️ Optional hardening
- proc = subprocess.Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - cwd=cwd, - env=compiler_env, - ) - out, _ = proc.communicate() + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=cwd, + env=compiler_env, + ) + try: + out, _ = proc.communicate(timeout=30 * 60) + except subprocess.TimeoutExpired: + proc.kill() + out, _ = proc.communicate() + raise RuntimeError( + "MSVC compilation timed out after 30 minutes.\n" + "Command line: " + " ".join(cmd) + ) if proc.returncode != 0:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` around lines 275 - 287, The current subprocess invocation in msvc.py uses subprocess.Popen and proc.communicate() with no timeout, which can hang; modify the call in the block that creates proc (the variables cmd, cwd, compiler_env and the proc variable) to use a timeout (e.g., communicate(timeout=1800) or replace with subprocess.run(..., timeout=1800, capture_output=True)) and catch subprocess.TimeoutExpired to kill/terminate the child process, collect any partial output (use proc.stdout/communicate after terminate if needed), and raise a RuntimeError that includes the timeout indication plus any captured output and the command line (same message structure as the existing RuntimeError using py_str(out) and " ".join(cmd)). Ensure the timeout handling is applied around the proc.communicate()/subprocess.run call and that processes are properly terminated to avoid zombie processes.cmake/FindPipCUDAToolkit.cmake (1)
168-258: Minor:_tilelang_sdk_binis prepended toENV{PATH}twice within one configure.Line 170 adds
_tilelang_sdk_binto$ENV{PATH}once, then lines 243–258 re-prepend system path + re-prepend_tilelang_sdk_bina second time. End result:_tilelang_sdk_binappears twice inPATH. Harmless but noisy and makesfind_programdiagnostics harder to read.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/FindPipCUDAToolkit.cmake` around lines 168 - 258, The _tilelang_sdk_bin path is being prepended to ENV{PATH} twice via the early set(ENV{PATH} "${_tilelang_sdk_bin};$ENV{PATH}") and again near the end; remove the first prepend (the one immediately after the message(STATUS) call) or guard both with a conditional that checks whether "$ENV{PATH}" already MATCHES "(^|;)\Q${_tilelang_sdk_bin}\E($|;)" so _tilelang_sdk_bin is only added once; key symbols: _tilelang_sdk_bin, ENV{PATH}, _tilelang_system_path_prefix/_tilelang_system_path_candidates.CMakeLists.txt (1)
28-78: Import-library helper looks solid.Good structure overall — per-DLL subdirectory keying (Line 39) makes the on-disk
.libcache self-invalidate when the source DLL's basename changes, the machine-type mapping covers the realistic MSVC generator platforms, and the tool-missing / non-zero exit paths fail loudly with actionable messages. Theif(NOT EXISTS "${_tilelang_import_lib}")short-circuit (Line 42) is a nice configure-time perf win.Two very small suggestions, both optional:
- Line 42: if the
.defexists but the.libdoesn't (e.g., previous run was killed between the twoexecute_processcalls), you skip regenerating the.defbut still runlib.exe. That's actually fine, but consider caching via the.libtimestamp vs. the source DLL's timestamp if you want to rebuild when the DLL is updated in place.- Line 63-67:
CMAKE_GENERATOR_PLATFORMis only set for multi-config Visual Studio generators; Ninja users on ARM64 would get the X64 default. If ARM64 Windows is a supported target, consider also consultingCMAKE_SYSTEM_PROCESSOR/CMAKE_HOST_SYSTEM_PROCESSORas a fallback.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@CMakeLists.txt` around lines 28 - 78, The tilelang_generate_windows_import_library helper should (1) regenerate the .def when the source DLL is newer than the cached .def (avoid a stale .def when a previous run was interrupted): compare timestamps of "${DLL_PATH}" vs "${_tilelang_def_file}" and run the generate step if the DLL is newer or the .def is missing before invoking lib.exe (refer to _tilelang_def_file and _tilelang_import_lib), and (2) make the machine selection more robust by falling back to CMAKE_SYSTEM_PROCESSOR / CMAKE_HOST_SYSTEM_PROCESSOR when CMAKE_GENERATOR_PLATFORM is unset so ARM64 Ninja builds get "ARM64" instead of defaulting to "X64" (update the block that sets _tilelang_machine and checks CMAKE_GENERATOR_PLATFORM to consult those fallback variables).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cmake/find_pip_cuda.py`:
- Around line 177-184: The _library_dir function returns cu_dir/"lib"/"x64" on
Windows without checking existence, causing missing paths to flow into
CMAKE_LIBRARY_PATH; update _library_dir to check that cu_dir/"lib"/"x64" exists
when os.name == "nt" (similar to the lib64 is_dir() check on Linux) and if it
doesn't exist either try alternative paths or raise/return a clear error/None so
the caller (_tilelang_run_find_pip_cuda / FindPipCUDAToolkit logic) can surface
a descriptive message rather than propagating a non-existent path; reference the
_library_dir function and the cu_dir/"lib"/"x64" and cu_dir/"lib64" path checks
when making the change.
In `@cmake/FindPipCUDAToolkit.cmake`:
- Around line 276-286: The linker /LIBPATH flags are being prepended into the
cached CMAKE_{EXE,SHARED,MODULE}_LINKER_FLAGS on every configure via
${_tilelang_libpath_flags}, causing unbounded growth; modify the logic (around
variables _tilelang_libpath_flags and the CMAKE_*_LINKER_FLAGS cache sets) to be
idempotent—either add a guard variable like _TILELANG_LINK_FLAGS_APPLIED
(mirroring _TILELANG_INCLUDE_FLAGS_APPLIED) or check whether the cached flags
already contain the /LIBPATH entries before setting; only set the cached
CMAKE_EXE_LINKER_FLAGS, CMAKE_SHARED_LINKER_FLAGS and CMAKE_MODULE_LINKER_FLAGS
once (or append only missing entries) so repeated cmake re-configures do not
duplicate /LIBPATH fragments.
In `@CMakeLists.txt`:
- Around line 383-394: Update the NVRTC discovery block to use the documented
CUDAToolkit_BIN_DIR instead of CUDAToolkit_ROOT: change the file(GLOB ...)
patterns to use "${CUDAToolkit_BIN_DIR}/bin/..." (refer to the if(WIN32 AND NOT
TILELANG_USE_CUDA_STUBS) block and the _tilelang_nvrtc_candidates variable),
keep the existing list(SORT ...), list(GET ... _tilelang_nvrtc_library) and
tilelang_generate_windows_import_library(...) logic, and if no
_tilelang_nvrtc_candidates are found emit a status message (e.g., message(STATUS
"... no NVRTC DLL found, continuing without setting CUDA_NVRTC_LIBRARY")) so
callers know fallback discovery will run.
- Around line 570-601: On Windows the tilelang objects are archived into
tilelang_archive and linked into tvm.dll, but because tvm is marked
EXCLUDE_FROM_ALL it may not get built during a plain cmake --build ., leaving
tilelang_cython_wrapper to fail at runtime; add a build dependency so tvm.dll is
always built for local dev builds by adding
add_dependencies(tilelang_cython_wrapper tvm) (or alternatively ensure a
tilelang SHARED target is created on Windows and depends on tvm) so the tvm
target is built before tilelang_cython_wrapper is built.
In `@pyproject.toml`:
- Line 108: The global setting cmake.args = ["-G", "Ninja"] forces Ninja
everywhere and overrides CMAKE_GENERATOR; remove this unconditional entry from
pyproject.toml and instead only set the Ninja generator for Windows
builds—either by configuring CMAKE_GENERATOR in your CI (cibuildwheel
environment-pass) or by applying a platform-conditional override so that
cmake.args (the "-G", "Ninja" pair) is only used when sys_platform == "win32"
(or when building with MSVC); keep other platforms relying on CMAKE_GENERATOR or
scikit-build-core's selection.
In `@src/target/codegen_cuda.cc`:
- Around line 3591-3595: The tl::__tan() branch currently uses CUDAFastMath
(variable CUDAFastMath math_func) which emits lower-precision __tanf for
float32; switch this to use CUDAFastMathTan instead so the tangent call
preserves the higher-precision implementation already defined (i.e., replace
CUDAFastMath usage with CUDAFastMathTan in the tl::__tan() handling and call its
function-name generator the same way you did previously to produce the correct
tan wrapper).
In `@src/target/stubs/dynlib.h`:
- Around line 104-107: The header-level definition of _GNU_SOURCE in dynlib.h is
too late for headers included earlier (breaking RTLD_DEFAULT/RTLD_NEXT usage in
dynlib_find_loaded); remove the `#ifndef/_GNU_SOURCE/`#endif block from dynlib.h
and instead add a project-wide compile definition via CMake for the affected
stub targets (use target_compile_definitions(<stub-target> PRIVATE
_GNU_SOURCE=1) for the targets that build the stubs so _GNU_SOURCE is defined
before any system headers are processed), then rebuild to ensure
dynlib_find_loaded() sees RTLD_DEFAULT and RTLD_NEXT.
In `@src/transform/layout_inference.cc`:
- Around line 756-761: The code currently assumes attr::kLayoutMap is Map<Var,
Layout> and ICHECKs that, which will abort when the annotation is a different
map type (e.g., Map<Buffer, Layout> written by this pass); change the logic in
the block that reads op->annotations.Get(attr::kLayoutMap) so it attempts to
as<Map<Var, Layout>>() and if map_opt.has_value() is false simply skip/ignore
the annotation (do not ICHECK/abort) — i.e., only proceed to use the result when
map_opt has a value, otherwise return/continue from the function so
already-materialized layout maps are safely ignored; keep references to
attr::kLayoutMap, op->annotations.Get, and the Map<Var, Layout> cast to locate
the change.
In `@tilelang/contrib/cc.py`:
- Around line 54-60: get_cc() currently performs MSVC discovery at import and
isn't cached, causing slow shell-out on Windows and repeated probes; add
functools.cache to get_cc() (matching get_cplus_compiler()) and change the eager
default usage in create_shared.get_target_triple from os.environ.get("CXX",
get_cc()) to os.environ.get("CXX") or get_cc() so get_cc() is only invoked when
CXX is unset/empty; keep the MSVC helper imports (get_msvc_subprocess_env,
get_env_path) and shutil.which usage inside get_cc() so subprocess discovery is
deferred and subsequent calls use the cache.
In `@tilelang/contrib/msvc.py`:
- Around line 113-145: In get_msvc_subprocess_env(), when VS setup fails (no
vsdevcmd found, _import_vsdevcmd_environment returns None, or the resulting
compiler_env still lacks cl.exe), set _MSVC_ENV_ERROR as you're already doing
but change the cached result to None (set _MSVC_ENV_CACHE = None) and return
None instead of returning base_env; keep the early fast-path that returns
base_env when cl.exe is already present in the unmodified environment. Update
the failure branches that reference _find_vsdevcmd,
_import_vsdevcmd_environment, and the final shutil.which(check with
get_env_path(compiler_env)) to follow this behavior so callers can distinguish
“no MSVC setup” via a None return while still exposing errors via
get_msvc_environment_error().
---
Outside diff comments:
In `@tilelang/jit/adapter/wrapper.py`:
- Around line 60-86: The CPU backend's exported call symbol is not marked with
TL_EXPORT, causing missing export on Windows; change the CPU wrapper emission so
the call signature uses the same export macro as CUDA/HIP by replacing the plain
'extern "C" int32_t call(...)' emitted by TLCPUSourceWrapper.CALL_PREFIX with
the TL_EXPORT-enabled prefix (as defined by PREDEF_HOST_FUNC / PREDEF_INIT_FUNC
where TL_EXPORT is already defined), ensuring the CPU shared lib exposes call
and the ctypes adapter can resolve self.lib.call.
---
Nitpick comments:
In `@cmake/find_pip_cuda.py`:
- Around line 186-210: The two silent failure exits in main() should emit a
short diagnostic to stderr before exiting: when _find_cu_dir() returns None
print a one-line message (e.g. "no CUDA directory candidate found") to
sys.stderr, and when _find_nvcc(cu_dir) returns None print a one-line message
(e.g. "nvcc not found under <cu_dir>") to sys.stderr; update the main() function
to call print(..., file=sys.stderr) with these messages immediately prior to the
existing sys.exit(1) calls so the JSON output contract is unchanged but failures
are debuggable via stderr.
- Around line 87-120: The helper add inside _candidate_nvidia_dirs is
implemented as a single-item generator which obscures control flow; change
add(path: pathlib.Path) from a generator to a regular helper that resolves the
path inside contextlib.suppress(OSError) and returns the resolved pathlib.Path
or None, then replace usages `for candidate in add(...)` with a simple check
like `candidate = add(...)` followed by `if candidate: yield candidate`; keep
all callers (the nvidia.__path__ branch, sys.path loop, site.getsitepackages()
and getusersitepackages() branches) using the returned value and yielding when
non-None to preserve behavior.
In `@cmake/FindPipCUDAToolkit.cmake`:
- Around line 168-258: The _tilelang_sdk_bin path is being prepended to
ENV{PATH} twice via the early set(ENV{PATH} "${_tilelang_sdk_bin};$ENV{PATH}")
and again near the end; remove the first prepend (the one immediately after the
message(STATUS) call) or guard both with a conditional that checks whether
"$ENV{PATH}" already MATCHES "(^|;)\Q${_tilelang_sdk_bin}\E($|;)" so
_tilelang_sdk_bin is only added once; key symbols: _tilelang_sdk_bin, ENV{PATH},
_tilelang_system_path_prefix/_tilelang_system_path_candidates.
In `@cmake/pypi-z3/FindZ3.cmake`:
- Around line 36-40: The if(WIN32)/else block is redundant because both branches
call the same add_library(z3::libz3 SHARED IMPORTED GLOBAL); remove the
conditional and replace it with a single call to add_library with the identifier
z3::libz3 (keep the SHARED IMPORTED GLOBAL properties intact) so the unnecessary
if/else wrapper around add_library is eliminated.
- Around line 19-30: The Z3 import-lib pre-seeding is misleading because
_z3_import_dir and Z3_IMPORT_LIBRARY are set but never used by
tilelang_generate_windows_import_library; remove the unused pre-seed and
directory creation or defer to the helper: delete the file(MAKE_DIRECTORY
"${_z3_import_dir}") and the set(Z3_IMPORT_LIBRARY
"${_z3_import_dir}/libz3.lib") / if(NOT EXISTS ...) block and instead call
tilelang_generate_windows_import_library when Z3_DLL is set so that the helper
sets Z3_IMPORT_LIBRARY via its OUTPUT_VAR (or alternatively update
tilelang_generate_windows_import_library to accept an explicit output path and
write to _z3_import_dir), referencing the symbols Z3_DLL, _z3_import_dir,
Z3_IMPORT_LIBRARY, and tilelang_generate_windows_import_library to locate code
to change.
In `@CMakeLists.txt`:
- Around line 28-78: The tilelang_generate_windows_import_library helper should
(1) regenerate the .def when the source DLL is newer than the cached .def (avoid
a stale .def when a previous run was interrupted): compare timestamps of
"${DLL_PATH}" vs "${_tilelang_def_file}" and run the generate step if the DLL is
newer or the .def is missing before invoking lib.exe (refer to
_tilelang_def_file and _tilelang_import_lib), and (2) make the machine selection
more robust by falling back to CMAKE_SYSTEM_PROCESSOR /
CMAKE_HOST_SYSTEM_PROCESSOR when CMAKE_GENERATOR_PLATFORM is unset so ARM64
Ninja builds get "ARM64" instead of defaulting to "X64" (update the block that
sets _tilelang_machine and checks CMAKE_GENERATOR_PLATFORM to consult those
fallback variables).
In `@scripts/windows/uv-sync.cmd`:
- Line 8: Replace the bare %1 check with the quote-stripped form so the
terminator comparison handles quoted passthrough arguments correctly: change the
conditional that compares "%1" to "" to use %~1 (the stripped-quote form) in the
if test that jumps to the args_done label (the existing if ... goto args_done
check), ensuring the check uses %~1 for robustness with quoted arguments.
In `@src/runtime/logging.cc`:
- Around line 108-124: Remove the unreachable "return settings;" statements that
follow the LOG(FATAL) calls in the level-parsing block (the checks for
ill-formed token, empty level, and invalid numeric level) because LOG(FATAL)
routes to LogFatalImpl which is [[noreturn]]; keep behavior consistent with the
earlier LOG(FATAL) that already omits a trailing return and leave only the
LOG(FATAL) calls in place.
- Around line 3-4: The translation unit uses std::istringstream and
std::stringstream in ParseSpec and VLogContext::str() but doesn't include
<sstream>, relying on transitive includes; add an explicit `#include` <sstream> at
the top of the file so ParseSpec and VLogContext::str() have the required header
and the compilation is self-contained.
In `@src/target/stubs/nvrtc.cc`:
- Around line 32-45: The Windows nvRTC candidate array kLibNvrtcPaths in
nvrtc.cc stops at CUDA 13 names so systems with CUDA 14 may fail to load; update
kLibNvrtcPaths to include "nvrtc64_140_0.dll" (preferably prepended to the
Windows list) or add a clear TODO comment above kLibNvrtcPaths noting to add
future nvrtc64_XXX_0.dll entries for new CUDA major versions so the loader finds
newer libraries without silent failures.
In `@src/transform/layout_reducer.cc`:
- Around line 174-175: Update the function docstring to remove the stale
reference to analyzer_->const_int_bound and instead document that it reads and
relies on thread_var_, var_to_buffer_, and the GetThreadConstIntBound(...)
helper; specifically replace mentions of analyzer_->const_int_bound with
GetThreadConstIntBound(...) and briefly state the expected contract/return
(e.g., that it provides per-thread constant integer bounds) so the doc matches
the implementation.
- Around line 391-403: ReducerInfoDetector currently continues traversing even
after has_reducer_info is true; modify its VisitStmt_ override(s) (at least
VisitStmt_(const BlockNode *op)) to early-return immediately when
has_reducer_info is true before examining annotations or recursing via
StmtVisitor::VisitStmt_, and apply the same guard in any other
VisitStmt_/visitor overrides in ReducerInfoDetector to prevent further recursion
once the flag is set.
In `@tilelang/__init__.py`:
- Around line 125-126: The except block swallowing all exceptions around
tvm_ffi.libinfo.find_libtvm_ffi() should log the exception so failures are
discoverable; update the except Exception: pass in tilelang/__init__.py to catch
Exception as e and call logger.debug (or logger.exception) with a clear message
referencing find_libtvm_ffi and the exception, so if the runtime path isn't
prepended and import tvm fails the root cause is visible (logger is available at
module scope).
- Around line 177-191: The CUDA DLL directory list is duplicated here; instead
of rebuilding the same tuple, call tilelang.libinfo.get_dll_directories() (or
add a small helper in tilelang.libinfo like get_cuda_dll_directories() if you
want CUDA-only behavior) from the __init__.py block that runs on Windows (the
code using _prepend_windows_runtime_paths, _dll_handles, cuda_home); use the
returned iterable to iterate and call os.add_dll_directory for each existing
path, and remove the hard-coded tuple so discovery logic is centralized and not
duplicated (also keep the existing use of _load_tile_lang_lib which already
relies on libinfo.get_dll_directories()).
- Around line 54-57: The file currently sets os.RTLD_LAZY = 0 on Windows which
mutates the global os module; instead remove that import-time assignment and
update the loader to use a local fallback flag (e.g., read rtld_flag =
getattr(os, "RTLD_LAZY", 0)) when building the CDLL flags inside the
_lazy_load_lib function so you no longer inject RTLD_LAZY into os; locate
references to os.RTLD_LAZY and replace them with a local getattr fallback in
_lazy_load_lib (or compute a module-level constant like _RTLD_LAZY = getattr(os,
"RTLD_LAZY", 0) without assigning to os) to preserve current behavior without
mutating the os module.
In `@tilelang/contrib/cc.py`:
- Around line 54-60: The Windows-specific MSVC discovery logic duplicated in
get_cc and get_cplus_compiler should be extracted into a helper (e.g.,
_find_msvc_cl) that performs the same steps: import get_msvc_subprocess_env and
get_env_path from tilelang.contrib.msvc, call get_msvc_subprocess_env(), compute
the env path via get_env_path(compiler_env or {}), locate "cl.exe" with
shutil.which(..., path=...), and return the discovered path; then replace the
duplicated blocks in get_cc and get_cplus_compiler with a call to
_find_msvc_cl() so behavior remains identical and both functions stay in sync.
In `@tilelang/contrib/msvc.py`:
- Around line 174-179: The MD5 call that generates prefix (variable prefix via
hashlib.md5(...).hexdigest()) should explicitly mark the hash as non-security by
passing usedforsecurity=False to hashlib.md5 to silence Ruff S324 and document
intent; update the hashlib.md5 invocation in this function/block that computes
prefix so it becomes hashlib.md5(<data>, usedforsecurity=False) (Python 3.9+).
- Around line 275-287: The current subprocess invocation in msvc.py uses
subprocess.Popen and proc.communicate() with no timeout, which can hang; modify
the call in the block that creates proc (the variables cmd, cwd, compiler_env
and the proc variable) to use a timeout (e.g., communicate(timeout=1800) or
replace with subprocess.run(..., timeout=1800, capture_output=True)) and catch
subprocess.TimeoutExpired to kill/terminate the child process, collect any
partial output (use proc.stdout/communicate after terminate if needed), and
raise a RuntimeError that includes the timeout indication plus any captured
output and the command line (same message structure as the existing RuntimeError
using py_str(out) and " ".join(cmd)). Ensure the timeout handling is applied
around the proc.communicate()/subprocess.run call and that processes are
properly terminated to avoid zombie processes.
In `@tilelang/jit/adapter/libgen.py`:
- Around line 152-158: The subprocess.run invocation that assigns ret (the call
using command, timeout, env=compiler_env, capture_output=True, text=True) can
raise UnicodeDecodeError on Windows; update that call to explicitly pass
encoding and errors (e.g., encoding='utf-8' and errors='replace' or
errors='surrogateescape') so decoder falls back instead of throwing, and keep
using ret.stdout/ret.stderr for logging; modify the subprocess.run call in
libgen.py where ret is set to include these parameters.
In `@tilelang/libinfo.py`:
- Around line 49-56: The nested for…else is redundant: iterate TL_LIBS and
lib_names, return lib_dll_path on first existing file, and if none found raise a
RuntimeError after both loops; remove the for...else construct and place the
error raise after the loops. Update the error message to list the full candidate
paths (each os.path.join(lib_root, lib_name)) instead of calling TL_LIBS "List
of candidates", referencing TL_LIBS, lib_names and lib_dll_path so the message
shows the exact searched file paths.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a0964f23-707e-4b99-8a1d-ed07241fe60c
📒 Files selected for processing (35)
.gitattributes.gitignoreCMakeLists.txtcmake/FindPipCUDAToolkit.cmakecmake/find_pip_cuda.pycmake/generate_windows_import_lib.pycmake/pypi-z3/FindZ3.cmakedocs/get_started/Installation.mdpyproject.tomlscripts/windows/uv-sync.cmdsrc/layout/gemm_layouts.ccsrc/op/operator.hsrc/op/parallel.ccsrc/runtime/logging.ccsrc/target/codegen_cuda.ccsrc/target/codegen_cuda.hsrc/target/stubs/cuda.ccsrc/target/stubs/cuda.hsrc/target/stubs/cudart.ccsrc/target/stubs/dynlib.hsrc/target/stubs/nvrtc.ccsrc/transform/common/pipeline_utils.hsrc/transform/layout_inference.ccsrc/transform/layout_reducer.cctilelang/__init__.pytilelang/cache/kernel_cache.pytilelang/carver/arch/driver/cuda_driver.pytilelang/contrib/cc.pytilelang/contrib/msvc.pytilelang/contrib/nvcc.pytilelang/env.pytilelang/jit/adapter/libgen.pytilelang/jit/adapter/tvm_ffi.pytilelang/jit/adapter/wrapper.pytilelang/libinfo.py
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tilelang/jit/adapter/wrapper.py (1)
751-758:⚠️ Potential issue | 🟡 MinorCPU backend
callis missing the export annotation for Windows compatibility.
PREDEF_INIT_FUNCdefinesTL_EXPORT(as__declspec(dllexport)on Windows) and applies it toget_last_errorandinit, whilePREDEF_HOST_FUNCapplies it to thecallfunction in GPU backends. However,TLCPUSourceWrapper.CALL_PREFIXdeclarescallwith onlyextern "C"and noTL_EXPORTannotation. On Windows with MSVC, without this annotation, the symbol won't be exported from the generated DLL, causing symbol resolution failures when attempting to load the CPU-compiled kernel.Suggested fix
CALL_PREFIX = textwrap.dedent(""" `#ifdef` __cplusplus - extern "C" + extern "C" TL_EXPORT `#endif` int32_t call({}) {{ return {}; }} """)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/jit/adapter/wrapper.py` around lines 751 - 758, The CALL_PREFIX in TLCPUSourceWrapper is missing the TL_EXPORT annotation used by PREDEF_INIT_FUNC and PREDEF_HOST_FUNC, so update CALL_PREFIX to apply TL_EXPORT to the exported call symbol (keep extern "C" and the int32_t signature) so the declaration matches the pattern used for get_last_error/init and GPU call exports; ensure the TL_EXPORT macro is placed alongside extern "C" (e.g., extern "C" TL_EXPORT int32_t call(...)) and that the replacement still fills in the function parameters/return expression as before.tilelang/jit/adapter/libgen.py (1)
140-149:⚠️ Potential issue | 🟠 MajorWindows NVCC invocation is missing the MSVC environment — inconsistent with the fix in
compile_cuda.
subprocess.run(command, timeout=timeout)inherits the parent process environment. On Windows, NVCC needs to locatecl.exe(the host compiler), which is the reasontilelang/contrib/nvcc.py::compile_cudawas updated to passenv=get_nvcc_subprocess_env(). The same fix is needed here — otherwise JIT library compilation will fail for any user whose Python process wasn't launched from a VS Developer Command Prompt, even whenget_msvc_subprocess_env()can successfully bootstrap MSVC.🛠️ Proposed fix
-from tilelang.contrib.nvcc import ( - get_nvcc_compiler, - get_target_arch, - get_target_compute_version, -) +from tilelang.contrib.nvcc import ( + get_nvcc_compiler, + get_nvcc_subprocess_env, + get_target_arch, + get_target_compute_version, +)try: if verbose: print(f"compile_lib compilation command: {' '.join(command)}") - ret = subprocess.run(command, timeout=timeout) + compile_env = get_nvcc_subprocess_env() if is_cuda_target(target) else None + ret = subprocess.run(command, timeout=timeout, env=compile_env) except Exception as e: raise RuntimeError(f"Compile kernel failed because of {e}") from e🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/jit/adapter/libgen.py` around lines 140 - 149, The subprocess.run call in libgen.compile_lib (where command is executed and assigned to ret) fails on Windows because it inherits the parent env and NVCC cannot find MSVC; update that subprocess.run(...) to pass env=get_nvcc_subprocess_env() (the same helper used in tilelang/contrib/nvcc.py::compile_cuda) when running on Windows so NVCC can locate cl.exe and other MSVC tools; ensure you import or reference get_nvcc_subprocess_env() in this module and only apply it for Windows invocations.
♻️ Duplicate comments (2)
src/transform/layout_inference.cc (1)
755-761:⚠️ Potential issue | 🟠 MajorHard
ICHECKonMap<Var, Layout>cast will crash on re-entry.
LayoutInferencer::VisitStmt_(BlockNode)(Lines 1229-1230) writesresult_.layout_map(aMap<Buffer, Layout>) back intoattr::kLayoutMap. IfLayoutInferenceis ever run a second time on the transformed IR (or any other pass reuses the same annotation key with aBuffer-keyed map), Line 760'sICHECK(map_opt.has_value()) << "layout map must be Map<Var, Layout>"will abort instead of gracefully skipping the already-materialized layout map. Prefer an earlyreturnwhen the cast fails so this pass is idempotent.Suggested fix
if (op->annotations.count(attr::kLayoutMap)) { auto layout_map_ref = op->annotations.Get(attr::kLayoutMap); - ICHECK(layout_map_ref.has_value()); + if (!layout_map_ref.has_value()) { + return; + } auto map_opt = layout_map_ref.value().as<Map<Var, Layout>>(); - ICHECK(map_opt.has_value()) << "layout map must be Map<Var, Layout>"; - auto map = map_opt.value(); + if (!map_opt.has_value()) { + // LayoutInferencer emits Map<Buffer, Layout> on output blocks; + // ignore already-materialized annotation forms here. + return; + } + auto map = map_opt.value();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/transform/layout_inference.cc` around lines 755 - 761, Replace the hard ICHECK when casting the annotation value to Map<Var, Layout> with a safe check and early return so the pass is idempotent: in the block handling op->annotations.Get(attr::kLayoutMap) (where you currently call map_opt.value() and ICHECK(map_opt.has_value())), detect when the annotation cannot be cast to Map<Var, Layout> (e.g., map_opt.has_value() is false) and simply skip/apply no-op (return/continue) instead of aborting; this ensures a second run (or a layout_map keyed by Buffer from LayoutInferencer::VisitStmt_(BlockNode) writing result_.layout_map) won't crash the pass.src/target/codegen_cuda.cc (1)
3591-3595:⚠️ Potential issue | 🟠 Major
tl::__tan()still usesCUDAFastMath— precision regression persists after the extraction.For float32,
CUDAFastMathemits__tanf, but this file intentionally definesCUDAFastMathTan(lines 167-185) with an inline comment explaining that__tanfdeviates too much from numpy'stan. Moving the handler intoHandleLateIntrinsicCallcarried over the incorrect functor, so the fp32 tangent path still falls back to the less accurate intrinsic.Proposed fix
} else if (op->op.same_as(tl::__tan())) { - CUDAFastMath math_func; + CUDAFastMathTan math_func; std::string func_name = math_func(op->dtype, "tan"); os << func_name << "(" << PrintExpr(op->args[0]) << ")"; return true;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/codegen_cuda.cc` around lines 3591 - 3595, The tl::__tan() branch is still constructing CUDAFastMath (which emits __tanf for float32) causing a precision regression; update the handler in HandleLateIntrinsicCall (the tl::__tan() case) to use CUDAFastMathTan instead of CUDAFastMath so fp32 paths call the higher-accuracy implementation declared as CUDAFastMathTan (preserve the current call pattern that builds func_name and prints the arg via PrintExpr(op->args[0])). Ensure you replace the CUDAFastMath instantiation with CUDAFastMathTan and keep the emitted function name logic identical so other types are unaffected.
🧹 Nitpick comments (9)
src/layout/gemm_layouts.cc (1)
14-20: LGTM — portable unreachable hint.The macro correctly maps to the canonical builtins per compiler and falls back safely. Ordering
_MSC_VERbefore__clang__also handles clang-cl correctly (which defines both and supports__assume(0)).Two optional nits you can ignore:
- Since this lives in a
.ccand is only used locally, you could#undef TILELANG_COMPILER_UNREACHABLEat the bottom of the file to avoid leaking the macro if this header-style block ever gets relocated.- On the
#elsefallback,((void)0)gives no UB hint, so some compilers may warn about a missing return at the call sites. If that becomes an issue on an exotic toolchain, considerstd::abort()(requires<cstdlib>) as a safer fallback.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/layout/gemm_layouts.cc` around lines 14 - 20, Macro TILELANG_COMPILER_UNREACHABLE is fine but please add a cleanup and an optional safer fallback: at the end of src/layout/gemm_layouts.cc undefine TILELANG_COMPILER_UNREACHABLE to avoid leaking the macro, and if you hit missing-return warnings on exotic toolchains replace the `#else` fallback ((void)0) with a call to std::abort() (add `#include` <cstdlib>), keeping the existing branches for _MSC_VER and __GNUC__/__clang__ unchanged.src/runtime/logging.cc (1)
126-128: Duplicatefilename=<level>entries silently keep the first value.
std::unordered_map::emplacedoes not overwrite on existing keys, so specs likeTVM_LOG_DEBUG="foo.cc=1,foo.cc=3"silently ignore the3and keep1. This matches the upstream TVM behavior, so it's likely intentional, but if last-writer-wins is preferred here, usingoperator[]orinsert_or_assignwould be clearer. Also consider emitting a warning (or at least avoiding theLOG(INFO)for the shadowed entry) to avoid misleading users who see both assignments acknowledged.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/runtime/logging.cc` around lines 126 - 128, The current code uses settings.vlog_level_map_.emplace(name, level_val) which silently keeps the first value for duplicate keys (e.g., "foo.cc=1,foo.cc=3"); change this to insert_or_assign(name, level_val) (or settings.vlog_level_map_[name] = level_val) so the last-writer wins, and add a LOG(WARNING) (or avoid LOG(INFO) for shadowed entries) when a duplicate key is detected to inform users that an earlier value was overwritten; reference settings.vlog_level_map_ and the code block handling TVM_LOG_DEBUG parsing in logging.cc to locate and apply the change.src/target/stubs/cuda.h (1)
231-231: Nit: trailing#endifcomment understates the condition.The guard opened at line 159 is
#if !defined(_WIN32) || defined(__CYGWIN__), so the closing comment// !_WIN32is slightly misleading — under Cygwin (where__CYGWIN__is defined) the block is still compiled regardless of_WIN32. A tiny tweak keeps future readers from assuming the guard is purely "non-Windows".📝 Proposed comment tweak
-#endif // !_WIN32 +#endif // !_WIN32 || __CYGWIN__🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/stubs/cuda.h` at line 231, The closing preprocessor comment misstates the opening guard: update the trailing comment on the closing '#endif // !_WIN32' to match the opener '#if !defined(_WIN32) || defined(__CYGWIN__)' so readers see the full condition (for example: change to '#endif // !defined(_WIN32) || defined(__CYGWIN__)'), referencing the original '#if !defined(_WIN32) || defined(__CYGWIN__)' and the existing closing '#endif'.docs/get_started/Installation.md (1)
77-90: Option A installs the CUDA toolkit twice.
.[nvcc]already resolves tocuda-toolkit[cccl,nvcc,nvrtc]>=13(seepyproject.tomllines 55-64), so the preceding explicitpip install "cuda-toolkit[cccl,nvcc,nvrtc]>=13"is redundant and just slows first-time installs. Consider dropping it and letting the extra handle it, which also keeps the version constraint in a single place.♻️ Proposed simplification
pip install -r requirements-dev.txt -pip install "cuda-toolkit[cccl,nvcc,nvrtc]>=13" pip install .[nvcc] -v🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/get_started/Installation.md` around lines 77 - 90, Remove the redundant explicit installation of the CUDA toolkit: the step `pip install "cuda-toolkit[cccl,nvcc,nvrtc]>=13"` is unnecessary because the package extra `.[nvcc]` already requires `cuda-toolkit[cccl,nvcc,nvrtc]>=13` (see the dependency block in pyproject.toml). Update the Installation.md Option A instructions to drop that explicit pip install line and keep only `pip install .[nvcc] -v` (after installing requirements-dev.txt), ensuring the version constraint stays centralized in pyproject.toml.scripts/windows/uv-sync.cmd (1)
24-38: Optional: de-duplicate the two vswhere lookups.The PF(x86) and PF probes are identical aside from the root env var. Folding them into a small
:find_vsdevcmdsubroutine (or iterating over both roots) keeps behavior the same while halving the block that has to be kept in sync if the vswhere invocation ever changes (e.g., adding-prerelease, tweaking-requires, etc.). Non-blocking — the current form works.♻️ Sketch
-if not defined VSDEVCMD ( - if exist "%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe" ( - for /f "usebackq delims=" %%I in (`"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe" -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -find Common7\Tools\VsDevCmd.bat`) do ( - if not defined VSDEVCMD set "VSDEVCMD=%%~fI" - ) - ) -) - -if not defined VSDEVCMD ( - if exist "%ProgramFiles%\Microsoft Visual Studio\Installer\vswhere.exe" ( - for /f "usebackq delims=" %%I in (`"%ProgramFiles%\Microsoft Visual Studio\Installer\vswhere.exe" -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -find Common7\Tools\VsDevCmd.bat`) do ( - if not defined VSDEVCMD set "VSDEVCMD=%%~fI" - ) - ) -) +for %%R in ("%ProgramFiles(x86)%" "%ProgramFiles%") do ( + if not defined VSDEVCMD if exist "%%~R\Microsoft Visual Studio\Installer\vswhere.exe" ( + for /f "usebackq delims=" %%I in (`"%%~R\Microsoft Visual Studio\Installer\vswhere.exe" -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -find Common7\Tools\VsDevCmd.bat`) do ( + if not defined VSDEVCMD set "VSDEVCMD=%%~fI" + ) + ) +)Worth a quick sanity-run since combining
%%Rexpansion with nested%%Iand%ProgramFiles(x86)%paren-in-quotes is a known batch gotcha — please confirm the refactored form still resolves VsDevCmd on a machine that only has VS in%ProgramFiles%(x64-only install).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/windows/uv-sync.cmd` around lines 24 - 38, Refactor the duplicated vswhere lookup into a single reusable step: create a small batch subroutine (e.g., :find_vsdevcmd) or loop over the two roots ("%ProgramFiles(x86)%" and "%ProgramFiles%") that runs the same vswhere invocation and sets the VSDEVCMD variable when it finds Common7\Tools\VsDevCmd.bat; ensure you still check if VSDEVCMD is defined before calling vswhere and preserve the exact vswhere arguments (-latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -find Common7\Tools\VsDevCmd.bat). After implementing, sanity-run the script on a machine with VS installed only in %ProgramFiles% to confirm the subroutine/loop correctly resolves VsDevCmd and sets VSDEVCMD.tilelang/libinfo.py (1)
49-56: Simplify thefor…else.The inner body uses
return, so the outer loop neverbreaks; theelseclause always executes when nothing was returned, making it equivalent to putting theraiseafter the loop. The current form is easy to misread.♻️ Proposed refactor
for lib_root in TL_LIBS: for lib_name in lib_names: lib_dll_path = os.path.join(lib_root, lib_name) if os.path.exists(lib_dll_path) and os.path.isfile(lib_dll_path): return lib_dll_path - else: - message = f"Cannot find libraries: {', '.join(lib_names)}\n" + "List of candidates:\n" + "\n".join(TL_LIBS) - raise RuntimeError(message) + message = f"Cannot find libraries: {', '.join(lib_names)}\n" + "List of candidates:\n" + "\n".join(TL_LIBS) + raise RuntimeError(message)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/libinfo.py` around lines 49 - 56, The for...else is misleading because the inner loop returns on success and the else always runs when nothing is returned; refactor the search by removing the for...else and placing the error construction and RuntimeError raise after the nested loops: iterate TL_LIBS and lib_names, build lib_dll_path and return when os.path.exists and os.path.isfile succeed (same checks using lib_dll_path), and if no match is found after both loops construct the same message (using lib_names and TL_LIBS) and raise RuntimeError; this keeps symbols TL_LIBS, lib_names, lib_dll_path and the same behavior but eliminates the confusing else clause.src/target/stubs/dynlib.h (1)
62-95: Optional: replace manualmalloc/realloc/freewithstd::vector<HMODULE>.The current code is correct but has several
freepaths that are easy to regress on future edits. Astd::vector<HMODULE>withresize()gives you RAII-cleanup and removes the three explicitfree(modules)bail-outs.♻️ Proposed refactor sketch
inline void *dynlib_find_loaded(const char *symbol_name, void *exclude = nullptr) { HANDLE process = GetCurrentProcess(); - DWORD buf_bytes = 1024 * sizeof(HMODULE); - HMODULE *modules = static_cast<HMODULE *>(malloc(buf_bytes)); - if (!modules) - return nullptr; + std::vector<HMODULE> modules(1024); + DWORD buf_bytes = static_cast<DWORD>(modules.size() * sizeof(HMODULE)); DWORD needed = 0; - if (!EnumProcessModules(process, modules, buf_bytes, &needed)) { - free(modules); + if (!EnumProcessModules(process, modules.data(), buf_bytes, &needed)) return nullptr; - } if (needed > buf_bytes) { - buf_bytes = needed; - HMODULE *resized = static_cast<HMODULE *>(realloc(modules, buf_bytes)); - if (!resized) { - free(modules); - return nullptr; - } - modules = resized; - if (!EnumProcessModules(process, modules, buf_bytes, &needed)) { - free(modules); + modules.resize(needed / sizeof(HMODULE)); + buf_bytes = static_cast<DWORD>(modules.size() * sizeof(HMODULE)); + if (!EnumProcessModules(process, modules.data(), buf_bytes, &needed)) return nullptr; - } } DWORD count = needed / sizeof(HMODULE); for (DWORD i = 0; i < count; i++) { void *sym = reinterpret_cast<void *>(GetProcAddress(modules[i], symbol_name)); - if (sym != nullptr && sym != exclude) { - free(modules); + if (sym != nullptr && sym != exclude) return reinterpret_cast<void *>(modules[i]); - } } - free(modules); return nullptr; }(Requires adding
#include <vector>at the top of the Windows section.)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/stubs/dynlib.h` around lines 62 - 95, Replace the manual malloc/realloc/free buffer handling in the function that enumerates modules (uses EnumProcessModules, modules, needed, buf_bytes) with a std::vector<HMODULE>: add `#include` <vector>, create a vector<HMODULE> v; v.resize(initial_size); call EnumProcessModules(&v[0], v.size() * sizeof(HMODULE), &needed) and if needed > current bytes resize the vector and retry; compute count = needed / sizeof(HMODULE) and iterate using v[i] when calling GetProcAddress(symbol_name) and comparing to exclude, returning the module handle as before—this removes all explicit free() calls and guarantees RAII cleanup on every early return.cmake/pypi-z3/FindZ3.cmake (1)
36-40: Deadif(WIN32)aroundadd_library— both branches are identical.Both the Windows and non-Windows branches call
add_library(z3::libz3 SHARED IMPORTED GLOBAL)with the same arguments, so the condition can be dropped.🧹 Proposed simplification
-if(WIN32) - add_library(z3::libz3 SHARED IMPORTED GLOBAL) -else() - add_library(z3::libz3 SHARED IMPORTED GLOBAL) -endif() +add_library(z3::libz3 SHARED IMPORTED GLOBAL)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/pypi-z3/FindZ3.cmake` around lines 36 - 40, The if(WIN32) conditional around add_library is redundant because both branches call the same command; remove the conditional block and leave a single add_library(z3::libz3 SHARED IMPORTED GLOBAL) call in FindZ3.cmake so the z3::libz3 target is defined unconditionally (referencing the add_library invocation and the z3::libz3 target).tilelang/contrib/cc.py (1)
53-60: Consider extracting the MSVCcl.exeprobe into a helper to avoid duplication.The 5-line Windows branch is copy-pasted between
get_cc()andget_cplus_compiler(). A small private helper (e.g.,_probe_msvc_cl()) that returns the resolved path orNonewould keep both functions in sync if the discovery logic ever evolves.Also applies to: 84-91
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/cc.py` around lines 53 - 60, Extract the duplicated Windows MSVC probe into a single private helper (e.g., _probe_msvc_cl()) and use it from both get_cc() and get_cplus_compiler(); move the logic that imports tilelang.contrib.msvc, calls get_msvc_subprocess_env(), computes get_env_path(compiler_env or {}), and resolves cl via shutil.which("cl.exe", path=...) into that helper which should return the resolved path or None, then replace the duplicated blocks in get_cc() and get_cplus_compiler() with calls to _probe_msvc_cl().
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@CMakeLists.txt`:
- Around line 13-78: tilelang_generate_windows_import_library currently defaults
_tilelang_machine to "X64" and only checks CMAKE_GENERATOR_PLATFORM for
"Win32"/"ARM64", which misses single-config generators like Ninja; update the
machine-detection logic in tilelang_generate_windows_import_library to also
consult CMAKE_SYSTEM_PROCESSOR and CMAKE_VS_PLATFORM_NAME (and/or
CMAKE_SIZEOF_VOID_P == 4) as fallbacks before defaulting to X64 so
cross-compiles to x86/ARM64 with Ninja are correctly handled; modify the block
that sets _tilelang_machine (and any conditional branches referencing
CMAKE_GENERATOR_PLATFORM) to prefer CMAKE_GENERATOR_PLATFORM, then
CMAKE_VS_PLATFORM_NAME, then CMAKE_SYSTEM_PROCESSOR (normalize values like
"x86"/"AMD64"/"arm64"/"aarch64"), and finally default to X64.
In `@docs/get_started/Installation.md`:
- Around line 122-132: Update the Option C example for WITH_PIP_CUDA_TOOLCHAIN
to include Windows equivalents: add instructions showing how to set the
environment variable from a Developer Command Prompt using the `set` syntax and
from PowerShell using the `$env:` syntax, and clarify typical Windows path
formats like .\.venv and .\.venv\Lib\site-packages; modify the Installation.md
Option C block and the trailing note so readers on Windows have direct, runnable
commands for both cmd.exe and PowerShell in addition to the existing bash
example.
In `@src/op/parallel.cc`:
- Around line 85-92: The annotation type for attr::kReducerInfo is inconsistent:
either make it uniformly Map<Var, ReducerInfo> or keep it as Map<Var,
Map<String,String>>; audit and update all readers/writers accordingly.
Concretely, ensure the writer that sets inside_reducer_range_ (the code that
constructs the annotation in layout_reducer.cc) and places it into
op->annotations uses the same type as readers that call
annotations.Get(attr::kReducerInfo) (e.g., the cast in parallel.cc using
reducer_info_map_, the BlockNode reader in layout_reducer.cc,
verify_parallel_loop.cc, and lower_tile_op.cc). If the intended canonical type
is Map<Var, ReducerInfo>, change the casts and subsequent handling in the
readers to use Map<Var, ReducerInfo> and adapt field access to ReducerInfo;
otherwise revert the writer to emit Map<Var, Map<String,String>> and in readers
that need ReducerInfo perform an explicit field-by-field conversion from
Map<String,String> to ReducerInfo (mirroring the existing BlockNode conversion
pattern).
In `@src/target/stubs/cudart.cc`:
- Around line 107-111: Update the error message emitted in cudart.cc where
fprintf(stderr, "...") and abort() are called: change the wording to mention
Windows DLL names (cudart64_*.dll), pip-installed CUDA wheels, and standalone
CUDA toolkit installs as possible sources of missing symbols rather than only
saying "install PyTorch with CUDA"; locate the fprintf in the cudart
symbols-not-found branch and replace the message with a concise multi-platform
guidance that suggests checking CUDA toolkit installation, PyTorch CUDA builds,
or pip CUDA wheels and verifying DLL/so paths/resolvability before aborting.
In `@tilelang/contrib/msvc.py`:
- Around line 174-179: The MD5 call used to produce the deterministic filename
prefix (the hashlib.md5(...) call that assigns prefix) should be explicitly
marked as non-security-sensitive by passing usedforsecurity=False; update the
call that builds prefix so it uses hashlib.md5(..., usedforsecurity=False). This
affects the prefix generation (variable name prefix) and ensures the intent is
clear to linters and FIPS systems while leaving the rest of the patched_path
creation and file write (patched_path, tmp_dir, patched) unchanged.
---
Outside diff comments:
In `@tilelang/jit/adapter/libgen.py`:
- Around line 140-149: The subprocess.run call in libgen.compile_lib (where
command is executed and assigned to ret) fails on Windows because it inherits
the parent env and NVCC cannot find MSVC; update that subprocess.run(...) to
pass env=get_nvcc_subprocess_env() (the same helper used in
tilelang/contrib/nvcc.py::compile_cuda) when running on Windows so NVCC can
locate cl.exe and other MSVC tools; ensure you import or reference
get_nvcc_subprocess_env() in this module and only apply it for Windows
invocations.
In `@tilelang/jit/adapter/wrapper.py`:
- Around line 751-758: The CALL_PREFIX in TLCPUSourceWrapper is missing the
TL_EXPORT annotation used by PREDEF_INIT_FUNC and PREDEF_HOST_FUNC, so update
CALL_PREFIX to apply TL_EXPORT to the exported call symbol (keep extern "C" and
the int32_t signature) so the declaration matches the pattern used for
get_last_error/init and GPU call exports; ensure the TL_EXPORT macro is placed
alongside extern "C" (e.g., extern "C" TL_EXPORT int32_t call(...)) and that the
replacement still fills in the function parameters/return expression as before.
---
Duplicate comments:
In `@src/target/codegen_cuda.cc`:
- Around line 3591-3595: The tl::__tan() branch is still constructing
CUDAFastMath (which emits __tanf for float32) causing a precision regression;
update the handler in HandleLateIntrinsicCall (the tl::__tan() case) to use
CUDAFastMathTan instead of CUDAFastMath so fp32 paths call the higher-accuracy
implementation declared as CUDAFastMathTan (preserve the current call pattern
that builds func_name and prints the arg via PrintExpr(op->args[0])). Ensure you
replace the CUDAFastMath instantiation with CUDAFastMathTan and keep the emitted
function name logic identical so other types are unaffected.
In `@src/transform/layout_inference.cc`:
- Around line 755-761: Replace the hard ICHECK when casting the annotation value
to Map<Var, Layout> with a safe check and early return so the pass is
idempotent: in the block handling op->annotations.Get(attr::kLayoutMap) (where
you currently call map_opt.value() and ICHECK(map_opt.has_value())), detect when
the annotation cannot be cast to Map<Var, Layout> (e.g., map_opt.has_value() is
false) and simply skip/apply no-op (return/continue) instead of aborting; this
ensures a second run (or a layout_map keyed by Buffer from
LayoutInferencer::VisitStmt_(BlockNode) writing result_.layout_map) won't crash
the pass.
---
Nitpick comments:
In `@cmake/pypi-z3/FindZ3.cmake`:
- Around line 36-40: The if(WIN32) conditional around add_library is redundant
because both branches call the same command; remove the conditional block and
leave a single add_library(z3::libz3 SHARED IMPORTED GLOBAL) call in
FindZ3.cmake so the z3::libz3 target is defined unconditionally (referencing the
add_library invocation and the z3::libz3 target).
In `@docs/get_started/Installation.md`:
- Around line 77-90: Remove the redundant explicit installation of the CUDA
toolkit: the step `pip install "cuda-toolkit[cccl,nvcc,nvrtc]>=13"` is
unnecessary because the package extra `.[nvcc]` already requires
`cuda-toolkit[cccl,nvcc,nvrtc]>=13` (see the dependency block in
pyproject.toml). Update the Installation.md Option A instructions to drop that
explicit pip install line and keep only `pip install .[nvcc] -v` (after
installing requirements-dev.txt), ensuring the version constraint stays
centralized in pyproject.toml.
In `@scripts/windows/uv-sync.cmd`:
- Around line 24-38: Refactor the duplicated vswhere lookup into a single
reusable step: create a small batch subroutine (e.g., :find_vsdevcmd) or loop
over the two roots ("%ProgramFiles(x86)%" and "%ProgramFiles%") that runs the
same vswhere invocation and sets the VSDEVCMD variable when it finds
Common7\Tools\VsDevCmd.bat; ensure you still check if VSDEVCMD is defined before
calling vswhere and preserve the exact vswhere arguments (-latest -products *
-requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -find
Common7\Tools\VsDevCmd.bat). After implementing, sanity-run the script on a
machine with VS installed only in %ProgramFiles% to confirm the subroutine/loop
correctly resolves VsDevCmd and sets VSDEVCMD.
In `@src/layout/gemm_layouts.cc`:
- Around line 14-20: Macro TILELANG_COMPILER_UNREACHABLE is fine but please add
a cleanup and an optional safer fallback: at the end of
src/layout/gemm_layouts.cc undefine TILELANG_COMPILER_UNREACHABLE to avoid
leaking the macro, and if you hit missing-return warnings on exotic toolchains
replace the `#else` fallback ((void)0) with a call to std::abort() (add `#include`
<cstdlib>), keeping the existing branches for _MSC_VER and __GNUC__/__clang__
unchanged.
In `@src/runtime/logging.cc`:
- Around line 126-128: The current code uses
settings.vlog_level_map_.emplace(name, level_val) which silently keeps the first
value for duplicate keys (e.g., "foo.cc=1,foo.cc=3"); change this to
insert_or_assign(name, level_val) (or settings.vlog_level_map_[name] =
level_val) so the last-writer wins, and add a LOG(WARNING) (or avoid LOG(INFO)
for shadowed entries) when a duplicate key is detected to inform users that an
earlier value was overwritten; reference settings.vlog_level_map_ and the code
block handling TVM_LOG_DEBUG parsing in logging.cc to locate and apply the
change.
In `@src/target/stubs/cuda.h`:
- Line 231: The closing preprocessor comment misstates the opening guard: update
the trailing comment on the closing '#endif // !_WIN32' to match the opener '#if
!defined(_WIN32) || defined(__CYGWIN__)' so readers see the full condition (for
example: change to '#endif // !defined(_WIN32) || defined(__CYGWIN__)'),
referencing the original '#if !defined(_WIN32) || defined(__CYGWIN__)' and the
existing closing '#endif'.
In `@src/target/stubs/dynlib.h`:
- Around line 62-95: Replace the manual malloc/realloc/free buffer handling in
the function that enumerates modules (uses EnumProcessModules, modules, needed,
buf_bytes) with a std::vector<HMODULE>: add `#include` <vector>, create a
vector<HMODULE> v; v.resize(initial_size); call EnumProcessModules(&v[0],
v.size() * sizeof(HMODULE), &needed) and if needed > current bytes resize the
vector and retry; compute count = needed / sizeof(HMODULE) and iterate using
v[i] when calling GetProcAddress(symbol_name) and comparing to exclude,
returning the module handle as before—this removes all explicit free() calls and
guarantees RAII cleanup on every early return.
In `@tilelang/contrib/cc.py`:
- Around line 53-60: Extract the duplicated Windows MSVC probe into a single
private helper (e.g., _probe_msvc_cl()) and use it from both get_cc() and
get_cplus_compiler(); move the logic that imports tilelang.contrib.msvc, calls
get_msvc_subprocess_env(), computes get_env_path(compiler_env or {}), and
resolves cl via shutil.which("cl.exe", path=...) into that helper which should
return the resolved path or None, then replace the duplicated blocks in get_cc()
and get_cplus_compiler() with calls to _probe_msvc_cl().
In `@tilelang/libinfo.py`:
- Around line 49-56: The for...else is misleading because the inner loop returns
on success and the else always runs when nothing is returned; refactor the
search by removing the for...else and placing the error construction and
RuntimeError raise after the nested loops: iterate TL_LIBS and lib_names, build
lib_dll_path and return when os.path.exists and os.path.isfile succeed (same
checks using lib_dll_path), and if no match is found after both loops construct
the same message (using lib_names and TL_LIBS) and raise RuntimeError; this
keeps symbols TL_LIBS, lib_names, lib_dll_path and the same behavior but
eliminates the confusing else clause.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1871ce26-d4a5-43bd-be60-1a9c9b709986
📒 Files selected for processing (35)
.gitattributes.gitignoreCMakeLists.txtcmake/FindPipCUDAToolkit.cmakecmake/find_pip_cuda.pycmake/generate_windows_import_lib.pycmake/pypi-z3/FindZ3.cmakedocs/get_started/Installation.mdpyproject.tomlscripts/windows/uv-sync.cmdsrc/layout/gemm_layouts.ccsrc/op/operator.hsrc/op/parallel.ccsrc/runtime/logging.ccsrc/target/codegen_cuda.ccsrc/target/codegen_cuda.hsrc/target/stubs/cuda.ccsrc/target/stubs/cuda.hsrc/target/stubs/cudart.ccsrc/target/stubs/dynlib.hsrc/target/stubs/nvrtc.ccsrc/transform/common/pipeline_utils.hsrc/transform/layout_inference.ccsrc/transform/layout_reducer.cctilelang/__init__.pytilelang/cache/kernel_cache.pytilelang/carver/arch/driver/cuda_driver.pytilelang/contrib/cc.pytilelang/contrib/msvc.pytilelang/contrib/nvcc.pytilelang/env.pytilelang/jit/adapter/libgen.pytilelang/jit/adapter/tvm_ffi.pytilelang/jit/adapter/wrapper.pytilelang/libinfo.py
✅ Files skipped from review due to trivial changes (3)
- .gitignore
- .gitattributes
- cmake/find_pip_cuda.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tilelang/env.py
- src/transform/layout_reducer.cc
- tilelang/cache/kernel_cache.py
- pyproject.toml
- src/target/stubs/nvrtc.cc
- src/target/stubs/cuda.cc
- cmake/FindPipCUDAToolkit.cmake
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/jit/adapter/libgen.py (1)
146-154:⚠️ Potential issue | 🟠 MajorWindows NVCC invocation is missing the MSVC host environment.
subprocess.run(command, timeout=timeout)does not passenv=, so on Windows NVCC inherits only the parent process environment and will not see the MSVC toolchain (PATH,INCLUDE,LIB,WindowsSdkDir, …) discovered byget_nvcc_subprocess_env()/get_msvc_subprocess_env(). Unless the Python process was itself launched from a Visual Studio developer prompt, NVCC will fail to locatecl.exeor the Windows SDK headers and the whole Windows JIT path breaks.Mirror the treatment in
tilelang/contrib/nvcc.py::compile_cuda(which now threadsenv=compiler_envthroughsubprocess.Popen).🛠️ Proposed fix
+ from tilelang.contrib.nvcc import get_nvcc_subprocess_env + compiler_env = get_nvcc_subprocess_env() if is_cuda_target(target) else None try: if verbose: print(f"compile_lib compilation command: {' '.join(command)}") - ret = subprocess.run(command, timeout=timeout) + ret = subprocess.run(command, timeout=timeout, env=compiler_env) except Exception as e: raise RuntimeError(f"Compile kernel failed because of {e}") from e(Adjust imports so the helper is only imported for the CUDA branch.)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/jit/adapter/libgen.py` around lines 146 - 154, The subprocess.run call that invokes NVCC must be passed the proper MSVC/NVCC environment so cl.exe and Windows SDK headers are visible; obtain the environment via the existing helpers (get_nvcc_subprocess_env()/get_msvc_subprocess_env()) and pass it as env=compiler_env to subprocess.run(command, timeout=timeout, env=compiler_env) (mirror tilelang/contrib/nvcc.py::compile_cuda). Only import those helpers for the CUDA branch if needed and ensure the same env is used for the return-code check (ret.returncode) and error messages.
♻️ Duplicate comments (1)
src/target/codegen_cuda.cc (1)
3591-3595:⚠️ Potential issue | 🟠 MajorUse
CUDAFastMathTanfortl::__tan().The
tanbranch still routes throughCUDAFastMath, which emits__tanffor float32. TheCUDAFastMathTanspecialization exists precisely because__tanfwas found too inaccurate vs.tanf, so this keeps the prior precision regression in place after the extraction.Proposed fix
} else if (op->op.same_as(tl::__tan())) { - CUDAFastMath math_func; + CUDAFastMathTan math_func; std::string func_name = math_func(op->dtype, "tan"); os << func_name << "(" << PrintExpr(op->args[0]) << ")"; return true;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/codegen_cuda.cc` around lines 3591 - 3595, The tan operator branch currently uses CUDAFastMath (producing __tanf for float32) which is too imprecise; change the branch handling tl::__tan() to use the CUDAFastMathTan specialization instead of CUDAFastMath so the emitted function respects the higher-precision tan implementation; update the code that builds func_name (when op->op.same_as(tl::__tan())) to call CUDAFastMathTan(op->dtype, "tan") rather than CUDAFastMath, keeping the existing usage of PrintExpr(op->args[0]) and returning true.
🧹 Nitpick comments (12)
src/layout/gemm_layouts.cc (1)
14-20: LGTM on the compiler-dispatch macro.The MSVC/GCC/Clang/fallback arms are correct, and the surrounding call sites (Lines 870, 899) only reach the macro after an
ICHECK(0) << ...that is guaranteed to throw, so the no-op fallback on unknown compilers is safe w.r.t. correctness. Minor optional nit: if this macro starts being needed elsewhere in the codebase, consider hoisting it into a shared header (e.g., atl_compiler.h) so the definition isn't duplicated per TU.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/layout/gemm_layouts.cc` around lines 14 - 20, Macro TILELANG_COMPILER_UNREACHABLE is implemented correctly for MSVC/GCC/Clang with a safe no-op fallback, but consider hoisting it to a shared header to avoid duplication; move the TILELANG_COMPILER_UNREACHABLE definition out of src/layout/gemm_layouts.cc into a common header (e.g., tl_compiler.h) and replace occurrences in this TU and any others with an `#include` of that header so callers like the sites that follow ICHECK(0) can reuse the same, single definition.tilelang/contrib/msvc.py (1)
152-161:-std=...is silently dropped with no MSVC equivalent emitted.Callers (e.g., TVM/TileLang JIT FFI glue) commonly pass
-std=c++17. Dropping it without appending/std:c++17leavescl.exeat its default (C++14 pre-VS2022 defaults), which can cause generated FFI sources that rely on C++17 features (structured bindings,if constexpr,std::string_view) to fail to compile on Windows. Similarly,-g→/Z7or/Zishould be considered if debug symbols matter for kernel cache debugging.♻️ Proposed mapping
def _normalize_option(option: str) -> str | None: if option.startswith("-I"): return "/I" + option[2:] if option.startswith("-D"): return "/D" + option[2:] - if option in ("-g", "-fPIC"): - return None - if option.startswith("-std="): - return None + if option == "-fPIC": + return None + if option == "-g": + return "/Z7" + if option.startswith("-std="): + return "/std:" + option[len("-std="):] return option🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` around lines 152 - 161, The _normalize_option function currently drops -std=... and -g options silently; update _normalize_option to translate -std=<ver> into MSVC's /std:<ver> (e.g., -std=c++17 -> /std:c++17) and map -g to an appropriate MSVC debug flag (e.g., /Zi or /Z7) instead of returning None; locate and edit the _normalize_option function in msvc.py and ensure it parses the suffix after "-std=" and returns "/std:" + suffix, and replace the -g branch to return "/Zi" (or "/Z7" if you prefer full debug info) so callers like the TVM/TileLang JIT FFI glue that pass "-std=c++17" get equivalent MSVC flags emitted.tilelang/jit/adapter/libgen.py (1)
63-66: Path extension replacement is fragile on Windows.
src.name.replace(".cu", _lib_ext)replaces the first occurrence of".cu"anywhere in the path, not just the extension. If the tempfile lands under a directory that happens to contain.cu(e.g., a user project path, a cached kernel directory named.cuda_cache, etc.), the outputlibpathwill be mangled. This pre-existed but is worth tightening alongside the new Windows extension logic.♻️ Proposed fix
- _lib_ext = ".dll" if sys.platform == "win32" else ".so" - src = tempfile.NamedTemporaryFile(mode="w", suffix=".cu", delete=False) # noqa: SIM115 - target_arch = get_target_arch(get_target_compute_version(target)) - libpath = src.name.replace(".cu", _lib_ext) + _lib_ext = ".dll" if sys.platform == "win32" else ".so" + src = tempfile.NamedTemporaryFile(mode="w", suffix=".cu", delete=False) # noqa: SIM115 + target_arch = get_target_arch(get_target_compute_version(target)) + libpath = os.path.splitext(src.name)[0] + _lib_extThe same concern applies to the
.cpp→.sopaths at lines 106 and 124.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/jit/adapter/libgen.py` around lines 63 - 66, The code currently builds library paths by calling src.name.replace(".cu", _lib_ext) which can accidentally replace ".cu" anywhere in the temp path; update the logic to replace only the file extension using os.path.splitext on src.name (e.g., use os.path.splitext(src.name)[0] + _lib_ext) for the .cu→.dll/.so case (around the src = tempfile.NamedTemporaryFile(...) / libpath = ... block) and apply the same change to the other occurrences where .cpp is converted to .so (the .cpp → .so conversions later in the file) so extension replacement is robust on Windows and other platforms.src/target/stubs/dynlib.h (1)
59-95: Replace manualmalloc/reallocwithstd::vector<HMODULE>.This is a C++ header (uses
namespace,reinterpret_cast,thread_local), so prefer RAII over pairedmalloc/freeacross multiple early-return branches — it removes the risk of leaks if the code is extended and makes the grow-and-retry logic more idiomatic.♻️ Proposed refactor
inline void *dynlib_find_loaded(const char *symbol_name, void *exclude = nullptr) { HANDLE process = GetCurrentProcess(); - DWORD buf_bytes = 1024 * sizeof(HMODULE); - HMODULE *modules = static_cast<HMODULE *>(malloc(buf_bytes)); - if (!modules) - return nullptr; + std::vector<HMODULE> modules(1024); + DWORD buf_bytes = static_cast<DWORD>(modules.size() * sizeof(HMODULE)); DWORD needed = 0; - if (!EnumProcessModules(process, modules, buf_bytes, &needed)) { - free(modules); + if (!EnumProcessModules(process, modules.data(), buf_bytes, &needed)) { return nullptr; } if (needed > buf_bytes) { - buf_bytes = needed; - HMODULE *resized = static_cast<HMODULE *>(realloc(modules, buf_bytes)); - if (!resized) { - free(modules); - return nullptr; - } - modules = resized; - if (!EnumProcessModules(process, modules, buf_bytes, &needed)) { - free(modules); + modules.resize(needed / sizeof(HMODULE)); + buf_bytes = static_cast<DWORD>(modules.size() * sizeof(HMODULE)); + if (!EnumProcessModules(process, modules.data(), buf_bytes, &needed)) { return nullptr; } } DWORD count = needed / sizeof(HMODULE); for (DWORD i = 0; i < count; i++) { void *sym = reinterpret_cast<void *>(GetProcAddress(modules[i], symbol_name)); if (sym != nullptr && sym != exclude) { - free(modules); return reinterpret_cast<void *>(modules[i]); } } - free(modules); return nullptr; }Remember to
#include <vector>in the Windows branch.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/stubs/dynlib.h` around lines 59 - 95, The function dynlib_find_loaded uses manual malloc/realloc/free and early returns; replace that with a std::vector<HMODULE> to own the buffer and avoid leaks: include <vector>, create a vector<HMODULE> modules; set an initial size using buf_bytes / sizeof(HMODULE) (or reserve/resize), pass modules.data() to EnumProcessModules, check needed and if greater than current capacity resize the vector (resize to needed/sizeof(HMODULE)) and call EnumProcessModules again, then iterate modules[0..count) and call GetProcAddress as before; remove all malloc/realloc/free and ensure returned module pointer logic and exclude comparison remain unchanged.tilelang/carver/arch/driver/cuda_driver.py (1)
55-65: Minor:cudart64_11.dllis not a real CUDA runtime library name.CUDA 11.x ships
cudart64_110.dll(which is already in the list); there is nocudart64_11.dll. The extra"11"entry is harmless but dead. Consider also covering future versions or deriving the suffix fromCUDA_PATH/ discovered nvcc to be more robust.♻️ Proposed tweak
- if sys.platform == "win32": - for ver in ("13", "12", "110", "11"): - try: - return ctypes.windll.LoadLibrary(f"cudart64_{ver}.dll") - except OSError: - continue - raise OSError("Cannot find cudart64_*.dll") + if sys.platform == "win32": + for ver in ("13", "12", "110"): + try: + return ctypes.windll.LoadLibrary(f"cudart64_{ver}.dll") + except OSError: + continue + raise OSError("Cannot find cudart64_*.dll")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/carver/arch/driver/cuda_driver.py` around lines 55 - 65, The Windows DLL search in _load_cudart currently includes a dead "11" suffix; remove the "11" entry from the for-loop and replace the hardcoded list with a more robust strategy: either (a) enumerate matching files using a glob like "cudart64_*.dll" and pick the highest/first match, or (b) derive the proper suffix from environment info (CUDA_PATH or nvcc --version) and try that suffix first, falling back to a glob scan; ensure the function still falls back to raising OSError if nothing is found.cmake/FindPipCUDAToolkit.cmake (1)
243-258: DuplicatePATH/PATHEXT/SDK-bin prepends inside_tilelang_activate_msvc_env.Lines 243–250 re-prepend
_tilelang_system_path_candidatestoENV{PATH}despite lines 68–75 already doing this. Similarly, lines 251–255 re-apply the samePATHEXTlogic as lines 77–81, and line 257 re-prepends_tilelang_sdk_binafter line 170 already did. This causesPATHto accumulate duplicate SDK/system entries on every configure and effectively pushes MSVC include paths away from the front as the function runs.Consider removing the duplicated second block (lines 243–258) or, if it was intended to re-assert precedence after the
INCLUDE/env block mutatedPATH, de-duplicate the resulting PATH list explicitly.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/FindPipCUDAToolkit.cmake` around lines 243 - 258, The function _tilelang_activate_msvc_env is prepending _tilelang_system_path_candidates, PATHEXT logic, and _tilelang_sdk_bin twice which causes PATH/PATHEXT duplication; remove the second block (the repeated handling of _tilelang_system_path_candidates, PATHEXT, and _tilelang_sdk_bin) or, if you intended to ensure precedence after the INCLUDE/env modifications, replace the second block with a de-dup step that normalizes ENV{PATH} and ENV{PATHEXT} (split on ';', remove duplicate entries while preserving order, then re-join) so that ENV{PATH} and ENV{PATHEXT} contain each entry only once; update code referencing _tilelang_system_path_candidates, _tilelang_default_pathext, and _tilelang_sdk_bin accordingly.tilelang/env.py (1)
47-53: Standardizesys.platformcomparison for consistency.The code correctly prepends
TL_LIBStoPATHfor child subprocesses on Windows. However, line 47 usessys.platform.startswith("win32")while line 112 in the same file usessys.platform == "win32". Sincesys.platformon Windows is exactly"win32", both are functionally equivalent, but standardizing to== "win32"(orstartswith()) improves readability. Note:os.add_dll_directory()is already invoked separately intilelang/__init__.py(lines 202–203) vialibinfo.get_dll_directories()for Python 3.8+ DLL resolution, so no additional changes needed there.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/env.py` around lines 47 - 53, Change the Windows platform check to use equality for consistency: replace the condition using sys.platform.startswith("win32") with sys.platform == "win32" in the block that builds path_entries/normalized_entries/prepend_entries and updates os.environ["PATH"] so it matches the other usage in this file (and continues to prepend TL_LIBS for subprocesses).tilelang/__init__.py (1)
115-146: Optional: log the swallowedtvm_ffiimport failure at debug level.
_prepend_windows_runtime_pathsis best-effort, which is fine, but silently swallowing thetvm_ffiimport/lookup failure (lines 125-126) makes "why isn't libtvm_ffi on PATH?" hard to diagnose on Windows. A singlelogger.debug(...)call would give users something to grep without changing behavior. Static analysis flags S110/BLE001 here for the same reason.🛠️ Suggested change
try: from tvm_ffi import libinfo as tvm_ffi_libinfo runtime_paths.append(os.path.dirname(tvm_ffi_libinfo.find_libtvm_ffi())) - except Exception: - pass + except Exception as exc: # noqa: BLE001 - best-effort PATH setup at import time + logger.debug("tvm_ffi libinfo discovery failed: %r", exc)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/__init__.py` around lines 115 - 146, In _prepend_windows_runtime_paths(), catch the broad Exception from the tvm_ffi import/lookup and add a logger.debug(...) call that includes the exception details (e.g., referencing tvm_ffi_libinfo.find_libtvm_ffi and the caught exception) so the failure is recorded at debug level without changing behavior; keep the existing pass semantics but emit a single debug log entry when the import/lookup fails to aid diagnostics.tilelang/libinfo.py (1)
8-22: Optional: de-duplicate the Windows CUDA path list withtilelang/__init__.py.The CUDA subdirectory enumeration (
bin,bin/x86_64,lib/x64,nvvm/bin) is copy-pasted into__init__.pylines 182-188. As the list grows, keeping them in lockstep is easy to forget. Consider having__init__.pyconsumeget_dll_directories()(or a new helper that returns just the CUDA dirs) as the single source of truth.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/libinfo.py` around lines 8 - 22, Extract the CUDA subdirectory list into a single helper (e.g., new function get_cuda_dll_subdirs or extend get_dll_directories to expose the CUDA-only list) and have both tilelang/libinfo.py (get_dll_directories) and tilelang/__init__.py call that helper instead of duplicating the literal list; modify get_dll_directories to call the helper when sys.platform startswith "win32" and update __init__.py to import and reuse that helper so the CUDA paths are defined in one place (reference functions: get_dll_directories and the new helper name you add).src/target/stubs/nvrtc.cc (1)
47-71: Consider mirroring the explanatory comment fromcudart.cchere.
cudart.cc(lines 80-81) documents why Windows prefers a freshdynlib_openoverdynlib_find_loadedfirst (some already-loaded modules re-export a subset of cudart symbols). The same inverted ordering appears here but without any comment, which will make the divergence between the Windows and POSIX branches look accidental to future readers.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/stubs/nvrtc.cc` around lines 47 - 71, The Windows branch in TryLoadLibNvrtc uses dynlib_open before dynlib_find_loaded (iterating kLibNvrtcPaths) but lacks the explanatory comment present in cudart.cc; update nvrtc.cc by adding a brief comment above the Windows-first block (near TryLoadLibNvrtc, kLibNvrtcPaths, dynlib_open, dynlib_find_loaded, nvrtcVersion) mirroring cudart.cc's note: explain that on Windows we prefer a fresh dynlib_open because some already-loaded modules may re-export a subset of NVRTC symbols, which would make dynlib_find_loaded return an incomplete handle, so we attempt explicit opens first.tilelang/cache/kernel_cache.py (2)
91-111: Stamp is correct but order-sensitive; worth a brief note.
combined_hashmixes inname.encode()and file bytes in the iteration order oflib_dirs × lib_names. Two installs that share identical library contents but differ inTL_LIBSdirectory ordering will produce different stamps (and therefore different cache namespaces). That's safe (no false hits), just slightly wasteful. SinceTL_LIBSis constructed fromtilelang.env, this is usually stable, but sortinglib_dirs/lib_namesbefore hashing would make the stamp fully deterministic across arbitraryTL_LIBSpermutations. Optional.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/cache/kernel_cache.py` around lines 91 - 111, The stamp generation is order-sensitive because combined_hash mixes name.encode() and file bytes in the nested loop over lib_dirs and lib_names; make it deterministic by iterating over sorted(lib_dirs) and sorted(lib_names) (instead of their original order) when computing combined_hash and appended to found, so the resulting stamp from combined_hash and '+'.join(found) is stable across TL_LIBS permutations; update the loops in kernel_cache.py where combined_hash, lib_dirs, lib_names, and found are used to use sorted(...) accordingly.
413-417: Temp path still uses.sosuffix on Windows.
_safe_write_executablewrites to{pid}_{uuid}.sobeforeos.replace-ing into the finalkernel_lib.soslot. With MSVCfcompilenow wired in on Windows,export_libraryinvokesmsvc_create_shared, andlink.exealso emits a sidecar import lib named after the output base (<...>.lib) next to it inTILELANG_TMP_DIR. The final atomic rename won't clean that sidecar up, so you can accumulate orphaned.lib/.expfiles in the tmp dir over time. Consider using a platform-appropriate suffix on Windows (and/or explicitly cleaning up sidecars), e.g.:♻️ Proposed adjustment
`@classmethod` def _safe_write_executable(cls, executable: Executable, path: str): - temp_path = os.path.join(env.TILELANG_TMP_DIR, f"{os.getpid()}_{uuid.uuid4()}.so") - executable.export_library(temp_path, **cls._get_compile_args()) - os.replace(temp_path, path) + suffix = ".dll" if sys.platform == "win32" else ".so" + temp_path = os.path.join(env.TILELANG_TMP_DIR, f"{os.getpid()}_{uuid.uuid4()}{suffix}") + try: + executable.export_library(temp_path, **cls._get_compile_args()) + os.replace(temp_path, path) + finally: + if sys.platform == "win32": + base, _ = os.path.splitext(temp_path) + for ext in (".lib", ".exp"): + try: + os.remove(base + ext) + except OSError: + pass🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/cache/kernel_cache.py` around lines 413 - 417, _safe_write_executable currently always writes a temp file with a .so suffix which on Windows causes link.exe to emit orphaned .lib/.exp sidecars; change the temp filename generation in _safe_write_executable to use a platform-appropriate suffix (e.g. '.dll' on Windows, '.so' on Unix) and after calling Executable.export_library(temp_path, **cls._get_compile_args()) scan the same temp directory for sidecars with the same stem (e.g. f"{stem}.lib", f"{stem}.exp") and remove them (or move them alongside the temp file) before calling os.replace(temp_path, path) so no orphaned sidecars remain in env.TILELANG_TMP_DIR. Ensure you still use uuid/os.getpid() to make the temp name unique and reference env.TILELANG_TMP_DIR, _safe_write_executable, Executable.export_library and os.replace in your change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cmake/generate_windows_import_lib.py`:
- Around line 10-58: The EXPORT_RE currently captures the last token, which
loses the real export name for forwarded or data exports; update the regex
EXPORT_RE to capture the exported symbol name as a separate group and an
optional trailing forwarded/data suffix (e.g.
r"^\s+\d+\s+[0-9A-Fa-f]+\s+[0-9A-Fa-f]+\s+(\S+?)(?:\s*\((forwarded to
.*)\))?(?:,DATA)?\s*$"), then in main where match is used extract the real
symbol from match.group(1) and detect if the original line included ",DATA" or a
forwarded suffix (via match.groups or by inspecting the original line) so you
can (a) skip or correctly handle forwarded exports (emit the original name, not
the forwarded target) and (b) append " DATA" to that export when writing the
.def; update the exports list construction and the def_path.write_text call to
include the " DATA" marker for data exports or otherwise document/raise when
encountering forwarded/data exports if you choose not to support them.
In `@src/target/stubs/nvrtc.cc`:
- Around line 73-76: The error message in nvrtc.cc currently tells users to
"Make sure PyTorch with CUDA is installed" which is misleading for Windows/pip
or standalone CUDA NVRTC installs; update the fprintf in the nvrtc
symbols-not-found branch (the fprintf(stderr, ...) followed by abort()) to use
broader wording similar to cudart.cc—e.g., mention that CUDA NVRTC symbols were
not found and suggest installing a CUDA-enabled PyTorch or a
system/properly-installed CUDA toolkit (or the nvidia-cuda-nvrtc wheel) and/or
verifying CUDA installation and PATH/LD_LIBRARY_PATH; keep the same stderr +
abort() behavior but replace the hardcoded PyTorch-only guidance with the
inclusive message.
In `@src/transform/layout_reducer.cc`:
- Around line 131-139: The code currently ICHECKs that
p_result->annotations.Get(attr::kLayoutMap) can be cast to Map<Var, Layout>,
which will abort when the stored map uses non-Var keys (e.g., Map<Buffer,
Layout> written by LayoutInferencer::VisitStmt_(BlockNode)); change the logic in
ReducerLayoutAnnotator so after obtaining layout_map_ref and maybe_layout_map
you do not ICHECK on maybe_layout_map.has_value(), but instead treat a failed
cast as "nothing to merge" and skip/early-return: if
maybe_layout_map.has_value() assign layout_map = maybe_layout_map.value() else
leave layout_map empty and continue; keep the presence check of attr::kLayoutMap
but avoid aborting on wrong-keyed maps to safely handle Map<Buffer, Layout>.
In `@tilelang/contrib/nvcc.py`:
- Around line 104-105: The added Windows-only compiler option uses the
MSVC-specific flag "/Zc:preprocessor" which requires MSVC 19.25+; update
tilelang/contrib/nvcc.py so you only append ["-Xcompiler", "/Zc:preprocessor"]
to cmd when the host toolchain supports it. Concretely: detect Windows (os.name
== "nt") as before, then probe cl.exe version (e.g., run cl.exe to read its
reported version string or use an environment/toolchain API) and only add the
flag when the parsed major/minor version is >= 19.25 (or else skip the flag);
alternatively add a clear comment and a fallback path that documents the minimum
MSVC requirement if you prefer not to probe. Ensure you reference and modify the
code that builds cmd so the unsafe flag is conditionally applied.
---
Outside diff comments:
In `@tilelang/jit/adapter/libgen.py`:
- Around line 146-154: The subprocess.run call that invokes NVCC must be passed
the proper MSVC/NVCC environment so cl.exe and Windows SDK headers are visible;
obtain the environment via the existing helpers
(get_nvcc_subprocess_env()/get_msvc_subprocess_env()) and pass it as
env=compiler_env to subprocess.run(command, timeout=timeout, env=compiler_env)
(mirror tilelang/contrib/nvcc.py::compile_cuda). Only import those helpers for
the CUDA branch if needed and ensure the same env is used for the return-code
check (ret.returncode) and error messages.
---
Duplicate comments:
In `@src/target/codegen_cuda.cc`:
- Around line 3591-3595: The tan operator branch currently uses CUDAFastMath
(producing __tanf for float32) which is too imprecise; change the branch
handling tl::__tan() to use the CUDAFastMathTan specialization instead of
CUDAFastMath so the emitted function respects the higher-precision tan
implementation; update the code that builds func_name (when
op->op.same_as(tl::__tan())) to call CUDAFastMathTan(op->dtype, "tan") rather
than CUDAFastMath, keeping the existing usage of PrintExpr(op->args[0]) and
returning true.
---
Nitpick comments:
In `@cmake/FindPipCUDAToolkit.cmake`:
- Around line 243-258: The function _tilelang_activate_msvc_env is prepending
_tilelang_system_path_candidates, PATHEXT logic, and _tilelang_sdk_bin twice
which causes PATH/PATHEXT duplication; remove the second block (the repeated
handling of _tilelang_system_path_candidates, PATHEXT, and _tilelang_sdk_bin)
or, if you intended to ensure precedence after the INCLUDE/env modifications,
replace the second block with a de-dup step that normalizes ENV{PATH} and
ENV{PATHEXT} (split on ';', remove duplicate entries while preserving order,
then re-join) so that ENV{PATH} and ENV{PATHEXT} contain each entry only once;
update code referencing _tilelang_system_path_candidates,
_tilelang_default_pathext, and _tilelang_sdk_bin accordingly.
In `@src/layout/gemm_layouts.cc`:
- Around line 14-20: Macro TILELANG_COMPILER_UNREACHABLE is implemented
correctly for MSVC/GCC/Clang with a safe no-op fallback, but consider hoisting
it to a shared header to avoid duplication; move the
TILELANG_COMPILER_UNREACHABLE definition out of src/layout/gemm_layouts.cc into
a common header (e.g., tl_compiler.h) and replace occurrences in this TU and any
others with an `#include` of that header so callers like the sites that follow
ICHECK(0) can reuse the same, single definition.
In `@src/target/stubs/dynlib.h`:
- Around line 59-95: The function dynlib_find_loaded uses manual
malloc/realloc/free and early returns; replace that with a std::vector<HMODULE>
to own the buffer and avoid leaks: include <vector>, create a vector<HMODULE>
modules; set an initial size using buf_bytes / sizeof(HMODULE) (or
reserve/resize), pass modules.data() to EnumProcessModules, check needed and if
greater than current capacity resize the vector (resize to
needed/sizeof(HMODULE)) and call EnumProcessModules again, then iterate
modules[0..count) and call GetProcAddress as before; remove all
malloc/realloc/free and ensure returned module pointer logic and exclude
comparison remain unchanged.
In `@src/target/stubs/nvrtc.cc`:
- Around line 47-71: The Windows branch in TryLoadLibNvrtc uses dynlib_open
before dynlib_find_loaded (iterating kLibNvrtcPaths) but lacks the explanatory
comment present in cudart.cc; update nvrtc.cc by adding a brief comment above
the Windows-first block (near TryLoadLibNvrtc, kLibNvrtcPaths, dynlib_open,
dynlib_find_loaded, nvrtcVersion) mirroring cudart.cc's note: explain that on
Windows we prefer a fresh dynlib_open because some already-loaded modules may
re-export a subset of NVRTC symbols, which would make dynlib_find_loaded return
an incomplete handle, so we attempt explicit opens first.
In `@tilelang/__init__.py`:
- Around line 115-146: In _prepend_windows_runtime_paths(), catch the broad
Exception from the tvm_ffi import/lookup and add a logger.debug(...) call that
includes the exception details (e.g., referencing
tvm_ffi_libinfo.find_libtvm_ffi and the caught exception) so the failure is
recorded at debug level without changing behavior; keep the existing pass
semantics but emit a single debug log entry when the import/lookup fails to aid
diagnostics.
In `@tilelang/cache/kernel_cache.py`:
- Around line 91-111: The stamp generation is order-sensitive because
combined_hash mixes name.encode() and file bytes in the nested loop over
lib_dirs and lib_names; make it deterministic by iterating over sorted(lib_dirs)
and sorted(lib_names) (instead of their original order) when computing
combined_hash and appended to found, so the resulting stamp from combined_hash
and '+'.join(found) is stable across TL_LIBS permutations; update the loops in
kernel_cache.py where combined_hash, lib_dirs, lib_names, and found are used to
use sorted(...) accordingly.
- Around line 413-417: _safe_write_executable currently always writes a temp
file with a .so suffix which on Windows causes link.exe to emit orphaned
.lib/.exp sidecars; change the temp filename generation in
_safe_write_executable to use a platform-appropriate suffix (e.g. '.dll' on
Windows, '.so' on Unix) and after calling Executable.export_library(temp_path,
**cls._get_compile_args()) scan the same temp directory for sidecars with the
same stem (e.g. f"{stem}.lib", f"{stem}.exp") and remove them (or move them
alongside the temp file) before calling os.replace(temp_path, path) so no
orphaned sidecars remain in env.TILELANG_TMP_DIR. Ensure you still use
uuid/os.getpid() to make the temp name unique and reference
env.TILELANG_TMP_DIR, _safe_write_executable, Executable.export_library and
os.replace in your change.
In `@tilelang/carver/arch/driver/cuda_driver.py`:
- Around line 55-65: The Windows DLL search in _load_cudart currently includes a
dead "11" suffix; remove the "11" entry from the for-loop and replace the
hardcoded list with a more robust strategy: either (a) enumerate matching files
using a glob like "cudart64_*.dll" and pick the highest/first match, or (b)
derive the proper suffix from environment info (CUDA_PATH or nvcc --version) and
try that suffix first, falling back to a glob scan; ensure the function still
falls back to raising OSError if nothing is found.
In `@tilelang/contrib/msvc.py`:
- Around line 152-161: The _normalize_option function currently drops -std=...
and -g options silently; update _normalize_option to translate -std=<ver> into
MSVC's /std:<ver> (e.g., -std=c++17 -> /std:c++17) and map -g to an appropriate
MSVC debug flag (e.g., /Zi or /Z7) instead of returning None; locate and edit
the _normalize_option function in msvc.py and ensure it parses the suffix after
"-std=" and returns "/std:" + suffix, and replace the -g branch to return "/Zi"
(or "/Z7" if you prefer full debug info) so callers like the TVM/TileLang JIT
FFI glue that pass "-std=c++17" get equivalent MSVC flags emitted.
In `@tilelang/env.py`:
- Around line 47-53: Change the Windows platform check to use equality for
consistency: replace the condition using sys.platform.startswith("win32") with
sys.platform == "win32" in the block that builds
path_entries/normalized_entries/prepend_entries and updates os.environ["PATH"]
so it matches the other usage in this file (and continues to prepend TL_LIBS for
subprocesses).
In `@tilelang/jit/adapter/libgen.py`:
- Around line 63-66: The code currently builds library paths by calling
src.name.replace(".cu", _lib_ext) which can accidentally replace ".cu" anywhere
in the temp path; update the logic to replace only the file extension using
os.path.splitext on src.name (e.g., use os.path.splitext(src.name)[0] +
_lib_ext) for the .cu→.dll/.so case (around the src =
tempfile.NamedTemporaryFile(...) / libpath = ... block) and apply the same
change to the other occurrences where .cpp is converted to .so (the .cpp → .so
conversions later in the file) so extension replacement is robust on Windows and
other platforms.
In `@tilelang/libinfo.py`:
- Around line 8-22: Extract the CUDA subdirectory list into a single helper
(e.g., new function get_cuda_dll_subdirs or extend get_dll_directories to expose
the CUDA-only list) and have both tilelang/libinfo.py (get_dll_directories) and
tilelang/__init__.py call that helper instead of duplicating the literal list;
modify get_dll_directories to call the helper when sys.platform startswith
"win32" and update __init__.py to import and reuse that helper so the CUDA paths
are defined in one place (reference functions: get_dll_directories and the new
helper name you add).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4ee9b92c-2c4d-4b61-9ca4-4fcf36d442e7
📒 Files selected for processing (35)
.gitattributes.gitignoreCMakeLists.txtcmake/FindPipCUDAToolkit.cmakecmake/find_pip_cuda.pycmake/generate_windows_import_lib.pycmake/pypi-z3/FindZ3.cmakedocs/get_started/Installation.mdpyproject.tomlscripts/windows/uv-sync.cmdsrc/layout/gemm_layouts.ccsrc/op/operator.hsrc/op/parallel.ccsrc/runtime/logging.ccsrc/target/codegen_cuda.ccsrc/target/codegen_cuda.hsrc/target/stubs/cuda.ccsrc/target/stubs/cuda.hsrc/target/stubs/cudart.ccsrc/target/stubs/dynlib.hsrc/target/stubs/nvrtc.ccsrc/transform/common/pipeline_utils.hsrc/transform/layout_inference.ccsrc/transform/layout_reducer.cctilelang/__init__.pytilelang/cache/kernel_cache.pytilelang/carver/arch/driver/cuda_driver.pytilelang/contrib/cc.pytilelang/contrib/msvc.pytilelang/contrib/nvcc.pytilelang/env.pytilelang/jit/adapter/libgen.pytilelang/jit/adapter/tvm_ffi.pytilelang/jit/adapter/wrapper.pytilelang/libinfo.py
✅ Files skipped from review due to trivial changes (3)
- .gitignore
- .gitattributes
- cmake/find_pip_cuda.py
🚧 Files skipped from review as they are similar to previous changes (5)
- src/transform/common/pipeline_utils.h
- tilelang/jit/adapter/tvm_ffi.py
- cmake/pypi-z3/FindZ3.cmake
- pyproject.toml
- src/runtime/logging.cc
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/runtime/logging.cc (1)
3-9:⚠️ Potential issue | 🟡 MinorAdd missing
#include <sstream>forstd::istringstreamandstd::stringstreamusage.
ParseSpec(line 85) usesstd::istringstream, andVLogContext::str()(line 150) usesstd::stringstream, but<sstream>is not explicitly included. While this currently compiles via transitive inclusion (likely from<tvm/runtime/logging.h>), relying on transitive includes is fragile and can break across different toolchains—particularly relevant for the MSVC compatibility work in this PR.Proposed fix
`#include` <algorithm> `#include` <cstdlib> `#include` <ctime> `#include` <iomanip> `#include` <iostream> +#include <sstream> `#include` <string>🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/runtime/logging.cc` around lines 3 - 9, Add an explicit include for <sstream> because ParseSpec uses std::istringstream and VLogContext::str() uses std::stringstream; update the includes at the top of src/runtime/logging.cc to add `#include` <sstream> so the file no longer relies on transitive includes and compiles reliably across toolchains (references: ParseSpec and VLogContext::str()).tilelang/jit/adapter/libgen.py (1)
146-154:⚠️ Potential issue | 🟠 MajorWindows: NVCC subprocess is launched without the MSVC host-compiler environment.
subprocess.run(command, ...)is invoked without an explicitenv, so on Windows the child process inherits whateverPATHthe current Python process has. Unless the user started Python from a Visual Studio developer prompt,cl.exewill not be resolvable and NVCC will fail with "Cannot find compiler 'cl.exe' in PATH" — even thoughtilelang.contrib.nvcc.compile_cudacarefully threadsget_nvcc_subprocess_env()through for exactly this reason.Mirror the pattern from
tilelang/contrib/nvcc.pyso the MSVC toolchain (located viaVsDevCmd.bat) is wired up for this NVCC invocation as well.🔧 Proposed fix
+ compiler_env = None + if is_cuda_target(target) and sys.platform == "win32": + from tilelang.contrib.nvcc import get_nvcc_subprocess_env + compiler_env = get_nvcc_subprocess_env() + try: if verbose: print(f"compile_lib compilation command: {' '.join(command)}") - ret = subprocess.run(command, timeout=timeout) + ret = subprocess.run(command, timeout=timeout, env=compiler_env) except Exception as e: raise RuntimeError(f"Compile kernel failed because of {e}") from e🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/jit/adapter/libgen.py` around lines 146 - 154, The NVCC invocation in the try block (the subprocess.run call in libgen.py) does not pass the MSVC host-compiler environment; obtain the proper env by calling the same helper used in tilelang.contrib.nvcc (get_nvcc_subprocess_env) and pass it to subprocess.run as env=env so NVCC can locate cl.exe on Windows; update the try block around the compile command (the code that prints "compile_lib compilation command" and assigns ret = subprocess.run(...)) to compute env = get_nvcc_subprocess_env() (importing that helper if needed) and call subprocess.run(command, timeout=timeout, env=env).
♻️ Duplicate comments (4)
src/target/codegen_cuda.cc (1)
3591-3595:⚠️ Potential issue | 🟠 Major
tl::__tan()now usesCUDAFastMathand reintroduces the__tanfprecision regression.This branch previously used
CUDAFastMathTan, which deliberately falls back totanfforfloat32(see the comment on lines 173–174 inCUDAFastMathTan::operator()). WithCUDAFastMathhere,tan(float32)will emit__tanf, which was found to deviate too much from the reference implementation — that is exactly whyCUDAFastMathTanexists as a separate functor.🛠️ Suggested fix
} else if (op->op.same_as(tl::__tan())) { - CUDAFastMath math_func; + CUDAFastMathTan math_func; std::string func_name = math_func(op->dtype, "tan"); os << func_name << "(" << PrintExpr(op->args[0]) << ")"; return true;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/codegen_cuda.cc` around lines 3591 - 3595, The tl::__tan() branch incorrectly uses CUDAFastMath (which emits __tanf for float32) causing a precision regression; replace CUDAFastMath with the specialized CUDAFastMathTan functor in the tl::__tan() branch so the existing CUDAFastMathTan::operator() behavior (falling back to tanf for float32 per its comment) is preserved — update the instantiation and func_name generation to use CUDAFastMathTan instead of CUDAFastMath and keep the PrintExpr(op->args[0]) usage and return true logic intact.src/target/stubs/nvrtc.cc (1)
73-76:⚠️ Potential issue | 🟡 MinorError message still assumes PyTorch-with-CUDA is the only install path.
With this PR adding Windows and pip-installed CUDA toolkit support (nvidia-cuda-nvrtc wheels, standalone toolkits), reaching
abort()here and telling users to "install PyTorch with CUDA" is misleading on those paths. Consider the broadened wording suggested on the previous review round.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/stubs/nvrtc.cc` around lines 73 - 76, The error message in nvrtc.cc (the fprintf and abort path) incorrectly assumes PyTorch-with-CUDA is the only install path; update the fprintf text in the nvrtc symbols-not-found branch to use broader wording that mentions missing NVRTC/CUDA installation or missing nvrtc symbols (e.g., "NVRTC symbols not found. Ensure a CUDA toolkit or compatible NVRTC runtime is installed and accessible (for example via PyTorch with CUDA, pip-installed nvidia-cuda-nvrtc wheels, or a system CUDA toolkit)"), so the message covers Windows/pip standalone toolkits and other install methods instead of only suggesting "install PyTorch with CUDA". Ensure the change is made where fprintf(...) and the following abort() are called.src/target/stubs/cudart.cc (1)
107-110:⚠️ Potential issue | 🟡 MinorError message still only references PyTorch.
Same issue as previously flagged: on Windows and with pip-installed CUDA toolkit support added in this PR, this
abort()path is also reachable without PyTorch, so the "install PyTorch with CUDA" hint is misleading. Broaden the wording as suggested in the previous review.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/stubs/cudart.cc` around lines 107 - 110, The error path prints a message that only mentions PyTorch but abort() can be reached even when PyTorch is not involved; update the fprintf call in cudart.cc (the call that prints "TileLang Error: cudart symbols not found..." before abort()) to a broader, accurate hint — e.g., instruct the user to ensure the CUDA runtime/toolkit is available (mentioning options like installing PyTorch with CUDA or the CUDA toolkit/pip CUDA support) and include the same context string "cudart symbols not found" so the location is clear.src/target/stubs/dynlib.h (1)
104-107:⚠️ Potential issue | 🔴 Critical
_GNU_SOURCEis defined too late for TUs that include other system headers beforedynlib.h.Same concern as the previous review round: because
<features.h>is latched by the first system header that includes it, any TU that transitively includes<stdlib.h>(etc.) beforedynlib.hwill seeRTLD_DEFAULT/RTLD_NEXTundeclared even though they're used indynlib_find_loaded(). The per-man-page rule is explicit: "The _GNU_SOURCE feature test macro must be defined in order to obtain the definitions of RTLD_DEFAULT and RTLD_NEXT from <dlfcn.h>". Prefer defining_GNU_SOURCEviatarget_compile_definitions(... PRIVATE _GNU_SOURCE=1)on the stub targets in CMake so it's active before any system header is processed.
🧹 Nitpick comments (5)
src/runtime/logging.cc (1)
94-129: Inconsistentreturn settings;afterLOG(FATAL).
LogFatalImplis[[noreturn]](line 62), so thereturn settings;statements following theLOG(FATAL)branches at lines 110, 115, and 123 are unreachable. The matching branch at lines 100–102 correctly omits the redundant return. Consider dropping the dead returns for consistency (or, conversely, guard with a comment if some callers catch the thrownInternalErrorand you want a deterministic fallback — in which case the behavior should be the same across all four FATAL branches, including the "empty filename" one).🧹 Proposed cleanup
if (!std::getline(spec_stream, level, ',')) { LOG(FATAL) << "TVM_LOG_DEBUG ill-formed at position " << tell_pos(level) << ": expecting \"=<level>\" after \"" << name << "\""; - return settings; } if (level.empty()) { LOG(FATAL) << "TVM_LOG_DEBUG ill-formed at position " << tell_pos(level) << ": empty level after \"" << name << "\""; - return settings; } char *end_of_level = nullptr; int level_val = static_cast<int>(strtol(level.c_str(), &end_of_level, 10)); if (end_of_level != level.c_str() + level.size()) { LOG(FATAL) << "TVM_LOG_DEBUG ill-formed at position " << tell_pos(level) << ": invalid level: \"" << level << "\""; - return settings; }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/runtime/logging.cc` around lines 94 - 129, The LOG(FATAL) calls inside the while loop invoke LogFatalImpl which is [[noreturn]], so remove the unreachable "return settings;" statements that follow the LOG(FATAL) branches in the loop (the returns after the "expecting \"=<level>\"", "empty level after", and "invalid level" checks) to match the earlier branch that omits a return; alternatively, if you prefer explicit clarity, replace those returns with a brief comment referencing LogFatalImpl's [[noreturn]] behavior (e.g., "// LogFatalImpl is [[noreturn]] - no return needed") so all FATAL branches are handled consistently.tilelang/libinfo.py (1)
49-56: Minor:for…elsehere is functionally equivalent to placing the raise unindented.Since the inner loop uses
return(notbreak), the outerforalways completes "normally" when no match is found, so theelseclause is entered on every miss. This is correct, but theelseadds no semantic distinction from a plain trailingraise, and it's one of the more commonly misread Python constructs. Optional cleanup for readability.♻️ Suggestion
for lib_root in TL_LIBS: for lib_name in lib_names: lib_dll_path = os.path.join(lib_root, lib_name) if os.path.exists(lib_dll_path) and os.path.isfile(lib_dll_path): return lib_dll_path - else: - message = f"Cannot find libraries: {', '.join(lib_names)}\n" + "List of candidates:\n" + "\n".join(TL_LIBS) - raise RuntimeError(message) + message = ( + f"Cannot find libraries: {', '.join(lib_names)}\n" + "List of candidates:\n" + "\n".join(TL_LIBS) + ) + raise RuntimeError(message)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/libinfo.py` around lines 49 - 56, The for...else after iterating TL_LIBS and lib_names is unnecessary and potentially confusing because the inner loop returns on a match; replace the for...else by unindenting the raise so it follows the nested loops (i.e., after the two loops complete) — locate the loops over TL_LIBS and lib_names and move the RuntimeError raise (which builds the message from lib_names and TL_LIBS and references lib_dll_path/lib_root) out of the else block and place it directly after the loops.tilelang/__init__.py (2)
180-190: Consider delegating CUDA DLL directory enumeration tolibinfo.get_dll_directories()to remove duplication.The CUDA directory list here (
cuda_home,bin,bin/x86_64,lib/x64,nvvm/bin) duplicates the logic intilelang/libinfo.py::get_dll_directories()(lines 8–22), which is also invoked later in_load_tile_lang_lib(). Any future change (e.g., addinglibnvvp, adjusting for CUDA 13 layouts) would need to be made in two places. Importinglibinfoearlier and reusing its output would keep the two callsites in sync and avoid maintaining parallel path lists.♻️ Sketch
if sys.platform.startswith("win32") and sys.version_info >= (3, 8): _prepend_windows_runtime_paths() + from . import libinfo as _libinfo_for_dll _dll_handles = [] - cuda_home = getattr(env, "CUDA_HOME", "") - if cuda_home: - for path in ( - cuda_home, - os.path.join(cuda_home, "bin"), - os.path.join(cuda_home, "bin", "x86_64"), - os.path.join(cuda_home, "lib", "x64"), - os.path.join(cuda_home, "nvvm", "bin"), - ): - if os.path.isdir(path): - _dll_handles.append(os.add_dll_directory(path)) + for path in _libinfo_for_dll.get_dll_directories(): + _dll_handles.append(os.add_dll_directory(path))Note that
_load_tile_lang_lib()at lines 201–203 callsos.add_dll_directory(path)without retaining the returned cookies, so those entries are eligible for garbage collection as soon as the loop exits; the handles kept in_dll_handleshere are what actually keep the CUDA paths alive. Consolidating the enumeration would also make that ownership explicit.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/__init__.py` around lines 180 - 190, Replace the duplicated CUDA path list with a call to libinfo.get_dll_directories(): import tilelang.libinfo (or the module where get_dll_directories lives), call get_dll_directories() instead of constructing the tuple of cuda_home paths, and iterate its returned paths appending each os.add_dll_directory return value to the existing _dll_handles list; also update _load_tile_lang_lib() so it does not call os.add_dll_directory(path) without retaining the returned handle—reuse or extend _dll_handles there to maintain ownership.
115-147: Silenttry/except Exception: passontvm_ffipath lookup will mask real failures.When
from tvm_ffi import libinfoorfind_libtvm_ffi()raises something unexpected (e.g., a packaging/version skew that installs a brokentvm_ffi), Windows users will silently end up withruntime_pathsmissing thelibtvm_ffidirectory and no indication of why. A debug-level log would preserve the non-fatal behavior while giving a breadcrumb when DLL loading later fails.🔧 Suggestion
try: from tvm_ffi import libinfo as tvm_ffi_libinfo runtime_paths.append(os.path.dirname(tvm_ffi_libinfo.find_libtvm_ffi())) - except Exception: - pass + except Exception as exc: # noqa: BLE001 + logger.debug("tvm_ffi libinfo unavailable for PATH setup: %r", exc)As per static analysis hint Ruff
S110/BLE001on lines 125–126.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/__init__.py` around lines 115 - 147, The silent except in _prepend_windows_runtime_paths around "from tvm_ffi import libinfo as tvm_ffi_libinfo" and find_libtvm_ffi() should be changed to capture and log the exception at debug level so failures aren't masked: add a module logger via logging.getLogger(__name__) and replace "except Exception: pass" with "except Exception as e:" and call logger.debug with a descriptive message (including the exception and exc_info=True) referencing tvm_ffi and tvm_ffi_libinfo.find_libtvm_ffi to help trace why the runtime path lookup failed while preserving the non-fatal behavior.tilelang/contrib/msvc.py (1)
288-290: Addget_target_tripleattribute to match thefcompilecontract used bytilelang/contrib/cc.py.The
cc.py::create_sharedfunction sets bothoutput_formatandget_target_tripleattributes. While TVM'scross_compiler()defensively checks for this attribute viahasattr()before accessing it, and the MSVCcreate_sharedis also sometimes wrapped through that path, matching this interface pattern ensures consistency and prevents future issues if direct attribute access occurs without defensive checks.Optional addition
create_shared.output_format = "dll" +create_shared.get_target_triple = lambda: None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` around lines 288 - 290, The create_shared callable in the MSVC backend lacks the get_target_triple attribute expected by the fcompile contract; add a get_target_triple attribute on create_shared (matching the pattern used in cc.py::create_shared) that returns the target triple from the passed config/context so consumers can call create_shared.get_target_triple(cfg) consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@scripts/windows/uv-sync.cmd`:
- Around line 7-15: The argument forwarding in the :parse_args block currently
builds SYNC_ARGS using %1 which loses original quoting and breaks values with
spaces; update the concatenation to append the quoted, unescaped token such as
"%~1" instead of %1 so SYNC_ARGS preserves spaces and quotes when passed to the
downstream `uv sync` command, leaving the rest of the parsing (NO_NVCC handling
and shift/goto parse_args) unchanged.
In `@src/target/stubs/dynlib.h`:
- Around line 122-136: dynlib_find_loaded currently returns the raw handle
(NULL) when dlsym(RTLD_DEFAULT, ...) succeeds on glibc, which callers interpret
as "not found"; change dynlib_find_loaded to return a non-null sentinel (e.g. a
unique pointer value like SENTINEL_RTLD_DEFAULT) when dlsym(RTLD_DEFAULT, ...)
finds the symbol (and likewise a different sentinel for RTLD_NEXT if you check
it), and update dynlib_sym to translate that sentinel back to RTLD_DEFAULT (or
RTLD_NEXT) when resolving the symbol so the real dlsym is invoked with the
correct special handle; reference dynlib_find_loaded and dynlib_sym to locate
where to insert the sentinel logic and the translation logic.
In `@tilelang/carver/arch/driver/cuda_driver.py`:
- Around line 55-64: Remove the bogus "11" candidate and stop using
ctypes.windll for loading the CUDA runtime in _load_cudart(); update the Windows
candidate tuple (e.g., keep "13", "12", "110", etc. but remove "11") and call
ctypes.cdll.LoadLibrary(...) instead of ctypes.windll.LoadLibrary so the loader
uses the cdecl convention and matches non-Windows behavior; ensure the function
still falls back to raising OSError if no candidate loads and retains the
non-Windows path that uses ctypes.cdll.LoadLibrary("libcudart.so").
In `@tilelang/jit/adapter/libgen.py`:
- Around line 64-66: The temp file path is being mangled by
src.name.replace(".cu", _lib_ext) which can replace ".cu" occurrences in parent
directories; change to using os.path.splitext on src.name to swap the extension
(e.g., base, _ = os.path.splitext(src.name); libpath = base + _lib_ext) so only
the filename extension is replaced. Update the same pattern in the hip and cpu
branches (the other replace usages around the sections handling those targets)
and keep references to src (NamedTemporaryFile), libpath, and _lib_ext when
making the change.
---
Outside diff comments:
In `@src/runtime/logging.cc`:
- Around line 3-9: Add an explicit include for <sstream> because ParseSpec uses
std::istringstream and VLogContext::str() uses std::stringstream; update the
includes at the top of src/runtime/logging.cc to add `#include` <sstream> so the
file no longer relies on transitive includes and compiles reliably across
toolchains (references: ParseSpec and VLogContext::str()).
In `@tilelang/jit/adapter/libgen.py`:
- Around line 146-154: The NVCC invocation in the try block (the subprocess.run
call in libgen.py) does not pass the MSVC host-compiler environment; obtain the
proper env by calling the same helper used in tilelang.contrib.nvcc
(get_nvcc_subprocess_env) and pass it to subprocess.run as env=env so NVCC can
locate cl.exe on Windows; update the try block around the compile command (the
code that prints "compile_lib compilation command" and assigns ret =
subprocess.run(...)) to compute env = get_nvcc_subprocess_env() (importing that
helper if needed) and call subprocess.run(command, timeout=timeout, env=env).
---
Duplicate comments:
In `@src/target/codegen_cuda.cc`:
- Around line 3591-3595: The tl::__tan() branch incorrectly uses CUDAFastMath
(which emits __tanf for float32) causing a precision regression; replace
CUDAFastMath with the specialized CUDAFastMathTan functor in the tl::__tan()
branch so the existing CUDAFastMathTan::operator() behavior (falling back to
tanf for float32 per its comment) is preserved — update the instantiation and
func_name generation to use CUDAFastMathTan instead of CUDAFastMath and keep the
PrintExpr(op->args[0]) usage and return true logic intact.
In `@src/target/stubs/cudart.cc`:
- Around line 107-110: The error path prints a message that only mentions
PyTorch but abort() can be reached even when PyTorch is not involved; update the
fprintf call in cudart.cc (the call that prints "TileLang Error: cudart symbols
not found..." before abort()) to a broader, accurate hint — e.g., instruct the
user to ensure the CUDA runtime/toolkit is available (mentioning options like
installing PyTorch with CUDA or the CUDA toolkit/pip CUDA support) and include
the same context string "cudart symbols not found" so the location is clear.
In `@src/target/stubs/nvrtc.cc`:
- Around line 73-76: The error message in nvrtc.cc (the fprintf and abort path)
incorrectly assumes PyTorch-with-CUDA is the only install path; update the
fprintf text in the nvrtc symbols-not-found branch to use broader wording that
mentions missing NVRTC/CUDA installation or missing nvrtc symbols (e.g., "NVRTC
symbols not found. Ensure a CUDA toolkit or compatible NVRTC runtime is
installed and accessible (for example via PyTorch with CUDA, pip-installed
nvidia-cuda-nvrtc wheels, or a system CUDA toolkit)"), so the message covers
Windows/pip standalone toolkits and other install methods instead of only
suggesting "install PyTorch with CUDA". Ensure the change is made where
fprintf(...) and the following abort() are called.
---
Nitpick comments:
In `@src/runtime/logging.cc`:
- Around line 94-129: The LOG(FATAL) calls inside the while loop invoke
LogFatalImpl which is [[noreturn]], so remove the unreachable "return settings;"
statements that follow the LOG(FATAL) branches in the loop (the returns after
the "expecting \"=<level>\"", "empty level after", and "invalid level" checks)
to match the earlier branch that omits a return; alternatively, if you prefer
explicit clarity, replace those returns with a brief comment referencing
LogFatalImpl's [[noreturn]] behavior (e.g., "// LogFatalImpl is [[noreturn]] -
no return needed") so all FATAL branches are handled consistently.
In `@tilelang/__init__.py`:
- Around line 180-190: Replace the duplicated CUDA path list with a call to
libinfo.get_dll_directories(): import tilelang.libinfo (or the module where
get_dll_directories lives), call get_dll_directories() instead of constructing
the tuple of cuda_home paths, and iterate its returned paths appending each
os.add_dll_directory return value to the existing _dll_handles list; also update
_load_tile_lang_lib() so it does not call os.add_dll_directory(path) without
retaining the returned handle—reuse or extend _dll_handles there to maintain
ownership.
- Around line 115-147: The silent except in _prepend_windows_runtime_paths
around "from tvm_ffi import libinfo as tvm_ffi_libinfo" and find_libtvm_ffi()
should be changed to capture and log the exception at debug level so failures
aren't masked: add a module logger via logging.getLogger(__name__) and replace
"except Exception: pass" with "except Exception as e:" and call logger.debug
with a descriptive message (including the exception and exc_info=True)
referencing tvm_ffi and tvm_ffi_libinfo.find_libtvm_ffi to help trace why the
runtime path lookup failed while preserving the non-fatal behavior.
In `@tilelang/contrib/msvc.py`:
- Around line 288-290: The create_shared callable in the MSVC backend lacks the
get_target_triple attribute expected by the fcompile contract; add a
get_target_triple attribute on create_shared (matching the pattern used in
cc.py::create_shared) that returns the target triple from the passed
config/context so consumers can call create_shared.get_target_triple(cfg)
consistently.
In `@tilelang/libinfo.py`:
- Around line 49-56: The for...else after iterating TL_LIBS and lib_names is
unnecessary and potentially confusing because the inner loop returns on a match;
replace the for...else by unindenting the raise so it follows the nested loops
(i.e., after the two loops complete) — locate the loops over TL_LIBS and
lib_names and move the RuntimeError raise (which builds the message from
lib_names and TL_LIBS and references lib_dll_path/lib_root) out of the else
block and place it directly after the loops.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 306edc71-c1f8-489a-8454-1a75b1561156
📒 Files selected for processing (32)
.gitattributes.gitignoreCMakeLists.txtcmake/FindPipCUDAToolkit.cmakecmake/find_pip_cuda.pycmake/generate_windows_import_lib.pycmake/pypi-z3/FindZ3.cmakedocs/get_started/Installation.mdpyproject.tomlscripts/windows/uv-sync.cmdsrc/layout/gemm_layouts.ccsrc/op/operator.hsrc/runtime/logging.ccsrc/target/codegen_cuda.ccsrc/target/codegen_cuda.hsrc/target/stubs/cuda.ccsrc/target/stubs/cuda.hsrc/target/stubs/cudart.ccsrc/target/stubs/dynlib.hsrc/target/stubs/nvrtc.ccsrc/transform/layout_reducer.cctilelang/__init__.pytilelang/cache/kernel_cache.pytilelang/carver/arch/driver/cuda_driver.pytilelang/contrib/cc.pytilelang/contrib/msvc.pytilelang/contrib/nvcc.pytilelang/env.pytilelang/jit/adapter/libgen.pytilelang/jit/adapter/tvm_ffi.pytilelang/jit/adapter/wrapper.pytilelang/libinfo.py
✅ Files skipped from review due to trivial changes (5)
- .gitattributes
- .gitignore
- cmake/pypi-z3/FindZ3.cmake
- pyproject.toml
- CMakeLists.txt
🚧 Files skipped from review as they are similar to previous changes (4)
- src/target/codegen_cuda.h
- tilelang/jit/adapter/wrapper.py
- cmake/FindPipCUDAToolkit.cmake
- src/target/stubs/cuda.cc
SiriusNEO
left a comment
SiriusNEO
left a comment
There was a problem hiding this comment.
Overall looks good, but we don't have a Windows CI to test it. Maybe we could test it locally? And left some comments.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/cache/kernel_cache.py (1)
91-107:⚠️ Potential issue | 🟡 MinorLibrary stamp still hashes only a single DLL despite expanded
lib_nameslist.The loop returns on the first match (line 106), so adding
tvm.dllandtvm_ffi.dlltolib_namesdoesn't actually include them in the stamp whentilelang.dllexists — which is the common case. If the intent (per the AI summary) is to invalidate the cache when any oftilelang/tvm/tvm_ffichanges, the stamp needs to hash all found libraries and combine the digests.♻️ Suggested change
- for lib_dir in lib_dirs: - for name in lib_names: - path = os.path.join(lib_dir, name) - if os.path.exists(path): - file_hash = sha256() - with open(path, "rb") as f: - for chunk in iter(lambda: f.read(1 << 20), b""): - file_hash.update(chunk) - return f"{name}:{file_hash.hexdigest()}" - return None + parts: list[str] = [] + for name in lib_names: + for lib_dir in lib_dirs: + path = os.path.join(lib_dir, name) + if os.path.exists(path): + file_hash = sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + file_hash.update(chunk) + parts.append(f"{name}:{file_hash.hexdigest()}") + break + return "|".join(parts) if parts else NoneIf only
tilelang.dllis meant to stamp the cache (andtvm.dll/tvm_ffi.dllare only fallbacks when TileLang isn't located), please document that intent — the current list ordering makes the added names effectively dead code in the normal install layout.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/cache/kernel_cache.py` around lines 91 - 107, The current loop over lib_names returns on the first match so the stamp only reflects a single library; change the logic in the kernel cache routine that builds the library stamp (use the existing lib_names, lib_dirs, sha256 variables) to collect all matching library files instead of returning immediately, compute each file's SHA-256 (as done with file_hash.update(...)) and combine their hex digests (e.g., concatenate with a separator or hash the concatenation) into a single stamp string (replace the current return f"{name}:{file_hash.hexdigest()}" with a combined digest for all found libs), and still return None if no libraries are found; ensure the symbol names lib_names, lib_dirs, sha256, and file_hash.hexdigest() are used to locate and update the code.
♻️ Duplicate comments (4)
tilelang/carver/arch/driver/cuda_driver.py (1)
55-64:⚠️ Potential issue | 🟡 MinorBogus
"11"DLL candidate andctypes.windllusage still unresolved.Previously flagged concerns still apply:
cudart64_11.dlldoes not exist in any CUDA release — Windows CUDA runtime DLLs use patterns likecudart64_110.dll,cudart64_12.dll,cudart64_130_96.dll, etc. The"11"entry is dead code.ctypes.windll.LoadLibraryreturns aWinDLL(stdcall default prototype), while the CUDA Runtime API is__cdecl. Harmless on x64 today, butctypes.cdll.LoadLibraryis semantically correct and keeps Windows/Linux paths consistent.Additionally, note the error-handling contract mismatch with the C++ side:
src/target/stubs/cudart.cc(TryLoadLibCudart) aborts on failure, whereas this Python helper raisesOSError. The callerget_device_attributeswallows it via the broadexcept Exceptionon Line 85, so behavior is fine here — just worth being aware of if this helper is reused elsewhere.♻️ Proposed change
def _load_cudart(): """Load the CUDA runtime library, searching for the version-suffixed DLL on Windows.""" if sys.platform == "win32": - for ver in ("13", "12", "110", "11"): + for ver in ("13", "12", "110"): try: - return ctypes.windll.LoadLibrary(f"cudart64_{ver}.dll") + return ctypes.cdll.LoadLibrary(f"cudart64_{ver}.dll") except OSError: continue raise OSError("Cannot find cudart64_*.dll") return ctypes.cdll.LoadLibrary("libcudart.so")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/carver/arch/driver/cuda_driver.py` around lines 55 - 64, The _load_cudart function currently tries a bogus "11" DLL name and uses ctypes.windll (WinDLL) which is semantically incorrect for the cdecl CUDA runtime; update _load_cudart to (1) replace the Windows candidate list with real CUDA runtime names (e.g. "cudart64_130_96.dll", "cudart64_130.dll", "cudart64_12.dll", "cudart64_110.dll" or similar valid suffixes) instead of "11", and (2) call ctypes.cdll.LoadLibrary on Windows as well (use ctypes.cdll.LoadLibrary(f"...") inside the Windows loop) so the runtime is loaded with cdecl semantics; keep raising OSError on not found so callers like get_device_attribute that catch exceptions continue to behave as before.tilelang/contrib/msvc.py (2)
175-175:⚠️ Potential issue | 🟡 MinorPass
usedforsecurity=Falsetohashlib.md5().The digest is only a deterministic filename prefix; marking it non-security-sensitive silences Ruff
S324and documents intent on FIPS-restricted systems wheremd5is otherwise blocked.🛠️ Proposed fix
- prefix = hashlib.md5(os.path.dirname(os.path.abspath(path)).encode()).hexdigest()[:8] + prefix = hashlib.md5( + os.path.dirname(os.path.abspath(path)).encode(), usedforsecurity=False + ).hexdigest()[:8]As per static-analysis hint Ruff
S324on line 175.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` at line 175, The MD5 usage for generating the filename prefix should be marked non-security-sensitive: modify the call to hashlib.md5(...) used when computing the variable prefix so it passes usedforsecurity=False to hashlib.md5 (i.e., update the hashlib.md5(...) invocation in the prefix assignment in msvc.py) to silence Ruff S324 and document intent on FIPS-restricted systems.
113-145:⚠️ Potential issue | 🟠 MajorStill caching/returning
base_envon MSVC setup failure.Lines 129, 135, 140 continue to cache and return
base_envwhen VsDevCmd discovery/import fails orcl.exeis still not on the resulting PATH. Callers (create_shared) then re-probe viashutil.which("cl.exe", path=get_env_path(compiler_env or {})), so on a machine wherecl.exehappens to be on the user's default PATH despite VsDevCmd failure, compilation proceeds with an environment that is missingINCLUDE/LIB/SDK variables and fails later with a far less clear error. ReturningNonehere (and leaving the explanation in_MSVC_ENV_ERROR) lets the caller surface the detail in itsRuntimeErrormessage.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` around lines 113 - 145, The function get_msvc_subprocess_env currently caches and returns base_env on failures, which masks missing MSVC vars; instead ensure that on any VsDevCmd discovery/import failure or when cl.exe is not found in the compiler_env you set _MSVC_ENV_CACHE to None (not base_env) and set _MSVC_ENV_ERROR accordingly, then return None so callers (e.g., create_shared) can surface the detailed error; update the branches that currently assign _MSVC_ENV_CACHE = base_env (the lines handling missing vsdevcmd, import failure, and missing cl.exe) to assign None and return None while preserving the _MSVC_ENV_ERROR message.cmake/FindPipCUDAToolkit.cmake (1)
283-285:⚠️ Potential issue | 🟠 MajorLinker-flag cache still grows on every re-configure.
CMAKE_{EXE,SHARED,MODULE}_LINKER_FLAGSare still being prepended and re-cached withFORCEeach time_tilelang_activate_msvc_envruns. Unlike the include flags (guarded by_TILELANG_INCLUDE_FLAGS_APPLIEDon line 224), these three still lack an idempotency guard, so repeatedcmakeinvocations keep accumulating/LIBPATH:…fragments until MSVC's command-line length limits start biting.🔧 Proposed guard
- set(CMAKE_EXE_LINKER_FLAGS_INIT "${_tilelang_libpath_flags} ${CMAKE_EXE_LINKER_FLAGS_INIT}") - set(CMAKE_SHARED_LINKER_FLAGS_INIT "${_tilelang_libpath_flags} ${CMAKE_SHARED_LINKER_FLAGS_INIT}") - set(CMAKE_MODULE_LINKER_FLAGS_INIT "${_tilelang_libpath_flags} ${CMAKE_MODULE_LINKER_FLAGS_INIT}") - set(CMAKE_EXE_LINKER_FLAGS "${_tilelang_libpath_flags} ${CMAKE_EXE_LINKER_FLAGS}" CACHE STRING "Executable linker flags" FORCE) - set(CMAKE_SHARED_LINKER_FLAGS "${_tilelang_libpath_flags} ${CMAKE_SHARED_LINKER_FLAGS}" CACHE STRING "Shared linker flags" FORCE) - set(CMAKE_MODULE_LINKER_FLAGS "${_tilelang_libpath_flags} ${CMAKE_MODULE_LINKER_FLAGS}" CACHE STRING "Module linker flags" FORCE) + if(NOT DEFINED _TILELANG_LIBPATH_FLAGS_APPLIED) + set(CMAKE_EXE_LINKER_FLAGS_INIT "${_tilelang_libpath_flags} ${CMAKE_EXE_LINKER_FLAGS_INIT}") + set(CMAKE_SHARED_LINKER_FLAGS_INIT "${_tilelang_libpath_flags} ${CMAKE_SHARED_LINKER_FLAGS_INIT}") + set(CMAKE_MODULE_LINKER_FLAGS_INIT "${_tilelang_libpath_flags} ${CMAKE_MODULE_LINKER_FLAGS_INIT}") + set(CMAKE_EXE_LINKER_FLAGS "${_tilelang_libpath_flags} ${CMAKE_EXE_LINKER_FLAGS}" CACHE STRING "Executable linker flags" FORCE) + set(CMAKE_SHARED_LINKER_FLAGS "${_tilelang_libpath_flags} ${CMAKE_SHARED_LINKER_FLAGS}" CACHE STRING "Shared linker flags" FORCE) + set(CMAKE_MODULE_LINKER_FLAGS "${_tilelang_libpath_flags} ${CMAKE_MODULE_LINKER_FLAGS}" CACHE STRING "Module linker flags" FORCE) + set(_TILELANG_LIBPATH_FLAGS_APPLIED TRUE CACHE INTERNAL "") + endif()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/FindPipCUDAToolkit.cmake` around lines 283 - 285, The linker flag lines in _tilelang_activate_msvc_env repeatedly prepend and FORCE-cache into CMAKE_EXE_LINKER_FLAGS, CMAKE_SHARED_LINKER_FLAGS and CMAKE_MODULE_LINKER_FLAGS on every configure, causing accumulation; add an idempotency guard (e.g. check/set a variable like _TILELANG_LINKER_FLAGS_APPLIED) around those three set(...) calls so they only run once (mirroring the existing _TILELANG_INCLUDE_FLAGS_APPLIED pattern used for include flags) and prevent repeated prepends on subsequent cmake runs.
🧹 Nitpick comments (10)
src/layout/gemm_layouts.cc (1)
14-21: Consider centralizing the unreachable macro in a shared header.If other translation units need the same "unreachable" semantics (the PR description mentions MSVC portability fixes across several files), defining
TILELANG_COMPILER_UNREACHABLElocally here will lead to duplication. Consider moving it to a common internal header (e.g., atl_common.hor similar) so it can be reused consistently across the codebase.Minor note: the
((void)0)fallback does not suppress "control reaches end of non-void function" warnings on exotic compilers that define none of_MSC_VER,__GNUC__, or__clang__— acceptable given those are effectively all mainstream toolchains today.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/layout/gemm_layouts.cc` around lines 14 - 21, The TILELANG_COMPILER_UNREACHABLE macro is being defined locally in gemm_layouts.cc which duplicates behavior across translation units; move this definition into a shared internal header (e.g., tl_common.h or internal/common.h) and replace the local definition with an `#include` of that header in gemm_layouts.cc and other files that need it. Ensure the macro is wrapped with include guards or a `#ifndef` TILELANG_COMPILER_UNREACHABLE check in the shared header to avoid redefinition, and preserve the same platform checks (__MSC_VER, __GNUC__, __clang__) and fallback semantics when migrating.src/op/parallel.cc (1)
85-93: Defensive optional/cast guarding is correct and consistent with the writer.The cast target
Map<Var, ReducerInfo>matches whatReducerLayoutAnnotator::VisitStmt_(ForNode)insrc/transform/layout_reducer.ccwrites toattr::kReducerInfo(inside_reducer_range_is declared asMap<Var, ReducerInfo>). The two-step guard (outerOptional<Any>then inneras<Map<Var, ReducerInfo>>()) avoids the MSVC debugstd::optionalassertion when the annotation is absent, and also tolerates the annotation being present with a different payload type.One small tradeoff worth noting: a silent fallback on a failed cast will mask an unexpected payload type rather than surface it. If in practice only this writer can produce the For-node annotation, consider either (a) leaving a
DLOG(WARNING)in the else-branch of the innerif, or (b) keeping anICHECKon the inner cast to catch regressions, while still guarding the outer optional. Feel free to ignore if the tolerant behavior is intentional (matching thekLayoutMaphandling inlayout_reducer.cc).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/op/parallel.cc` around lines 85 - 93, The current two-step guard around op->annotations.Get(attr::kReducerInfo) correctly avoids MSVC optional assertions but silently ignores a mismatched payload type; update the inner branch that attempts reducer_info_ref.value().as<Map<Var, ReducerInfo>>() so that if the cast fails you emit a diagnostic (e.g., DLOG(WARNING) with the annotation contents or type) rather than silently skipping it—leave the outer optional guard in place and still populate p->reducer_info_map_.Set(buffer, info) when the cast succeeds; alternatively, if you prefer failing fast during development, replace the silent path with an ICHECK on the inner cast to surface regressions (this mirrors the writer in ReducerLayoutAnnotator::VisitStmt_(ForNode) which produces Map<Var, ReducerInfo>).tilelang/libinfo.py (1)
49-56:for ... elsewith an innerreturnis harder to read than a plain post-loop raise.Since the inner loop never
breaks, theelse:clause runs exactly when the nested loops complete without a match — equivalent to placing theraisedirectly after the outerfor. Dropping theelsemakes the control flow more obvious.♻️ Proposed simplification
for lib_root in TL_LIBS: for lib_name in lib_names: lib_dll_path = os.path.join(lib_root, lib_name) if os.path.exists(lib_dll_path) and os.path.isfile(lib_dll_path): return lib_dll_path - else: - message = f"Cannot find libraries: {', '.join(lib_names)}\n" + "List of candidates:\n" + "\n".join(TL_LIBS) - raise RuntimeError(message) + message = ( + f"Cannot find libraries: {', '.join(lib_names)}\n" + "List of candidates:\n" + "\n".join(TL_LIBS) + ) + raise RuntimeError(message)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/libinfo.py` around lines 49 - 56, The for...else is unnecessary because the inner loop returns on match and never breaks, so simplify control flow by removing the else and placing the RuntimeError raise after the outer loop; specifically in tilelang.libinfo (the loop iterating TL_LIBS and lib_names that builds lib_dll_path), keep the nested iteration and the existing return when os.path.exists/os.path.isfile succeed, then after both loops complete raise the RuntimeError with the same message using TL_LIBS and lib_names.cmake/find_pip_cuda.py (2)
186-197: Silentsys.exit(1)on failure hides the root cause from CMake.
main()distinguishes "cu_dir not resolved" (line 192) from "nvcc missing inside cu_dir" (line 196) but emits no diagnostic in either case._tilelang_run_find_pip_cudaincmake/FindPipCUDAToolkit.cmakeonly seesresult_EQUAL 0vs non-zero, so whenWITH_PIP_CUDA_TOOLCHAINis set to a path wherecu_dirwas resolved butnvcchappens to be missing, the user sees the generic "no pip-installed CUDA toolkit could be resolved" message rather than the actionable "resolvedcu<ver>directory X butbin/nvccis missing". Consider writing a short message tosys.stderrin each branch so it surfaces viaexecute_processstderr.♻️ Proposed diff
- if cu_dir is None: - sys.exit(1) - - nvcc = _find_nvcc(cu_dir) - if nvcc is None: - sys.exit(1) + if cu_dir is None: + print("find_pip_cuda: could not resolve a pip-installed CUDA cu<ver> directory", file=sys.stderr) + sys.exit(1) + + nvcc = _find_nvcc(cu_dir) + if nvcc is None: + print(f"find_pip_cuda: resolved {cu_dir} but nvcc was not found under bin/", file=sys.stderr) + sys.exit(1)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/find_pip_cuda.py` around lines 186 - 197, The main() currently exits silently on failures; update it to write explicit error diagnostics to sys.stderr before calling sys.exit(1): when cu_dir is None (from _find_cu_dir()) print a clear message like "could not resolve cu<ver> directory" with context, and when nvcc is None (after calling _find_nvcc(cu_dir)) print a message that the resolved cu_dir exists but bin/nvcc is missing (include the cu_dir path). Ensure messages are short, go to sys.stderr so CMake's execute_process will surface them, and keep using sys.exit(1) after printing.
28-41: Redundant fallback probe in_find_nvcc.On each platform, one of the two fallback branches checks the same filename as
candidate(e.g., on Windows_nvcc_name()is alreadynvcc.exe, making lines 37‑39 redundant; on POSIX_nvcc_name()isnvcc, making lines 33‑35 redundant). The logic still works but the duplication is easy to misread.♻️ Proposed simplification
def _find_nvcc(cu_dir: pathlib.Path) -> pathlib.Path | None: - candidate = cu_dir / "bin" / _nvcc_name() - if candidate.is_file(): - return candidate - - fallback = cu_dir / "bin" / "nvcc" - if fallback.is_file(): - return fallback - - fallback_exe = cu_dir / "bin" / "nvcc.exe" - if fallback_exe.is_file(): - return fallback_exe - - return None + bin_dir = cu_dir / "bin" + for name in (_nvcc_name(), "nvcc", "nvcc.exe"): + candidate = bin_dir / name + if candidate.is_file(): + return candidate + return None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/find_pip_cuda.py` around lines 28 - 41, The _find_nvcc function currently probes the same filename multiple times because _nvcc_name() already returns platform-specific names; simplify by building an ordered list of unique filenames to check (start with _nvcc_name(), then add "nvcc" and "nvcc.exe" only if they differ) and iterate that list returning the first existing file; update _find_nvcc to use that de-duplicated probe sequence so candidate/fallback/fallback_exe logic is removed but behavior remains the same.tilelang/contrib/msvc.py (2)
216-224: Narrow the blanketexcept Exceptionwhen probingtvm_ffi.libinfo.Catching
Exceptionhides real bugs (for example a typo infind_windows_implibwould silently degrade link behavior). Narrow to the concrete failure modes (ImportError,AttributeError,OSError) and keep the warning.♻️ Proposed fix
- try: - from tvm_ffi import libinfo as tvm_ffi_libinfo - - add_import_lib(tvm_ffi_libinfo.find_windows_implib()) - except Exception as e: - import logging - - logging.getLogger(__name__).warning("Could not locate tvm_ffi import lib: %s", e) + try: + from tvm_ffi import libinfo as tvm_ffi_libinfo + + add_import_lib(tvm_ffi_libinfo.find_windows_implib()) + except (ImportError, AttributeError, OSError) as e: + import logging + + logging.getLogger(__name__).warning("Could not locate tvm_ffi import lib: %s", e)As per static-analysis hint Ruff
BLE001on line 220.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` around lines 216 - 224, The current try/except around importing tvm_ffi.libinfo and calling tvm_ffi_libinfo.find_windows_implib() uses a blanket except Exception which hides real bugs; change the except to catch only the concrete failure modes (ImportError, AttributeError, OSError) and keep the existing logging.getLogger(__name__).warning(...) call so failures still warn but other unexpected exceptions propagate; reference the import of tvm_ffi, the tvm_ffi_libinfo.find_windows_implib() call, and the add_import_lib() invocation when making the change.
152-161:_normalize_optionsilently passes through unrecognized GCC flags tocl.exe.Only
-I,-D,-g,-fPIC, and-std=are translated or dropped. Unrecognized flags like-L<path>,-l<name>,-Wl,…,-shared,-pthread,-fvisibility=…,-Wno-…, or-O3will be passed tocl.exeunchanged, causing it to error out or misinterpret them.Consider extending the mapping for common flags that may be passed by TVM's Windows runtime, or silently dropping any unrecognized flags prefixed with
-.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/contrib/msvc.py` around lines 152 - 161, The _normalize_option function currently only handles -I, -D, -g, -fPIC, and -std= and will accidentally pass other GCC-style flags to cl.exe; update _normalize_option to (1) explicitly handle common flags used on Windows by mapping or dropping them (e.g. treat options starting with "-L" and "-l" as linker directives to be dropped here or converted for the linker, drop "-shared", "-pthread", "-fvisibility=", "-Wl,", "-Wno-", and translate or drop optimization flags like "-O*"), and (2) as a safe fallback, if an option starts with "-" and is not one of the handled cases, return None so the flag is not forwarded to cl.exe; keep the function name _normalize_option and its return semantics.cmake/FindPipCUDAToolkit.cmake (3)
130-137: Dead branches in arch normalization.
_tilelang_sdk_archis seeded from_tilelang_target_arch, which is already one ofx64/x86/arm64per lines 83-88. Theif/elifbranches just re-assign the same value; theelsefallback tox64is unreachable given the upstream initialization. Collapsing this to the single initialset()removes the confusion.♻️ Suggested simplification
- set(_tilelang_sdk_arch "${_tilelang_target_arch}") - if(_tilelang_sdk_arch STREQUAL "arm64") - set(_tilelang_sdk_arch "arm64") - elseif(_tilelang_sdk_arch STREQUAL "x86") - set(_tilelang_sdk_arch "x86") - else() - set(_tilelang_sdk_arch "x64") - endif() + set(_tilelang_sdk_arch "${_tilelang_target_arch}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/FindPipCUDAToolkit.cmake` around lines 130 - 137, _tilelang_sdk_arch is redundantly normalized from _tilelang_target_arch using unreachable branches; remove the if/elseif/else block and keep a single assignment of _tilelang_sdk_arch from _tilelang_target_arch (i.e., just set _tilelang_sdk_arch "${_tilelang_target_arch}"), ensuring no fallback to "x64" remains and preserving the original variable names _tilelang_sdk_arch and _tilelang_target_arch.
243-258: DuplicatePATH/PATHEXTprefix handling.Lines 59-81 already prepend
_tilelang_system_path_candidatestoENV{PATH}and install thePATHEXTdefault, and lines 169-171 already prepend_tilelang_sdk_bin. Re-running the exact same blocks here (244-250, 251-255, 256-258) prepends those entries a second time within the same configure, makingENV{PATH}grow and duplicatingPATHEXThandling for no functional reason.♻️ Suggested change
- if(_tilelang_system_path_candidates) - list(JOIN _tilelang_system_path_candidates ";" _tilelang_system_path_prefix) - if(DEFINED ENV{PATH} AND NOT "$ENV{PATH}" STREQUAL "") - set(ENV{PATH} "${_tilelang_system_path_prefix};$ENV{PATH}") - else() - set(ENV{PATH} "${_tilelang_system_path_prefix}") - endif() - endif() - if(NOT DEFINED ENV{PATHEXT} OR "$ENV{PATHEXT}" STREQUAL "") - set(ENV{PATHEXT} "${_tilelang_default_pathext}") - elseif(NOT "$ENV{PATHEXT}" MATCHES "(^|;)[.]EXE($|;)") - set(ENV{PATHEXT} "$ENV{PATHEXT};${_tilelang_default_pathext}") - endif() - if(NOT _tilelang_sdk_bin STREQUAL "") - set(ENV{PATH} "${_tilelang_sdk_bin};$ENV{PATH}") - endif()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/FindPipCUDAToolkit.cmake` around lines 243 - 258, This diff duplicates earlier PATH/PATHEXT prepending logic; remove or guard the repeated block that re-prepends _tilelang_system_path_candidates and _tilelang_sdk_bin and re-applies PATHEXT defaults. Locate the second occurrence around the block manipulating ENV{PATH} and ENV{PATHEXT} (references: _tilelang_system_path_candidates, _tilelang_default_pathext, _tilelang_sdk_bin) and either delete this duplicated section or wrap it with a one-time guard flag (e.g. set a variable like _tilelang_paths_prepended and check NOT _tilelang_paths_prepended before modifying ENV{PATH}/ENV{PATHEXT}) so PATH/PATHEXT are only modified once per configure.
464-494:_tilelang_run_find_pip_cudacannot distinguish the two failure modes offind_pip_cuda.py.Both "no
cu<ver>dir found" and "cu<ver>resolved butnvccmissing" exit with code 1 and no stderr. If the first Python candidate lands in the second bucket, the loop silently tries the next interpreter and the finalmessage(STATUS "FindPipCUDAToolkit: pip-installed CUDA toolkit not found")(line 510) buries the real cause. Pairing this with a short stderr message infind_pip_cuda.py(see my note on itsmain()) and forwardingOUTPUT_STRIP_TRAILING_WHITESPACE+ capturingERROR_VARIABLEhere would make misconfiguration diagnosable without re-running manually.♻️ Suggested diff
- execute_process( - COMMAND "${_python}" "${CMAKE_CURRENT_LIST_DIR}/find_pip_cuda.py" "${ARGV2}" - OUTPUT_VARIABLE _candidate_output - OUTPUT_STRIP_TRAILING_WHITESPACE - RESULT_VARIABLE _candidate_result - ) + execute_process( + COMMAND "${_python}" "${CMAKE_CURRENT_LIST_DIR}/find_pip_cuda.py" "${ARGV2}" + OUTPUT_VARIABLE _candidate_output + ERROR_VARIABLE _candidate_err + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _candidate_result + ) + if(NOT _candidate_result EQUAL 0 AND NOT "${_candidate_err}" STREQUAL "") + message(STATUS "FindPipCUDAToolkit: ${_python} -> ${_candidate_err}") + endif()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmake/FindPipCUDAToolkit.cmake` around lines 464 - 494, Modify function _tilelang_run_find_pip_cuda to capture execute_process stderr and surface it back to the caller: add ERROR_VARIABLE _candidate_error to both execute_process calls (keep OUTPUT_STRIP_TRAILING_WHITESPACE), and when a candidate returns non‑zero but _candidate_error is non‑empty, set _result_json to _candidate_error and _result_python to the current ${_python} and break the loop so the caller can log the real stderr; keep the existing behavior for successful (RESULT 0) runs.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/target/stubs/dynlib.h`:
- Around line 59-95: In dynlib_find_loaded, avoid an out-of-bounds read when
EnumProcessModules reports a larger 'needed' after resizing by computing the
module count from the actual allocated buffer size (buf_bytes) rather than the
reported 'needed'; clamp or recompute DWORD count = (buf_bytes /
sizeof(HMODULE)) (or take min(needed/sizeof(HMODULE),
buf_bytes/sizeof(HMODULE))) before iterating, so the for-loop over modules[i]
cannot exceed the allocated array; ensure this change covers both the initial
and post-realloc paths and still frees modules on all early returns.
---
Outside diff comments:
In `@tilelang/cache/kernel_cache.py`:
- Around line 91-107: The current loop over lib_names returns on the first match
so the stamp only reflects a single library; change the logic in the kernel
cache routine that builds the library stamp (use the existing lib_names,
lib_dirs, sha256 variables) to collect all matching library files instead of
returning immediately, compute each file's SHA-256 (as done with
file_hash.update(...)) and combine their hex digests (e.g., concatenate with a
separator or hash the concatenation) into a single stamp string (replace the
current return f"{name}:{file_hash.hexdigest()}" with a combined digest for all
found libs), and still return None if no libraries are found; ensure the symbol
names lib_names, lib_dirs, sha256, and file_hash.hexdigest() are used to locate
and update the code.
---
Duplicate comments:
In `@cmake/FindPipCUDAToolkit.cmake`:
- Around line 283-285: The linker flag lines in _tilelang_activate_msvc_env
repeatedly prepend and FORCE-cache into CMAKE_EXE_LINKER_FLAGS,
CMAKE_SHARED_LINKER_FLAGS and CMAKE_MODULE_LINKER_FLAGS on every configure,
causing accumulation; add an idempotency guard (e.g. check/set a variable like
_TILELANG_LINKER_FLAGS_APPLIED) around those three set(...) calls so they only
run once (mirroring the existing _TILELANG_INCLUDE_FLAGS_APPLIED pattern used
for include flags) and prevent repeated prepends on subsequent cmake runs.
In `@tilelang/carver/arch/driver/cuda_driver.py`:
- Around line 55-64: The _load_cudart function currently tries a bogus "11" DLL
name and uses ctypes.windll (WinDLL) which is semantically incorrect for the
cdecl CUDA runtime; update _load_cudart to (1) replace the Windows candidate
list with real CUDA runtime names (e.g. "cudart64_130_96.dll",
"cudart64_130.dll", "cudart64_12.dll", "cudart64_110.dll" or similar valid
suffixes) instead of "11", and (2) call ctypes.cdll.LoadLibrary on Windows as
well (use ctypes.cdll.LoadLibrary(f"...") inside the Windows loop) so the
runtime is loaded with cdecl semantics; keep raising OSError on not found so
callers like get_device_attribute that catch exceptions continue to behave as
before.
In `@tilelang/contrib/msvc.py`:
- Line 175: The MD5 usage for generating the filename prefix should be marked
non-security-sensitive: modify the call to hashlib.md5(...) used when computing
the variable prefix so it passes usedforsecurity=False to hashlib.md5 (i.e.,
update the hashlib.md5(...) invocation in the prefix assignment in msvc.py) to
silence Ruff S324 and document intent on FIPS-restricted systems.
- Around line 113-145: The function get_msvc_subprocess_env currently caches and
returns base_env on failures, which masks missing MSVC vars; instead ensure that
on any VsDevCmd discovery/import failure or when cl.exe is not found in the
compiler_env you set _MSVC_ENV_CACHE to None (not base_env) and set
_MSVC_ENV_ERROR accordingly, then return None so callers (e.g., create_shared)
can surface the detailed error; update the branches that currently assign
_MSVC_ENV_CACHE = base_env (the lines handling missing vsdevcmd, import failure,
and missing cl.exe) to assign None and return None while preserving the
_MSVC_ENV_ERROR message.
---
Nitpick comments:
In `@cmake/find_pip_cuda.py`:
- Around line 186-197: The main() currently exits silently on failures; update
it to write explicit error diagnostics to sys.stderr before calling sys.exit(1):
when cu_dir is None (from _find_cu_dir()) print a clear message like "could not
resolve cu<ver> directory" with context, and when nvcc is None (after calling
_find_nvcc(cu_dir)) print a message that the resolved cu_dir exists but bin/nvcc
is missing (include the cu_dir path). Ensure messages are short, go to
sys.stderr so CMake's execute_process will surface them, and keep using
sys.exit(1) after printing.
- Around line 28-41: The _find_nvcc function currently probes the same filename
multiple times because _nvcc_name() already returns platform-specific names;
simplify by building an ordered list of unique filenames to check (start with
_nvcc_name(), then add "nvcc" and "nvcc.exe" only if they differ) and iterate
that list returning the first existing file; update _find_nvcc to use that
de-duplicated probe sequence so candidate/fallback/fallback_exe logic is removed
but behavior remains the same.
In `@cmake/FindPipCUDAToolkit.cmake`:
- Around line 130-137: _tilelang_sdk_arch is redundantly normalized from
_tilelang_target_arch using unreachable branches; remove the if/elseif/else
block and keep a single assignment of _tilelang_sdk_arch from
_tilelang_target_arch (i.e., just set _tilelang_sdk_arch
"${_tilelang_target_arch}"), ensuring no fallback to "x64" remains and
preserving the original variable names _tilelang_sdk_arch and
_tilelang_target_arch.
- Around line 243-258: This diff duplicates earlier PATH/PATHEXT prepending
logic; remove or guard the repeated block that re-prepends
_tilelang_system_path_candidates and _tilelang_sdk_bin and re-applies PATHEXT
defaults. Locate the second occurrence around the block manipulating ENV{PATH}
and ENV{PATHEXT} (references: _tilelang_system_path_candidates,
_tilelang_default_pathext, _tilelang_sdk_bin) and either delete this duplicated
section or wrap it with a one-time guard flag (e.g. set a variable like
_tilelang_paths_prepended and check NOT _tilelang_paths_prepended before
modifying ENV{PATH}/ENV{PATHEXT}) so PATH/PATHEXT are only modified once per
configure.
- Around line 464-494: Modify function _tilelang_run_find_pip_cuda to capture
execute_process stderr and surface it back to the caller: add ERROR_VARIABLE
_candidate_error to both execute_process calls (keep
OUTPUT_STRIP_TRAILING_WHITESPACE), and when a candidate returns non‑zero but
_candidate_error is non‑empty, set _result_json to _candidate_error and
_result_python to the current ${_python} and break the loop so the caller can
log the real stderr; keep the existing behavior for successful (RESULT 0) runs.
In `@src/layout/gemm_layouts.cc`:
- Around line 14-21: The TILELANG_COMPILER_UNREACHABLE macro is being defined
locally in gemm_layouts.cc which duplicates behavior across translation units;
move this definition into a shared internal header (e.g., tl_common.h or
internal/common.h) and replace the local definition with an `#include` of that
header in gemm_layouts.cc and other files that need it. Ensure the macro is
wrapped with include guards or a `#ifndef` TILELANG_COMPILER_UNREACHABLE check in
the shared header to avoid redefinition, and preserve the same platform checks
(__MSC_VER, __GNUC__, __clang__) and fallback semantics when migrating.
In `@src/op/parallel.cc`:
- Around line 85-93: The current two-step guard around
op->annotations.Get(attr::kReducerInfo) correctly avoids MSVC optional
assertions but silently ignores a mismatched payload type; update the inner
branch that attempts reducer_info_ref.value().as<Map<Var, ReducerInfo>>() so
that if the cast fails you emit a diagnostic (e.g., DLOG(WARNING) with the
annotation contents or type) rather than silently skipping it—leave the outer
optional guard in place and still populate p->reducer_info_map_.Set(buffer,
info) when the cast succeeds; alternatively, if you prefer failing fast during
development, replace the silent path with an ICHECK on the inner cast to surface
regressions (this mirrors the writer in
ReducerLayoutAnnotator::VisitStmt_(ForNode) which produces Map<Var,
ReducerInfo>).
In `@tilelang/contrib/msvc.py`:
- Around line 216-224: The current try/except around importing tvm_ffi.libinfo
and calling tvm_ffi_libinfo.find_windows_implib() uses a blanket except
Exception which hides real bugs; change the except to catch only the concrete
failure modes (ImportError, AttributeError, OSError) and keep the existing
logging.getLogger(__name__).warning(...) call so failures still warn but other
unexpected exceptions propagate; reference the import of tvm_ffi, the
tvm_ffi_libinfo.find_windows_implib() call, and the add_import_lib() invocation
when making the change.
- Around line 152-161: The _normalize_option function currently only handles -I,
-D, -g, -fPIC, and -std= and will accidentally pass other GCC-style flags to
cl.exe; update _normalize_option to (1) explicitly handle common flags used on
Windows by mapping or dropping them (e.g. treat options starting with "-L" and
"-l" as linker directives to be dropped here or converted for the linker, drop
"-shared", "-pthread", "-fvisibility=", "-Wl,", "-Wno-", and translate or drop
optimization flags like "-O*"), and (2) as a safe fallback, if an option starts
with "-" and is not one of the handled cases, return None so the flag is not
forwarded to cl.exe; keep the function name _normalize_option and its return
semantics.
In `@tilelang/libinfo.py`:
- Around line 49-56: The for...else is unnecessary because the inner loop
returns on match and never breaks, so simplify control flow by removing the else
and placing the RuntimeError raise after the outer loop; specifically in
tilelang.libinfo (the loop iterating TL_LIBS and lib_names that builds
lib_dll_path), keep the nested iteration and the existing return when
os.path.exists/os.path.isfile succeed, then after both loops complete raise the
RuntimeError with the same message using TL_LIBS and lib_names.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 571ca3d0-455e-4379-9037-2fbd23c761bc
📒 Files selected for processing (32)
.gitattributes.gitignoreCMakeLists.txtcmake/FindPipCUDAToolkit.cmakecmake/find_pip_cuda.pycmake/generate_windows_import_lib.pycmake/pypi-z3/FindZ3.cmakedocs/get_started/Installation.mdpyproject.tomlsrc/layout/gemm_layouts.ccsrc/op/operator.hsrc/op/parallel.ccsrc/runtime/logging.ccsrc/target/codegen_cuda.ccsrc/target/codegen_cuda.hsrc/target/stubs/cuda.ccsrc/target/stubs/cuda.hsrc/target/stubs/cudart.ccsrc/target/stubs/dynlib.hsrc/target/stubs/nvrtc.ccsrc/transform/layout_reducer.cctilelang/__init__.pytilelang/cache/kernel_cache.pytilelang/carver/arch/driver/cuda_driver.pytilelang/contrib/cc.pytilelang/contrib/msvc.pytilelang/contrib/nvcc.pytilelang/env.pytilelang/jit/adapter/libgen.pytilelang/jit/adapter/tvm_ffi.pytilelang/jit/adapter/wrapper.pytilelang/libinfo.py
✅ Files skipped from review due to trivial changes (3)
- .gitignore
- .gitattributes
- CMakeLists.txt
🚧 Files skipped from review as they are similar to previous changes (9)
- src/target/codegen_cuda.h
- tilelang/jit/adapter/tvm_ffi.py
- tilelang/jit/adapter/libgen.py
- cmake/pypi-z3/FindZ3.cmake
- src/target/stubs/cuda.h
- src/runtime/logging.cc
- tilelang/jit/adapter/wrapper.py
- src/target/stubs/cudart.cc
- src/target/stubs/cuda.cc
|
At this point, we have enabled support for most operators in both compile and JIT modes on Windows through |
d8b6dd8 to
2f5cede
Compare
|
Yes. Our ROCm CI seems to have some problems, but it's not an issue with this PR. |
The last available torch version is |
|
Update: ccache is now supported on Windows builds, including when PEP 517 build isolation is enabled. This should significantly speed up incremental rebuilds in isolated environments. |
|
Parallel autotune is now available on Windows as well, though it may be worth further investigation into why |
82d7d73 to
baa4400
Compare
All tests passed, with the exception of those related to tmcgen05 or cutlass-dsl. |
|
No more comments, looks good to me. cc @LeiWang1999. Since this PR is quite large, it may take some time to review and merge. |
|
@SiriusNEO I think this pr is ready, would you mind helping resolve the conflict and run perf regression test, then we can let it in. |
…ues (tile-ai#2093) * [Feature] Add Windows support with cross-platform minimal diff * [Feature] Add ccache support for Windows * [Feature] Add parallel autotune support for Windows * [Feature] Add nvrtc support for Windows * [Bugfix] Fix tilelibrary support for Windows * [Bugfix] Fix alignment for Windows * [Feature] Add clang-cl support for Windows * [Bugfix] Fix numeric limit for Windows * [Bugfix] Fix rand for Windows * [Misc] Use dynamic test workload to avoid hang * [Bugfix] Fix TMA alignment for Windows and sm120 * [Bugfix] Add unix-like hex formatter for Windows * [CI] Add Windows wheel job * [Feature] Flexible CUDA loader * fix rebase --------- Co-authored-by: SiriusNEO <chaofan@deepseek.com>
3 participants
Summary
This PR adds end-to-end Windows support for TileLang across build, packaging, runtime loading, and developer setup.
Along the way, it also fixes several cross-platform issues uncovered during the Windows bring-up, especially around CUDA toolkit discovery, MSVC compatibility, runtime DLL loading, logging, and layout/reducer handling.
Main changes
Bug fixes
TVM_DLLimport/export behaviorTVM_LOG_DEBUGTODO / Follow-ups
dynload) polyfillKnown issues
cutlass-dsldoes not officially support Windows, so any code paths depending on it are currently unavailable in the Windows build.__declspec(align(...))if is not newer than17.10. As a workaround, we use a torch-style fallback that degrades to the largest alignment MSVC accepts; this matches PyTorch's behavior but means the strict CUDA 13 alignment guarantee is not preserved on MSVC.Summary by CodeRabbit
Release Notes