MLIR  16.0.0git
IndexingUtils.h
Go to the documentation of this file.
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/IR/Builders.h"
18 #include "mlir/Support/LLVM.h"
19 #include "llvm/ADT/ArrayRef.h"
20 #include "llvm/ADT/SmallVector.h"
21 
22 namespace mlir {
23 class ArrayAttr;
24 
25 /// Computes and returns the linearized index of 'offsets' w.r.t. 'basis'.
26 int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);
27 
28 /// Given the strides together with a linear index in the dimension
29 /// space, returns the vector-space offsets in each dimension for a
30 /// de-linearized index.
31 SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> strides,
32  int64_t linearIndex);
33 
34 /// Apply the permutation defined by `permutation` to `inVec`.
35 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
36 /// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector
37 /// `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a', 'b']`.
38 template <typename T, unsigned N>
40  ArrayRef<int64_t> permutation) {
41  SmallVector<T, N> auxVec(inVec.size());
42  for (const auto &en : enumerate(permutation))
43  auxVec[en.index()] = inVec[en.value()];
44  inVec = auxVec;
45 }
46 
47 /// Helper that returns a subset of `arrayAttr` as a vector of int64_t.
48 SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr,
49  unsigned dropFront = 0,
50  unsigned dropBack = 0);
51 
52 /// Computes and returns linearized affine expression w.r.t. `basis`.
54 
55 /// Given the strides in the dimension space, returns the affine expressions for
56 /// vector-space offsets in each dimension for a de-linearized index.
59 
60 } // namespace mlir
61 
62 #endif // MLIR_DIALECT_UTILS_INDEXINGUTILS_H
Include the generated interface declarations.
void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)
Apply the permutation defined by permutation to inVec.
Definition: IndexingUtils.h:39
SmallVector< int64_t, 4 > getI64SubArray(ArrayAttr arrayAttr, unsigned dropFront=0, unsigned dropBack=0)
Helper that returns a subset of arrayAttr as a vector of int64_t.
mlir::AffineExpr getLinearAffineExpr(ArrayRef< int64_t > basis, mlir::Builder &b)
Computes and returns linearized affine expression w.r.t. basis.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:232
Base type for affine expression.
Definition: AffineExpr.h:68
SmallVector< int64_t, 4 > delinearize(ArrayRef< int64_t > strides, int64_t linearIndex)
Given the strides together with a linear index in the dimension space, returns the vector-space offse...
SmallVector< mlir::AffineExpr, 4 > getDelinearizedAffineExpr(ArrayRef< int64_t > strides, mlir::Builder &b)
Given the strides in the dimension space, returns the affine expressions for vector-space offsets in ...
void dropFront(int64_t arr[N], int64_t *res)
Definition: CRunnerUtils.h:118
This class is a general helper class for creating context-global objects like types, attributes, and affine expressions.
Definition: Builders.h:49
int64_t linearize(ArrayRef< int64_t > offsets, ArrayRef< int64_t > basis)
Computes and returns the linearized index of &#39;offsets&#39; w.r.t. &#39;basis&#39;.