[mlir][vector] Wrapping populateFlattenVectorTransferPatterns as a transform pass.

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

@@ -539,4 +539,22 @@ def ApplySinkVectorMemPatternsOp : Op<Transform_Dialect,
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyFlattenVectorTransferOpsPatternsOp : Op<Transform_Dialect,
+    "apply_patterns.vector.flatten_vector_transfer_ops",
+    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Collect patterns to rewrite contiguous row-major vector.transfer_read or 
+    vector.transfer_write operations to a 1D operation.
+  }];
+
+  let arguments = (ins
+  DefaultValuedAttr<UI32Attr,
+    "std::numeric_limits<unsigned>::max()">:$target_vector_bitwidth
+  );
+
+  let assemblyFormat = [{
+    (`target_vector_bitwidth` `=` $target_vector_bitwidth^)? attr-dict
+  }];
+}
+
 #endif // VECTOR_TRANSFORM_OPS

mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td

@@ -71,6 +71,5 @@ def ApplyShuffleVectorFMAOpsPatternsOp : Op<Transform_Dialect,
   let assemblyFormat = "attr-dict";
 }
 
-
 #endif // X86VECTOR_TRANSFORM_OPS
 

mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

@@ -227,6 +227,12 @@ void transform::ApplySinkVectorMemPatternsOp::populatePatterns(
   vector::populateSinkVectorMemOpsPatterns(patterns);
 }
 
+void mlir::transform::ApplyFlattenVectorTransferOpsPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  vector::populateFlattenVectorTransferPatterns(patterns,
+                                                getTargetVectorBitwidth());
+}
+
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp

@@ -21,6 +21,7 @@
 using namespace mlir;
 using namespace mlir::x86vector;
 using namespace mlir::transform;
+using namespace mlir::vector;
 
 void mlir::transform::ApplyVectorContractToFMAPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {

mlir/test/Dialect/Vector/flatten-vector-transfer-ops.mlir

@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s
+
+func.func @flatten_transfer_ops(%arg0: memref<16x16xf32>, %arg1: vector<8xf32>) -> vector<8xf32> {
+  %c0 = arith.constant 0 : index
+  %c8 = arith.constant 8 : index
+  %b0 = ub.poison : f32
+  %0 = vector.transfer_read %arg0[%c0, %c0], %b0 {in_bounds = [true, true]} : memref<16x16xf32>, vector<1x8xf32>
+  %1 = vector.transfer_read %arg0[%c0, %c8], %b0 {in_bounds = [true, true]} : memref<16x16xf32>, vector<1x8xf32>
+  %2 = vector.shape_cast %0 : vector<1x8xf32> to vector<8xf32>
+  %3 = vector.shape_cast %1 : vector<1x8xf32> to vector<8xf32>
+  %4 = vector.fma %2, %3, %arg1 : vector<8xf32>
+  return %4 : vector<8xf32>
+}
+
+// CHECK-LABEL: @flatten_transfer_ops
+// CHECK-NOT: vector.transfer_read {{.*}},  vector<1x8xf32>
+// CHECK-NOT: vector.transfer_read {{.*}},  vector<1x8xf32>
+// CHECK: vector.transfer_read {{.*}},  vector<8xf32> 
+// CHECK-NEXT: vector.transfer_read {{.*}},  vector<8xf32>
+// CHECK-NOT: vector.shape_cast
+// CHECK-NOT: vector.shape_cast
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %func {
+      transform.apply_patterns.vector.flatten_vector_transfer_ops
+    } : !transform.any_op
+    transform.yield
+  }
+}

mlir/test/python/dialects/transform_vector_ext.py

@@ -54,6 +54,8 @@ def non_configurable_patterns():
     vector.ApplyLowerScanPatternsOp()
     # CHECK: transform.apply_patterns.vector.lower_shape_cast
     vector.ApplyLowerShapeCastPatternsOp()
+    # CHECK: transform.apply_patterns.vector.flatten_vector_transfer_ops
+    vector.ApplyFlattenVectorTransferOpsPatternsOp()
 
 
 @run_apply_patterns