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

[arun-thmn]

This PR covers the mlir::vector::populateFlattenVectorTransferPatterns as a transform pass.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Arun Thangamani (arun-thmn)

Changes

This PR covers the mlir::vector::populateFlattenVectorTransferPatterns as a transform pass.

---

Full diff: https://github.com/llvm/llvm-project/pull/178134.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td (+11)
• (modified) mlir/include/mlir/Dialect/X86Vector/Transforms.h (+5)
• (modified) mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp (+7)

diff --git a/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td b/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
index 891829fca017f..4c953abc125bc 100644
--- a/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
@@ -71,6 +71,17 @@ def ApplyShuffleVectorFMAOpsPatternsOp : Op<Transform_Dialect,
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyFlattenVectorTransferOpsPatternsOp : Op<Transform_Dialect,
+    "apply_patterns.x86vector.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 assemblyFormat = "attr-dict";
+}
+
 
 #endif // X86VECTOR_TRANSFORM_OPS
 
diff --git a/mlir/include/mlir/Dialect/X86Vector/Transforms.h b/mlir/include/mlir/Dialect/X86Vector/Transforms.h
index aadca92708908..3c73a2e172487 100644
--- a/mlir/include/mlir/Dialect/X86Vector/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86Vector/Transforms.h
@@ -104,6 +104,11 @@ void populateSinkVectorProducerOpsPatterns(RewritePatternSet &patterns);
 // grouped with respect to odd/even packed index.
 void populateShuffleVectorFMAOpsPatterns(RewritePatternSet &patterns);
 
+// Rewrites contiguous row-major vector.transfer_read or vector.transfer_write
+// operations by inserting  a memref.collapse_shape on the source,
+// transforming the operation to use a 1D source.
+void populateFlattenVectorTransferOpsPatterns(RewritePatternSet &patterns);
+
 //===----------------------------------------------------------------------===//
 /// Helpers extracted from:
 ///   - clang/lib/Headers/avxintrin.h
diff --git a/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp b/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
index c6be69305da50..cfc6c4194f5af 100644
--- a/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
@@ -12,6 +12,7 @@
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/X86Vector/Transforms.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 
@@ -21,6 +22,7 @@
 using namespace mlir;
 using namespace mlir::x86vector;
 using namespace mlir::transform;
+using namespace mlir::vector;
 
 void mlir::transform::ApplyVectorContractToFMAPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
@@ -47,6 +49,11 @@ void mlir::transform::ApplyShuffleVectorFMAOpsPatternsOp::populatePatterns(
   x86vector::populateShuffleVectorFMAOpsPatterns(patterns);
 }
 
+void mlir::transform::ApplyFlattenVectorTransferOpsPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  vector::populateFlattenVectorTransferPatterns(patterns);
+}
+
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//

[arun-thmn]

@adam-smnk @rengolin Small patch to expose vector::populateFlattenVectorTransferPatterns as a transform pattern. This is the test-vector-transfer-flatten-patterns flag that we are discussing.

[rolfmorel]

Looks good. Though can we get at least one test to ensure that this little bit of code is exercised as well?

[arun-thmn]

@adam-smnk Any more comments? Can it be Merged?

[github-actions]

🪟 Windows x64 Test Results

• 3456 tests passed
• 413 tests skipped

✅ The build succeeded and all tests passed.

[github-actions]

🐧 Linux x64 Test Results

• 7465 tests passed
• 600 tests skipped

✅ The build succeeded and all tests passed.