tilelang icon indicating copy to clipboard operation
tilelang copied to clipboard

[Feat] Integrate Z3 in TVM Arith Analyzer

Open kurisu6912 opened this issue 3 weeks ago • 2 comments

This pr integrate z3 in tvm arith analyzer

kurisu6912 avatar Dec 02 '25 10:12 kurisu6912

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 narrowed cumsum_fragment signature.

Possibly related PRs

  • tile-ai/tilelang#1426 — touches cumsum_fragment handling; opposite direction of signature change (related to Buffer/Region handling).
  • tile-ai/tilelang#1358 — modifies cumsum/reduce paths in reduce.py; overlaps on same function.
  • tile-ai/tilelang#1334 — also changes cumsum_fragment and 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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

coderabbitai[bot] avatar Dec 02 '25 10:12 coderabbitai[bot]

👋 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! 🚀

github-actions[bot] avatar Dec 02 '25 10:12 github-actions[bot]

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:

  1. Link Z3 statically: causing binary size 10~20% larger
  2. Force only 1 version of Z3: may cause conflict with other libraries using Z3
  3. Manually hack Z3 to make it portable: large engineering effort

I think the best choice may be 1, @LeiWang1999

kurisu6912 avatar Dec 09 '25 04:12 kurisu6912

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.

oraluben avatar Dec 09 '25 06:12 oraluben

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.

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!

kurisu6912 avatar Dec 10 '25 08:12 kurisu6912

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.

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.

oraluben avatar Dec 10 '25 11:12 oraluben

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.

image

kurisu6912 avatar Dec 12 '25 05:12 kurisu6912

And in MACOS, the libz3.dylib is placed in the bin directory. The inconsistent library path causes an error.

We might be able to set different INSTALL_RPATH for macos to resolve this?

oraluben avatar Dec 12 '25 13:12 oraluben

@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.

kurisu6912 avatar Dec 13 '25 16:12 kurisu6912

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 avatar Dec 15 '25 00:12 oraluben

@oraluben, can you take a look at the macOS build?

kurisu6912 avatar Dec 15 '25 13:12 kurisu6912

Macos pending on: https://github.com/tile-ai/tvm/pull/15 cc @kurisu6912

oraluben avatar Dec 15 '25 14:12 oraluben

This PR brings performance improvements for mha_fwd_bhsd_wgmma_pipelined, as well as other cases such as warp_specialize_gemm_copy1_gemm_0. speedup_bar

yyttt6 avatar Dec 17 '25 11:12 yyttt6