https://github.com/sergei-grechanik requested changes to this pull request. I'm not sure about handling affine.apply like this because it's recommended to fold affine.apply ops into load/store ops anyway. https://github.com/llvm/llvm-project/pull/77968