chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,272 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <ostream>
|
||||
#include <string>
|
||||
#include <variant>
|
||||
#include <vector>
|
||||
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/common/overloaded.h"
|
||||
#include "paddle/pir/include/core/dll_decl.h"
|
||||
#include "paddle/pir/include/core/utils.h"
|
||||
|
||||
namespace symbol {
|
||||
|
||||
#define SYMBOL_NOT_IMPLEMENTED \
|
||||
PADDLE_THROW(common::errors::Unimplemented("Not Implemented"))
|
||||
|
||||
template <typename T>
|
||||
struct UnaryDimExpr {
|
||||
explicit UnaryDimExpr(const T& d) : data(std::make_shared<Data>(d)) {}
|
||||
struct Data {
|
||||
explicit Data(const T& d) : data(d) {}
|
||||
T data;
|
||||
};
|
||||
|
||||
const Data& operator*() const { return *data; }
|
||||
Data& operator*() { return *data; }
|
||||
const Data* operator->() const { return data.get(); }
|
||||
Data* operator->() { return data.get(); }
|
||||
|
||||
std::shared_ptr<Data> data;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct BinaryDimExpr {
|
||||
explicit BinaryDimExpr(const T& l, const T& r)
|
||||
: data(std::make_shared<Data>(l, r)) {}
|
||||
|
||||
struct Data {
|
||||
explicit Data(const T& l, const T& r) : lhs(l), rhs(r) {}
|
||||
T lhs;
|
||||
T rhs;
|
||||
};
|
||||
|
||||
const Data& operator*() const { return *data; }
|
||||
Data& operator*() { return *data; }
|
||||
const Data* operator->() const { return data.get(); }
|
||||
Data* operator->() { return data.get(); }
|
||||
|
||||
std::shared_ptr<Data> data;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class List final {
|
||||
public:
|
||||
List(const List&) = default;
|
||||
List(List&&) = default;
|
||||
List& operator=(const List&) = default;
|
||||
List& operator=(List&&) = default;
|
||||
|
||||
using value_type = T;
|
||||
|
||||
explicit List() : vector_(std::make_shared<std::vector<T>>()) {}
|
||||
|
||||
template <
|
||||
typename Arg,
|
||||
std::enable_if_t<!std::is_same_v<std::decay_t<Arg>, List>, bool> = true>
|
||||
explicit List(Arg&& arg)
|
||||
: vector_(std::make_shared<std::vector<T>>(
|
||||
std::vector<T>{std::forward<Arg>(arg)})) {}
|
||||
|
||||
template <typename Arg0, typename Arg1, typename... Args>
|
||||
List(Arg0&& arg0, Arg1&& arg1, Args&&... args)
|
||||
: vector_(std::make_shared<std::vector<T>>(
|
||||
std::vector<T>{std::forward<Arg0>(arg0),
|
||||
std::forward<Arg1>(arg1),
|
||||
std::forward<Args>(args)...})) {}
|
||||
|
||||
bool operator==(const List& other) const {
|
||||
if (&vector() == &other.vector()) {
|
||||
return true;
|
||||
}
|
||||
return vector() == other.vector();
|
||||
}
|
||||
|
||||
bool operator!=(const List& other) const { return !(*this == other); }
|
||||
|
||||
std::vector<T>& operator*() const { return *vector_; }
|
||||
std::vector<T>* operator->() const { return vector_.get(); }
|
||||
|
||||
const std::vector<T>& vector() const { return *vector_; }
|
||||
|
||||
const auto& Get(std::size_t idx) const { return vector_->at(idx); }
|
||||
|
||||
private:
|
||||
std::shared_ptr<std::vector<T>> vector_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Negative final : public UnaryDimExpr<T> {
|
||||
using UnaryDimExpr<T>::UnaryDimExpr;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Add final {
|
||||
List<T> operands;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Mul final {
|
||||
List<T> operands;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Div final : public BinaryDimExpr<T> {
|
||||
using BinaryDimExpr<T>::BinaryDimExpr;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Max final {
|
||||
List<T> operands;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Min final {
|
||||
List<T> operands;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Broadcast final {
|
||||
List<T> operands;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Equal final : public BinaryDimExpr<T> {
|
||||
using BinaryDimExpr<T>::BinaryDimExpr;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Broadcastable final : public BinaryDimExpr<T> {
|
||||
using BinaryDimExpr<T>::BinaryDimExpr;
|
||||
bool operator==(const Broadcastable& other) const {
|
||||
return (this->data->lhs == other.data->lhs &&
|
||||
this->data->rhs == other.data->rhs) ||
|
||||
(this->data->lhs == other.data->rhs &&
|
||||
this->data->rhs == other.data->lhs);
|
||||
}
|
||||
};
|
||||
|
||||
class DimExpr;
|
||||
|
||||
// DimExpr = std::int64_t
|
||||
// | std::string
|
||||
// | Negative DimExpr
|
||||
// | Add DimExpr
|
||||
// | Mul DimExpr
|
||||
// | Div DimExpr
|
||||
// | Max DimExpr
|
||||
// | Min DimExpr
|
||||
// | Broadcast DimExpr
|
||||
using DimExprBase = std::variant<std::int64_t,
|
||||
std::string,
|
||||
Negative<DimExpr>,
|
||||
Add<DimExpr>,
|
||||
Mul<DimExpr>,
|
||||
Div<DimExpr>,
|
||||
Max<DimExpr>,
|
||||
Min<DimExpr>,
|
||||
Broadcast<DimExpr>>;
|
||||
|
||||
class IR_API DimExpr : public DimExprBase {
|
||||
public:
|
||||
using DimExprBase::DimExprBase;
|
||||
|
||||
template <typename T>
|
||||
bool isa() const {
|
||||
return std::holds_alternative<T>(*this);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const T& dyn_cast() const {
|
||||
return std::get<T>(*this);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool Has() const {
|
||||
return std::holds_alternative<T>(*this);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const T& Get() const {
|
||||
return std::get<T>(*this);
|
||||
}
|
||||
|
||||
const DimExprBase& variant() const {
|
||||
return static_cast<const DimExprBase&>(*this);
|
||||
}
|
||||
|
||||
DEFINE_MATCH_METHOD();
|
||||
|
||||
DimExpr operator+(const DimExpr& other) const;
|
||||
DimExpr operator-(const DimExpr& other) const;
|
||||
DimExpr operator*(const DimExpr& other) const;
|
||||
DimExpr operator/(const DimExpr& other) const;
|
||||
bool operator==(const DimExpr& other) const;
|
||||
bool operator!=(const DimExpr& other) const;
|
||||
};
|
||||
|
||||
// DimExprConstraint = Equal DimExpr
|
||||
// | Broadcastable DimExpr
|
||||
using DimExprConstraint = std::variant<Equal<DimExpr>, Broadcastable<DimExpr>>;
|
||||
|
||||
IR_API std::string ToString(const DimExpr& dim_expr);
|
||||
|
||||
IR_API std::ostream& operator<<(std::ostream&, const DimExpr& dim_expr);
|
||||
|
||||
IR_API std::ostream& operator<<(std::ostream&,
|
||||
const std::vector<DimExpr>& dim_exprs);
|
||||
|
||||
IR_API std::size_t GetHashValue(const DimExpr& dim_expr);
|
||||
|
||||
} // namespace symbol
|
||||
|
||||
namespace std {
|
||||
|
||||
template <>
|
||||
struct hash<symbol::DimExpr> {
|
||||
std::size_t operator()(const symbol::DimExpr& dim_expr) const {
|
||||
return symbol::GetHashValue(dim_expr);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct hash<std::vector<symbol::DimExpr>> {
|
||||
std::size_t operator()(const std::vector<symbol::DimExpr>& dim_exprs) const {
|
||||
std::size_t hash_value = 0;
|
||||
const auto hash_func = std::hash<symbol::DimExpr>();
|
||||
for (const auto& dim_expr : dim_exprs) {
|
||||
hash_value = pir::detail::hash_combine(hash_value, hash_func(dim_expr));
|
||||
}
|
||||
return hash_value;
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct hash<symbol::Broadcastable<symbol::DimExpr>> {
|
||||
std::size_t operator()(
|
||||
const symbol::Broadcastable<symbol::DimExpr>& broadcastable) const {
|
||||
return pir::detail::hash_combine(GetHashValue(broadcastable.data->lhs),
|
||||
GetHashValue(broadcastable.data->rhs));
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace std
|
||||
Reference in New Issue
Block a user