[ModuloSchedule] Pass loop block explicitly to kernel rewriter.

This change is NFC upstream. We pass in the loop's block to the kernel
rewriter explicitly, instead of assuming it's the loop's top block. This
change is made for downstream targets where this assumption doesn't hold.

Differential Revision: https://reviews.llvm.org/D104811