llvm/mlir/lib/Dialect/Tensor/Transforms/ExtractSliceFromReshapeUtils.cpp

//===- ExtractSliceFromReshapeUtils.cpp - Slice reshape rewrites ----------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements rewrites that replace slices of reshape results with
// aggregated slices of the reshape source.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/TransformUtils.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "llvm/ADT/STLExtras.h"

usingnamespacemlir;
usingnamespacemlir::affine;
usingnamespacemlir::tensor;

/// A tuple that represents (dimension number, dimension value).
DimAndIndex;

/// Transform `dimAndIndex` from the output index space of a (non-rank-reducing)
/// slice described by `sliceParams` into the input index space.
static DimAndIndex invertSliceIndexing(OpBuilder &b, Location loc,
                                       ArrayRef<Range> sliceParams,
                                       const DimAndIndex &dimAndIndex) { … }

/// Transform `dimAndIndex` from the result tensor index space of a
/// CollapseShapeOp to the source tensor index space.
static ValueRange invertCollapseShapeIndexing(
    OpBuilder &b, Location loc, ArrayRef<ReassociationIndices> reassociation,
    ArrayRef<OpFoldResult> reshapeSourceShape, const DimAndIndex &dimAndIndex) { … }

FailureOr<ExtractSliceFromCollapseHelper>
tensor::ExtractSliceFromCollapseHelper::create(
    OpBuilder &b, tensor::CollapseShapeOp collapseOp,
    tensor::ExtractSliceOp extractOp) { … }

FailureOr<ExtractSliceFromCollapseHelper>
tensor::ExtractSliceFromCollapseHelper::create(OpBuilder &b,
                                               tensor::CollapseShapeOp op,
                                               ArrayRef<Range> sliceParams) { … }

std::pair<Value, SmallVector<Range>>
tensor::ExtractSliceFromCollapseHelper::emitLoopNestBody(
    OpBuilder &builder, Location loc, ValueRange tileInductionVars) { … }

FailureOr<Operation *>
tensor::simplifyCollapseShapeWithRankReducingExtractSlice(
    tensor::CollapseShapeOp op, RewriterBase &rewriter) { … }