Lunderberg commented on issue #12852:
URL: https://github.com/apache/tvm/issues/12852#issuecomment-1253808768
I can reproduce the error, and made a simplified test script below. Playing
around a bit with the shape, it looks like it only occurs when there's a
dimension of size 1.
```python
#!/usr/bin/env python3
import tvm
import tvm.testing
from tvm.script import tir as T
shape = tvm.testing.parameter(
(1, 1, 256, 512),
(1, 2, 256, 512),
(2, 1, 256, 512),
(2, 2, 256, 512),
)
def test_rewrite_layout(shape):
n, c, h, w = shape
@tvm.script.ir_module
class mod:
@T.prim_func
def main(a: T.handle):
A = T.match_buffer(a, shape, "float32")
T.func_attr({"layout_free_buffers": [0]})
for ax0, ax1, ax2, ax3 in T.grid(n, c, h, w):
with T.block():
v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2,
ax3])
T.evaluate(A[v0, v1, v2, v3])
sch = tvm.tir.Schedule(mod)
tvm.meta_schedule.postproc.RewriteLayout().apply(sch)
if __name__ == "__main__":
tvm.testing.main()
```
--
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]
