[Feat] Integrate Z3 in TVM Arith Analyzer
This pr integrate z3 in tvm arith analyzer
Walkthrough
Adds 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 tir.Buffer, updates TVM submodule pointer, and adds z3-solver to requirements.
Changes
| Cohort / File(s) | Summary |
|---|---|
Z3 Integration (Build & CMake) CMakeLists.txt, cmake/pypi-z3/FindZ3.cmake |
Adds USE_Z3 and USE_PYPI_Z3 CMake options; conditionally loads cmake/pypi-z3 and find_package(Z3) when opted in; adjusts BUILD_RPATH and TILELANG_INSTALL_RPATH for Z3 libraries on Apple/UNIX; sets per-target INSTALL_RPATH for tilelang, tilelang_module, tvm, tvm_runtime; introduces imported target z3::libz3. |
Packaging & Dependencies pyproject.toml, requirements.txt, requirements-dev.txt, requirements-test.txt |
Adds z3-solver>=4.13.0 to dependencies and build-system requirements; adds patchelf to build requires (platform-specific); updates manylinux image and wheel repair exclusions to omit libz3.so. |
Arithmetic Analysis Tests testing/python/arith/test_arith_hard.py, testing/python/arith/test_arith_intset.py, testing/python/arith/test_arith_iter_affine_map.py, testing/python/arith/test_arith_simplify.py |
Adds four comprehensive test modules exercising Analyzer, int-set reasoning, iter-affine-map analysis, simplification and SMT-LIB2 / Z3-backed proofs, including many new test helpers and entrypoints. |
CUDA GEMM Tests Parameterization testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py |
Parameterizes SM90/SM80 tests with explicit trans_A/trans_B args, requires CUDA and compute-version guards, and threads new params through run_gemm_sp_* call chain. |
Tests Cleanup / Minor Fixes testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py, tilelang/utils/sparse.py |
Comments out an issue-specific test and wrapper; removes a non-functional comment in sparse.py except block. |
API Signature Restriction tilelang/language/reduce.py |
Narrows cumsum_fragment macro signature: src and dst now typed as tir.Buffer (previously allowed BufferRegion/BufferLoad). |
Submodule Update 3rdparty/tvm |
Advances TVM submodule pointer from 2b1ead1a... to 4d3ec925... (no code-level changes). |
Estimated code review effort
🎯 3 (Moderate) | ⏱️ ~35 minutes
- Review hotspots:
CMakeLists.txt/cmake/pypi-z3/FindZ3.cmake: RPATH logic, imported target properties, and conditional find_package behavior.- Packaging changes in
pyproject.toml: build-system requires and wheel repair exclusions. - New test modules: scan for flaky assumptions, Z3 dependency usage, and test entrypoints.
tilelang/language/reduce.py: verify callers are compatible with narrowedcumsum_fragmentsignature.
Possibly related PRs
- tile-ai/tilelang#1426 — touches
cumsum_fragmenthandling; opposite direction of signature change (related to Buffer/Region handling). - tile-ai/tilelang#1358 — modifies
cumsum/reduce paths inreduce.py; overlaps on same function. - tile-ai/tilelang#1334 — also changes
cumsum_fragmentand region-to-buffer conversions; likely to conflict or require coordination.
Suggested reviewers
- LeiWang1999
- XuehaiPan
Poem
🐰 I hopped into CMake with a cheerful click,
I fetched Z3 from PyPI — quick, quick, quick!
Tests grew bold proofs, SMT whispers at night,
Buffers got picky, now signatures are tight.
Submodule bumped, wheels tuned just right — hooray!
Pre-merge checks and finishing touches
❌ Failed checks (1 warning)
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage | ⚠️ Warning | Docstring coverage is 10.64% which is insufficient. The required threshold is 80.00%. | You can run @coderabbitai generate docstrings to improve docstring coverage. |
✅ Passed checks (2 passed)
| Check name | Status | Explanation |
|---|---|---|
| Description Check | ✅ Passed | Check skipped - CodeRabbit’s high-level summary is enabled. |
| Title check | ✅ Passed | The title '[Feat] Integrate Z3 in TVM Arith Analyzer' clearly and specifically describes the main feature addition—integrating Z3 SMT solver into TVM's Arith Analyzer—which aligns with the core purpose of this PR as confirmed by the objectives and file changes. |
✨ Finishing touches
- [ ] 📝 Generate docstrings
🧪 Generate unit tests (beta)
- [ ] Create PR with unit tests
- [ ] Post copyable unit tests in a comment
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 help to get the list of available commands and usage tips.
👋 Hi! Thank you for contributing to the TileLang project.
Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.
We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀
Sad, Z3 has a hard-encoded SONAME libz3.so.x.xx. The linker will make TVM dependent on only 1 version of Z3. We have to choose:
- Link Z3 statically: causing binary size 10~20% larger
- Force only 1 version of Z3: may cause conflict with other libraries using Z3
- Manually hack Z3 to make it portable: large engineering effort
I think the best choice may be 1, @LeiWang1999
Sad, Z3 has a hard-encoded SONAME
libz3.so.x.xx. The linker will make TVM dependent on only 1 version of Z3. We have to choose:
libz3.so in z3-solver==4.13.0 has a SONAME "libz3.so", 4.14 and 4.15 have version number in their SONAMEs. We could build tvm with the specific version and then use a looser z3 in runtime.
Also we could use patchelf --replace-needed to modify the DL_NEEDED part of libtvm (libz3.so.x.xx -> libz3.so). This is rather common in python packaging. I could help if needed.
Sad, Z3 has a hard-encoded SONAME
libz3.so.x.xx. The linker will make TVM dependent on only 1 version of Z3. We have to choose:
libz3.soinz3-solver==4.13.0has a SONAME"libz3.so", 4.14 and 4.15 have version number in their SONAMEs. We could build tvm with the specific version and then use a looser z3 in runtime.Also we could use
patchelf --replace-neededto modify theDL_NEEDEDpart of libtvm (libz3.so.x.xx->libz3.so). This is rather common in python packaging. I could help if needed.
Oh, that's soo cool!
I will check z3-solver==4.13.0 in the build environment, and call you if needed!
Thanks very much!
Also we could use
patchelf --replace-neededto modify theDL_NEEDEDpart of libtvm (libz3.so.x.xx->libz3.so). This is rather common in python packaging. I could help if needed.
Actually this looks more robust. Building against specific version of library does not make much sense, especially when you look at the arm & mac failure:
The z3-solver==4.13.0 on arm claims to be manylinux2014, but the libz3.so rely on glibc 2.29, causing the failure.
I didn't dig into the mac failure but it sounds like a similar issue.
For linux arm issue, we may have to embed libz3.so into the wheel if its future version stay buggy.
Z3 doesn't provide a lower version of glibc on ARM; the minimal version is 2.34. We can bundle a Z3 in TileLang. I'll try to choose static linkage or dynamic linkage.
And in MACOS, the libz3.dylib is placed in the bin directory. The inconsistent library path causes an error.
And in MACOS, the libz3.dylib is placed in the
bindirectory. The inconsistent library path causes an error.
We might be able to set different INSTALL_RPATH for macos to resolve this?
@oraluben yes, absolutely. But Z3 in Ubuntu ARM is still a problem. I finally chose to build a static linked library in libtvm.so, which can create a 20% larger library size.
We can post an issue to the Z3 community for their opinion on supporting a lower version of glibc. Once Z3 adds support, we can update the current linkage to use Z3 in the Python site-packages.
I finally chose to build a static linked library in
libtvm.so, which can create a 20% larger library size.
I think the best option is to embedding a shared libz3.so, to prevent the conflict between static lib and shared one, in case any other package depends the shared lib of libz3 (but that shouldn't be possible, right? the libz3.so on arm just don't work).
From perspective of maintenance, we could use different auditwheel command with [[tool.cibuildwheel.overrides]] to craft the lib into wheel.
@oraluben, can you take a look at the macOS build?
Macos pending on: https://github.com/tile-ai/tvm/pull/15 cc @kurisu6912
This PR brings performance improvements for mha_fwd_bhsd_wgmma_pipelined, as well as other cases such as warp_specialize_gemm_copy1_gemm_0.