[Bug] Failing Compute-At Primitive Usage

[zxybazh]

When scheduling a winograd conv2d workload I found that the comput-at schedule primitive might fail in some case.

Error message: The primitive requires all the consumer(s) of the given block to be present under the target loop. However, there are 1 consumer(s) not satisfying the constraint. List of the consumer(s):tir.Block#0

Script available here: https://gist.github.com/zxybazh/4e6b2ff3dca7ca66254ea1e553103cfd

[junrushao]

CC @spectrometerHBH

[zxybazh]

Updated with another example with triggers the failed check.
https://gist.github.com/zxybazh/81a76e6d5502e8a200eed190ec361b28

Now the script is working.

[junrushao]

seems like the same type of error. BTW, let's use vthread.x instead of vthread

[zxybazh]

Fixed.

[Hzfengsy]

The TVMScript is incorrect.

https://gist.github.com/zxybazh/81a76e6d5502e8a200eed190ec361b28#file-test_compute_at_cuda-py-L56
The block "input_tile" does not have block vars

[junrushao]

@zxybazh could you fix the TVMScript accordingly?

[zxybazh]

I have fixed the scripts and the cuda script is working while the cpu script is throwing the same error.

[Hzfengsy]

Can you provide the CPU script?

[zxybazh]

Hi Siyuan, the cpu script is here https://gist.github.com/zxybazh/4e6b2ff3dca7ca66254ea1e553103cfd

[Hzfengsy]

https://gist.github.com/zxybazh/4e6b2ff3dca7ca66254ea1e553103cfd#file-test_compute_at-py-L178

Here we don't need to get consumer but directly

b78 = sch.get_block(name="input_tile")
l80 = sch.sample_compute_location(block=b78, decision=xxx)
sch.compute_at(block=b78, loop=l80, preserve_unit_loops=True)

[zxybazh]

Got it. It was caused because of change of meaning for sample_compute_location, it automatically finds consumers and sample the compute locations. Now the script can run without any problem. Thanks everyone.