1 //===- IndexingUtils.h - Helpers related to index computations --*- 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 // This header file defines utilities and common canonicalization patterns for
10 // reshape operations.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #ifndef MLIR_DIALECT_UTILS_INDEXINGUTILS_H
15 #define MLIR_DIALECT_UTILS_INDEXINGUTILS_H
16
17 #include "mlir/Support/LLVM.h"
18 #include "llvm/ADT/ArrayRef.h"
19 #include "llvm/ADT/SmallVector.h"
20
21 namespace mlir {
22 class ArrayAttr;
23
24 /// Computes and returns the linearized index of 'offsets' w.r.t. 'basis'.
25 int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);
26
27 /// Given the strides together with a linear index in the dimension
28 /// space, returns the vector-space offsets in each dimension for a
29 /// de-linearized index.
30 SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> strides,
31 int64_t linearIndex);
32
33 /// Apply the permutation defined by `permutation` to `inVec`.
34 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
35 /// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector
36 /// `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a', 'b']`.
37 template <typename T, unsigned N>
applyPermutationToVector(SmallVector<T,N> & inVec,ArrayRef<int64_t> permutation)38 void applyPermutationToVector(SmallVector<T, N> &inVec,
39 ArrayRef<int64_t> permutation) {
40 SmallVector<T, N> auxVec(inVec.size());
41 for (const auto &en : enumerate(permutation))
42 auxVec[en.index()] = inVec[en.value()];
43 inVec = auxVec;
44 }
45
46 /// Helper that returns a subset of `arrayAttr` as a vector of int64_t.
47 SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr,
48 unsigned dropFront = 0,
49 unsigned dropBack = 0);
50 } // namespace mlir
51
52 #endif // MLIR_DIALECT_UTILS_INDEXINGUTILS_H
53