wehub-resource-sync
370 lines
14 KiB
C++
370 lines
14 KiB
C++
// Copyright (c) 2025 CINN Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#pragma once
|
|
#include <map>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "paddle/cinn/common/integer_set.h"
|
|
#include "paddle/cinn/ir/ir.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
/*!
|
|
* \brief Apply func `fleaf` into each leaf node of `expr`.
|
|
* which leaf node is the most outside node that has TNode type.
|
|
* \param expr The expression to be applied.
|
|
* \param fleaf The function to be applied.
|
|
*/
|
|
template <typename TNode, typename FLeaf>
|
|
inline void UnpackReduction(const ir::IndexExpr &expr, FLeaf fleaf) {
|
|
if (const TNode *node = expr.As<TNode>()) {
|
|
UnpackReduction<TNode, FLeaf>(node->a(), fleaf);
|
|
UnpackReduction<TNode, FLeaf>(node->b(), fleaf);
|
|
} else {
|
|
fleaf(expr);
|
|
}
|
|
}
|
|
|
|
/*!
|
|
* \brief Flatten the expression into a vector of expressions splited by `Add`
|
|
* or `Mul`.
|
|
*
|
|
* For example (Add):
|
|
* 1. `S0 + S1` ==> {S0, S1}
|
|
* 2. `S0 + S1 * S2` ==> {S0, S1 * S2}
|
|
* 3. `S0 + S1 * (S2 + S3)` ==> {S0, S1 * (S2 + S3)}
|
|
*
|
|
* \param lhs The left hand side expression to be compared.
|
|
* \param rhs The right hand side expression to be compared.
|
|
* \return A boolean value indicating whether the priority of `lhs` is higher
|
|
* than `rhs`.
|
|
*/
|
|
template <typename T>
|
|
inline std::vector<ir::IndexExpr> GetFlattenExprs(const ir::IndexExpr &expr) {
|
|
std::vector<ir::IndexExpr> result;
|
|
auto fcollect = [&](ir::IndexExpr val) { result.push_back(val); };
|
|
UnpackReduction<T>(expr, fcollect);
|
|
return result;
|
|
}
|
|
|
|
/*!
|
|
* \brief Compare the priority of the two expressions. This function follows the
|
|
* above rules:
|
|
* 1. if lhs = var, rhs = const, return 1 (lhs > rhs);
|
|
* 2. if lhs = const, rhs = var, return -1 (lhs < rhs);
|
|
* 3. if lhs = var, rhs = var, return comparison result of lhs_var_name and
|
|
* rhs_var_name (0 if equal, -1 if lhs < rhs, 1 if lhs > rhs);
|
|
* 4. if lhs.length > rhs.length, return 1 (lhs > rhs);
|
|
* 5. if lhs.length == rhs.length, return comparison result of lhs_type and
|
|
* rhs_type (Add < Mul < Div < Mod, 0 if equal, -1 if lhs < rhs, 1 if lhs >
|
|
* rhs);
|
|
* 6. if lhs.length < rhs.length return -1 (lhs < rhs);
|
|
*
|
|
* For example:
|
|
* 1. `ComparePriority(S0, 2)` return 1 (lhs > rhs);
|
|
* 2. `ComparePriority(S0, S0)` return 0 (equal);
|
|
* 3. `ComparePriority(S0, S1)` return -1 (lhs < rhs) if S0 < S1;
|
|
* 4. `ComparePriority(S0, S1 + 1)` return -1 (lhs < rhs);
|
|
* 5. `ComparePriority(S0 % 2, S1 + 1)` return -1 (lhs < rhs);
|
|
*
|
|
* \param lhs The left hand side expression to be compared.
|
|
* \param rhs The right hand side expression to be compared.
|
|
* \return An integer value indicating the comparison result:
|
|
* - 1: lhs has strictly higher priority than rhs
|
|
* - 0: lhs and rhs have equal priority
|
|
* - -1: lhs has strictly lower priority than rhs
|
|
*/
|
|
int ComparePriority(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
|
|
|
|
/*!
|
|
* \brief Comparison function for sorting expressions by priority. This function
|
|
* follows the strict weak ordering requirement for std::sort by calling
|
|
* ComparePriority and converting its result to a boolean.
|
|
*
|
|
* This function implements the ordering such that:
|
|
* - If ComparePriority(lhs, rhs) returns 1, returns true (lhs should come
|
|
* before rhs)
|
|
* - If ComparePriority(lhs, rhs) returns 0 or -1, returns false (lhs should not
|
|
* come before rhs)
|
|
*
|
|
* This ensures that expressions are sorted in descending priority order, with
|
|
* higher priority expressions coming first in the sorted sequence.
|
|
*
|
|
* \param lhs The left hand side expression to be compared.
|
|
* \param rhs The right hand side expression to be compared.
|
|
* \return A boolean value indicating whether lhs should come before rhs in the
|
|
* sorted sequence according to the priority rules.
|
|
*/
|
|
bool SortComparePriority(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
|
|
|
|
/*!
|
|
* \brief Determines whether there are sub-parts in the `expr` that can be
|
|
* simplified by `Add` operation with the input `symbol`. If true is returned,
|
|
* the operation will be attempted on each subpart in outer
|
|
* `SimplifySymbolicAdd` function.
|
|
*
|
|
* For example:
|
|
* 1. `IsSumPartialBySymbol(5, S0)` return false;
|
|
* 2. `IsSumPartialBySymbol(S0, S0)` return true;
|
|
* 3. `IsSumPartialBySymbol(S0 + S1, S1)` return true;
|
|
* 4. `IsSumPartialBySymbol(S0 * 5 + S1, S0)` return true;
|
|
* 5. `IsSumPartialBySymbol(S0 / 3, S0)` return true;
|
|
* 6. `IsSumPartialBySymbol(S0 / 3 + S1, S0)` return true;
|
|
* 7. `IsSumPartialBySymbol(S0 % 3, S0)` return false;
|
|
*
|
|
* Note: For performance reasons, special patterns will not be matched here.
|
|
* This can be allowed for extreme optimization.
|
|
* For example:
|
|
* `IsSumPartialBySymbol((S0 + S1 / 5 * 25) / 5, S1 % 5)` return false;
|
|
*
|
|
* \param expr The expression to be checked.
|
|
* \param symbol The symbol to be checked.
|
|
* \return True means there are sub-parts in the `expr` that can be simplified.
|
|
*/
|
|
bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
|
|
const ir::IndexExpr &symbol);
|
|
|
|
/*!
|
|
* \brief Simplify the `lhs` by symbol `sym`. Usually run after
|
|
* `IsSumPartialBySymbol`
|
|
*
|
|
* \param lhs The expression to be simplified.
|
|
* \param sym The symbol to be checked.
|
|
* it may be `i, j ..` or `S0, S1 ..` or other symbolic expr.
|
|
* \param outer_mul_factor The scale of symbolic expr.
|
|
* e.g. `S0 * 4` ===> sym == S0, outer_mul_factor == 4
|
|
* \return The expr after simplification.
|
|
*/
|
|
ir::IndexExpr SimplifySymbolicAdd(
|
|
const ir::IndexExpr &lhs,
|
|
const ir::IndexExpr &sym,
|
|
const ir::IndexExpr &outer_mul_factor = ir::IndexExpr(1));
|
|
|
|
/*!
|
|
* \brief Determines whether there are sub-parts in the `expr` that can be
|
|
* simplified by `Div` operation with the input `symbol`. If true is returned,
|
|
* the operation will be attempted on each subpart in outer
|
|
* `SimplifySymbolicDivide` function.
|
|
*
|
|
* For example:
|
|
* 1. `IsDivisibleBySymbol(5, S0, div)` return false;
|
|
* 2. `IsDivisibleBySymbol(S0, S0, div)` return true;
|
|
* 3. `IsDivisibleBySymbol(S0 + S1, S1, div)` return false;
|
|
* 4. `IsDivisibleBySymbol(S0 * 5 + S1 * S2, S0, div)` return true;
|
|
* 5. `IsDivisibleBySymbol(S0 / 3, S0, div)` return true;
|
|
* 6. `IsDivisibleBySymbol(S0 * 4 / 3, S0, div)` return true;
|
|
* 7. `IsDivisibleBySymbol(S0 % 3, S0, div)` return false;
|
|
* 8. `IsDivisibleBySymbol(S0 / 3, S0, mod)` return false;
|
|
*
|
|
* \param expr The expression to be checked.
|
|
* \param symbol The symbol to be checked.
|
|
* \param ty ty is `Mod` or `Div`.
|
|
* \return True means there are sub-parts in the `expr` that can be simplified.
|
|
* \note this func dont deal the corner case, please use `ProveDivisible` for
|
|
* exact result. e.g. `IsDivisibleBySymbol(f % S0 - f, S0, div)` is false
|
|
*/
|
|
bool IsDivisibleBySymbol(const ir::IndexExpr &expr,
|
|
const ir::IndexExpr &symbol,
|
|
const ir::IrNodeTy &ty);
|
|
|
|
/*!
|
|
* \brief Simplify the `lhs` by symbol `sym`. Usually run after
|
|
* `IsDivisibleBySymbol`
|
|
*
|
|
* \param lhs The expression to be simplified.
|
|
* \param sym The symbol to be checked.
|
|
* it may be `i, j ..` or `S0, S1 ..` or other symbolic expr.
|
|
* \param ty ty is `Mod` or `Div`.
|
|
* \return The expr after simplification.
|
|
*/
|
|
ir::IndexExpr SimplifySymbolicDivide(const ir::IndexExpr &lhs,
|
|
const ir::IndexExpr &sym,
|
|
const ir::IrNodeTy &ty);
|
|
|
|
/*!
|
|
* \brief Determine whether `lhs` is divisible by `rhs`, regardless of whether
|
|
* `rhs` is a constant or a symbol.
|
|
* \param lhs lhs is dividend.
|
|
* \param rhs rhs is divisor.
|
|
* \return A boolean value indicating whether the `lhs` is divisible by `rhs`
|
|
*/
|
|
bool ProveDivisible(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
|
|
|
|
/*!
|
|
* \brief Judge whether `candidate` is a negated index expression.
|
|
* \param candidate The expression to be checked.
|
|
* \param expr The positive part
|
|
* \return A boolean value indicating whether `candidate` is negative.
|
|
*/
|
|
bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
|
|
ir::IndexExpr &expr); // NOLINT
|
|
|
|
/*!
|
|
* \brief Construct index expression by node type with or without simplify.
|
|
* \param ty The node type of index expression.
|
|
* \param lhs left operand.
|
|
* \param rhs right operand.
|
|
* \param simplify_flag Whether to simplify the result.
|
|
* \return The constructed index expression.
|
|
*/
|
|
ir::IndexExpr ConstructIndexExprByNodeType(const ir::IrNodeTy &ty,
|
|
const ir::IndexExpr &lhs,
|
|
const ir::IndexExpr &rhs,
|
|
bool simplify_flag = true);
|
|
|
|
/*!
|
|
* \brief Change the sequence of `Div` and `Mod` in index expression.
|
|
* Mathematical formula: `(a / b) % c = (a % (b * c)) / b`
|
|
* For example:
|
|
* 1. i / 4 % 8 => i % 32 / 4
|
|
* 2. i / S0 % S1 => i % (S0 * S1) / S0
|
|
* 3. (i * 32 + j) / 4 % 8 => (i * 32 + j) % 32 / 4
|
|
*
|
|
* \param expr The `IndexExpr` to be change
|
|
* \return `IndexExpr` after change.
|
|
*/
|
|
ir::IndexExpr ChangeSeqOfDivMod(const ir::IndexExpr &expr);
|
|
|
|
/*!
|
|
* \brief Judge type of `expr` is valid type of `IndexExpr` or not.
|
|
* \param expr The expression to be checked.
|
|
* \return A enum IndexType value indicating whether the type of `expr` is valid
|
|
* IndexExpr.
|
|
*
|
|
* Note: Although load and cast are also legal IndexExpr, we need to know this
|
|
* information in some scenarios.
|
|
*/
|
|
ir::IndexExpr::IndexType VerifyIndex(const ir::Expr &expr);
|
|
|
|
/*!
|
|
* \brief The multiplication in rhs is broken down and each sub-part is
|
|
* independently determined to be divisible.
|
|
* \param lhs The dividend.
|
|
* \param rhs The divisor.
|
|
* \param ty ty is `Mod` or `Div`.
|
|
* \return A optional index expression indicating whether the `lhs`
|
|
* is divisible, nullopt indicating not divisible.
|
|
*
|
|
* For example:
|
|
* 1. i * S0 * S1 * S2 / (S0 * S1) ==> i / S2
|
|
* 2. i * S0 * S1 / S0 ==> i * S1
|
|
* 3. i * S0 / (S0 + 1) ==> nullopt
|
|
*/
|
|
std::optional<ir::IndexExpr> DivByPartMul(const ir::IndexExpr &lhs,
|
|
const ir::IndexExpr &rhs,
|
|
ir::IrNodeTy ty);
|
|
|
|
/*!
|
|
* \brief Simplify complex modulo expressions.
|
|
* \param lhs The dividend.
|
|
* \param rhs The divisor.
|
|
* \return A optional index expression indicating whether simplified
|
|
*
|
|
* For example:
|
|
* 1. (i / S0 * S0 + i % (S0 * S1)) % S0 ==> i % S0
|
|
* 2. (i / S0 * S0 * S1 + i % (S0 * S1 * S2)) % (S0 * S1) ==> i % (S0 * S1)
|
|
* 3. i % (S0 * S1) % S0 ==> i % S0
|
|
* 4. i * S0 * S1 % (S0 * S1) ==> 0
|
|
*/
|
|
std::optional<ir::IndexExpr> SimplifyComplexMod(const ir::IndexExpr &lhs,
|
|
const ir::IndexExpr &rhs);
|
|
|
|
/*!
|
|
* \brief Check whether the expression matches the pattern.
|
|
* \param expr The expression to be checked.
|
|
* \param pattern The pattern to be matched. which includes some variables.
|
|
* \param map return the matched variables.
|
|
* \return A boolean value indicating whether `expr` is matched.
|
|
*
|
|
* For example:
|
|
* 1. (i / S0 * S0 + i % (S0 * S1)) % S0 matched by a / b * b + a % (b * c)
|
|
* with map = {a: i, b: S0, c: S1}
|
|
* 2. S0 + 5 matched by a + 5 with map = {a: S0, b: 5}
|
|
*
|
|
* Note: a * b and b * a is two different pattern.
|
|
*/
|
|
bool CheckPattern(const ir::IndexExpr &expr,
|
|
const ir::IndexExpr &pattern,
|
|
std::unordered_map<std::string, ir::IndexExpr> *map);
|
|
|
|
// TODO(liujinnan): Delete historical `simplify func` related files, temporary
|
|
// placement of tool functions that are still in use, remove it in the future.
|
|
bool IsPureMath(Expr expr);
|
|
|
|
/*!
|
|
* \brief Parse the expression from string to Expr.
|
|
* \param expr_str The expression to be checked.
|
|
* \return A Expr parsed from string.
|
|
*
|
|
* For example:
|
|
* 1. ParseExpressionFromString("a + b * c") return Add(Var(a), Mul(Var(b),
|
|
* Var(c)))
|
|
* 2. ParseExpressionFromString("a + 10") return Add(Var(a), IntImm(10)))
|
|
*/
|
|
ir::Expr ParseExpressionFromString(const std::string &expr_str);
|
|
|
|
/*!
|
|
* \brief Check whether the expression matches the pattern.
|
|
* \param expr The expression to be checked.
|
|
* \param pattern_str The pattern string to be matched.
|
|
* \param condition A optional condition function to further check the matched
|
|
* \return A optional map indicating the matched variables.
|
|
*
|
|
* For example:
|
|
* 1. i / S0 * S0 + i % (S0 * S1) matched by a / b * b + a % (b * c)
|
|
* return map = {a: i, b: S0, c: S1}
|
|
* 2. i / (S0 * S1) * S0 + i % (S0 * S1) / S1 matched by a / f * b + a % f / c
|
|
* with optional condition f = a * b, return map = {a: i, b: S0, c: S1, f: S0 *
|
|
* S1}
|
|
* 3. S0 + 5 matched by a + 5 return map = {a: S0, b: 5}
|
|
*/
|
|
std::optional<std::unordered_map<std::string, ir::IndexExpr>> MatchPattern(
|
|
const ir::IndexExpr &expr,
|
|
const std::string &pattern_str,
|
|
const std::function<bool(
|
|
const std::unordered_map<std::string, ir::IndexExpr> &)> &condition =
|
|
nullptr);
|
|
|
|
/*!
|
|
* \brief Simplify IndexExpr with bound information.
|
|
* For example:
|
|
* x % S0 ==> x if x < S0
|
|
* x / S0 ==> 0 if x < S0
|
|
*
|
|
* \param expr The `IndexExpr` to be simplified.
|
|
* \return `IndexExpr` after simplification.
|
|
*/
|
|
ir::IndexExpr BoundSimplify(const ir::IndexExpr &expr);
|
|
|
|
/*!
|
|
* \brief Simplify IndexExpr with broadcastable information.
|
|
* For example:
|
|
* x % cinn_max(cinn_max(cinn_max(S0, S10), S20), S30))
|
|
* % cinn_max(S10, S30) ==> x % cinn_max(S10, S30),
|
|
* if broadcastable(S0, S10, S20, S30).
|
|
* Note: The following conditions must be met:
|
|
* 1. Two consecutive modular operations.
|
|
* 2. The first modulus is broadcastable.
|
|
* 3. The second modulus is a subset of the first modulus.
|
|
*
|
|
* \param expr The `IndexExpr` to be simplified.
|
|
* \return `IndexExpr` after simplification.
|
|
*/
|
|
ir::IndexExpr BroadcastSimplify(const ir::IndexExpr &expr);
|
|
} // namespace optim
|
|
} // namespace cinn
|