From 41a8152d967eca2d3d313dbbc41bc98c9d158c49 Mon Sep 17 00:00:00 2001
From: Han-Chung Wang <hanhan0912@gmail.com>
Date: Mon, 13 Nov 2023 12:10:22 -0800
Subject: [PATCH] [CPU] Optimize vector transfers for pack/unpack  before
 vector lowering (#15465)

Fixes https://github.com/openxla/iree/issues/15349
---
 compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
index d299285a1f2d..4310da3073e6 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -646,6 +646,9 @@ void addCPUDataTilingPipeline(OpPassManager &passManager,
   addBufferizePasses(nestedModulePM);
 
   {
+    // Fold unit dims before vector lowering.
+    nestedModulePM.addNestedPass<func::FuncOp>(
+        createOptimizeVectorTransferPass(/*flatten=*/false));
     LLVMCPUVectorLoweringPassOptions options;
     options.splitVectorTransfersTo = "linalg-copy";
     nestedModulePM.addNestedPass<func::FuncOp>(