[Mlir-commits] [mlir] [mlir][VectorOps] Add `vector.interleave` operation (PR #80315)

[Benjamin Maxwell]

================
@@ -1734,6 +1734,70 @@ struct VectorSplatNdOpLowering : public ConvertOpToLLVMPattern<SplatOp> {
   }
 };
 
+struct VectorInterleaveOpLowering
+    : public ConvertOpToLLVMPattern<vector::InterleaveOp> {
+  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+  void initialize() {
+    // This pattern recursively unpacks one dimension at a time. The recursion
+    // bounded as the rank is strictly decreasing.
+    setHasBoundedRewriteRecursion();
+  }
+
+  LogicalResult
+  matchAndRewrite(vector::InterleaveOp interleaveOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    VectorType resultType = interleaveOp.getResultVectorType();
+
+    // If the result is rank 1, then this directly maps to LLVM.
+    if (resultType.getRank() == 1) {
+      if (resultType.isScalable()) {
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_interleave2>(
+            interleaveOp, typeConverter->convertType(resultType),
+            adaptor.getLhs(), adaptor.getRhs());
+        return success();
+      }
+      // Lower fixed-size interleaves to a shufflevector. While the
+      // vector.interleave2 intrinsic supports fixed and scalable vectors, the
+      // langref still recommends fixed-vectors use shufflevector, see:
+      // https://llvm.org/docs/LangRef.html#id876.
+      int64_t resultVectorSize = resultType.getNumElements();
+      SmallVector<int32_t> interleaveShuffleMask;
+      interleaveShuffleMask.reserve(resultVectorSize);
+      for (int i = 0; i < resultVectorSize / 2; i++) {
+        interleaveShuffleMask.push_back(i);
+        interleaveShuffleMask.push_back((resultVectorSize / 2) + i);
+      }
+      rewriter.replaceOpWithNewOp<LLVM::ShuffleVectorOp>(
+          interleaveOp, adaptor.getLhs(), adaptor.getRhs(),
+          interleaveShuffleMask);
+      return success();
+    }
+
+    // It's not possible to unroll a scalable dimension.
+    if (resultType.getScalableDims().front())
----------------
MacDue wrote:

This is correct :)
This just looks at the scalable flag for the front dimension (which is the dimension that is unrolled below), and checks if it is true (i.e. scalable). If it is scalable it can't be unrolled. 



https://github.com/llvm/llvm-project/pull/80315