[Feat] Integrate Z3 in TVM Arith Analyzer #1367
Conversation
WalkthroughUpdates the TVM submodule pointer, integrates Z3 (build/runtime and Python dependency), adds TileLang arith sources to CMake, extends RPATH handling for z3, introduces many new arithmetic and parameterized GEMM tests, removes one obsolete test helper, and makes a small utility cleanup. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
⏰ 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). (5)
🔇 Additional comments (2)
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
There was a problem hiding this comment.
Please double-check the modification on the build system
The other modification on the build system can be found in 3rdparty/tvm/CMakeLists.txt.
There was a problem hiding this comment.
The build system changes LGTM (including 3rdparty/tvm/CMakeLists.txt).
BTW now tvm add python and z3 libs as a non-optional build time dependency, is that what we want?
There was a problem hiding this comment.
Now, Z3 is an optional dependency of TVM (but I don't test building TVM without Tilelang). And Tilelang depends on Z3 by setting set(USE_Z3 "pypi")
|
|
||
| # Collect source files | ||
| file(GLOB TILE_LANG_SRCS | ||
| src/arith/*.cc |
There was a problem hiding this comment.
This should be removed
There was a problem hiding this comment.
The build system changes LGTM (including 3rdparty/tvm/CMakeLists.txt).
BTW now tvm add python and z3 libs as a non-optional build time dependency, is that what we want?
pyproject.toml
Outdated
| "torch>=2.7; platform_system == 'Darwin'", | ||
| "tqdm>=4.62.3", | ||
| "typing-extensions>=4.10.0", | ||
| "z3-solver>=4.13.0" |
There was a problem hiding this comment.
trivial update: add comma after each line
| "z3-solver>=4.13.0" | |
| "z3-solver>=4.13.0", |
![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)
136-136: Clarify or address the past review comment.There's an existing comment from kurisu6912 stating "This should be removed" on this line. If the
src/arith/*.ccsources are intentionally needed for Z3 integration, please resolve or dismiss that comment. If it was meant for a different change, the comment should be clarified.
🧹 Nitpick comments (8)
testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py (2)
41-69: Consider deleting or updating this code rather than commenting it out.The comment from the previous review indicates that "After integrating Z3, this access can be proved safe." If this is the case, consider one of the following approaches:
- Delete the commented code entirely if the test is no longer needed
- Update the test to verify that Z3 can now prove the access is safe (remove the boundary check from the expected output)
Leaving commented-out code adds noise and maintenance burden.
148-152: Same concern: prefer deletion over commenting out test code.This test invocation should either be deleted along with the helper function, or updated to test the improved Z3 behavior.
testing/python/arith/test_arith_iter_affine_map.py (2)
733-735: Remove commented-out dead code.The commented variable
zis unused and should be removed.x = tvm.tir.Var("x", "int32") y = tvm.tir.Var("y", "int32") - # z = tvm.tir.Var("z", "int32")
1145-1150: Confusing variable naming:aandbreuse names "x" and "y".The variables
aandbare constructed with names"x"and"y"respectively, which matches the names of the earlierxandyvariables. While this may be intentional for the test case (testing overlapped iterations with same-named variables), it could cause confusion during debugging.Consider using distinct names:
- a = tvm.tir.Var("x", "int32") - b = tvm.tir.Var("y", "int32") + a = tvm.tir.Var("a", "int32") + b = tvm.tir.Var("b", "int32")testing/python/arith/test_arith_simplify.py (2)
18-21: Missing explicitimport tvm.testing.Line 39 uses
tvm.testing.parameter(...), buttvm.testingis not explicitly imported. While this may work iftilelang.tvmre-exports the testing module, it differs fromtest_arith_intset.pywhich explicitly importsimport tvm.testing. For consistency and clarity, add the explicit import.from tilelang import tvm import tilelang.testing from tvm import tir import tvm.ir +import tvm.testing
39-51: Float dtypes included in parameterized tests.The PR description states Z3 integration targets integer expressions only. Including float dtypes (float16, float32, float64) is fine for general analyzer coverage, but these cases will not exercise the Z3 backend. Consider adding a comment clarifying this distinction, or if Z3-specific testing is intended, use a separate integer-only parameter set.
testing/python/arith/test_arith_intset.py (2)
177-179: Minor: prefer generator expression over list comprehension inall().Using a generator expression is slightly more efficient as it short-circuits without building the full list.
if result is None: - assert all([_ is None for _ in expect]) + assert all(e is None for e in expect) return
372-378: Consider explicit variable assignment for clarity.The pattern
i, j = var_domunpacks dict keys (theVarobjects) relying on insertion order. While correct in Python 3.7+, explicit assignment would be clearer and less surprising to readers.def test_region_lower_bound_unfusable(): + i = tvm.tir.Var("i", "int32") + j = tvm.tir.Var("j", "int32") var_dom = { - tvm.tir.Var("i", "int32"): tvm.ir.Range(8), - tvm.tir.Var("j", "int32"): tvm.ir.Range(4), + i: tvm.ir.Range(8), + j: tvm.ir.Range(4), } - i, j = var_dom check_region_bound({(i + j) // 2: (0, 6)}, var_dom, mode="lowerbound")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
3rdparty/tvm(1 hunks)CMakeLists.txt(3 hunks)pyproject.toml(2 hunks)testing/python/arith/test_arith_hard.py(1 hunks)testing/python/arith/test_arith_intset.py(1 hunks)testing/python/arith/test_arith_iter_affine_map.py(1 hunks)testing/python/arith/test_arith_simplify.py(1 hunks)testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py(4 hunks)testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py(5 hunks)testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py(2 hunks)tilelang/utils/sparse.py(0 hunks)
💤 Files with no reviewable changes (1)
- tilelang/utils/sparse.py
🧰 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/transform/test_tilelang_transform_legalize_safe_memory_access.pytesting/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py
🧬 Code graph analysis (4)
testing/python/arith/test_arith_hard.py (1)
tilelang/tools/Analyzer.py (1)
Analyzer(33-218)
testing/python/arith/test_arith_simplify.py (2)
tilelang/tools/Analyzer.py (1)
Analyzer(33-218)tilelang/language/ast/ir.py (1)
Range(1716-1728)
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py (3)
tilelang/testing/__init__.py (3)
requires_cuda_compute_version(38-102)requires_cuda_compute_version_ge(105-106)requires_cuda_compute_version_le(121-122)tilelang/tileop/gemm/__init__.py (2)
trans_A(103-104)trans_B(107-108)tilelang/tileop/gemm_sp/gemm_sp_base.py (2)
trans_A(45-46)trans_B(49-50)
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py (2)
maint/gemm_v2/correctness_evaluation.py (3)
run_gemm_rs(181-213)run_gemm_sr(267-300)run_gemm_rr(357-390)testing/python/amd/test_tilelang_test_amd.py (1)
run_gemm_rs(179-223)
🪛 Ruff (0.14.7)
testing/python/arith/test_arith_hard.py
23-23: Avoid specifying long messages outside the exception class
(TRY003)
95-95: Use raise without specifying exception name
Remove exception name
(TRY201)
⏰ 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). (5)
- GitHub Check: Quick Lint
- GitHub Check: Build SDist
- 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.1
🔇 Additional comments (30)
testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py (3)
1-1: LGTM!The pytest import is correctly added to support the parameterized test decorators.
235-294: Excellent parameterization for SM90 tests.The test cases comprehensively cover multiple data types (float16, float8_e4m3, int8), block configurations, staging options, and transposition combinations. The
trans_Aandtrans_Bparameters are correctly integrated into both the kernel construction and execution.
297-299: Manual verification needed for decorator combination behavior.The combination of
requires_cuda_compute_version_ge(8, 0)andrequires_cuda_compute_version_le(8, 9)is intended to create a version range [8.0, 8.9] for SM80 tests. However, the implementation of these decorators in thetilelang.testingmodule could not be inspected to confirm that pytest combines multiple decorators using AND logic (test runs only if both conditions are true) rather than OR logic. This pattern should be verified manually by reviewing the decorator source code intilelang.testingand confirming how pytest chains them together.testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py (4)
110-159: Well-structured parameterization for skipped SS tests.The test is properly parameterized and will be easy to re-enable once the GEMM SS issues mentioned in the skip reason are resolved. The parameter coverage includes multiple data types and all transposition combinations.
270-319: Well-structured parameterization for skipped RS tests.Similar to
test_gemm_ss, this test is properly parameterized and documented with a skip marker. The test will be straightforward to re-enable once GEMM RS issues are resolved.
430-479: Active SR tests with documented coverage gap.The parameterization is correct, and the TODO comment at line 443 appropriately documents the known float32 coverage issue. This test is active (unlike SS and RS), which indicates SR functionality is working.
594-644: Comprehensive parameterization for RR tests.This test has the most extensive coverage, including bfloat16 in addition to the other data types. The 16 test cases systematically cover various configurations, block sizes, and transposition combinations. The active status (no skip marker) indicates this functionality is stable.
3rdparty/tvm (1)
1-1: Verify the submodule commit reference and z3 branch merge completion.The commit
68f5e91e2c46223ea8b1270711a143fa6799cbc5should be verified to confirm:
- It exists in the TVM repository and is reachable from
tilelang_main- It represents the successful completion of the z3 branch merge
- No breaking API changes affect TileLang consumers
Regarding the z3-solver>=4.13.0 dependency in pyproject.toml: no widely documented TVM-specific incompatibilities exist with recent z3-solver versions, though general installation and platform-specific issues can occur. Verify this specific version pairing has been tested in your build environment.
Please confirm the status of outstanding work items before merging:
- Vectorization bug fix for the
AtomicAddop (referenced as #1366)- Full verification that build configuration changes in CMakeLists.txt and pyproject.toml are complete and compatible
CMakeLists.txt (2)
225-226: LGTM: Z3 integration enabled via PyPI.Setting
USE_Z3to"pypi"aligns with thez3-solverdependency added topyproject.tomland enables the TVM Z3 integration for arithmetic proving.
291-299: LGTM: RPATH updated for Z3 library discovery.The RPATH additions for both Apple (
@loader_path/../../z3/lib) and UNIX ($ORIGIN/../../z3/lib) platforms ensure the z3 shared libraries can be found at runtime alongside the existing tvm_ffi paths.pyproject.toml (2)
54-54: LGTM: z3-solver build-time dependency added.Adding
z3-solver>=4.13.0to build requirements ensures Z3 is available during the build process for TVM's Z3 integration.
43-43: LGTM: z3-solver runtime dependency added.The
z3-solver>=4.13.0dependency enables the new Z3-based arithmetic proving capabilities at runtime. The version constraint is appropriate, allowing flexibility to use versions up to the latest (4.15.4.0).testing/python/arith/test_arith_hard.py (4)
1-5: LGTM: Imports are appropriate for Z3 arithmetic testing.The imports from
tvm.arith,tvm.ir.expr, andtvm.tir.exprare correct for testing the new Z3-based arithmetic analyzer.
8-9: LGTM: Logical implication helper.The
implies(x, y)helper correctly implements logical implication as¬x ∨ y.
12-54: LGTM: Comprehensive arithmetic proof tests.The
test_hard_provefunction tests complex arithmetic implications including:
- Integer division/modulo properties
- Comparison transitivity (with caveats for
d)- Bit manipulation patterns
The
check_exprhelper provides useful debugging output via SMT-LIB2 on failure.
78-95: LGTM: Range binding test.The
test_bindfunction correctly tests the analyzer's ability to prove expressions with bound variables. The exception handling provides useful SMT-LIB2 debugging output.testing/python/arith/test_arith_iter_affine_map.py (4)
1-20: LGTM: License header and imports.The Apache 2.0 license header is appropriate for TVM-derived test code, and the imports correctly reference the tilelang and TVM modules needed for iter-map testing.
23-48: LGTM: Well-documented helper functions.The utility functions
ifuse,isplit,var_dom, andconvert_iter_exprare clearly documented and provide a clean API for constructing test cases.
51-117: LGTM: Comprehensive assertion helpers.The
assert_iter_sum_pattern,assert_iter_map_simplify, andassert_iter_sum_failurefunctions provide thorough validation of iter-map detection results with support for various check levels and options.
120-1422: LGTM: Comprehensive iter-map test coverage.The test suite provides thorough coverage of TVM's arithmetic iter-map analysis including:
- Trivial cases, fusion, and splitting
- Compound patterns and predicates
- Subspace division
- Normalization and inverse affine mappings
- Padding scenarios
- Symbolic cases and edge conditions
This is valuable test coverage for the TVM arithmetic integration.
testing/python/arith/test_arith_simplify.py (4)
86-94: LGTM - regression test pattern is appropriate.The test intentionally has no assertion since it verifies that
rewrite_simplifycompletes without infinite recursion. The test passes by not hanging, which is the correct approach for this type of regression test.
121-122: Inconsistent test runner with sibling test file.This file uses
tilelang.testing.main()whiletest_arith_intset.pyusestvm.testing.main(). For consistency across test files in the same directory, consider aligning the test runner usage.
24-36: LGTM - well-structured index simplification test.The test correctly verifies that the analyzer simplifies a flattened index expression back to its original form, validating the reshape pattern simplification logic.
68-83: LGTM - comprehensive symbolic bound proof tests.Good coverage of symbolic comparison proofs with
ProofStrength.SYMBOLIC_BOUND, testing various inequality forms and boundary conditions.testing/python/arith/test_arith_intset.py (6)
17-21: LGTM - imports are well-organized.Imports are complete and consistent. The explicit
import tvm.testingaligns with thetvm.testing.main()usage at the end of the file.
24-36: LGTM - clean helper class design.The
IntSetCheckerclass provides a reusable pattern for int set verification with informative error messages for debugging test failures.
58-63: LGTM - appropriate vscale handling test.This test validates that scalable vectors with
vscale()correctly producepos_inf()for the max value, which is the expected behavior since vscale is runtime-determined. This aligns with the PR noting that vscales are not supported in Z3 integration.
92-122: LGTM - thorough modulo operation testing.Good coverage of both
truncmodandfloormodoperations with various interval ranges. The use of separateIntSetCheckerinstances (ck vs ck1) correctly isolates different binding contexts.
381-392: LGTM - good edge case coverage for union bounds.The test correctly validates union behavior with infinite bounds, including the edge case where an empty interval (pos_inf to neg_inf) is included in the union.
205-240: LGTM - comprehensive dependent region bound tests.Good coverage of scenarios where region bounds cannot be computed (dependent indices) for lowerbound mode, with corresponding upperbound fallbacks.
| solver = z3.Solver() | ||
| solver.from_string(smtlib2) | ||
| assert solver.check() |
There was a problem hiding this comment.
Potential issue: solver.check() returns CheckSatResult, not a boolean.
In Z3, solver.check() returns z3.sat, z3.unsat, or z3.unknown, not a boolean. While z3.sat is truthy, z3.unknown is also truthy and would pass this assertion unexpectedly.
Consider being explicit about the expected result:
- assert solver.check()
+ assert solver.check() == z3.sat, f"Expected sat, got {solver.check()}"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| solver = z3.Solver() | |
| solver.from_string(smtlib2) | |
| assert solver.check() | |
| solver = z3.Solver() | |
| solver.from_string(smtlib2) | |
| assert solver.check() == z3.sat, f"Expected sat, got {solver.check()}" |
🤖 Prompt for AI Agents
In testing/python/arith/test_arith_hard.py around lines 73 to 75, the code
asserts solver.check() directly but Z3 returns a CheckSatResult (z3.sat /
z3.unsat / z3.unknown), so unknown would incorrectly pass; change the assertion
to explicitly compare the result to z3.sat (e.g. result = solver.check(); assert
result == z3.sat) and only call solver.model() or use the model when result is
z3.sat.
|
Sad, Z3 has a hard-encoded SONAME
I think the best choice may be 1, @LeiWang1999 |
Also we could use |
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! |
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 For linux arm issue, we may have to embed |
3 participants
This pr integrates Z3 into the TVM arith analyzer.
TODO
AtomicAddop ([BUG] Auto Vectorization on AtomicAdd Ignores BufferLayout and LowerArgs #1366)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
Dependencies
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.