1//===- VectorInterfaces.td - Vector interfaces -------------*- tablegen -*-===//
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// Defines the interface for operations on vectors.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef MLIR_INTERFACES_VECTORINTERFACES
14#define MLIR_INTERFACES_VECTORINTERFACES
15
16include "mlir/IR/OpBase.td"
17
18def VectorUnrollOpInterface : OpInterface<"VectorUnrollOpInterface"> {
19  let description = [{
20    Encodes properties of an operation on vectors that can be unrolled.
21  }];
22  let cppNamespace = "::mlir";
23
24  let methods = [
25    InterfaceMethod<
26      /*desc=*/[{
27        Return the shape ratio of unrolling to the target vector shape
28        `targetShape`. Return `None` if the op cannot be unrolled to the target
29        vector shape.
30      }],
31      /*retTy=*/"::llvm::Optional<::llvm::SmallVector<int64_t, 4>>",
32      /*methodName=*/"getShapeForUnroll",
33      /*args=*/(ins),
34      /*methodBody=*/"",
35      /*defaultImplementation=*/[{
36        assert($_op->getNumResults() == 1);
37        auto vt = $_op.getResult().getType().
38          template dyn_cast<::mlir::VectorType>();
39        if (!vt)
40          return ::mlir::None;
41        ::llvm::SmallVector<int64_t, 4> res(vt.getShape().begin(), vt.getShape().end());
42        return res;
43      }]
44    >,
45  ];
46}
47
48def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
49  let description = [{
50    Encodes properties of a transfer read or write operation.
51  }];
52  let cppNamespace = "::mlir";
53
54  let methods = [
55    StaticInterfaceMethod<
56      /*desc=*/"Return the `in_bounds` attribute name.",
57      /*retTy=*/"::mlir::StringRef",
58      /*methodName=*/"getInBoundsAttrStrName",
59      /*args=*/(ins),
60      /*methodBody=*/"",
61      /*defaultImplementation=*/ [{ return "in_bounds"; }]
62    >,
63    StaticInterfaceMethod<
64      /*desc=*/"Return the `permutation_map` attribute name.",
65      /*retTy=*/"::mlir::StringRef",
66      /*methodName=*/"getPermutationMapAttrStrName",
67      /*args=*/(ins),
68      /*methodBody=*/"",
69      /*defaultImplementation=*/ [{ return "permutation_map"; }]
70    >,
71    InterfaceMethod<
72      /*desc=*/[{ Return `true` if dimension `dim` is in-bounds. Return `false`
73                 otherwise. }],
74      /*retTy=*/"bool",
75      /*methodName=*/"isDimInBounds",
76      /*args=*/(ins "unsigned":$dim),
77      /*methodBody=*/"",
78      /*defaultImplementation=*/[{
79        return $_op.isBroadcastDim(dim)
80            || ($_op.getInBounds()
81                && $_op.getInBounds()->template cast<::mlir::ArrayAttr>()[dim]
82                                    .template cast<::mlir::BoolAttr>().getValue());
83      }]
84    >,
85    InterfaceMethod<
86      /*desc=*/"Return the memref or ranked tensor operand.",
87      /*retTy=*/"::mlir::Value",
88      /*methodName=*/"source",
89      /*args=*/(ins),
90      /*methodBody=*/"return $_op.getSource();"
91      /*defaultImplementation=*/
92    >,
93    InterfaceMethod<
94      /*desc=*/"Return the vector operand or result.",
95      /*retTy=*/"::mlir::Value",
96      /*methodName=*/"vector",
97      /*args=*/(ins),
98      /*methodBody=*/"return $_op.getVector();"
99      /*defaultImplementation=*/
100    >,
101    InterfaceMethod<
102      /*desc=*/"Return the indices operands.",
103      /*retTy=*/"::mlir::ValueRange",
104      /*methodName=*/"indices",
105      /*args=*/(ins),
106      /*methodBody=*/"return $_op.getIndices();"
107      /*defaultImplementation=*/
108    >,
109    InterfaceMethod<
110      /*desc=*/"Return the permutation map.",
111      /*retTy=*/"::mlir::AffineMap",
112      /*methodName=*/"permutation_map",
113      /*args=*/(ins),
114      /*methodBody=*/"return $_op.getPermutationMap();"
115      /*defaultImplementation=*/
116    >,
117    InterfaceMethod<
118      /*desc=*/[{ Returns true if the specified dimension is a broadcast. }],
119      /*retTy=*/"bool",
120      /*methodName=*/"isBroadcastDim",
121      /*args=*/(ins "unsigned":$idx),
122      /*methodBody=*/"",
123      /*defaultImplementation=*/[{
124        auto expr = $_op.getPermutationMap().getResult(idx);
125        return expr.template isa<::mlir::AffineConstantExpr>() &&
126               expr.template dyn_cast<::mlir::AffineConstantExpr>().getValue() == 0;
127      }]
128    >,
129    InterfaceMethod<
130      /*desc=*/[{ Returns true if at least one of the dimensions in the
131                  permutation map is a broadcast.}],
132      /*retTy=*/"bool",
133      /*methodName=*/"hasBroadcastDim",
134      /*args=*/(ins),
135      /*methodBody=*/"",
136      /*defaultImplementation=*/[{
137        for (unsigned i = 0, rank = getTransferRank(); i < rank; ++i) {
138          if ($_op.isBroadcastDim(i))
139            return true;
140        }
141        return false;
142      }]
143    >,
144    InterfaceMethod<
145      /*desc=*/"Return the `in_bounds` boolean ArrayAttr.",
146      /*retTy=*/"::llvm::Optional<::mlir::ArrayAttr>",
147      /*methodName=*/"in_bounds",
148      /*args=*/(ins),
149      /*methodBody=*/"return $_op.getInBounds();"
150      /*defaultImplementation=*/
151    >,
152    InterfaceMethod<
153      /*desc=*/"Return the ShapedType.",
154      /*retTy=*/"::mlir::ShapedType",
155      /*methodName=*/"getShapedType",
156      /*args=*/(ins),
157      /*methodBody=*/"",
158      /*defaultImplementation=*/
159        "return $_op.getSource().getType().template cast<::mlir::ShapedType>();"
160    >,
161    InterfaceMethod<
162      /*desc=*/"Return the VectorType.",
163      /*retTy=*/"::mlir::VectorType",
164      /*methodName=*/"getVectorType",
165      /*args=*/(ins),
166      /*methodBody=*/"",
167      /*defaultImplementation=*/[{
168        return $_op.getVector().getType().template dyn_cast<::mlir::VectorType>();
169        }]
170    >,
171    InterfaceMethod<
172      /*desc=*/"Return the mask type if the op has a mask.",
173      /*retTy=*/"::mlir::VectorType",
174      /*methodName=*/"getMaskType",
175      /*args=*/(ins),
176      /*methodBody=*/"",
177      /*defaultImplementation=*/[{
178        return $_op.getMask()
179            ? ::mlir::vector::detail::transferMaskType(
180                $_op.getVectorType(), $_op.getPermutationMap())
181            : ::mlir::VectorType();
182      }]
183    >,
184    InterfaceMethod<
185      /*desc=*/[{ Return the number of dimensions that participate in the
186                  permutation map.}],
187      /*retTy=*/"unsigned",
188      /*methodName=*/"getTransferRank",
189      /*args=*/(ins),
190      /*methodBody=*/"",
191      /*defaultImplementation=*/
192        "return $_op.getPermutationMap().getNumResults();"
193    >,
194    InterfaceMethod<
195      /*desc=*/[{ Return the number of leading shaped dimensions that do not
196                  participate in the permutation map.}],
197      /*retTy=*/"unsigned",
198      /*methodName=*/"getLeadingShapedRank",
199      /*args=*/(ins),
200      /*methodBody=*/"",
201      /*defaultImplementation=*/
202        "return $_op.getShapedType().getRank() - $_op.getTransferRank();"
203    >,
204    InterfaceMethod<
205      /*desc=*/[{ Returns true if at least one of the dimensions may be
206                  out-of-bounds.}],
207      /*retTy=*/"bool",
208      /*methodName=*/"hasOutOfBoundsDim",
209      /*args=*/(ins),
210      /*methodBody=*/"",
211      /*defaultImplementation=*/[{
212        for (unsigned idx = 0, e = $_op.getTransferRank(); idx < e; ++idx)
213          if (!$_op.isDimInBounds(idx))
214            return true;
215        return false;
216      }]
217    >,
218    InterfaceMethod<
219      /*desc=*/[{
220      Helper function to account for the fact that `permutationMap` results and
221      `op.indices` sizes may not match and may not be aligned. The first
222      `getLeadingShapedRank()` indices may just be indexed and not
223      transferred from/into the vector.
224      For example:
225      ```
226         vector.transfer %0[%i, %j, %k, %c0] :
227           memref<?x?x?x?xf32>, vector<2x4xf32>
228      ```
229      with `permutation_map = (d0, d1, d2, d3) -> (d2, d3)`.
230      Provide a zip function to coiterate on 2 running indices: `resultIdx` and
231      `indicesIdx` which accounts for this misalignment.
232      }],
233      /*retTy=*/"void",
234      /*methodName=*/"zipResultAndIndexing",
235      /*args=*/(ins "::llvm::function_ref<void(int64_t, int64_t)>":$fun),
236      /*methodBody=*/"",
237      /*defaultImplementation=*/[{
238        for (int64_t resultIdx = 0,
239                   indicesIdx = $_op.getLeadingShapedRank(),
240                   eResult = $_op.getTransferRank();
241           resultIdx < eResult;
242           ++resultIdx, ++indicesIdx)
243        fun(resultIdx, indicesIdx);
244      }]
245    >,
246    InterfaceMethod<
247      /*desc=*/[{
248      Return an upper-bound shape accessed by the transfer op within the
249      tensor/memref operand.
250      For example:
251      ```
252        vector.transfer %w0[%i, %j] {
253          permutation_map = affine_map<(d0, d1) -> (d1, d0, 0)>} :
254          tensor<?x?xf32>, vector<4x2x6xf32>
255      ```
256      returns a shape [2, 4].
257      }],
258      /*retTy=*/"SmallVector<int64_t>",
259      /*methodName=*/"getTransferChunkAccessed",
260      /*args=*/(ins),
261      /*methodBody=*/"",
262      /*defaultImplementation=*/[{
263        SmallVector<int64_t> dimSizes($_op.getPermutationMap().getNumDims(), 0);
264        for (auto vecDims : llvm::zip($_op.getPermutationMap().getResults(),
265                                      $_op.getVectorType().getShape())) {
266          AffineExpr dim = std::get<0>(vecDims);
267          int64_t size = std::get<1>(vecDims);
268          // Skip broadcast.
269          if (dim.isa<AffineConstantExpr>())
270            continue;
271          dimSizes[dim.cast<AffineDimExpr>().getPosition()] = size;
272        }
273        return dimSizes;
274      }]
275    >,
276  ];
277}
278
279#endif // MLIR_INTERFACES_VECTORINTERFACES
280