We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 613e662 commit 25564e3Copy full SHA for 25564e3
python/tvm/topi/arm_cpu/injective.py
@@ -16,7 +16,6 @@
16
# under the License.
17
# pylint: disable=invalid-name, unused-variable
18
"""Schedule for pooling operators"""
19
-import numpy as np
20
import tvm
21
from tvm import te
22
from ..utils import is_empty_shape
@@ -69,7 +68,8 @@ def schedule_injective(outs):
69
68
if list(s[x].op.axis):
70
# do not vectorize for broadcast
71
dtype = "uint16" if x.dtype == "bfloat16" else x.dtype
72
- (io, ii) = s[x].split(list(s[x].op.axis)[-1], 16 // (tvm.DataType(dtype).bits // 8))
+ itemsize = max(1, tvm.DataType(dtype).bits // 8)
+ (io, ii) = s[x].split(list(s[x].op.axis)[-1], 16 // itemsize)
73
s[x].vectorize(ii)
74
tvm.te.schedule.AutoInlineInjective(s)
75
0 commit comments