419 lines
17 KiB
C++
419 lines
17 KiB
C++
//===- CommonFolders.h - Common Operation Folders----------------*- 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This header file declares various common operation folders. These folders
|
|
// are intended to be used by dialects to support common folding behavior
|
|
// without requiring each dialect to provide its own implementation.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_DIALECT_COMMONFOLDERS_H
|
|
#define MLIR_DIALECT_COMMONFOLDERS_H
|
|
|
|
#include "mlir/IR/Attributes.h"
|
|
#include "mlir/IR/BuiltinAttributeInterfaces.h"
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypeInterfaces.h"
|
|
#include "mlir/IR/Types.h"
|
|
#include "llvm/ADT/ArrayRef.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
|
|
#include <cassert>
|
|
#include <cstddef>
|
|
#include <optional>
|
|
|
|
namespace mlir {
|
|
namespace ub {
|
|
class PoisonAttr;
|
|
}
|
|
/// Performs constant folding `calculate` with element-wise behavior on the two
|
|
/// attributes in `operands` and returns the result if possible.
|
|
/// Uses `resultType` for the type of the returned attribute.
|
|
/// Optional PoisonAttr template argument allows to specify 'poison' attribute
|
|
/// which will be directly propagated to result.
|
|
template <class LAttrElementT, class RAttrElementT = LAttrElementT,
|
|
class LElementValueT = typename LAttrElementT::ValueType,
|
|
class RElementValueT = typename RAttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class ResultAttrElementT = LAttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT = function_ref<std::optional<ResultElementValueT>(
|
|
LElementValueT, RElementValueT)>>
|
|
Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
|
|
Type resultType,
|
|
CalculationT &&calculate) {
|
|
assert(operands.size() == 2 && "binary op takes two operands");
|
|
static_assert(
|
|
std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
|
|
"PoisonAttr is undefined, either add a dependency on UB dialect or pass "
|
|
"void as template argument to opt-out from poison semantics.");
|
|
if constexpr (!std::is_void_v<PoisonAttr>) {
|
|
if (isa_and_nonnull<PoisonAttr>(operands[0]))
|
|
return operands[0];
|
|
|
|
if (isa_and_nonnull<PoisonAttr>(operands[1]))
|
|
return operands[1];
|
|
}
|
|
|
|
if (!resultType || !operands[0] || !operands[1])
|
|
return {};
|
|
|
|
if (isa<LAttrElementT>(operands[0]) && isa<RAttrElementT>(operands[1])) {
|
|
auto lhs = cast<LAttrElementT>(operands[0]);
|
|
auto rhs = cast<RAttrElementT>(operands[1]);
|
|
if constexpr (std::is_same_v<LElementValueT, RElementValueT>)
|
|
if (lhs.getType() != rhs.getType())
|
|
return {};
|
|
|
|
auto calRes = calculate(lhs.getValue(), rhs.getValue());
|
|
|
|
if (!calRes)
|
|
return {};
|
|
|
|
return ResultAttrElementT::get(resultType, *calRes);
|
|
}
|
|
|
|
if (isa<SplatElementsAttr>(operands[0]) &&
|
|
isa<SplatElementsAttr>(operands[1])) {
|
|
// Both operands are splats so we can avoid expanding the values out and
|
|
// just fold based on the splat value.
|
|
auto lhs = cast<SplatElementsAttr>(operands[0]);
|
|
auto rhs = cast<SplatElementsAttr>(operands[1]);
|
|
if constexpr (std::is_same_v<LElementValueT, RElementValueT>)
|
|
if (lhs.getType() != rhs.getType())
|
|
return {};
|
|
|
|
auto elementResult = calculate(lhs.getSplatValue<LElementValueT>(),
|
|
rhs.getSplatValue<RElementValueT>());
|
|
if (!elementResult)
|
|
return {};
|
|
|
|
return DenseElementsAttr::get(cast<ShapedType>(resultType), *elementResult);
|
|
}
|
|
|
|
if (isa<ElementsAttr>(operands[0]) && isa<ElementsAttr>(operands[1])) {
|
|
// Operands are ElementsAttr-derived; perform an element-wise fold by
|
|
// expanding the values.
|
|
auto lhs = cast<ElementsAttr>(operands[0]);
|
|
auto rhs = cast<ElementsAttr>(operands[1]);
|
|
if constexpr (std::is_same_v<LElementValueT, RElementValueT>)
|
|
if (lhs.getType() != rhs.getType())
|
|
return {};
|
|
|
|
auto maybeLhsIt = lhs.try_value_begin<LElementValueT>();
|
|
auto maybeRhsIt = rhs.try_value_begin<RElementValueT>();
|
|
if (!maybeLhsIt || !maybeRhsIt)
|
|
return {};
|
|
auto lhsIt = *maybeLhsIt;
|
|
auto rhsIt = *maybeRhsIt;
|
|
SmallVector<ResultElementValueT, 4> elementResults;
|
|
elementResults.reserve(lhs.getNumElements());
|
|
for (size_t i = 0, e = lhs.getNumElements(); i < e; ++i, ++lhsIt, ++rhsIt) {
|
|
auto elementResult = calculate(*lhsIt, *rhsIt);
|
|
if (!elementResult)
|
|
return {};
|
|
elementResults.push_back(*elementResult);
|
|
}
|
|
|
|
return DenseElementsAttr::get(cast<ShapedType>(resultType), elementResults);
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Performs constant folding `calculate` with element-wise behavior on the two
|
|
/// attributes in `operands` and returns the result if possible.
|
|
/// Uses the operand element type for the element type of the returned
|
|
/// attribute.
|
|
/// Optional PoisonAttr template argument allows to specify 'poison' attribute
|
|
/// which will be directly propagated to result.
|
|
template <class LAttrElementT, class RAttrElementT = LAttrElementT,
|
|
class LElementValueT = typename LAttrElementT::ValueType,
|
|
class RElementValueT = typename RAttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class ResultAttrElementT = LAttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT = function_ref<std::optional<ResultElementValueT>(
|
|
LElementValueT, RElementValueT)>>
|
|
Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
|
|
CalculationT &&calculate) {
|
|
assert(operands.size() == 2 && "binary op takes two operands");
|
|
static_assert(
|
|
std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
|
|
"PoisonAttr is undefined, either add a dependency on UB dialect or pass "
|
|
"void as template argument to opt-out from poison semantics.");
|
|
if constexpr (!std::is_void_v<PoisonAttr>) {
|
|
if (isa_and_nonnull<PoisonAttr>(operands[0]))
|
|
return operands[0];
|
|
|
|
if (isa_and_nonnull<PoisonAttr>(operands[1]))
|
|
return operands[1];
|
|
}
|
|
|
|
auto getAttrType = [](Attribute attr) -> Type {
|
|
if (auto typed = dyn_cast_or_null<TypedAttr>(attr))
|
|
return typed.getType();
|
|
return {};
|
|
};
|
|
|
|
Type lhsType = getAttrType(operands[0]);
|
|
Type rhsType = getAttrType(operands[1]);
|
|
if (!lhsType || !rhsType)
|
|
return {};
|
|
if constexpr (std::is_same_v<LElementValueT, RElementValueT>)
|
|
if (lhsType != rhsType)
|
|
return {};
|
|
|
|
return constFoldBinaryOpConditional<
|
|
LAttrElementT, RAttrElementT, LElementValueT, RElementValueT, PoisonAttr,
|
|
ResultAttrElementT, ResultElementValueT, CalculationT>(
|
|
operands, lhsType, std::forward<CalculationT>(calculate));
|
|
}
|
|
|
|
template <class LAttrElementT, class RAttrElementT = LAttrElementT,
|
|
class LElementValueT = typename LAttrElementT::ValueType,
|
|
class RElementValueT = typename RAttrElementT::ValueType,
|
|
class PoisonAttr = void, //
|
|
class ResultAttrElementT = LAttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT =
|
|
function_ref<ResultElementValueT(LElementValueT, RElementValueT)>>
|
|
Attribute constFoldBinaryOp(ArrayRef<Attribute> operands, Type resultType,
|
|
CalculationT &&calculate) {
|
|
return constFoldBinaryOpConditional<LAttrElementT, RAttrElementT,
|
|
LElementValueT, RElementValueT,
|
|
PoisonAttr, ResultAttrElementT>(
|
|
operands, resultType,
|
|
[&](LElementValueT a, RElementValueT b)
|
|
-> std::optional<ResultElementValueT> { return calculate(a, b); });
|
|
}
|
|
|
|
template <class LAttrElementT, class RAttrElementT = LAttrElementT,
|
|
class LElementValueT = typename LAttrElementT::ValueType,
|
|
class RElementValueT = typename RAttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class ResultAttrElementT = LAttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT =
|
|
function_ref<ResultElementValueT(LElementValueT, RElementValueT)>>
|
|
Attribute constFoldBinaryOp(ArrayRef<Attribute> operands,
|
|
CalculationT &&calculate) {
|
|
return constFoldBinaryOpConditional<LAttrElementT, RAttrElementT,
|
|
LElementValueT, RElementValueT,
|
|
PoisonAttr, ResultAttrElementT>(
|
|
operands,
|
|
[&](LElementValueT a, RElementValueT b)
|
|
-> std::optional<ResultElementValueT> { return calculate(a, b); });
|
|
}
|
|
|
|
/// Performs constant folding `calculate` with element-wise behavior on the one
|
|
/// attributes in `operands` and returns the result if possible.
|
|
/// Uses `resultType` for the type of the returned attribute.
|
|
/// Optional PoisonAttr template argument allows to specify 'poison' attribute
|
|
/// which will be directly propagated to result.
|
|
template <class AttrElementT, //
|
|
class ElementValueT = typename AttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class ResultAttrElementT = AttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT =
|
|
function_ref<std::optional<ResultElementValueT>(ElementValueT)>>
|
|
Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
|
|
Type resultType,
|
|
CalculationT &&calculate) {
|
|
if (!resultType || !llvm::getSingleElement(operands))
|
|
return {};
|
|
|
|
static_assert(
|
|
std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
|
|
"PoisonAttr is undefined, either add a dependency on UB dialect or pass "
|
|
"void as template argument to opt-out from poison semantics.");
|
|
if constexpr (!std::is_void_v<PoisonAttr>) {
|
|
if (isa<PoisonAttr>(operands[0]))
|
|
return operands[0];
|
|
}
|
|
|
|
if (isa<AttrElementT>(operands[0])) {
|
|
auto op = cast<AttrElementT>(operands[0]);
|
|
|
|
auto res = calculate(op.getValue());
|
|
if (!res)
|
|
return {};
|
|
return ResultAttrElementT::get(resultType, *res);
|
|
}
|
|
if (isa<SplatElementsAttr>(operands[0])) {
|
|
// Both operands are splats so we can avoid expanding the values out and
|
|
// just fold based on the splat value.
|
|
auto op = cast<SplatElementsAttr>(operands[0]);
|
|
|
|
auto elementResult = calculate(op.getSplatValue<ElementValueT>());
|
|
if (!elementResult)
|
|
return {};
|
|
return DenseElementsAttr::get(cast<ShapedType>(resultType), *elementResult);
|
|
} else if (isa<ElementsAttr>(operands[0])) {
|
|
// Operands are ElementsAttr-derived; perform an element-wise fold by
|
|
// expanding the values.
|
|
auto op = cast<ElementsAttr>(operands[0]);
|
|
|
|
auto maybeOpIt = op.try_value_begin<ElementValueT>();
|
|
if (!maybeOpIt)
|
|
return {};
|
|
auto opIt = *maybeOpIt;
|
|
SmallVector<ResultElementValueT> elementResults;
|
|
elementResults.reserve(op.getNumElements());
|
|
for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt) {
|
|
auto elementResult = calculate(*opIt);
|
|
if (!elementResult)
|
|
return {};
|
|
elementResults.push_back(*elementResult);
|
|
}
|
|
return DenseElementsAttr::get(cast<ShapedType>(resultType), elementResults);
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Performs constant folding `calculate` with element-wise behavior on the one
|
|
/// attributes in `operands` and returns the result if possible.
|
|
/// Uses the operand element type for the element type of the returned
|
|
/// attribute.
|
|
/// Optional PoisonAttr template argument allows to specify 'poison' attribute
|
|
/// which will be directly propagated to result.
|
|
template <class AttrElementT, //
|
|
class ElementValueT = typename AttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class ResultAttrElementT = AttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT =
|
|
function_ref<std::optional<ResultElementValueT>(ElementValueT)>>
|
|
Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
|
|
CalculationT &&calculate) {
|
|
if (!llvm::getSingleElement(operands))
|
|
return {};
|
|
|
|
static_assert(
|
|
std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
|
|
"PoisonAttr is undefined, either add a dependency on UB dialect or pass "
|
|
"void as template argument to opt-out from poison semantics.");
|
|
if constexpr (!std::is_void_v<PoisonAttr>) {
|
|
if (isa<PoisonAttr>(operands[0]))
|
|
return operands[0];
|
|
}
|
|
|
|
auto getAttrType = [](Attribute attr) -> Type {
|
|
if (auto typed = dyn_cast_or_null<TypedAttr>(attr))
|
|
return typed.getType();
|
|
return {};
|
|
};
|
|
|
|
Type operandType = getAttrType(operands[0]);
|
|
if (!operandType)
|
|
return {};
|
|
|
|
return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
|
|
ResultAttrElementT, ResultElementValueT,
|
|
CalculationT>(
|
|
operands, operandType, std::forward<CalculationT>(calculate));
|
|
}
|
|
|
|
template <class AttrElementT, //
|
|
class ElementValueT = typename AttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class ResultAttrElementT = AttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT = function_ref<ResultElementValueT(ElementValueT)>>
|
|
Attribute constFoldUnaryOp(ArrayRef<Attribute> operands, Type resultType,
|
|
CalculationT &&calculate) {
|
|
return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
|
|
ResultAttrElementT>(
|
|
operands, resultType,
|
|
[&](ElementValueT a) -> std::optional<ResultElementValueT> {
|
|
return calculate(a);
|
|
});
|
|
}
|
|
|
|
template <class AttrElementT, //
|
|
class ElementValueT = typename AttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class ResultAttrElementT = AttrElementT,
|
|
class ResultElementValueT = typename ResultAttrElementT::ValueType,
|
|
class CalculationT = function_ref<ResultElementValueT(ElementValueT)>>
|
|
Attribute constFoldUnaryOp(ArrayRef<Attribute> operands,
|
|
CalculationT &&calculate) {
|
|
return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
|
|
ResultAttrElementT>(
|
|
operands, [&](ElementValueT a) -> std::optional<ResultElementValueT> {
|
|
return calculate(a);
|
|
});
|
|
}
|
|
|
|
template <
|
|
class AttrElementT, class TargetAttrElementT,
|
|
class ElementValueT = typename AttrElementT::ValueType,
|
|
class TargetElementValueT = typename TargetAttrElementT::ValueType,
|
|
class PoisonAttr = ub::PoisonAttr,
|
|
class CalculationT = function_ref<TargetElementValueT(ElementValueT, bool)>>
|
|
Attribute constFoldCastOp(ArrayRef<Attribute> operands, Type resType,
|
|
CalculationT &&calculate) {
|
|
if (!llvm::getSingleElement(operands))
|
|
return {};
|
|
|
|
static_assert(
|
|
std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
|
|
"PoisonAttr is undefined, either add a dependency on UB dialect or pass "
|
|
"void as template argument to opt-out from poison semantics.");
|
|
if constexpr (!std::is_void_v<PoisonAttr>) {
|
|
if (isa<PoisonAttr>(operands[0]))
|
|
return operands[0];
|
|
}
|
|
|
|
if (isa<AttrElementT>(operands[0])) {
|
|
auto op = cast<AttrElementT>(operands[0]);
|
|
bool castStatus = true;
|
|
auto res = calculate(op.getValue(), castStatus);
|
|
if (!castStatus)
|
|
return {};
|
|
return TargetAttrElementT::get(resType, res);
|
|
}
|
|
if (isa<SplatElementsAttr>(operands[0])) {
|
|
// The operand is a splat so we can avoid expanding the values out and
|
|
// just fold based on the splat value.
|
|
auto op = cast<SplatElementsAttr>(operands[0]);
|
|
bool castStatus = true;
|
|
auto elementResult =
|
|
calculate(op.getSplatValue<ElementValueT>(), castStatus);
|
|
if (!castStatus)
|
|
return {};
|
|
auto shapedResType = cast<ShapedType>(resType);
|
|
if (!shapedResType.hasStaticShape())
|
|
return {};
|
|
return DenseElementsAttr::get(shapedResType, elementResult);
|
|
}
|
|
if (auto op = dyn_cast<ElementsAttr>(operands[0])) {
|
|
// Operand is ElementsAttr-derived; perform an element-wise fold by
|
|
// expanding the value.
|
|
bool castStatus = true;
|
|
auto maybeOpIt = op.try_value_begin<ElementValueT>();
|
|
if (!maybeOpIt)
|
|
return {};
|
|
auto opIt = *maybeOpIt;
|
|
SmallVector<TargetElementValueT> elementResults;
|
|
elementResults.reserve(op.getNumElements());
|
|
for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt) {
|
|
auto elt = calculate(*opIt, castStatus);
|
|
if (!castStatus)
|
|
return {};
|
|
elementResults.push_back(elt);
|
|
}
|
|
|
|
return DenseElementsAttr::get(cast<ShapedType>(resType), elementResults);
|
|
}
|
|
return {};
|
|
}
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_DIALECT_COMMONFOLDERS_H
|