MLIR  18.0.0git
VectorUtils.h
Go to the documentation of this file.
1 //===- VectorUtils.h - Vector Utilities -------------------------*- C++ -*-===//
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 #ifndef MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
10 #define MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
11 
14 #include "mlir/Support/LLVM.h"
15 
16 #include "llvm/ADT/DenseMap.h"
17 
18 namespace mlir {
19 
20 // Forward declarations.
21 class AffineMap;
22 class Block;
23 class Location;
24 class OpBuilder;
25 class Operation;
26 class ShapedType;
27 class Value;
28 class VectorType;
29 class VectorTransferOpInterface;
30 
31 namespace affine {
32 class AffineApplyOp;
33 class AffineForOp;
34 } // namespace affine
35 
36 namespace vector {
37 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
38 /// the type of `source`.
39 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
40 
41 /// Returns two dims that are greater than one if the transposition is applied
42 /// on a 2D slice. Otherwise, returns a failure.
43 FailureOr<std::pair<int, int>> isTranspose2DSlice(vector::TransposeOp op);
44 
45 } // namespace vector
46 
47 /// Constructs a permutation map of invariant memref indices to vector
48 /// dimension.
49 ///
50 /// If no index is found to be invariant, 0 is added to the permutation_map and
51 /// corresponds to a vector broadcast along that dimension.
52 ///
53 /// The implementation uses the knowledge of the mapping of loops to
54 /// vector dimension. `loopToVectorDim` carries this information as a map with:
55 /// - keys representing "vectorized enclosing loops";
56 /// - values representing the corresponding vector dimension.
57 /// Note that loopToVectorDim is a whole function map from which only enclosing
58 /// loop information is extracted.
59 ///
60 /// Prerequisites: `indices` belong to a vectorizable load or store operation
61 /// (i.e. at most one invariant index along each AffineForOp of
62 /// `loopToVectorDim`). `insertPoint` is the insertion point for the vectorized
63 /// load or store operation.
64 ///
65 /// Example 1:
66 /// The following MLIR snippet:
67 ///
68 /// ```mlir
69 /// affine.for %i3 = 0 to %0 {
70 /// affine.for %i4 = 0 to %1 {
71 /// affine.for %i5 = 0 to %2 {
72 /// %a5 = load %arg0[%i4, %i5, %i3] : memref<?x?x?xf32>
73 /// }}}
74 /// ```
75 ///
76 /// may vectorize with {permutation_map: (d0, d1, d2) -> (d2, d1)} into:
77 ///
78 /// ```mlir
79 /// affine.for %i3 = 0 to %0 step 32 {
80 /// affine.for %i4 = 0 to %1 {
81 /// affine.for %i5 = 0 to %2 step 256 {
82 /// %4 = vector.transfer_read %arg0, %i4, %i5, %i3
83 /// {permutation_map: (d0, d1, d2) -> (d2, d1)} :
84 /// (memref<?x?x?xf32>, index, index) -> vector<32x256xf32>
85 /// }}}
86 /// ```
87 ///
88 /// Meaning that vector.transfer_read will be responsible for reading the slice:
89 /// `%arg0[%i4, %i5:%15+256, %i3:%i3+32]` into vector<32x256xf32>.
90 ///
91 /// Example 2:
92 /// The following MLIR snippet:
93 ///
94 /// ```mlir
95 /// %cst0 = arith.constant 0 : index
96 /// affine.for %i0 = 0 to %0 {
97 /// %a0 = load %arg0[%cst0, %cst0] : memref<?x?xf32>
98 /// }
99 /// ```
100 ///
101 /// may vectorize with {permutation_map: (d0) -> (0)} into:
102 ///
103 /// ```mlir
104 /// affine.for %i0 = 0 to %0 step 128 {
105 /// %3 = vector.transfer_read %arg0, %c0_0, %c0_0
106 /// {permutation_map: (d0, d1) -> (0)} :
107 /// (memref<?x?xf32>, index, index) -> vector<128xf32>
108 /// }
109 /// ````
110 ///
111 /// Meaning that vector.transfer_read will be responsible of reading the slice
112 /// `%arg0[%c0, %c0]` into vector<128xf32> which needs a 1-D vector broadcast.
113 ///
114 AffineMap
115 makePermutationMap(Block *insertPoint, ArrayRef<Value> indices,
116  const DenseMap<Operation *, unsigned> &loopToVectorDim);
117 AffineMap
118 makePermutationMap(Operation *insertPoint, ArrayRef<Value> indices,
119  const DenseMap<Operation *, unsigned> &loopToVectorDim);
120 
121 namespace matcher {
122 
123 /// Matches vector.transfer_read, vector.transfer_write and ops that return a
124 /// vector type that is a multiple of the sub-vector type. This allows passing
125 /// over other smaller vector types in the function and avoids interfering with
126 /// operations on those.
127 /// This is a first approximation, it can easily be extended in the future.
128 /// TODO: this could all be much simpler if we added a bit that a vector type to
129 /// mark that a vector is a strict super-vector but it still does not warrant
130 /// adding even 1 extra bit in the IR for now.
131 bool operatesOnSuperVectorsOf(Operation &op, VectorType subVectorType);
132 
133 } // namespace matcher
134 } // namespace mlir
135 
136 #endif // MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
static AffineMap makePermutationMap(ArrayRef< Value > indices, const DenseMap< Operation *, unsigned > &enclosingLoopToVectorDim)
Constructs a permutation map from memref indices to vector dimension.
FailureOr< std::pair< int, int > > isTranspose2DSlice(vector::TransposeOp op)
Returns two dims that are greater than one if the transposition is applied on a 2D slice.
Definition: VectorUtils.cpp:80
Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim)
Helper function that creates a memref::DimOp or tensor::DimOp depending on the type of source.
Definition: VectorUtils.cpp:37
Include the generated interface declarations.