[Mlir-commits] [mlir] [mlir][vector] Add vector.to_elements unrolling (PR #157142)

[Jakub Kuderski]

================
@@ -0,0 +1,83 @@
+//===- LowerVectorToElements.cpp - Lower 'vector.to_elements' op ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements target-independent rewrites and utilities to lower the
+// 'vector.to_elements' operation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
+
+#define DEBUG_TYPE "lower-vector-to-elements"
+
+using namespace mlir;
+
+namespace {
+
+struct UnrollToElements : OpRewritePattern<vector::ToElementsOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ToElementsOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> vectors;
+    LogicalResult match =
+        mlir::vector::unrollVectorValue(op.getSource(), rewriter, vectors);
+    if (failed(match)) {
+      return match;
+    }
+
+    // May be large vector.
+    std::vector<Value> results;
+    for (const auto &vector : vectors) {
+      // we need to replace the current result
+      auto subElements =
+          rewriter.create<vector::ToElementsOp>(op.getLoc(), vector);
+      results.insert(results.end(), subElements.getResults().begin(),
+                     subElements.getResults().end());
+    }
+    rewriter.replaceOp(op, results);
+    return success();
+  }
+};
+
+/// Flattens 2 or more dimensional `vector.to_elements` ops by
+/// `vector.shape_cast` + `vector.to_elements`.
+struct FlattenToElements : OpRewritePattern<vector::ToElementsOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ToElementsOp op,
+                                PatternRewriter &rewriter) const override {
+    VectorType vecType = op.getSource().getType();
+    if (vecType.getRank() <= 1)
+      return rewriter.notifyMatchFailure(
+          op, "the rank is already less than or equal to 1");
+    if (vecType.getNumScalableDims() > 0)
+      return rewriter.notifyMatchFailure(
+          op, "scalable vector is not yet supported");
----------------
kuhar wrote:

Does to_elements support scalable vectors? If not, this should probably be an assertion

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