llvm / llvm-project / d62b4b08af03a9fc25274ed0e380d9d052fe251b / . / mlir / include / mlir / Dialect / Vector / VectorUtils.h

//===- VectorUtils.h - Vector Utilities -------------------------*- C++ -*-===// | |

// | |

// 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 | |

// | |

//===----------------------------------------------------------------------===// | |

#ifndef MLIR_DIALECT_VECTOR_VECTORUTILS_H_ | |

#define MLIR_DIALECT_VECTOR_VECTORUTILS_H_ | |

#include "mlir/IR/BuiltinAttributes.h" | |

#include "mlir/Support/LLVM.h" | |

#include "llvm/ADT/DenseMap.h" | |

namespace mlir { | |

// Forward declarations. | |

class AffineApplyOp; | |

class AffineForOp; | |

class AffineMap; | |

class Block; | |

class Location; | |

class OpBuilder; | |

class Operation; | |

class ShapedType; | |

class Value; | |

class VectorType; | |

class VectorTransferOpInterface; | |

namespace vector { | |

class TransferWriteOp; | |

class TransferReadOp; | |

/// Helper function that creates a memref::DimOp or tensor::DimOp depending on | |

/// the type of `source`. | |

Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim); | |

} // namespace vector | |

/// Return the number of elements of basis, `0` if empty. | |

int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis); | |

/// Given the shape and sizes of a vector, returns the corresponding | |

/// strides for each dimension. | |

/// TODO: needs better doc of how it is used. | |

SmallVector<int64_t, 4> computeStrides(ArrayRef<int64_t> shape, | |

ArrayRef<int64_t> sizes); | |

/// Computes and returns the linearized index of 'offsets' w.r.t. 'basis'. | |

int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis); | |

/// Given the strides together with a linear index in the dimension | |

/// space, returns the vector-space offsets in each dimension for a | |

/// de-linearized index. | |

SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> strides, | |

int64_t linearIndex); | |

/// Given the target sizes of a vector, together with vector-space offsets, | |

/// returns the element-space offsets for each dimension. | |

SmallVector<int64_t, 4> | |

computeElementOffsetsFromVectorSliceOffsets(ArrayRef<int64_t> sizes, | |

ArrayRef<int64_t> vectorOffsets); | |

/// Computes and returns the multi-dimensional ratio of `superShape` to | |

/// `subShape`. This is calculated by performing a traversal from minor to major | |

/// dimensions (i.e. in reverse shape order). If integral division is not | |

/// possible, returns None. | |

/// The ArrayRefs are assumed (and enforced) to only contain > 1 values. | |

/// This constraint comes from the fact that they are meant to be used with | |

/// VectorTypes, for which the property holds by construction. | |

/// | |

/// Examples: | |

/// - shapeRatio({3, 4, 5, 8}, {2, 5, 2}) returns {3, 2, 1, 4} | |

/// - shapeRatio({3, 4, 4, 8}, {2, 5, 2}) returns None | |

/// - shapeRatio({1, 2, 10, 32}, {2, 5, 2}) returns {1, 1, 2, 16} | |

Optional<SmallVector<int64_t, 4>> shapeRatio(ArrayRef<int64_t> superShape, | |

ArrayRef<int64_t> subShape); | |

/// Computes and returns the multi-dimensional ratio of the shapes of | |

/// `superVector` to `subVector`. If integral division is not possible, returns | |

/// None. | |

/// Assumes and enforces that the VectorTypes have the same elemental type. | |

Optional<SmallVector<int64_t, 4>> shapeRatio(VectorType superVectorType, | |

VectorType subVectorType); | |

/// Constructs a permutation map of invariant memref indices to vector | |

/// dimension. | |

/// | |

/// If no index is found to be invariant, 0 is added to the permutation_map and | |

/// corresponds to a vector broadcast along that dimension. | |

/// | |

/// The implementation uses the knowledge of the mapping of loops to | |

/// vector dimension. `loopToVectorDim` carries this information as a map with: | |

/// - keys representing "vectorized enclosing loops"; | |

/// - values representing the corresponding vector dimension. | |

/// Note that loopToVectorDim is a whole function map from which only enclosing | |

/// loop information is extracted. | |

/// | |

/// Prerequisites: `indices` belong to a vectorizable load or store operation | |

/// (i.e. at most one invariant index along each AffineForOp of | |

/// `loopToVectorDim`). `insertPoint` is the insertion point for the vectorized | |

/// load or store operation. | |

/// | |

/// Example 1: | |

/// The following MLIR snippet: | |

/// | |

/// ```mlir | |

/// affine.for %i3 = 0 to %0 { | |

