[CPU] Optimize transposed masks (#14141)

Make sure we don't transpose a mask but create a transposed mask
instead.

compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp
index 5aec874..45f88ae 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp
@@ -106,14 +106,10 @@
 
   // Make sure we remove redundant vector ops (e.g., vector tranposes) before we
   // lower them and can't be optimized away anymore.
+  // TODO (dcaballe): We should run full canonicalization here.
   {
     RewritePatternSet patterns(ctx);
-    SmallVector<Dialect *> dialects;
-    dialects.push_back(ctx->getLoadedDialect<vector::VectorDialect>());
-    dialects.push_back(ctx->getLoadedDialect<memref::MemRefDialect>());
-    dialects.push_back(ctx->getLoadedDialect<linalg::LinalgDialect>());
-    for (auto &dialect : dialects)
-      dialect->getCanonicalizationPatterns(patterns);
+    vector::TransposeOp::getCanonicalizationPatterns(patterns, ctx);
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   }
 

compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir
index 7ef532a..28b86fc 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir
@@ -125,3 +125,25 @@
 
 // CHECK-LABEL:       func.func @matmul_scalar_load
 // CHECK-COUNT-128:     memref.load
+
+// -----
+
+// Make sure we don't transpose a mask but create a transposed mask instead.
+
+func.func @transpose_mask() {
+  %a = arith.constant 4 : index
+  %b = arith.constant 8 : index
+  %c0 = arith.constant 0 : index
+  %3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) : memref<4x2xi1>
+  %mask = vector.create_mask %a, %b : vector<2x4xi1>
+  %transpose_mask = vector.transpose %mask, [1, 0] : vector<2x4xi1> to vector<4x2xi1>
+  vector.transfer_write %transpose_mask, %3[%c0, %c0] {in_bounds = [true, true]} : vector<4x2xi1>, memref<4x2xi1>
+  return
+}
+
+// CHECK-LABEL: func.func @transpose_mask
+//   CHECK-NOT:   vector.create_mask
+//   CHECK-NOT:   vector.constant_mask [2, 4]
+//   CHECK-NOT:   vector.transpose
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   vector.constant_mask [4, 2] : vector<4x2xi1>