MasterJH5574 opened a new pull request #10175:
URL: https://github.com/apache/tvm/pull/10175
This PR fixes a bug in TVMScript parser. The bug will be triggered by the
following code.
```python
import numpy as np
from tvm.script import tir as T
indptr = np.array([0, 1, 2, 3])
@T.prim_func
def func(A_data: T.Buffer[(13056,), "float32"], B_data: T.Buffer[(131072,),
"float32"], C_data: T.Buffer[(131072,), "float32"], J_indptr: T.Buffer[(33,),
"int32"], J_indices: T.Buffer[(51,), "int32"]) -> None:
for v_vi, v_vbi, v_vf in T.grid(32, 16, 256):
with T.block("bsrmm0"):
vi, vbi, vf = T.axis.remap("SSS", [v_vi, v_vbi, v_vf])
...
```
The code above leads to error
```
Traceback (most recent call last):
File "tensorir.py", line 111, in <module>
def func(A_data: T.Buffer[(13056,), "float32"], B_data:
T.Buffer[(131072,), "float32"], C_data: T.Buffer[(131072,), "float32"],
J_indptr: T.Buffer[(33,), "int32"], J_indices: T.Buffer[(51,), "int32"]) ->
None:
File "/home/rhlai/tvm/python/tvm/script/tir/prim_func.py", line 40, in
prim_func
result = from_source(input_func)
File "/home/rhlai/tvm/python/tvm/script/parser.py", line 1209, in
from_source
namespace = [key for key in env.keys() if env[key] == tir]
File "/home/rhlai/tvm/python/tvm/script/parser.py", line 1209, in
<listcomp>
namespace = [key for key in env.keys() if env[key] == tir]
ValueError: The truth value of an array with more than one element is
ambiguous. Use a.any() or a.all()
```
, which is because previously we used operator `==` to recognize the
`tvm.script.tir` submodule. When `key` is `indptr`, `env[key]` will be
`numpy.ndarray([0, 1, 2, 3])`. Comparing a `numpy.ndarray` with `tir` submodule
leads to the error.
To avoid the error, we are supposed to use operator `is` for the comparison.
---
However, I'm not sure how to add a suitable regression test 😐, and would
like to ask for some idea to add a test.
cc @junrushao1994 @vinx13 @Hzfengsy
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