/// affine.for %i4 = 0 to %1 { | |

/// affine.for %i5 = 0 to %2 { | |

/// %a5 = load %arg0[%i4, %i5, %i3] : memref<?x?x?xf32> | |

/// }}} | |

/// ``` | |

/// | |

/// may vectorize with {permutation_map: (d0, d1, d2) -> (d2, d1)} into: | |

/// | |

/// ```mlir | |

/// affine.for %i3 = 0 to %0 step 32 { | |

/// affine.for %i4 = 0 to %1 { | |

/// affine.for %i5 = 0 to %2 step 256 { | |

/// %4 = vector.transfer_read %arg0, %i4, %i5, %i3 | |

/// {permutation_map: (d0, d1, d2) -> (d2, d1)} : | |

/// (memref<?x?x?xf32>, index, index) -> vector<32x256xf32> | |

/// }}} | |

/// ``` | |

/// | |

/// Meaning that vector.transfer_read will be responsible for reading the slice: | |

/// `%arg0[%i4, %i5:%15+256, %i3:%i3+32]` into vector<32x256xf32>. | |

/// | |

/// Example 2: | |

/// The following MLIR snippet: | |

/// | |

/// ```mlir | |

/// %cst0 = arith.constant 0 : index | |

/// affine.for %i0 = 0 to %0 { | |

/// %a0 = load %arg0[%cst0, %cst0] : memref<?x?xf32> | |

/// } | |

/// ``` | |

/// | |

/// may vectorize with {permutation_map: (d0) -> (0)} into: | |

/// | |

/// ```mlir | |

/// affine.for %i0 = 0 to %0 step 128 { | |

/// %3 = vector.transfer_read %arg0, %c0_0, %c0_0 | |

/// {permutation_map: (d0, d1) -> (0)} : | |

/// (memref<?x?xf32>, index, index) -> vector<128xf32> | |

/// } | |

/// ```` | |

/// | |

/// Meaning that vector.transfer_read will be responsible of reading the slice | |

/// `%arg0[%c0, %c0]` into vector<128xf32> which needs a 1-D vector broadcast. | |

/// | |

AffineMap | |

makePermutationMap(Block *insertPoint, ArrayRef<Value> indices, | |

const DenseMap<Operation *, unsigned> &loopToVectorDim); | |

AffineMap | |

makePermutationMap(Operation *insertPoint, ArrayRef<Value> indices, | |

const DenseMap<Operation *, unsigned> &loopToVectorDim); | |

/// Build the default minor identity map suitable for a vector transfer. This | |

/// also handles the case memref<... x vector<...>> -> vector<...> in which the | |

/// rank of the identity map must take the vector element type into account. | |

AffineMap getTransferMinorIdentityMap(ShapedType shapedType, | |

VectorType vectorType); | |

/// Return true if we can prove that the transfer operations access disjoint | |

/// memory. | |

bool isDisjointTransferSet(VectorTransferOpInterface transferA, | |

VectorTransferOpInterface transferB); | |

/// Same behavior as `isDisjointTransferSet` but doesn't require the operations | |

/// to have the same tensor/memref. This allows comparing operations accessing | |

/// different tensors. | |

bool isDisjointTransferIndices(VectorTransferOpInterface transferA, | |

VectorTransferOpInterface transferB); | |

/// Return true if the transfer_write fully writes the data accessed by the | |

/// transfer_read. | |

bool checkSameValueRAW(vector::TransferWriteOp defWrite, | |

vector::TransferReadOp read); | |

/// Return true if the write op fully over-write the priorWrite transfer_write | |

/// op. | |

bool checkSameValueWAW(vector::TransferWriteOp write, | |

vector::TransferWriteOp priorWrite); | |

// Helper that returns a subset of `arrayAttr` as a vector of int64_t. | |

SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr, | |

unsigned dropFront = 0, | |

unsigned dropBack = 0); | |

namespace matcher { | |

/// Matches vector.transfer_read, vector.transfer_write and ops that return a | |

/// vector type that is a multiple of the sub-vector type. This allows passing | |

/// over other smaller vector types in the function and avoids interfering with | |

/// operations on those. | |

/// This is a first approximation, it can easily be extended in the future. | |

/// TODO: this could all be much simpler if we added a bit that a vector type to | |

/// mark that a vector is a strict super-vector but it still does not warrant | |

/// adding even 1 extra bit in the IR for now. | |

bool operatesOnSuperVectorsOf(Operation &op, VectorType subVectorType); | |

} // end namespace matcher | |

} // end namespace mlir | |

#endif // MLIR_DIALECT_VECTOR_VECTORUTILS_H_ |