[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/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

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

mlir/test/Dialect/Vector/transform-vector.mlir

@@ -137,3 +137,35 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+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

@@ -67,6 +67,9 @@ def configurable_patterns():
     # CHECK-SAME: max_transfer_rank = 3
     # CHECK-SAME: full_unroll = true
     vector.ApplyTransferToScfPatternsOp(max_transfer_rank=3, full_unroll=True)
+    # CHECK: transform.apply_patterns.vector.flatten_vector_transfer_ops
+    # CHECK-SAME: target_vector_bitwidth = 1
+    vector.ApplyFlattenVectorTransferOpsPatternsOp(target_vector_bitwidth=1)
 
 
 @run_apply_patterns