topi/cuda/pooling.py

21 from .. import generic
26 @generic.schedule_adaptive_pool.register(["cuda", "gpu"])
59         s[Out].bind(ty, thread_y)
60         s[Out].bind(tx, thread_x)
61         s[Out].bind(by, block_y)
62         s[Out].bind(bx, block_x)
92 @generic.schedule_pool.register(["cuda", "gpu"])
124         s[Out].bind(bx, tvm.thread_axis("blockIdx.x"))
125         s[Out].bind(tx, tvm.thread_axis("threadIdx.x"))
156 @generic.schedule_pool_grad.register(['cuda', 'gpu'])
182         s[out].bind(bx, tvm.thread_axis("blockIdx.x"))
183         s[out].bind(tx, tvm.thread_axis("threadIdx.x"))