| //===- VectorInterfaces.td - Vector interfaces -------------*- tablegen -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Defines the interface for operations on vectors. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef MLIR_INTERFACES_VECTORINTERFACES |
| #define MLIR_INTERFACES_VECTORINTERFACES |
| |
| include "mlir/IR/OpBase.td" |
| |
| def VectorUnrollOpInterface : OpInterface<"VectorUnrollOpInterface"> { |
| let description = [{ |
| Encodes properties of an operation on vectors that can be unrolled. |
| }]; |
| let cppNamespace = "::mlir"; |
| |
| let methods = [ |
| InterfaceMethod< |
| /*desc=*/[{ |
| Return the shape ratio of unrolling to the target vector shape |
| `targetShape`. Return `None` if the op cannot be unrolled to the target |
| vector shape. |
| }], |
| /*retTy=*/"::mlir::Optional<::mlir::SmallVector<int64_t, 4>>", |
| /*methodName=*/"getShapeForUnroll", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| assert($_op->getNumResults() == 1); |
| auto vt = $_op.getResult().getType(). |
| template dyn_cast<::mlir::VectorType>(); |
| if (!vt) |
| return ::mlir::None; |
| ::mlir::SmallVector<int64_t, 4> res(vt.getShape().begin(), vt.getShape().end()); |
| return res; |
| }] |
| >, |
| ]; |
| } |
| |
| def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> { |
| let description = [{ |
| Encodes properties of a transfer read or write operation. |
| }]; |
| let cppNamespace = "::mlir"; |
| |
| let methods = [ |
| StaticInterfaceMethod< |
| /*desc=*/"Return the `in_bounds` attribute name.", |
| /*retTy=*/"::mlir::StringRef", |
| /*methodName=*/"getInBoundsAttrName", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/ [{ return "in_bounds"; }] |
| >, |
| StaticInterfaceMethod< |
| /*desc=*/"Return the `permutation_map` attribute name.", |
| /*retTy=*/"::mlir::StringRef", |
| /*methodName=*/"getPermutationMapAttrName", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/ [{ return "permutation_map"; }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ Return `true` if dimension `dim` is in-bounds. Return `false` |
| otherwise. }], |
| /*retTy=*/"bool", |
| /*methodName=*/"isDimInBounds", |
| /*args=*/(ins "unsigned":$dim), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return $_op.isBroadcastDim(dim) |
| || ($_op.in_bounds() |
| && $_op.in_bounds()->template cast<::mlir::ArrayAttr>()[dim] |
| .template cast<::mlir::BoolAttr>().getValue()); |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the memref or ranked tensor operand.", |
| /*retTy=*/"::mlir::Value", |
| /*methodName=*/"source", |
| /*args=*/(ins), |
| /*methodBody=*/"return $_op.source();" |
| /*defaultImplementation=*/ |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the vector operand or result.", |
| /*retTy=*/"::mlir::Value", |
| /*methodName=*/"vector", |
| /*args=*/(ins), |
| /*methodBody=*/"return $_op.vector();" |
| /*defaultImplementation=*/ |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the indices operands.", |
| /*retTy=*/"::mlir::ValueRange", |
| /*methodName=*/"indices", |
| /*args=*/(ins), |
| /*methodBody=*/"return $_op.indices();" |
| /*defaultImplementation=*/ |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the permutation map.", |
| /*retTy=*/"::mlir::AffineMap", |
| /*methodName=*/"permutation_map", |
| /*args=*/(ins), |
| /*methodBody=*/"return $_op.permutation_map();" |
| /*defaultImplementation=*/ |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Returns true if op involves a 0-d tensor/memref and a vector |
| of shape {1}. This is temporary until we have 0-d vectors. |
| // TODO: turn this into 0-d vectors + empty permutation_map. |
| }], |
| /*retTy=*/"bool", |
| /*methodName=*/"isZeroD", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| if (getShapedType().getRank() > 0) |
| return false; |
| if (getVectorType().getShape() != ArrayRef<int64_t>{1}) |
| return false; |
| AffineMap map = AffineMap::get( |
| /*numDims=*/0, /*numSymbols=*/0, |
| getAffineConstantExpr(0, $_op->getContext())); |
| if ($_op.permutation_map() != map) |
| return false; |
| return true; |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ Returns true if the specified dimension is a broadcast. }], |
| /*retTy=*/"bool", |
| /*methodName=*/"isBroadcastDim", |
| /*args=*/(ins "unsigned":$idx), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| auto expr = $_op.permutation_map().getResult(idx); |
| return expr.template isa<::mlir::AffineConstantExpr>() && |
| expr.template dyn_cast<::mlir::AffineConstantExpr>().getValue() == 0; |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ Returns true if at least one of the dimensions in the |
| permutation map is a broadcast.}], |
| /*retTy=*/"bool", |
| /*methodName=*/"hasBroadcastDim", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| // 0-d transfers are not considered broadcasts but they need to be |
| // represented with a vector<1xt> until we have 0-d vectors. |
| if ($_op.isZeroD()) return false; |
| for (unsigned i = 0; i < $_op.permutation_map().getNumResults(); ++i) { |
| if ($_op.isBroadcastDim(i)) |
| return true; |
| } |
| return false; |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the `in_bounds` boolean ArrayAttr.", |
| /*retTy=*/"::mlir::Optional<::mlir::ArrayAttr>", |
| /*methodName=*/"in_bounds", |
| /*args=*/(ins), |
| /*methodBody=*/"return $_op.in_bounds();" |
| /*defaultImplementation=*/ |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the ShapedType.", |
| /*retTy=*/"::mlir::ShapedType", |
| /*methodName=*/"getShapedType", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/ |
| "return $_op.source().getType().template cast<::mlir::ShapedType>();" |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the VectorType.", |
| /*retTy=*/"::mlir::VectorType", |
| /*methodName=*/"getVectorType", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return $_op.vector().getType().template dyn_cast<::mlir::VectorType>(); |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/"Return the mask type if the op has a mask.", |
| /*retTy=*/"::mlir::VectorType", |
| /*methodName=*/"getMaskType", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return $_op.mask() |
| ? ::mlir::vector::detail::transferMaskType( |
| $_op.getVectorType(), $_op.permutation_map()) |
| : ::mlir::VectorType(); |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ Return the number of dimensions that participate in the |
| permutation map.}], |
| /*retTy=*/"unsigned", |
| /*methodName=*/"getTransferRank", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/ |
| "return $_op.permutation_map().getNumResults();" |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ Return the number of leading shaped dimensions that do not |
| participate in the permutation map.}], |
| /*retTy=*/"unsigned", |
| /*methodName=*/"getLeadingShapedRank", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/ |
| "return $_op.getShapedType().getRank() - $_op.getTransferRank();" |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ Returns true if at least one of the dimensions may be |
| out-of-bounds.}], |
| /*retTy=*/"bool", |
| /*methodName=*/"hasOutOfBoundsDim", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| for (unsigned idx = 0, e = $_op.getTransferRank(); idx < e; ++idx) |
| if (!$_op.isDimInBounds(idx)) |
| return true; |
| return false; |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Helper function to account for the fact that `permutationMap` results and |
| `op.indices` sizes may not match and may not be aligned. The first |
| `getLeadingShapedRank()` indices may just be indexed and not |
| transferred from/into the vector. |
| For example: |
| ``` |
| vector.transfer %0[%i, %j, %k, %c0] : |
| memref<?x?x?x?xf32>, vector<2x4xf32> |
| ``` |
| with `permutation_map = (d0, d1, d2, d3) -> (d2, d3)`. |
| Provide a zip function to coiterate on 2 running indices: `resultIdx` and |
| `indicesIdx` which accounts for this misalignment. |
| }], |
| /*retTy=*/"void", |
| /*methodName=*/"zipResultAndIndexing", |
| /*args=*/(ins "::llvm::function_ref<void(int64_t, int64_t)>":$fun), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| for (int64_t resultIdx = 0, |
| indicesIdx = $_op.getLeadingShapedRank(), |
| eResult = $_op.getTransferRank(); |
| resultIdx < eResult; |
| ++resultIdx, ++indicesIdx) |
| fun(resultIdx, indicesIdx); |
| }] |
| >, |
| ]; |
| } |
| |
| #endif // MLIR_INTERFACES_VECTORINTERFACES |