IndexingUtils.cpp | Explore in Territory

//===- IndexingUtils.cpp - Helpers related to index computations ----------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/MLIRContext.h"
#include "llvm/ADT/STLExtras.h"
#include <numeric>
#include <optional>

usingnamespacemlir;

template <typename ExprType>
SmallVector<ExprType> computeSuffixProductImpl(ArrayRef<ExprType> sizes,
                                               ExprType unit) { … }

template <typename ExprType>
SmallVector<ExprType> computeElementwiseMulImpl(ArrayRef<ExprType> v1,
                                                ArrayRef<ExprType> v2) { … }

template <typename ExprType>
ExprType linearizeImpl(ArrayRef<ExprType> offsets, ArrayRef<ExprType> basis,
                       ExprType zero) { … }

template <typename ExprType, typename DivOpTy>
SmallVector<ExprType> delinearizeImpl(ExprType linearIndex,
                                      ArrayRef<ExprType> strides,
                                      DivOpTy divOp) { … }

//===----------------------------------------------------------------------===//
// Utils that operate on static integer values.
//===----------------------------------------------------------------------===//

SmallVector<int64_t> mlir::computeSuffixProduct(ArrayRef<int64_t> sizes) { … }

SmallVector<int64_t> mlir::computeElementwiseMul(ArrayRef<int64_t> v1,
                                                 ArrayRef<int64_t> v2) { … }

int64_t mlir::computeSum(ArrayRef<int64_t> basis) { … }

int64_t mlir::computeProduct(ArrayRef<int64_t> basis) { … }

int64_t mlir::linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis) { … }

SmallVector<int64_t> mlir::delinearize(int64_t linearIndex,
                                       ArrayRef<int64_t> strides) { … }

std::optional<SmallVector<int64_t>>
mlir::computeShapeRatio(ArrayRef<int64_t> shape, ArrayRef<int64_t> subShape) { … }

//===----------------------------------------------------------------------===//
// Utils that operate on AffineExpr.
//===----------------------------------------------------------------------===//

SmallVector<AffineExpr> mlir::computeSuffixProduct(ArrayRef<AffineExpr> sizes) { … }

SmallVector<AffineExpr> mlir::computeElementwiseMul(ArrayRef<AffineExpr> v1,
                                                    ArrayRef<AffineExpr> v2) { … }

AffineExpr mlir::computeSum(MLIRContext *ctx, ArrayRef<AffineExpr> basis) { … }

AffineExpr mlir::computeProduct(MLIRContext *ctx, ArrayRef<AffineExpr> basis) { … }

AffineExpr mlir::linearize(MLIRContext *ctx, ArrayRef<AffineExpr> offsets,
                           ArrayRef<AffineExpr> basis) { … }

AffineExpr mlir::linearize(MLIRContext *ctx, ArrayRef<AffineExpr> offsets,
                           ArrayRef<int64_t> basis) { … }

SmallVector<AffineExpr> mlir::delinearize(AffineExpr linearIndex,
                                          ArrayRef<AffineExpr> strides) { … }

SmallVector<AffineExpr> mlir::delinearize(AffineExpr linearIndex,
                                          ArrayRef<int64_t> strides) { … }

//===----------------------------------------------------------------------===//
// Permutation utils.
//===----------------------------------------------------------------------===//

SmallVector<int64_t>
mlir::invertPermutationVector(ArrayRef<int64_t> permutation) { … }

bool mlir::isIdentityPermutation(ArrayRef<int64_t> permutation) { … }

bool mlir::isPermutationVector(ArrayRef<int64_t> interchange) { … }

SmallVector<int64_t>
mlir::computePermutationVector(int64_t permSize, ArrayRef<int64_t> positions,
                               ArrayRef<int64_t> desiredPositions) { … }

SmallVector<int64_t> mlir::dropDims(ArrayRef<int64_t> inputPerm,
                                    ArrayRef<int64_t> dropPositions) { … }

SmallVector<int64_t> mlir::getI64SubArray(ArrayAttr arrayAttr,
                                          unsigned dropFront,
                                          unsigned dropBack) { … }

// TODO: do we have any common utily for this?
static MLIRContext *getContext(OpFoldResult val) { … }

std::pair<AffineExpr, SmallVector<OpFoldResult>>
mlir::computeLinearIndex(OpFoldResult sourceOffset,
                         ArrayRef<OpFoldResult> strides,
                         ArrayRef<OpFoldResult> indices) { … }

std::pair<AffineExpr, SmallVector<OpFoldResult>>
mlir::computeLinearIndex(OpFoldResult sourceOffset, ArrayRef<int64_t> strides,
                         ArrayRef<Value> indices) { … }

//===----------------------------------------------------------------------===//
// TileOffsetRange
//===----------------------------------------------------------------------===//

/// Apply left-padding by 1 to the tile shape if required.
static SmallVector<int64_t> padTileShapeToSize(ArrayRef<int64_t> tileShape,
                                               unsigned paddedSize) { … }

mlir::detail::TileOffsetRangeImpl::TileOffsetRangeImpl(
    ArrayRef<int64_t> shape, ArrayRef<int64_t> tileShape,
    ArrayRef<int64_t> loopOrder)
    : … { … }

SmallVector<int64_t> mlir::detail::TileOffsetRangeImpl::getStaticTileOffsets(
    int64_t linearIndex) const { … }

SmallVector<AffineExpr>
mlir::detail::TileOffsetRangeImpl::getDynamicTileOffsets(
    AffineExpr linearIndex) const { … }
llvm/mlir/lib/Dialect/Utils/IndexingUtils.cpp
