tvm
tvm copied to clipboard
[Bug] rfactor:use split&rfactor for same stage, will coredump in tvm 0.6
trafficstars
I have a problem today, maybe is a bug in tvm 0.6, my english is poor, i will try my best to describe the problem clearly ^ ^
the test py :
shape_in = (4, 64, 32, 16)
dtype = "float16"
ph_0 = tvm.placeholder(shape_in, dtype=dtype, name="ph_0")
reduce_axis_1 = tvm.reduce_axis((0, 64), name="reduce_axis_1")
reduce_axis_3 = tvm.reduce_axis((0, 16), name="reduce_axis_3")
shape_out = (4, 1, 32, 1)
reduce_0 = tvm.compute(shape_out,
lambda *i: tvm.sum(ph_0[i[0], reduce_axis_1, i[2], reduce_axis_3],
(reduce_axis_1, reduce_axis_3)),
name="reduce_0")
sch = tvm.create_schedule([reduce_0.op])
cache_read_buffer = sch.cache_read(ph_0, "local.UB", [reduce_0])
blk_outer, blk_inner = sch[reduce_0].split(reduce_0.op.axis[0], factor=2)
rfactor_buffer = sch.rfactor(reduce_0, reduce_0.op.reduce_axis[-1], -1)
print(tvm.lower(sch, [ph_0], simple_mode=True))
the error message: terminate called after throwing an instance of 'std::out_of_range' what(): _Map_base::at Aborted (core dumped)
the reason is: in rfactor function, we create .rf stage, this stage has origin stage's root_axis and relations, but not has origin stage's leaf_axis. so, if origin stage's relations has a SplitNode, rf stage will also has. then, in inferbound's passup phase, it will coredump. because we use up_state.at(r->outer), but up_state dont has r->outer.