[Bug][TIR] Incorrect transform for the MakePackedAPI pass
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 19h
- Merged PRs (30d)
- 111
Description
### Actual behavior
```
Traceback (most recent call last):
File "/data/qshenaf/remote_pc/TirFuzz/bugs/bug1.py", line 10, in
mod = tir.transform.FP8StorageLegalize()(mod)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/python/tvm/ir/transform.py", line 238, in __call__
return _ffi_transform_api.RunPass(self, mod)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "tvm/_ffi/_cython/./packed_func.pxi", line 339, in tvm._ffi._cy3.core.PackedFuncBase.__call__
File "tvm/_ffi/_cython/./packed_func.pxi", line 270, in tvm._ffi._cy3.core.FuncCall
File "tvm/_ffi/_cython/./packed_func.pxi", line 259, in tvm._ffi._cy3.core.FuncCall3
File "tvm/_ffi/_cython/./base.pxi", line 185, in tvm._ffi._cy3.core.CHECK_CALL
File "/data/qshenaf/envs/tvm/python/tvm/_ffi/base.py", line 468, in raise_last_ffi_error
raise py_err
File "/data/qshenaf/envs/tvm/src/tir/ir/transform.cc", line 121, in tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
func = pass_func(std::move(func), mod, pass_ctx);
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/tir/transforms/unsupported_dtype_legalize.cc", line 796, in operator()
return FP8StorageLegalizer().Legalize(f);
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/tir/transforms/unsupported_dtype_legalize.cc", line 519, in tvm::tir::StorageLegalizer::Legalize(tvm::tir::PrimFunc)
ICHECK_EQ(func->buffer_map.size(), 0) << "This pass must be called after MakePackedAPI";
^^^^^^^^^^^^^^^^^^^^^^^^^^^
tvm.error.InternalError: Traceback (most recent call last):
2: tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
at /data/qshenaf/envs/tvm/src/tir/ir/transform.cc:121
1: operator()
at /data/qshenaf/envs/tvm/src/tir/transforms/unsupported_dtype_legalize.cc:796
0: tvm::tir::StorageLegalizer::Legalize(tvm::tir::PrimFunc)
at /data/qshenaf/envs/tvm/src/tir/transforms/unsupported_dtype_legalize.cc:519
File "/data/qshenaf/envs/tvm/src/tir/transforms/unsupported_dtype_legalize.cc", line 519
InternalError: Check failed: func->buffer_map.size() == 0 (2 vs. 0) : This pass must be called after MakePackedAPI
```
### Environment
TVM-0.21.dev0 (latest)
### Steps to reproduce
```
import tvm
from tvm import te, topi, tir
data = te.placeholder((1, 3, 224, 224), dtype='float32', name='data')
op_output = topi.nn.adaptive_pool(data, output_size=(112, 112), pool_type='max', layout='NCHW')
sch = tir.Schedule(te.create_prim_func([data, op_output]).with_attr("target", tvm.target.Target("llvm")))
mod = tir.transform.MakePackedAPI()(sch.mod)
print(mod)
mod = tir.transform.FP8StorageLegalize()(mod)
```
The IRs after executing the ``.
```
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(data: T.Buffer((1, 3, 224, 224), "float32"), adaptive_pool_max: T.Buffer((1, 3, 112, 112), "float32")):
T.func_attr({"target": T.target({"keys": ["cpu"], "kind": "llvm", "mtriple": "x86_64-unknown-linux-gnu", "tag": ""}), "tir.noalias": T.bool(True)})
# with T.block("root"):
for ax0, ax1, ax2, ax3, rv0, rv1 in T.grid(1, 3, 112, 112, 2, 2):
with T.block("adaptive_pool_max"):
v_ax0, v_ax1, v_ax2, v_ax3, v_rv0, v_rv1 = T.axis.remap("SSSSRR", [ax0, ax1, ax2, ax3, rv0, rv1])
data_1 = T.Buffer((150528,), data=data.data)
T.reads(data_1[v_ax0 * 150528 + v_ax1 * 50176 + v_ax2 * 448 + v_rv0 * 224 + v_ax3 * 2 + v_rv1:v_ax0 * 150528+ v_ax1 * 50176 + v_ax2 * 448 + v_rv0 * 224 + v_ax3 * 2 + v_rv1 + (v_ax0 * 150528 + v_ax1 * 50176 + v_ax2 * 448 + v_rv0 * 224 + v_ax3 * 2 + v_rv1 + 1 - (v_ax0 * 150528 + v_ax1 * 50176 + v_ax2 * 448 + v_rv0 * 224 + v_ax3 * 2 + v_rv1))])
adaptive_pool_max_1 = T.Buffer((37632,), data=adaptive_pool_max.data)
T.writes(adaptive_pool_max_1[v_ax0 * 37632 + v_ax1 * 12544 + v_ax2 * 112 + v_ax3:v_ax0 * 37632 + v_ax1 * 12544 + v_ax2 * 112 + v_ax3 + (v_ax0 * 37632 + v_ax1 * 12544 + v_ax2 * 112 + v_ax3 + 1 - (v_ax0 * 37632 + v_ax1 * 12544 + v_ax2 * 112 + v_ax3))])
T.block_attr({"schedule_rule": "meta_schedule.adaptive_pool_max"})
with T.init():
adaptive_pool_max_1[v_ax0 * 37632 + v_ax1 * 12544 + v_ax2 * 112 + v_ax3] = T.float32(-340282346638528859811704183484516925440.0)
adaptive_pool_max_1[v_ax0 * 37632 + v_ax1 * 12544 + v_ax2 * 112 + v_ax3] = T.max(adaptive_pool_max_1[v_ax0 *37632 + v_ax1 * 12544 + v_ax2 * 112 + v_ax3], data_1[v_ax0 * 150528 + v_ax1 * 50176 + v_ax2 * 448 + v_rv0 * 224 + v_ax3 * 2 + v_rv1])
```
### Triage
Please refer to the list of label tags [here](https://github.com/apache/tvm/wiki/Issue-Triage-Labels) to find the relevant tags and add them below in a bullet format (example below).
* needs-triage
cc @Hzfengsy @junrushao @quic-sanirudh @shingjan
Contributor guide
No contributing guide indexed for this repository
Research direction
Reproduce the sequence in the issue, applying MakePackedAPI before FP8StorageLegalize. Read src/tir/transforms/unsupported_dtype_legalize.cc around StorageLegalizer::Legalize and the FP8StorageLegalize entry point, then inspect how the generated IR retains its buffer map. Done means the reported sequence no longer raises the buffer_map assertion, with regression coverage for this case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100