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 
9 #include "mlir/Dialect/Utils/IndexingUtils.h"
10 
11 #include "mlir/IR/BuiltinAttributes.h"
12 
linearize(ArrayRef<int64_t> offsets,ArrayRef<int64_t> basis)13 int64_t mlir::linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis) {
14   assert(offsets.size() == basis.size());
15   int64_t linearIndex = 0;
16   for (unsigned idx = 0, e = basis.size(); idx < e; ++idx)
17     linearIndex += offsets[idx] * basis[idx];
18   return linearIndex;
19 }
20 
delinearize(ArrayRef<int64_t> sliceStrides,int64_t index)21 llvm::SmallVector<int64_t, 4> mlir::delinearize(ArrayRef<int64_t> sliceStrides,
22                                                 int64_t index) {
23   int64_t rank = sliceStrides.size();
24   SmallVector<int64_t, 4> vectorOffsets(rank);
25   for (int64_t r = 0; r < rank; ++r) {
26     assert(sliceStrides[r] > 0);
27     vectorOffsets[r] = index / sliceStrides[r];
28     index %= sliceStrides[r];
29   }
30   return vectorOffsets;
31 }
32 
getI64SubArray(ArrayAttr arrayAttr,unsigned dropFront,unsigned dropBack)33 llvm::SmallVector<int64_t, 4> mlir::getI64SubArray(ArrayAttr arrayAttr,
34                                                    unsigned dropFront,
35                                                    unsigned dropBack) {
36   assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds");
37   auto range = arrayAttr.getAsRange<IntegerAttr>();
38   SmallVector<int64_t, 4> res;
39   res.reserve(arrayAttr.size() - dropFront - dropBack);
40   for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;
41        it != eit; ++it)
42     res.push_back((*it).getValue().getSExtValue());
43   return res;
44 }
45