The limitation on iter_args introduced with D108806 is too restricting. Changes of the runtime type should be allowed. Extends the dim op canonicalization with a simple analysis to determine when it is safe to canonicalize. Differential Revision: https://reviews.llvm.org/D109125