fix(tir): convert bool dtype to uint1 in BlockRealize predicates#341
Open
cipher982 wants to merge 5 commits intomlc-ai:mlcfrom
Open
fix(tir): convert bool dtype to uint1 in BlockRealize predicates#341cipher982 wants to merge 5 commits intomlc-ai:mlcfrom
cipher982 wants to merge 5 commits intomlc-ai:mlcfrom
Conversation
MLC local ci setup. Also CI for Windows and macOS building, which may take 90-100 mins. Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
- Revert "[CMake][MSVC] Disable permissive mode for MSVC builds (#16343)" - Skip MSC tests - Disable NNPack and TFLite - Tweak CMAKE_CUDA_ARCHITECTURES
TVM C++ runtime rejects 'bool' dtype with:
CHECK(dtype.is_int() || dtype.is_uint())
Extended fix to handle ALL forms of bool predicates:
- Python bool literals (True/False) → IntImm("uint1", value)
- Constant expressions (const, IntImm) → IntImm("uint1", value)
- Dynamic expressions (EQ, Not, Or, And) → Cast("uint1", predicate)
Critical: Dynamic predicates MUST be cast, not replaced with constant
True. This preserves runtime logic for guarded blocks that depend on
loop indices or runtime conditions.
Implementation note:
str(predicate.dtype) is used because TVM's dtype attribute returns a
DataType object rather than a string.
Fixes compilation failures for WebGPU/WASM targets in MLC-LLM when
BlockRealize is constructed with boolean predicates.
Testing:
- Comprehensive regression tests with pytest parameterization
- Tests cover Python bool, tir.const, tir.IntImm (constants)
- Tests cover EQ, Not, Or, And (dynamic predicates)
- Verifies dynamic predicates are cast, not replaced
16a61fc to
033e0cb
Compare
Author
851d153 to
da3714d
Compare
90ab5d3 to
b3d4fe9
Compare
2ef52fe to
bfd7787
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
TVM's C++ runtime rejects
"bool"as a dtype for IntImm expressions:However,
tir.BlockRealize.__init__acceptspredicate: Union[PrimExpr, bool]per the documented API, and various code paths create predicates with dtype"bool":True/Falsetir.const(True, "bool")tir.IntImm("bool", 1)tir.EQ(i, j),tir.Not(flag), etc.This causes runtime errors during compilation, particularly for WebGPU/WASM targets in MLC-LLM.
Solution
Modified
BlockRealize.__init__inpython/tvm/tir/stmt.pyto convert any"bool"dtype predicate to TVM's canonical boolean representation (uint1):IntImm("uint1", value)Cast("uint1", predicate)to preserve logicTesting
Added comprehensive regression tests in
tests/python/tir-base/test_tir_block_realize_predicate.py: