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/UB/IR/UBOps.h" |
15 | #include "mlir/Dialect/Vector/IR/VectorOps.h" |
16 | #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h" |
17 | #include "mlir/Dialect/Vector/Utils/VectorUtils.h" |
18 | #include "mlir/IR/BuiltinTypes.h" |
19 | #include "mlir/IR/PatternMatch.h" |
20 | |
21 | #define DEBUG_TYPE "vector-bitcast-lowering" |
22 | |
23 | using namespace mlir; |
24 | using namespace mlir::vector; |
25 | |
26 | namespace { |
27 | |
28 | /// A one-shot unrolling of vector.bitcast to the `targetRank`. |
29 | /// |
30 | /// Example: |
31 | /// |
32 | /// vector.bitcast %a, %b : vector<1x2x3x4xi64> to vector<1x2x3x8xi32> |
33 | /// |
34 | /// Would be unrolled to: |
35 | /// |
36 | /// %result = ub.poison : vector<1x2x3x8xi32> |
37 | /// %0 = vector.extract %a[0, 0, 0] ─┐ |
38 | /// : vector<4xi64> from vector<1x2x3x4xi64> | |
39 | /// %1 = vector.bitcast %0 | - Repeated 6x for |
40 | /// : vector<4xi64> to vector<8xi32> | all leading positions |
41 | /// %2 = vector.insert %1, %result [0, 0, 0] | |
42 | /// : vector<8xi64> into vector<1x2x3x8xi32> ─┘ |
43 | /// |
44 | /// Note: If any leading dimension before the `targetRank` is scalable the |
45 | /// unrolling will stop before the scalable dimension. |
46 | class UnrollBitCastOp final : public OpRewritePattern<vector::BitCastOp> { |
47 | public: |
48 | UnrollBitCastOp(int64_t targetRank, MLIRContext *context, |
49 | PatternBenefit benefit = 1) |
50 | : OpRewritePattern(context, benefit), targetRank(targetRank) {}; |
51 | |
52 | LogicalResult matchAndRewrite(vector::BitCastOp op, |
53 | PatternRewriter &rewriter) const override { |
54 | VectorType resultType = op.getResultVectorType(); |
55 | auto unrollIterator = vector::createUnrollIterator(vType: resultType, targetRank); |
56 | if (!unrollIterator) |
57 | return failure(); |
58 | |
59 | auto unrollRank = unrollIterator->getRank(); |
60 | ArrayRef<int64_t> shape = resultType.getShape().drop_front(unrollRank); |
61 | ArrayRef<bool> scalableDims = |
62 | resultType.getScalableDims().drop_front(unrollRank); |
63 | auto bitcastResType = |
64 | VectorType::get(shape, resultType.getElementType(), scalableDims); |
65 | |
66 | Location loc = op.getLoc(); |
67 | Value result = rewriter.create<ub::PoisonOp>(loc, 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 | |
87 | void mlir::vector::populateVectorBitCastLoweringPatterns( |
88 | RewritePatternSet &patterns, int64_t targetRank, PatternBenefit benefit) { |
89 | patterns.add<UnrollBitCastOp>(arg&: targetRank, args: patterns.getContext(), args&: benefit); |
90 | } |
91 | |