tinytile/Vectorization.cpp at main · Groverkss/tinytile · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
#include "Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir::tutorial {

#define GEN_PASS_DEF_TUTORIALVECTORIZATION
#include "Passes.h.inc"

namespace {

class TutorialVectorization final
    : public impl::TutorialVectorizationBase<TutorialVectorization> {
  using TutorialVectorizationBase::TutorialVectorizationBase;
  void runOnOperation() override;
};

}  // namespace

void TutorialVectorization::runOnOperation() {
  func::FuncOp funcOp = getOperation();
  IRRewriter rewriter(funcOp);

  SmallVector<linalg::GenericOp> candidates;
  funcOp.walk([&](linalg::GenericOp op) { candidates.push_back(op); });

  for (linalg::GenericOp candidate : candidates) {
    (void)linalg::vectorize(rewriter, candidate);
  }

  // Cleanup.
  {
    RewritePatternSet vectorizationPatterns(funcOp.getContext());
    vector::populateVectorTransferPermutationMapLoweringPatterns(
        vectorizationPatterns);
    vector::populateVectorReductionToContractPatterns(vectorizationPatterns);
    vector::populateSinkVectorOpsPatterns(vectorizationPatterns);
    vector::TransferReadOp::getCanonicalizationPatterns(vectorizationPatterns,
                                                        funcOp.getContext());
    vector::TransferWriteOp::getCanonicalizationPatterns(vectorizationPatterns,
                                                         funcOp.getContext());
    if (failed(
            applyPatternsGreedily(funcOp, std::move(vectorizationPatterns)))) {
      return signalPassFailure();
    }
  }
}

}  // namespace mlir::tutorial