This re-applies f905bf3e1e, which was reverted in c861c1a046 due to compiler errors, with a fix for MLIR.
f905bf3e1e
c861c1a046