wehub-resource-sync
340 lines
12 KiB
C++
340 lines
12 KiB
C++
/*
|
|
* Licensed to the Apache Software Foundation (ASF) under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you 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.
|
|
*/
|
|
|
|
/*!
|
|
* \file tvm/arith/int_set.h
|
|
* \brief Integer set
|
|
*/
|
|
#ifndef TVM_ARITH_INT_SET_H_
|
|
#define TVM_ARITH_INT_SET_H_
|
|
|
|
#include <tvm/ir/expr.h>
|
|
#include <tvm/tirx/expr.h>
|
|
|
|
#include <unordered_map>
|
|
|
|
namespace tvm {
|
|
namespace arith {
|
|
|
|
using tirx::IterVar;
|
|
using tirx::Var;
|
|
using tirx::VarNode;
|
|
|
|
class AnalyzerObj;
|
|
class Analyzer;
|
|
|
|
//-----------------------------------------------
|
|
// Integer set data structure.
|
|
//
|
|
// This is a API build on top of the base
|
|
// integer analysis API to provide set analysis.
|
|
//------------------------------------------------
|
|
/*!
|
|
* \brief Sign type of an integer expression.
|
|
*/
|
|
enum SignType { kPositive, kNegative, kZero, kUnknown };
|
|
|
|
/*!
|
|
* \brief Base class of all Integer set containers.
|
|
* represent a set of integers in one dimension.
|
|
* \sa IntSet
|
|
*/
|
|
class IntSetNode : public ffi::Object {
|
|
public:
|
|
TVM_FFI_DECLARE_OBJECT_INFO("ir.IntSet", IntSetNode, ffi::Object);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to IntSetNode.
|
|
* \sa IntSetNode
|
|
*/
|
|
class IntSet : public ffi::ObjectRef {
|
|
public:
|
|
/*!
|
|
* \brief Find a range that covers the region.
|
|
* \param max_range The range to be covered.
|
|
* \return The covering range.
|
|
*/
|
|
Range CoverRange(Range max_range) const;
|
|
/*! \return Lower bound of the set */
|
|
PrimExpr min() const;
|
|
/*! \return upper bound of the set */
|
|
PrimExpr max() const;
|
|
/*! \return The sign of the elements in the integer set */
|
|
SignType GetSignType() const;
|
|
/*! \return Whether the set represent nothing */
|
|
bool IsNothing() const;
|
|
/*! \return Whether the set represent everything */
|
|
bool IsEverything() const;
|
|
/*! \return Whether the set is a single point */
|
|
bool IsSinglePoint() const;
|
|
/*!
|
|
* \brief Check if we can prove it is a single point.
|
|
*
|
|
* Unlike IsSinglePoint, which only checks ptr equality
|
|
* this function will invoke analyzer to do stonger proofs
|
|
* but also takes longer time.
|
|
*
|
|
* Use this function in some of the primitives but do not
|
|
* use it in the inner loop of simplification.
|
|
*
|
|
* \param ana Analyzer used in the proof.
|
|
* \return Whether we can prove it is a single point
|
|
*/
|
|
bool CanProveSinglePoint(const Analyzer& ana) const;
|
|
// TODO(tvm-team): update all CanProve to explicitly take
|
|
// analyzer to encourage more analyzer reuse
|
|
/*! \return Whether the set is proved to be bigger than 0 */
|
|
bool CanProvePositive() const;
|
|
/*! \return Whether the set is proved to be smaller than 0 */
|
|
bool CanProveNegative() const;
|
|
/*! \return Whether the set is proved to be smaller than or equal to 0 */
|
|
bool CanProveNonPositive() const;
|
|
/*! \return Whether the set is proved to be larger than or equal to 0 */
|
|
bool CanProveNonNegative() const;
|
|
/*! \return Whether the set has upper bound. */
|
|
bool HasUpperBound() const;
|
|
/*! \return Whether the set has lower bound. */
|
|
bool HasLowerBound() const;
|
|
|
|
/*!
|
|
* \brief The single point value, call only if IsSinglePoint is true
|
|
* \return The point value.
|
|
*/
|
|
PrimExpr PointValue() const;
|
|
/*!
|
|
* \brief Try to match IntSet with range r.
|
|
*
|
|
* \note It is guanrateed that IntSet::FromRange(r).MatchRange(r) == true
|
|
* \return true if we can prove they are the same.
|
|
*/
|
|
bool MatchRange(const tvm::Range& r) const;
|
|
/*! \return The set contains nothing */
|
|
static IntSet Nothing();
|
|
/*! \return The set contains everything */
|
|
static IntSet Everything();
|
|
/*!
|
|
* \brief construct a point set.
|
|
* \param point The point in the set.
|
|
* \return construct a single point set
|
|
*/
|
|
static IntSet SinglePoint(PrimExpr point);
|
|
/*!
|
|
* \brief construct a integer set from vector expression.
|
|
* \param vec The vector expression, can also be single point.
|
|
* \return The result set containing the indices in the vector.
|
|
*/
|
|
static IntSet Vector(PrimExpr vec);
|
|
/*!
|
|
* \brief Construct a set representing a range [min, min + extent).
|
|
* \param min The minimum of the range range
|
|
* \param extent The extent of the range.
|
|
* \return The constructed set.
|
|
*/
|
|
static IntSet FromMinExtent(PrimExpr min, PrimExpr extent);
|
|
/*!
|
|
* \brief Construct a set representing a range.
|
|
* \param r The range
|
|
* \return The constructed set.
|
|
*/
|
|
static IntSet FromRange(tvm::Range r);
|
|
/*!
|
|
* \brief Construct a set representing a interval.
|
|
* \param min The minimum value of the interval.
|
|
* \param max The maximum value of the interval.
|
|
* \return The constructed set.
|
|
*/
|
|
static IntSet Interval(PrimExpr min, PrimExpr max);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntSet, ffi::ObjectRef, IntSetNode);
|
|
};
|
|
|
|
//-----------------------------------------------
|
|
// Integer set legacy API.
|
|
//------------------------------------------------
|
|
/*!
|
|
* \brief Convert std::unordered_map<const VarNode*, IntSet> to ffi::Map<Var, IntSet>
|
|
*
|
|
* \param dom_map The domain map to convert.
|
|
* \return The converted map.
|
|
*/
|
|
ffi::Map<Var, IntSet> ConvertDomMap(const std::unordered_map<const VarNode*, IntSet>& dom_map);
|
|
/*!
|
|
* \brief Find an symbolic integer set that contains all possible values of
|
|
* e given the domain of each iteration variables.
|
|
*
|
|
* \param e The expression to be evaluated.
|
|
* \param dom_map The domain of each variable.
|
|
* \return An integer set that can cover all the possible values of e.
|
|
*/
|
|
IntSet EvalSet(PrimExpr e, const ffi::Map<IterVar, IntSet>& dom_map);
|
|
/*!
|
|
* \brief Find an symbolic integer set that contains all possible values of
|
|
* e given the domain of each variables.
|
|
*
|
|
* \param e The expression to be evaluated.
|
|
* \param dom_map The domain of each variable.
|
|
* \return An integer set that can cover all the possible values of e.
|
|
*/
|
|
IntSet EvalSet(PrimExpr e, const ffi::Map<Var, IntSet>& dom_map);
|
|
/*!
|
|
* \brief Same as EvalSet, but takes unordered_map
|
|
*
|
|
* \param e The expression to be evaluated.
|
|
* \param dom_map The domain of each variable.
|
|
* \return An integer set that can cover all the possible values of e.
|
|
*/
|
|
IntSet EvalSet(PrimExpr e, const std::unordered_map<const tirx::VarNode*, IntSet>& dom_map);
|
|
/*!
|
|
* \brief Find an symbolic integer set that contains is union over
|
|
* all the possible conditional values in dom_map.
|
|
*
|
|
* \param r The initial range.
|
|
* \param dom_map The domain of each variable.
|
|
* \return An integer set that can cover all the possible values.
|
|
*/
|
|
IntSet EvalSet(Range r, const ffi::Map<IterVar, IntSet>& dom_map);
|
|
|
|
/*!
|
|
* \brief Find an symbolic integer set that contains is union over
|
|
* all the possible conditional values in dom_map.
|
|
*
|
|
* \param s The initial set.
|
|
* \param dom_map The domain of each variable.
|
|
* \return An integer set that can cover all the possible values.
|
|
*/
|
|
IntSet EvalSet(IntSet s, const std::unordered_map<const VarNode*, IntSet>& dom_map);
|
|
/*!
|
|
* \brief Same as EvalSet, but takes unordered_map
|
|
*
|
|
* \param r The range to be evaluated.
|
|
* \param dom_map The domain of each variable.
|
|
* \return An integer set that can cover all the possible values of e.
|
|
*/
|
|
IntSet EvalSet(Range r, const std::unordered_map<const VarNode*, IntSet>& dom_map);
|
|
/*!
|
|
* \brief Same as EvalSet, but takes ffi::Array<Range>
|
|
*
|
|
* \param region The range to be evaluated.
|
|
* \param dom_map The domain of each variable.
|
|
* \return An array of integer sets that can cover all the possible values.
|
|
*/
|
|
ffi::Array<IntSet> EvalSet(const ffi::Array<Range>& region, const ffi::Map<Var, IntSet>& dom_map);
|
|
/*! \brief Map from Expr to IntSet */
|
|
using ExprIntSetMap = std::unordered_map<PrimExpr, IntSet, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
|
|
/*!
|
|
* \brief Find the integer set of every sub-expression, given the
|
|
* domain of each iteration variables.
|
|
*
|
|
* \param e The expression to be evaluated.
|
|
* \param dom_map The domain of each variable.
|
|
* \return the map from the expression to its possible value.
|
|
*/
|
|
ExprIntSetMap EvalSetForEachSubExpr(PrimExpr e,
|
|
const std::unordered_map<const VarNode*, IntSet>& dom_map);
|
|
|
|
/*!
|
|
* \brief Create a union set of all sets, possibly relaxed
|
|
* \param sets The sets to be combined
|
|
* \return the set after union
|
|
*/
|
|
IntSet Union(const ffi::Array<IntSet>& sets);
|
|
|
|
/*!
|
|
* \brief The union of N-dimensional integer sets
|
|
* \param nd_int_sets A list of N-dimensional integer sets
|
|
* \return An N-dimensional integer set as the result of union
|
|
*/
|
|
ffi::Array<IntSet> UnionRegion(const ffi::Array<ffi::Array<IntSet>>& nd_int_sets);
|
|
|
|
/*!
|
|
* \brief Create a lower-bound of union set, where some of the segments may be dropped
|
|
* \param sets The sets to be combined
|
|
* \return the set after union
|
|
*/
|
|
IntSet UnionLowerBound(const ffi::Array<IntSet>& sets);
|
|
|
|
/*!
|
|
* \brief The union of N-dimensional integer sets
|
|
* \param nd_int_sets A list of N-dimensional integer sets
|
|
* \return An N-dimensional integer set as the result of union
|
|
*/
|
|
ffi::Array<IntSet> UnionRegionLowerBound(const ffi::Array<ffi::Array<IntSet>>& nd_int_sets);
|
|
|
|
/*!
|
|
* \brief Create an intersected set of all sets
|
|
* \param sets The sets to be intersected
|
|
* \return the set after intersected
|
|
*/
|
|
IntSet Intersect(const ffi::Array<IntSet>& sets);
|
|
|
|
/*!
|
|
* \brief Converts the Ranges to IntSets
|
|
* \param var_dom The ranges of variables
|
|
* \return The integer sets of the variables
|
|
*/
|
|
ffi::Map<Var, arith::IntSet> AsIntSet(const ffi::Map<Var, Range>& var_dom);
|
|
|
|
/*!
|
|
* \brief Analyze the region with affine map, given the domain of variables and their predicate.
|
|
* The result should be strict, i.e. no region is discarded or relaxed.
|
|
* \param region The region to be analyzed
|
|
* \param var_dom The ranges of the variables
|
|
* \param predicate The predicate for the affine map
|
|
* \param analyzer The analyzer used
|
|
* \return std::nullopt if the detection fails, or an array of arith::IntSet as the result of
|
|
* analysis
|
|
*/
|
|
TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(
|
|
const ffi::Array<Range>& region, const ffi::Map<Var, Range>& var_dom, const PrimExpr& predicate,
|
|
const arith::Analyzer& analyzer);
|
|
|
|
/*!
|
|
* \brief Analyze the region with affine map, given the domain of variables and their predicate.
|
|
* Some subregion may be discarded during the lower-bound analysis.
|
|
* \param region The region to be analyzed
|
|
* \param var_dom The ranges of the variables
|
|
* \param predicate The predicate for the affine map
|
|
* \param analyzer The analyzer used
|
|
* \return std::nullopt if the detection fails, or an array of arith::IntSet as the result of
|
|
* analysis
|
|
*/
|
|
TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionLowerBound(
|
|
const ffi::Array<Range>& region, const ffi::Map<Var, Range>& var_dom, const PrimExpr& predicate,
|
|
const arith::Analyzer& analyzer);
|
|
|
|
/*!
|
|
* \brief Analyze the region with affine map, given the domain of variables and their predicate
|
|
* Relaxation of the region may be used in upper-bound analysis, i.e. some extra region may be added
|
|
* to the result.
|
|
* \param region The region to be analyzed
|
|
* \param var_dom The ranges of the variables
|
|
* \param predicate The predicate for the affine map
|
|
* \param analyzer The analyzer used
|
|
* \return an array of arith::IntSet as the result of analysis
|
|
*/
|
|
TVM_DLL ffi::Array<IntSet> EstimateRegionUpperBound(const ffi::Array<Range>& region,
|
|
const ffi::Map<Var, Range>& var_dom,
|
|
const PrimExpr& predicate,
|
|
const arith::Analyzer& analyzer);
|
|
|
|
} // namespace arith
|
|
} // namespace tvm
|
|
#endif // TVM_ARITH_INT_SET_H_
|