[mlir] NFC Replace some code snippets with equivalent method calls

Replace some code snippets With scf::ForOp methods. Additionally,
share a listener at one more point (although this pattern is still
not safe to roll back currently)

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