wehub-resource-sync
603 lines
24 KiB
C++
603 lines
24 KiB
C++
// Copyright (c) 2023 PaddlePaddle 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.
|
|
|
|
#include <optional>
|
|
#include <typeinfo>
|
|
|
|
#include "paddle/cinn/adt/equation.h"
|
|
#include "paddle/cinn/adt/equation_value_match_trait.h"
|
|
#include "paddle/cinn/adt/get_sub_reshape_dim_ranges.h"
|
|
#include "paddle/cinn/adt/index_expr_infer_context.h"
|
|
#include "paddle/cinn/adt/match.h"
|
|
#include "paddle/cinn/adt/simplify_value.h"
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/dim_expr_util.h"
|
|
namespace cinn::adt {
|
|
|
|
namespace {
|
|
|
|
template <typename T, typename ExprT>
|
|
ExprT MatchAndRewrite(const ExprT& expr, const IndexExprInferContext& ctx) {
|
|
if (cinn::adt::Match<typename T::source_pattern_type>(expr)) {
|
|
return T().MatchAndRewrite(expr, ctx);
|
|
} else {
|
|
return expr;
|
|
}
|
|
}
|
|
|
|
struct SimplifyBroadcastedIterator {
|
|
using source_pattern_type = BroadcastedIterator<Value, std::int64_t>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [iterator, dim] =
|
|
value.Get<BroadcastedIterator<Value, DimExpr>>().tuple();
|
|
if (dim.Get<std::int64_t>() == 1) {
|
|
return DimExpr{std::int64_t(0)};
|
|
} else {
|
|
return iterator;
|
|
}
|
|
}
|
|
};
|
|
|
|
struct SimplifyRedundantBroadcastedIterator {
|
|
using source_pattern_type =
|
|
BroadcastedIterator<BroadcastedIterator<Value, DimExpr>, DimExpr>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [outer_iterator, outer_dim] =
|
|
value.Get<BroadcastedIterator<Value, DimExpr>>().tuple();
|
|
const auto& [inner_iterator, inner_dim] =
|
|
outer_iterator.Get<BroadcastedIterator<Value, DimExpr>>().tuple();
|
|
|
|
if (outer_dim == inner_dim) {
|
|
return SimplifyValue(outer_iterator, ctx);
|
|
} else {
|
|
const auto& bd = MakeBroadcastedDim(outer_dim, inner_dim);
|
|
const auto& simplified_bd = DimExpr{symbol::SimplifyDimExpr(bd)};
|
|
return BroadcastedIterator<Value, DimExpr>{inner_iterator, simplified_bd};
|
|
}
|
|
PADDLE_THROW(::common::errors::Fatal("Dead code"));
|
|
}
|
|
};
|
|
|
|
struct SimplifyList {
|
|
using source_pattern_type = List<Value>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
List<Value> ret{};
|
|
for (const auto& v : *value.Get<List<Value>>()) {
|
|
ret->emplace_back(SimplifyValue(v, ctx));
|
|
}
|
|
return ret;
|
|
}
|
|
};
|
|
|
|
struct SimplifyDotUndot {
|
|
using source_pattern_type =
|
|
IndexDotValue<List<ListGetItem<IndexUnDotValue<Value, List<std::int64_t>>,
|
|
std::int64_t>>,
|
|
List<std::int64_t>>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [list_get_item_values, dot_dims] =
|
|
value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& list_get_items = list_get_item_values.Get<List<Value>>();
|
|
std::optional<Value> pre_index_undot{std::nullopt};
|
|
for (std::size_t i = 0; i < list_get_items->size(); ++i) {
|
|
const auto& [index_undot_value, constant_idx] =
|
|
list_get_items.Get(i).Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
if (constant_idx.Get<std::int64_t>() != i) {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
}
|
|
if (pre_index_undot.has_value()) {
|
|
if (!(pre_index_undot.value() == index_undot_value)) {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
} else {
|
|
// do nothing
|
|
}
|
|
} else {
|
|
pre_index_undot = index_undot_value;
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(pre_index_undot.has_value(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"pre_index_undot should not be null"));
|
|
const auto& [index_value, undot_dims] =
|
|
pre_index_undot.value()
|
|
.Get<IndexUnDotValue<Value, List<DimExpr>>>()
|
|
.tuple();
|
|
if (dot_dims == undot_dims) {
|
|
return index_value;
|
|
}
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
}
|
|
};
|
|
|
|
struct SimplifyUndotDot {
|
|
using source_pattern_type = ListGetItem<
|
|
IndexUnDotValue<IndexDotValue<List<Value>, List<std::int64_t>>,
|
|
List<std::int64_t>>,
|
|
std::int64_t>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [index_undot_value, constant_idx] =
|
|
value.Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
const auto& [index_value, undot_dims] =
|
|
index_undot_value.Get<IndexUnDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& [index_dot_values, dot_dims] =
|
|
index_value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& iter_values = index_dot_values.Get<List<Value>>();
|
|
if (dot_dims == undot_dims) {
|
|
return iter_values.Get(constant_idx.Get<std::int64_t>());
|
|
} else {
|
|
return ListGetItem<Value, DimExpr>{SimplifyValue(index_undot_value, ctx),
|
|
constant_idx};
|
|
}
|
|
}
|
|
};
|
|
|
|
struct SimplifyListGetItem {
|
|
using source_pattern_type = ListGetItem<Value, DimExpr>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [list_values, constant_idx] =
|
|
value.Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
return ListGetItem<Value, DimExpr>{SimplifyValue(list_values, ctx),
|
|
constant_idx};
|
|
}
|
|
};
|
|
|
|
struct SimplifyListGetItemList {
|
|
using source_pattern_type = ListGetItem<List<Value>, std::int64_t>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [list_values, constant_idx] =
|
|
value.Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
const auto& iter_values = list_values.Get<List<Value>>();
|
|
return iter_values.Get(constant_idx.Get<std::int64_t>());
|
|
}
|
|
};
|
|
|
|
struct SimplifyGcdShape {
|
|
using source_pattern_type = ListGetItem<
|
|
IndexUnDotValue<IndexDotValue<List<Value>, List<std::int64_t>>,
|
|
List<std::int64_t>>,
|
|
std::int64_t>;
|
|
|
|
bool IsConstantListAllPositiveInt64(const List<DimExpr>& constants) {
|
|
for (const auto& constant : *constants) {
|
|
if (!constant.Has<std::int64_t>() || constant.Get<std::int64_t>() <= 0) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [index_undot_value, constant_idx] =
|
|
value.Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
const auto& [index_value, undot_dims] =
|
|
index_undot_value.Get<IndexUnDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& [index_dot_values, dot_dims] =
|
|
index_value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& iter_values = index_dot_values.Get<List<Value>>();
|
|
const auto& undot_dim_values = undot_dims;
|
|
const auto& dot_dim_values = dot_dims;
|
|
PADDLE_ENFORCE_EQ(IsConstantListAllPositiveInt64(undot_dim_values),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"The undot_dim_values should be all positive int64"));
|
|
PADDLE_ENFORCE_EQ(IsConstantListAllPositiveInt64(dot_dim_values),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"The dot_dim_values should be all positive int64"));
|
|
const auto& sub_reshape_dim_ranges =
|
|
GetSubReshapeDimRanges(undot_dim_values, dot_dim_values);
|
|
if (!sub_reshape_dim_ranges.has_value()) {
|
|
return ListGetItem<Value, DimExpr>{SimplifyValue(index_undot_value, ctx),
|
|
constant_idx};
|
|
}
|
|
const auto& [undot_dim_ranges, dot_dim_ranges] =
|
|
sub_reshape_dim_ranges.value();
|
|
if (undot_dim_ranges.size() >= 1) {
|
|
const auto& [sub_range_idx, sub_range_item_idx] = GetSubRangeItemIdx(
|
|
undot_dim_ranges, constant_idx.Get<std::int64_t>());
|
|
List<DimExpr> sub_range_undot_dims = GetSubRangeDotDims(
|
|
undot_dim_values, undot_dim_ranges.at(sub_range_idx));
|
|
List<Value> sub_range_dot_iterators = GetSubRangeDotIterators(
|
|
iter_values, dot_dim_ranges.at(sub_range_idx));
|
|
List<DimExpr> sub_range_dot_dims =
|
|
GetSubRangeDotDims(dot_dim_values, dot_dim_ranges.at(sub_range_idx));
|
|
if (sub_range_dot_dims == sub_range_undot_dims) {
|
|
return sub_range_dot_iterators.Get(sub_range_item_idx);
|
|
} else {
|
|
IndexDotValue<Value, List<DimExpr>> sub_range_dot{
|
|
sub_range_dot_iterators, sub_range_dot_dims};
|
|
if (sub_range_undot_dims->size() == 1) {
|
|
PADDLE_ENFORCE_EQ(
|
|
sub_range_item_idx,
|
|
0UL,
|
|
::common::errors::InvalidArgument(
|
|
"The sub_range_item_idx should be 0, but got %d.",
|
|
sub_range_item_idx));
|
|
return sub_range_dot;
|
|
} else {
|
|
IndexUnDotValue<Value, List<DimExpr>> sub_range_undot{
|
|
sub_range_dot, sub_range_undot_dims};
|
|
return ListGetItem<Value, DimExpr>{sub_range_undot,
|
|
sub_range_item_idx};
|
|
}
|
|
}
|
|
}
|
|
return ListGetItem<Value, DimExpr>{SimplifyValue(index_undot_value, ctx),
|
|
constant_idx};
|
|
}
|
|
|
|
std::pair<int, int> GetSubRangeItemIdx(
|
|
const std::vector<std::pair<int, int>>& ranges,
|
|
std::int64_t index) const {
|
|
for (std::size_t i = 0; i < ranges.size(); ++i) {
|
|
const auto& [begin, end] = ranges.at(i);
|
|
if (index >= begin && index < end) {
|
|
return std::pair<int, int>{i, index - begin};
|
|
}
|
|
}
|
|
}
|
|
|
|
List<Value> GetSubRangeDotIterators(const List<Value>& iterators,
|
|
const std::pair<int, int>& range) const {
|
|
return GetSubRange<List<Value>>(iterators, range);
|
|
}
|
|
|
|
List<DimExpr> GetSubRangeDotDims(const List<DimExpr>& dims,
|
|
const std::pair<int, int>& range) const {
|
|
return GetSubRange<List<DimExpr>>(dims, range);
|
|
}
|
|
|
|
template <typename ContainerT>
|
|
ContainerT GetSubRange(const ContainerT& container,
|
|
const std::pair<int, int>& range) const {
|
|
CheckRange(container, range);
|
|
ContainerT ret{};
|
|
ret->assign(std::next(container->begin(), range.first),
|
|
std::next(container->begin(), range.second));
|
|
return ret;
|
|
}
|
|
|
|
template <typename ContainerT>
|
|
void CheckRange(const ContainerT& container,
|
|
const std::pair<int, int>& range) const {
|
|
PADDLE_ENFORCE_GE(
|
|
range.first,
|
|
0UL,
|
|
::common::errors::InvalidArgument(
|
|
"The range.first should be greater than or equal to 0, "
|
|
"but got %d.",
|
|
range.first));
|
|
PADDLE_ENFORCE_GE(
|
|
range.second,
|
|
0UL,
|
|
::common::errors::InvalidArgument(
|
|
"The range.second should be greater than or equal to 0, "
|
|
"but got %d.",
|
|
range.second));
|
|
PADDLE_ENFORCE_LE(range.first,
|
|
container->size(),
|
|
::common::errors::InvalidArgument(
|
|
"The range.first should be less than or equal to the "
|
|
"size of the container, but got range.first = %d, "
|
|
"container size = %d.",
|
|
range.first,
|
|
container->size()));
|
|
PADDLE_ENFORCE_LE(
|
|
range.second,
|
|
container->size(),
|
|
::common::errors::InvalidArgument(
|
|
"The range.second should be less than or equal to the "
|
|
"size of the container, but got range.second = %d, "
|
|
"container size = %d.",
|
|
range.second,
|
|
container->size()));
|
|
PADDLE_ENFORCE_LT(range.first,
|
|
range.second,
|
|
::common::errors::InvalidArgument(
|
|
"The range.first should be less than range.second, "
|
|
"but got range.first = %d, range.second = %d.",
|
|
range.first,
|
|
range.second));
|
|
}
|
|
};
|
|
|
|
struct SimplifyDotDot {
|
|
using source_pattern_type = IndexDotValue<List<Value>, List<std::int64_t>>;
|
|
|
|
std::int64_t Product(const List<DimExpr>& dims) {
|
|
std::int64_t ret = 1;
|
|
for (const auto& dim : *dims) {
|
|
PADDLE_ENFORCE_EQ(
|
|
dim.Has<std::int64_t>(),
|
|
true,
|
|
::common::errors::InvalidArgument("dim should have std::int64_t"));
|
|
ret *= dim.Get<std::int64_t>();
|
|
}
|
|
return ret;
|
|
}
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [index_dot_values, dot_dims] =
|
|
value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
PADDLE_ENFORCE_EQ(
|
|
index_dot_values.Get<List<Value>>()->size(),
|
|
dot_dims->size(),
|
|
::common::errors::InvalidArgument(
|
|
"The size of index_dot_values and dot_dims should be equal, "
|
|
"but got index_dot_values size = %d, dot_dims size = %d.",
|
|
index_dot_values.Get<List<Value>>()->size(),
|
|
dot_dims->size()));
|
|
List<Value> new_dot_values{};
|
|
List<DimExpr> new_dot_dims{};
|
|
for (std::size_t i = 0; i < index_dot_values.Get<List<Value>>()->size();
|
|
++i) {
|
|
const auto& index_dot_value = index_dot_values.Get<List<Value>>()->at(i);
|
|
const auto& dot_dim = dot_dims->at(i).Get<std::int64_t>();
|
|
if (Match<source_pattern_type>(index_dot_value)) {
|
|
const auto& [sub_index_dot_values, sub_dot_dims] =
|
|
index_dot_value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& sub_dot_dim_values = sub_dot_dims;
|
|
std::int64_t dim_product = Product(sub_dot_dim_values);
|
|
if (dim_product == dot_dim) {
|
|
for (std::size_t j = 0;
|
|
j < sub_index_dot_values.Get<List<Value>>()->size();
|
|
++j) {
|
|
const auto& sub_index_dot_value =
|
|
sub_index_dot_values.Get<List<Value>>()->at(j);
|
|
const auto& sub_dot_dim = sub_dot_dim_values->at(j);
|
|
new_dot_values->emplace_back(sub_index_dot_value);
|
|
new_dot_dims->emplace_back(sub_dot_dim);
|
|
}
|
|
} else {
|
|
new_dot_values->emplace_back(index_dot_value);
|
|
new_dot_dims->emplace_back(dot_dim);
|
|
}
|
|
} else {
|
|
new_dot_values->emplace_back(index_dot_value);
|
|
new_dot_dims->emplace_back(dot_dim);
|
|
}
|
|
}
|
|
return IndexDotValue<Value, List<DimExpr>>{new_dot_values, new_dot_dims};
|
|
}
|
|
};
|
|
|
|
struct SymbolicDim_SimplifyDotUndot {
|
|
using source_pattern_type = IndexDotValue<
|
|
List<ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>,
|
|
List<DimExpr>>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [list_get_item_values, dot_dims] =
|
|
value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& list_get_items = list_get_item_values.Get<List<Value>>();
|
|
std::optional<Value> pre_index_undot{std::nullopt};
|
|
for (std::size_t i = 0; i < list_get_items->size(); ++i) {
|
|
const auto& [index_undot_value, constant_idx] =
|
|
list_get_items.Get(i).Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
if (constant_idx.Get<std::int64_t>() != i) {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
}
|
|
if (pre_index_undot.has_value()) {
|
|
if (!(pre_index_undot.value() == index_undot_value)) {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
} else {
|
|
// do nothing
|
|
}
|
|
} else {
|
|
pre_index_undot = index_undot_value;
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(pre_index_undot.has_value(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"pre_index_undot should not be null"));
|
|
const auto& [index_value, undot_dims] =
|
|
pre_index_undot.value()
|
|
.Get<IndexUnDotValue<Value, List<DimExpr>>>()
|
|
.tuple();
|
|
const auto& dot_dim_values = dot_dims;
|
|
const auto& undot_dim_values = undot_dims;
|
|
if (ctx.DimsEqual(dot_dim_values, undot_dim_values)) {
|
|
return index_value;
|
|
} else {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
}
|
|
PADDLE_THROW(::common::errors::Fatal("Dead code"));
|
|
}
|
|
};
|
|
|
|
struct SymbolicDim_SimplifyDotUndot_DimExpr {
|
|
using source_pattern_type = IndexDotValue<
|
|
List<ListGetItem<
|
|
IndexUnDotValue<Value, List<Union<DimExpr, std::int64_t>>>,
|
|
std::int64_t>>,
|
|
List<Union<DimExpr, std::int64_t>>>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [list_get_item_values, dot_dims] =
|
|
value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& list_get_items = list_get_item_values.Get<List<Value>>();
|
|
std::optional<Value> pre_index_undot{std::nullopt};
|
|
for (std::size_t i = 0; i < list_get_items->size(); ++i) {
|
|
const auto& [index_undot_value, constant_idx] =
|
|
list_get_items.Get(i).Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
if (constant_idx.Get<std::int64_t>() != i) {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
}
|
|
if (pre_index_undot.has_value()) {
|
|
if (!(pre_index_undot.value() == index_undot_value)) {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
} else {
|
|
// do nothing
|
|
}
|
|
} else {
|
|
pre_index_undot = index_undot_value;
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(pre_index_undot.has_value(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"pre_index_undot should not be null"));
|
|
const auto& [index_value, undot_dims] =
|
|
pre_index_undot.value()
|
|
.Get<IndexUnDotValue<Value, List<DimExpr>>>()
|
|
.tuple();
|
|
const auto& dot_dim_values = dot_dims;
|
|
const auto& undot_dim_values = undot_dims;
|
|
if (dot_dim_values == undot_dim_values) {
|
|
return index_value;
|
|
} else {
|
|
return IndexDotValue<Value, List<DimExpr>>{
|
|
SimplifyValue(list_get_item_values, ctx), dot_dims};
|
|
}
|
|
PADDLE_THROW(::common::errors::Fatal("Dead code"));
|
|
}
|
|
};
|
|
|
|
struct SymbolicDim_SimplifyUndotDot {
|
|
using source_pattern_type = ListGetItem<
|
|
IndexUnDotValue<IndexDotValue<List<Value>, List<DimExpr>>, List<DimExpr>>,
|
|
std::int64_t>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [index_undot_value, constant_idx] =
|
|
value.Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
const auto& [index_value, undot_dims] =
|
|
index_undot_value.Get<IndexUnDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& [index_dot_values, dot_dims] =
|
|
index_value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& iter_values = index_dot_values.Get<List<Value>>();
|
|
if (ctx.DimsEqual(dot_dims, undot_dims)) {
|
|
return iter_values.Get(constant_idx.Get<std::int64_t>());
|
|
} else {
|
|
return ListGetItem<Value, DimExpr>{SimplifyValue(index_undot_value, ctx),
|
|
constant_idx};
|
|
}
|
|
}
|
|
};
|
|
|
|
struct SymbolicDim_SimplifyUndotDot_DimExpr {
|
|
using source_pattern_type = ListGetItem<
|
|
IndexUnDotValue<
|
|
IndexDotValue<List<Value>, List<Union<DimExpr, std::int64_t>>>,
|
|
List<Union<DimExpr, std::int64_t>>>,
|
|
std::int64_t>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [index_undot_value, constant_idx] =
|
|
value.Get<ListGetItem<Value, DimExpr>>().tuple();
|
|
const auto& [index_value, undot_dims] =
|
|
index_undot_value.Get<IndexUnDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& [index_dot_values, dot_dims] =
|
|
index_value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& iter_values = index_dot_values.Get<List<Value>>();
|
|
if (dot_dims == undot_dims) {
|
|
return iter_values.Get(constant_idx.Get<std::int64_t>());
|
|
} else {
|
|
return ListGetItem<Value, DimExpr>{SimplifyValue(index_undot_value, ctx),
|
|
constant_idx};
|
|
}
|
|
}
|
|
};
|
|
|
|
struct SymbolicDim_SimplifyDotDot {
|
|
using source_pattern_type = IndexDotValue<List<Value>, List<DimExpr>>;
|
|
|
|
Value MatchAndRewrite(const Value& value, const IndexExprInferContext& ctx) {
|
|
const auto& [index_dot_values, dot_dims] =
|
|
value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
PADDLE_ENFORCE_EQ(
|
|
index_dot_values.Get<List<Value>>()->size(),
|
|
dot_dims->size(),
|
|
::common::errors::InvalidArgument(
|
|
"The size of index_dot_values and dot_dims should be equal, "
|
|
"but got index_dot_values size = %d, dot_dims size = %d.",
|
|
index_dot_values.Get<List<Value>>()->size(),
|
|
dot_dims->size()));
|
|
List<Value> new_dot_values{};
|
|
List<DimExpr> new_dot_dims{};
|
|
for (std::size_t i = 0; i < index_dot_values.Get<List<Value>>()->size();
|
|
++i) {
|
|
const auto& index_dot_value = index_dot_values.Get<List<Value>>()->at(i);
|
|
DimExpr dot_dim = dot_dims->at(i);
|
|
if (Match<source_pattern_type>(index_dot_value)) {
|
|
const auto& [sub_index_dot_values, sub_dot_dims] =
|
|
index_dot_value.Get<IndexDotValue<Value, List<DimExpr>>>().tuple();
|
|
const auto& sub_dot_dim_values = sub_dot_dims;
|
|
if (ctx.ProductEqual(sub_dot_dim_values, dot_dim)) {
|
|
for (std::size_t j = 0;
|
|
j < sub_index_dot_values.Get<List<Value>>()->size();
|
|
++j) {
|
|
const auto& sub_index_dot_value =
|
|
sub_index_dot_values.Get<List<Value>>()->at(j);
|
|
const auto& sub_dot_dim = sub_dot_dim_values->at(j);
|
|
new_dot_values->emplace_back(sub_index_dot_value);
|
|
new_dot_dims->emplace_back(sub_dot_dim);
|
|
}
|
|
} else {
|
|
new_dot_values->emplace_back(index_dot_value);
|
|
new_dot_dims->emplace_back(dot_dim);
|
|
}
|
|
} else {
|
|
new_dot_values->emplace_back(index_dot_value);
|
|
new_dot_dims->emplace_back(dot_dim);
|
|
}
|
|
}
|
|
return IndexDotValue<Value, List<DimExpr>>{new_dot_values, new_dot_dims};
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
// Only simplify top-layer of value
|
|
Value SimplifyValue(Value value, const IndexExprInferContext& ctx) {
|
|
value = MatchAndRewrite<SimplifyList>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyListGetItem>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyBroadcastedIterator>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyRedundantBroadcastedIterator>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyDotUndot>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyUndotDot>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyListGetItemList>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyGcdShape>(value, ctx);
|
|
value = MatchAndRewrite<SimplifyDotDot>(value, ctx);
|
|
// For symbolic dim simplification
|
|
value = MatchAndRewrite<SymbolicDim_SimplifyDotUndot>(value, ctx);
|
|
value = MatchAndRewrite<SymbolicDim_SimplifyUndotDot>(value, ctx);
|
|
// value = MatchAndRewrite<SymbolicDim_SimplifyGcdShape>(value, ctx);
|
|
value = MatchAndRewrite<SymbolicDim_SimplifyDotDot>(value, ctx);
|
|
// For DimExpr
|
|
value = MatchAndRewrite<SymbolicDim_SimplifyDotUndot_DimExpr>(value, ctx);
|
|
value = MatchAndRewrite<SymbolicDim_SimplifyUndotDot_DimExpr>(value, ctx);
|
|
|
|
return value;
|
|
}
|
|
|
|
} // namespace cinn::adt
|