MLIR  17.0.0git
IndexingUtils.cpp
Go to the documentation of this file.
1 //===- IndexingUtils.cpp - Helpers related to index computations ----------===//
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 
10 
11 #include "mlir/IR/AffineExpr.h"
12 #include "mlir/IR/Builders.h"
14 
15 #include <numeric>
16 #include <optional>
17 
18 using namespace mlir;
19 
21  SmallVector<int64_t> strides(sizes.size(), 1);
22  for (int64_t r = strides.size() - 2; r >= 0; --r)
23  strides[r] = strides[r + 1] * sizes[r + 1];
24  return strides;
25 }
26 
28  ArrayRef<int64_t> v2) {
29  SmallVector<int64_t> result;
30  for (auto it : llvm::zip(v1, v2))
31  result.push_back(std::get<0>(it) * std::get<1>(it));
32  return result;
33 }
34 
35 std::optional<SmallVector<int64_t>>
37  if (shape.size() < subShape.size())
38  return std::nullopt;
39  assert(llvm::all_of(shape, [](int64_t s) { return s > 0; }) &&
40  "shape must be nonnegative");
41  assert(llvm::all_of(subShape, [](int64_t s) { return s > 0; }) &&
42  "subShape must be nonnegative");
43 
44  // Starting from the end, compute the integer divisors.
45  std::vector<int64_t> result;
46  result.reserve(shape.size());
47  for (auto [size, subSize] :
48  llvm::zip(llvm::reverse(shape), llvm::reverse(subShape))) {
49  // If integral division does not occur, return and let the caller decide.
50  if (size % subSize != 0)
51  return std::nullopt;
52  result.push_back(size / subSize);
53  }
54  // At this point we computed the ratio (in reverse) for the common size.
55  // Fill with the remaining entries from the shape (still in reverse).
56  int commonSize = subShape.size();
57  std::copy(shape.rbegin() + commonSize, shape.rend(),
58  std::back_inserter(result));
59  // Reverse again to get it back in the proper order and return.
60  return SmallVector<int64_t>{result.rbegin(), result.rend()};
61 }
62 
64  assert(offsets.size() == basis.size());
65  int64_t linearIndex = 0;
66  for (unsigned idx = 0, e = basis.size(); idx < e; ++idx)
67  linearIndex += offsets[idx] * basis[idx];
68  return linearIndex;
69 }
70 
72  int64_t index) {
73  int64_t rank = sliceStrides.size();
74  SmallVector<int64_t> vectorOffsets(rank);
75  for (int64_t r = 0; r < rank; ++r) {
76  assert(sliceStrides[r] > 0);
77  vectorOffsets[r] = index / sliceStrides[r];
78  index %= sliceStrides[r];
79  }
80  return vectorOffsets;
81 }
82 
84  if (basis.empty())
85  return 0;
86  return std::accumulate(basis.begin(), basis.end(), 1,
87  std::multiplies<int64_t>());
88 }
89 
92  SmallVector<int64_t> inversion(permutation.size());
93  for (const auto &pos : llvm::enumerate(permutation)) {
94  inversion[pos.value()] = pos.index();
95  }
96  return inversion;
97 }
98 
100  llvm::SmallDenseSet<int64_t, 4> seenVals;
101  for (auto val : interchange) {
102  if (seenVals.count(val))
103  return false;
104  seenVals.insert(val);
105  }
106  return seenVals.size() == interchange.size();
107 }
108 
110  unsigned dropFront,
111  unsigned dropBack) {
112  assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds");
113  auto range = arrayAttr.getAsRange<IntegerAttr>();
115  res.reserve(arrayAttr.size() - dropFront - dropBack);
116  for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;
117  it != eit; ++it)
118  res.push_back((*it).getValue().getSExtValue());
119  return res;
120 }
121 
123  mlir::Builder &b) {
124  AffineExpr resultExpr = b.getAffineDimExpr(0);
125  resultExpr = resultExpr * basis[0];
126  for (unsigned i = 1; i < basis.size(); i++)
127  resultExpr = resultExpr + b.getAffineDimExpr(i) * basis[i];
128  return resultExpr;
129 }
130 
133  AffineExpr resultExpr = b.getAffineDimExpr(0);
134  int64_t rank = strides.size();
135  SmallVector<AffineExpr> vectorOffsets(rank);
136  vectorOffsets[0] = resultExpr.floorDiv(strides[0]);
137  resultExpr = resultExpr % strides[0];
138  for (unsigned i = 1; i < rank; i++) {
139  vectorOffsets[i] = resultExpr.floorDiv(strides[i]);
140  resultExpr = resultExpr % strides[i];
141  }
142  return vectorOffsets;
143 }
void dropFront(int64_t arr[N], int64_t *res)
Definition: CRunnerUtils.h:118
static void copy(Location loc, Value dst, Value src, Value size, OpBuilder &builder)
Copies the given number of bytes from src to dst pointers.
Base type for affine expression.
Definition: AffineExpr.h:68
AffineExpr floorDiv(uint64_t v) const
Definition: AffineExpr.cpp:778
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:50
AffineExpr getAffineDimExpr(unsigned position)
Definition: Builders.cpp:327
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:223
Include the generated interface declarations.
SmallVector< int64_t > computeElementwiseMul(ArrayRef< int64_t > v1, ArrayRef< int64_t > v2)
Return a vector containing llvm::zip of v1 and v2 multiplied elementwise.
SmallVector< int64_t > computeStrides(ArrayRef< int64_t > sizes)
Given a set of sizes, compute and return the strides (i.e.
SmallVector< mlir::AffineExpr > getDelinearizedAffineExpr(ArrayRef< int64_t > strides, mlir::Builder &b)
Given the strides in the dimension space, returns the affine expressions for vector-space offsets in ...
SmallVector< int64_t > 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< int64_t > 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.
SmallVector< int64_t > invertPermutationVector(ArrayRef< int64_t > permutation)
Helper method to apply to inverse a permutation.
int64_t computeMaxLinearIndex(ArrayRef< int64_t > basis)
Return the number of elements of basis (i.e.
int64_t linearize(ArrayRef< int64_t > offsets, ArrayRef< int64_t > basis)
Computes and returns the linearized index of 'offsets' w.r.t. 'basis'.
std::optional< SmallVector< int64_t > > computeShapeRatio(ArrayRef< int64_t > shape, ArrayRef< int64_t > subShape)
Compute and return the multi-dimensional integral ratio of subShape to the trailing dimensions of sha...
bool isPermutationVector(ArrayRef< int64_t > interchange)
Method to check if an interchange vector is a permutation.