1 //===- LowerVectorBitCast.cpp - Lower 'vector.bitcast' operation ----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements target-independent rewrites and utilities to lower the
10 // 'vector.bitcast' operation.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "mlir/Dialect/Vector/IR/VectorOps.h"
15 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
16 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
17 #include "mlir/IR/BuiltinTypes.h"
18 #include "mlir/IR/PatternMatch.h"
19
20 #define DEBUG_TYPE "vector-bitcast-lowering"
21
22 using namespace mlir;
23 using namespace mlir::vector;
24
25 namespace {
26
27 /// A one-shot unrolling of vector.bitcast to the `targetRank`.
28 ///
29 /// Example:
30 ///
31 /// vector.bitcast %a, %b : vector<1x2x3x4xi64> to vector<1x2x3x8xi32>
32 ///
33 /// Would be unrolled to:
34 ///
35 /// %result = arith.constant dense<0> : vector<1x2x3x8xi32>
36 /// %0 = vector.extract %a[0, 0, 0] ─┐
37 /// : vector<4xi64> from vector<1x2x3x4xi64> |
38 /// %1 = vector.bitcast %0 | - Repeated 6x for
39 /// : vector<4xi64> to vector<8xi32> | all leading positions
40 /// %2 = vector.insert %1, %result [0, 0, 0] |
41 /// : vector<8xi64> into vector<1x2x3x8xi32> ─┘
42 ///
43 /// Note: If any leading dimension before the `targetRank` is scalable the
44 /// unrolling will stop before the scalable dimension.
45 class UnrollBitCastOp final : public OpRewritePattern<vector::BitCastOp> {
46 public:
UnrollBitCastOp(int64_t targetRank,MLIRContext * context,PatternBenefit benefit=1)47 UnrollBitCastOp(int64_t targetRank, MLIRContext *context,
48 PatternBenefit benefit = 1)
49 : OpRewritePattern(context, benefit), targetRank(targetRank) {};
50
matchAndRewrite(vector::BitCastOp op,PatternRewriter & rewriter) const51 LogicalResult matchAndRewrite(vector::BitCastOp op,
52 PatternRewriter &rewriter) const override {
53 VectorType resultType = op.getResultVectorType();
54 auto unrollIterator = vector::createUnrollIterator(resultType, targetRank);
55 if (!unrollIterator)
56 return failure();
57
58 auto unrollRank = unrollIterator->getRank();
59 ArrayRef<int64_t> shape = resultType.getShape().drop_front(unrollRank);
60 ArrayRef<bool> scalableDims =
61 resultType.getScalableDims().drop_front(unrollRank);
62 auto bitcastResType =
63 VectorType::get(shape, resultType.getElementType(), scalableDims);
64
65 Location loc = op.getLoc();
66 Value result = rewriter.create<arith::ConstantOp>(
67 loc, resultType, rewriter.getZeroAttr(resultType));
68 for (auto position : *unrollIterator) {
69 Value extract =
70 rewriter.create<vector::ExtractOp>(loc, op.getSource(), position);
71 Value bitcast =
72 rewriter.create<vector::BitCastOp>(loc, bitcastResType, extract);
73 result =
74 rewriter.create<vector::InsertOp>(loc, bitcast, result, position);
75 }
76
77 rewriter.replaceOp(op, result);
78 return success();
79 }
80
81 private:
82 int64_t targetRank = 1;
83 };
84
85 } // namespace
86
populateVectorBitCastLoweringPatterns(RewritePatternSet & patterns,int64_t targetRank,PatternBenefit benefit)87 void mlir::vector::populateVectorBitCastLoweringPatterns(
88 RewritePatternSet &patterns, int64_t targetRank, PatternBenefit benefit) {
89 patterns.add<UnrollBitCastOp>(targetRank, patterns.getContext(), benefit);
90 }
91