llvm/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

//===- TosaCanonicalizations.cpp - Canonicalization patterns & folders ----===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// \file
// TOSA canonicalization patterns and folders.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Quant/IR/Quant.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
#include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/InliningUtils.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/TypeSwitch.h"

#include <functional>

usingnamespacemlir;
usingnamespacemlir::tosa;

//===----------------------------------------------------------------------===//
// Operator Canonicalizers.
//===----------------------------------------------------------------------===//

struct ConcatOptimization : public OpRewritePattern<tosa::ConcatOp> { … };

void ConcatOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) { … }

LogicalResult SelectOp::canonicalize(SelectOp op, PatternRewriter &rewriter) { … }

struct ConsolidateTransposeOptimization
    : public OpRewritePattern<tosa::TransposeOp> { … };

// Determines the case when tosa.transpose is a tosa.reshape operation.
struct TransposeIsReshape : public OpRewritePattern<tosa::TransposeOp> { … };

void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) { … }

struct MaterializePadValue : public OpRewritePattern<tosa::PadOp> { … };

void PadOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                        MLIRContext *context) { … }

struct MaxPool2dIsNoOp : public OpRewritePattern<tosa::MaxPool2dOp> { … };

void MaxPool2dOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) { … }

struct ClampIsNoOp : public OpRewritePattern<tosa::ClampOp> { … };

struct ClampClampOptimization : public OpRewritePattern<tosa::ClampOp> { … };

void ClampOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) { … }

struct ConcatSliceOptimization : public OpRewritePattern<tosa::SliceOp> { … };

void SliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) { … }

//===----------------------------------------------------------------------===//
// Operator Folders.
//===----------------------------------------------------------------------===//

template <typename IntFolder, typename FloatFolder>
DenseElementsAttr binaryFolder(DenseElementsAttr lhs, DenseElementsAttr rhs,
                               RankedTensorType returnTy) { … }

static bool isSplatZero(Type elemType, DenseElementsAttr val) { … }

static bool isSplatOne(Type elemType, DenseElementsAttr val, int64_t shift) { … }

OpFoldResult AddOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult ArgMaxOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult IntDivOp::fold(FoldAdaptor adaptor) { … }

namespace {
DenseElementsAttr mulBinaryFolder(DenseElementsAttr lhs, DenseElementsAttr rhs,
                                  RankedTensorType ty, int32_t shift) { … }
} // namespace

OpFoldResult MulOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult SubOp::fold(FoldAdaptor adaptor) { … }

namespace {
template <typename Cmp>
struct ComparisonFold { … };

struct APIntFoldGreater { … };

struct APIntFoldGreaterEqual { … };
} // namespace

OpFoldResult GreaterOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult GreaterEqualOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult EqualOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult CastOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult ConstOp::fold(FoldAdaptor adaptor) { … }

#define REDUCE_FOLDER …

REDUCE_FOLDER(ReduceAllOp)
REDUCE_FOLDER(ReduceAnyOp)
REDUCE_FOLDER(ReduceMaxOp)
REDUCE_FOLDER(ReduceMinOp)
REDUCE_FOLDER(ReduceProdOp)
REDUCE_FOLDER(ReduceSumOp)
#undef REDUCE_FOLDER

OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult PadOp::fold(FoldAdaptor adaptor) { … }

// Fold away cases where a tosa.resize operation returns a copy
// of the input image.
OpFoldResult ResizeOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult ReverseOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult SliceOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult tosa::SelectOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult TileOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult tosa::LogOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult tosa::ExpOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult tosa::NegateOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult tosa::AbsOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult ConcatOp::fold(FoldAdaptor adaptor) { … }

OpFoldResult tosa::ReciprocalOp::fold(FoldAdaptor adaptor) { … }