wehub-resource-sync
136 lines
4.5 KiB
C++
136 lines
4.5 KiB
C++
// Copyright (c) 2023 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 <atomic>
|
|
#include <functional>
|
|
#include <string>
|
|
#include <type_traits>
|
|
|
|
#include "paddle/cinn/adt/adt.h"
|
|
#include "paddle/cinn/adt/dim_expr.h"
|
|
#include "paddle/cinn/adt/equation_variable.h"
|
|
#include "paddle/cinn/adt/tags.h"
|
|
#include "paddle/cinn/common/equation_graph_topo_walker.h"
|
|
|
|
namespace cinn::adt {
|
|
|
|
template <typename OutT, typename InT>
|
|
struct Identity;
|
|
|
|
// Identity (tOut Iterator) (tIn Iterator)
|
|
template <>
|
|
struct Identity<tOut<Iterator>, tIn<Iterator>>
|
|
: public Tuple<tOut<Iterator>, tIn<Iterator>> {
|
|
using Tuple<tOut<Iterator>, tIn<Iterator>>::Tuple;
|
|
};
|
|
|
|
// Identity (tOut Index) (tIn Index)
|
|
template <>
|
|
struct Identity<tOut<Index>, tIn<Index>>
|
|
: public Tuple<tOut<Index>, tIn<Index>> {
|
|
using Tuple<tOut<Index>, tIn<Index>>::Tuple;
|
|
};
|
|
|
|
template <typename DimT, typename OutT, typename InT>
|
|
struct IndexDot;
|
|
|
|
// IndexDot [DimExpr] (tOut Index) (tIn [Iterator])
|
|
template <>
|
|
struct IndexDot<List<DimExpr>, tOut<Index>, tIn<List<Iterator>>>
|
|
: public Tuple<List<DimExpr>, tOut<Index>, tIn<List<Iterator>>> {
|
|
using Tuple<List<DimExpr>, tOut<Index>, tIn<List<Iterator>>>::Tuple;
|
|
};
|
|
|
|
template <typename DimT, typename OutT, typename InT>
|
|
struct IndexUnDot;
|
|
|
|
// IndexUnDot [DimExpr] (tOut [Iterator]) (tIn Index)
|
|
template <>
|
|
struct IndexUnDot<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>
|
|
: public Tuple<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>> {
|
|
using Tuple<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>::Tuple;
|
|
};
|
|
|
|
// OpArgIndexes = (tIn [Index], tOut [Index])
|
|
template <typename OutIndexT>
|
|
struct OpArgIndexes final
|
|
: public Tuple<tIn<List<Index>>, tOut<List<OutIndexT>>> {
|
|
using Tuple<tIn<List<Index>>, tOut<List<OutIndexT>>>::Tuple;
|
|
};
|
|
|
|
template <typename FakeOpT, typename OutT, typename InT>
|
|
struct InMsg2OutMsg;
|
|
|
|
// InMsg2OutMsg (tOut FakeOpPlaceHolder) (tOut (tOutMsg OpArgIndexes))
|
|
// (tIn (tInMsg OpArgIndexes))
|
|
template <>
|
|
struct InMsg2OutMsg<tOut<FakeOpPlaceHolder>,
|
|
tOut<OpArgIndexes<std::optional<Index>>>,
|
|
tIn<OpArgIndexes<Index>>>
|
|
: public Tuple<tOut<FakeOpPlaceHolder>,
|
|
tOut<OpArgIndexes<std::optional<Index>>>,
|
|
tIn<OpArgIndexes<Index>>> {
|
|
using Tuple<tOut<FakeOpPlaceHolder>,
|
|
tOut<OpArgIndexes<std::optional<Index>>>,
|
|
tIn<OpArgIndexes<Index>>>::Tuple;
|
|
};
|
|
|
|
template <typename T0, typename T1>
|
|
struct ConstantFunction;
|
|
|
|
template <>
|
|
struct ConstantFunction<tOut<Iterator>, tIn<Index>> final
|
|
: public Tuple<tOut<Iterator>, tIn<Index>, DimExpr> {
|
|
using Tuple<tOut<Iterator>, tIn<Index>, DimExpr>::Tuple;
|
|
};
|
|
|
|
template <typename DimT, typename OutT, typename InT>
|
|
struct GetBroadcastedIterator;
|
|
|
|
template <>
|
|
struct GetBroadcastedIterator<DimExpr, tOut<Iterator>, tIn<Iterator>>
|
|
: public Tuple<DimExpr, tOut<Iterator>, tIn<Iterator>> {
|
|
using Tuple<DimExpr, tOut<Iterator>, tIn<Iterator>>::Tuple;
|
|
};
|
|
|
|
// clang-format off
|
|
DEFINE_ADT_UNION(Equation,
|
|
Identity<tOut<Iterator>, tIn<Iterator>>,
|
|
Identity<tOut<Index>, tIn<Index>>,
|
|
GetBroadcastedIterator<DimExpr,
|
|
tOut<Iterator>, tIn<Iterator>>,
|
|
IndexDot<List<DimExpr>, tOut<Index>,
|
|
tIn<List<Iterator>>>,
|
|
IndexUnDot<List<DimExpr>,
|
|
tOut<List<Iterator>>, tIn<Index>>,
|
|
InMsg2OutMsg<tOut<FakeOpPlaceHolder>,
|
|
tOut<OpArgIndexes<std::optional<Index>>>,
|
|
tIn<OpArgIndexes<Index>>>,
|
|
ConstantFunction<tOut<Iterator>, tIn<Index>>);
|
|
// clang-format on
|
|
|
|
// Function = Equation
|
|
using Function = Equation;
|
|
|
|
using Equations = List<Equation>;
|
|
using GraphView = EquationGraphTopoWalker<Variable, const Equation*>;
|
|
|
|
std::string GetFunctionTypeName(const Function& function);
|
|
|
|
const void* GetFunctionDataPtr(const Function& function);
|
|
|
|
} // namespace cinn::adt
|