From 99fe7c6b6ad23db0d1ea57a6f7be654e76eff081 Mon Sep 17 00:00:00 2001
From: Diego Caballero <diegocaballero@google.com>
Date: Wed, 31 May 2023 17:55:45 -0700
Subject: [PATCH] [CPU] Add shape cast lowering to LLVMCPUMmt4dVectorLowering
 (#13874)

We missed the 'vector.shape_cast' lowering patterns in LLVMCPUMmt4dVectorLowering.
---
 .../Codegen/LLVMCPU/LLVMCPUMmt4dVectorLowering.cpp     | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMmt4dVectorLowering.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMmt4dVectorLowering.cpp
index 2159090e47cda..ec0c013ca22b3 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMmt4dVectorLowering.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMmt4dVectorLowering.cpp
@@ -162,6 +162,16 @@ void LLVMCPUMmt4dVectorLoweringPass::runOnOperation() {
       return signalPassFailure();
     }
   }
+
+  // 'vector.shape_cast' are very expensive operations that are even generated
+  // by some of the lowerings above (e.g., flatten transfer ops). There are
+  // chances to cancel them out if they are not lowered too early so we lower
+  // them at the very end of the pass.
+  {
+    RewritePatternSet patterns(&getContext());
+    vector::populateVectorShapeCastLoweringPatterns(patterns);
+    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  }
 }
 
 std::unique_ptr<OperationPass<func::FuncOp>>