| //===- 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/BuiltinAttributes.h" |
| #include "mlir/IR/BuiltinTypes.h" |
| #include "llvm/ADT/ArrayRef.h" |
| #include "llvm/ADT/STLExtras.h" |
| #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 AttrElementT, |
| class ElementValueT = typename AttrElementT::ValueType, |
| class PoisonAttr = ub::PoisonAttr, |
| class CalculationT = function_ref< |
| std::optional<ElementValueT>(ElementValueT, ElementValueT)>> |
| 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<AttrElementT>(operands[0]) && isa<AttrElementT>(operands[1])) { |
| auto lhs = cast<AttrElementT>(operands[0]); |
| auto rhs = cast<AttrElementT>(operands[1]); |
| if (lhs.getType() != rhs.getType()) |
| return {}; |
| |
| auto calRes = calculate(lhs.getValue(), rhs.getValue()); |
| |
| if (!calRes) |
| return {}; |
| |
| return AttrElementT::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 (lhs.getType() != rhs.getType()) |
| return {}; |
| |
| auto elementResult = calculate(lhs.getSplatValue<ElementValueT>(), |
| rhs.getSplatValue<ElementValueT>()); |
| 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 (lhs.getType() != rhs.getType()) |
| return {}; |
| |
| auto maybeLhsIt = lhs.try_value_begin<ElementValueT>(); |
| auto maybeRhsIt = rhs.try_value_begin<ElementValueT>(); |
| if (!maybeLhsIt || !maybeRhsIt) |
| return {}; |
| auto lhsIt = *maybeLhsIt; |
| auto rhsIt = *maybeRhsIt; |
| SmallVector<ElementValueT, 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 AttrElementT, |
| class ElementValueT = typename AttrElementT::ValueType, |
| class PoisonAttr = ub::PoisonAttr, |
| class CalculationT = function_ref< |
| std::optional<ElementValueT>(ElementValueT, ElementValueT)>> |
| 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 getResultType = [](Attribute attr) -> Type { |
| if (auto typed = dyn_cast_or_null<TypedAttr>(attr)) |
| return typed.getType(); |
| return {}; |
| }; |
| |
| Type lhsType = getResultType(operands[0]); |
| Type rhsType = getResultType(operands[1]); |
| if (!lhsType || !rhsType) |
| return {}; |
| if (lhsType != rhsType) |
| return {}; |
| |
| return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr, |
| CalculationT>( |
| operands, lhsType, std::forward<CalculationT>(calculate)); |
| } |
| |
| template <class AttrElementT, |
| class ElementValueT = typename AttrElementT::ValueType, |
| class PoisonAttr = void, |
| class CalculationT = |
| function_ref<ElementValueT(ElementValueT, ElementValueT)>> |
| Attribute constFoldBinaryOp(ArrayRef<Attribute> operands, Type resultType, |
| CalculationT &&calculate) { |
| return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>( |
| operands, resultType, |
| [&](ElementValueT a, ElementValueT b) -> std::optional<ElementValueT> { |
| return calculate(a, b); |
| }); |
| } |
| |
| template <class AttrElementT, |
| class ElementValueT = typename AttrElementT::ValueType, |
| class PoisonAttr = ub::PoisonAttr, |
| class CalculationT = |
| function_ref<ElementValueT(ElementValueT, ElementValueT)>> |
| Attribute constFoldBinaryOp(ArrayRef<Attribute> operands, |
| CalculationT &&calculate) { |
| return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>( |
| operands, |
| [&](ElementValueT a, ElementValueT b) -> std::optional<ElementValueT> { |
| return calculate(a, b); |
| }); |
| } |
| |
| /// Performs constant folding `calculate` with element-wise behavior on the one |
| /// attributes in `operands` and returns the result if possible. |
| /// 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 CalculationT = |
| function_ref<std::optional<ElementValueT>(ElementValueT)>> |
| Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands, |
| CalculationT &&calculate) { |
| assert(operands.size() == 1 && "unary op takes one operands"); |
| if (!operands[0]) |
| 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 AttrElementT::get(op.getType(), *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(op.getType(), *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<ElementValueT> 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(op.getShapedType(), elementResults); |
| } |
| return {}; |
| } |
| |
| template <class AttrElementT, |
| class ElementValueT = typename AttrElementT::ValueType, |
| class PoisonAttr = ub::PoisonAttr, |
| class CalculationT = function_ref<ElementValueT(ElementValueT)>> |
| Attribute constFoldUnaryOp(ArrayRef<Attribute> operands, |
| CalculationT &&calculate) { |
| return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>( |
| operands, [&](ElementValueT a) -> std::optional<ElementValueT> { |
| 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) { |
| assert(operands.size() == 1 && "Cast op takes one operand"); |
| if (!operands[0]) |
| 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 {}; |
| return DenseElementsAttr::get(cast<ShapedType>(resType), 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 |
