[Feat] Integrate Z3 in TVM Arith Analyzer #1367
Conversation
WalkthroughAdds optional Z3 integration to the build and packaging (CMake + Python packaging), introduces multiple new arithmetic test modules (including Z3-backed proofs), parameterizes CUDA GEMM tests, tightens cumsum_fragment's signature to Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~35 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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 |
|
👋 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! 🚀 |
testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py
Show resolved
Hide resolved
|
Macos pending on: tile-ai/tvm#15 cc @kurisu6912 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
CMakeLists.txt (1)
313-314: Variable name mismatch:Python3_SITELIBis undefined; usePython_SITELIB.Line 279 calls
find_package(Python REQUIRED ...), which definesPython_SITELIB(notPython3_SITELIB). The BUILD_RPATH references will expand to empty, breaking Z3 library discovery at runtime during development builds.Apply this diff to fix the variable name:
- set_property(TARGET tvm APPEND PROPERTY BUILD_RPATH ${Python3_SITELIB}/z3/lib) - set_property(TARGET tvm APPEND PROPERTY BUILD_RPATH ${Python3_SITELIB}/z3/bin) + set_property(TARGET tvm APPEND PROPERTY BUILD_RPATH ${Python_SITELIB}/z3/lib) + set_property(TARGET tvm APPEND PROPERTY BUILD_RPATH ${Python_SITELIB}/z3/bin)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
CMakeLists.txt(2 hunks)pyproject.toml(4 hunks)requirements-dev.txt(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- requirements-dev.txt
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
- GitHub Check: Build wheels for Python 3.9 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.9 on macos-latest with Metal
- GitHub Check: Build wheels for Python 3.9 on ubuntu-latest with CUDA-12.8
- GitHub Check: Build SDist
🔇 Additional comments (5)
pyproject.toml (3)
46-46: Platform-conditional Z3 and patchelf dependencies are correctly scoped.The conditional placement (platform_system == 'Linux') in both dependencies and build-system.requires ensures Z3-related tooling is only pulled on Linux, avoiding unnecessary dependencies on macOS and Windows. This aligns with the Z3 packaging strategy outlined in the PR (avoiding SONAME issues across platforms).
Also applies to: 57-62
236-236: manylinux version bump reflects Z3 glibc requirement.The change from manylinux_2_28 (glibc 2.28) to manylinux_2_34 (glibc 2.34) on aarch64 aligns with Z3's minimum glibc requirement on ARM, as noted in the PR discussion. This ensures compatibility without linker errors on ARM builds.
265-265: libz3.so exclusion from wheel repair is correct.Excluding libz3.so from auditwheel's repair prevents the Z3 shared library from being bundled twice (once from the PyPI wheel, once from repair), which could cause symbol conflicts or packaging issues.
CMakeLists.txt (2)
318-335: INSTALL_RPATH logic correctly differentiates platforms and Z3 availability.The conditional setup (Apple with
@loader_pathpaths, UNIX with$ORIGINand optional Z3 library path) is well-structured and matches the packaging strategy. macOS uses loader_path for relocatability, and Linux respects both TVM FFI and Z3 library locations when Z3 is enabled.
341-362: Install-time Z3 soname patching is correctly guarded.The patchelf logic is properly gated behind both
if(USE_Z3 AND USE_PYPI_Z3)andif(UNIX AND NOT APPLE), ensuring it only runs on Linux when Z3 is actually enabled. The glob pattern*.so*is appropriate for Linux-only context. This approach aligns with the PR's strategy of patching Z3's custom SONAME to a stablelibz3.soto avoid runtime binding issues.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
CMakeLists.txt (1)
225-231: UnconditionalUSE_Z3=ONoverrides user configuration and surprises opt-out users.Setting USE_Z3 unconditionally forces Z3 support even for users explicitly building with
-DUSE_Z3=OFFor wanting a minimal CPU-only build. This was previously flagged as a major issue; a gating option is needed.Apply this diff to respect user configuration:
+# Gate Z3 support behind TileLang option (default OFF for minimal builds) +option(TILELANG_USE_Z3 "Enable Z3 support in TVM arith analyzer" OFF) + -set(USE_Z3 ON CACHE STRING "Use Z3 SMT solver for TileLang optimizations") -set(USE_PYPI_Z3 ON CACHE BOOL "Use Z3 provided by PyPI z3-solver package") - -if(USE_Z3 AND USE_PYPI_Z3) +if(TILELANG_USE_Z3) + set(USE_Z3 ON CACHE STRING "Use Z3 SMT solver for TileLang optimizations") + set(USE_PYPI_Z3 ON CACHE BOOL "Use Z3 provided by PyPI z3-solver package") list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}/cmake/pypi-z3") find_package(Z3 REQUIRED) endif()This allows users to opt-in via
-DTILELANG_USE_Z3=ONwhile respecting explicit-DUSE_Z3=OFFchoices.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
CMakeLists.txt(2 hunks)
🔇 Additional comments (1)
CMakeLists.txt (1)
305-323: RPATH configuration is well-structured and platform-aware.The INSTALL_RPATH setup correctly differentiates between platforms (macOS
@loader_path, Linux$ORIGIN), properly gates Z3-specific paths, and applies settings consistently to all relevant targets. Assuming thePython_SITELIBfix above is applied, this segment will function correctly.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py (2)
219-220: Consider removing decorators from helper function.The
@tilelang.testing.requires_cudaand@tilelang.testing.requires_cuda_compute_versiondecorators onrun_gemm_sp_sm90are redundant since this is a helper function, not a pytest-discovered test. Pytest markers only affect test functions (those starting withtest_). The effective decorators are ontest_gemm_sp_sm90(lines 310-311), so these can be safely removed for clarity.Apply this diff:
-@tilelang.testing.requires_cuda -@tilelang.testing.requires_cuda_compute_version(9, 0) def run_gemm_sp_sm90(
264-266: Consider removing decorators from helper function.Similar to
run_gemm_sp_sm90, these decorators onrun_gemm_sp_sm80are redundant. The effective compute-version guards are already ontest_gemm_sp_sm80(lines 334-336), so these can be removed for clarity.Apply this diff:
-@tilelang.testing.requires_cuda -@tilelang.testing.requires_cuda_compute_version_ge(8, 0) -@tilelang.testing.requires_cuda_compute_version_le(8, 9) def run_gemm_sp_sm80(
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py(4 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-11-14T07:56:11.098Z
Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.
Applied to files:
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py
🧬 Code graph analysis (1)
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py (3)
tilelang/testing/__init__.py (3)
requires_cuda_compute_version(37-101)requires_cuda_compute_version_ge(104-105)requires_cuda_compute_version_le(120-121)tilelang/tileop/gemm/__init__.py (2)
trans_A(101-102)trans_B(105-106)tilelang/tileop/gemm_sp/gemm_sp_base.py (2)
trans_A(45-46)trans_B(49-50)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
- GitHub Check: Build wheels for Python 3.9 on ubuntu-latest with CUDA-12.8
- GitHub Check: Build wheels for Python 3.9 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build SDist
- GitHub Check: Build wheels for Python 3.9 on macos-latest with Metal
🔇 Additional comments (1)
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py (1)
233-234: LGTM: Transpose parameter additions are correct.The
trans_Aandtrans_Bparameters are properly threaded through the entire call chain (kernel generation → run_gemm_sp → validation) and correctly control matrix transposition in the sparse GEMM tests. The test parameterization provides good coverage of transpose combinations.Also applies to: 279-280
|
This PR brings performance improvements for mha_fwd_bhsd_wgmma_pipelined, as well as other cases such as warp_specialize_gemm_copy1_gemm_0. |
5 participants
This pr integrates Z3 into the TVM arith analyzer.
TODO
AtomicAddop ([BUG] Auto Vectorization on AtomicAdd Ignores BufferLayout and LowerArgs #1366) (I don't know which pr fixed it, anyway, it is fixed)CMakeLists.txt,3rdparty/tvm/CMakeLists.txt, and build scripts, package dependencies3rdparty/tvm, mergez3branch intotilelang_mainAdded API
analyzer.get_smtlib2: Returns a SMTLIB2 format including all constraints, variables and debug informationsanalyzer.get_smtlib2(expr): Based onget_smtlib2and adding expr as the prove targetanalyzer.set_z3_timeout_ms(t): Set z3 timeout in millsecondsanalyzer.set_z3_max_step(step): Set z3 max stepanalyzer.get_z3_stats(): Return z3 statistic information (in str)Example
This example gives the proving result
True.analyzer.get_smtlib2():analyzer.get_smtlib2(try_to_prove):analyzer.get_z3_stats():Summary by CodeRabbit
New Features
Breaking Changes
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.