[stealth-rklopper]

Summary

Z3Prover translates bitwise_and, bitwise_or, bitwise_xor and bitwise_not through bit-vectors, also when the operands are booleans. TileLang writes a condition with &, such as (bx % 10 < 2) & (tk > 0), as a bitwise_and of two booleans. As a constraint, its bit-vector terms use up the rlimit budget of the solver, and an easy proof under it returns unknown.

ThreadSync of TileLang aborts the compile with LOG(FATAL) in two such cases:

• It cannot prove that a guard of 448 of 512 threads is warp-uniform under an enclosing & condition ([BUG][ice-on-valid-code] ThreadSync rejects a 448-thread partial barrier as not warp-uniform under an & condition instead of emitting it tilelang#3653).
• It cannot prove that the threads below 128 write disjoint words of shared memory after a binding inside = (slot >= 0) & (slot < pages). It then places a barrier inside the guard tx < 128, and it cannot prove that this guard is warp-uniform either.

With T.And in place of &, the proofs hold. This PR translates a bitwise operation of booleans to its logical operation, so both forms give the same proofs.

Reproduction and cause

import tvm
from tvm import tirx

tx1, tx2, pages = (tirx.Var(name, "int32") for name in ("tx1", "tx2", "pages"))
bindings = []
for name in ("1", "2"):
    slot = tirx.Var("slot" + name, "int32")
    inside = tirx.Var("inside" + name, "bool")
    bindings.append(inside == tirx.bitwise_and(slot >= 0, slot < pages))
condition = tirx.all(tx1 != tx2, 0 <= tx1, tx1 < 128, 0 <= tx2, tx2 < 128, *bindings)
ana = tvm.arith.Analyzer()
with ana.constraint_scope(condition):
    print(ana.can_prove(tx1 // 32 * 256 + tx1 % 32 != tx2 // 32 * 256 + tx2 % 32 + 32))
# Before this change: False. With tirx.And in place of tirx.bitwise_and: True.

The binding of each thread instance is the form in which ThreadSync passes the constraints of a statement to the solver. VisitBitwiseOp converts each operand with int2bv and the result with bv2int(..., true), so each boolean becomes a 1-bit vector. The claim itself has no bitwise operation, but the solver gives up on it within its budget of 10,000 steps. The prover sets that budget with SetRLimit(1e4), so a longer timeout gives the same result.

Changes

• VisitExpr_(const CallNode *) translates bitwise_and, bitwise_or, bitwise_xor and bitwise_not with a scalar boolean result to &&, ||, != and ! of the boolean operands.
• Integer bitwise operations keep their bit-vector translation.
• test_z3_proves_under_bitwise_and_of_booleans in tests/python/arith/test_arith_simplify.py checks the proof above. It fails before this change and passes after it.

Validation

• python -m pytest tests/python/arith -q gives 1019 passed, 21 failed, 1 skipped and 8 xfailed. Without this change it gives 1018 passed and 22 failed. The 21 failures are the same on both builds: the vscale and SVE cases, test_domain_touched, test_domain_touched_vector and one case of TestFloormodIndex. The 22nd failure is the new test.
• With this change in 3rdparty/tvm, TileLang lowers the reproduction of [BUG][ice-on-valid-code] ThreadSync rejects a 448-thread partial barrier as not warp-uniform under an & condition instead of emitting it tilelang#3653. Without it, the reproduction fails with the error of the issue. The 203 tests of testing/python/transform/test_tilelang_transform_thread_sync.py and testing/python/arith pass with it. Both TileLang builds are 80364b3 of main with [BugFix][CUDA] Pack a float32 broadcast to 8 lanes without the address of a temporary tilelang#3681, [BugFix][Transform] Fix divergent and missing barriers in thread_storage_sync tilelang#3682 and thesm-os/tilelang@fc5c79b merged. These commits change three files of TileLang and do not change 3rdparty/tvm.
• ruff check and ruff format --check of ruff 0.12.3 pass on test_arith_simplify.py. clang-format 20.1.8 leaves the added lines of z3_prover_on.cc unchanged. It also reformats many lines of the file that this PR does not touch, and the PR leaves those lines as they are.
• git diff --check passes.

A TileLang PR follows with the bump of 3rdparty/tvm and regression tests for both cases of ThreadSync.