https://github.com/akroviakov approved this pull request. LGTM for end-to-end compilation. Semantically, multiplying unit matrices with unit scales is not very telling about layout/distribution correctness. https://github.com/llvm/llvm-project/pull/201680