188 lines
7.6 KiB
C++
188 lines
7.6 KiB
C++
//===- HomomorphismSimplification.h -----------------------------*- C++ -*-===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_TRANSFORMS_SIMPLIFY_HOMOMORPHISM_H_
|
|
#define MLIR_TRANSFORMS_SIMPLIFY_HOMOMORPHISM_H_
|
|
|
|
#include "mlir/IR/IRMapping.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "mlir/Support/LLVM.h"
|
|
#include "llvm/ADT/SmallVector.h"
|
|
#include "llvm/Support/Casting.h"
|
|
#include <iterator>
|
|
#include <optional>
|
|
#include <type_traits>
|
|
#include <utility>
|
|
|
|
namespace mlir {
|
|
|
|
// If `h` is an homomorphism with respect to the source algebraic structure
|
|
// induced by function `s` and the target algebraic structure induced by
|
|
// function `t`, transforms `s(h(x1), h(x2) ..., h(xn))` into
|
|
// `h(t(x1, x2, ..., xn))`.
|
|
//
|
|
// Functors:
|
|
// ---------
|
|
// `GetHomomorphismOpOperandFn`: `(Operation*) -> OpOperand*`
|
|
// Returns the operand relevant to the homomorphism.
|
|
// There may be other operands that are not relevant.
|
|
//
|
|
// `GetHomomorphismOpResultFn`: `(Operation*) -> OpResult`
|
|
// Returns the result relevant to the homomorphism.
|
|
//
|
|
// `GetSourceAlgebraicOpOperandsFn`: `(Operation*, SmallVector<OpOperand*>&) ->
|
|
// void` Populates into the vector the operands relevant to the homomorphism.
|
|
//
|
|
// `GetSourceAlgebraicOpResultFn`: `(Operation*) -> OpResult`
|
|
// Return the result of the source algebraic operation relevant to the
|
|
// homomorphism.
|
|
//
|
|
// `GetTargetAlgebraicOpResultFn`: `(Operation*) -> OpResult`
|
|
// Return the result of the target algebraic operation relevant to the
|
|
// homomorphism.
|
|
//
|
|
// `IsHomomorphismOpFn`: `(Operation*, std::optional<Operation*>) -> bool`
|
|
// Check if the operation is an homomorphism of the required type.
|
|
// Additionally if the optional is present checks if the operations are
|
|
// compatible homomorphisms.
|
|
//
|
|
// `IsSourceAlgebraicOpFn`: `(Operation*) -> bool`
|
|
// Check if the operation is an operation of the algebraic structure.
|
|
//
|
|
// `CreateTargetAlgebraicOpFn`: `(Operation*, IRMapping& operandsRemapping,
|
|
// PatternRewriter &rewriter) -> Operation*`
|
|
template <typename GetHomomorphismOpOperandFn,
|
|
typename GetHomomorphismOpResultFn,
|
|
typename GetSourceAlgebraicOpOperandsFn,
|
|
typename GetSourceAlgebraicOpResultFn,
|
|
typename GetTargetAlgebraicOpResultFn, typename IsHomomorphismOpFn,
|
|
typename IsSourceAlgebraicOpFn, typename CreateTargetAlgebraicOpFn>
|
|
struct HomomorphismSimplification : public RewritePattern {
|
|
template <typename GetHomomorphismOpOperandFnArg,
|
|
typename GetHomomorphismOpResultFnArg,
|
|
typename GetSourceAlgebraicOpOperandsFnArg,
|
|
typename GetSourceAlgebraicOpResultFnArg,
|
|
typename GetTargetAlgebraicOpResultFnArg,
|
|
typename IsHomomorphismOpFnArg, typename IsSourceAlgebraicOpFnArg,
|
|
typename CreateTargetAlgebraicOpFnArg,
|
|
typename... RewritePatternArgs>
|
|
HomomorphismSimplification(
|
|
GetHomomorphismOpOperandFnArg &&getHomomorphismOpOperand,
|
|
GetHomomorphismOpResultFnArg &&getHomomorphismOpResult,
|
|
GetSourceAlgebraicOpOperandsFnArg &&getSourceAlgebraicOpOperands,
|
|
GetSourceAlgebraicOpResultFnArg &&getSourceAlgebraicOpResult,
|
|
GetTargetAlgebraicOpResultFnArg &&getTargetAlgebraicOpResult,
|
|
IsHomomorphismOpFnArg &&isHomomorphismOp,
|
|
IsSourceAlgebraicOpFnArg &&isSourceAlgebraicOp,
|
|
CreateTargetAlgebraicOpFnArg &&createTargetAlgebraicOpFn,
|
|
RewritePatternArgs &&...args)
|
|
: RewritePattern(std::forward<RewritePatternArgs>(args)...),
|
|
getHomomorphismOpOperand(std::forward<GetHomomorphismOpOperandFnArg>(
|
|
getHomomorphismOpOperand)),
|
|
getHomomorphismOpResult(std::forward<GetHomomorphismOpResultFnArg>(
|
|
getHomomorphismOpResult)),
|
|
getSourceAlgebraicOpOperands(
|
|
std::forward<GetSourceAlgebraicOpOperandsFnArg>(
|
|
getSourceAlgebraicOpOperands)),
|
|
getSourceAlgebraicOpResult(
|
|
std::forward<GetSourceAlgebraicOpResultFnArg>(
|
|
getSourceAlgebraicOpResult)),
|
|
getTargetAlgebraicOpResult(
|
|
std::forward<GetTargetAlgebraicOpResultFnArg>(
|
|
getTargetAlgebraicOpResult)),
|
|
isHomomorphismOp(std::forward<IsHomomorphismOpFnArg>(isHomomorphismOp)),
|
|
isSourceAlgebraicOp(
|
|
std::forward<IsSourceAlgebraicOpFnArg>(isSourceAlgebraicOp)),
|
|
createTargetAlgebraicOpFn(std::forward<CreateTargetAlgebraicOpFnArg>(
|
|
createTargetAlgebraicOpFn)) {}
|
|
|
|
LogicalResult matchAndRewrite(Operation *op,
|
|
PatternRewriter &rewriter) const override {
|
|
SmallVector<OpOperand *> algebraicOpOperands;
|
|
if (failed(matchOp(op, algebraicOpOperands))) {
|
|
return failure();
|
|
}
|
|
return rewriteOp(op, algebraicOpOperands, rewriter);
|
|
}
|
|
|
|
private:
|
|
LogicalResult
|
|
matchOp(Operation *sourceAlgebraicOp,
|
|
SmallVector<OpOperand *> &sourceAlgebraicOpOperands) const {
|
|
if (!isSourceAlgebraicOp(sourceAlgebraicOp)) {
|
|
return failure();
|
|
}
|
|
sourceAlgebraicOpOperands.clear();
|
|
getSourceAlgebraicOpOperands(sourceAlgebraicOp, sourceAlgebraicOpOperands);
|
|
if (sourceAlgebraicOpOperands.empty()) {
|
|
return failure();
|
|
}
|
|
|
|
Operation *firstHomomorphismOp =
|
|
sourceAlgebraicOpOperands.front()->get().getDefiningOp();
|
|
if (!firstHomomorphismOp ||
|
|
!isHomomorphismOp(firstHomomorphismOp, std::nullopt)) {
|
|
return failure();
|
|
}
|
|
OpResult firstHomomorphismOpResult =
|
|
getHomomorphismOpResult(firstHomomorphismOp);
|
|
if (firstHomomorphismOpResult != sourceAlgebraicOpOperands.front()->get()) {
|
|
return failure();
|
|
}
|
|
|
|
for (auto operand : sourceAlgebraicOpOperands) {
|
|
Operation *homomorphismOp = operand->get().getDefiningOp();
|
|
if (!homomorphismOp ||
|
|
!isHomomorphismOp(homomorphismOp, firstHomomorphismOp)) {
|
|
return failure();
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
rewriteOp(Operation *sourceAlgebraicOp,
|
|
const SmallVector<OpOperand *> &sourceAlgebraicOpOperands,
|
|
PatternRewriter &rewriter) const {
|
|
IRMapping irMapping;
|
|
for (auto operand : sourceAlgebraicOpOperands) {
|
|
Operation *homomorphismOp = operand->get().getDefiningOp();
|
|
irMapping.map(operand->get(),
|
|
getHomomorphismOpOperand(homomorphismOp)->get());
|
|
}
|
|
Operation *targetAlgebraicOp =
|
|
createTargetAlgebraicOpFn(sourceAlgebraicOp, irMapping, rewriter);
|
|
|
|
irMapping.clear();
|
|
assert(!sourceAlgebraicOpOperands.empty());
|
|
Operation *firstHomomorphismOp =
|
|
sourceAlgebraicOpOperands[0]->get().getDefiningOp();
|
|
irMapping.map(getHomomorphismOpOperand(firstHomomorphismOp)->get(),
|
|
getTargetAlgebraicOpResult(targetAlgebraicOp));
|
|
Operation *newHomomorphismOp =
|
|
rewriter.clone(*firstHomomorphismOp, irMapping);
|
|
rewriter.replaceAllUsesWith(getSourceAlgebraicOpResult(sourceAlgebraicOp),
|
|
getHomomorphismOpResult(newHomomorphismOp));
|
|
return success();
|
|
}
|
|
|
|
GetHomomorphismOpOperandFn getHomomorphismOpOperand;
|
|
GetHomomorphismOpResultFn getHomomorphismOpResult;
|
|
GetSourceAlgebraicOpOperandsFn getSourceAlgebraicOpOperands;
|
|
GetSourceAlgebraicOpResultFn getSourceAlgebraicOpResult;
|
|
GetTargetAlgebraicOpResultFn getTargetAlgebraicOpResult;
|
|
IsHomomorphismOpFn isHomomorphismOp;
|
|
IsSourceAlgebraicOpFn isSourceAlgebraicOp;
|
|
CreateTargetAlgebraicOpFn createTargetAlgebraicOpFn;
|
|
};
|
|
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_TRANSFORMS_SIMPLIFY_HOMOMORPHISM_H_
|