[RFC] Stride caching in registers on the GPU

[moderato]

Hello! I encountered a problem about scheduling stride caching in registers on the GPU. I discussed it in the following post with @maplegu.

How to avoid allocating the wrong amount of registers in CUDA scheduling in this example?

Hello, I found sometimes registers are allocated more than expected in CUDA code generation. It usually happens when data is cached in registers and read to shared memory chunk by chunk. Here’s a simple example resembling an actual problem I recently encountered. I’m aiming at reading the whole A tensor and caching the data in registers, then calculate B in the shared memory through blocking with this data, and write to C as the output. import tvm def schedule(A, B, C): s = tvm.create_sch…

Looks like the "warp” option for cache_read is not fully working. I’m also thinking if we can use “local” + virtual threads here as well. Currently this also doesn’t work as it infers wrong array indices. Not sure if this is a bug or I’m using the schedule primitives in a wrong way.

Any advice on how to address this problem is appreciated!

@tqchen @Laurawly @FrozenGene @Huyuwei @YuanLin