wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Waiting to run
Details
CI / MacOS (push) Waiting to run
Details
CI / Windows (push) Waiting to run
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tvm/arith/analyzer.h
+993
View File
@@ -0,0 +1,993 @@
/*
* 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/analyzer.h
* \brief Algebra expression simplifications.
*/
#ifndef TVM_ARITH_ANALYZER_H_
#define TVM_ARITH_ANALYZER_H_
#include <tvm/arith/int_set.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/with_context.h>
#include <limits>
#include <memory>
#include <unordered_map>
#include <utility>
#include <vector>
namespace tvm {
/*! \brief namespace of arithmetic analysis. */
namespace arith {
//-------------------------------------------------------
// Base integer analysis API.
//
// We have multiple type of analyzers to do relaxed
// integer set analysis(bound analysis, modulo) and
// equivalence checking and simplification.
//
// Importantly, each analyzer may need result from
// another analyzer.
//-------------------------------------------------------
// Forward declare the analyzer object and its reference handle.
class AnalyzerObj;
class Analyzer;
class ConstraintContext;
using tirx::Var;
enum DivMode {
/*! \brief Truncated division. */
kTruncDiv,
/*! \brief Floor division. */
kFloorDiv
};
/*!
* \brief The strength used in top-level condition proves
* \note The higher, the more time consuming it can be.
*
* Do not use level beyond kDefault in internal recursive rewriting in arith
* analysis and only use it at top-level simplification to avoid speed issues.
*/
enum class ProofStrength : int {
/*! \brief default strength, can be used in. */
kDefault = 0,
/*!
* \brief Prove using symbolic bound analysis
*/
kSymbolicBound = 1,
};
/*!
* \brief Constant integer up and lower bound(inclusive).
* Useful for value bound analysis.
*
* set = [min_value, max_value]
*/
class ConstIntBoundNode : public ffi::Object {
public:
int64_t min_value;
int64_t max_value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ConstIntBoundNode>()
.def_ro("min_value", &ConstIntBoundNode::min_value)
.def_ro("max_value", &ConstIntBoundNode::max_value);
}
/*! \brief Number to represent +inf */
static const constexpr int64_t kPosInf = std::numeric_limits<int64_t>::max();
/*!
* \brief Number to represent -inf
* \note We can make use the of fact that -kPosInf == kNegInf in the project.
*/
static const constexpr int64_t kNegInf = -kPosInf;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.ConstIntBound", ConstIntBoundNode, ffi::Object);
};
/*!
* \brief reference class to ConstIntBoundNode
* \sa ConstIntBoundNode
*/
class ConstIntBound : public ffi::ObjectRef {
public:
/*!
* \brief constructor by fields.
* \param min_value The mininum value.
* \param max_value The maximum value.
*/
TVM_DLL ConstIntBound(int64_t min_value, int64_t max_value);
static const constexpr int64_t kPosInf = ConstIntBoundNode::kPosInf;
static const constexpr int64_t kNegInf = ConstIntBoundNode::kNegInf;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ConstIntBound, ffi::ObjectRef, ConstIntBoundNode);
};
/*!
* \brief Analyzer to get constant integer bound over expression.
*/
class ConstIntBoundAnalyzer {
public:
using BoundMapType =
std::unordered_map<PrimExpr, ConstIntBound, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL ConstIntBound operator()(const PrimExpr& expr) const;
/*!
* \brief analyze the expr with the intermediate memorized to avoid redundant computation
* \param expr The expression of interest.
* \param bound The lookup table to store the intermediate results
* \return the result of the analysis.
*/
TVM_DLL ConstIntBound operator()(const PrimExpr& expr, BoundMapType* bound);
/*!
* \brief Update constant int bound information of var.
*
* \param var The variable of interest.
* \param info The bound information.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const ConstIntBound& info, bool allow_override = false);
/*!
* \brief Bind variable to a range.
*
* \param var The variable.
* \param range The range we bind to.
* \param allow_override Whether we allow overriding an existing var's range.
*/
TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Check if a variable is bound to a range.
* \param var The variable.
* \return Whether the variable is bound to a range.
*/
TVM_DLL bool IsBound(const Var& var) const;
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit ConstIntBoundAnalyzer(AnalyzerObj* parent);
TVM_DLL ~ConstIntBoundAnalyzer();
void CopyFrom(const ConstIntBoundAnalyzer& other);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
std::function<void()> EnterConstraint(const PrimExpr& constraint);
struct Entry;
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*!
* \brief Range of a linear integer function.
* Use to do specify the possible index values.
*
* set = { coeff * x + base | x in Z }
*
* When coeff != 0, it can also be written as
* set = { n | n % coeff == base }
*
* This is useful to decide if the index is dividable by certain value.
* For example, if index = 0 + 4 x, then we know it can be divided by 4.
*/
class ModularSetNode : public ffi::Object {
public:
/*! \brief linear co-efficient */
int64_t coeff;
/*! \brief The base */
int64_t base;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ModularSetNode>()
.def_ro("coeff", &ModularSetNode::coeff)
.def_ro("base", &ModularSetNode::base);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.ModularSet", ModularSetNode, ffi::Object);
};
/*!
* \brief reference of ModularSetNode
* \sa ModularSetNode
*/
class ModularSet : public ffi::ObjectRef {
public:
TVM_DLL ModularSet(int64_t coeff, int64_t base);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ModularSet, ffi::ObjectRef, ModularSetNode);
};
/*!
* \brief Analyzer to get modular information over expression.
*/
class ModularSetAnalyzer {
public:
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL ModularSet operator()(const PrimExpr& expr);
/*!
* \brief Update constant int bound information of var.
*
* \param var The variable of interest.
* \param info The bound information.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const ModularSet& info, bool allow_override = false);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit ModularSetAnalyzer(AnalyzerObj* parent);
TVM_DLL ~ModularSetAnalyzer();
void CopyFrom(const ModularSetAnalyzer& other);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
std::function<void()> EnterConstraint(const PrimExpr& constraint);
struct Entry;
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*!
* \brief Rewrite-rule based simplifier.
*/
class RewriteSimplifier {
public:
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL PrimExpr operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_expr
* \param allow_override Whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);
/*! \brief Flags to enable more computationally-intensive simplifications
*
* These simplifications may be required for specific schedules, but
* would impose too high a compile-time cost to enable by default.
* They can be enabled on an as-needed basis by calling
* `RewriteSimplifier::SetEnabledExtensions` prior to using
* `RewriteSimplifier::operator()`.
*
* Flags are defined as powers of two to allow future expansion. To
* enable multiple extensions, a user should pass a bitwise OR of the
* flags for each desired extension.
*/
enum Extension {
// No extensions enabled
kNone = 0,
/* When simplifying an inequality, attempt to use scope-based knowns.
*
* Example:
* if_then_else(i<j && j<k, i<k, false) => if_then_else(i<j && j<k, true, false)
*/
kTransitivelyProveInequalities = (1 << 0),
/* When simplifying a boolean expression, convert to an AND of ORs
* (conjunctive normal form).
*
* Example:
* (a && b) || c => (a || c) && (b || c)
*/
kConvertBooleanToAndOfOrs = (1 << 1),
/* When simplifying a boolean AND or a boolean OR, simplify each
* branch under the assumption that the other branch does not
* already dominate the result. That is, simplify each branch of
* (A && B) under the assumption that the other branch is true,
* and simplify each branch of (A || B) under the assumption that
* the other branch is false.
*
* Example:
* (n < 10) && (n < 5) => (n < 10)
* (n < 10) || (n < 5) => (n < 5)
*/
kApplyConstraintsToBooleanBranches = (1 << 2),
/* Special handling for expressions `(A+B)*C < (A*B)*D`
*
* Expressions of the form `(A+B)*C < (A*B)*D` can occur occur
* when comparing the number of operations required for two
* different orderings in which matrix multiplications can be
* performed. Proving or disproving this conditional allows an
* optimal order of execution to be selected, even for dynamic
* argument shapes.
*
* The default behavior of `ConstIntBounds` assumes that each term
* in an expression is independent, and is insufficient to prove
* these inequalities. For example, the maximum value of `(A+B)*C
* - (A*B)*D` is determined by taking the maximum value of
* `(A+B)*C` and subtracting the minimum value of `(A*B)*D`.
* While this algorithm can be applied in all cases, the bound it
* provides is looser than strictly required.
*
* This extension adds a check for this case. When `A`, `B`, `C`,
* and `D` are all positive values, as is the case for tensor
* shapes, the inequality can be written as `1/A + 1/B < D/C`. If
* this inequality holds for the minimum values of `A`, `B`, and
* `D`, along with the maximum value of `C`, then the inequality
* holds for all values.
*
* This extension requires little to no performance overhead, and
* may be enabled by default in future releases.
*/
kComparisonOfProductAndSum = (1 << 3),
};
/*! \brief Enable an optional extension or extensions
*
* \param flags A bitwise OR of all optional extensions that should
* be enabled.
*/
TVM_DLL void SetEnabledExtensions(Extension flags);
/*! \brief Return the currently enabled extensions */
TVM_DLL Extension GetEnabledExtensions() const;
/*! \brief Return the statistics counters */
TVM_DLL ffi::ObjectRef GetStatsCounters() const;
/*! \brief Reset the statistics counters */
TVM_DLL void ResetStatsCounters();
/*! \brief Set the maximum allowed number of rewrite steps
*
* By default, the simplifier may perform as many steps as are
* required. If a positive limit is set, then the simplifier will
* throw an exception when exceeding that number of rewrite steps.
* This allows tests to guard against performance regressions.
*
* Note: To maintain accurate usage counters, `Analyzer` instances
* should be re-used wherever possible. For example, TIR
* transformations should declare a single `Analyzer` that is used
* throughout the pass. Internal helper functions that only borrow
* the analyzer temporarily may receive the underlying `AnalyzerObj*`
* from their calling scope.
*/
TVM_DLL void SetMaximumRewriteSteps(int64_t maximum);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
friend class CanonicalSimplifier;
explicit RewriteSimplifier(AnalyzerObj* parent);
TVM_DLL ~RewriteSimplifier();
void CopyFrom(const RewriteSimplifier& other);
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*!
* \brief Canonical-form based simplifier.
*/
class CanonicalSimplifier {
public:
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL PrimExpr operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_expr
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit CanonicalSimplifier(AnalyzerObj* parent);
TVM_DLL ~CanonicalSimplifier();
void CopyFrom(const CanonicalSimplifier& other);
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*! \brief Structure for representing result of known
*
* Values are assigned to allow these flags to be used in bitwise
* operations.
*/
enum class CompareResult : int {
kInconsistent = 0,
kEQ = 1,
kLT = 2,
kLE = 3,
kGT = 4,
kGE = 5,
kNE = 6,
kUnknown = 7
};
inline constexpr CompareResult operator&(CompareResult lhs, CompareResult rhs) {
return CompareResult(static_cast<int>(lhs) & static_cast<int>(rhs));
}
inline constexpr CompareResult operator|(CompareResult lhs, CompareResult rhs) {
return CompareResult(static_cast<int>(lhs) | static_cast<int>(rhs));
}
/*!
* \brief Using previously specified knowns, compare the expressions provided
*
* Given known expressions [(a OP b), (b OP c), ..., (y OP z)], search
* for a known result for `(a OP z)`.
*/
class TransitiveComparisonAnalyzer {
public:
/* \brief Using previously specified knowns, compare the expressions provided
*
* \param lhs The left-hand side of the comparison
*
* \param rhs The right-hand side of the comparison
*
* \param propagate_inequalities If true, attempt to find a sequence
* of transitive inequalities that allow the lhs and rhs to be
* compared. If false, only use the known comparison that have been
* directly provided. Using `propagate_inequalities = false` is
* roughly equivalent to comparing against all known inequality
* expressions using `ExprDeepEqual`, but also allows for constant
* offsets on either side of the inequality.
*
* \return The most specific result that can be proven about the
* comparison. If nothing can be proven, returns kUnknown.
*/
TVM_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs,
bool propagate_inequalities = true);
/*! \brief Bind a variable as being equal to a known expression
*
* \param var The variable of interest.
* \param expr The bound expression
* \param allow_override Whether to allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*! \brief Bind a variable as being within a specified range
*
* \param var The variable of interest.
* \param range The known range
* \param allow_override Whether to allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
TransitiveComparisonAnalyzer();
TVM_DLL ~TransitiveComparisonAnalyzer();
void CopyFrom(const TransitiveComparisonAnalyzer& other);
class Impl;
/*! \brief Internal impl */
std::unique_ptr<Impl> impl_;
};
/*!
* \brief Integer set analyzer.
*/
class IntSetAnalyzer {
public:
/*!
* \brief Find a symbolic integer set that contains all possible values of
* expr given the domain of each variables.
*
* \param expr The expression of interest.
* \param dom_map The domain map to indicate which variable to relax.
* \return the result of the analysis.
*/
TVM_DLL IntSet operator()(const PrimExpr& expr, const ffi::Map<Var, IntSet>& dom_map);
/*!
* \brief Find a symbolic integer set that contains all possible
* values of expr given the domain of each variables, using
* the domain map defined by bound variables.
*
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL IntSet operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_interval_set The set of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const IntSet& new_interval_set, bool allow_override = false);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_range The range of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
std::function<void()> EnterConstraint(const PrimExpr& constraint);
private:
friend class AnalyzerObj;
explicit IntSetAnalyzer(AnalyzerObj* parent);
TVM_DLL ~IntSetAnalyzer();
void CopyFrom(const IntSetAnalyzer& other);
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
class Z3Prover {
public:
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_range The range of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param expr The bound expression.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*!
* \brief Whether the Z3 backend is compiled into this build (USE_Z3=ON).
*
* \return true if the real Z3 prover is available, false for the stub.
*/
TVM_DLL bool IsEnabled() const;
/*!
* \brief Whether can we prove expr is always true.
*
* \param expr The expression.
* \return Whether we can prove it.
*/
TVM_DLL bool CanProve(const PrimExpr& expr);
/*!
* \brief Update the internal state to enter constraint.
*
* \param constraint A constraint expression.
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
std::function<void()> EnterConstraint(const PrimExpr& constraint);
/*!
* \brief Get the SMTLIB2 representation of the current context.
*
* \param expr The optional expression to check.
* \return The SMTLIB2 string.
*/
ffi::String GetSMTLIB2(const ffi::Optional<PrimExpr> expr);
/*!
* \brief Get statistics about Z3 prover.
*
* \return The statistics string.
*/
ffi::String GetStats();
/*!
* \brief Set timeout in milliseconds for Z3 prover.
*
* \param timeout_ms The timeout in milliseconds.
*/
void SetTimeoutMs(unsigned timeout_ms);
/*!
* \brief Set resource limitation for Z3 prover.
*
* \param rlimit the resource limitation.
*/
void SetRLimit(unsigned rlimit);
/*!
* \brief Get the Z3 model for the given expression if satisfiable.
*
* \param expr The expression to get the model for.
* \return The model as a string.
*/
ffi::String GetModel(const PrimExpr& expr);
/*!
* \brief Count the number of integer values that satisfy the current constraints.
*
* This method uses Z3's model enumeration to count how many distinct values of
* the given variable satisfy all current constraints.
*
* \param var The variable to count satisfying values for.
* \param max_count Maximum number of solutions to enumerate.
* \param min_consecutive Minimum consecutive count requirement.
* \return The number of distinct values that satisfy the constraints, or a negative error code.
*/
TVM_DLL int64_t CountSatisfyingValues(const Var& var, int64_t max_count = 2048,
int64_t min_consecutive = 1);
private:
friend class AnalyzerObj;
friend class Analyzer;
explicit Z3Prover(AnalyzerObj* parent);
TVM_DLL ~Z3Prover();
void CopyFrom(const Z3Prover& other);
class Impl;
Impl* impl_;
};
/*!
* \brief Analyzer that contains bunch of sub-analyzers.
*
* Each sub-analyzer can make use of another sub-analyzer
* by weak reference of this.
*
* NOTE for sub-analyzer developers:
* If the analyzer uses memoization, we need to clear the internal
* cache when information about a Var has been overridden.
*/
class TVM_DLL AnalyzerObj : public ffi::Object {
public:
/*! \brief sub-analyzer: const integer bound */
ConstIntBoundAnalyzer const_int_bound;
/*! \brief sub-analyzer: modular set */
ModularSetAnalyzer modular_set;
/*! \brief sub-analyzer rewrite simplify */
RewriteSimplifier rewrite_simplify;
/*! \brief sub-analyzer canonical simplify */
CanonicalSimplifier canonical_simplify;
/*! \brief sub-analyzer: int set */
IntSetAnalyzer int_set;
/*! \brief sub-analyzer transitive comparisons */
TransitiveComparisonAnalyzer transitive_comparisons;
/*! \brief sub-analyzer using Z3 */
Z3Prover z3_prover;
/*! \brief constructor */
AnalyzerObj();
/*!
* \brief Mark the value as non-negative value globally in analyzer.
*
* Only call this function if the non-neg condition is global and
* not context-dependent.
*
* This function does best-effort propagations to the sub-analyzers
*
* A canonical use of MarkGlobalNonNegValue is to record a non-negativity
* fact at a Var's definition site. Because each Var identity is defined
* exactly once in canonical IR, the fact is globally valid for that identity.
*
* \note We expose this function because non-negative global values,
* such as symbolic buffer shapes in function arguments are really
* important to ensure the best simplification, and usually they
* can be handled in a simpler way than the generic constraints.
*
* This function may call into the Update function of the sub-analyzers.
*/
void MarkGlobalNonNegValue(const PrimExpr& value);
/*!
* \brief Notify all the sub-analyzers that var
* is created and binded to expr.
*
* Each var can only be bound once.
*
* \param var The variable.
* \param expr The expression we bind to.
* \param allow_override Whether we allow overriding an existing var's
* expression. This option should not be used if there is any dependency
* between variables.
*/
void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*!
* \brief Notify all the sub-analyzers that var
* is created and bound to a range.
*
* Each var can only be bound once.
*
* \param var The variable.
* \param range The range we bind to.
* \param allow_override Whether we allow overriding an existing var's
* expression. This option should not be used if there is any dependency
* between variables.
*/
void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Bind all the vars in the Map
*
* \param variables The {variable -> range} map.
* \param allow_override Whether we allow overriding an existing var's
* expression. This option should not be used if there is any dependency
* between variables.
*/
void Bind(const ffi::Map<Var, Range>& variables, bool allow_override = false);
/*!
* \brief Whether can we prove expr >= val.
* Non-negative proof is very useful in integer analysis
* to lower divisions and mods given difference in trunc and ceil mode.
*
* \param expr The expression.
* \param lower_bound The lower bound.
* \return Whether we can prove it.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
bool CanProveGreaterEqual(const PrimExpr& expr, int64_t lower_bound);
/*!
* \brief Whether can we prove expr < val.
* Non-negative proof is very useful in integer analysis
* to lower divisions and mods given difference in trunc and ceil mode.
*
* \param expr The expression.
* \param upper_bound The upper bound.
* \return Whether we can prove it.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
bool CanProveLess(const PrimExpr& expr, int64_t upper_bound);
/*!
* \brief Whether can we prove lhs == rhs.
*
* \param lhs The input lhs.
* \param rhs The input rhs.
* \return Whether we can prove lhs == rhs.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
bool CanProveEqual(const PrimExpr& lhs, const PrimExpr& rhs);
/*!
* \brief Whether we can prove lhs is smaller than possibly symbolic shape.
*
* By calling this function, the caller gives an extra hint that shape > 0,
* because it appeared in buffer shape.
*
* This is useful to prove condition such as 32 <= 32 * n where the 32 * n
* is known to be a shape. Use this routine to reduce the symbolic comparisons
* in buffer compaction.
*
* The underlying analyzer will use the kSymbolicBound proof.
*
* \param lhs The input lhs.
* \param shape The symbolic shape.
* \return Whether we can prove lhs <= shape.
*/
bool CanProveLessEqualThanSymbolicShapeValue(const PrimExpr& lhs, const PrimExpr& shape);
/*!
* \brief Whether can we prove condition.
*
* \param cond The expression to be proved.
* \param strength the strength of the prove.
*
* \return The result.
*
* \note Analyzer will call into sub-analyzers to get the result.
* Do not use strength beyond default in sub-analyzers and
* only use it in top-level predicate analysis.
*/
bool CanProve(const PrimExpr& cond, ProofStrength strength = ProofStrength::kDefault);
/*!
* \brief Simplify expr.
*
* \param expr The expression to be simplified.
* \param steps The simplification runs in the order of
* rewrite_simplify (step 1) -> canonical_simplify (step 2) ->
* rewrite_simplify (step 3) -> canonical_simplify (step 4) -> ...
* param steps controls how many steps to run.
* Default is 2, i.e., rewrite_simplify + canonical_simplify.
* \return The result.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
PrimExpr Simplify(const PrimExpr& expr, int steps = 2);
/*!
* \brief Deep-copy this analyzer into a new, independent Analyzer.
*
* The returned analyzer carries the same accumulated facts (variable
* bounds, modular sets, rewrite/canonical bindings, integer-set domains,
* literal constraints and transitive comparisons) as this one, but owns
* its own state: binding or simplifying on either analyzer afterwards does
* not affect the other. This is the deep copy that handle-copying an
* Analyzer does not provide.
*
* \note Do not call this while a `With<ConstraintContext>` scope is active
* on this analyzer. The clone would inherit the scoped constraints
* but not the recovery functions that pop them on scope exit, so the
* constraints would leak as if they were global facts. Clone at a
* point where no constraint scope is in effect.
*
* \return A new Analyzer holding an independent copy of the facts.
*/
Analyzer Clone() const;
/*!
* \brief Analyzer methods update facts, constraints, caches, and stats.
*
* Marking the object mutable makes the `Analyzer` ObjectRef expose a
* non-const `operator->`, so APIs can take `const Analyzer&` while still
* allowing calls such as `analyzer->Bind(...)`.
* `const Analyzer&` keeps the handle itself from being rebound; it does
* not make the underlying AnalyzerObj immutable.
*/
static constexpr bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.Analyzer", AnalyzerObj, ffi::Object);
};
/*!
* \brief Managed reference to AnalyzerObj.
*
* Analyzer is a lightweight, reference-counted handle around a heap-allocated
* AnalyzerObj. Because it is now a first-class FFI object, an Analyzer can be
* passed across the tvm-ffi boundary (e.g. handed from Python into a C++ pass)
* and shared, so that accumulated bindings/constraints persist across calls.
* Copying an Analyzer copies the handle, and both handles share the same
* mutable AnalyzerObj state.
* This is not a deep copy of analyzer facts or caches.
*
* \sa AnalyzerObj
*/
class Analyzer : public ffi::ObjectRef {
public:
/*! \brief Default-construct a fresh analyzer (allocates an AnalyzerObj). */
Analyzer() : Analyzer(ffi::make_object<AnalyzerObj>()) {}
explicit Analyzer(ffi::ObjectPtr<AnalyzerObj> n) : ffi::ObjectRef(std::move(n)) {
TVM_FFI_ICHECK(this->get() != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Analyzer, ffi::ObjectRef, AnalyzerObj);
};
/*!
* \brief Constraint context.
*
* \code
*
* Var x("x");
* arith::Analyzer analyzer;
* {
* With<arith::ConstraintContext> scope(analyzer, tvm::floormod(x, 3) == 0);
* TVM_FFI_ICHECK_EQ(analyzer->modular_set(x)->coeff, 3);
* }
* // constraint no longer in effect.
* TVM_FFI_ICHECK_NE(analyzer->modular_set(x)->coeff, 3);
*
* \endcode
*/
class ConstraintContext {
private:
// declare friend to enable with.
friend class With<ConstraintContext>;
/*!
* \brief Construct a constraint context.
* \param analyzer The analyzer whose context is updated. The context
* keeps a reference to the analyzer while the scope is active.
* \param constraint The constraint to be applied.
*/
ConstraintContext(const Analyzer& analyzer, PrimExpr constraint)
: ConstraintContext(analyzer, std::move(constraint), false) {}
/*!
* \brief Construct a constraint context.
* \param analyzer The analyzer whose context is updated. The context
* keeps a reference to the analyzer while the scope is active.
* \param constraint The constraint to be applied.
* \param is_assume Whether the constraint comes from an assumption.
*/
ConstraintContext(const Analyzer& analyzer, PrimExpr constraint, bool is_assume)
: analyzer_(analyzer), constraint_(std::move(constraint)), is_assume_(is_assume) {}
/*!
* \brief Construct a constraint context from a borrowed analyzer object.
* \param analyzer The borrowed analyzer object.
* \param constraint The constraint to be applied.
*
* This overload is for internal callers that already operate on AnalyzerObj*.
*/
ConstraintContext(AnalyzerObj* analyzer, PrimExpr constraint)
: ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint), false) {}
/*!
* \brief Construct a constraint context from a borrowed analyzer object.
* \param analyzer The borrowed analyzer object.
* \param constraint The constraint to be applied.
* \param is_assume Whether the constraint comes from an assumption.
*/
ConstraintContext(AnalyzerObj* analyzer, PrimExpr constraint, bool is_assume)
: ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint), is_assume) {}
// enter the scope.
void EnterWithScope();
// exit the scope.
void ExitWithScope();
/*! \brief Analyzer kept alive while the context is active. */
Analyzer analyzer_;
/*! \brief The constraint */
PrimExpr constraint_;
/*! \brief functions to be called in recovery */
std::vector<std::function<void()>> recovery_functions_;
/*! \brief Whether the constraint comes from an assumption. */
bool is_assume_;
};
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_ANALYZER_H_
include/tvm/arith/bound.h
+85
View File
@@ -0,0 +1,85 @@
/*
* 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/bound.h
* \brief Bound deducers.
*/
#ifndef TVM_ARITH_BOUND_H_
#define TVM_ARITH_BOUND_H_
#include <tvm/arith/int_set.h>
#include <tvm/ir/expr.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/stmt.h>
#include <unordered_map>
namespace tvm {
namespace arith {
using tirx::Region;
using tirx::Stmt;
using tirx::Var;
using tirx::VarNode;
/*!
* \brief Deduce the bound of the target variable in a expression,
* give the domain of each variables. Return undefined IntSet to
* represent failure.
*
* \note The returned set may be smaller than set that
* contains all possible values of v that satisfies the bound.
*
* \param v The target variable to be deduced.
* \param cond The conditional expression.
* \param hint_map The domain of variable, used to help deduce.
* \param relax_map The domain of each variable, used to relax the domain,
* The deduce bound must implies e for all value in relax_map
* \return An integer set that always satisfies the condition.
*/
IntSet DeduceBound(PrimExpr v, PrimExpr cond, const ffi::Map<Var, IntSet>& hint_map,
const ffi::Map<Var, IntSet>& relax_map);
/*!
* \brief Same as DeduceBound with unordered_map signature.
*
* \param v The target variable to be deduced.
* \param cond The conditional expression.
* \param hint_map The domain of variable, used to help deduce.
* \param relax_map The domain of each variable, used to relax the domain,
* The deduce bound mush implies e for all value in relax_map
* \return An integer set that always satisfies the condition.
*/
IntSet DeduceBound(PrimExpr v, PrimExpr cond,
const std::unordered_map<const VarNode*, IntSet>& hint_map,
const std::unordered_map<const VarNode*, IntSet>& relax_map);
/*!
* \brief Infer a regular domain that covers all the calls or provides within the given statement.
* \param body The given statement.
* \param buffer The buffer to check the access info.
* \param consider_loads If loads are considered.
* \param consider_stores If stores are considered.
* \return The domain that covers all the calls or provides within the given statement.
*/
Region DomainTouched(const Stmt& body, const tirx::Buffer& buffer, bool consider_loads,
bool consider_stores);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_BOUND_H_
include/tvm/arith/int_set.h
+339
View File
@@ -0,0 +1,339 @@
/*
* 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_
include/tvm/arith/int_solver.h
+338
View File
@@ -0,0 +1,338 @@
/*
* 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_solver.h
* \brief integer constraints data structures and solvers
*/
#ifndef TVM_ARITH_INT_SOLVER_H_
#define TVM_ARITH_INT_SOLVER_H_
#include <tvm/ir/expr.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <unordered_map>
#include <utility>
#include <vector>
#include "analyzer.h"
namespace tvm {
namespace arith {
using tirx::IterVar;
using tirx::Var;
using tirx::VarNode;
// According to experiments two best simplifications orders were can->rw and rw->can->rw,
// but rw->can->rw is better for a couple of cases.
// Also we should end with rw because it factors multipliers out.
constexpr int kSimplifyRewriteCanonicalRewrite = 3;
/*!
* \brief Represent integer grouped bounds which are classified into
* lower bounds (inclusive), upper bounds (inclusive) and equalities.
* It also contains coefficient as a multiplier for the bounds, i.e.,
* coef * var >= lower
* coef * var == equal
* coef * var <= upper
* \sa IntGroupBounds
*/
class IntGroupBoundsNode : public ffi::Object {
public:
PrimExpr coef;
ffi::Array<PrimExpr> lower;
ffi::Array<PrimExpr> equal;
ffi::Array<PrimExpr> upper;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IntGroupBoundsNode>()
.def_ro("coef", &IntGroupBoundsNode::coef)
.def_ro("lower", &IntGroupBoundsNode::lower)
.def_ro("equal", &IntGroupBoundsNode::equal)
.def_ro("upper", &IntGroupBoundsNode::upper);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntGroupBounds", IntGroupBoundsNode, ffi::Object);
};
/*!
* \brief Managed reference to IntGroupBoundsNode.
* \sa IntGroupBoundsNode
*/
class IntGroupBounds : public ffi::ObjectRef {
public:
/*!
* \brief Constructor by fields
* \param coef The coefficient. Must be integer.
* coef * var >= lower
* coef * var == equal
* coef * var >= upper
* \param lower the lower bounds (include)
* \param equal equalities
* \param upper the upper bounds (include)
*/
TVM_DLL IntGroupBounds(PrimExpr coef, ffi::Array<PrimExpr> lower, ffi::Array<PrimExpr> equal,
ffi::Array<PrimExpr> upper);
/*!
* \brief Construct bounds from a range.
* \param r The range
* \return constructed bounds.
*/
static IntGroupBounds FromRange(const Range& r);
/*!
* \brief Perform substitution on all components of the struct.
*/
IntGroupBounds Substitute(const ffi::Map<Var, PrimExpr>& subst) const;
/*!
* \brief Find the best range from the grouped bounds.
* \param vranges_addl additional variable ranges that help infer the best range.
* \return The best range (has the least difference between the lower bound and upper bound).
* undefined if (-inf, +inf).
*/
Range FindBestRange(const ffi::Map<Var, Range>& vranges_addl = {}) const;
/*!
* \brief Combine the bounds with another range.
* \param r range to be combined.
* \return combined bounds.
*/
IntGroupBounds operator+(const Range& r);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntGroupBounds, ffi::ObjectRef, IntGroupBoundsNode);
};
/*!
* \brief Represent integer constrains including (integer) variables, their ranges and
* the relations between them (either equations or inequalities).
* \sa LinearSystem
*/
class IntConstraintsNode : public ffi::Object {
public:
// e.g., \alpha, \beta, must be integers
ffi::Array<Var> variables;
// e.g., 1 <= \alpha <= N, etc.
// it is absolutely ok to include ranges for parameters
// (variables that are not in this->variables) in this map
ffi::Map<Var, Range> ranges;
// linear equalities or inequalities
// e.g., A \alpha = \beta or A \alpha <= \beta
ffi::Array<PrimExpr> relations;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IntConstraintsNode>()
.def_ro("variables", &IntConstraintsNode::variables)
.def_ro("ranges", &IntConstraintsNode::ranges)
.def_ro("relations", &IntConstraintsNode::relations);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntConstraints", IntConstraintsNode, ffi::Object);
};
/*!
* \brief Managed reference to IntConstraintsNode.
* \sa IntConstraintsNode
*/
class IntConstraints : public ffi::ObjectRef {
public:
/*!
* \brief Constructor by fields
* \param variables The variables in the constraints, must be integers.
* \param ranges The ranges of the variables.
* \param relations The linear relations between the variables
* (either equations or inequalities)
*/
TVM_DLL IntConstraints(ffi::Array<Var> variables, ffi::Map<Var, Range> ranges,
ffi::Array<PrimExpr> relations);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntConstraints, ffi::ObjectRef, IntConstraintsNode);
};
/*!
* \brief We can have different set of variables to represent the same constraints.
* For example, the following two systems are equivalent,
* {a + b = 0 | a >= 0, b >= 0} and
* {m - n = 0 | m >= 0, n <= 0}
* This data structure represents the transformation
* between two equivalent linear systems.
* In the above example,
* src : {a + b = 0 | a >= 0, b >= 0}
* dst : {m - n = 0 | m >= 0, n <= 0}
* src_to_dst : {a -> m, b -> -n}
* dst_to_src : {m -> a, n -> -b}
* \sa IntConstraintsTransform
*/
class IntConstraintsTransformNode : public ffi::Object {
public:
IntConstraints src;
IntConstraints dst;
ffi::Map<Var, PrimExpr> src_to_dst;
ffi::Map<Var, PrimExpr> dst_to_src;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IntConstraintsTransformNode>()
.def_ro("src", &IntConstraintsTransformNode::src)
.def_ro("dst", &IntConstraintsTransformNode::dst)
.def_ro("src_to_dst", &IntConstraintsTransformNode::src_to_dst)
.def_ro("dst_to_src", &IntConstraintsTransformNode::dst_to_src);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntConstraintsTransform", IntConstraintsTransformNode,
ffi::Object);
};
/*!
* \brief Managed reference to IntConstraintsTransformNode.
* \sa IntConstraintsTransformNode
*/
class IntConstraintsTransform : public ffi::ObjectRef {
public:
/*!
* \brief Constructor by fields
* \param src source integer constraints, e.g., {a + b = 0 | a >= 0, b >= 0}
* \param dst integer constraints equivalent to the source,
* e.g., {m - n = 0 | m >= 0, n <= 0}
* \param src_to_dst mapping from variables in the \p src to the variables in the \p dst,
* e.g., {a -> m, b -> -n}
* \param dst_to_src mapping from variables in the \p dst to the variables in the \p src,
* e.g., {m -> a, n -> -b}
*/
TVM_DLL IntConstraintsTransform(IntConstraints src, IntConstraints dst,
ffi::Map<Var, PrimExpr> src_to_dst,
ffi::Map<Var, PrimExpr> dst_to_src);
/*!
* \brief Chain-compose two IntConstraintsTransform together.
* this->dst must be the same as other->src.
* @param other another IntConstraintsTransform whose src is same as this->dst.
* @return composed IntConstraintsTransform(this->src, other->dst)
* with its variables and ranges are properly modified.
*/
IntConstraintsTransform operator+(const IntConstraintsTransform& other) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntConstraintsTransform, ffi::ObjectRef,
IntConstraintsTransformNode);
};
typedef std::pair<ffi::Map<Var, IntGroupBounds>, ffi::Array<PrimExpr>> PartialSolvedInequalities;
/*!
* \brief Obtain Smith Normal Form of linear equation A x = y.
* Smith Normal Form of matrix A_{mxn} is S_{mxn} = U_{mxm} A_{mxn} V_{nxn},
* in which S_{mxn} is diag(s1, s2, ..., sr, 0, ..., 0) and r is the rank of A.
* NOTE: Although in standard Smith Normal Form the diagonal elements satisfy
* s_i | s_{i+1} (| means divides), the implement here does not guarantee it.
* TODO(yzhliu): From sergei-grechanik:
* computing the proper Smith normal form may improve stability of automatic
* differentiation (generating the same gradient code for slightly different but equivalent input
* code U_{mxm} and V_{nxn} are invertible matrices. This function modifies \p S to be S_{mxn}, \p V
* to be V_{nxn}, \p y to be U_{mxm} y_{mx1} and \p x to be V^{-1} x. \param S the original
* A_{mxn}, it will be modified to S_{mxn} \param V an identity matrix, it will be modified to
* V_{nxn} \param x the x in A x = y. it will be modified to V^{-1}_{nxn} x_{nx1} \param y the y
* in A x = y. it will be modified to U_{mxm} y_{mx1}
*/
void SmithNormalFormDiag(std::vector<std::vector<int64_t>>* S, std::vector<std::vector<int64_t>>* V,
std::vector<PrimExpr>* x, std::vector<PrimExpr>* y);
/*!
* \brief Solve linear equations.
* \param system_to_solve the variables to solve, their ranges, and a list of equations.
* \return A new linear system, with less variables (if \p system_to_solve is NOT of full rank),
* or no variable (if \p system_to_solve is of full rank),
* or an empty linear system (if \p system_to_solve is unsolvable).
* It also provides the ranges of the variables in the new system,
* as well as inequalities inferred from the \p system_to_solve.
* You can get the mapping from the original variables to the solution via ret->src_to_dst.
*/
IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_solve);
/*!
* \brief Solve linear inequalities.
* \param system_to_solve the variables to solve, their ranges, and a list of inequalities.
* The inequalities are rewritten using Fourier-Motzkin elimination.
* This function takes an array of (in)equalities and an array of variables, and essentially
* rewrites the (in)equalities into an array of (in)equalities of the following form,
*
* x0 >= f0(x1, x2, ..., xn)
* x0 <= g0(x1, x2, ..., xn)
* x1 >= f1(x2, ..., xn)
* x1 <= g1(x2, ..., xn)
* ...
* xn >= fn() // just a constant
* xn <= gn() // just a constant
*
* \return A map of variables and their solved bounds,
* and constrains that cannot be solved to bounds.
*/
PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_to_solve);
/*!
* \brief Combine the information into an array of (in)equalities.
* \param variables The variables in \p bounds.
* It is used to determine the iteration order to avoid indeterministic results.
* \param bounds grouped boundary of the variables.
* \param relations other relations.
*/
ffi::Array<PrimExpr> AsConditions(const ffi::Array<Var>& variables,
const ffi::Map<Var, IntGroupBounds>& bounds,
const ffi::Array<PrimExpr>& relations);
/*!
* \brief Solve linear inequalities and infer the range of each variable.
* \param system_to_solve the variables to solve, their ranges, and a list of inequalities.
* \return The result ranges for each variables.
* The returned IntConstraints(variables, ranges, relations) contains,
* 1. variables - the variables that have been solved.
* 2. ranges - the best range of each variable.
* 3. relations - constraints that cannot be transformed to
* Range will be stored in relations.
*/
IntConstraints SolveInequalitiesToRange(const IntConstraints& system_to_solve);
/*!
* \brief Solve linear inequalities and deskew the ranges towards zero.
* \param system_to_solve the variables to solve, their ranges, and a list of inequalities.
* \return A transform (src IntConstraints -> dst IntConstraints)
* from original variables to a set of new variables.
* The ranges of new variables always start from zero,
* their extents are solved from \p system_to_solve.
* src IntConstraints is the same as \p system_to_solve.
* dst IntConstraints(variables, ranges, relations) contains,
* 1. variables - the variables that have been solved.
* 2. ranges - the best range (start from zero) of each variable.
* 3. relations - constraints that cannot be transformed to
* Range will be stored in relations.
* Variable mapping can be obtained from
* IntConstraintsTransform.src_to_dst and IntConstraintsTransform.dst_to_src.
*/
IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& system_to_solve);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_INT_SOLVER_H_
include/tvm/arith/iter_affine_map.h
+425
View File
@@ -0,0 +1,425 @@
/*
* 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/iter_affine_map.h
* \brief Iterator quasi-affine mapping patterns.
*
* This file defines a collection of mapping patterns
* maps a collection of independent iterators to another
* collection of independent iterators.
*
* There are two main kinds of mapping patterns:
*
* - Fuse: fuse a collection of iterators into a single one
*
* domain(x0) = [0, 4), domain(x1) = [0, 3), domain(x2) = [0, 2)
* fuse(x0, x1, x2): y = x2 * 12 + x1 * 4 + x0
* domain(y) = [0, 24)
*
* - Split: split an iterator into multiple ones
*
* domain(x) = [0, 24)
* split(x, 3, 12): [y0, y1, y2] = [x % 3, (x % 12) / 3, x / 12]
* domain(y0) = [0, 3), domain(y1) = [0, 4), domain(y2) = [0, 2)
*
* We use the name "(quasi)affine" to be consistent with
* the terminology used in the polyhedral compilation.
* Notably, fuse is an affine transformation,
* while split corresponds to additional floordiv/mod operations
* that can appear in quasi-affine transformations.
*/
#ifndef TVM_ARITH_ITER_AFFINE_MAP_H_
#define TVM_ARITH_ITER_AFFINE_MAP_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/tirx/var.h>
namespace tvm {
namespace arith {
/*!
* \brief Base class of all iter map expressions.
*
* An IterMapExpr is a special expression to store
* the result of IterMapDetection.
* It should not appear in a legal TIR PrimFunc.
*/
class IterMapExprNode : public ExprNode {
public:
static constexpr const uint32_t _type_child_slots = 2;
TVM_FFI_DECLARE_OBJECT_INFO("arith.IterMapExpr", IterMapExprNode, ExprNode);
};
/*!
* \brief Managed reference to IterMapExprNode.
* \sa IterMapExprNode
*/
class IterMapExpr : public PrimExpr {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMapExpr, PrimExpr, IterMapExprNode);
static constexpr bool _type_container_is_exact = true;
};
/*!
* \brief Mark the source as an iterator in [0, extent).
*
* IterMark is used to mark source expression as a valid
* iterator to make future analysis easy.
*/
class IterMarkNode : public ffi::Object {
public:
/*!
* \brief The source expression, can either be
* a IterSumExpr or a Var.
*/
PrimExpr source;
/*!
* \brief The extent of the iteration.
*/
PrimExpr extent;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterMarkNode>()
.def_ro("source", &IterMarkNode::source)
.def_ro("extent", &IterMarkNode::extent);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMark", IterMarkNode, ffi::Object);
};
/*!
* \brief Managed reference to IterMarkExprNode.
* \sa IterMarkExprNode
*/
class IterMark : public ffi::ObjectRef {
public:
/*!
* \brief constructor.
* \param source The source expression.
* \param extent The extent of the iterator.
*/
TVM_DLL IterMark(PrimExpr source, PrimExpr extent);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMark, ffi::ObjectRef, IterMarkNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterMarkNode);
};
/*!
* \brief Split of an iterator.
*
* result = floormod(floordiv(source, lower_factor), extent) * scale
*/
class IterSplitExprNode : public IterMapExprNode {
public:
/*! \brief The source marked iterator. */
IterMark source;
/*! \brief The lower factor to split the source. */
PrimExpr lower_factor;
/*! \brief The extent of the split. */
PrimExpr extent;
/*! \brief Additional scale. */
PrimExpr scale;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterSplitExprNode>()
.def_ro("source", &IterSplitExprNode::source)
.def_ro("lower_factor", &IterSplitExprNode::lower_factor)
.def_ro("extent", &IterSplitExprNode::extent)
.def_ro("scale", &IterSplitExprNode::scale);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSplitExpr", IterSplitExprNode, IterMapExprNode);
};
/*!
* \brief Managed reference to IterSplitExprNode.
* \sa IterSplitExprNode
*/
class IterSplitExpr : public IterMapExpr {
public:
/*!
* \brief constructor from just source.
* \param source The source expression.
*/
TVM_DLL explicit IterSplitExpr(IterMark source);
/*!
* \brief constructor from just source.
* \param source The source expression.
* \param scale The additional scaling factor.
*/
TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr scale);
/*!
* \brief constructor
* \param source The source expression.
* \param lower_factor The lower factor to split the source.
* \param extent The extent of the split.
* \param scale The additional scaling factor.
*/
TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent,
PrimExpr scale);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSplitExpr, IterMapExpr, IterSplitExprNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSplitExprNode);
};
/*!
* \brief Fuse multiple iterators by summing them with scaling.
*
* result = sum(args) + base
*/
class IterSumExprNode : public IterMapExprNode {
public:
/*! \brief The args to the sum. */
ffi::Array<IterSplitExpr> args;
/*! \brief The base offset. */
PrimExpr base;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterSumExprNode>()
.def_ro("args", &IterSumExprNode::args)
.def_ro("base", &IterSumExprNode::base);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSumExpr", IterSumExprNode, IterMapExprNode);
};
/*!
* \brief Managed reference to IterSumExprNode.
* \sa IterSumExprNode
*/
class IterSumExpr : public IterMapExpr {
public:
/*!
* \brief constructor.
* \param args The args to the sum.
* \param base The base offset.
*/
TVM_DLL IterSumExpr(ffi::Array<IterSplitExpr> args, PrimExpr base);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSumExpr, IterMapExpr, IterSumExprNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSumExprNode);
};
} // namespace arith
namespace ffi {
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, arith::IterSplitExprNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, arith::IterSumExprNode> = true;
} // namespace ffi
namespace arith {
/*! \brief Mapping level for iterators. */
enum IterMapLevel {
// Require the mapping to be bijective.
Bijective = 0,
// Require the mapping to be surjective.
Surjective = 1,
// No mapping safety check.
NoCheck = 3
};
/*!
* \brief Result of DetectIterMap.
*/
class IterMapResultNode : public ffi::Object {
public:
// The detected pattern if a match exists.
ffi::Array<IterSumExpr> indices;
// Any errors that occurred while converting the input indices. If
// the array is empty, the conversion was successful.
ffi::Array<ffi::String> errors;
/*! \brief Boolean expression indicating if a specific value w
*
* `padding_predicate` evaluates to true for a set of indices that
* are outside the bounds of the provided index iterators, but
* inside the bounds of the returned index iterators. This
* expression is in terms of the variables provided in
* `input_iters`.
*/
PrimExpr padding_predicate;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterMapResultNode>()
.def_ro("indices", &IterMapResultNode::indices)
.def_ro("errors", &IterMapResultNode::errors)
.def_ro("padding_predicate", &IterMapResultNode::padding_predicate);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMapResult", IterMapResultNode, ffi::Object);
};
/*!
* \brief Managed reference to IterMapResultNode.
* \sa IterMapResultNode
*/
class IterMapResult : public ffi::ObjectRef {
public:
// constructor
IterMapResult() { data_ = ffi::make_object<IterMapResultNode>(); }
/*! \return mutable pointers to the node. */
IterMapResultNode* operator->() const { return static_cast<IterMapResultNode*>(get_mutable()); }
};
/*!
* \brief Detect if indices can be written as
* [y_0 + c_0, y_1 + c_1, ..., y_n + c_n]
*
* Here y = some-quasi-affine-iter-map(input_iters)
* and c are symbolic constants.
*
* We also requires that y_i and y_j to be independent for i != j.
*
* For returned value rv, the following is always true:
* - rv[i]->args.size() <=1: only one iterator per element.
*
* \param indices The indices to detect pattern for.
* \param input_iters Map from variable to iterator's range.
* \param predicate The predicate constraints on the input iterators
* \param check_level The iter mapping checking level.
* \param analyzer Analyzer used to get context information.
* \param simplify_trivial_iterators If true, iterators with extent of
* 1 will be replaced with a constant value.
*
* \return The detected iteration result.
* The return object's .indices is empty on failure.
*/
IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
const ffi::Map<Var, Range>& input_iters, const PrimExpr& predicate,
IterMapLevel check_level, const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
/*!
* \brief Use IterVarMap detector to rewrite and simplify the indices
*
* \param indices The indices to detect pattern for.
* \param input_iters Map from variable to iterator's range.
* \param input_pred The predicate constraints on the input iterators
* \param check_level The iter mapping checking level.
* \param analyzer Analyzer used to get context information.
* \param simplify_trivial_iterators If true, iterators with unit extents are simplified
* \return The indices after rewrite
*/
ffi::Array<PrimExpr> IterMapSimplify(const ffi::Array<PrimExpr>& indices,
const ffi::Map<Var, Range>& input_iters,
const PrimExpr& input_pred, IterMapLevel check_level,
const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
/*!
* \brief Apply the inverse of the affine transformation to the outputs.
*
* Similar to the back-propagation, starting from the outputs, it visits the DAG of the expressions
* in reverse topology order and applies the inverse of the affine transformation until it reaches
* the input. The affine iter map is required to be bijective.
*
* For example, iter_map = [l0 // 16, l0 % 16], outputs = [output_0, output_1],
* the affine transformation specified by `iter_map` will be applied to `outputs` and the result
* will be {l0: ((output_0*16) + output_1)}.
*
* The range of `outputs` should be the same as the output range of the affine transmation.
*
* \sa DetectIterMap
*
* \param iter_map The bijective affine iter map.
* \param outputs The outputs of the affine transformation.
*
* \return The map from the input to the transformed result.
*/
ffi::Map<Var, PrimExpr> InverseAffineIterMap(const ffi::Array<IterSumExpr>& iter_map,
const ffi::Array<PrimExpr> outputs);
/*!
* \brief Detect if bindings can be written as
* [a_0*e_0 + b_0 + c_0, a_1*e_1 + b_1, ..., a_n*e_n + b_n]
*
* where a = some-quasi-affine-iter-map(input_iters set_minus sub_iters)
* b = some-quasi-affine-iter-map(sub_iters)
* c is constant symbols
* e is the extent of b
*
* For example, z*12 + y*3 + x + c = (z*4+y)*3 + x, if sub_iters={x}
*
* \param bindings The input bindings
* \param input_iters Map from variable to iterator's range.
* \param sub_iters Iterators of subspace.
* \param predicate The predicate constraints on the input iterators
* \param check_level The iter mapping checking level.
* \param analyzer Analyzer used to get context information.
* \param simplify_trivial_iterators If true, iterators with extent of
* 1 will be replaced with a constant value.
*
* \return The result list has length len(bindings) + 1
[0, len(bindings)): The iter map matching result. The inner list is of length 2.
The first expr is the basis of the quotient space.
The second expr is the basis of the subspace.
len(bindings): the predicate of outer space and inner space
Empty array if no match can be found.
*/
ffi::Array<ffi::Array<IterMark>> SubspaceDivide(const ffi::Array<PrimExpr>& bindings,
const ffi::Map<Var, Range>& input_iters,
const ffi::Array<Var>& sub_iters,
const PrimExpr& predicate, IterMapLevel check_level,
const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
/*!
* \brief Given an expression that may contain IterMapExpr, transform it to normal PrimExpr.
* \param expr The input expression, which may contain IterMapExpr.
* \return The corresponding normal PrimExpr.
*/
PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr);
/*!
* \brief Rewrite index as IterSumExpr
*
* ((i0 // b0) % a0) * s0 + ((i0 // b1) % a1) * s1 ... + base
*
* The iterators are ordered such that s0 > s1 ...
* if we can prove the relation.
*
* Note that base may contain expressions that cannot be detected
* as the right pattern.
*
* \param index The input index
* \param input_iters The input iterators.
* \param analyzer The input analyzer.
* \note This function is useful to detect iterator stride patterns.
*/
IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<Var, Range>& input_iters,
const arith::Analyzer& analyzer);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_ITER_AFFINE_MAP_H_
include/tvm/arith/pattern.h
+54
View File
@@ -0,0 +1,54 @@
/*
* 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/pattern.h
* \brief Expression pattern detectors.
*/
#ifndef TVM_ARITH_PATTERN_H_
#define TVM_ARITH_PATTERN_H_
#include <tvm/ir/expr.h>
#include <tvm/tirx/expr.h>
namespace tvm {
namespace arith {
/*!
* \brief Detect if e can be rewritten as e = sum_{i=0}^{n-1} var[i] * coeff[i] + coeff[n]
* Where coeff[i] and base are invariant of var[j] for all i and j.
*
* \param e The expression to be detected.
* \param vars List of variables to be used in detection.
* \return [coeff[i]] if it is possible, empty array if it is not.
*/
ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<tirx::Var>& vars);
/*!
* \brief Detect if expression corresponds to clip bound of the vars
*
* \param e The expression to be detected.
* \param vars List of variables to be used in detection.
* \return concat([min_value[i], max_value[i]]), None is returned if there is no min or max value
* return empty if the e does not match the pattern.
*/
ffi::Array<PrimExpr> DetectClipBound(const PrimExpr& e, const ffi::Array<tirx::Var>& vars);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_PATTERN_H_
include/tvm/ir/attr_registry_map.h
+144
View File
@@ -0,0 +1,144 @@
/*
* 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/ir/attr_registry_map.h
* \brief Attribute map used in registry.
*/
#ifndef TVM_IR_ATTR_REGISTRY_MAP_H_
#define TVM_IR_ATTR_REGISTRY_MAP_H_
#include <tvm/ffi/string.h>
#include <utility>
#include <vector>
namespace tvm {
/*!
* \brief Generic attribute map.
* \tparam KeyType the type of the key.
*/
template <typename KeyType>
class AttrRegistryMapContainerMap {
public:
/*!
* \brief Check if the map has key.
* \param key The key to the map
* \return 1 if key is contained in map, 0 otherwise.
*/
int count(const KeyType& key) const {
if (key.defined()) {
const uint32_t idx = key->AttrRegistryIndex();
return idx < data_.size() ? (data_[idx].second != 0) : 0;
} else {
return 0;
}
}
/*!
* \brief get the corresponding value element at key.
* \param key The key to the map
* \return the const reference to the content value.
*/
const ffi::Any& operator[](const KeyType& key) const {
TVM_FFI_ICHECK(key.defined());
const uint32_t idx = key->AttrRegistryIndex();
TVM_FFI_ICHECK(idx < data_.size() && data_[idx].second != 0)
<< "Attribute " << attr_name_ << " has not been registered for " << key->name;
return data_[idx].first;
}
/*!
* \brief get the corresponding value element at key with default value.
* \param key The key to the map
* \param def_value The default value when the key does not exist.
* \return the const reference to the content value.
* \tparam ValueType The content value type.
*/
template <typename ValueType>
ValueType get(const KeyType& key, ValueType def_value) const {
TVM_FFI_ICHECK(key.defined());
const uint32_t idx = key->AttrRegistryIndex();
if (idx < data_.size() && data_[idx].second != 0) {
if constexpr (std::is_same_v<ValueType, ffi::Any>) {
return data_[idx].first;
} else {
return data_[idx].first.template cast<ValueType>();
}
} else {
return def_value;
}
}
private:
/*! \brief The name of the attr field */
ffi::String attr_name_;
/*! \brief The internal data. */
std::vector<std::pair<ffi::Any, int>> data_;
/*! \brief The constructor */
AttrRegistryMapContainerMap() = default;
template <typename, typename>
friend class AttrRegistry;
friend class OpRegEntry;
};
/*!
* \brief ffi::Map<Key, ValueType> used to store meta-data.
* \tparam KeyType The type of the key
* \tparam ValueType The type of the value stored in map.
*/
template <typename KeyType, typename ValueType>
class AttrRegistryMap {
public:
/*!
* \brief constructor
* \param map The internal map.
*/
explicit AttrRegistryMap(const AttrRegistryMapContainerMap<KeyType>& map) : map_(map) {}
/*!
* \brief Check if the map has op as key.
* \param key The key to the map
* \return 1 if op is contained in map, 0 otherwise.
*/
int count(const KeyType& key) const { return map_.count(key); }
/*!
* \brief get the corresponding value element at key.
* \param key The key to the map
* \return the const reference to the content value.
*/
ValueType operator[](const KeyType& key) const {
if constexpr (std::is_same_v<ValueType, ffi::Any>) {
return map_[key];
} else {
return map_[key].template cast<ValueType>();
}
}
/*!
* \brief get the corresponding value element at key with default value.
* \param key The key to the map
* \param def_value The default value when the key does not exist.
* \return the const reference to the content value.
*/
ValueType get(const KeyType& key, ValueType def_value) const { return map_.get(key, def_value); }
protected:
/*! \brief The internal map field */
const AttrRegistryMapContainerMap<KeyType>& map_;
};
} // namespace tvm
#endif // TVM_IR_ATTR_REGISTRY_MAP_H_
include/tvm/ir/attrs.h
+325
View File
@@ -0,0 +1,325 @@
/*
* 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/ir/attrs.h
* \brief Helpers for attribute objects.
*
* This module enables declaration of named attributes
* which support default value setup and bound checking.
*
* \sa AttrsNode
*/
#ifndef TVM_IR_ATTRS_H_
#define TVM_IR_ATTRS_H_
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
namespace tvm {
/*!
* \brief Base class of all attribute class
* \sa Attrs
*/
class AttrsNode : public ffi::Object {
public:
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("ir.Attrs", AttrsNode, ffi::Object);
};
/*!
* \brief Managed reference to AttrsNode.
* \sa AttrsNode
*/
class Attrs : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Attrs, ffi::ObjectRef, AttrsNode);
};
/*!
* \brief Specialized attribute type that is backed by a map.
* The DictAttrsNode implements the Attrs behavior,
* its fields are directly accessible via object.field_name
* like other normal nodes.
*/
class DictAttrsNode : public AttrsNode {
public:
/*! \brief internal attrs map */
ffi::Map<ffi::String, ffi::Any> dict;
static void RegisterReflection() {
namespace rfl = ffi::reflection;
rfl::ObjectDef<DictAttrsNode>().def_ro("__dict__", &DictAttrsNode::dict);
}
// type info
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DictAttrs", DictAttrsNode, AttrsNode);
};
/*!
* \brief Managed reference to DictAttrsNode
* \sa DictAttrsNode.
*
* \note DictAttrs is NOTNULLABLE: every instance must hold a backing
* DictAttrsNode. The class enforces this end-to-end by:
* - the default constructor (no args) allocating an empty backing,
* - the copy/move ctors and assignments leaving the moved-from
* instance in a defined-but-empty state rather than null,
* - the FFI type traits rejecting None at deserialization boundaries
* (since `_type_is_nullable == false`), and
* - the FFI lambda for ``ir.IRModule`` explicitly normalizing a
* missing/None attrs argument to ``DictAttrs()`` before forwarding
* to the C++ constructor.
* Callers (including third-party code via templates like ``WithAttr``)
* can therefore rely on ``attrs->dict`` being safe to dereference
* without a ``.defined()`` guard.
*/
class DictAttrs : public Attrs {
public:
/*!
* \brief Construct a DictAttrs backed by DictAttrsNode.
*
* The no-argument form constructs an empty (but always defined) DictAttrs.
* \param dict The attributes.
*/
explicit DictAttrs(ffi::Map<ffi::String, Any> dict = {}) {
ffi::ObjectPtr<DictAttrsNode> n = ffi::make_object<DictAttrsNode>();
n->dict = std::move(dict);
data_ = std::move(n);
}
/*!
* \brief Move constructor that leaves the source in a defined-but-empty
* state rather than null, preserving the NOTNULLABLE invariant
* even after `std::move`.
*/
DictAttrs(DictAttrs&& other) noexcept : Attrs(ffi::UnsafeInit{}) {
data_ = std::move(other.data_);
other.data_ = ffi::make_object<DictAttrsNode>();
}
/*!
* \brief Move assignment that leaves the source in a defined-but-empty
* state rather than null, preserving the NOTNULLABLE invariant
* even after `std::move`.
*/
DictAttrs& operator=(DictAttrs&& other) noexcept {
if (this != &other) {
data_ = std::move(other.data_);
other.data_ = ffi::make_object<DictAttrsNode>();
}
return *this;
}
// Explicit copy ctor/assign defaults. Declaring the move members above
// would otherwise suppress the implicit copy members.
DictAttrs(const DictAttrs& other) = default;
DictAttrs& operator=(const DictAttrs& other) = default;
// Utils for accessing attributes
/*!
* \brief Get a function attribute.
*
* \param attr_key The attribute key.
* \param default_value The default value if the key does not exist, defaults to nullptr.
*
* \return The result
*
* \tparam TOBjectRef the expected object type.
* \throw Error if the key exists but the value does not match TObjectRef
*
* \code
*
* void GetAttrExample(const BaseFunc& f) {
* auto value = f->attrs.GetAttr<int64_t>("AttrKey", 0);
* }
*
* \endcode
*/
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetAttr(
const std::string& attr_key,
ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
const DictAttrsNode* node = get();
auto it = node->dict.find(attr_key);
if (it != node->dict.end()) {
return (*it).second.cast<TObjectRef>();
} else {
return default_value;
}
}
// variant that uses TObjectRef to enable implicit conversion to default value.
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
}
/*!
* \brief Check whether the function has an non-zero integer attr.
*
* This function can be used to check whether an optional
* attribute mark(e.g. inline) exists.
*
* \param attr_key The key to the attribute.
* \return The check result.
*
* \code
*
* void HasNonzeroAttrExample(const BaseFunc& f) {
* if (f->HasNonzeroAttr(attr::kInline)) {
* // inline the function.
* }
* }
*
* \endcode
*/
bool HasNonzeroAttr(const std::string& attr_key) const {
return GetAttr<int64_t>(attr_key, 0).value_or(0) != 0;
}
// Inline-expand TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE here, minus
// the default copy/move it normally injects (we define our own move members
// above so the moved-from instance stays defined-but-empty).
explicit DictAttrs(::tvm::ffi::UnsafeInit tag) : Attrs(tag) {}
using __PtrType =
std::conditional_t<DictAttrsNode::_type_mutable, DictAttrsNode*, const DictAttrsNode*>;
__PtrType operator->() const { return static_cast<__PtrType>(data_.get()); }
__PtrType get() const { return static_cast<__PtrType>(data_.get()); }
static constexpr bool _type_is_nullable = false;
using ContainerType = DictAttrsNode;
TVM_DEFINE_OBJECT_REF_COW_METHOD(DictAttrsNode);
};
/*!
* \brief Copy the function or module, but overrides
* the attribute value key with the value.
*
* \param input The thing to annotate (BaseFunc or IRModule)
* \param attr_key The attribute key.
* \param attr_value The value attribute value.
*
* \tparam TFunc The corresponding function or module type.
*
* \returns The new function or module with updated attributes.
*
* \note This function performs copy on write optimization for func and module.
* If we move a uniquely referenced func or module into WithAttr,
* then no additional copy will be performed.
*
* This is also why we make it as a function instead of a member function
* and why we pass by value in the first argument.
*
* \code
*
* // Recommended way to trigger copy on write
* func = WithAttr(std::move(func), "key1", value1);
* func = WithAttr(std::move(func), "key2", value2);
*
* \endcode
*/
template <typename TFunc>
inline TFunc WithAttr(TFunc input, const std::string& attr_key, Any attr_value) {
using TNode = typename TFunc::ContainerType;
static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
TNode* node = input.CopyOnWrite();
// node->attrs is NOTNULLABLE by contract, but defend against a caller
// that left a moved-from DictAttrs in place by re-initializing here.
if (!node->attrs.defined()) node->attrs = DictAttrs();
node->attrs.CopyOnWrite()->dict.Set(attr_key, std::move(attr_value));
return input;
}
/*!
* \brief Copy the function or module, but overrides the attributes with the entries from \p attrs.
*
* \param input The thing to annotate (BaseFunc or IRModule)
* \param attrs Key/values attributes to add to \p input.
*
* \tparam TFunc The corresponding function or module type.
*
* \returns The new function or module with updated attributes.
*/
template <typename TFunc>
inline TFunc WithAttrs(TFunc input, ffi::Map<ffi::String, Any> attrs) {
using TNode = typename TFunc::ContainerType;
static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
if (attrs.empty()) return input;
TNode* node = input.CopyOnWrite();
// node->attrs is NOTNULLABLE by contract, but defend against a caller
// that left a moved-from DictAttrs in place by re-initializing here.
if (!node->attrs.defined()) node->attrs = DictAttrs();
auto* dict_node = node->attrs.CopyOnWrite();
for (const auto& [k, v] : attrs) {
dict_node->dict.Set(k, v);
}
return input;
}
/*!
* \brief Copy the function or module, but removes the specified
* attribute.
*
* \param input The thing to annotate (BaseFunc or IRModule)
* \param attr_key The attribute key.
*
* \tparam TFunc The corresponding function or module type.
*
* \returns The new function or module with removed attribute.
*
* \note This function performs copy on write optimization for func and module.
* If we move a uniquely referenced func or module into WithoutAttr,
* then no additional copy will be performed.
*
* This is also why we make it as a function instead of a member function
* and why we pass by value in the first argument.
*
* \code
*
* // Recommended way to trigger copy on write
* func = WithoutAttr(std::move(func), "key1");
* func = WithoutAttr(std::move(func), "key2");
*
* \endcode
*/
template <typename TFunc>
inline TFunc WithoutAttr(TFunc input, const std::string& attr_key) {
using TNode = typename TFunc::ContainerType;
static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
TNode* node = input.CopyOnWrite();
// node->attrs is NOTNULLABLE by contract, but defend against a caller
// that left a moved-from DictAttrs in place; nothing to erase from an
// empty dict.
if (!node->attrs.defined()) {
node->attrs = DictAttrs();
return input;
}
node->attrs.CopyOnWrite()->dict.erase(attr_key);
return input;
}
} // namespace tvm
#endif // TVM_IR_ATTRS_H_
include/tvm/ir/base_expr.h
+521
View File
@@ -0,0 +1,521 @@
/*
* 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/ir/base_expr.h
* \brief Base expression and primitive type nodes.
*/
#ifndef TVM_IR_BASE_EXPR_H_
#define TVM_IR_BASE_EXPR_H_
#include <tvm/ffi/cast.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/source_map.h>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <type_traits>
namespace tvm {
/*!
* \brief Type is the base type of all types.
*
* TVM's type system contains following subclasses:
*
* - PrimType: type of primitive type values used in the low-level IR.
* - FuncType: type of a function.
* - TensorType: type of certain Tensor values in the expression.
*
* There are also advanced types to support generic(polymorphic types).
* \sa Type
*/
class TypeNode : public ffi::Object {
public:
/*!
* \brief Span that points to the original source code.
* Reserved debug information.
*/
mutable Span span;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
// span do not participate in structural equal and hash.
refl::ObjectDef<TypeNode>().def_ro("span", &TypeNode::span, refl::DefaultValue(Span()),
refl::AttachFieldFlag::SEqHashIgnore());
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
static constexpr const uint32_t _type_child_slots = 14;
TVM_FFI_DECLARE_OBJECT_INFO("ir.Type", TypeNode, ffi::Object);
};
/*!
* \brief Managed reference to TypeNode.
* \sa TypeNode
*/
class Type : public ffi::ObjectRef {
public:
/*! \brief Sentinel for a type that has not been populated yet. */
TVM_DLL static Type Missing();
/*! \return whether this is the missing-type sentinel. */
TVM_DLL bool IsMissing() const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef, TypeNode);
};
/*!
* \brief Primitive data types used in the low-level IR.
*
* PrimType represents primitive POD values and the void sentinel.
*
* \sa PrimType
*/
class PrimTypeNode final : public TypeNode {
public:
/*!
* \brief The raw DLPack dtype represented by this primitive type.
*/
DLDataType dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PrimTypeNode>().def_ro("dtype", &PrimTypeNode::dtype);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PrimType", PrimTypeNode, TypeNode);
};
/*
* \brief Managed reference to PrimTypeNode.
* \sa PrimTypeNode
*/
class PrimType final : public Type {
public:
/*!
* \brief Construct from a raw DLPack dtype.
* \param dtype The corresponding DLPack dtype.
*/
TVM_DLL explicit PrimType(DLDataType dtype);
/*!
* \brief Construct from DLPack dtype fields.
* \param code The DLPack dtype code.
* \param bits The scalar bit width.
* \param lanes The fixed lane count.
*/
TVM_DLL PrimType(DLDataTypeCode code, int bits, int lanes = 1);
/*! \brief Construct a signed integer type with fixed lanes. */
TVM_DLL static PrimType Int(int bits, int lanes = 1);
/*! \brief Construct an unsigned integer type with fixed lanes. */
TVM_DLL static PrimType UInt(int bits, int lanes = 1);
/*! \brief Construct a floating-point type with fixed lanes. */
TVM_DLL static PrimType Float(int bits, int lanes = 1);
/*! \brief Construct a bfloat type with fixed lanes. */
TVM_DLL static PrimType BFloat(int bits, int lanes = 1);
/*! \brief Construct a boolean type with fixed lanes. */
TVM_DLL static PrimType Bool(int lanes = 1);
/*! \brief Construct the void sentinel type, encoded as handle(0, 0). */
TVM_DLL static PrimType Void();
/*!
* \brief Construct a scalable vector type.
* \param code The DLPack dtype code.
* \param bits The scalar bit width.
* \param lanes The positive vscale factor to encode in the DLPack lane field.
*/
TVM_DLL static PrimType ScalableVector(DLDataTypeCode code, int bits, int lanes);
/*! \return The DLPack dtype code. */
TVM_FFI_INLINE DLDataTypeCode code() const {
return static_cast<DLDataTypeCode>(static_cast<int>(get()->dtype.code));
}
/*! \return The scalar bit width. */
TVM_FFI_INLINE int32_t bits() const { return get()->dtype.bits; }
/*!
* \return The fixed lane count.
* \note Throws on scalable vector types, where the encoded lane field stores a vscale factor.
*/
TVM_FFI_INLINE int32_t lanes() const {
int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
if (TVM_FFI_PREDICT_FALSE(encoded_lanes < 0)) {
TVM_FFI_THROW(InternalError)
<< "Can't fetch the lanes of a scalable vector at a compile time.";
}
return encoded_lanes;
}
/*!
* \brief Check the scalar element code and bit width.
* \note Lane count and scalable-vector encoding are intentionally ignored.
*/
TVM_FFI_INLINE bool MatchesElementType(DLDataTypeCode code, int bits) const {
DLDataType dtype = get()->dtype;
return dtype.code == static_cast<uint8_t>(code) && dtype.bits == bits;
}
/*!
* \brief Check whether the dtype code matches any of the provided DLPack codes.
* \note Bit width and lanes are intentionally ignored.
*/
template <typename... Codes>
TVM_FFI_INLINE bool MatchesCode(Codes... codes) const {
uint8_t dtype_code = get()->dtype.code;
return ((dtype_code == static_cast<uint8_t>(codes)) || ...);
}
/*! \brief Whether this type is a scalar, excluding fixed and scalable vectors. */
TVM_FFI_INLINE bool IsScalar() const {
int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
return encoded_lanes == 1;
}
/*! \brief Whether this type is the void sentinel `handle(0, 0)`. */
TVM_FFI_INLINE bool IsVoid() const {
DLDataType dtype = get()->dtype;
return dtype.code == static_cast<uint8_t>(DLDataTypeCode::kDLOpaqueHandle) && dtype.bits == 0 &&
static_cast<int16_t>(dtype.lanes) == 0;
}
/*! \brief Whether this type is a scalable vector. */
TVM_FFI_INLINE bool IsScalableVector() const {
return static_cast<int16_t>(get()->dtype.lanes) < -1;
}
/*! \brief Whether this type is a fixed-length vector. */
TVM_FFI_INLINE bool IsFixedLengthVector() const {
return static_cast<int16_t>(get()->dtype.lanes) > 1;
}
/*!
* \brief Return the number of bytes needed to store one value of this type.
*
* This uses the same packed sub-byte dtype sizing rule as runtime tensors.
* Scalable vector types have no compile-time storage size and are rejected.
*/
TVM_FFI_INLINE size_t StorageBytes() const {
DLDataType dtype = get()->dtype;
int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
if (TVM_FFI_PREDICT_FALSE(encoded_lanes < 0)) {
TVM_FFI_THROW(InternalError)
<< "Cannot compute compile-time storage bytes for non-fixed vector type " << dtype;
}
return static_cast<size_t>(
(static_cast<uint64_t>(dtype.bits) * static_cast<uint64_t>(dtype.lanes) + 7) / 8);
}
/*! \brief Return the same type with a different dtype code, preserving bits and lanes. */
TVM_FFI_INLINE PrimType WithCode(DLDataTypeCode code) const {
DLDataType dtype = get()->dtype;
int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
if (encoded_lanes < -1) {
return ScalableVector(code, dtype.bits, -encoded_lanes);
}
return PrimType(code, dtype.bits, encoded_lanes);
}
/*! \brief Return the same type with a different scalar bit width, preserving code and lanes. */
TVM_FFI_INLINE PrimType WithBits(int bits) const {
DLDataType dtype = get()->dtype;
int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
if (encoded_lanes < -1) {
return ScalableVector(this->code(), bits, -encoded_lanes);
}
return PrimType(this->code(), bits, encoded_lanes);
}
/*! \brief Return the same scalar element type with a fixed lane count. */
TVM_FFI_INLINE PrimType WithLanes(int lanes) const {
return PrimType(this->code(), this->bits(), lanes);
}
/*! \return The vscale factor encoded in a scalable vector type. */
TVM_FFI_INLINE int32_t VScaleFactor() const {
int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
if (encoded_lanes >= -1) {
TVM_FFI_THROW(InternalError) << "A fixed length vector doesn't have a vscale factor.";
}
return -encoded_lanes;
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimType, Type, PrimTypeNode);
};
inline bool operator==(const PrimType& lhs, const PrimType& rhs) {
return lhs->dtype == rhs->dtype;
}
inline bool operator!=(const PrimType& lhs, const PrimType& rhs) { return !(lhs == rhs); }
/*!
* \brief Base type of all the expressions.
* \sa Expr
*/
class ExprNode : public ffi::Object {
public:
/*!
* \brief Span that points to the original source code.
* Reserved debug information.
*/
mutable Span span;
/*!
* \brief The deduced or annotated type of the expression.
*
* Type::Missing() denotes type information that will be populated by
* later analysis passes instead of expression constructors.
*/
mutable Type ty = Type::Missing();
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
// span does not participate in structural equal and hash.
refl::ObjectDef<ExprNode>()
.def_ro("span", &ExprNode::span, refl::DefaultValue(Span()),
refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("ty", &ExprNode::ty, refl::DefaultValue(Type::Missing()));
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
static constexpr const uint32_t _type_child_slots = 64;
TVM_FFI_DECLARE_OBJECT_INFO("ir.Expr", ExprNode, ffi::Object);
};
/*!
* \brief Managed reference to ExprNode.
* \sa ExprNode
*/
class Expr : public ffi::ObjectRef {
public:
// Expressions do not implicitly compare by object identity or address. Callers must name
// whether they intend object identity, structural equality, or primitive symbolic comparison.
bool operator==(const Expr& other) const = delete;
bool operator!=(const Expr& other) const = delete;
bool operator<(const Expr& other) const = delete;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Expr, ffi::ObjectRef, ExprNode);
};
class Call;
/*!
* \brief Typed reference/view over an expression whose result type is a
* specific Type subtype.
* \tparam ExpectedType The expected expression result type.
*/
template <typename ExpectedType>
class TypedExpr : public Expr {
public:
/*! \return the typed result of this expression. */
ExpectedType ty() const {
const auto* node = get();
TVM_FFI_DCHECK(node != nullptr);
const auto* ty_node = node->ExprNode::ty.template as<typename ExpectedType::ContainerType>();
TVM_FFI_DCHECK(ty_node != nullptr);
return ffi::GetRef<ExpectedType>(ty_node);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TypedExpr, Expr, ExprNode);
static constexpr bool _type_container_is_exact = false;
};
/*!
* \brief Typed reference/view over any Expr whose `ExprNode::ty` is PrimType.
*
* PrimExpr is a type category rather than a dedicated runtime node category.
* It can contain intrinsic primitive nodes such as IntImmNode and FloatImmNode,
* or a general ExprNode such as CallNode, when that expression's `ty` field is
* a PrimType. This keeps primitive-only APIs explicit while allowing shared
* Expr nodes for cross-dialect values with richer result types when needed.
*/
class PrimExpr : public TypedExpr<PrimType> {
public:
using TypedExpr<PrimType>::ty;
/*!
* \brief Construct from a call after checking that its result type is
* PrimType.
* \param call The call to view as a primitive expression.
*/
TVM_DLL PrimExpr(Call call); // NOLINT(*)
/*!
* \brief construct from integer.
* \param value The value to be constructed.
*/
TVM_DLL PrimExpr(int32_t value); // NOLINT(*)
/*!
* \brief construct from float.
* \param value The value to be constructed.
*/
TVM_DLL PrimExpr(float value); // NOLINT(*)
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExpr, TypedExpr<PrimType>, ExprNode);
static constexpr bool _type_container_is_exact = false;
/*!
* \brief construct from string to form a StringImm.
* \param value The value to be constructed.
*/
TVM_DLL static PrimExpr ConvertFallbackValue(ffi::String value); // NOLINT(*)
};
/*!
* \brief Base class for other IR constructs that can be converted to PrimExpr.
* This is useful for the FFI to convert the expressions to PrimExpr.
* \sa PrimExpr
*/
class PrimExprConvertibleNode : public ffi::Object {
public:
virtual ~PrimExprConvertibleNode() {}
virtual PrimExpr ToPrimExpr() const = 0;
TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertible", PrimExprConvertibleNode, ffi::Object);
};
/*!
* \brief Managed reference to PrimExprConvertibleNode.
* \sa PrimExprConvertibleNode
*/
class PrimExprConvertible : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExprConvertible, ffi::ObjectRef,
PrimExprConvertibleNode);
};
namespace ffi {
template <>
inline constexpr bool use_default_type_traits_v<PrimType> = false;
template <>
struct TypeTraits<PrimType> : public ObjectRefWithFallbackTraitsBase<PrimType, DLDataType> {
TVM_FFI_INLINE static PrimType ConvertFallbackValue(DLDataType dtype) { return PrimType(dtype); }
};
template <typename ExpectedType>
inline constexpr bool use_default_type_traits_v<TypedExpr<ExpectedType>> = false;
template <typename ExpectedType>
struct TypeTraits<TypedExpr<ExpectedType>>
: public ObjectRefTypeTraitsBase<TypedExpr<ExpectedType>> {
using Base = ObjectRefTypeTraitsBase<TypedExpr<ExpectedType>>;
using Base::CopyFromAnyViewAfterCheck;
using Base::CopyToAnyView;
using Base::GetMismatchTypeInfo;
using Base::MoveFromAnyAfterCheck;
using Base::MoveToAny;
using Base::TypeSchema;
using Base::TypeStr;
TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
if (src->type_index == TypeIndex::kTVMFFINone) {
return TypedExpr<ExpectedType>::_type_is_nullable;
}
if (src->type_index < TypeIndex::kTVMFFIStaticObjectBegin ||
!details::IsObjectInstance<ExprNode>(src->type_index)) {
return false;
}
const auto* expr = static_cast<const ExprNode*>(
details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(src->v_obj).get());
return details::AnyUnsafe::CheckAnyStrict<ExpectedType>(expr->ty);
}
TVM_FFI_INLINE static std::optional<TypedExpr<ExpectedType>> TryCastFromAnyView(
const TVMFFIAny* src) {
if (CheckAnyStrict(src)) {
if (src->type_index == TypeIndex::kTVMFFINone) {
return details::ObjectUnsafe::ObjectRefFromObjectPtr<TypedExpr<ExpectedType>>(nullptr);
}
return details::ObjectUnsafe::ObjectRefFromObjectPtr<TypedExpr<ExpectedType>>(
details::ObjectUnsafe::ObjectPtrFromUnowned<ExprNode>(src->v_obj));
}
return std::nullopt;
}
};
template <>
inline constexpr bool use_default_type_traits_v<PrimExpr> = false;
template <typename ObjectRefType, typename ExpectedType, typename... FallbackTypes>
struct TypedExprWithFallbackTraitsBase
: public ObjectRefWithFallbackTraitsBase<ObjectRefType, FallbackTypes...> {
using Base = ObjectRefWithFallbackTraitsBase<ObjectRefType, FallbackTypes...>;
TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
return TypeTraits<TypedExpr<ExpectedType>>::CheckAnyStrict(src);
}
TVM_FFI_INLINE static std::optional<ObjectRefType> TryCastFromAnyView(const TVMFFIAny* src) {
if (TypeTraits<TypedExpr<ExpectedType>>::TryCastFromAnyView(src)) {
return details::ObjectUnsafe::ObjectRefFromObjectPtr<ObjectRefType>(
details::ObjectUnsafe::ObjectPtrFromUnowned<ExprNode>(src->v_obj));
}
return Base::template TryFallbackTypes<FallbackTypes...>(src);
}
};
// define automatic conversion from bool, int64_t, double, ffi::String to PrimExpr
// These functions are declared early to avoid circular dependency
template <>
struct TypeTraits<PrimExpr>
: public TypedExprWithFallbackTraitsBase<PrimExpr, PrimType, StrictBool, int64_t, double,
ffi::String, PrimExprConvertible> {
using Base = TypedExprWithFallbackTraitsBase<PrimExpr, PrimType, StrictBool, int64_t, double,
ffi::String, PrimExprConvertible>;
using Base::CheckAnyStrict;
using Base::CopyFromAnyViewAfterCheck;
using Base::CopyToAnyView;
using Base::GetMismatchTypeInfo;
using Base::MoveFromAnyAfterCheck;
using Base::MoveToAny;
using Base::TryCastFromAnyView;
using Base::TypeSchema;
using Base::TypeStr;
TVM_DLL static PrimExpr ConvertFallbackValue(StrictBool value);
TVM_DLL static PrimExpr ConvertFallbackValue(int64_t value);
TVM_DLL static PrimExpr ConvertFallbackValue(double value);
TVM_FFI_INLINE static PrimExpr ConvertFallbackValue(ffi::String value) {
return PrimExpr::ConvertFallbackValue(value);
}
TVM_FFI_INLINE static PrimExpr ConvertFallbackValue(PrimExprConvertible value) {
return value->ToPrimExpr();
}
};
template <>
inline constexpr bool use_default_type_traits_v<Expr> = false;
// Allow generic Expr arguments to use the primitive-literal conversions
// already defined by PrimExpr.
template <>
struct TypeTraits<Expr> : public ObjectRefWithFallbackTraitsBase<Expr, PrimExpr> {
TVM_FFI_INLINE static Expr ConvertFallbackValue(PrimExpr value) { return value; }
};
} // namespace ffi
} // namespace tvm
#endif // TVM_IR_BASE_EXPR_H_
include/tvm/ir/config_schema.h
+242
View File
@@ -0,0 +1,242 @@
/*
* 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/ir/config_schema.h
* \brief Minimal schema for dynamic config canonicalization and validation.
*
* This utility is intended for dynamic map-like configs (e.g. Target options),
* where we still want type checking, optional defaulting, and canonicalization.
*/
#ifndef TVM_IR_CONFIG_SCHEMA_H_
#define TVM_IR_CONFIG_SCHEMA_H_
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <sstream>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>
namespace tvm {
namespace ir {
/*!
* \brief Dynamic config schema for map-like options.
*
* The schema supports:
* - Option declaration (`def_option<T>`)
* - Optional canonicalizer (`set_canonicalizer`)
* - Resolution (`Resolve`) that performs validation/defaulting, unknown-key policy,
* and canonicalization (last).
*/
class ConfigSchema {
public:
using ConfigMap = ffi::Map<ffi::String, ffi::Any>;
using Canonicalizer = ffi::TypedFunction<ConfigMap(ConfigMap)>;
/*! \brief Schema entry for one declared option. */
struct OptionEntry {
/*! \brief Option key. */
ffi::String key;
/*! \brief Type string for this option. */
ffi::String type_str;
/*! \brief Per-option validator/coercer (Any -> Any). */
ffi::TypedFunction<ffi::Any(ffi::Any)> validator;
/*! \brief Whether this option has a default value or factory. */
bool has_default = false;
/*! \brief Whether the default is produced by a factory function. */
bool default_from_factory = false;
/*! \brief Default value, or factory function () -> Any when default_from_factory is true. */
ffi::Any default_value;
};
/*!
* \brief Declare a typed option.
*
* Validation/coercion is implicitly generated from `T`.
* Additional optional traits may be supplied (e.g. `refl::DefaultValue`, `const char*` doc).
*
* \tparam T The canonical value type of this option.
* \tparam Traits Optional metadata/traits.
* \param key Option key.
* \param traits Optional traits.
* \return Reference to `*this` for chaining.
*/
template <typename T, typename... Traits>
ConfigSchema& def_option(const ffi::String& key, Traits&&... traits) {
std::string skey(key);
if (key_to_index_.count(skey)) {
TVM_FFI_THROW(ValueError) << "Duplicate config option key: '" << key << "'";
}
key_to_index_[skey] = options_.size();
options_.push_back(MakeEntry<T>(key, std::forward<Traits>(traits)...));
return *this;
}
/*! \brief Set whole-object canonicalizer. */
void set_canonicalizer(Canonicalizer f) { canonicalizer_ = std::move(f); }
/*! \brief Trait to set a custom validator for a config option. */
struct AttrValidator {
ffi::TypedFunction<ffi::Any(ffi::Any)> func;
explicit AttrValidator(ffi::TypedFunction<ffi::Any(ffi::Any)> f) : func(std::move(f)) {}
};
/*! \brief Set whether unknown keys trigger an error. */
void set_error_on_unknown(bool value) { error_on_unknown_ = value; }
/*!
* \brief Default/validate, then canonicalize a config object.
*
* Resolve flow:
* 1) Validate/coerce declared options in declaration order.
* 2) Materialize defaults and enforce required options.
* 3) Apply unknown-key policy.
* 4) Run canonicalizer as final step.
*
* \param config Input config object.
* \return Canonical validated config object.
* \throws ValueError/TypeError with option context.
*/
ConfigMap Resolve(ConfigMap config) const {
ConfigMap result;
// Step 1: validate/coerce and materialize options in declaration order
for (const auto& e : options_) {
auto it = config.find(e.key);
if (it != config.end()) {
result.Set(e.key, e.validator((*it).second));
} else if (e.has_default) {
if (e.default_from_factory) {
result.Set(e.key, e.default_value.cast<ffi::Function>()());
} else {
result.Set(e.key, e.default_value);
}
}
// else: missing non-required option, stays absent
}
// Step 2: unknown-key policy
if (error_on_unknown_) {
for (const auto& kv : config) {
if (!key_to_index_.count(std::string(kv.first))) {
std::ostringstream os;
os << "Unknown config option '" << kv.first << "'. Known options: ";
bool first = true;
for (const auto& e : options_) {
if (!first) os << ", ";
os << "'" << e.key << "'";
first = false;
}
TVM_FFI_THROW(ValueError) << os.str();
}
}
}
// Step 3: whole-object canonicalization (last)
if (canonicalizer_ != nullptr) {
result = canonicalizer_(result);
}
return result;
}
/*!
* \brief List declared options in declaration order.
* \return Const reference to the option entries vector.
*/
const std::vector<OptionEntry>& ListOptions() const { return options_; }
/*! \brief Check if an option with the given key exists. */
bool HasOption(const ffi::String& key) const { return key_to_index_.count(std::string(key)) > 0; }
private:
template <typename T>
static ffi::TypedFunction<ffi::Any(ffi::Any)> MakeValidator(const ffi::String& key) {
return ffi::TypedFunction<ffi::Any(ffi::Any)>([key](ffi::Any val) -> ffi::Any {
auto opt = val.try_cast<T>();
if (!opt.has_value()) {
TVM_FFI_THROW(TypeError) << "Option '" << key << "': expected type '"
<< ffi::TypeTraits<T>::TypeStr() << "' but got '"
<< val.GetTypeKey() << "'";
}
return ffi::Any(opt.value());
});
}
template <typename Trait>
static void ApplyTrait(OptionEntry* entry, ffi::reflection::FieldInfoBuilder* info,
Trait&& trait) {
using T = std::decay_t<Trait>;
if constexpr (std::is_same_v<T, AttrValidator>) {
entry->validator = std::move(trait.func);
} else if constexpr (std::is_base_of_v<ffi::reflection::InfoTrait, T>) {
trait.Apply(info);
} else if constexpr (std::is_same_v<T, const char*> || std::is_same_v<T, char*>) {
const char* doc = trait;
if (doc != nullptr && doc[0] != '\0') {
info->doc = TVMFFIByteArray{doc, std::char_traits<char>::length(doc)};
}
}
}
template <typename T, typename... Traits>
OptionEntry MakeEntry(const ffi::String& key, Traits&&... traits) {
OptionEntry e;
e.key = key;
e.type_str = ffi::String(ffi::TypeTraits<T>::TypeStr());
e.validator = MakeValidator<T>(key);
// Apply traits through a temporary FieldInfoBuilder so existing
// reflection traits (notably refl::DefaultValue) are reused unchanged.
ffi::reflection::FieldInfoBuilder info{};
info.flags = 0;
info.default_value_or_factory = ffi::AnyView(nullptr).CopyToTVMFFIAny();
info.doc = TVMFFIByteArray{nullptr, 0};
(ApplyTrait(&e, &info, std::forward<Traits>(traits)), ...);
if (info.flags & kTVMFFIFieldFlagBitMaskHasDefault) {
e.has_default = true;
e.default_from_factory = (info.flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory) != 0;
e.default_value = ffi::AnyView::CopyFromTVMFFIAny(info.default_value_or_factory);
// Release the extra ref created by CopyToTVMFFIAny in Apply
if (info.default_value_or_factory.type_index >= TVMFFITypeIndex::kTVMFFIStaticObjectBegin) {
ffi::details::ObjectUnsafe::DecRefObjectHandle(info.default_value_or_factory.v_obj);
}
}
return e;
}
/*! \brief Declared options in declaration order. */
std::vector<OptionEntry> options_;
/*! \brief Map from key string to index in options_. */
std::unordered_map<std::string, size_t> key_to_index_;
/*! \brief Optional whole-config canonicalizer. */
Canonicalizer canonicalizer_{nullptr};
/*! \brief Whether unknown keys trigger an error. */
bool error_on_unknown_ = true;
};
} // namespace ir
} // namespace tvm
#endif // TVM_IR_CONFIG_SCHEMA_H_
include/tvm/ir/cow.h
+70
View File
@@ -0,0 +1,70 @@
/*
* 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/ir/cow.h
* \brief Copy-on-write helper macro for IR ffi::ObjectRef types.
*/
#ifndef TVM_IR_COW_H_
#define TVM_IR_COW_H_
#include <tvm/ffi/object.h>
#include <utility>
namespace tvm {
/*!
* \brief Define CopyOnWrite function in an ffi::ObjectRef.
* \param ObjectName The Type of the Node.
*
* CopyOnWrite will generate a unique copy of the internal node.
* The node will be copied if it is referenced by multiple places.
* The function returns the raw pointer to the node to allow modification
* of the content.
*
* \code
*
* MyCOWObjectRef ref, ref2;
* ref2 = ref;
* ref.CopyOnWrite()->value = new_value;
* assert(ref2->value == old_value);
* assert(ref->value == new_value);
*
* \endcode
*/
#ifndef TVM_DEFINE_OBJECT_REF_COW_METHOD
#define TVM_DEFINE_OBJECT_REF_COW_METHOD(ObjectName) \
static_assert(ObjectName::_type_final, \
"TVM's CopyOnWrite may only be used for " \
"Object types that are declared as final, " \
"using the TVM_FFI_DECLARE_OBJECT_INFO_FINAL macro."); \
ObjectName* CopyOnWrite() { \
TVM_FFI_ICHECK(data_ != nullptr); \
if (!data_.unique()) { \
auto n = ::tvm::ffi::make_object<ObjectName>(*(operator->())); \
::tvm::ffi::ObjectPtr<::tvm::ffi::Object>(std::move(n)).swap(data_); \
} \
return static_cast<ObjectName*>(data_.get()); \
}
#endif // TVM_DEFINE_OBJECT_REF_COW_METHOD
} // namespace tvm
#endif // TVM_IR_COW_H_
include/tvm/ir/env_func.h
+162
View File
@@ -0,0 +1,162 @@
/*
* 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/ir/env_func.h
* \brief Serializable global function used in IR.
*/
#ifndef TVM_IR_ENV_FUNC_H_
#define TVM_IR_ENV_FUNC_H_
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/base.h>
#include <string>
#include <utility>
namespace tvm {
/*!
* \brief A serializable function backed by TVM's global environment.
*
* This is a wrapper to enable serializable global ffi::Function.
* An EnvFunc is saved by its name in the global registry
* under the assumption that the same function is registered during load.
* \sa EnvFunc
*/
class EnvFuncNode : public ffi::Object {
public:
/*! \brief Unique name of the global function */
ffi::String name;
/*! \brief The internal packed function */
ffi::Function func;
/*! \brief constructor */
EnvFuncNode() {}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
// func do not participate in structural equal and hash.
refl::ObjectDef<EnvFuncNode>()
.def_ro("name", &EnvFuncNode::name)
.def_ro("func", &EnvFuncNode::func, refl::AttachFieldFlag::SEqHashIgnore());
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.EnvFunc", EnvFuncNode, ffi::Object);
};
/*!
* \brief Managed reference to EnvFuncNode.
* \sa EnvFuncNode
*/
class EnvFunc : public ffi::ObjectRef {
public:
EnvFunc() {}
explicit EnvFunc(ffi::ObjectPtr<ffi::Object> n) : ffi::ObjectRef(n) {}
/*!
* \brief constructor with UnsafeInit
*/
explicit EnvFunc(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
/*! \return The internal global function pointer */
const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
/*!
* \brief Invoke the function.
* \param args The arguments
* \returns The return value.
*/
template <typename... Args>
ffi::Any operator()(Args&&... args) const {
const EnvFuncNode* n = operator->();
TVM_FFI_ICHECK(n != nullptr);
return n->func(std::forward<Args>(args)...);
}
/*!
* \brief Get a global function based on the name.
* \param name The name of the global function.
* \return The created global function.
* \note The function can be unique
*/
TVM_DLL static EnvFunc Get(const ffi::String& name);
/*! \brief specify container node */
using ContainerType = EnvFuncNode;
};
/*!
* \brief Please refer to \ref TypedEnvFuncAnchor "TypedEnvFunc<R(Args..)>"
*/
template <typename FType>
class TypedEnvFunc;
/*!
* \anchor TypedEnvFuncAnchor
* \brief A typed version of EnvFunc.
* It is backed by a GlobalFuncNode internally.
*
* \tparam R The return value of the function.
* \tparam Args The argument signature of the function.
* \sa EnvFunc
*/
template <typename R, typename... Args>
class TypedEnvFunc<R(Args...)> : public ffi::ObjectRef {
public:
/*! \brief short hand for this function type */
using TSelf = TypedEnvFunc<R(Args...)>;
TypedEnvFunc() {}
explicit TypedEnvFunc(ffi::ObjectPtr<ffi::Object> n) : ffi::ObjectRef(n) {}
/*!
* \brief constructor with UnsafeInit
*/
explicit TypedEnvFunc(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
/*!
* \brief Assign global function to a TypedEnvFunc
* \param other Another global function.
* \return reference to self.
*/
TSelf& operator=(const EnvFunc& other) {
ffi::ObjectRef::operator=(other);
return *this;
}
/*! \return The internal global function pointer */
const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
/*!
* \brief Invoke the function.
* \param args The arguments
* \returns The return value.
*/
R operator()(Args... args) const {
const EnvFuncNode* n = operator->();
TVM_FFI_ICHECK(n != nullptr);
if constexpr (std::is_same_v<R, void>) {
n->func(std::forward<Args>(args)...);
} else {
ffi::Any res = n->func(std::forward<Args>(args)...);
if constexpr (std::is_same_v<R, ffi::Any>) {
return res;
} else {
return std::move(res).cast<R>();
}
}
}
/*! \brief specify container node */
using ContainerType = EnvFuncNode;
};
} // namespace tvm
#endif // TVM_IR_ENV_FUNC_H_
include/tvm/ir/expr.h
+568
View File
@@ -0,0 +1,568 @@
/*
* 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/ir/expr.h
* \brief Base expr nodes in TVM.
*/
#ifndef TVM_IR_EXPR_H_
#define TVM_IR_EXPR_H_
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/base_expr.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/source_map.h>
#include <algorithm>
#include <functional>
#include <limits>
#include <optional>
#include <string>
#include <type_traits>
namespace tvm {
// Forward-declare VirtualDevice to avoid circular imports.
class VirtualDevice;
/*!
* \brief add operator
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator+(PrimExpr a, PrimExpr b);
/*!
* \brief subtraction operator
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator-(PrimExpr a, PrimExpr b);
/*!
* \brief negation.
*
* \param a input.
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator-(PrimExpr a);
/*!
* \brief multiplication operator
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator*(PrimExpr a, PrimExpr b);
/*!
* \brief division operator
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator/(PrimExpr a, PrimExpr b);
/*!
* \brief left shift operator
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b);
/*!
* \brief right shift operator
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b);
/*!
* \brief greater
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator>(PrimExpr a, PrimExpr b);
/*!
* \brief greater_equal
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b);
/*!
* \brief less
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator<(PrimExpr a, PrimExpr b);
/*!
* \brief less_equal
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b);
/*!
* \brief equal
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator==(PrimExpr a, PrimExpr b);
/*!
* \brief not_equal
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b);
/*!
* \brief and
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note This operator does eager constant folding.
*/
TVM_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b);
/*!
* \brief or
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note This operator does eager constant folding.
*/
TVM_DLL PrimExpr operator||(PrimExpr a, PrimExpr b);
/*!
* \brief not
*
* \param a left operand
* \return The result expression.
* \note This operator does eager constant folding.
*/
TVM_DLL PrimExpr operator!(PrimExpr a);
/*!
* \brief take bitwise and of two values
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator&(PrimExpr a, PrimExpr b);
/*!
* \brief take bitwise or of two values
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator|(PrimExpr a, PrimExpr b);
/*!
* \brief take bitwise xor of two values
*
* \param a left operand
* \param b right operand
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
/*!
* \brief take bitwise negation of two values
*
* \param a the input expression.
* \return The result expression.
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr operator~(PrimExpr a);
class GlobalVar;
/*!
* \brief Global variable that lives in the top-level module.
*
* A GlobalVar only refers to function definitions.
* This is used to enable recursive calls between function.
*
* \sa GlobalVarNode
*/
class GlobalVarNode : public ExprNode {
public:
/*! \brief The name of the variable, this only acts as a hint. */
ffi::String name_hint;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GlobalVarNode>().def_ro("name_hint", &GlobalVarNode::name_hint);
// A GlobalVar identifies a module-level symbol. Its type is derived from the
// corresponding function definition and is not part of the symbol identity.
refl::TypeAttrDef<GlobalVarNode>()
.def("__s_equal__", &GlobalVarNode::SEqual)
.def("__s_hash__", &GlobalVarNode::SHash);
}
bool SEqual(const GlobalVarNode* other,
ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const {
return equal(name_hint, other->name_hint, false, "name_hint");
}
int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const {
return hash(name_hint, init_hash, false);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.GlobalVar", GlobalVarNode, ExprNode);
};
/*!
* \brief Managed reference to GlobalVarNode.
* \sa GlobalVarNode
*/
class GlobalVar : public Expr {
public:
TVM_DLL explicit GlobalVar(ffi::String name_hint, Span span = {});
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GlobalVar, Expr, GlobalVarNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(GlobalVarNode);
};
/*!
* \brief Call corresponds to callable invocation.
*/
class CallNode : public ExprNode {
public:
/*!
* \brief The operator/function being invoked.
*
* It can be an Op, a GlobalVar, a local function value, or another callable
* expression.
*/
Expr op;
/*! \brief The arguments of the call. */
ffi::Array<Expr> args;
/*! \brief The additional attributes. */
Attrs attrs;
/*! \brief The type information arguments passed to the callee. */
ffi::Array<Type> ty_args;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<CallNode>()
.def_ro("op", &CallNode::op)
.def_ro("args", &CallNode::args)
.def_ro("attrs", &CallNode::attrs)
.def_ro("ty_args", &CallNode::ty_args);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Call", CallNode, ExprNode);
};
/*!
* \brief Managed reference to CallNode.
*/
class Call : public Expr {
public:
TVM_DLL Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs = Attrs(),
ffi::Array<Type> ty_args = ffi::Array<Type>(), Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Call, Expr, CallNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
};
/*!
* \brief Constant integer literals in the program.
* \sa IntImm
*/
class IntImmNode : public ExprNode {
public:
/*! \brief the Internal value. */
int64_t value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IntImmNode>().def_ro("value", &IntImmNode::value);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmNode, ExprNode);
};
/*!
* \brief Managed reference class to IntImmNode.
*
* \sa IntImmNode
*/
class IntImm : public PrimExpr {
public:
/*!
* \brief Constructor.
* \param value_ty The primitive type of the value.
* \param value The internal value.
* \param span The location of this object in the source code.
*/
TVM_DLL IntImm(PrimType value_ty, int64_t value, Span span = Span());
/*!
* \brief Construct a scalar boolean constant.
* \param value The boolean value.
* \param span The location of this object in the source code.
*/
static IntImm Bool(bool value, Span span = Span()) {
return IntImm(PrimType::Bool(), value, span);
}
/*!
* \brief Construct a scalar int32 constant.
* \param value The integer value.
* \param span The location of this object in the source code.
*/
static IntImm Int32(int64_t value, Span span = Span()) {
return IntImm(PrimType::Int(32), value, span);
}
/*!
* \brief Construct a scalar int64 constant.
* \param value The integer value.
* \param span The location of this object in the source code.
*/
static IntImm Int64(int64_t value, Span span = Span()) {
return IntImm(PrimType::Int(64), value, span);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntImm, PrimExpr, IntImmNode);
static constexpr bool _type_container_is_exact = true;
TVM_DEFINE_OBJECT_REF_COW_METHOD(IntImmNode);
};
/*!
* \brief Constant floating point literals in the program.
* \sa FloatImm
*/
class FloatImmNode : public ExprNode {
public:
/*! \brief The constant value content. */
double value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FloatImmNode>().def_ro("value", &FloatImmNode::value);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.FloatImm", FloatImmNode, ExprNode);
};
/*!
* \brief Managed reference class to FloatImmNode.
*
* \sa FloatImmNode
*/
class FloatImm : public PrimExpr {
public:
/*!
* \brief Constructor.
* \param value_ty The primitive type of the value.
* \param value The internal value.
* \param span The location in the source code.
*/
TVM_DLL FloatImm(PrimType value_ty, double value, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FloatImm, PrimExpr, FloatImmNode);
static constexpr bool _type_container_is_exact = true;
TVM_DEFINE_OBJECT_REF_COW_METHOD(FloatImmNode);
};
/*! \brief range over one dimension */
class RangeNode : public ffi::Object {
public:
/*! \brief beginning of the node */
PrimExpr min;
/*! \brief the extend of range */
PrimExpr extent;
/*! \brief the location of this range in the source */
mutable Span span;
/*! \brief constructor */
RangeNode() {}
RangeNode(PrimExpr min, PrimExpr extent, Span span = Span())
: min(min), extent(extent), span(span) {}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RangeNode>()
.def_ro("min", &RangeNode::min)
.def_ro("extent", &RangeNode::extent)
.def_ro("span", &RangeNode::span, refl::AttachFieldFlag::SEqHashIgnore());
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Range", RangeNode, ffi::Object);
};
/*! \brief Range container */
class Range : public ffi::ObjectRef {
public:
/*!
* \brief constructor by begin and end
* \param begin The begin of the range.
* \param end The end of the range.
* \param span The location of the Range in the source.
*/
TVM_DLL Range(PrimExpr begin, PrimExpr end, Span span = Span());
/*!
* \brief construct a new range with min and extent
* The corresponding constructor is removed,
* because that is counter convention of tradition meaning
* of range(begin, end)
*
* \param min The minimum range.
* \param extent The extent of the range.
* \param span The location of the Range in the source.
*/
TVM_DLL static Range FromMinExtent(PrimExpr min, PrimExpr extent, Span span = Span());
// declare range.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Range, ffi::ObjectRef, RangeNode);
};
namespace ffi {
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, IntImmNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, FloatImmNode> = true;
// Type traits to enable automatic conversion into IntImm, Integer, and Bool
// when called through the FFI
template <>
inline constexpr bool use_default_type_traits_v<IntImm> = false;
// specialize to enable implicit conversion from const char*
template <>
struct TypeTraits<IntImm> : public ObjectRefWithFallbackTraitsBase<IntImm, int64_t> {
TVM_FFI_INLINE static IntImm ConvertFallbackValue(int64_t value) {
auto value_ty =
(value > std::numeric_limits<int>::max() || value < std::numeric_limits<int>::min())
? PrimType::Int(64)
: PrimType::Int(32);
return IntImm(value_ty, value);
}
};
template <>
inline constexpr bool use_default_type_traits_v<FloatImm> = false;
template <>
struct TypeTraits<FloatImm> : public ObjectRefWithFallbackTraitsBase<FloatImm, double> {
TVM_FFI_INLINE static FloatImm ConvertFallbackValue(double value) {
return FloatImm(PrimType::Float(32), value);
}
};
} // namespace ffi
} // namespace tvm
/* \brief Allow tvm.GLobalVar as key in STL tables
*
* For most IR expressions, it would be ambiguous whether the
* expression should follow reference equality or structural equality.
* This is not the case for variables, which do not contain nested
* internal structure, and are frequently used as keys in lookup
* tables.
*
* Providing `std::hash` and `std::equal_to` specializations for
* `tvm::GlobalVar` allows it to be used as a key in STL tables. For
* other IR expressions, the user must specify the type of equality
* used (e.g. `std::unordered_set<T, StructuralHash, StructuralEqual>`
* or `std::unordered_set<T, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>`).
*/
template <>
struct std::hash<tvm::GlobalVar> {
std::size_t operator()(const tvm::GlobalVar& var) const { return tvm::ffi::ObjectPtrHash()(var); }
};
template <>
struct std::equal_to<tvm::GlobalVar> {
bool operator()(const tvm::GlobalVar& var_a, const tvm::GlobalVar& var_b) const {
return tvm::ffi::ObjectPtrEqual()(var_a, var_b);
}
};
#endif // TVM_IR_EXPR_H_
include/tvm/ir/function.h
+256
View File
@@ -0,0 +1,256 @@
/*
* 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/ir/function.h
* \brief Function nodes.
*/
#ifndef TVM_IR_FUNCTION_H_
#define TVM_IR_FUNCTION_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/expr.h>
#include <string>
#include <type_traits>
namespace tvm {
/*!
* \brief Possible Calling conventions.
*
* NOTE: The calling convention also implies
* the way we implement the function during lowering.
*/
enum class CallingConv : int {
/*!
* \brief Default calling convention.
*
* - Uses the native calling convention of the target.
* - Implementation: specified by the native target.
*/
kDefault = 0,
/*!
* \brief ffi::Function that exposes a Cffi::Function signature.
*
* - Calling by ffi::Function calling convention.
* - Implementation: Expose a function with the Cffi::Function signature.
*/
kCPackedFunc = 1,
/*!
* \brief Device kernel launch
*
* - Call by ffi::Function calling convention.
* - Implementation: defined by device runtime(e.g. runtime/cuda)
*/
kDeviceKernelLaunch = 2,
};
/*!
* \brief Supported linkage types.
*/
enum class LinkageType : int {
/*!
* \brief Internal linkage.
*/
kInternal = 0,
/*!
* \brief External linkage.
- Function with external linkage should have a global symbol attached to it.
*/
kExternal = 1
};
/*!
* \brief Generic attribute names that can be attached to any function.
*
* \sa tvm::tirx::attr, tvm::relax::attr
*/
namespace attr {
/*!
* \brief Indicates the special calling convention.
*
* Type: IntImm
*
* \sa tvm::CallingConv
*/
constexpr const char* kCallingConv = "calling_conv";
/*!
* \brief Compilation target of the function.
*
* Type: Target
*
* \sa tvm::Target
*/
constexpr const char* kTarget = "target";
/*!
* \brief Global linker symbol of the function in generated code.
*
* This option forces the code generator to name the
* function with the given.
*
* For example, we could set a global_symbol of a function
* early to make sure that we can always refer to it by
* the symbol name in the generated DLL.
*
* We should not set the attribute for local functions,
* so that the compiler can freely rename them.
*
* A unique global symbol will be automatically assigned
* to each function in the module before the target code
* generation phase.
*
* Type: String
*/
constexpr const char* kGlobalSymbol = "global_symbol";
/*!
* \brief The function uses s_tir (apache-derived TIR) semantics:
* parser fills layout=None, ScriptComplete wraps body in a root SBlock,
* and printer emits `s_tir=True` on the decorator.
* Default (attr absent or False) is tirx semantics.
*
* Type: IntImm (bool dtype)
*/
constexpr const char* kSTir = "s_tir";
/*!
* \brief Number of inputs of the Primfunc
*
* Type: Int
*/
constexpr const char* kNumInputs = "num_inputs";
} // namespace attr
/*!
* \brief Base node of all functions.
*
* We support several variants of functions throughout the stack.
* All of the functions share the same type system
* to support cross variant calls.
*
* \sa BaseFunc
*/
class BaseFuncNode : public ExprNode {
public:
/*! \brief Additional attributes storing the meta-data */
DictAttrs attrs;
/*!
* \brief Get a function attribute.
*
* \param attr_key The attribute key.
* \param default_value The default value if the key does not exist, defaults to nullptr.
*
* \return The result
*
* \tparam TOBjectRef the expected object type.
* \throw Error if the key exists but the value does not match TObjectRef
*
* \code
*
* void GetAttrExample(const BaseFunc& f) {
* auto value = f->GetAttr<int64_t>("AttrKey", 0);
* }
*
* \endcode
*/
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key,
ffi::Optional<TObjectRef> default_value = std::nullopt) const {
return attrs.GetAttr(attr_key, default_value);
}
// variant that uses TObjectRef to enable implicit conversion to default value.
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
}
/*!
* \brief Check whether the function has an non-zero integer attr.
*
* This function can be used to check whether an optional
* attribute mark(e.g. inline) exists.
*
* \param attr_key The key to the attribute.
* \return The check result.
*
* \code
*
* void HasNonzeroAttrExample(const BaseFunc& f) {
* if (f->HasNonzeroAttr(attr::kInline)) {
* // inline the function.
* }
* }
*
* \endcode
*/
bool HasNonzeroAttr(const std::string& attr_key) const { return attrs.HasNonzeroAttr(attr_key); }
/*!
* \brief Get the type of the linkage.
*
* Currently, we only consider external/internal linkage.
* This can be extended in the future when necessary.
*
* \return Linkage type.
*
* \code
*
* void Example(const BaseFunc& f) {
* if (f->GetLinkageType() == tvm::LinkageType::kExternal) {
* // Do not remove a function with external linkage
* }
* }
*
* \endcode
*/
LinkageType GetLinkageType() const {
if (GetAttr<ffi::String>(attr::kGlobalSymbol))
return LinkageType::kExternal;
else
return LinkageType::kInternal;
}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BaseFuncNode>().def_ro("attrs", &BaseFuncNode::attrs);
}
static constexpr const uint32_t _type_child_slots = 2;
TVM_FFI_DECLARE_OBJECT_INFO("ir.BaseFunc", BaseFuncNode, ExprNode);
};
/*!
* \brief Managed reference to BaseFuncNode.
* \sa BaseFuncNode
*/
class BaseFunc : public Expr {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BaseFunc, Expr, BaseFuncNode);
};
} // namespace tvm
#endif // TVM_IR_FUNCTION_H_
include/tvm/ir/global_info.h
+117
View File
@@ -0,0 +1,117 @@
/*
* 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/ir/global_info.h
* \brief GlobalInfo are globally static object that are referred by the IR itself.
*/
#ifndef TVM_IR_GLOBAL_INFO_H_
#define TVM_IR_GLOBAL_INFO_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/target/target.h>
namespace tvm {
/*!
* \brief Abstract label for an area of memory.
*/
using MemoryScope = ffi::String;
/*!
* \brief GlobalInfo are globally static object that are referred by the IR itself.
* Base node for all global info that can appear in the IR
*/
class GlobalInfoNode : public ffi::Object {
public:
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("ir.GlobalInfo", GlobalInfoNode, ffi::Object);
};
/*!
* \brief Managed reference to GlobalInfoNode.
* \sa GlobalInfoNode
*/
class GlobalInfo : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GlobalInfo, ffi::ObjectRef, GlobalInfoNode);
};
/*!
* \brief A global info subclass for virtual devices.
*/
class VDeviceNode : public GlobalInfoNode {
public:
/*! \brief The \p Target describing how to compile for the virtual device. */
Target target;
/*! \brief The device identifier for the virtual device. This enables us to
* differentiate between distinct devices with same Target, such as multiple GPUs.
*/
int vdevice_id;
MemoryScope memory_scope;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<VDeviceNode>()
.def_ro("target", &VDeviceNode::target)
.def_ro("vdevice_id", &VDeviceNode::vdevice_id)
.def_ro("memory_scope", &VDeviceNode::memory_scope);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.VDevice", VDeviceNode, GlobalInfoNode);
};
/*!
* \brief Managed reference to VDeviceNode.
* \sa VDeviceNode
*/
class VDevice : public GlobalInfo {
public:
TVM_DLL explicit VDevice(Target tgt, int dev_id, MemoryScope mem_scope);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VDevice, GlobalInfo, VDeviceNode);
};
/*!
* \brief A dummy global info sub-class for testing purpose.
*/
class DummyGlobalInfoNode : public GlobalInfoNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DummyGlobalInfoNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DummyGlobalInfo", DummyGlobalInfoNode, GlobalInfoNode);
};
/*!
* \brief Managed reference to DummyGlobalInfoNode.
* \sa DummyGlobalInfoNode
*/
class DummyGlobalInfo : public GlobalInfo {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DummyGlobalInfo, GlobalInfo, DummyGlobalInfoNode);
};
} // namespace tvm
#endif // TVM_IR_GLOBAL_INFO_H_
include/tvm/ir/instrument.h
+158
View File
@@ -0,0 +1,158 @@
/*
* 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/ir/instrument.h
*
* This file introduces a pass instrument infrastructure, inspired by LLVM and MLIR.
* It inserts instrumentation points around passes.
*/
#ifndef TVM_IR_INSTRUMENT_H_
#define TVM_IR_INSTRUMENT_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <utility>
#include <vector>
namespace tvm {
class IRModule;
// Forward class for PassInstrumentNode methods
namespace transform {
class PassInfo;
} // namespace transform
namespace instrument {
/*!
* \brief PassInstrumentNode forms an instrument implementation.
* It provides API for users to register callbacks at different instrumentation points.
*
* Within a PassContext, call sequence of a PassInstrument implementation is like:
*
* with PassContext(instruments=[pi]): # pi = a PassInstrument implementation
* pi.EnterPassContext()
*
* if pi.ShouldRun(Pass1):
* pi.RunBeforePass()
* Pass1()
* pi.RunAfterPass()
*
* if pi.ShouldRun(Pass2):
* pi.RunBeforePass()
* Pass2()
* pi.RunAfterPass()
*
* pi.ExitPassContext()
*
* `EnterPassContext` and `ExitPassContext` are only called once when entering/exiting a
* PassContext. `ShouldRun`, `RunBeforePass` and `RunAfterPass` are called multiple times depending
* on how many passes.
*
* If there are multiple pass instrumentations provided, the instrument points are the same.
* PassInstrument implementations' callbacks are called in order:
*
* with PassContext(instruments=[pi1, pi2]): # pi1, pi2 = two distinct PassInstrument impls
* pi.EnterPassContext() for pi in instruments
*
* should_run = all([pi.ShoudRun(Pass1) for pi in instruments)])
* if (should_run)
* pi.RunBeforePass() for pi in instruments
* Pass1()
* pi.RunAfterPass() for pi in instruments
*
* should_run = all([pi.ShouldRun(Pass2) for pi in instruments)])
* if (should_run)
* pi.RunBeforePass() for pi in instruments
* Pass2()
* pi.RunAfterPass() for pi in instruments
*
* pi.ExitPassContext() for pi in instruments
*
* Note:
* 1. Assume there is no dependency between PassInstrument implementations in `instruments` .
* 2. `EnterPassContext` and `ExitPassContext` have `with` behavior (see PassContext and its FFI):
* If there is any exception raised in `ShouldRun()`, `RunBeforePass()`, `RunAfterPass()` and
* `Pass()`, `ExitPassContext()` is still called.
* 3. In mutiple PassInstrument instances scenario, callbacks are called in order:
* If one throws exceptions, remainings will not be called.
*
* \sa PassInstrument
* \sa src/ir/transform.cc
*/
class PassInstrumentNode : public ffi::Object {
public:
/*! \brief Name of this pass instrument object. */
ffi::String name;
virtual ~PassInstrumentNode() {}
/*! \brief Instrument when entering PassContext. Called once within a PassContext. */
virtual void EnterPassContext() const = 0;
/*! \brief Instrument when exiting PassContext. Called once within a PassContext. */
virtual void ExitPassContext() const = 0;
/*!
* \brief Determine whether to run the pass or not. Called multiple times depend on number of
* passes.
* \param mod The module that an optimization pass runs on.
* \param info The pass information.
*
* \return true to run the pass; false to skip the pass.
*/
virtual bool ShouldRun(const IRModule& mod, const transform::PassInfo& info) const = 0;
/*!
* \brief Instrument before pass run. Called multiple times depend on number of passes.
* \param mod The module that an optimization pass runs on.
* \param info The pass information.
*/
virtual void RunBeforePass(const IRModule& mod, const transform::PassInfo& info) const = 0;
/*!
* \brief Instrument after pass run. Called multiple time depend on number of passes.
* \param mod The module that an optimization pass runs on.
* \param info The pass information.
*/
virtual void RunAfterPass(const IRModule& mod, const transform::PassInfo& info) const = 0;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PassInstrumentNode>().def_ro("name", &PassInstrumentNode::name);
}
TVM_FFI_DECLARE_OBJECT_INFO("instrument.PassInstrument", PassInstrumentNode, ffi::Object);
};
/*!
* \brief Managed reference class for PassInstrumentNode
* \sa PassInstrumentNode
*/
class PassInstrument : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInstrument, ffi::ObjectRef, PassInstrumentNode);
};
} // namespace instrument
} // namespace tvm
#endif // TVM_IR_INSTRUMENT_H_
include/tvm/ir/module.h
+366
View File
@@ -0,0 +1,366 @@
/*
* 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/ir/module.h
* \brief IRModule that holds the functions and type definitions.
*/
#ifndef TVM_IR_MODULE_H_
#define TVM_IR_MODULE_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
#include <tvm/ir/global_info.h>
#include <tvm/ir/source_map.h>
#include <tvm/ir/type.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
namespace tvm {
class IRModule;
/*!
* \brief IRModule that holds functions and type definitions.
*
* IRModule is the basic unit for all IR transformations across the stack.
*
* Many operations require access to the global IRModule.
* We pass the IRModule by value in a functional style as an explicit argument,
* but we mutate the Module while optimizing programs.
* \sa IRModule
*/
class IRModuleNode : public ffi::Object {
public:
/*! \brief A map from ids to all global functions. */
ffi::Map<GlobalVar, BaseFunc> functions;
/*! \brief The source map for the module. */
SourceMap source_map;
/* \brief Additional attributes storing meta-data about the module. */
DictAttrs attrs;
/*! \brief Globally static object that are referred by the IR itself */
ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
/*!
* \brief A map from string names to global variables that
* ensures global uniqueness.
*/
ffi::Map<ffi::String, GlobalVar> global_var_map_;
/*!
* \brief Get a module attribute.
*
* \param attr_key The attribute key.
* \param default_value The default value if the key does not exist, defaults to nullptr.
*
* \return The result
*
* \tparam TObjectRef the expected object type.
* \throw Error if the key exists but the value does not match TObjectRef
*
* \code
*
* void GetAttrExample(const IRModule& mod) {
* auto value = f->GetAttr<int64_t>("AttrKey", 0);
* }
*
* \endcode
*/
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetAttr(
const std::string& attr_key,
ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
return attrs.GetAttr(attr_key, default_value);
}
// variant that uses TObjectRef to enable implicit conversion to default value.
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
}
/*!
* \brief Get the metadata attributes.
* \returns The additional meta-data attributes
*/
DictAttrs GetAttrs() const { return attrs; }
/*!
* \brief Check whether the module has an non-zero integer attr.
*
* This function can be used to check whether an optional
* attribute mark(e.g. inline) exists.
*
* \param attr_key The key to the attribute.
* \return The check result.
*
* \code
*
* void HasNonzeroAttrExample(const IRModule& mod) {
* if (mod->HasNonzeroAttr(attr::kInline)) {
* // inline the function.
* }
* }
*
* \endcode
*/
bool HasNonzeroAttr(const std::string& attr_key) const { return attrs.HasNonzeroAttr(attr_key); }
IRModuleNode() : source_map() {}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IRModuleNode>()
.def_ro("functions", &IRModuleNode::functions)
.def_ro("global_var_map_", &IRModuleNode::global_var_map_)
.def_ro("source_map", &IRModuleNode::source_map)
.def_ro("attrs", &IRModuleNode::attrs)
.def_ro("global_infos", &IRModuleNode::global_infos);
// register custom structural equal and hash.
refl::TypeAttrDef<IRModuleNode>()
.def("__s_equal__", &IRModuleNode::SEqual)
.def("__s_hash__", &IRModuleNode::SHash);
}
TVM_DLL bool SEqual(const IRModuleNode* other,
ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
TVM_DLL int64_t SHash(int64_t init_hash,
ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const;
/*!
* \brief Add a function to the global environment.
* \param var The var of the global function.
* \param func The function.
* \param update Controls whether you can replace a definition in the
* environment.
*/
TVM_DLL void Add(const GlobalVar& var, const BaseFunc& func, bool update = false);
/*!
* \brief Add a function to the global environment.
* \param var The name of the global function.
* \param func The function.
*
* It does not do type inference as Add does.
*/
TVM_DLL void AddUnchecked(const GlobalVar& var, const BaseFunc& func);
/*!
* \brief Update a function in the global environment.
* \param var The name of the global function to update.
* \param func The new function.
*/
TVM_DLL void Update(const GlobalVar& var, const BaseFunc& func);
/*!
* \brief Update an array of global infos in the global environment.
* \param name The name of the global info.
* \param info The new array of global infos.
*/
TVM_DLL void UpdateGlobalInfo(const ffi::String& name, const ffi::Array<GlobalInfo>& info);
/*!
* \brief Remove a function from the global environment.
* \param var The name of the global function to update.
*/
TVM_DLL void Remove(const GlobalVar& var);
/*!
* \brief Check if the global_var_map_ contains a global variable.
* \param name The variable name.
* \returns true if contains, otherise false.
*/
TVM_DLL bool ContainGlobalVar(const ffi::String& name) const;
/*!
* \brief Lookup a global function by its variable.
* \param str The unique string specifying the global variable.
* \returns The global variable.
*/
TVM_DLL GlobalVar GetGlobalVar(const ffi::String& str) const;
/*!
* \brief Collect all global vars defined in this module, ordered by
* the global variable name.
* \returns An array of global vars
*/
TVM_DLL ffi::Array<GlobalVar> GetGlobalVars() const;
/*!
* \brief Look up a global function by its variable.
* \param var The global var to lookup.
* \returns The function named by the variable argument.
*/
TVM_DLL BaseFunc Lookup(const GlobalVar& var) const;
/*!
* \brief Look up a global function by its string name
* \param name The name of the function.
* \returns The function named by the argument.
*/
TVM_DLL BaseFunc Lookup(const ffi::String& name) const;
/*!
* \brief Update the functions inside this environment by
* functions in another environment.
* \param other The other environment.
*/
TVM_DLL void Update(const IRModule& other);
/*!
* \brief Create a shallow copy of this IRModule.
* \returns The shallow copy of the IRModule.
*/
TVM_DLL IRModule ShallowCopy();
/*!
* \brief The set of imported files.
*/
TVM_DLL std::unordered_set<ffi::String> Imports() const;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IRModule", IRModuleNode, ffi::Object);
private:
friend class IRModule;
};
/*!
* \brief Managed reference class to IRModuleNode.
* \sa IRModuleNode
*/
class IRModule : public ffi::ObjectRef {
public:
/*!
* \brief constructor
* \param functions Functions in the module.
* \param map The module source map.
* \param attrs The module meta-data attributes.
* \param global_infos Global infos in the module.
*/
TVM_DLL explicit IRModule(ffi::Map<GlobalVar, BaseFunc> functions, SourceMap map = {},
DictAttrs attrs = DictAttrs(),
ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos = {});
/*! \brief default constructor */
IRModule() : IRModule(ffi::Map<GlobalVar, BaseFunc>({})) {}
/*!
* \brief constructor
* \param n The object pointer.
*/
explicit IRModule(ffi::ObjectPtr<IRModuleNode> n) : ffi::ObjectRef(n) {}
/*!
* \brief constructor with UnsafeInit
*/
explicit IRModule(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
/*! \return mutable pointers to the node. */
IRModuleNode* operator->() const {
auto* ptr = get_mutable();
TVM_FFI_ICHECK(ptr != nullptr);
return static_cast<IRModuleNode*>(ptr);
}
/*!
* \brief As for \p FromExprInContext, but assuming \p expr is bound to 'main' and no
* imports.
*/
TVM_DLL static IRModule FromExpr(const Expr& expr,
const ffi::Map<GlobalVar, BaseFunc>& global_funcs = {});
/*!
* \brief Create a shallow copy of an IRModule.
* \param mod The module to copy.
* \return The copied module.
*/
IRModule ShallowCopyIRModule(IRModule mod);
/*! \brief Declare the container type. */
using ContainerType = IRModuleNode;
// allow copy on write.
TVM_DEFINE_OBJECT_REF_COW_METHOD(IRModuleNode);
};
namespace attr {
// Following are attributes for IRModule only.
/*!
* \brief Name of the module
*
* Type: String
*/
constexpr const char* kModuleName = "mod_name";
/*!
* \brief All the runtime::Modules accumulated during compilation by external codegen. These
* modules must be either directly linked or captured in the final compilation artifact.
*
* Type: ffi::Array<runtime::Module>
*/
constexpr const char* kExternalMods = "external_mods";
/*!
* \brief A prefix for generating C symbols system lib creation.
*
* This prefix guides passes that creates global_symbol for internal functions
* that may have c linkage (e.g. TIR functions and some BYOC functions). It also affects
* the symbol of the fat bin blob during module export.
*
* This attribute is used to avoid symbol conflict when we
* generate and combine multiple system libs that get linked into one.
*
* Rationale: mechanisms like BYOC rely on the common global symbol
* and each external compiler also has its own mechanism of mangling.
* As a result, we cannot rely on other mechanisms on setting a global_symbol and then renaming,
* because the external compiler already agreed on the name.
*
* system_lib_prefix provides a way to hint at the passes to allow names to
* avoid name conflict at the beginning.
*
* Note that users can still directly specify global symbols that may conflict.
* It is up to the downstream toolchain to manage those external-facing functions.
*
* This does not affect non-C linkage functions it is less of an issue because
* they will be embedded into fatbin that in different symbols,
* The system lib loader can pick the right prefix for a given prefix.
*
* Having this attribute implies system lib generation linkage.
*/
constexpr const char* kSystemLibPrefix = "system_lib_prefix";
/*!
* \brief All the named runtime::Tensors accumulated during compilation by external codegen.
* Generally the associated runtime::Module will indicate it requires bindings for these names,
* and during module initialization these bindings will be recovered from a ConstLoaderModule.
*
* Type: ffi::Map<ffi::String, runtime::Tensor>
*/
constexpr const char* kConstNameToConstant = "const_name_to_constant";
} // namespace attr
} // namespace tvm
#endif // TVM_IR_MODULE_H_
include/tvm/ir/node_functor.h
+194
View File
@@ -0,0 +1,194 @@
/*
* 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/ir/node_functor.h
* \brief Defines the Functor data structures.
*/
#ifndef TVM_IR_NODE_FUNCTOR_H_
#define TVM_IR_NODE_FUNCTOR_H_
#include <tvm/ffi/error.h>
#include <cstring>
#include <type_traits>
#include <utility>
#include <vector>
namespace tvm {
/*!
* \brief A dynamically dispatched functor on the type of the first argument.
*
* This is a class that is useful to construct polymorphic dispatching
* base on the AST/IR node's type.
*
* \code
* NodeFunctor<std::string (const ffi::ObjectRef& n, std::string prefix)> tostr;
* tostr.set_dispatch<Add>([](const ffi::ObjectRef& op, std::string prefix) {
* return prefix + "Add";
* });
* tostr.set_dispatch<IntImm>([](const ffi::ObjectRef& op, std::string prefix) {
* return prefix + "IntImm"
* });
*
* Expr x = MakeConst(1);
* Expr y = x + x;
* // dispatch to IntImm, outputs "MyIntImm"
* LOG(INFO) << tostr(x, "My");
* // dispatch to IntImm, outputs "MyAdd"
* LOG(INFO) << tostr(y, "My");
* \endcode
*
* \tparam FType function signiture
* This type if only defined for FType with function signature
*/
template <typename FType>
class NodeFunctor;
template <typename R, typename... Args>
class NodeFunctor<R(const ffi::ObjectRef& n, Args...)> {
private:
/*! \brief internal function pointer type */
typedef R (*FPointer)(const ffi::ObjectRef& n, Args...);
/*! \brief refer to itself. */
using TSelf = NodeFunctor<R(const ffi::ObjectRef& n, Args...)>;
/*! \brief internal function table */
std::vector<FPointer> func_;
/*! \brief start range of func index */
uint32_t begin_type_index_{0};
public:
/*! \brief the result type of this functor */
using result_type = R;
/*!
* \brief Whether the functor can dispatch the corresponding Node
* \param n The node to be dispatched
* \return Whether dispatching function is registered for n's type.
*/
bool can_dispatch(const ffi::ObjectRef& n) const {
uint32_t type_index = n->type_index();
if (type_index < begin_type_index_) return false;
type_index -= begin_type_index_;
return type_index < func_.size() && func_[type_index] != nullptr;
}
/*!
* \brief invoke the functor, dispatch on type of n
* \param n The Node argument
* \param args The additional arguments
* \return The result.
*/
R operator()(const ffi::ObjectRef& n, Args... args) const {
TVM_FFI_ICHECK(can_dispatch(n))
<< "NodeFunctor calls un-registered function on type " << n->GetTypeKey();
return (*func_[n->type_index() - begin_type_index_])(n, std::forward<Args>(args)...);
}
/*!
* \brief set the dispatcher for type TNode
* \param f The function to be set.
* \tparam TNode the type of Node to be dispatched.
* \return reference to self.
*/
template <typename TNode>
TSelf& set_dispatch(FPointer f) { // NOLINT(*)
uint32_t tindex = TNode::RuntimeTypeIndex();
if (func_.size() <= tindex) {
func_.resize(tindex + 1, nullptr);
}
TVM_FFI_ICHECK(func_[tindex] == nullptr)
<< "Dispatch for " << TNode::_type_key << " is already set";
TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << " Cannot call set_dispatch after calling Finalize";
func_[tindex] = f;
return *this;
}
/*!
* \brief unset the dispatcher for type TNode
*
* \tparam TNode the type of Node to be dispatched.
* \return reference to self.
*/
template <typename TNode>
TSelf& clear_dispatch() { // NOLINT(*)
uint32_t tindex = TNode::RuntimeTypeIndex();
TVM_FFI_ICHECK_LT(tindex, func_.size()) << "clear_dispatch: index out of range";
TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << " Cannot call clear_dispatch after calling Finalize";
func_[tindex] = nullptr;
return *this;
}
/*!
* \brief Finalize the functor after calling sequence of set_dispatch
* This function will attempt to find the min type index that is not null
* and optimize the space of the func table so it is more compact
*/
void Finalize() {
TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << "Can only call Finalize once";
while (begin_type_index_ < func_.size() && func_[begin_type_index_] == nullptr) {
++begin_type_index_;
}
// shift up the function value
size_t new_ftable_size = func_.size() - begin_type_index_;
if (begin_type_index_ != 0) {
std::memmove(func_.data(), func_.data() + begin_type_index_,
new_ftable_size * sizeof(FPointer));
}
func_.resize(new_ftable_size);
func_.shrink_to_fit();
}
};
#define TVM_REG_FUNC_VAR_DEF(ClsName) [[maybe_unused]] static auto& __make_functor##_##ClsName
/*!
* \brief Useful macro to set NodeFunctor dispatch in a global static field.
*
* \code
* // Use NodeFunctor to implement TVMScriptPrinter similar to Visitor Pattern.
* // vtable allows easy patch of new Node types, without changing
* // the interface of TVMScriptPrinter.
*
* class TVMScriptPrinter {
* public:
* // the dispatch function.
* static std::string Script(const ffi::ObjectRef& node, const PrinterConfig& cfg) {
* return vtable()(node, cfg);
* }
* using FType = NodeFunctor<std::string(const ffi::ObjectRef&, const PrinterConfig&)>;
* // function to return global function table
* static FType& vtable();
* };
*
* // in cpp/cc file
* TVMScriptPrinter::FType& TVMScriptPrinter::vtable() {
* static FType inst; return inst;
* }
*
* TVM_STATIC_IR_FUNCTOR(TVMScriptPrinter, vtable)
* .set_dispatch<AddNode>([](const ffi::ObjectRef& ref, const PrinterConfig& cfg) {
* auto* n = static_cast<const AddNode*>(ref.get());
* return Script(n->a, cfg) + " + " + Script(n->b, cfg);
* });
*
* \endcode
*
* \param ClsName The name of the class
* \param FField The static function that returns a singleton of NodeFunctor.
*/
#define TVM_STATIC_IR_FUNCTOR(ClsName, FField) \
TVM_FFI_STR_CONCAT(TVM_REG_FUNC_VAR_DEF(ClsName), __COUNTER__) = ClsName::FField()
} // namespace tvm
#endif // TVM_IR_NODE_FUNCTOR_H_
include/tvm/ir/op.h
+423
View File
@@ -0,0 +1,423 @@
/*
* 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/ir/op.h
* \brief Primitive operators(builtin intrinsics)
* and registry for them.
*/
#ifndef TVM_IR_OP_H_
#define TVM_IR_OP_H_
#include <tvm/ffi/error.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attr_registry_map.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/type.h>
#include <string>
#include <utility>
#include <vector>
namespace tvm {
// forward declare name.
template <typename>
class OpAttrMap;
/*!
* \brief Information about an input field of an Op (name, type, description).
*
* Populated via OpRegEntry::add_argument and consumed both by
* internal sanity checks / error messages and by external tooling
* that wants to introspect an Op's argument schema.
*/
class ArgumentInfoNode : public ffi::Object {
public:
/*! \brief name of the field */
ffi::String name;
/*! \brief type docstring information in str. */
ffi::String type_info;
/*! \brief detailed description of the type */
ffi::String description;
static void RegisterReflection() {
namespace rfl = ffi::reflection;
rfl::ObjectDef<ArgumentInfoNode>()
.def_ro("name", &ArgumentInfoNode::name)
.def_ro("type_info", &ArgumentInfoNode::type_info)
.def_ro("description", &ArgumentInfoNode::description);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.ArgumentInfo", ArgumentInfoNode, ffi::Object);
};
/*! \brief Managed reference to ArgumentInfoNode. */
class ArgumentInfo : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ArgumentInfo, ffi::ObjectRef, ArgumentInfoNode);
};
// TODO(tvm-team): migrate low-level intrinsics to use Op
/*!
* \brief Primitive Op(builtin intrinsics)
*
* This data structure stores the meta-data
* about primitive operators that can be invoked via Call.
*
* Low-level IR intrinsics(such as libc.expf) are also
* implemented via Op.
*
* \sa Op
*/
class OpNode : public ExprNode {
public:
/*! \brief name of the operator */
ffi::String name;
/*!
* \brief detailed description of the operator
* This can be used to generate docstring automatically for the operator.
*/
ffi::String description;
/* \brief Information of input arguments to the operator */
ffi::Array<ArgumentInfo> arguments;
/*!
* \brief The type key of the attribute field
* This can be empty, in which case it defaults to anything.
*/
ffi::String attrs_type_key;
/*!
* \brief attribute type index,
* this field varies in each run and is not exposed to frontend.
*/
uint32_t attrs_type_index{0};
/*!
* \brief number of input arguments to the operator,
* -1 means it is variable length
*/
int32_t num_inputs = -1;
/*!
* \brief support level of the operator,
* The lower the more priority it contains.
* This is in analogies to BLAS levels.
*/
int32_t support_level = 10;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<OpNode>()
.def_ro("name", &OpNode::name)
.def_ro("description", &OpNode::description, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("arguments", &OpNode::arguments, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("attrs_type_key", &OpNode::attrs_type_key, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("num_inputs", &OpNode::num_inputs, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("support_level", &OpNode::support_level, refl::AttachFieldFlag::SEqHashIgnore());
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Op", OpNode, ExprNode);
private:
/*! \return the internal attr registry index. */
uint32_t AttrRegistryIndex() const { return index_; }
/*! \brief repr to be printed in registry*/
std::string AttrRegistryName() const { return name; }
// friend class
template <typename>
friend class AttrRegistryMapContainerMap;
template <typename, typename>
friend class AttrRegistry;
friend class OpRegEntry;
// Program internal unique index of operator.
// Used to help index the program.
uint32_t index_{0};
};
/*!
* \brief Managed reference class to OpNode.
* \sa OpNode
*/
class Op : public Expr {
public:
explicit Op(ffi::ObjectPtr<OpNode> n) : Expr(std::move(n)) {
TVM_FFI_CHECK(defined(), ValueError) << "Op expects a defined OpNode";
}
/*!
* \brief Get additional registered attribute about operators.
* If nothing has been registered, an empty OpAttrMap will be returned.
* \param attr_name The name of the attribute.
* \return An OpAttrMap of specified attr_name.
* \tparam ValueType The type of the attribute.
*/
template <typename ValueType>
inline static OpAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
/*!
* \brief Checks if an attr map is present in the registry.
* \param attr_name The name of the attribute.
* \return bool True if the attr is present.
*/
TVM_DLL static bool HasAttrMap(const ffi::String& attr_name);
/*!
* \brief Get an Op for a given operator name.
* Will raise an error if the op has not been registered.
* \param op_name Name of the operator.
* \return Pointer to a Op, valid throughout program lifetime.
*/
TVM_DLL static const Op& Get(const ffi::String& op_name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Op, Expr, OpNode);
private:
/*!
* \brief Get generic attrmap given attr name
* \param key The attribute key
* \return The attr map.
*/
TVM_DLL static const AttrRegistryMapContainerMap<Op>& GetAttrMapContainer(const ffi::String& key);
};
/*!
* \brief Helper structure to register operators
* \sa TVM_REGISTER_OP
*/
class OpRegEntry {
public:
/*! \return the operator */
const Op& op() const { return op_; }
/*!
* \brief setter function during registration
* Set the description of operator
* \param descr the description string.
* \return reference to self.
*/
inline OpRegEntry& describe(const std::string& descr); // NOLINT(*)
/*!
* \brief Add argument information to the function.
* \param name Name of the argument.
* \param type Type of the argument.
* \param description Description of the argument.
* \return reference to self.
*/
inline OpRegEntry& add_argument(const std::string& name, const std::string& type,
const std::string& description);
/*!
* \brief Set the attrs type key and index to be AttrsType.
* \tparam AttrsType the attribute type to b set.
* \return reference to self.
*/
template <typename AttrsType>
inline OpRegEntry& set_attrs_type();
/*!
* \brief Set the attrs type key and index to be AttrsType.
* \param key The attribute type key to be set.
* \return reference to self.
*/
inline OpRegEntry& set_attrs_type_key(const ffi::String& key);
/*!
* \brief Set the num_inputs
* \param n The number of inputs to be set.
* \return reference to self.
*/
inline OpRegEntry& set_num_inputs(int32_t n); // NOLINT(*)
/*!
* \brief Set the support level of op.
* \param level The support level.
* \return reference to self.
*/
inline OpRegEntry& set_support_level(int32_t level); // NOLINT(*)
/*!
* \brief Register additional attributes to operator.
* \param attr_name The name of the attribute.
* \param value The value to be set.
* \param plevel The priority level of this set,
* an higher priority level attribute
* will replace lower priority level attribute.
* Must be bigger than 0.
*
* Cannot set with same plevel twice in the code.
*
* \tparam ValueType The type of the value to be set.
*/
template <typename ValueType>
inline OpRegEntry& set_attr(const std::string& attr_name, // NOLINT(*)
const ValueType& value, int plevel = 10);
/*!
* \brief Resets an attr of the registry.
* \param attr_name The name of the attribute.
*/
inline void reset_attr(const std::string& attr_name);
// set the name of the op to be the same as registry
inline OpRegEntry& set_name() { // NOLINT(*)
if (get()->name.length() == 0) {
get()->name = name;
}
return *this;
}
/*!
* \brief Register or get a new entry.
* \param name The name of the operator.
* \return the corresponding entry.
*/
TVM_DLL static OpRegEntry& RegisterOrGet(const ffi::String& name);
private:
template <typename, typename>
friend class AttrRegistry;
// the name
std::string name;
/*! \brief The operator */
Op op_;
/*! \brief Construct the non-null Op for this registry entry. */
static Op MakeOp(uint32_t reg_index);
// private constructor
TVM_DLL OpRegEntry(uint32_t reg_index);
// return internal pointer to op.
inline OpNode* get();
// update the attribute OpAttrMap
TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
};
/*!
* \brief ffi::Map<Op,ValueType> used to store meta-information about Op.
* \tparam ValueType The type of the value stored in map.
*/
template <typename ValueType>
class OpAttrMap : public AttrRegistryMap<Op, ValueType> {
public:
/*!
* \brief get the corresponding value element at op with default value.
* \param expr The key to the map
* \param def_value The default value when the key does not exist
* or if expr is not an Op.
* \return the const reference to the content value.
*/
inline ValueType get(const Expr& expr, ValueType def_value) const;
using TParent = AttrRegistryMap<Op, ValueType>;
using TParent::count;
using TParent::get;
using TParent::operator[];
private:
friend class Op;
// constructor
explicit OpAttrMap(const AttrRegistryMapContainerMap<Op>& map) : TParent(map) {}
};
// internal macros to make
#define TVM_OP_REGISTER_VAR_DEF [[maybe_unused]] static ::tvm::OpRegEntry& __make_##Op
/*!
* \def TVM_REGISTER_OP
* \brief Register a new operator, or set attribute of the corresponding op.
*
* \param OpName The name of registry
*
* \code
*
* TVM_REGISTER_OP("add")
* .describe("add two inputs together")
* .set_num_inputs(2)
* .set_attr<OpKernel>("gpu_kernel", AddKernel);
*
* \endcode
*/
#define TVM_REGISTER_OP(OpName) \
TVM_FFI_STR_CONCAT(TVM_OP_REGISTER_VAR_DEF, __COUNTER__) = \
::tvm::OpRegEntry::RegisterOrGet(OpName).set_name()
// implementations
template <typename ValueType>
inline OpAttrMap<ValueType> Op::GetAttrMap(const ffi::String& key) {
return OpAttrMap<ValueType>(Op::GetAttrMapContainer(key));
}
inline OpNode* OpRegEntry::get() { return const_cast<OpNode*>(op_.operator->()); }
inline OpRegEntry& OpRegEntry::describe(const std::string& descr) { // NOLINT(*)
get()->description = descr;
return *this;
}
inline OpRegEntry& OpRegEntry::add_argument(const std::string& name, const std::string& type,
const std::string& description) {
auto n = ffi::make_object<ArgumentInfoNode>();
n->name = name;
n->type_info = type;
n->description = description;
get()->arguments.push_back(ArgumentInfo(n));
return *this;
}
inline OpRegEntry& OpRegEntry::set_num_inputs(int32_t n) { // NOLINT(*)
get()->num_inputs = n;
return *this;
}
template <typename AttrsType>
inline OpRegEntry& OpRegEntry::set_attrs_type() { // NOLINT(*)
get()->attrs_type_key = AttrsType::_type_key;
get()->attrs_type_index = AttrsType::RuntimeTypeIndex();
return *this;
}
inline OpRegEntry& OpRegEntry::set_attrs_type_key(const ffi::String& key) { // NOLINT(*)
get()->attrs_type_key = key;
get()->attrs_type_index = tvm::ffi::TypeKeyToIndex(key.c_str());
return *this;
}
inline OpRegEntry& OpRegEntry::set_support_level(int32_t n) { // NOLINT(*)
get()->support_level = n;
return *this;
}
template <typename ValueType>
inline OpRegEntry& OpRegEntry::set_attr( // NOLINT(*)
const std::string& attr_name, const ValueType& value, int plevel) {
TVM_FFI_ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
UpdateAttr(attr_name, Any(value), plevel);
return *this;
}
// member functions of OpAttrMap
template <typename ValueType>
inline ValueType OpAttrMap<ValueType>::get(const Expr& expr, ValueType def_value) const {
TVM_FFI_ICHECK(expr.defined());
if (const OpNode* op = expr.as<OpNode>()) {
return this->map_.get(ffi::GetRef<Op>(op), def_value);
} else {
return def_value;
}
}
} // namespace tvm
#endif // TVM_IR_OP_H_
include/tvm/ir/scope_stack.h
+123
View File
@@ -0,0 +1,123 @@
/*
* 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/ir/scope_stack.h
* \brief A generic scope stack for managing hierarchical state during IR visiting.
*/
#ifndef TVM_IR_SCOPE_STACK_H_
#define TVM_IR_SCOPE_STACK_H_
#include <tvm/ffi/error.h>
#include <deque>
#include <type_traits>
namespace tvm {
/*!
* \brief A scope stack for maintaining hierarchical state during IR visiting.
*
* During IR tree traversal, visitors often need to track scope-local state
* (e.g., active constraints, variable bindings) that should be automatically
* cleaned up when leaving a scope. ScopeStack provides this via WithNewScope,
* which pushes a new element on entry and pops it on exit.
*
* \code
* ScopeStack<WithGroup<ConstraintContext>> constraints;
*
* // In VisitStmt_(ForNode):
* return constraints.WithNewScope([&]() -> Stmt {
* constraints.Current().Emplace(analyzer, condition);
* return StmtExprMutator::VisitStmt_(op);
* });
* \endcode
*
* \tparam T The element type stored on the stack. Must be default-constructible.
*/
template <typename T>
class ScopeStack {
public:
/*! \brief Construct with one initial scope level. */
ScopeStack() { stack_.emplace_back(); }
/*! \brief Return the number of active scopes. */
size_t size() const { return stack_.size(); }
/*! \brief Return true if no scopes are active. */
bool empty() const { return stack_.empty(); }
/*!
* \brief Access the current (innermost) scope element.
*
* The returned reference is stable across push_back/pop_back because
* std::deque guarantees pointer stability for these operations.
*
* \return Mutable reference to the top element.
*/
T& Current() {
TVM_FFI_ICHECK(!stack_.empty());
return stack_.back();
}
/*! \brief Const access to the current (innermost) scope element. */
const T& Current() const {
TVM_FFI_ICHECK(!stack_.empty());
return stack_.back();
}
/*!
* \brief Execute body within a new scope.
*
* Pushes a new T onto the stack, executes the body, then pops it.
*
* \param body A callable to execute within the scope.
* \return The return value of body(), if non-void.
*/
template <typename F>
auto WithNewScope(F&& body) -> decltype(body()) {
stack_.emplace_back();
struct Guard {
std::deque<T>* stack;
~Guard() noexcept(false) { stack->pop_back(); }
} guard{&stack_};
if constexpr (std::is_void_v<decltype(body())>) {
body();
} else {
return body();
}
}
private:
/*!
* \brief The scope stack.
*
* We use std::deque rather than std::vector for pointer stability:
* references returned by Current() remain valid across push/pop operations.
* This is critical because methods called on Current() (e.g., Emplace on
* a WithGroup) may trigger re-entrant code that pushes new scopes onto
* the same stack. With std::vector the internal buffer reallocation would
* invalidate the reference, causing use-after-free.
*/
std::deque<T> stack_;
};
} // namespace tvm
#endif // TVM_IR_SCOPE_STACK_H_
include/tvm/ir/source_map.h
+226
View File
@@ -0,0 +1,226 @@
/*
* 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 source_map.h
* \brief A map from source names to source code.
*/
#ifndef TVM_IR_SOURCE_MAP_H_
#define TVM_IR_SOURCE_MAP_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/base.h>
#include <fstream>
#include <string>
#include <utility>
#include <vector>
namespace tvm {
/*!
* \brief The source name in the Span
* \sa SourceNameNode, Span
*/
class SourceName;
/*!
* \brief The name of a source fragment.
*/
class SourceNameNode : public ffi::Object {
public:
/*! \brief The source name. */
ffi::String name;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SourceNameNode>().def_ro("name", &SourceNameNode::name);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SourceName", SourceNameNode, ffi::Object);
};
/*!
* \brief The source name of a file span.
* \sa SourceNameNode, Span
*/
class SourceName : public ffi::ObjectRef {
public:
/*!
* \brief Get an SourceName for a given operator name.
* Will raise an error if the source name has not been registered.
* \param name Name of the operator.
* \return SourceName valid throughout program lifetime.
*/
TVM_DLL static SourceName Get(const ffi::String& name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SourceName, ffi::ObjectRef, SourceNameNode);
};
/*!
* \brief Span information for debugging purposes
*/
class Span;
/*!
* \brief Stores locations in frontend source that generated a node.
*/
class SpanNode : public ffi::Object {
public:
/*! \brief The source name. */
SourceName source_name;
/*! \brief The line number. */
int line;
/*! \brief The column offset. */
int column;
/*! \brief The end line number. */
int end_line;
/*! \brief The end column number. */
int end_column;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SpanNode>()
.def_ro("source_name", &SpanNode::source_name)
.def_ro("line", &SpanNode::line)
.def_ro("column", &SpanNode::column)
.def_ro("end_line", &SpanNode::end_line)
.def_ro("end_column", &SpanNode::end_column);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("ir.Span", SpanNode, ffi::Object);
};
class Span : public ffi::ObjectRef {
public:
TVM_DLL Span(SourceName source_name, int line, int end_line, int column, int end_column);
/*! \brief Merge two spans into one which captures the combined regions. */
TVM_DLL Span Merge(const Span& other) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Span, ffi::ObjectRef, SpanNode);
};
/*!
* \brief Store a list of spans for an expr generated from mulitple source exprs
*/
class SequentialSpanNode : public SpanNode {
public:
/*! \brief The original source list of spans to construct a sequential span. */
ffi::Array<Span> spans;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SequentialSpanNode>().def_ro("spans", &SequentialSpanNode::spans);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SequentialSpan", SequentialSpanNode, SpanNode);
};
/*!
* \brief Reference class of SequentialSpanNode.
* \sa SequentialSpanNode
*/
class SequentialSpan : public Span {
public:
TVM_DLL SequentialSpan(ffi::Array<Span> spans);
TVM_DLL SequentialSpan(std::initializer_list<Span> init);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SequentialSpan, Span, SequentialSpanNode);
};
/*! \brief A program source in any language.
*
* Could represent the source from an ML framework or a source
* representing a tvm::IRModule.
*/
class Source;
class SourceNode : public ffi::Object {
public:
/*! \brief The source name. */
SourceName source_name;
/*! \brief The raw source. */
ffi::String source;
/*! \brief A mapping of line breaks into the raw source. */
std::vector<std::pair<int, int>> line_map;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SourceNode>()
.def_ro("source_name", &SourceNode::source_name)
.def_ro("source", &SourceNode::source);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Source", SourceNode, ffi::Object);
};
class Source : public ffi::ObjectRef {
public:
TVM_DLL Source(SourceName src_name, std::string source);
TVM_DLL tvm::ffi::String GetLine(int line);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Source, ffi::ObjectRef, SourceNode);
};
/*!
* \brief A mapping from a unique source name to source fragment.
*/
class SourceMap;
/*!
* \brief Stores locations in frontend source that generated a node.
*/
class SourceMapObj : public ffi::Object {
public:
/*! \brief The source mapping. */
ffi::Map<SourceName, Source> source_map;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SourceMapObj>().def_ro("source_map", &SourceMapObj::source_map);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SourceMap", SourceMapObj, ffi::Object);
};
class SourceMap : public ffi::ObjectRef {
public:
explicit SourceMap(ffi::Map<SourceName, Source> source_map);
explicit SourceMap(std::initializer_list<std::pair<SourceName, Source>> source_map)
: SourceMap(ffi::Map<SourceName, Source>(source_map)) {}
SourceMap() : SourceMap(ffi::Map<SourceName, Source>()) {}
void Add(const Source& source);
SourceMapObj* operator->() {
TVM_FFI_ICHECK(get() != nullptr);
return static_cast<SourceMapObj*>(get_mutable());
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SourceMap, ffi::ObjectRef, SourceMapObj);
};
} // namespace tvm
#endif // TVM_IR_SOURCE_MAP_H_
include/tvm/ir/transform.h
+584
View File
@@ -0,0 +1,584 @@
/*
* 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/ir/transform.h
*
* This file implements a pass manager. The pass manager manages a sequence
* of IRModule -> IRModule transformation passes over a particlar unit of AST. The
* design is largely inspired from LLVM's pass manager and modern deep learning
* frameworks that perform tensor->tensor transformations.
*
* The responsibilities of a traditional compiler pass manager usually involves:
* - Organizing the execution order of optimization passes though not
* necessarily in the optimal sequence.
* - Collecting required analysis information and keep them up-to-date.
* - Reducing the effort required to implement new passes for compiler
* developers, etc.
*
* Similar to LLVM's pass manager, we designed the Relax pass manager to work
* different granularity, i.e. module level, function level, and even sequential
* passe that contains a host of passes.
*
* However, we also extend the functionality of the traditional pass manager
* with the consideration of requirements/convention from deep learning
* frameworks, such as Pytorch and Gluon, etc. Each pass in the Relax pass
* manager performs the IRModule -> IRModule transformation. All
* different types of passes, including the sequential-level pass object, are
* essentially pass objects. This design, therefore, effectively provides users
* a consistent and convenient interface, i.e. Pass, to play with. It offers a
* means to ease the development and testing of Relax passes. For example, with
* the pass manager, external users will be able to have custom passes correctly
* scheduled without having to modify a single handcrafted pass order.
*
* In the future we need to describe constraints between passes. For example,
* we may want to preserve dependencies between different passes and validate
* them on the completion of a certain pass.
*
* We also need to store side information and import the error reporting system.
*/
#ifndef TVM_IR_TRANSFORM_H_
#define TVM_IR_TRANSFORM_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/creator.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/instrument.h>
#include <tvm/ir/module.h>
#include <tvm/ir/with_context.h>
#include <string>
#include <type_traits>
#include <utility>
namespace tvm {
namespace transform {
/*!
* \brief PassContextNode contains the information that a pass can rely on,
* such as analysis results.
* \sa PassContext
*/
class PassContextNode : public ffi::Object {
public:
/*! \brief The default optimization level. */
int opt_level{2};
/*! \brief The list of required passes. */
ffi::Array<ffi::String> required_pass;
/*! \brief The list of disabled passes. */
ffi::Array<ffi::String> disabled_pass;
/*! \brief Pass specific configurations. */
ffi::Map<ffi::String, Any> config;
/*! \brief A list of pass instrument implementations. */
ffi::Array<instrument::PassInstrument> instruments;
PassContextNode() = default;
/*!
* \brief Get a config value from the pass context.
*
* \param key The config key.
* \param default_value The default value if the key does not exist, defaults to nullptr.
*
* \return The result
*
* \tparam TOBjectRef the expected object type.
* \throw Error if the key exists but the value does not match TObjectRef.
*/
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetConfig(
const std::string& key,
ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
if (!config.defined()) return default_value;
auto it = config.find(key);
if (it != config.end()) {
return (*it).second.as_or_throw<ffi::Optional<TObjectRef>>();
} else {
return default_value;
}
}
// variant that uses TObjectRef to enable implicit conversion to default value.
template <typename TObjectRef>
ffi::Optional<TObjectRef> GetConfig(const std::string& key, TObjectRef default_value) const {
return GetConfig<TObjectRef>(key, ffi::Optional<TObjectRef>(default_value));
}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PassContextNode>()
.def_ro("opt_level", &PassContextNode::opt_level)
.def_ro("required_pass", &PassContextNode::required_pass)
.def_ro("disabled_pass", &PassContextNode::disabled_pass)
.def_ro("instruments", &PassContextNode::instruments)
.def_ro("config", &PassContextNode::config);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.PassContext", PassContextNode, ffi::Object);
};
/*!
* \brief PassContext that is used to configure the pass behavior.
*
* \code
*
* auto new_ctx = PassContext::Create();
* ctx->opt_level = 2;
* With<PassContext> scope(ctx);
* // pass context in effect.
*
* \endcode
* \sa PassContextNode
*/
class PassContext : public ffi::ObjectRef {
public:
PassContext() {}
/*!
* \brief constructor with UnsafeInit
*/
explicit PassContext(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
/*!
* \brief constructor with ffi::ObjectPtr
*/
explicit PassContext(ffi::ObjectPtr<PassContextNode> n) : ffi::ObjectRef(n) {}
/*!
* \brief const accessor.
* \return const access pointer.
*/
const PassContextNode* operator->() const {
TVM_FFI_ICHECK(get() != nullptr);
return static_cast<const PassContextNode*>(get());
}
/*!
* \brief mutable accessor.
* \return mutable access pointer.
*/
PassContextNode* operator->() {
TVM_FFI_ICHECK(get() != nullptr);
return static_cast<PassContextNode*>(get_mutable());
}
/*!
* \brief Construct a PassContext containing the default configurations.
* \return The new PassContext.
*/
TVM_DLL static PassContext Create();
/*!
* \brief Get the default pass context in the current scope.
* \return The pass context.
*/
TVM_DLL static PassContext Current();
/*!
* \brief Get all supported configuration names and metadata, registered within the PassContext.
* \return Map indexed by the config name, pointing to the metadata map as key-value
*/
TVM_DLL static ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> ListConfigs();
/*!
* \brief Call instrument implementations' callbacks when entering PassContext.
* The callbacks are called in order, and if one raises an exception, the rest will not be
* called.
*/
TVM_DLL void InstrumentEnterPassContext();
/*!
* \brief Call instrument implementations' callbacks when exiting PassContext.
* The callbacks are called in order, and if one raises an exception, the rest will not be
* called.
*/
TVM_DLL void InstrumentExitPassContext();
/*!
* \brief Call instrument implementations' callbacks before a pass run.
* The callbacks are called in order, and if one raises an exception, the rest will not be
* called.
*
* \param mod The module that an optimization pass runs on.
* \param info The pass information.
*
* \return false: the pass is skipped; true: the pass runs.
*/
TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
/*!
* \brief Call instrument implementations callbacks after a pass run.
* The callbacks are called in order, and if one raises an exception, the rest will not be
* called.
*
* \param mod The module that an optimization pass runs on.
* \param info The pass information.
*/
TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
/*!
* \brief Check whether a pass is enabled.
* \param info The pass information.
* \return true if the pass is enabled. Otherwise, false.
*/
TVM_DLL bool PassEnabled(const PassInfo& info) const;
/*!
* \brief Register a valid configuration option and its ValueType for validation.
*
* \param key The configuration key.
* \tparam ValueType The value type to be registered
*/
template <typename ValueType>
static int32_t RegisterConfigOption(const char* key) {
// NOTE: we could further update the function later.
if constexpr (std::is_base_of_v<ffi::ObjectRef, ValueType>) {
int32_t tindex = ffi::TypeToRuntimeTypeIndex<ValueType>::v();
auto type_key = ffi::TypeIndexToTypeKey(tindex);
auto legalization = [=](ffi::Any value) -> ffi::Any {
if (auto opt_map = value.try_cast<ffi::Map<ffi::String, ffi::Any>>()) {
return ffi::reflection::ObjectCreator(type_key)(opt_map.value());
} else {
auto opt_val = value.try_cast<ValueType>();
if (!opt_val.has_value()) {
TVM_FFI_THROW(AttributeError)
<< "Expect config " << key << " to have type " << type_key << ", but instead get "
<< ffi::details::AnyUnsafe::GetMismatchTypeInfo<ValueType>(value);
}
return *opt_val;
}
};
RegisterConfigOption(key, type_key, legalization);
} else {
// non-object type, do not support implicit conversion from map
std::string type_str = ffi::TypeTraits<ValueType>::TypeStr();
auto legalization = [=](ffi::Any value) -> ffi::Any {
auto opt_val = value.try_cast<ValueType>();
if (!opt_val.has_value()) {
TVM_FFI_THROW(AttributeError)
<< "Expect config " << key << " to have type " << type_str << ", but instead get "
<< ffi::details::AnyUnsafe::GetMismatchTypeInfo<ValueType>(value);
} else {
return *opt_val;
}
};
RegisterConfigOption(key, type_str, legalization);
}
return 0;
}
// accessor.
using ContainerType = PassContextNode;
class Internal;
private:
// The entry of a pass context scope.
TVM_DLL void EnterWithScope();
// The exit of a pass context scope.
TVM_DLL void ExitWithScope();
// Register configuration key value type.
TVM_DLL static void RegisterConfigOption(const char* key, ffi::String value_type_str,
std::function<ffi::Any(ffi::Any)> legalization);
// Classes to get the Python `with` like syntax.
friend class Internal;
friend class With<PassContext>;
};
/*!
* \brief Create a pass-config object with all default values, using the
* reflection defaults.
* \tparam TConfig the ObjectRef type to be created.
* \return An instance with all reflection-defined default values applied.
*/
template <typename TConfig>
inline TConfig PassConfigWithDefaults() {
static_assert(std::is_base_of_v<ffi::ObjectRef, TConfig>,
"Can only create ObjectRef-derived types");
using ContainerType = typename TConfig::ContainerType;
static auto finit_object = ffi::Function::GetGlobalRequired("ffi.MakeObjectFromPackedArgs");
ffi::AnyView packed_args[1];
packed_args[0] = ContainerType::RuntimeTypeIndex();
ffi::Any rv;
finit_object.CallPacked(ffi::PackedArgs(packed_args, 1), &rv);
return rv.cast<TConfig>();
}
#define TVM_PASS_CTX_CONFIG_VAR_DEF [[maybe_unused]] static uint32_t __make_PassContext_tid
/*!
* \brief Helper macro to register the object type to runtime.
* Makes sure that the runtime type table is correctly populated.
*
* Use this macro in the cc file for each terminal class.
*/
#define TVM_REGISTER_PASS_CONFIG_OPTION(Key, ValueType) \
TVM_FFI_STR_CONCAT(TVM_PASS_CTX_CONFIG_VAR_DEF, __COUNTER__) = \
::tvm::transform::PassContext::RegisterConfigOption<ValueType>(Key)
/*!
* \brief Meta data that will be used to help optimization and analysis.
* \sa PassInfo
*/
class PassInfoNode : public ffi::Object {
public:
/*! \brief The minimal optimization level that this pass will be enabled. */
int opt_level;
/*! \brief The name of an optimization/analysis pass. */
ffi::String name;
/*! \brief Boolean that tells whether this pass will be traced or not. */
bool traceable;
/*! \brief The passes that are required to perform the current pass. */
ffi::Array<ffi::String> required;
PassInfoNode() = default;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PassInfoNode>()
.def_ro("opt_level", &PassInfoNode::opt_level)
.def_ro("name", &PassInfoNode::name)
.def_ro("required", &PassInfoNode::required)
.def_ro("traceable", &PassInfoNode::traceable);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.PassInfo", PassInfoNode, ffi::Object);
};
/*!
* \brief Managed reference class for PassInfoNode
* \sa PassInfoNode
*/
class PassInfo : public ffi::ObjectRef {
public:
/*!
* \brief Constructor
* \param opt_level The optimization level
* \param name Name of the pass.
* \param required The passes that are required to perform the current pass.
* \param traceable Boolean that tells whether the pass is traceable.
*/
TVM_DLL PassInfo(int opt_level, ffi::String name, ffi::Array<ffi::String> required,
bool traceable);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInfo, ffi::ObjectRef, PassInfoNode);
};
/*!
* \brief PassNode is the base type of differnt types of optimization passes.
* It is designed as a pure class and implemented by different pass subclasses
* at different granularity of Relax nodes.
*/
class PassNode : public ffi::Object {
public:
virtual ~PassNode() {}
/*!
* \brief Get the pass information/meta data. */
virtual PassInfo Info() const = 0;
/*!
* \brief Transform mod using the default PassContext in the current scope.
*
* \param mod The module that an optimization pass runs on.
*
* \return The transformed module.
*/
IRModule operator()(IRModule mod) const {
return this->operator()(std::move(mod), PassContext::Current());
}
/*!
* \brief Transform mod using a functor under a given pass context.
*
* \param mod The module that an optimization pass runs on.
* \param pass_ctx The pass context that can provide information for the optimization.
*
* \return The transformed module.
*/
virtual IRModule operator()(IRModule mod, const PassContext& pass_ctx) const = 0;
TVM_FFI_DECLARE_OBJECT_INFO("transform.Pass", PassNode, ffi::Object);
};
class Pass : public ffi::ObjectRef {
public:
/*!
* \brief Transform mod using the default PassContext in the current scope.
*
* \code
*
* // If you do no longer need the input module
* // it is recommended to use std::move to move your input module.
* mod = pass(std::move(mod));
*
* \endcode
*
* \param mod The module that an optimization pass runs on.
*
* \return The transformed module.
*/
IRModule operator()(IRModule mod) const;
/*!
* \brief Transform mod using a functor under a given pass context.
*
* \param mod The module that an optimization pass runs on.
* \param pass_ctx The pass context that can provide information for the optimization.
*
* \return The transformed module.
*/
IRModule operator()(IRModule mod, const PassContext& pass_ctx) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Pass, ffi::ObjectRef, PassNode);
private:
IRModule static AssertImmutableModule(const IRModule& mod, const PassNode* node,
const PassContext& pass_ctx);
};
/*!
* \brief The SequentialNode contains a set of passes that transform Relax
* programs from one AST to another semantically equivalent one.
*
* One example of this level of pass is that the pass manager needs to correctly
* perform a host of optimizations with a given optimization level and disabled
* passes.
*/
class SequentialNode : public PassNode {
public:
/* \brief The pass meta data.*/
PassInfo pass_info;
/*! \brief A list of passes that used to compose a sequential pass. */
tvm::ffi::Array<Pass> passes;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SequentialNode>()
.def_ro("pass_info", &SequentialNode::pass_info)
.def_ro("passes", &SequentialNode::passes);
}
/*!
* \brief Get the pass information/meta data.
*/
PassInfo Info() const override { return pass_info; }
/*!
* \brief Resolve the pass dependency. It globs all required passes by
* a given pass and executes them.
*
* \param mod The module that an optimization pass runs on.
*
* TODO(zhiics) Build a dependency graph among the passes using provided
* metadata, i.e. required_passes. Likely, we can have a data structure, i.e.
* PassInfo, to store the relevant information including the parent passes.
*/
void ResolveDependency(const IRModule& mod);
/*!
* \brief Perform optimizations on a series of passes. The aforementioned
* typical pass manager jobs could be done by it. This function could
* be overloaded to focus on different metrics, i.e. performance,
* memory footprint, etc.
*
* \param mod The module that these passes are applied on.
* \param pass_ctx The context that these passes execute on.
*
* \return Return the updated module.
*/
IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.Sequential", SequentialNode, PassNode);
};
class Sequential : public Pass {
public:
/*!
* \brief The constructor of `Sequential`.
*
* \param passes The passes to apply.
* \param pass_info The pass metadata.
*/
TVM_DLL Sequential(ffi::Array<Pass> passes, PassInfo pass_info);
/*!
* \brief The constructor of `Sequential`.
*
* \param passes The passes to apply.
* \param name The name of a sequential pass. It's defaulted to "sequential".
* This allows users to only provide a list of passes and execute them
* under a given context.
*/
TVM_DLL Sequential(ffi::Array<Pass> passes, ffi::String name = "sequential");
Sequential() = default;
explicit Sequential(ffi::ObjectPtr<SequentialNode> n) : Pass(n) {}
const SequentialNode* operator->() const;
using ContainerType = SequentialNode;
};
/*
* \brief Create a module pass.
*
* \param pass_func The packed function that contains the optimization.
* \param opt_level The optimization level of the module pass.
* \param name The name of the module pass.
* \param required The list of the passes that the module pass is dependent on.
*
* \return The created module pass.
*/
TVM_DLL Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
int opt_level, ffi::String name, ffi::Array<ffi::String> required,
bool traceable = false);
/*!
* \brief A special trace pass that prints the header and IR to LOG(INFO).
* \param header The header to be attached to the output.
* \return The pass.
*/
TVM_DLL Pass PrintIR(ffi::String header = "");
/*!
* \brief Enrich a pass-time error with a TVMScript-rendered, underlined source
* location derived from the error's embedded VisitErrorContext.
*
* Returns an ffi::Error that preserves err's kind, original message, and
* backtrace, and appends the failing pass name plus the offending location
* rendered as TVMScript (the whole \p mod, or local to \p func when provided).
* The returned error drops the VisitErrorContext payload, so an outer catch
* that re-enriches finds no context and returns the error unchanged.
*
* Pure and total: never throws; returns \p err unchanged when there is no
* context, the path is unresolvable, or rendering fails.
*
* \param err The error thrown by the pass body.
* \param mod The IRModule the pass ran on (the access-path root, or the
* container of \p func when \p func is provided).
* \param pass_name The name of the failing pass, shown in the message.
* \param func When set, resolve and render the location local to
* \p mod->functions[func]; otherwise use the whole module.
* \return The enriched (or, on any fallback, the original) error.
*/
TVM_DLL ffi::Error EnrichPassErrorWithContext(
const ffi::Error& err, const IRModule& mod, ffi::String pass_name,
ffi::Optional<GlobalVar> func = ffi::Optional<GlobalVar>(std::nullopt));
} // namespace transform
} // namespace tvm
#endif // TVM_IR_TRANSFORM_H_
include/tvm/ir/type.h
+207
View File
@@ -0,0 +1,207 @@
/*
* 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/ir/type.h
* \brief IR/AST nodes for TVM types shared across IR variants.
*/
#ifndef TVM_IR_TYPE_H_
#define TVM_IR_TYPE_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/base_expr.h>
#include <tvm/ir/source_map.h>
#include <string>
namespace tvm {
/*!
* \brief Low-level raw pointer type.
*
* PointerType represents type hints in the TIR to be
* passed to the final code generator.
*
* PointerType should not occur in the high-level analysis.
*
* \sa PointerType
*/
class PointerTypeNode : public TypeNode {
public:
/*!
* \brief The type of the element which the pointer points to.
*/
Type element_type = PrimType::Void();
/*!
* \brief The storage scope of the pointer
*/
ffi::String storage_scope;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PointerTypeNode>()
.def_ro("element_type", &PointerTypeNode::element_type)
.def_ro("storage_scope", &PointerTypeNode::storage_scope);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PointerType", PointerTypeNode, TypeNode);
};
/*
* \brief Managed reference to PointerTypeNode.
* \sa PointerTypeNode
*/
class PointerType : public Type {
public:
/*!
* \brief Constructor
* \param element_type The type of the element which the pointer points to.
* \param storage_scope The storage scope into which the pointer addresses
*/
TVM_DLL explicit PointerType(Type element_type, ffi::String storage_scope = "");
/*! \brief Construct an opaque pointer with void element type. */
TVM_DLL static PointerType VoidPointerTy(ffi::String storage_scope = "");
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PointerType, Type, PointerTypeNode);
};
/*!
* \brief The type of tuple values.
* \sa TupleType
*/
class TupleTypeNode : public TypeNode {
public:
/*! \brief The type of each field in the tuple. */
ffi::Array<Type> fields;
TupleTypeNode() {}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TupleTypeNode>().def_ro("fields", &TupleTypeNode::fields);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TupleType", TupleTypeNode, TypeNode);
};
/*!
* \brief Managed reference to TupleTypeNode.
* \sa TupleTypeNode.
*/
class TupleType : public Type {
public:
/*!
* \brief Constructor
* \param fields Fields in the tuple.
* \param span The span of the type.
*/
TVM_DLL explicit TupleType(ffi::Array<Type> fields, Span span = Span());
/*!
* \brief Create an empty tuple type that constains nothing.
* \return A empty tuple type.
*/
TVM_DLL TupleType static Empty();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleType, Type, TupleTypeNode);
};
/*!
* \return a type that represents void.
*/
inline Type VoidType() { return TupleType::Empty(); }
/*!
* \brief Check whether the tyep represents void.
* \return The check result.
*/
inline bool IsVoidType(const Type& type) {
auto* n = type.as<TupleTypeNode>();
return n && n->fields.size() == 0;
}
/*!
* \brief Function type.
*
* We support polymorphic function type.
* This can be roughly viewed as template function in C++.
*
* \sa FuncType, TypeConstraint
*/
class FuncTypeNode : public TypeNode {
public:
/*! \brief type type of arguments */
ffi::Array<Type> arg_types;
/*! \brief The type of return value. */
Type ret_type = VoidType();
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FuncTypeNode>()
.def_ro("arg_types", &FuncTypeNode::arg_types)
.def_ro("ret_type", &FuncTypeNode::ret_type);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.FuncType", FuncTypeNode, TypeNode);
};
/*!
* \brief Managed reference to FuncTypeNode.
* \sa FuncTypeNode
*/
class FuncType : public Type {
public:
/*!
* \brief Constructor
* \param arg_types The types of the arguments.
* \param ret_type The type of the return value.
* \param span The span information.
* \sa FuncTypeNode for more docs about these fields.
*/
TVM_DLL FuncType(ffi::Array<Type> arg_types, Type ret_type, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
};
/*!
* \brief The type of tensor map.
* \sa TensorMapType
*/
class TensorMapTypeNode : public TypeNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TensorMapTypeNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TensorMapType", TensorMapTypeNode, TypeNode);
};
/*!
* \brief Managed reference to TensorMapTypeNode.
* \sa TensorMapTypeNode
*/
class TensorMapType : public Type {
public:
TVM_DLL TensorMapType(Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorMapType, Type, TensorMapTypeNode);
};
} // namespace tvm
#endif // TVM_IR_TYPE_H_
include/tvm/ir/unique_name_supply.h
+143
View File
@@ -0,0 +1,143 @@
/*
* 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/ir/unique_name_supply.h
* \brief UniqueNameSupply that can be used to generate unique variable names.
*/
#ifndef TVM_IR_UNIQUE_NAME_SUPPLY_H_
#define TVM_IR_UNIQUE_NAME_SUPPLY_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <cstdint>
#include <string>
#include <utility>
namespace tvm {
/*!
* \brief UniqueNameSupply can be used to generate unique names.
*/
class UniqueNameSupplyNode : public ffi::Object {
public:
/*!
* \brief Empty constructor. Needed by the TVM_REGISTER_NODE_TYPE macro.
*/
UniqueNameSupplyNode() = default;
/*!
* \brief Constructor.
* \param prefix The prefix to be used with this UniqueNameSupply.
* \param name_map The map used to guarantee uniqueness.
*/
UniqueNameSupplyNode(const ffi::String& prefix, ffi::Map<ffi::String, int64_t> name_map)
: prefix_(prefix), name_map(std::move(name_map)) {}
/*!
* \brief Generates a unique name from this UniqueNameSupply.
* \param name The name from which the generated name is derived.
* \param add_prefix If set to true, then the prefix of this UniqueNameSupply will be prepended to
* the name.
* \param add_underscore If set to true, add '_' between prefix and a digit.
* \return A unique name.
*/
ffi::String FreshName(const ffi::String& name, bool add_prefix = true,
bool add_underscore = true);
/*!
* \brief Reserves an existing name with this UniqueNameSupply.
* \param name The name to be reserved.
* \param add_prefix If set to true, then the prefix of this UniqueNameSupply will be prepended to
* the name before reserving it. \return The name that was reserved with the UniqueNameSupply. It
* can be different if a prefix is added.
*/
ffi::String ReserveName(const ffi::String& name, bool add_prefix = true);
/*!
* \brief Checks if this UniqueNameSupply already generated a name.
* \param name The name to check.
* \param add_prefix If set to true, then the prefix of this UniqueNameSupply will be prepended to
* the name before checking for it. \return True if the name has already been generated. False
* otherwise.
*/
bool ContainsName(const ffi::String& name, bool add_prefix = true);
// Prefix for all GlobalVar names. It can be empty.
std::string prefix_;
static constexpr const bool _type_mutable = true;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<UniqueNameSupplyNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.UniqueNameSupply", UniqueNameSupplyNode, ffi::Object);
private:
/*! \brief Helper function to add the UniqueNameSupply prefix to the name. */
ffi::String AddPrefixToName(const ffi::String& name);
/*!
* \brief Function that will generate a unique name.
* \param name The name to be used as a base.
* \param add_underscore If set to true, add '_' between prefix and a digit.
* \return A unique name.
*/
std::string GetUniqueName(std::string name, bool add_underscore = true);
/*! \brief A map that is used to generate unique names. */
ffi::Map<ffi::String, int64_t> name_map;
};
/*!
* \brief Managed reference class to UniqueNameSupplyNode.
* \sa UniqueNameSupplyNode
*/
class UniqueNameSupply : public ffi::ObjectRef {
public:
/*!
* \brief Constructor.
* \param prefix The prefix to be used with this UniqueNameSupply.
* \param name_map An optional map.
*/
TVM_DLL explicit UniqueNameSupply(const ffi::String& prefix = "",
ffi::Map<ffi::String, int64_t> name_map = {});
/*!
* \brief Construct UniqueNameSupply by reserving names from the given iterator range.
*
* The functor should return the name of the dereferenced object.
*/
template <typename Iter, typename Lambda>
TVM_DLL UniqueNameSupply(Iter begin, Iter end, Lambda f) : UniqueNameSupply("") {
for (auto it = begin; it != end; ++it) {
this->operator->()->ReserveName(f(*it), false);
}
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(UniqueNameSupply, ffi::ObjectRef,
UniqueNameSupplyNode);
};
} // namespace tvm
#endif // TVM_IR_UNIQUE_NAME_SUPPLY_H_
include/tvm/ir/with_context.h
+168
View File
@@ -0,0 +1,168 @@
/*
* 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/ir/with_context.h
* \brief RAII wrapper function to enter and exit a context object
* similar to python's with syntax.
*/
#ifndef TVM_IR_WITH_CONTEXT_H_
#define TVM_IR_WITH_CONTEXT_H_
#include <exception>
#include <functional>
#include <memory>
#include <utility>
#include <vector>
namespace tvm {
/*!
* \brief RAII wrapper function to enter and exit a context object
* similar to python's with syntax.
*
* \code
* // context class
* class MyContext {
* private:
* friend class With<MyContext>;
MyContext(arguments);
* void EnterWithScope();
* void ExitWithScope();
* };
*
* {
* With<MyContext> scope(arguments);
* // effect take place.
* }
* \endcode
*
* \tparam ContextType Type of the context object.
*/
template <typename ContextType>
class With {
public:
/*!
* \brief constructor.
* Enter the scope of the context.
*/
template <typename... Args>
explicit With(Args&&... args) : ctx_(std::forward<Args>(args)...) {
ctx_.EnterWithScope();
}
/*! \brief destructor, leaves the scope of the context. */
~With() noexcept(false) { ctx_.ExitWithScope(); }
// Disable copy and move construction. `With` is intended only for
// use in nested contexts that are exited in the reverse order of
// entry. Allowing context to be copied or moved would break this
// expectation.
With(const With& other) = delete;
With& operator=(const With& other) = delete;
With(With&& other) = delete;
With& operator=(With&& other) = delete;
ContextType* get() { return &ctx_; }
const ContextType* get() const { return &ctx_; }
ContextType* operator->() { return get(); }
const ContextType* operator->() const { return get(); }
ContextType& operator*() { return *get(); }
const ContextType* operator*() const { return *get(); }
ContextType operator()() { return ctx_; }
private:
/*! \brief internal context type. */
ContextType ctx_;
};
/*!
* \brief A group of RAII contexts managed together.
*
* Allows dynamically emplacing multiple context objects that are
* all exited (in reverse order) when the group is destroyed.
* ContextType must declare `friend class With<ContextType>`
* and provide EnterWithScope() / ExitWithScope() methods.
*
* \code
* WithGroup<ConstraintContext> group;
* group.Emplace(analyzer, cond1); // constructs and enters
* group.Emplace(analyzer, cond2); // constructs and enters
* // destructor: exits cond2, then cond1
* \endcode
*
* \tparam ContextType The context type with EnterWithScope/ExitWithScope.
*/
template <typename ContextType>
class WithGroup {
public:
WithGroup() = default;
WithGroup(WithGroup&&) = default;
WithGroup& operator=(WithGroup&&) = default;
WithGroup(const WithGroup&) = delete;
WithGroup& operator=(const WithGroup&) = delete;
/*!
* \brief Construct a context and enter its scope.
* \param args Arguments forwarded to ContextType constructor.
*/
template <typename... Args>
void Emplace(Args&&... args) {
entries_.push_back(std::make_unique<With<ContextType>>(std::forward<Args>(args)...));
}
/*! \brief Number of active contexts in this group. */
size_t size() const { return entries_.size(); }
/*!
* \brief Destructor — exits all contexts in reverse order.
*
* On normal exit: if any ExitWithScope throws, the remaining
* contexts are still cleaned up, then the first exception
* is re-thrown.
*
* During stack unwinding: all exceptions are swallowed
* to avoid std::terminate.
*/
~WithGroup() noexcept(false) {
bool unwinding = std::uncaught_exceptions() > 0;
std::exception_ptr first_exc;
while (!entries_.empty()) {
// Move the last entry out of the vector first, then destroy it.
// This ensures entries_ shrinks even if ~With() throws.
auto entry = std::move(entries_.back());
entries_.pop_back();
try {
entry.reset(); // calls ~With<ContextType>() -> ExitWithScope()
} catch (...) {
if (!unwinding && !first_exc) {
first_exc = std::current_exception();
}
}
}
if (first_exc) std::rethrow_exception(first_exc);
}
private:
std::vector<std::unique_ptr<With<ContextType>>> entries_;
};
} // namespace tvm
#endif // TVM_IR_WITH_CONTEXT_H_
include/tvm/relax/analysis.h
+685
View File
@@ -0,0 +1,685 @@
/*
* 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/relax/analysis.h
* \brief The set of Relax specific analysis on IR.
*/
#ifndef TVM_RELAX_ANALYSIS_H_
#define TVM_RELAX_ANALYSIS_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
#include <tvm/relax/type.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/index_map.h>
#include <functional>
#include <set>
#include <utility>
namespace tvm {
namespace relax {
//-----------------------------------
// Shape expression analysis
//----------------------------------
/*!
* \brief Can prove the two symbolic shape arrays equals to each other.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \param ana The analyzer used for integer analysis.
* \return The prove result.
*
* \note This function does best effort prove, which means
* if result is false, there is still possibility that
* two shapes equals to each other during runtime.
*/
TVM_DLL bool CanProveShapeEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Array<PrimExpr>& rhs,
const arith::Analyzer& ana);
/*!
* \brief Can prove the two symbolic shape expressions equals to each other.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \param ana The analyzer used for integer analysis.
*
* \note This function does best effort prove, which means
* if result is false, there is still possibility that
* two shapes equals to each other during runtime.
*/
TVM_DLL bool CanProveShapeEqual(const Expr& lhs, const Expr& rhs, const arith::Analyzer& ana);
//-----------------------------------
// Foundational Type analysis
//-----------------------------------
/*!
* \brief Get the corresponding static type from a given type.
* \param info The type.
* \return the corresponding static type.
*/
TVM_DLL Type GetStaticType(const Type& info);
/*!
* \brief Get the corresponding type from static type.
* \param type The input type
* \return the corresponding type.
*/
TVM_DLL Type TypeFromStaticType(const Type& type);
/*!
* \return Derive the call's ret value type from inputs.
* \param finfo The function type.
* \param call The call expression to be derived.
* \param ctx The builder context.
* \return The derived type of the call.
* \note call->op field is ignored during derivation and we only rely on information
* presented by func_ty.
*/
TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx);
/*!
* \brief Derive the call's ret value type using a caller-provided analyzer.
* \param finfo The function type.
* \param call The call expression to be derived.
* \param ctx The builder context.
* \param ana Context analyzer to prove symbolic expression equality.
* \return The derived type of the call.
*/
TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx,
const arith::Analyzer& ana);
/*!
* \brief Erase the info to a corresponding more coarse grained
* type that is still well-defined(with all the vars in scope).
*
* When we are returning a Type to another scope,
* it is important to remember that Type may carry
* dependencies on var that is not defined the other scope.
*
* In such cases, it is important to call EraseToWellDefined to get
* another Type that **only** contains the vars that are defined
* in the target scope.
*
* For example, consider the following function
*
* \code
*
* @R.function
* def f(x: R.Tensor[(n, m)]):
* k = tirx.Var("k", "int64")
* v0 = opaque_fn(x)
* v1 = match_cast(v0, R.Tensor[(n, k)])
* v2 : R.Tensor[(n + 1, k + 2)] = pad(v1)
* return v2
*
* \endcode
*
* In the above code, the return value y have shape `(n + 1, k + 2)`,
* However, at the level of function signature, only n, m are defined,
* k is undefined here.
*
* When we call EraseToWellDefined(R.Tensor[(n + 1, k + 2)], fshape_var_map={n: n, m: m}),
* we will obtain R.Tensor(ndim=2), which is an erased info that does not depend
* on k(which is undefined from parameter signature).
*
* However, if we call EraseToWellDefined(R.Tensor[(n + 1, m)], fshape_var_map={n: n, m: m}),
* Then the return value will be R.Tensor[(n + 1, m)], because both n and m are defined.
*
* We can also make these var map to return a different expression.
* For example, EraseToWellDefined(R.Tensor[(n + 1, m)], fshape_var_map={n: 2, m: m})
* will give us R.Tensor[(3, m)], where n get replaced by 2.
*
* Use this function in the following scenarios:
* - Decide the ty of expr with sub-scopes, such as If, SeqExpr
* - Decide the deduced return ty of a function that can be fully decided by params.
*
* \param info The type.
* \param f_shape_var_map callback function to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param f_var_map callback function to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
*
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(
const Type& info,
std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map = nullptr,
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map = nullptr);
/*!
* \brief EraseToWellDefined overload using a caller-provided analyzer.
* \param info The type.
* \param f_shape_var_map callback function to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param f_var_map callback function to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
* \param ana Context analyzer to prove symbolic expression equality.
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(
const Type& info, std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map,
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map, const arith::Analyzer& ana);
/*!
* \brief EraseToWellDefined variant with map.
* \param info The type.
* \param shape_var_map map to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param var_map map to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
*
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
ffi::Map<Var, Expr> var_map);
/*!
* \brief EraseToWellDefined map overload using a caller-provided analyzer.
* \param info The type.
* \param shape_var_map map to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param var_map map to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
* \param ana Context analyzer to prove symbolic expression equality.
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
ffi::Map<Var, Expr> var_map, const arith::Analyzer& ana);
/*!
* \brief Fine grained result of base check.
*
* This analysis comes with different levels of checking failures
* that can help to customize the compilation decisions.
*
* For a given pair of lhs_ty, rhs_ty. We adopt
* the following terminology:
* - LSet = {value | value matches lhs_ty}
* - RSet = {value | value matches rhs_ty}
*
* See the definition of each level below.
*/
enum class BaseCheckResult {
/*!
* \brief The two value sets have no intersection at all: Interset(LSet, RSet) = empty
*/
kFailL0 = 0,
/*!
* \brief LSet is not superset of RSet by only looking at static information.
*
* \note This level will trigger static type checking error when lhs is param and rhs is arg.
*/
kFailL1 = 1,
/*!
* \brief WLSet is not superset of RSet because of mismatch in value information.
*
* L1-level mismatches in params of FuncType is categorized as
* If lhs is FuncType, then L1-level mismatch in its params
* is categorized as L2-level mismatch for lhs.
*
* Design considerations for functions:
* - (a) We want to be able to erase type/value in function signature
* when we unify function type and preserve simpler representations.
* - (b) We automatically insert match_cast at function boundary, so
* we can erase (int)->int argument as (object)->int.
* The input shape/type mismatch will be detected by runtime checks at function boundary.
* This behavior is also consistent with the ffi::Function behavior.
*
* \note This level means there is no problem about static known information.
* It is OK for the checker to do best effort and return this value.
*/
kFailL2 = 2,
/*! \brief LSet is superset of RSet. */
kPass = 3
};
/*!
* \brief Run a base check to see if base subsumes derived.
*
* This function returns fine-grained base-check result on reasons of failure.
*
* \param base The base type.
* \param derived The derived type.
* \return Whether the relation holds.
*
* \sa BaseCheckResult
*/
TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived);
/*!
* \brief Run a base check using a caller-provided analyzer.
* \param base The base type.
* \param derived The derived type.
* \param ana Context analyzer to prove symbolic expression equality.
* \return Whether the relation holds.
*
* \sa BaseCheckResult
*/
TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived,
const arith::Analyzer& ana);
/*!
* \brief Check the relation of two type to see if one subsumes another one.
*
* \param base The base type.
* \param derived The derived type.
* \return Whether the relation holds.
*/
TVM_DLL bool IsBaseOf(const Type& base, const Type& derived);
/*!
* \brief Check whether one type subsumes another using a caller-provided analyzer.
* \param base The base type.
* \param derived The derived type.
* \param ana Context analyzer to prove symbolic expression equality.
* \return Whether the relation holds.
*/
TVM_DLL bool IsBaseOf(const Type& base, const Type& derived, const arith::Analyzer& ana);
/*!
* \brief Return the condition for which base is a superset of derived
*
* This function returns finer-grained conditions for kFailL2 than TypeBaseCheck
*
* If the returned expression is true, or simplifies to true, then
* base is a superset of derived. If the returned expression is
* false, or simplifies to false, then base is not a superset of
* derived.
*
* If the returned expression is neither true nor false, it is an
* expression in terms of the symbolic variables available in `base`
* and `derived`.
*
* \param base The base type.
* \param derived The derived type.
* \return Whether base is a base of derived.
*
* \sa BaseCheckResult
*/
TVM_DLL PrimExpr TypeBaseCheckPrecondition(const Type& base, const Type& derived);
/*!
* \brief Unify the two type to their least common ancestor.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \return The unified information.
*/
TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs);
/*!
* \brief Unify two types using a caller-provided analyzer.
* \param lhs The left operand.
* \param rhs The right operand.
* \param ana Context analyzer to prove symbolic expression equality.
* \return The unified information.
*/
TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs, const arith::Analyzer& ana);
/*!
* \brief Get the TIR variables that appear in the input type.
* The returned list is deduplicated - each TIR variable will appear at most once.
* \param ty The type object to be analyzed.
* \return The list of TIR variables that appear in the input type.
*/
TVM_DLL ffi::Array<tirx::Var> TIRVarsInType(const Type& ty);
/*!
* \brief Get the TIR variables that appear in the input type.
*
* Returns all symbolic variables that are definable based on, and
* used within, the Type.
*
* \param ty The type object to be analyzed.
*
* \return A tuple of (definable,used) TIR variables. Both lists are
* deduplicated, each TIR variable will appear at most once, and in
* order of occurrence.
*/
TVM_DLL ffi::Array<tirx::Var> DefinableTIRVarsInType(const Type& ty);
/*! \brief Collect expressions whose usage requires them to be non-negative
*
* Any PrimExpr that is used as a tensor shape, or as an element in a
* ShapeExpr, may not be negative. This utility function can be used
* to generate assertions prior to calling a kernel, or to provide
* assumptions within a kernel that may be useful for simplification.
*
* \param ty The type to be analyzed
*
* \return A list of non-negative expressions.
*/
TVM_DLL ffi::Array<PrimExpr> CollectNonNegativeExpressions(const Type& ty);
/*!
* \brief Get the TIR variables that defined in the input function.
* The returned list is deduplicated - each TIR variable will appear at most once.
* \param expr The relax expression (e.g. a Function) to be analyzed.
* \return The list of TIR variables that are defined in the input function.
*/
TVM_DLL ffi::Array<tirx::Var> DefinedSymbolicVars(const Expr& expr);
/*!
* \brief Get the TIR variables that are used but not defined in the input function.
* The returned list is deduplicated - each TIR variable will appear at most once.
* \param expr The relax expression (e.g. a Function) to be analyzed.
* \return The list of TIR variables that are used but not defined in the input function.
*/
TVM_DLL ffi::Array<tirx::Var> FreeSymbolicVars(const Expr& expr);
//-----------------------------------
// General IR analysis
//-----------------------------------
/*!
* \brief Get all bound variables from expression expr.
*
* Bound variables are all variables that are declared in the expr.
* They only have meaning inside that expr, and can only be used in it.
*
* \param expr the expression.
*
* \return List of bound vars, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<Var> BoundVars(const Expr& expr);
/*!
* \brief Get free type parameters from expression expr.
*
* Free variables are variables that are not bound by a
* varbinding or a function parameter in the context.
*
* \param expr the expression.
*
* \return List of free vars, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<Var> FreeVars(const Expr& expr);
/*!
* \brief Get all variables from expression expr.
*
* \param expr the expression.
*
* \return List of all vars, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<Var> AllVars(const Expr& expr);
/*!
* \brief Get all global variables from expression expr.
*
* AllVars is a superset of BoundVars and FreeVars.
* The union of BoundVars and FreeVars is Allvars.
*
* \param expr the expression.
*
* \return List of all global variables, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<GlobalVar> AllGlobalVars(const Expr& expr);
/*!
* \brief Find all sets of recursive or mutually recursive functions in the module.
*
* Two or more functions are mutually recursive if there is some cycle of references
* among them. For example, if there are two functions A and B, they are
* mutually recursive if A calls B and B calls A. Another case would be with
* three functions A, B, and C, where A calls B, B calls C, and C calls A.
*
* (Note that functions do not have to call each other to reference each other.
* For example, if a function returns another function, that is still a reference
* that could potentially be recursive, even without a call.)
*
* If a function is simply recursive and not mutually recursive with any other,
* it will be reported as a group by itself.
*
* \param m The module
*
* \return List of all groups of mutually recursive functions.
* Each member of the result is a list of functions in the module
* that are all mutually recursive.
* If a function is simply recursive and not mutually recursive with any other,
* then it will be listed as a group by itself.
*/
TVM_DLL tvm::ffi::Array<tvm::ffi::Array<GlobalVar>> DetectRecursion(const IRModule& m);
/*!
* \brief Analyze var -> value mapping from VarBindings.
*
* \param m The IRModule to check.
* \return Var -> Value (Expr)
*/
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const IRModule& m);
/*!
* \brief Analyze var -> value mapping from VarBindings.
*
* \param expr The expression to check.
* \return Var -> Value (Expr)
*/
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const Expr& expr);
/*!
* \brief Analyze var -> value mapping from VarBindings.
*
* \param dfb The dataflow block to check.
* \return Var -> Value (Expr)
*/
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const DataflowBlock& dfb);
/*!
* \brief Return a mapping from variable name to its Bindings.
*
* \param fn The function to be analyzed.
* \return A mapping from variable name to its Bindings.
*/
TVM_DLL ffi::Map<ffi::String, ffi::Array<Binding>> NameToBinding(const Function& fn);
/*!
* \brief Get the use-def chain of variables inside a dataflow block.
*
* \param dfb The dataflow block to be analyzed.
* \return A map mapping variable definitions to a set of uses.
*/
TVM_DLL ffi::Map<Var, ffi::Array<Var>> DataflowBlockUseDef(const DataflowBlock& dfb);
/*!
* \brief Get the use-def chain of variables inside a function.
*
* \param expr The expression to be analyzed.
*
* \return A tuple of variable usage and variable outputs. The first
* element is a map from variable definitions to the set of downstream
* users of that definition. The second element is a list of
* variables whose usage occurs outside of any variable binding,
* typically the output body of a relax::Function or a relax::SeqExpr.
*/
std::pair<ffi::Map<Var, ffi::Array<Var>>, ffi::Array<Var>> FunctionUseDef(const Expr& expr);
/*! \brief A utility struct returned by CollectVarUsage
*/
struct VarUsageInfo {
/* \brief A map from variables to the bound expression.
*
* This is equivalent to the output of AnalyzeVar2Value
*/
ffi::Map<Var, Expr> bound_values;
/* \brief The map from variables to downstream usages of the variable
*
* This is equivalent to the first output of FunctionUseDef.
*/
ffi::Map<Var, ffi::Array<Var>> downstream_usage;
/* \brief A list of variables produced as output
*
* This is equivalent to the second output of FunctionUseDef
*/
ffi::Array<Var> outputs;
};
/*! \brief Collect variable bindings and usage
*
* This function is equivalent to calling both FunctionUseDef and
* AnalyzeVar2Value, but requires only a single traversal of the
* expression.
*
* \param expr The expression to analyze
*
* \return The collected information
*/
VarUsageInfo CollectVarUsage(const Expr& expr);
/*!
* \brief Get the used variables in an expression.
*
* This function collects all variables that are referenced within the given expression.
*
* \param expr The expression to analyze
*
* \return A set of variable nodes that are used in the expression
*/
TVM_DLL std::set<const VarNode*> GetUsedVars(const Expr& expr);
/*!
* \brief Remove unused statements inside DataflowBlocks.
*
* \param expr The expression (typically a relax::Function) from which
* to remove unused statements.
*
* \return The updated function with no unused statements in DataflowBlock.
*/
TVM_DLL Expr RemoveAllUnused(Expr expr);
/*!
* \brief Annotate Op Pattern Kind for PrimFunc, which is used in relax FuseOps.
*
* \param func The PrimFunc to be analyzed.
* \return The Op Pattern Kind.
*
* \note This analysis applies on TIR function but is primarily used by relax passes.
* As a result we place it under the relax namespace.
*/
TVM_DLL OpPatternKind AnalyzeOpPatternKind(const tirx::PrimFunc& func);
/*!
* \brief Check if the given PrimFunc is essentially doing a reshape operation.
* The reshape operation also includes expand_dims, squeeze, flatten, etc.
* \details Here the allowed reshape pattern is: for example, assume the operation is
* `B[l_0, l_1, ..., l_b] = A[r_0, r_1, ..., r_a]`, we check if we can prove that the flattened
* index of l_0, ..., l_b under buffer B equals to the flattened index of r_0, ..., r_a under
* buffer A.
* \param func The function to be examined.
* \return A boolean indicating if the given PrimFunc is doing a reshape.
* \note According to the description above, the returned result can only be false-negative and
* cannot be false-positive, since whenever we cannot prove the equality, we return false. This
* property guarantees the safety of this function.
*/
TVM_DLL bool HasReshapePattern(const tirx::PrimFunc& func);
/*!
* \brief Check if the given expression (likely a function body) contains any impure calls.
* \param expr The expression to be examined. If expr is a function, we check the body.
* \param own_name (Optional.) If we are checking a recursive function body,
* the caller can pass the function's name so recursive calls
* can be ignored in the check (must be a Var or GlobalVar).
* \return The impure expression, if one exists within the given
* expression. Otherwise, std::nullopt.
* \note Relies on Type annotations, so ensure that the module has been normalized first.
* Also, an impure call in a *nested* function does *not* mean that the outer expression contains
* an impure call--it only does if the nested function is *later called*.
*/
TVM_DLL ffi::Optional<Expr> FindImpureCall(
const Expr& expr, const ffi::Optional<Expr>& own_name = ffi::Optional<Expr>(std::nullopt));
/*!
* \brief Check if the given expression (likely a function body) contains any impure calls.
* \param expr The expression to be examined. If expr is a function, we check the body.
* \param own_name (Optional.) If we are checking a recursive function body,
* the caller can pass the function's name so recursive calls
* can be ignored in the check (must be a Var or GlobalVar).
* \return A boolean indicating if the expression contains any impure calls.
* \note Relies on Type annotations, so ensure that the module has been normalized first.
* Also, an impure call in a *nested* function does *not* mean that the outer expression contains
* an impure call--it only does if the nested function is *later called*.
*/
TVM_DLL bool ContainsImpureCall(
const Expr& expr, const ffi::Optional<Expr>& own_name = ffi::Optional<Expr>(std::nullopt));
/*!
* \brief Check if an IRModule or Function is well-formed.
*
* Throws an ffi::Error on the first well-formedness violation. The error is
* seeded with the offending node so a pass runner can resolve a precise access
* path. Use \ref CheckWellFormed for a boolean answer.
*
* \param obj The IRModule or relax::Function to check.
* \param check_ty If true, verify that every Expr has ty populated.
* \note By default the type information is always checked. It is only in test cases
* where `check_ty` might be false, so that other well-formed requirements
* will be well tested and will not be blocked by not having type information.
*/
TVM_DLL void WellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
/*!
* \brief Return whether an IRModule or Function is well-formed.
*
* Wraps \ref WellFormed, returning false instead of throwing on the first
* violation.
*
* \param obj The IRModule or relax::Function to check.
* \param check_ty If true, verify that every Expr has ty populated.
* \return true if the object is well-formed, false otherwise.
*/
TVM_DLL bool CheckWellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
/*!
* \brief Using the layout transforms on the outputs, suggest layout transformation on the blocks
* and buffers for the PrimFunc.
*
* \param fn The PrimFunc to be analyzed.
* \param write_buffer_transformations Array of IndexMap transformations on PrimFunc outputs.
* \return Suggested transforms per block in `fn`. For each block the returned value is a map
* from the object (block or buffer) to it's index map transformation.
*/
TVM_DLL ffi::Map<tirx::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
const Function& fn, ffi::Array<tirx::IndexMap> write_buffer_transformations);
/* \brief Collect variables whose value can be computed at compile-time
*
* If a function has the `kNumInput` attribute, then the first
* `kNumInput` parameters are provided at run-time, while all
* remaining parameters may be known at compile-time. This utility
* collects all variable bindings that only depend, directly or
* indirectly, on the parameters known at compile-time.
*
* \param func The relax::Function to analyze
*
* \return The set of variables that can be computed at compile-time,
* in order of their occurrence within the function.
*/
TVM_DLL ffi::Array<Var> ComputableAtCompileTime(const Function& func);
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ANALYSIS_H_
include/tvm/relax/attrs/ccl.h
+91
View File
@@ -0,0 +1,91 @@
/*
* 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/relax/attrs/ccl.h
* \brief Attributes for ccl operators.
*/
#ifndef TVM_RELAX_ATTRS_CCL_H_
#define TVM_RELAX_ATTRS_CCL_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in allreduce operators */
struct AllReduceAttrs : public tvm::AttrsNode {
ffi::String op_type;
bool in_group;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AllReduceAttrs>()
.def_ro("op_type", &AllReduceAttrs::op_type,
"The type of reduction operation to be applied to the input data. Now only sum is "
"supported.")
.def_ro("in_group", &AllReduceAttrs::in_group,
"Whether the reduction operation performs in group or globally or in group as "
"default.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllReduceAttrs", AllReduceAttrs, AttrsNode);
}; // struct AllReduceAttrs
/*! \brief Attributes used in allgather operators */
struct AllGatherAttrs : public tvm::AttrsNode {
int num_workers;
bool in_group;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AllGatherAttrs>()
.def_ro("num_workers", &AllGatherAttrs::num_workers,
"The number of workers, also the number of parts the given buffer should be "
"chunked into.")
.def_ro("in_group", &AllGatherAttrs::in_group,
"Whether the allgather operation performs in group or globally or in group as "
"default.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllGatherAttrs", AllGatherAttrs, AttrsNode);
}; // struct AllGatherAttrs
/*! \brief Attributes used in scatter operators */
struct ScatterCollectiveAttrs : public tvm::AttrsNode {
int num_workers;
int axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ScatterCollectiveAttrs>()
.def_ro("num_workers", &ScatterCollectiveAttrs::num_workers,
"The number of workers, also the number of parts the given buffer should be "
"chunked into.")
.def_ro("axis", &ScatterCollectiveAttrs::axis,
"The axis of the tensor to be scattered. The tensor will be chunked along "
"this axis.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterCollectiveAttrs", ScatterCollectiveAttrs,
AttrsNode);
}; // struct ScatterCollectiveAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_CCL_H_
include/tvm/relax/attrs/create.h
+60
View File
@@ -0,0 +1,60 @@
/*
* 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/relax/attrs/create.h
* \brief Attributes for tensor creation operators.
*/
#ifndef TVM_RELAX_ATTRS_CREATE_H_
#define TVM_RELAX_ATTRS_CREATE_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in full/full_like, ones/ones_like, and zeros/zeros_like operators */
struct InitAttrs : public AttrsNode {
ffi::Optional<DLDataType> dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<InitAttrs>().def_ro("dtype", &InitAttrs::dtype,
"The data type of the created tensor.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InitAttrs", InitAttrs, AttrsNode);
}; // struct InitAttrs
/*! \brief Attributes used in tril and triu operator */
struct TriluAttrs : public AttrsNode {
int k;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TriluAttrs>().def_ro(
"k", &TriluAttrs::k,
"The number of diagonals above or below the main diagonal to exclude or include.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TriluAttrs", TriluAttrs, AttrsNode);
}; // struct TriluAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_CREATE_H_
include/tvm/relax/attrs/datatype.h
+57
View File
@@ -0,0 +1,57 @@
/*
* 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/relax/attrs/datatype.h
* \brief Attributes for datatype operators.
*/
#ifndef TVM_RELAX_ATTRS_DATATYPE_H_
#define TVM_RELAX_ATTRS_DATATYPE_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in astype operator */
struct AstypeAttrs : public AttrsNode {
DLDataType dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AstypeAttrs>().def_ro("dtype", &AstypeAttrs::dtype, "Target data type");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AstypeAttrs", AstypeAttrs, AttrsNode);
}; // struct AstypeAttrs.
/*! \brief Attributes used in wrap_param operator */
struct WrapParamAttrs : public AttrsNode {
DLDataType dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<WrapParamAttrs>().def_ro("dtype", &WrapParamAttrs::dtype, "Target data type");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.WrapParamAttrs", WrapParamAttrs, AttrsNode);
}; // struct WrapParamAttrs.
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_DATATYPE_H_
include/tvm/relax/attrs/distributed.h
+53
View File
@@ -0,0 +1,53 @@
/*
* 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/relax/attrs/distributed.h
* \brief Attributes for redistribute and annotate_sharding operators.
*/
#ifndef TVM_RELAX_ATTRS_DISTRIBUTED_H_
#define TVM_RELAX_ATTRS_DISTRIBUTED_H_
#include <tvm/relax/distributed/global_info.h>
#include <tvm/relax/distributed/type.h>
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes for redistribute and annotate_sharding operator */
struct DistributionAttrs : public AttrsNode {
distributed::DeviceMesh device_mesh;
distributed::Placement placement;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DistributionAttrs>()
.def_ro("device_mesh", &DistributionAttrs::device_mesh,
"The device mesh of a tensor's distribution plan")
.def_ro("placement", &DistributionAttrs::placement,
"The placement of a tensor's distribution plan");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DistributionAttrs", DistributionAttrs, AttrsNode);
}; // struct DistributionAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_DISTRIBUTED_H_
include/tvm/relax/attrs/image.h
+168
View File
@@ -0,0 +1,168 @@
/*
* 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/relax/attrs/image.h
* \brief Attributes for image operators.
*/
#ifndef TVM_RELAX_ATTRS_IMAGE_H_
#define TVM_RELAX_ATTRS_IMAGE_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in image resize2d operator */
struct Resize2DAttrs : public AttrsNode {
ffi::Array<FloatImm> roi;
ffi::String layout;
ffi::String method;
ffi::String coordinate_transformation_mode;
ffi::String rounding_method;
double cubic_alpha;
int cubic_exclude;
double extrapolation_value;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Resize2DAttrs>()
.def_ro("roi", &Resize2DAttrs::roi,
"Region of Interest for coordinate transformation mode 'tf_crop_and_resize'")
.def_ro("layout", &Resize2DAttrs::layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Resize is applied on the 'H' and"
"'W' dimensions.")
.def_ro("method", &Resize2DAttrs::method,
"Specify the mode to use for scaling."
"nearest_neighbor - Nearest Neighbor"
"linear - Bilinear Interpolation"
"cubic - Bicubic Interpolation")
.def_ro("coordinate_transformation_mode", &Resize2DAttrs::coordinate_transformation_mode,
"Describes how to transform the coordinate in the resized tensor"
"to the coordinate in the original tensor."
"Refer to the ONNX Resize operator specification for details"
"Available options are half_pixel, align_corners and asymmetric")
.def_ro("rounding_method", &Resize2DAttrs::rounding_method,
"indicates how to find the \"nearest\" pixel in nearest_neighbor method"
"Available options are round, floor, and ceil.")
.def_ro("cubic_alpha", &Resize2DAttrs::cubic_alpha,
"Spline Coefficient for Bicubic Interpolation")
.def_ro("cubic_exclude", &Resize2DAttrs::cubic_exclude,
"Flag to exclude exterior of the image during bicubic interpolation")
.def_ro("extrapolation_value", &Resize2DAttrs::extrapolation_value,
"Value to return when roi is outside of the image")
.def_ro(
"out_dtype", &Resize2DAttrs::out_dtype,
"The dtype of the output tensor. It it is not specified, the output will have the same "
"dtype as input if not specified.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize2DAttrs", Resize2DAttrs, AttrsNode);
}; // struct Resize2dAttrs
/*! \brief Attributes used in image resize3d operator */
struct Resize3DAttrs : public AttrsNode {
ffi::Array<FloatImm> roi;
ffi::String layout;
ffi::String method;
ffi::String coordinate_transformation_mode;
ffi::String rounding_method;
double cubic_alpha;
int cubic_exclude;
double extrapolation_value;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Resize3DAttrs>()
.def_ro("roi", &Resize3DAttrs::roi,
"Region of Interest for coordinate transformation mode 'tf_crop_and_resize'")
.def_ro("layout", &Resize3DAttrs::layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Resize is applied on the 'D', 'H' and"
"'W' dimensions.")
.def_ro("method", &Resize3DAttrs::method,
"Specify the mode to use for scaling."
"nearest_neighbor - Nearest Neighbor"
"linear - Trilinear Interpolation"
"cubic - Tricubic Interpolation")
.def_ro("coordinate_transformation_mode", &Resize3DAttrs::coordinate_transformation_mode,
"Describes how to transform the coordinate in the resized tensor"
"to the coordinate in the original tensor."
"Refer to the ONNX Resize operator specification for details"
"Available options are half_pixel, align_corners and asymmetric")
.def_ro("rounding_method", &Resize3DAttrs::rounding_method,
"indicates how to find the \"nearest\" pixel in nearest_neighbor method"
"Available options are round, floor, and ceil.")
.def_ro("cubic_alpha", &Resize3DAttrs::cubic_alpha,
"Spline Coefficient for Tricubic Interpolation")
.def_ro("cubic_exclude", &Resize3DAttrs::cubic_exclude,
"Flag to exclude exterior of the image during tricubic interpolation")
.def_ro("extrapolation_value", &Resize3DAttrs::extrapolation_value,
"Value to return when roi is outside of the image")
.def_ro(
"out_dtype", &Resize3DAttrs::out_dtype,
"The dtype of the output tensor. It it is not specified, the output will have the same "
"dtype as input if not specified.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize3DAttrs", Resize3DAttrs, AttrsNode);
}; // struct Resize3DAttrs
/*! \brief Attributes used in image grid_sample operator */
struct GridSampleAttrs : public AttrsNode {
ffi::String method;
ffi::String layout;
ffi::String padding_mode;
bool align_corners;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GridSampleAttrs>()
.def_ro("method", &GridSampleAttrs::method,
"Interpolation method. Can be 'nearest', 'bilinear', or 'bicubic'.")
.def_ro("layout", &GridSampleAttrs::layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc.")
.def_ro("padding_mode", &GridSampleAttrs::padding_mode,
"Padding mode for outside grid values. Can be 'zeros', 'border', or 'reflection'.")
.def_ro("align_corners", &GridSampleAttrs::align_corners,
"If True, the corner pixels of the input and output tensors are aligned.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsNode);
}; // struct GridSampleAttrs
/*! \brief Attributes used in image affine_grid operator */
struct AffineGridAttrs : public AttrsNode {
bool align_corners;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AffineGridAttrs>().def_ro(
"align_corners", &AffineGridAttrs::align_corners,
"If True, normalized grid coordinates map to corner pixels; otherwise to pixel centers.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode);
}; // struct AffineGridAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_IMAGE_H_
include/tvm/relax/attrs/index.h
+65
View File
@@ -0,0 +1,65 @@
/*
* 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/relax/attrs/index.h
* \brief Attributes for indexing operators.
*/
#ifndef TVM_RELAX_ATTRS_INDEX_H_
#define TVM_RELAX_ATTRS_INDEX_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in take operator */
struct TakeAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
ffi::String mode;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TakeAttrs>()
.def_ro("axis", &TakeAttrs::axis, "The axis over which to select values.")
.def_ro("mode", &TakeAttrs::mode, "The mode for handling out-of-bounds indices.",
refl::DefaultValue("fast"));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeAttrs", TakeAttrs, AttrsNode);
}; // struct TakeAttrs
/*! \brief Attributes used in strided_slice operator */
struct StridedSliceAttrs : public AttrsNode {
bool assume_inbound;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<StridedSliceAttrs>().def_ro(
"assume_inbound", &StridedSliceAttrs::assume_inbound,
"Whether to assume the indices are in bound. If it is set to false, "
"out of bound indices will be clipped to the bound.",
refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StridedSliceAttrs", StridedSliceAttrs, AttrsNode);
}; // struct StridedSliceAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_INDEX_H_
include/tvm/relax/attrs/linear_algebra.h
+59
View File
@@ -0,0 +1,59 @@
/*
* 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/relax/attrs/linear_algebra.h
* \brief Attributes for linear algebra operators.
*/
#ifndef TVM_RELAX_ATTRS_LINEAR_ALGEBRA_H_
#define TVM_RELAX_ATTRS_LINEAR_ALGEBRA_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes for matmul operator */
struct MatmulAttrs : public AttrsNode {
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MatmulAttrs>().def_ro("out_dtype", &MatmulAttrs::out_dtype,
"The data type of the output tensor");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MatmulAttrs", MatmulAttrs, AttrsNode);
}; // struct MatmulAttrs
/*! \brief Attributes used in einsum operator */
struct EinsumAttrs : public AttrsNode {
ffi::String subscripts;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<EinsumAttrs>().def_ro("subscripts", &EinsumAttrs::subscripts,
"The einsum expression string");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.EinsumAttrs", EinsumAttrs, AttrsNode);
}; // struct EinsumAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_LINEAR_ALGEBRA_H_
include/tvm/relax/attrs/manipulate.h
+333
View File
@@ -0,0 +1,333 @@
/*
* 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/relax/attrs/manipulate.h
* \brief Attributes for tensor manipulation operators.
*/
#ifndef TVM_RELAX_ATTRS_MANIPULATE_H_
#define TVM_RELAX_ATTRS_MANIPULATE_H_
#include <tvm/relax/expr.h>
#include <tvm/tirx/index_map.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in concat operators */
struct ConcatAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ConcatAttrs>().def_ro("axis", &ConcatAttrs::axis,
"The axis at which the input arrays are concatenated."
"Should lie in range `[-ndim, ndim)`.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ConcatAttrs", ConcatAttrs, AttrsNode);
}; // struct ConcatAttrs
/*! \brief Attributes used in expand_dims operators */
struct ExpandDimsAttrs : public AttrsNode {
ffi::Array<int64_t> axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExpandDimsAttrs>().def_ro(
"axis", &ExpandDimsAttrs::axis,
"The axes at which the input array are expanded. "
"All values are required to lie in range `[-data.ndim - 1, data.ndim]`, "
"with the convention of negative indexing.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ExpandDimsAttrs", ExpandDimsAttrs, AttrsNode);
}; // struct ExpandDimsAttrs
/*! \brief Attributes used in layout_transform operator */
struct LayoutTransformAttrs : public AttrsNode {
tirx::IndexMap index_map;
// pad_value is chosen to be of PrimExpr type, as it represents constant TIR POD expression. This
// needs to be revisited in case PrimExpr is evolved to represent symbolic expression in future.
ffi::Optional<PrimExpr> pad_value;
/*!
* axis_separators between input axes when generating flattened output axes. For buffers
* representing flat 1-d memory (e.g. any buffer in RAM), this should be an empty array.
* For buffers representing non-flat memory, each entry in axis_separators should be the
* first input axis that is part of a new flattened axis.
*/
ffi::Optional<ffi::Array<IntImm>> axis_separators;
/*!
* axis_separators for input buffers.
* Needed to identify if the input buffer to layout_transform
* contains axis separator.
*/
ffi::Optional<ffi::Array<IntImm>> input_axis_separators;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<LayoutTransformAttrs>()
.def_ro("index_map", &LayoutTransformAttrs::index_map,
"The layout transformation to apply.")
.def_ro(
"pad_value", &LayoutTransformAttrs::pad_value,
"The specific value to be used to pad if the layout transform would result in implicit "
"padding. If not specified, the compiler is free to choose any value.")
.def_ro("axis_separators", &LayoutTransformAttrs::axis_separators,
"The separators between input axes when generating flat output axes")
.def_ro("input_axis_separators", &LayoutTransformAttrs::input_axis_separators,
"The separators between axes to regenerate output");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayoutTransformAttrs", LayoutTransformAttrs,
AttrsNode);
}; // struct LayoutTransformAttrs
/*! \brief Attributes used in permute_dims operator */
struct PermuteDimsAttrs : public AttrsNode {
ffi::Optional<ffi::Array<int64_t>> axes;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PermuteDimsAttrs>().def_ro(
"axes", &PermuteDimsAttrs::axes, "The target axes order, reverse order if not specified.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PermuteDimsAttrs", PermuteDimsAttrs, AttrsNode);
}; // struct PermuteDimsAttrs
/*! \brief Attributes used in split operator */
struct SplitAttrs : public AttrsNode {
ffi::ObjectRef indices_or_sections;
int axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SplitAttrs>()
.def_ro("indices_or_sections", &SplitAttrs::indices_or_sections,
"The input array of indices or the number of split sections.")
.def_ro("axis", &SplitAttrs::axis, "The axis to be splitted");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SplitAttrs", SplitAttrs, AttrsNode);
}; // struct SplitAttrs
/*! \brief Attributes used in squeeze operators */
struct SqueezeAttrs : public AttrsNode {
ffi::Optional<ffi::Array<int64_t>> axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SqueezeAttrs>().def_ro("axis", &SqueezeAttrs::axis,
"The axis to squeeze in the input tensor."
"If `axis = None`, all axis of dimension 1 get squeezed;"
"Else, the dimension in axes get squeezed."
"It is an error if an axis does not has dimension 1.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SqueezeAttrs", SqueezeAttrs, AttrsNode);
}; // struct SqueezeAttrs
/*! \brief Attributes used in stack operators */
struct StackAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<StackAttrs>().def_ro(
"axis", &StackAttrs::axis,
"The axis along which to stack the input tensors. "
"The axis will be inserted at this position in the output, "
"so it must be in range [-ndim-1, ndim] where ndim is the "
"number of dimensions of the input tensors.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StackAttrs", StackAttrs, AttrsNode);
}; // struct StackAttrs
/*! \brief Attributes used in repeat operators */
struct RepeatAttrs : public AttrsNode {
int repeats;
ffi::Optional<int64_t> axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RepeatAttrs>()
.def_ro("repeats", &RepeatAttrs::repeats, "The number of repetitions.")
.def_ro("axis", &RepeatAttrs::axis,
"The axis along which to repeat values. The negative numbers are interpreted "
"counting from the backward. By default, use the flattened input array, and "
"return a flat output array.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RepeatAttrs", RepeatAttrs, AttrsNode);
}; // struct RepeatAttrs
/*! \brief Attributes used in tile operators */
struct TileAttrs : public AttrsNode {
ffi::Array<int64_t> repeats;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TileAttrs>().def_ro("repeats", &TileAttrs::repeats,
"The number of repetitions of data along each axis.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TileAttrs", TileAttrs, AttrsNode);
}; // struct TileAttrs
/*! \brief Attributes used in flip operators */
struct FlipAttrs : public AttrsNode {
int64_t axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FlipAttrs>().def_ro("axis", &FlipAttrs::axis,
"The axis along which to flip over.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.FlipAttrs", FlipAttrs, AttrsNode);
}; // struct FlipAttrs
/*! \brief Attributes used in reverse_sequence operators */
struct ReverseSequenceAttrs : public AttrsNode {
int64_t seq_axis;
int64_t batch_axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ReverseSequenceAttrs>()
.def_ro("seq_axis", &ReverseSequenceAttrs::seq_axis,
"The axis along which to reverse variable length slices.")
.def_ro("batch_axis", &ReverseSequenceAttrs::batch_axis,
"The axis that indexes the batch.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ReverseSequenceAttrs", ReverseSequenceAttrs,
AttrsNode);
}; // struct ReverseSequenceAttrs
/*! \brief Attributes used in gather_elements operators */
struct GatherElementsAttrs : public AttrsNode {
int64_t axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GatherElementsAttrs>().def_ro("axis", &GatherElementsAttrs::axis,
"The axis along which to index.",
refl::DefaultValue(0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherElementsAttrs", GatherElementsAttrs,
AttrsNode);
}; // struct GatherElementsAttrs
/*! \brief Attributes used in gather_nd operators */
struct GatherNDAttrs : public AttrsNode {
int64_t batch_dims;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GatherNDAttrs>().def_ro("batch_dims", &GatherNDAttrs::batch_dims,
"The number of batch dims.", refl::DefaultValue(0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherNDAttrs", GatherNDAttrs, AttrsNode);
}; // struct GatherNDAttrs
/*! \brief Attributes used in index_put operator */
struct IndexPutAttrs : public AttrsNode {
bool accumulate;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IndexPutAttrs>().def_ro(
"accumulate", &IndexPutAttrs::accumulate,
"Whether to accumulate (add) values rather than replace. "
"If true, performs tensor[indices] += values, "
"otherwise performs tensor[indices] = values.",
refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.IndexPutAttrs", IndexPutAttrs, AttrsNode);
}; // struct IndexPutAttrs
/*! \brief Attribute used in meshgrid operator */
struct MeshgridAttrs : public AttrsNode {
ffi::Optional<ffi::String> indexing;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MeshgridAttrs>().def_ro("indexing", &MeshgridAttrs::indexing,
"Specifies how the grid dimensions are ordered.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MeshgridAttrs", MeshgridAttrs, AttrsNode);
};
/*! \brief Attributes used in scatter_elements operators */
struct ScatterElementsAttrs : public AttrsNode {
int64_t axis;
ffi::String reduction;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ScatterElementsAttrs>()
.def_ro("axis", &ScatterElementsAttrs::axis, "The axis over which to select values.",
refl::DefaultValue(0))
.def_ro("reduction", &ScatterElementsAttrs::reduction,
"Reduction mode of the scatter elements, "
"either \"update\", \"add\", \"mul\", \"mean\", \"min\" or \"max\".",
refl::DefaultValue("update"));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterElementsAttrs", ScatterElementsAttrs,
AttrsNode);
}; // struct ScatterElementsAttrs
/*! \brief Attributes used in scatter_nd operators */
struct ScatterNDAttrs : public AttrsNode {
ffi::String reduction;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ScatterNDAttrs>().def_ro(
"reduction", &ScatterNDAttrs::reduction,
"Accumulation mode of the ScatterND, "
"either \"update\", \"add\", \"mul\", \"min\" or \"max\".",
refl::DefaultValue("update"));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterNDAttrs", ScatterNDAttrs, AttrsNode);
}; // struct ScatterNDAttrs
/*! \brief Attributes used in slice_scatter operator */
struct SliceScatterAttrs : public AttrsNode {
int axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SliceScatterAttrs>().def_ro("axis", &SliceScatterAttrs::axis,
"the dimension to insert the slice into ",
refl::DefaultValue(0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SliceScatterAttrs", SliceScatterAttrs, AttrsNode);
}; // struct SliceScatterAttrs
/*! \brief Attributes used in one_hot operator */
struct OneHotAttrs : public AttrsNode {
int depth;
int axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<OneHotAttrs>()
.def_ro("depth", &OneHotAttrs::depth, "Depth of the one hot dimension.")
.def_ro("axis", &OneHotAttrs::axis, "Axis to fill.", refl::DefaultValue(-1));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.OneHotAttrs", OneHotAttrs, AttrsNode);
}; // struct OneHotAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_MANIPULATE_H_
include/tvm/relax/attrs/nn.h
+785
View File
@@ -0,0 +1,785 @@
/*
* 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/relax/attrs/nn.h
* \brief Attributes for neural network operators.
*/
#ifndef TVM_RELAX_ATTRS_NN_H_
#define TVM_RELAX_ATTRS_NN_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in Conv1d operator */
struct Conv1DAttrs : public AttrsNode {
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> dilation;
int groups;
ffi::String data_layout;
ffi::String kernel_layout;
ffi::String out_layout;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv1DAttrs>()
.def_ro("strides", &Conv1DAttrs::strides, "Specifies the strides of the convolution.")
.def_ro("padding", &Conv1DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on both sides"
"two int : padding width in the order of (left, right)")
.def_ro("dilation", &Conv1DAttrs::dilation,
"Specifies the dilation rate to use for dilated convolution.")
.def_ro("groups", &Conv1DAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
"output channels should be divisible by the number of groups.")
.def_ro("data_layout", &Conv1DAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, width"
"dimensions respectively. Convolution is applied on the 'W' dimensions.")
.def_ro("kernel_layout", &Conv1DAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
"'O', 'I', 'W' stands for num_filter, input_channel, and width"
"dimensions respectively.")
.def_ro("out_layout", &Conv1DAttrs::out_layout,
"Dimension ordering of output. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv1DAttrs::out_dtype,
"Output data type, set to explicit type under mixed precision setting");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DAttrs", Conv1DAttrs, AttrsNode);
}; // struct Conv1dAttrs
/*! \brief Attributes used in Conv2d operator */
struct Conv2DAttrs : public AttrsNode {
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> dilation;
int groups;
ffi::String data_layout;
ffi::String kernel_layout;
ffi::String out_layout;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv2DAttrs>()
.def_ro("strides", &Conv2DAttrs::strides, "Specifies the strides of the convolution.")
.def_ro("padding", &Conv2DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
"four int : padding width in the order of (top, left, bottom, right)")
.def_ro("dilation", &Conv2DAttrs::dilation,
"Specifies the dilation rate to use for dilated convolution.")
.def_ro("groups", &Conv2DAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
"output channels should be divisible by the number of groups.")
.def_ro("data_layout", &Conv2DAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Convolution is applied on the 'H' and"
"'W' dimensions.")
.def_ro("kernel_layout", &Conv2DAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
"'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
"dimensions respectively.")
.def_ro("out_layout", &Conv2DAttrs::out_layout,
"Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv2DAttrs::out_dtype,
"Output data type, set to explicit type under mixed precision setting");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DAttrs", Conv2DAttrs, AttrsNode);
}; // struct Conv2dAttrs
/*! \brief Attributes used in Conv3d operator */
struct Conv3DAttrs : public AttrsNode {
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> dilation;
int groups;
ffi::String data_layout;
ffi::String kernel_layout;
ffi::String out_layout;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv3DAttrs>()
.def_ro("strides", &Conv3DAttrs::strides, "Specifies the strides of the convolution.")
.def_ro(
"padding", &Conv3DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
"four int : padding width in the order of (forward, back, top, left, bottom, right)")
.def_ro("dilation", &Conv3DAttrs::dilation,
"Specifies the dilation rate to use for dilated convolution.")
.def_ro("groups", &Conv3DAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
"output channels should be divisible by the number of groups.")
.def_ro("data_layout", &Conv3DAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Convolution is applied on the 'D', 'H', and"
"'W' dimensions.")
.def_ro(
"kernel_layout", &Conv3DAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIDHW', 'OIDHW16o16i', etc."
"'O', 'I', 'D', 'H', 'W' stands for num_filter, input_channel, depth, height, and width"
"dimensions respectively.")
.def_ro("out_layout", &Conv3DAttrs::out_layout,
"Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv3DAttrs::out_dtype,
"Output data type, set to explicit type under mixed precision setting");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DAttrs", Conv3DAttrs, AttrsNode);
}; // struct Conv3dAttrs
/*! \brief Attributes used in Conv1DTranspose operator */
struct Conv1DTransposeAttrs : public AttrsNode {
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> output_padding;
ffi::Array<int64_t> dilation;
int groups;
ffi::String data_layout;
ffi::String kernel_layout;
ffi::String out_layout;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv1DTransposeAttrs>()
.def_ro("strides", &Conv1DTransposeAttrs::strides,
"Specifies the strides of the convolution.")
.def_ro("padding", &Conv1DTransposeAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on both sides"
"two int : padding width in the order of (left, right)")
.def_ro("output_padding", &Conv1DTransposeAttrs::output_padding,
"Used to disambiguate the output shape.")
.def_ro("dilation", &Conv1DTransposeAttrs::dilation,
"Specifies the dilation rate to use for dilated convolution.")
.def_ro("groups", &Conv1DTransposeAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
"output channels should be divisible by the number of groups.")
.def_ro("data_layout", &Conv1DTransposeAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, width"
"dimensions respectively. Convolution is applied on the 'W' dimensions.")
.def_ro("kernel_layout", &Conv1DTransposeAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
"'O', 'I', 'W' stands for num_filter, input_channel, and width"
"dimensions respectively.")
.def_ro("out_layout", &Conv1DTransposeAttrs::out_layout,
"Dimension ordering of output. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv1DTransposeAttrs::out_dtype,
"Output data type, set to explicit type under mixed precision setting");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DTransposeAttrs", Conv1DTransposeAttrs,
AttrsNode);
}; // struct Conv1DTransposeAttrs
/*! \brief Attributes used in Conv2d operator */
struct Conv2DTransposeAttrs : public AttrsNode {
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> output_padding;
ffi::Array<int64_t> dilation;
int groups;
ffi::String data_layout;
ffi::String kernel_layout;
ffi::String out_layout;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv2DTransposeAttrs>()
.def_ro("strides", &Conv2DTransposeAttrs::strides,
"Specifies the strides of the convolution.")
.def_ro("padding", &Conv2DTransposeAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
"four int : padding width in the order of (top, left, bottom, right)")
.def_ro("output_padding", &Conv2DTransposeAttrs::output_padding,
"Used to disambiguate the output shape.")
.def_ro("dilation", &Conv2DTransposeAttrs::dilation,
"Specifies the dilation rate to use for dilated convolution.")
.def_ro("groups", &Conv2DTransposeAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
"output channels should be divisible by the number of groups.")
.def_ro("data_layout", &Conv2DTransposeAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Convolution is applied on the 'H' and"
"'W' dimensions.")
.def_ro("kernel_layout", &Conv2DTransposeAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
"'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
"dimensions respectively.")
.def_ro("out_layout", &Conv2DTransposeAttrs::out_layout,
"Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv2DTransposeAttrs::out_dtype,
"Output data type, set to explicit type under mixed precision setting");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DTransposeAttrs", Conv2DTransposeAttrs,
AttrsNode);
}; // struct Conv2DTransposeAttrs
/*! \brief Attributes used in Conv3dTranspose operator */
struct Conv3DTransposeAttrs : public AttrsNode {
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> output_padding;
ffi::Array<int64_t> dilation;
int groups;
ffi::String data_layout;
ffi::String kernel_layout;
ffi::String out_layout;
ffi::Optional<DLDataType> out_dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv3DTransposeAttrs>()
.def_ro("strides", &Conv3DTransposeAttrs::strides,
"Specifies the strides of the convolution.")
.def_ro("padding", &Conv3DTransposeAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"three int : back/bottom/right will use same padding as front/top/left"
"six int : padding width in the order of (front, top, left, back, bottom, right)")
.def_ro("output_padding", &Conv3DTransposeAttrs::output_padding,
"Used to disambiguate the output shape.")
.def_ro("dilation", &Conv3DTransposeAttrs::dilation,
"Specifies the dilation rate to use for dilated convolution.")
.def_ro("groups", &Conv3DTransposeAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
"output channels should be divisible by the number of groups.")
.def_ro("data_layout", &Conv3DTransposeAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Convolution is applied on the 'D', 'H', and"
"'W' dimensions.")
.def_ro(
"kernel_layout", &Conv3DTransposeAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'IODHW', etc."
"'I', 'O', 'D', 'H', 'W' stands for input_channel, output_channel, depth, height, and "
"width"
"dimensions respectively.")
.def_ro("out_layout", &Conv3DTransposeAttrs::out_layout,
"Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv3DTransposeAttrs::out_dtype,
"Output data type, set to explicit type under mixed precision setting");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DTransposeAttrs", Conv3DTransposeAttrs,
AttrsNode);
}; // struct Conv3DTransposeAttrs
/*! \brief Attributes used in max_pool1d and avg_pool1d operator */
struct Pool1DAttrs : public AttrsNode {
ffi::Array<int64_t> pool_size;
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> dilation;
bool ceil_mode;
bool count_include_pad;
ffi::String layout;
ffi::String out_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Pool1DAttrs>()
.def_ro("pool_size", &Pool1DAttrs::pool_size, "Size of the pooling windows.")
.def_ro("strides", &Pool1DAttrs::strides, "Specifies the strides of the convolution.")
.def_ro("dilation", &Pool1DAttrs::dilation, "Specifies the dilation of the convolution.")
.def_ro("padding", &Pool1DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : padding width in the order of (left, right)")
.def_ro(
"ceil_mode", &Pool1DAttrs::ceil_mode,
"A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
"every element in the input tensor will be covered by a sliding window.")
.def_ro("count_include_pad", &Pool1DAttrs::count_include_pad,
"When true, will include padding to compute the average")
.def_ro("layout", &Pool1DAttrs::layout,
"Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, and width"
"dimensions respectively. Pooling is applied on the 'W' dimensions.",
refl::DefaultValue("NCW"))
.def_ro("out_layout", &Pool1DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, and width"
"dimensions respectively. Pooling is applied on the 'W' dimensions.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool1DAttrs", Pool1DAttrs, AttrsNode);
}; // struct Pool1dAttrs
/*! \brief Attributes used in max_pool2d and avg_pool2d operator */
struct Pool2DAttrs : public AttrsNode {
ffi::Array<int64_t> pool_size;
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> dilation;
bool ceil_mode;
bool count_include_pad;
ffi::String layout;
ffi::String out_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Pool2DAttrs>()
.def_ro("pool_size", &Pool2DAttrs::pool_size, "Size of the pooling windows.")
.def_ro("strides", &Pool2DAttrs::strides, "Specifies the strides of the convolution.")
.def_ro("dilation", &Pool2DAttrs::dilation, "Specifies the dilation of the convolution.")
.def_ro("padding", &Pool2DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
"four int : padding width in the order of (top, left, bottom, right)")
.def_ro(
"ceil_mode", &Pool2DAttrs::ceil_mode,
"A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
"every element in the input tensor will be covered by a sliding window.")
.def_ro("count_include_pad", &Pool2DAttrs::count_include_pad,
"When true, will include padding to compute the average")
.def_ro("layout", &Pool2DAttrs::layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Pooling is applied on the 'H' and"
"'W' dimensions.")
.def_ro("out_layout", &Pool2DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Pooling is applied on the 'H' and"
"'W' dimensions.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool2DAttrs", Pool2DAttrs, AttrsNode);
}; // struct Pool2dAttrs
/*! \brief Attributes used in max_pool3d and avg_pool3d operator */
struct Pool3DAttrs : public AttrsNode {
ffi::Array<int64_t> pool_size;
ffi::Array<int64_t> strides;
ffi::Array<int64_t> padding;
ffi::Array<int64_t> dilation;
bool ceil_mode;
bool count_include_pad;
ffi::String layout;
ffi::String out_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Pool3DAttrs>()
.def_ro("pool_size", &Pool3DAttrs::pool_size, "Size of the pooling windows.")
.def_ro("strides", &Pool3DAttrs::strides, "Specifies the strides of the convolution.")
.def_ro("dilation", &Pool3DAttrs::dilation, "Specifies the dilation of the convolution.")
.def_ro("padding", &Pool3DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"three int : back, bottom, right will use same padding as front, top, left"
"four int : padding width in the order of (front, top, left, back, bottom, right)")
.def_ro(
"ceil_mode", &Pool3DAttrs::ceil_mode,
"A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
"every element in the input tensor will be covered by a sliding window.")
.def_ro("count_include_pad", &Pool3DAttrs::count_include_pad,
"When true, will include padding to compute the average")
.def_ro("layout", &Pool3DAttrs::layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Pooling is applied on the 'D', 'H' and"
"'W' dimensions.")
.def_ro("out_layout", &Pool3DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Pooling is applied on the 'D', 'H' and"
"'W' dimensions.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool3DAttrs", Pool3DAttrs, AttrsNode);
}; // struct Pool3dAttrs
/*! \brief Attributes for 1d adaptive pool operator */
struct AdaptivePool1DAttrs : public AttrsNode {
ffi::Optional<ffi::Array<int64_t>> output_size;
ffi::String layout;
ffi::String out_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AdaptivePool1DAttrs>()
.def_ro("output_size", &AdaptivePool1DAttrs::output_size, "Output width.")
.def_ro("layout", &AdaptivePool1DAttrs::layout,
"Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel and width"
"dimensions respectively. Pooling is applied on the"
"'W' dimensions.")
.def_ro("out_layout", &AdaptivePool1DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel and width"
"dimensions respectively. Pooling is applied on the"
"'W' dimensions.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool1DAttrs", AdaptivePool1DAttrs,
AttrsNode);
}; // struct AdaptivePool1DAttrs
/*! \brief Attributes for 2d adaptive pool operator */
struct AdaptivePool2DAttrs : public AttrsNode {
ffi::Optional<ffi::Array<int64_t>> output_size;
ffi::String layout;
ffi::String out_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AdaptivePool2DAttrs>()
.def_ro("output_size", &AdaptivePool2DAttrs::output_size, "Output height and width.")
.def_ro("layout", &AdaptivePool2DAttrs::layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Pooling is applied on the 'H' and"
"'W' dimensions.")
.def_ro("out_layout", &AdaptivePool2DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Pooling is applied on the 'H' and"
"'W' dimensions.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool2DAttrs", AdaptivePool2DAttrs,
AttrsNode);
}; // struct AdaptivePool2DAttrs
/*! \brief Attributes for 3d adaptive pool operator */
struct AdaptivePool3DAttrs : public AttrsNode {
ffi::Optional<ffi::Array<int64_t>> output_size;
ffi::String layout;
ffi::String out_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AdaptivePool3DAttrs>()
.def_ro("output_size", &AdaptivePool3DAttrs::output_size, "Output depth, height and width.")
.def_ro("layout", &AdaptivePool3DAttrs::layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Pooling is applied on 'D', 'H' and"
"'W' dimensions.")
.def_ro("out_layout", &AdaptivePool3DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Pooling is applied on 'D', 'H' and"
"'W' dimensions.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool3DAttrs", AdaptivePool3DAttrs,
AttrsNode);
}; // struct AdaptivePool3DAttrs
/*! \brief Attributes used in softmax operators */
struct SoftmaxAttrs : public AttrsNode {
int axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SoftmaxAttrs>().def_ro("axis", &SoftmaxAttrs::axis,
"The axis to sum over when computing softmax.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftmaxAttrs", SoftmaxAttrs, AttrsNode);
};
/*! \brief Attributes used in softmax operators */
struct LeakyReluAttrs : public AttrsNode {
double alpha;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<LeakyReluAttrs>().def_ro("alpha", &LeakyReluAttrs::alpha,
"The slope of the negative part.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LeakyReluAttrs", LeakyReluAttrs, AttrsNode);
};
/*! \brief Attributes used in softplus operators */
struct SoftplusAttrs : public AttrsNode {
double beta;
double threshold;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SoftplusAttrs>()
.def_ro("beta", &SoftplusAttrs::beta,
"Scaling factor controlling the sharpness of the Softplus transition.")
.def_ro("threshold", &SoftplusAttrs::threshold,
"Value determining when to use linear approximation for numerical stability.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftplusAttrs", SoftplusAttrs, AttrsNode);
};
/*! \brief Attributes used in PReLU operator */
struct PReluAttrs : public AttrsNode {
int axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PReluAttrs>().def_ro("axis", &PReluAttrs::axis,
"The axis along which the alpha values are applied.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PReluAttrs", PReluAttrs, AttrsNode);
};
/*! \brief Attributes used in batch_norm operator */
struct BatchNormAttrs : public AttrsNode {
int axis;
double epsilon;
bool center;
bool scale;
double momentum;
bool training;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BatchNormAttrs>()
.def_ro("axis", &BatchNormAttrs::axis, "The axis along which the normalization is applied.")
.def_ro("epsilon", &BatchNormAttrs::epsilon,
"Small float added to variance to avoid dividing by zero")
.def_ro("center", &BatchNormAttrs::center,
"Indicating if the beta offset will be added to the normalized tensor.")
.def_ro("scale", &BatchNormAttrs::scale,
"Indicating if the gamma scale will be multiplied.")
.def_ro("momentum", &BatchNormAttrs::momentum,
"The value used for the moving_mean and moving_var update.")
.def_ro("training", &BatchNormAttrs::training,
"Whether we are training (i.e., not in eval mode).");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BatchNormAttrs", BatchNormAttrs, AttrsNode);
}; // struct BatchNormAttrs
/*! \brief Attributes used in layer_norm operator */
struct LayerNormAttrs : public AttrsNode {
ffi::Array<int64_t> axes;
double epsilon;
bool center;
bool scale;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<LayerNormAttrs>()
.def_ro("axes", &LayerNormAttrs::axes,
"The axes that along which the normalization is applied.")
.def_ro("epsilon", &LayerNormAttrs::epsilon,
"Small float added to variance to avoid dividing by zero")
.def_ro("center", &LayerNormAttrs::center,
"Indicating if the beta offset will be added to the normalized tensor.")
.def_ro("scale", &LayerNormAttrs::scale,
"Indicating if the gamma scale will be multiplied.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayerNormAttrs", LayerNormAttrs, AttrsNode);
}; // struct LayerNormAttrs
/*! \brief Attributes used in group_norm operator */
struct GroupNormAttrs : public AttrsNode {
int num_groups;
int channel_axis;
ffi::Array<int64_t> axes;
double epsilon;
bool center;
bool scale;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GroupNormAttrs>()
.def_ro("num_groups", &GroupNormAttrs::num_groups,
"The number of groups to separate the channels into.")
.def_ro("channel_axis", &GroupNormAttrs::channel_axis,
"The axis that represents the channel.")
.def_ro(
"axes", &GroupNormAttrs::axes,
"The axes that along which the normalization is applied (excluding the channel axis).")
.def_ro("epsilon", &GroupNormAttrs::epsilon,
"Small float added to variance to avoid dividing by zero")
.def_ro("center", &GroupNormAttrs::center,
"Indicating if the beta offset will be added to the normalized tensor.")
.def_ro("scale", &GroupNormAttrs::scale,
"Indicating if the gamma scale will be multiplied.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GroupNormAttrs", GroupNormAttrs, AttrsNode);
}; // struct GroupNormAttrs
/*! \brief Attributes used in instance_norm operator */
struct InstanceNormAttrs : public AttrsNode {
int channel_axis;
ffi::Array<int64_t> axes;
double epsilon;
bool center;
bool scale;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<InstanceNormAttrs>()
.def_ro("channel_axis", &InstanceNormAttrs::channel_axis,
"The axis that represents the channel.")
.def_ro("axes", &InstanceNormAttrs::axes,
"The axes that along which the normalization is applied.")
.def_ro("epsilon", &InstanceNormAttrs::epsilon,
"Small float added to variance to avoid dividing by zero")
.def_ro("center", &InstanceNormAttrs::center,
"Indicating if the beta offset will be added to the normalized tensor.")
.def_ro("scale", &InstanceNormAttrs::scale,
"Indicating if the gamma scale will be multiplied.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InstanceNormAttrs", InstanceNormAttrs, AttrsNode);
}; // struct InstanceNormAttrs
/*! \brief Attributes used in rms_norm operator */
struct RMSNormAttrs : public AttrsNode {
ffi::Array<int64_t> axes;
double epsilon;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RMSNormAttrs>()
.def_ro("axes", &RMSNormAttrs::axes,
"The axes that along which the normalization is applied.")
.def_ro("epsilon", &RMSNormAttrs::epsilon,
"Small float added to variance to avoid dividing by zero");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RMSNormAttrs", RMSNormAttrs, AttrsNode);
}; // struct RMSNormAttrs
/*! \brief Attributes used in nll_loss operator */
struct NLLLossAttrs : public AttrsNode {
ffi::String reduction;
int ignore_index;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<NLLLossAttrs>()
.def_ro("reduction", &NLLLossAttrs::reduction,
"The reduction method to apply to the output. Can be"
"'none', 'mean' or 'sum'.",
refl::DefaultValue("mean"))
.def_ro("ignore_index", &NLLLossAttrs::ignore_index, "The target value to ignore.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NLLLossAttrs", NLLLossAttrs, AttrsNode);
}; // struct NLLLossAttrs
/*! \brief Attributes used in dropout operator */
struct DropoutAttrs : public AttrsNode {
double rate;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DropoutAttrs>().def_ro(
"rate", &DropoutAttrs::rate,
"Fraction of the input that gets dropped out during training time");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DropoutAttrs", DropoutAttrs, AttrsNode);
}; // struct DropoutAttrs
/*! \brief Attributes used in Attention operator */
struct AttentionAttrs : public AttrsNode {
ffi::Optional<FloatImm> scale;
ffi::Optional<ffi::String> causal_mask;
ffi::Optional<IntImm> window_size;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AttentionAttrs>()
.def_ro(
"scale", &AttentionAttrs::scale,
"The custom scale applied before the softmax. The default value is 1 / sqrt(head_dim).")
.def_ro("causal_mask", &AttentionAttrs::causal_mask,
"The type of the causal mask, i.e. 'TopLeft' and 'BottomRight'.")
.def_ro("window_size", &AttentionAttrs::window_size,
"The size of the window for sliding-window attention.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AttentionAttrs", AttentionAttrs, AttrsNode);
}; // struct AttentionAttrs
/*! \brief Attributes used for the padding operator */
struct PadAttrs : public AttrsNode {
ffi::Array<int64_t> pad_width;
double pad_value = 0.0;
tvm::ffi::String pad_mode;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PadAttrs>()
.def_ro("pad_width", &PadAttrs::pad_width,
"Number of values padded to the edges of each axis, "
"in the format of (before_1, after_1, ..., before_N, after_N)")
.def_ro("pad_value", &PadAttrs::pad_value, "The value to fill in padded area with",
refl::DefaultValue(0.0))
.def_ro("pad_mode", &PadAttrs::pad_mode,
"Padding type to use. \"constant\" pads with constant_value, "
"\"edge\" pads using the edge values of the input array, "
"\"reflect\" pads by reflecting values with respect to the edges.",
refl::DefaultValue("constant"));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PadAttrs", PadAttrs, AttrsNode);
};
/*! \brief Attributes used for the pixel shuffle operator */
struct PixelShuffleAttrs : public AttrsNode {
int upscale_factor;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PixelShuffleAttrs>().def_ro("upscale_factor",
&PixelShuffleAttrs::upscale_factor,
"Scale factor for spatial upsampling.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsNode);
};
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_NN_H_
include/tvm/relax/attrs/op.h
+123
View File
@@ -0,0 +1,123 @@
/*
* 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/relax/attrs/op.h
* \brief Attributes for relax specific operators.
*/
#ifndef TVM_RELAX_ATTRS_OP_H_
#define TVM_RELAX_ATTRS_OP_H_
#include <tvm/ir/global_info.h>
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in call_tir_with_grad */
struct CallTIRWithGradAttrs : public AttrsNode {
ffi::String te_grad_name;
ffi::Map<ffi::String, Any> te_grad_kwargs;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<CallTIRWithGradAttrs>()
.def_ro(
"te_grad_name", &CallTIRWithGradAttrs::te_grad_name,
"The name of the te gradient function associated with this call_tir_with_grad node.")
.def_ro("te_grad_kwargs", &CallTIRWithGradAttrs::te_grad_kwargs,
"The keyword arguments passed to the te gradient function.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallTIRWithGradAttrs", CallTIRWithGradAttrs,
AttrsNode);
}; // struct CallTIRAttrs
/*! \brief Attributes used in call_tir_inplace */
struct CallTIRInplaceAttrs : public AttrsNode {
/*!
* \brief Indices that describe which input corresponds to which output.
*
* If the `i`th member has the value `k` >= 0, then that means that input `k` should be used to
* store the `i`th output. If an element has the value -1, that means a new tensor should be
* allocated for that output.
*/
ffi::Array<int64_t> inplace_indices;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<CallTIRInplaceAttrs>().def_ro("inplace_indices",
&CallTIRInplaceAttrs::inplace_indices);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallTIRInplaceAttrs", CallTIRInplaceAttrs,
AttrsNode);
}; // struct CallTIRInplaceAttrs
/*! \brief Attributes used in call_inplace_packed */
struct CallInplacePackedAttrs : public AttrsNode {
/*!
* \brief Indices that describe which input corresponds to which output.
*
* If the `i`th member has the value `k` >= 0, then that means that input `k` should be used to
* store the `i`th output. If an element has the value -1, that means the output will be newly
* allocated.
*/
ffi::Array<int64_t> inplace_indices;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<CallInplacePackedAttrs>().def_ro("inplace_indices",
&CallInplacePackedAttrs::inplace_indices);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallInplacePackedAttrs", CallInplacePackedAttrs,
AttrsNode);
}; // struct CallInplacePackedAttrs
/*! \brief Attributes used in to_vdevice */
struct ToVDeviceAttrs : public AttrsNode {
VDevice dst_vdevice;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ToVDeviceAttrs>().def_ro("dst_vdevice", &ToVDeviceAttrs::dst_vdevice,
"The destination device where the data is copied to.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ToVDeviceAttrs", ToVDeviceAttrs, AttrsNode);
}; // struct ToVDeviceAttrs
/*! \brief Attributes used in hint_on_device */
struct HintOnDeviceAttrs : public AttrsNode {
int32_t device_type;
int32_t index;
MemoryScope memory_scope;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<HintOnDeviceAttrs>()
.def_ro("device_type", &HintOnDeviceAttrs::device_type,
"The device type where the data is supposed to be executed.")
.def_ro("index", &HintOnDeviceAttrs::index, "The device id.")
.def_ro("memory_scope", &HintOnDeviceAttrs::memory_scope, "The device memory scope.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.HintOnDeviceAttrs", HintOnDeviceAttrs, AttrsNode);
}; // struct HintOnDeviceAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_OP_H_
include/tvm/relax/attrs/qdq.h
+52
View File
@@ -0,0 +1,52 @@
/*
* 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 include/tvm/relax/attrs/qdq.h
* \brief Attributes for quantize/dequantize operators.
*/
#ifndef TVM_RELAX_ATTRS_QDQ_H_
#define TVM_RELAX_ATTRS_QDQ_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes for relax.quantize/relax.dequantize operator */
struct QuantizeAttrs : public AttrsNode {
DLDataType out_dtype;
int axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<QuantizeAttrs>()
.def_ro("out_dtype", &QuantizeAttrs::out_dtype, "Output data type.")
.def_ro("axis", &QuantizeAttrs::axis,
"The output channel axis for channel wise quantization/dequantization. "
"Default value is -1, which corresponds to the last axis.",
refl::DefaultValue(-1));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.QuantizeAttrs", QuantizeAttrs, AttrsNode);
}; // QuantizeAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_QDQ_H_
include/tvm/relax/attrs/sampling.h
+49
View File
@@ -0,0 +1,49 @@
/*
* 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/relax/attrs/sampling.h
* \brief Attributes for sampling operators.
*/
#ifndef TVM_RELAX_ATTRS_SAMPLING_H_
#define TVM_RELAX_ATTRS_SAMPLING_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in multinomial_from_uniform operator */
struct MultinomialFromUniformAttrs : public AttrsNode {
DLDataType dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MultinomialFromUniformAttrs>().def_ro(
"dtype", &MultinomialFromUniformAttrs::dtype, "Data type of the output indices.",
refl::DefaultValue((DLDataType{kDLInt, 64, 1})));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultinomialFromUniformAttrs",
MultinomialFromUniformAttrs, AttrsNode);
}; // struct MultinomialFromUniformAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_SAMPLING_H_
include/tvm/relax/attrs/search.h
+69
View File
@@ -0,0 +1,69 @@
/*
* 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/relax/attrs/search.h
* \brief Attributes for search operators.
*/
#ifndef TVM_RELAX_ATTRS_SEARCH_H_
#define TVM_RELAX_ATTRS_SEARCH_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes for search operators */
struct ArgmaxArgminAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
bool keepdims;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ArgmaxArgminAttrs>()
.def_ro("axis", &ArgmaxArgminAttrs::axis,
"The axis along which to perform the argmin/argmax.")
.def_ro("keepdims", &ArgmaxArgminAttrs::keepdims,
"If this is set to `True`, the reduced axis is left in the result as dimension "
"with size "
"one.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgmaxArgminAttrs", ArgmaxArgminAttrs, AttrsNode);
}; // struct ArgmaxArgminAttrs
/*! \brief Attributes for bucketize operator */
struct BucketizeAttrs : public tvm::AttrsNode {
bool out_int32;
bool right;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BucketizeAttrs>()
.def_ro("out_int32", &BucketizeAttrs::out_int32,
"Indicate the output datatype, int32 if True, int64 otherwise.")
.def_ro("right", &BucketizeAttrs::right,
"Determines the behavior for values in boundaries");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BucketizeAttrs", BucketizeAttrs, AttrsNode);
}; // struct BucketizeAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_SEARCH_H_
include/tvm/relax/attrs/sorting.h
+107
View File
@@ -0,0 +1,107 @@
/*
* 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/relax/attrs/sorting.h
* \brief Attributes for sorting operators.
*/
#ifndef TVM_RELAX_ATTRS_SORTING_H_
#define TVM_RELAX_ATTRS_SORTING_H_
#include <tvm/relax/expr.h>
#include <tvm/tirx/index_map.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in sort operator */
struct SortAttrs : public AttrsNode {
int axis;
bool descending;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SortAttrs>()
.def_ro("axis", &SortAttrs::axis,
"Axis along which the sort is computed."
"The default the last axis is used.",
refl::DefaultValue(-1))
.def_ro("descending", &SortAttrs::descending,
"Whether to sort in descending order."
"If it is not specified, it defaults to the ascending order.",
refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SortAttrs", SortAttrs, AttrsNode);
}; // struct SortAttrs
/*! \brief Attributes used in argsort operator */
struct ArgsortAttrs : public AttrsNode {
int axis;
bool descending;
ffi::Optional<DLDataType> dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ArgsortAttrs>()
.def_ro("axis", &ArgsortAttrs::axis,
"Axis along which the argsort is computed."
"The default the last axis is used.",
refl::DefaultValue(-1))
.def_ro("descending", &ArgsortAttrs::descending,
"Whether to argsort in descending order."
"If it is not specified, it defaults to the ascending order.",
refl::DefaultValue(false))
.def_ro("dtype", &ArgsortAttrs::dtype, "DType of the output indices.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, AttrsNode);
}; // struct ArgsortAttrs
/*! \brief Attributes used in topk operator */
struct TopKAttrs : public AttrsNode {
int k;
int axis;
bool largest;
ffi::String ret_type;
ffi::Optional<DLDataType> dtype;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TopKAttrs>()
.def_ro("k", &TopKAttrs::k, "Number of top elements to select")
.def_ro("axis", &TopKAttrs::axis, "Axis along which to sort the input tensor.",
refl::DefaultValue(-1))
.def_ro("ret_type", &TopKAttrs::ret_type,
"The return type [both, values, indices]."
"both - return both top k data and indices."
"values - return top k data only."
"indices - return top k indices only.",
refl::DefaultValue("both"))
.def_ro("largest", &TopKAttrs::largest,
"Whether to return largest or smallest elements."
"By default, return the largest k elements.",
refl::DefaultValue(true))
.def_ro("dtype", &TopKAttrs::dtype, "Data type of the output indices.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, AttrsNode);
}; // struct TopKAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_SORTING_H_
include/tvm/relax/attrs/statistical.h
+74
View File
@@ -0,0 +1,74 @@
/*
* 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/relax/attrs/statistical.h
* \brief Attributes for statistical operators.
*/
#ifndef TVM_RELAX_ATTRS_STATISTICAL_H_
#define TVM_RELAX_ATTRS_STATISTICAL_H_
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes for statistical operators */
struct StatisticalAttrs : public AttrsNode {
ffi::Optional<ffi::Array<int64_t>> axis;
bool keepdims;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<StatisticalAttrs>()
.def_ro("axis", &StatisticalAttrs::axis,
"The axis or axes along which to perform the reduction.")
.def_ro("keepdims", &StatisticalAttrs::keepdims,
"If this is set to `True`, the reduced axes are left in the result as dimension "
"with size "
"one.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StatisticalAttrs", StatisticalAttrs, AttrsNode);
}; // struct StatisticalAttrs
/*! \brief Attributes used in scan operators like cumsum, cumprod */
struct ScanopAttrs : public AttrsNode {
ffi::Optional<int64_t> axis;
ffi::Optional<DLDataType> dtype;
bool exclusive = false;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ScanopAttrs>()
.def_ro("axis", &ScanopAttrs::axis,
"The axis along which to perform the scan computation."
"The default (None) is to compute over the flattened array.")
.def_ro("dtype", &ScanopAttrs::dtype,
"The output data type."
"If dtype is not specified, it defaults to the dtype of input data.")
.def_ro("exclusive", &ScanopAttrs::exclusive, "The first element is not included",
refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScanopAttrs", ScanopAttrs, AttrsNode);
}; // struct ScanopAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_STATISTICAL_H_
include/tvm/relax/attrs/vision.h
+182
View File
@@ -0,0 +1,182 @@
/*
* 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/relax/attrs/vision.h
* \brief Auxiliary attributes for vision operators.
*/
#ifndef TVM_RELAX_ATTRS_VISION_H_
#define TVM_RELAX_ATTRS_VISION_H_
#include <tvm/ffi/string.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/type.h>
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*! \brief Attributes used in AllClassNonMaximumSuppression operator */
struct AllClassNonMaximumSuppressionAttrs : public AttrsNode {
ffi::String output_format;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AllClassNonMaximumSuppressionAttrs>().def_ro(
"output_format", &AllClassNonMaximumSuppressionAttrs::output_format,
"Output format, onnx or tensorflow. Returns outputs in a way that can be easily "
"consumed by each frontend.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllClassNonMaximumSuppressionAttrs",
AllClassNonMaximumSuppressionAttrs, AttrsNode);
}; // struct AllClassNonMaximumSuppressionAttrs
/*! \brief Attributes used in ROIAlign operator */
struct ROIAlignAttrs : public AttrsNode {
ffi::Array<int64_t> pooled_size;
double spatial_scale;
int sample_ratio;
bool aligned;
ffi::String layout;
ffi::String mode;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ROIAlignAttrs>()
.def_ro("pooled_size", &ROIAlignAttrs::pooled_size, "Output size of roi align.")
.def_ro("spatial_scale", &ROIAlignAttrs::spatial_scale,
"Ratio of input feature map height (or width) to raw image height (or width).")
.def_ro("sample_ratio", &ROIAlignAttrs::sample_ratio,
"Optional sampling ratio of ROI align, using adaptive size by default.")
.def_ro("aligned", &ROIAlignAttrs::aligned,
"Whether to use the aligned ROIAlign semantics without the legacy 1-pixel clamp.")
.def_ro("layout", &ROIAlignAttrs::layout, "Dimension ordering of the input data.")
.def_ro("mode", &ROIAlignAttrs::mode, "Mode for ROI Align. Can be 'avg' or 'max'.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIAlignAttrs", ROIAlignAttrs, AttrsNode);
}; // struct ROIAlignAttrs
/*! \brief Attributes used in ROIPool operator */
struct ROIPoolAttrs : public AttrsNode {
ffi::Array<int64_t> pooled_size;
double spatial_scale;
ffi::String layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ROIPoolAttrs>()
.def_ro("pooled_size", &ROIPoolAttrs::pooled_size, "Output size of roi pool.")
.def_ro("spatial_scale", &ROIPoolAttrs::spatial_scale,
"Ratio of input feature map height (or width) to raw image height (or width).")
.def_ro("layout", &ROIPoolAttrs::layout, "Dimension ordering of the input data.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIPoolAttrs", ROIPoolAttrs, AttrsNode);
}; // struct ROIPoolAttrs
/*! \brief Attributes used in GetValidCounts operator */
struct GetValidCountsAttrs : public AttrsNode {
double score_threshold;
int id_index;
int score_index;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GetValidCountsAttrs>()
.def_ro("score_threshold", &GetValidCountsAttrs::score_threshold,
"Lower limit of score for valid bounding boxes.")
.def_ro("id_index", &GetValidCountsAttrs::id_index,
"Index of the class categories, -1 to disable.")
.def_ro("score_index", &GetValidCountsAttrs::score_index,
"Index of the scores/confidence of boxes.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GetValidCountsAttrs", GetValidCountsAttrs,
AttrsNode);
}; // struct GetValidCountsAttrs
/*! \brief Attributes used in NonMaximumSuppression operator */
struct NonMaximumSuppressionAttrs : public AttrsNode {
int max_output_size;
double iou_threshold;
bool force_suppress;
int top_k;
int coord_start;
int score_index;
int id_index;
bool return_indices;
bool invalid_to_bottom;
double soft_nms_sigma;
double score_threshold;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<NonMaximumSuppressionAttrs>()
.def_ro("max_output_size", &NonMaximumSuppressionAttrs::max_output_size,
"Max number of output valid boxes, -1 for no limit.")
.def_ro("iou_threshold", &NonMaximumSuppressionAttrs::iou_threshold,
"Non-maximum suppression IoU threshold.")
.def_ro("force_suppress", &NonMaximumSuppressionAttrs::force_suppress,
"Whether to suppress all detections regardless of class_id.")
.def_ro("top_k", &NonMaximumSuppressionAttrs::top_k,
"Keep maximum top k detections before nms, -1 for no limit.")
.def_ro("coord_start", &NonMaximumSuppressionAttrs::coord_start,
"Start index of the consecutive 4 coordinates.")
.def_ro("score_index", &NonMaximumSuppressionAttrs::score_index,
"Index of the scores/confidence of boxes.")
.def_ro("id_index", &NonMaximumSuppressionAttrs::id_index,
"Index of the class categories, -1 to disable.")
.def_ro("return_indices", &NonMaximumSuppressionAttrs::return_indices,
"Whether to return box indices in input data.")
.def_ro("invalid_to_bottom", &NonMaximumSuppressionAttrs::invalid_to_bottom,
"Whether to move all valid bounding boxes to the top.")
.def_ro("soft_nms_sigma", &NonMaximumSuppressionAttrs::soft_nms_sigma,
"Sigma for soft-NMS; 0.0 means standard hard NMS.")
.def_ro("score_threshold", &NonMaximumSuppressionAttrs::score_threshold,
"Score threshold for soft-NMS validity check; 0.0 when unused.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NonMaximumSuppressionAttrs",
NonMaximumSuppressionAttrs, AttrsNode);
}; // struct NonMaximumSuppressionAttrs
/*! \brief Attributes for multibox_transform_loc (SSD / TFLite-style box decode). */
struct MultiboxTransformLocAttrs : public AttrsNode {
bool clip;
double threshold;
ffi::Array<double> variances;
bool keep_background;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MultiboxTransformLocAttrs>()
.def_ro("clip", &MultiboxTransformLocAttrs::clip,
"Clip decoded ymin,xmin,ymax,xmax to [0,1].")
.def_ro("threshold", &MultiboxTransformLocAttrs::threshold,
"After softmax, zero scores strictly below this value.")
.def_ro("variances", &MultiboxTransformLocAttrs::variances,
"(x,y,w,h) scales = TFLite 1/x_scale,1/y_scale,1/w_scale,1/h_scale on "
"encodings. Very large w/h scales can overflow exp in decode.")
.def_ro("keep_background", &MultiboxTransformLocAttrs::keep_background,
"If false, force output scores[:,0,:] to 0 (background class).");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultiboxTransformLocAttrs",
MultiboxTransformLocAttrs, AttrsNode);
}; // struct MultiboxTransformLocAttrs
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ATTRS_VISION_H_
include/tvm/relax/backend.h
+51
View File
@@ -0,0 +1,51 @@
/*
* 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/relax/backend.h
* \brief Relax backend specific transformation passes.
*/
#ifndef TVM_RELAX_BACKEND_H_
#define TVM_RELAX_BACKEND_H_
#include <tvm/relax/transform.h>
namespace tvm {
namespace relax {
namespace transform {
/*!
* \brief Perform builtin lowering to map most of the op to VM builtin functions.
*
* \return The Pass.
*/
TVM_DLL Pass LowerRuntimeBuiltin();
/*!
* \brief Lower the shape expression in relax to VM shape heap and TIR functions.
*
* \return The Pass.
*/
TVM_DLL Pass VMShapeLower();
} // namespace transform
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_BACKEND_H_
include/tvm/relax/backend/adreno/transform.h
+67
View File
@@ -0,0 +1,67 @@
/*
* 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/relax/backend/adreno/transform.h
* \brief Adreno GPU specific transformation passes.
*/
#ifndef TVM_RELAX_BACKEND_ADRENO_TRANSFORM_H_
#define TVM_RELAX_BACKEND_ADRENO_TRANSFORM_H_
#include <tvm/relax/expr.h>
#include <tvm/relax/transform.h>
namespace tvm {
namespace relax {
namespace backend {
namespace adreno {
namespace transform {
using Pass = tvm::transform::Pass;
using PassInfo = tvm::transform::PassInfo;
using PassContext = tvm::transform::PassContext;
using Function = tvm::relax::Function;
using DataflowBlock = tvm::relax::DataflowBlock;
using tvm::relax::transform::CreateFunctionPass;
using tvm::transform::CreateModulePass;
/*!
* \brief This pass is designed to annotate the memory scope information via VDevice attribute.
* This pass need operator attrbutes which in general vanish aftre legalization.
* FuseOps and FuseTIR are modified to pass on the operator specific attributes and also
* op_pattern details as part of the PrimFunc. This pass is Adreno specific and annotates each
* BindingVar with appropriate HintInDevice. RealizeVDevice pass followed by handles these hints.
* Followed by this pass we also invoke SpecializePrimFuncBasedOnCallSite which updates the
* var_buffer_map based on this new VDevice information.
*/
TVM_DLL Pass AnnotateCustomMemoryScope(Target target);
/*
* \brief This is a texture specific pass that can optimize unnecessary to_device copies.
* Like texture_scope -> ToVDevice -> global scope. In this case the producer can directly
* store into global scope avoiding unnecessary device copy.
*/
TVM_DLL Pass FoldVDeviceScopeChange();
} // namespace transform
} // namespace adreno
} // namespace backend
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_BACKEND_ADRENO_TRANSFORM_H_
include/tvm/relax/binding_rewrite.h
+118
View File
@@ -0,0 +1,118 @@
/*
* 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/relax/binding_rewrite.h
* \brief An IR rewriter to easily add/remove/replace bindings (statements).
*/
#ifndef TVM_RELAX_BINDING_REWRITE_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr.h>
#include <map>
#include <set>
#include <type_traits>
#include <utility>
#include <vector>
namespace tvm {
namespace relax {
/*! \brief Statement rewriter for relax.DataflowBlock. */
class DataflowBlockRewriteNode : public ffi::Object {
public:
/*! \brief Replace all uses of old_var with new_var. */
void ReplaceAllUses(Var old_var, Var new_var);
/*! \brief Insert a Binding statement. */
void Add(Binding binding);
/*! \brief Insert an expression as VarBinding with variable name. */
void Add(ffi::String var_name, Expr expr, bool is_dfvar = false) {
auto var = is_dfvar ? DataflowVar(var_name, GetType(expr)) //
: Var(var_name, GetType(expr));
Add(VarBinding(std::move(var), std::move(expr)));
}
/*! \brief Insert an expression as VarBinding with automatic variable name. */
void Add(Expr expr, bool is_dfvar = false) {
Add(name_supply_->FreshName("tmp"), expr, is_dfvar);
}
/*! \brief Remove the definition statement of an unused variable. */
void RemoveUnused(Var unused, bool allow_undef = false);
/*! \brief Remove the definition statements of all unused variables. */
void RemoveAllUnused();
/*! \brief The rewritten dataflow block. */
DataflowBlock MutatedDataflowBlock() { return dfb_; }
/*! \brief The rewritten function. */
Function MutatedFunc() { return root_fn_.value(); }
/*! \brief The rewritten IRModule. */
IRModule MutateIRModule(IRModule irmod);
/*! \brief Visit attributes. */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DataflowBlockRewriteNode>()
.def_ro("dfb", &DataflowBlockRewriteNode::dfb_)
.def_ro("root_fn", &DataflowBlockRewriteNode::root_fn_);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DataflowBlockRewrite", DataflowBlockRewriteNode,
ffi::Object);
protected:
friend class DataflowBlockRewrite;
DataflowBlock dfb_; //!< The rewritten dataflow block.
ffi::Optional<Function> root_fn_; //!< The rewritten function.
const FunctionNode* original_fn_ptr_; //!< Pointer to the original function.
ffi::Map<Var, ffi::Array<Var>> to_users_; //!< Map from variable to its users.
ffi::Array<Var> fn_outputs_; //!< Variables required by function outputs.
private:
UniqueNameSupply name_supply_; //!< Unique name supply for tracking and generating unique names.
};
/*!
* \brief A statement rewriter for relax.DataflowBlock.
* \sa DataflowBlockRewriteNode
*/
class DataflowBlockRewrite : public ffi::ObjectRef {
public:
TVM_DLL explicit DataflowBlockRewrite(DataflowBlock dfb, Function root_fn);
/*!
* \brief mutable accessor.
* \return mutable access pointer.
*/
DataflowBlockRewriteNode* operator->() {
TVM_FFI_ICHECK(get() != nullptr);
return static_cast<DataflowBlockRewriteNode*>(get_mutable());
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowBlockRewrite, ffi::ObjectRef,
DataflowBlockRewriteNode);
};
} // namespace relax
} // namespace tvm
#define TVM_RELAX_BINDING_REWRITE_H_
#endif // TVM_RELAX_BINDING_REWRITE_H_
include/tvm/relax/block_builder.h
+320
View File
@@ -0,0 +1,320 @@
/*
* 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/relax/block_builder.h
* \brief The utility for constructing Relax binding blocks.
*/
#ifndef TVM_RELAX_BLOCK_BUILDER_H_
#define TVM_RELAX_BLOCK_BUILDER_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/utils.h>
#include <tvm/runtime/base.h>
namespace tvm {
namespace relax {
/*!
* \brief A builder to build Relax binding blocks.
*
* BlockBuilder provides the following three categories
* of main functionalities for IR building and transformations:
*
* - Global context management: manages the IRModule,
* allowing query, update the surrounding global context.
* Provide context tools for analysis.
* - Scope management:
* - Manages block scopes for bulding nested blocks.
* - Emit bindings to the current scope.
* - Construct blocks by calling EndScope.
* - Normalization: Take an Expr, normalize it
* to deduce shape/type, turn things into normal forms.
*
* Importantly, these three categories of features can be dependent
* on each other. For example, when we emit into scope we will call
* normalize to ensure the code is in normal form. Similarly, when we
* normalize we could choose to emit into the current context.
*
* We would encourage the developers to keep these three category
* in mind when using and developing BlockBuilder, we can group
* the code in a logically clean way.
*
* BlockBuilderNode is implemented as a virtual interface to
* allow logically grouped implementation and internal data
* structures that are hidden from the users.
*/
class BlockBuilderNode : public ffi::Object {
public:
//-------------------------------
// Global Context management
//-------------------------------
/*!
* \brief Get the unique name supply for generating unique names.
*
* \return The unique name supply.
*/
virtual UniqueNameSupply name_supply() = 0;
/*!
* \brief Get the context IRModule in this builder.
*
* \note The context
* \return The IRModule in this BlockBuilder.
*/
virtual IRModule GetContextIRModule() const = 0;
/*!
* \brief Finalize the building process and return the result IRModule. Possibly rename
* GlobalVars in the IRModule to ensure name uniqueness and the invariant:
* every public function has the same name as its "global_symbol" attribute.
*
* \note this method should be called only once at the end of the building process, since it may
* invalidate global vars previously returned by this builder. See also
* transform::NormalizeGlobalVar.
*
* \return The result IRModule.
*/
virtual IRModule Finalize() = 0;
/*!
* \brief Add a Relax function or a TIR PrimFunc to internal context module.
* \param func The function to be added.
* \param func_name_hint The name hint of the function to be added.
* \note If the function to be added already exists, return its
* GlobalVar directly.
* \return The global var bound to the added function.
*/
virtual GlobalVar AddFunction(const BaseFunc& func, ffi::String func_name_hint) = 0;
/*!
* \brief Update a Relax function or a TIR PrimFunc in the internal context module.
* \param gv The global var referring the function to be updated.
* \param function The updated function.
*/
virtual void UpdateFunction(const GlobalVar& gv, BaseFunc function) = 0;
//-------------------------------
// Scope management
//-------------------------------
/*!
* \brief Lookup the binding value that var binds to in the current emitted sequences.
* \param var The input var.
* \return The Expr bound to the input \p var.
* \note For function parameters, this function returns std::nullopt.
*/
virtual ffi::Optional<Expr> LookupBinding(const Var& var) = 0;
/*!
* \brief Begin a new scope, with optional parameters that
* are visible within the scope.
*
* Symbolic variables from the parent scope are not available.
*
* \param params Parameters that are visible within the scope.
*
* \note This function should be called when new scope is introduced
* (e.g. function bodies) to properly track the variable
* availability and help the best effort deduction.
*
* \sa EndScope
*/
virtual void BeginScope(ffi::Optional<ffi::Array<Var>> params) = 0;
/*!
* \brief Begin a new scope, which inherits visible parameters from
* its parent scope.
*
* Symbolic variables from the parent scope are available.
*
* \note This function should be called when an inner scope is
* introduced (e.g. conditional branches) to properly track
* the variable availability and help the best effort
* deduction.
*
* \sa EndScope
*/
virtual void BeginInnerScope() = 0;
/*!
* \brief Append a definition to the current scope.
*
* \param var A variable within the current scope.
*
* \note This function should be called when a new variable is
* defined that may impact struct inference (e.g. MatchCast)
* to properly track the variable availability and help the
* best effort deduction.
*
* \sa EndScope
*/
virtual void AddDefinitionToScope(Var var) = 0;
/*! \brief End the previously defined scope. */
virtual void EndScope() = 0;
/*! \brief Begin to build a DataflowBlock. */
virtual void BeginDataflowBlock() = 0;
/*! \brief Begin to build a BindingBlock. */
virtual void BeginBindingBlock() = 0;
/*!
* \brief End building a BindingBlock.
* \return The BindingBlock being built.
*/
virtual BindingBlock EndBlock() = 0;
/*!
* \brief Check if the block being built is DataflowBlock or not.
* \return A boolean that indicates if the block being built is DataflowBlock or not.
*/
virtual bool CurrentBlockIsDataFlow() = 0;
/*!
* \brief Emits an Expr, and returns the variable it is bound to.
* \param expr The Expr to be emitted.
* \param name_hint Name hint for the bound variable.
* \return The new variable that \p expr is bound to.
*
* \note This Emit function normalizes the \p expr, and
* performs shape and type deductions by calling Normalize.
*/
virtual Var Emit(Expr expr, ffi::String name_hint = "") = 0;
/*!
* \brief Emit a MatchCast.
* \param value The input value.
* \param ty The type to be matched.
* \param name_hint Name hint for the bound variable.
* \return The variable bound to the MatchCast.
*/
virtual Var EmitMatchCast(Expr value, Type ty, ffi::String name_hint = "") = 0;
/*!
* \brief Generate an output for the current dataflow block.
* \param output The output variable of the block.
* \param name_hint Name hint for the bound variable.
* \return The variable bound to \p output.
*/
virtual Var EmitOutput(Expr output, ffi::String name_hint = "") = 0;
/*!
* \brief Emit a binding that is already normalized.
*
* \param normalized_binding A binding whose value is already normalized.
*
* \note This function requires binding to be pre-normalized.
*/
virtual void EmitNormalized(Binding normalized_binding) = 0;
/*!
* \brief Convert an expression to normal form, and try to eagerly infer types and shapes.
* \param expr The input expression.
* \return The normalized expression.
*
* \note Invariant: If any of the sub expr have ty field.
* they must have already been normalized.
*/
virtual Expr Normalize(const Expr& expr) = 0;
/*!
* \brief Normalize argument to a call or another IRNode.
* \param expr The input expression.
* \return The normalized expression.
*
* \note This function will create a binding var for non-leaf expressions such as Call.
*/
virtual Expr NormalizeArgument(const Expr& expr) = 0;
/*!
* \brief Get the analyzer of the BlockBuilder.
* \return The BlockBuilder's arithmetic analyzer.
*/
virtual arith::Analyzer GetAnalyzer() = 0;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("relax.BlockBuilder", BlockBuilderNode, ffi::Object);
};
class BlockBuilder : public ffi::ObjectRef {
public:
/*!
* \brief Create a BlockBuilder.
*
* \param ctx_mod Optional before-transformation context module for rewriting.
*
* \return The created BlockBuilder.
*
* \note When rewriting an existing IRModule, it is important to pass it in as
* ctx_mod so you can lookup the context functions for cross function
* call analysis.
*/
TVM_DLL static BlockBuilder Create(ffi::Optional<IRModule> ctx_mod);
/*! \brief A marker struct to disable FNormalize
*
* This struct is used as a marker to disable the use of FNormalize
* by this block builder. This should only be used for TVMScript
* parsing, which may require producing un-normalized Relax IR for
* testing purposes, and to ensure that round-trips are unchanged.
*
* The name is deliberately verbose to draw attention during a code
* review. The explicit default constructor prevents aggregate
* initialization, ensuring that the full name of the marker struct
* appears at the callsite.
*
* This constructor is marked as no-lint to allow a zero-parameter
* constructor to be marked as explicit. The constructor must be
* explicit in order to disable aggregate initialization in C++17.
* While C++20 disables aggregate initialization when a
* user-declared constructor is present, C++17 only disables
* aggregate initialization when a user-defined constructor is
* present. Therefore, we need to mark the zero-parameter
* constructor as explicit in order to prevent aggregate
* initialization, and to ensure that the name appears at all
* callsites.
*/
struct DisableOperatorSpecificNormalizationForTVMScript {
explicit DisableOperatorSpecificNormalizationForTVMScript() = default; // NOLINT(*)
};
/*!
* \brief Create a BlockBuilder.
*
* \param ctx_mod Optional before-transformation context module for rewriting.
*
* \param tag An instance of DisableOperatorSpecificNormalizationForTVMScript
*
* \return The created BlockBuilder.
*
* \note When rewriting an existing IRModule, it is important to pass it in as
* ctx_mod so you can lookup the context functions for cross function
* call analysis.
*/
TVM_DLL static BlockBuilder Create(ffi::Optional<IRModule> ctx_mod,
DisableOperatorSpecificNormalizationForTVMScript tag);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BlockBuilder, ffi::ObjectRef, BlockBuilderNode);
};
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_BLOCK_BUILDER_H_
include/tvm/relax/dataflow_matcher.h
+107
View File
@@ -0,0 +1,107 @@
/*
* 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/relax/dataflow_matcher.h
* \brief A pattern matcher for matching dataflow properties.
*/
#ifndef TVM_RELAX_DATAFLOW_MATCHER_H_
#define TVM_RELAX_DATAFLOW_MATCHER_H_
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/relax/dataflow_pattern.h>
#include <memory>
namespace tvm {
namespace relax {
/**
* \brief Determine if a pattern matches an expression.
* \note The behavior of MatchExpr is to match a relax.Expr (`expr`) syntactically through
* one given pattern (`pattern`).
*
* \param pattern The pattern to match
* \param expr The expression to match
* \param bindings The mapping from relax.Var to relax.Expr
* \return true if matched
* \return false if unmatched
*/
bool MatchExpr(DFPattern pattern, Expr expr,
ffi::Optional<ffi::Map<Var, Expr>> bindings = std::nullopt);
/* \brief Similar to above, but return pairs of a matching pattern and an expression. */
ffi::Optional<ffi::Map<DFPattern, Expr>> ExtractMatchedExpr(
DFPattern pattern, Expr expr, ffi::Optional<ffi::Map<Var, Expr>> bindings = std::nullopt);
/**
* \brief Match a sub-graph in a DataflowBlock with a graph of patterns and return the mapping.
* \param ctx The graph-wise patterns.
* \param dfb The function to match.
* \return Matched patterns and corresponding bound variables
*/
TVM_DLL ffi::Optional<ffi::Map<DFPattern, Var>> MatchGraph(const PatternContext& ctx,
const DataflowBlock& dfb);
/**
* \brief Rewrite a function with the given pattern and the rewriter function.
* \param ctx The pattern constraint context under which rewriting takes place.
* \param rewriter The function to be called on a successful matching for rewriting.
Given the map of patterns and corresponding variables (bound variables or parameters),
it should return a map that specifies new values for matched bound variables.
* \param f The function to rewrite
* \return The rewritten or the input function, depending on the pattern matching result.
*/
TVM_DLL Function RewriteBindings(
const PatternContext& ctx,
ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var, Expr>)> rewriter,
Function f);
/**
* \brief Rewrite a function with the given pattern and the rewriter function.
*
* Pattern match and replace at an expression level. This level of
* granularity does not allow simultaneous replacement cannot be
* performed. In addition, removal of bindings cannot be performed
* explicitly, and is only done implicitly through RemoveAllUnused.
* See also `RewriteBindings`, which performs replacement on a
* block-level, and does not have these restrictions.
*
* \param pattern The pattern to be replaced
*
* \param rewriter The function to be called on a successful pattern
* matching. Given the matched expression and a map of sub-matches,
* it should return the replacement expression. If the expression
* doesn't require updating (e.g. replacement required checks beyond
* those expressed in the pattern), it should return the expression
* unmodified.
*
* \param func The function to rewrite
*
* \return The updated function, if any updates were applied.
*/
TVM_DLL Function RewriteCall(const DFPattern& pattern,
ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern, Expr>)> rewriter,
Function func);
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_DATAFLOW_MATCHER_H_
include/tvm/relax/dataflow_pattern.h
+978
View File
@@ -0,0 +1,978 @@
/*
* 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/relax/dataflow_pattern.h
* \brief A pattern language for matching dataflow properties.
*/
#ifndef TVM_RELAX_DATAFLOW_PATTERN_H_
#define TVM_RELAX_DATAFLOW_PATTERN_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/with_context.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
namespace tvm {
namespace arith {
class AnalyzerObj;
class Analyzer;
} // namespace arith
namespace relax {
class PatternSeq;
class CallPattern;
class OrPattern;
class AndPattern;
class NotPattern;
class ShapePattern;
class TypePattern;
class DataTypePattern;
class AttrPattern;
class SameShapeConstraint;
/*!
* \brief Create used-by relationship between lhs[-1] and rhs[0], with [*lhs, *rhs] returned.
*
* \param lhs Left hand side of the used-by relationship.
* \param rhs Right hand side of the used-by relationship.
* \param index lhs[-1] is used as the index'th argument of rhs[0].
* \return PatternSeq The concatenated sequence of [*lhs, *rhs].
*/
TVM_DLL PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index = -1);
/*! \brief Syntax sugar of UsedBy(lhs, rhs, -1). */
TVM_DLL PatternSeq operator^(const PatternSeq& lhs, const PatternSeq& rhs);
/*!
* \brief Create only-used-by relationship between lhs[-1] and rhs[0], with [*lhs, *rhs] returned.
*
* \param lhs Left hand side of the used-by relationship.
* \param rhs Right hand side of the used-by relationship.
* \param index lhs[-1] is used as the index'th argument of rhs[0].
* \return PatternSeq The concatenated sequence of [*lhs, *rhs].
*/
TVM_DLL PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index = -1);
/*! \brief Syntax sugar of OnlyUsedBy(lhs, rhs, -1). */
TVM_DLL PatternSeq operator>>(const PatternSeq& lhs, const PatternSeq& rhs);
/*!
* \brief Base type of all dataflow patterns.
* \sa DFPattern
*/
class DFPatternNode : public ffi::Object {
public:
static constexpr const uint32_t _type_child_slots = 21;
TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFPattern", DFPatternNode, ffi::Object);
};
/*!
* \brief Managed reference to dataflow patterns.
* \sa DFPatternNode
*/
class DFPattern : public ffi::ObjectRef {
public:
/*! \brief Syntatic Sugar for creating a CallPattern */
template <typename... Args>
CallPattern operator()(Args&&... args) const;
/*! \brief Syntatic Sugar for creating a CallPattern */
TVM_DLL CallPattern operator()(const std::vector<DFPattern>& args) const;
/*! \brief Syntatic Sugar for creating an OrPattern */
TVM_DLL OrPattern operator|(const DFPattern& other) const;
/*! \brief Syntatic Sugar for creating an AndPattern */
TVM_DLL AndPattern operator&(const DFPattern& other) const;
/*! \brief Syntatic Sugar for creating a NotPattern */
TVM_DLL NotPattern operator~() const;
/*! \brief Syntatic Sugar for creating an AttrPattern */
TVM_DLL AttrPattern HasAttr(const ffi::Map<ffi::String, Any>& attrs) const;
/*! \brief Syntatic Sugar for creating a TypePattern */
TVM_DLL TypePattern HasType(const Type& ty) const;
/*! \brief Syntatic Sugar for creating a DataTypePattern with a dtype */
TVM_DLL DataTypePattern HasDtype(DLDataType dtype) const;
/*! \brief Syntatic Sugar for creating a DataTypePattern with a data type's name */
TVM_DLL DataTypePattern HasDtype(const std::string& dtype) const;
/*! \brief Syntatic Sugar for creating a ShapePattern */
TVM_DLL ShapePattern HasShape(const ffi::Array<PrimExpr>& shape) const;
/*! \brief Syntatic Sugar for creating a ShapePattern */
TVM_DLL SameShapeConstraint HasSameShapeAs(const DFPattern& other) const;
/*! \brief Syntatic Sugar for duplicating the current pattern */
TVM_DLL DFPattern dup() const;
/*! \brief Implicit conversion from DFPattern to PatternSeq */
TVM_DLL operator PatternSeq() const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DFPattern, ffi::ObjectRef, DFPatternNode);
};
/*! \brief Constraint of a DFPattern edge (producer -> consumer) in graph-level matching */
struct PairCons {
/*! \brief Constraint types of the edge */
enum Type {
kUsedBy, /*!< producer ^ consumer */
kOnlyUsedBy, /*!< producer >> consumer */
} type = kUsedBy;
int index = -1; /*!< The argument index of the producer in the consumer caller site */
/*!
* \brief Construct a new PairCons object
*
* \param t The constraint type
* \param index The producer is called as the index'th argument of the consumer function.
*/
TVM_DLL explicit PairCons(Type t, int index = -1) : type(t), index(index) {}
bool operator==(const PairCons& other) const {
return type == other.type && index == other.index;
}
};
/*! \brief Additional constraints on the graph
*
* Unlike PairCons, these may relate nodes that are not directly
* connected by a DFPattern edge from producer to consumer. For
* example, constraining the two branches of an elementwise operation
* to have the same shape.
*/
class DFConstraintNode : public ffi::Object {
public:
/*! \brief Return the patterns on which the constraint depends */
virtual ffi::Array<DFPattern> GetDependentPatterns() const = 0;
/*! \brief Convert the constraint to a PrimExpr
*
* If the returned boolean parameter is true, then the returned
* expression is a necessary-and-sufficient condition for evaluating
* the constraint. In this case, the matcher may either mark the
* constraint as satisfied (no need to re-check later), or as failed
* (need to back-track).
*
* If the returned boolean parameter is false, then the returned
* expression is a necessary-but-not-sufficient condition for
* evaluating the constraint. In this case, the matcher may start
* backtracking as a result of a failed condition, but may not mark
* the constraint as satisfied. This typically occurs when the
* constraint involves a parameter that the matcher has not yet
* filled.
*
* \param match_state A function that can be called to check the
* current state of the match. The function takes as argument a
* pattern on which the constraint depends, and returns the relax
* variable matched by that pattern, or std::nullopt if the pattern
* has not yet been matched.
*
* \return A tuple of `PrimExpr` and `bool`. The first element is a
* necessary condition for the constraint to be satisfied. The
* second tuple element indicates whether the condition is also
* sufficient for the constraint to be satisfied.
*/
virtual std::tuple<PrimExpr, bool> AsCondition(
std::function<ffi::Optional<Var>(const DFPatternNode*)> match_state) const = 0;
static constexpr const uint32_t _type_child_slots = 1;
TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFConstraint", DFConstraintNode, ffi::Object);
};
class DFConstraint : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DFConstraint, ffi::ObjectRef, DFConstraintNode);
};
/*!
* \brief A sequence of DFPatterns that the previous DFPattern is connected to the next one.
* \sa PatternSeq
*/
class PatternSeqNode final : public ffi::Object {
public:
tvm::ffi::Array<DFPattern> patterns; /*!< The sequence of DFPatterns */
std::vector<PairCons> pair_constraints; /*!< Constraints between the previous and next patterns */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PatternSeqNode>().def_ro("patterns", &PatternSeqNode::patterns);
}
TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.PatternSeq", PatternSeqNode, ffi::Object);
};
/*!
* \brief Managed reference to pattern sequences.
* \sa PatternSeqNode
*/
class PatternSeq final : public ffi::ObjectRef {
public:
TVM_DLL explicit PatternSeq(DFPattern init_pattern);
TVM_DLL explicit PatternSeq(tvm::ffi::Array<DFPattern> patterns, bool only_used_by = false);
PatternSeq UsedBy(PatternSeq other, int index = -1) const;
PatternSeq OnlyUsedBy(PatternSeq other, int index = -1) const;
/*! \brief Syntatic Sugar for duplicating the current pattern sequence */
PatternSeq dup() const;
// friend functions
friend PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index);
friend PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PatternSeq, ffi::ObjectRef, PatternSeqNode);
};
/*!
* \brief A context to manage the graph-level pattern matching.
* \sa PatternContext
*/
class PatternContextNode : public ffi::Object {
public:
/*! \brief Constrainting matched graph with assertion to external uses */
enum ExternUse {
kMay, /*!< No constraints */
kMustNot, /*!< All nodes except outputs only have internal depedencies in the matched graph. */
} allow_extern_use = kMay;
// src node -> <dst node, constraint type> constraints.
// Dst nodes are kept in a vector to keep them ordered.
std::map<DFPattern, std::vector<std::pair<DFPattern, std::vector<PairCons>>>> edge_constraints;
// Underlying DFPattern nodes which the edge constraints may reference
// Kept as a separate vector of patterns to process constraints in a fixed order.
std::vector<DFPattern> src_ordered;
// Non-edge constraints
std::vector<DFConstraint> validation_constraints;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.PatternContext", PatternContextNode, ffi::Object);
};
/*!
* \brief Managed reference to a pattern context.
* \sa PatternContextNode
*/
class PatternContext : public ffi::ObjectRef {
public:
explicit PatternContext(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
TVM_DLL explicit PatternContext(ffi::ObjectPtr<ffi::Object> n) : ffi::ObjectRef(n) {}
TVM_DLL explicit PatternContext(bool incremental = false);
const PatternContextNode* operator->() const {
TVM_FFI_ICHECK(get() != nullptr);
return static_cast<const PatternContextNode*>(get());
}
PatternContextNode* operator->() {
TVM_FFI_ICHECK(get() != nullptr);
return static_cast<PatternContextNode*>(get_mutable());
}
/*!
* \brief Build an edge constraint between two patterns (producer and consumer).
*
* \param producer The pattern corresponding to the producer node.
* \param consumer The pattern corresponding to the consumer node.
* \param cons The constraint type. \sa PairCons
*/
void add_constraint(DFPattern producer, DFPattern consumer, PairCons cons) {
auto& pairs = (*this)->edge_constraints[producer];
auto it = std::find_if(pairs.begin(), pairs.end(),
[consumer](auto p) { return p.first == consumer; });
if (it == pairs.end()) {
pairs.emplace_back(consumer, std::vector{cons});
} else {
auto& vec = it->second;
TVM_FFI_ICHECK(std::find(vec.cbegin(), vec.cend(), cons) == vec.cend())
<< "Constraint already exists";
vec.push_back(cons);
}
auto& patterns = (*this)->src_ordered;
if (std::find(patterns.begin(), patterns.end(), producer) == patterns.end()) {
patterns.push_back(producer);
}
}
/*!
* \brief Add a validation constraint
*
* \param constraint The new constraint
*/
void add_constraint(DFConstraint constraint) {
(*this)->validation_constraints.push_back(constraint);
}
/*! \brief Get the constraint context object on the top of the stack */
TVM_DLL static ffi::Optional<PatternContext> Current();
/*! \brief The RAII-like entry of a constraint context scope */
TVM_DLL void EnterWithScope() const;
/*! \brief The RAII-like exit of a constraint context scope */
TVM_DLL void ExitWithScope() const;
private:
friend class With<PatternContext>;
};
/*!
* \brief Pattern for Relax Expression.
* \sa ExprPattern
*/
class ExprPatternNode : public DFPatternNode {
public:
Expr expr; /*!< The expression to match */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExprPatternNode>().def_ro("expr", &ExprPatternNode::expr);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExprPattern", ExprPatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to an ExprPattern.
* \sa ExprPatternNode
*/
class ExprPattern : public DFPattern {
public:
TVM_DLL explicit ExprPattern(Expr expr);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExprPattern, DFPattern, ExprPatternNode);
};
/*!
* \brief A Pattern to Match a Relax Variable.
* \note The name field matches any string if it is empty.
* \sa VarPattern
*/
class VarPatternNode : public DFPatternNode {
public:
ffi::String name;
const ffi::String& name_hint() const { return name; }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<VarPatternNode>().def_ro("name", &VarPatternNode::name);
}
static constexpr const uint32_t _type_child_slots = 1;
TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.VarPattern", VarPatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to a VarPattern.
* \sa VarPatternNode
*/
class VarPattern : public DFPattern {
public:
/*!
* \brief Create a pattern matching by variable name.
*
* \param name_hint Variable name to match. Any if empty ("").
*/
TVM_DLL VarPattern(ffi::String name_hint);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VarPattern, DFPattern, VarPatternNode);
};
/*!
* \brief A Pattern to Match a Relax Dataflow Variable
* \sa DataflowVarPattern
*/
class DataflowVarPatternNode : public VarPatternNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DataflowVarPatternNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataflowVarPattern", DataflowVarPatternNode,
VarPatternNode);
};
/*!
* \brief Managed reference to a DataflowVarPattern.
* \sa DataflowVarPatternNode
*/
class DataflowVarPattern : public DFPattern {
public:
/*! \sa VarPattern::VarPattern */
TVM_DLL DataflowVarPattern(ffi::String name_hint);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowVarPattern, DFPattern, DataflowVarPatternNode);
};
/*!
* \brief A Pattern to Match a Relax Global Variable
* \sa GlobalVarPattern
*/
class GlobalVarPatternNode : public VarPatternNode {
public:
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.GlobalVarPattern", GlobalVarPatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to a GlobalVarPattern.
* \sa GlobalVarPatternNode
*/
class GlobalVarPattern : public DFPattern {
public:
TVM_DLL GlobalVarPattern(ffi::String name_hint);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GlobalVarPattern, DFPattern, GlobalVarPatternNode);
};
/*!
* \brief A Pattern to Match a Relax Constant.
* \sa ConstantPattern
*/
class ConstantPatternNode : public DFPatternNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ConstantPatternNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ConstantPattern", ConstantPatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to a ConstantPattern.
* \sa ConstantPatternNode
*/
class ConstantPattern : public DFPattern {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ConstantPattern, DFPattern, ConstantPatternNode);
};
/*!
* \brief A pattern to match a callable node in Relax.
* \sa CallPattern
*/
class CallPatternNode : public DFPatternNode {
public:
/*!
* \note The op field can be:
* - relax::Op which corresponds to the primitive operators.
* - user defined functions (Function, GlobalVar, Var).
*/
DFPattern op; /*!< The operator (function) being invoked */
tvm::ffi::Array<DFPattern> args; /*!< The arguments of the function call */
/*!
* \note If varg_default_wildcard is true. Given args of [pA, pB], when matching a call whose
* arguments are [A, B, ...], the pattern will still match despite N(args) < N(call.args). That
* said, with varg_default_wildcard set to true, we match the args in the order we have, and
* regard the rest of the arguments as wildcards.
*/
bool varg_default_wildcard; /*!< N(args) can be < N(real args) by the padding of Wildcard */
// Todo(relax-team): Dataflow pattern for Type, and match ty_args
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<CallPatternNode>()
.def_ro("op", &CallPatternNode::op)
.def_ro("args", &CallPatternNode::args);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.CallPattern", CallPatternNode, DFPatternNode);
};
class CallPattern : public DFPattern {
public:
TVM_DLL CallPattern(DFPattern op, ffi::Array<DFPattern> args, bool varg_default_wildcard = false);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CallPattern, DFPattern, CallPatternNode);
};
/*!
* \brief A pattern to match an array of PrimExpr.
* \sa PrimArrPattern
* \note This is often used to match shapes specified as arguments to a function.
*/
class PrimArrPatternNode : public DFPatternNode {
public:
ffi::Array<PrimExpr> fields; /*!< The array to match */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PrimArrPatternNode>().def_ro("fields", &PrimArrPatternNode::fields);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.PrimArrPattern", PrimArrPatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to a PrimArrPattern.
* \sa PrimArrPatternNode
*/
class PrimArrPattern : public DFPattern {
public:
TVM_DLL PrimArrPattern(ffi::Array<PrimExpr> arr);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimArrPattern, DFPattern, PrimArrPatternNode);
};
/*!
* \brief A pattern to match a Relax Function
* \sa Function
* \sa FunctionPattern
*/
class FunctionPatternNode : public DFPatternNode {
public:
tvm::ffi::Array<DFPattern> params; /*!< The parameters of the function */
/*!
* \note Note that in Relax, the function body is a SeqExpr which contains
* 1) SeqExprNode::blocks, which is a list of blocks of statements; and 2)
* SeqExprNode::body, which is an Expr that can be anything. FunctionPattern
* only matches the body of the function (writing patterns to statements is tricky).
*/
DFPattern body; /*!< The body of the function */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FunctionPatternNode>()
.def_ro("params", &FunctionPatternNode::params)
.def_ro("body", &FunctionPatternNode::body);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.FunctionPattern", FunctionPatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to FunctionPatternNode.
* \sa FunctionPatternNode
*/
class FunctionPattern : public DFPattern {
public:
/*!
* \brief Constructor
* \param params The parameters of the function.
* \param body The body of the function.
*/
TVM_DLL FunctionPattern(tvm::ffi::Array<DFPattern> params, DFPattern body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FunctionPattern, DFPattern, FunctionPatternNode);
};
/*!
* \brief Pattern to match a tuple of ordered expressions.
* \sa TuplePattern
*/
class TuplePatternNode : public DFPatternNode {
public:
tvm::ffi::Array<DFPattern> fields; /*!< The fields of the tuple */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TuplePatternNode>().def_ro("fields", &TuplePatternNode::fields);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TuplePattern", TuplePatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to TuplePatternNode.
* \sa TuplePatternNode
*/
class TuplePattern : public DFPattern {
public:
TVM_DLL explicit TuplePattern(tvm::ffi::Array<DFPattern> fields);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TuplePattern, DFPattern, TuplePatternNode);
};
/*!
* \brief A pattern to match multiple expressions unorderedly.
* \sa UnorderedTuplePattern
*/
class UnorderedTuplePatternNode : public DFPatternNode {
public:
tvm::ffi::Array<DFPattern> fields; /*!< The fields of the tuple */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<UnorderedTuplePatternNode>().def_ro("fields",
&UnorderedTuplePatternNode::fields);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.UnorderedTuplePattern", UnorderedTuplePatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to UnorderedTuplePatternNode.
* \sa UnorderedTuplePatternNode
*/
class UnorderedTuplePattern : public DFPattern {
public:
TVM_DLL explicit UnorderedTuplePattern(tvm::ffi::Array<DFPattern> fields);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(UnorderedTuplePattern, DFPattern,
UnorderedTuplePatternNode);
};
/*!
* \brief A pattern to match n'th indexing to a tuple.
* \sa TupleGetItem
* \sa TupleGetItemPattern
*/
class TupleGetItemPatternNode : public DFPatternNode {
public:
DFPattern tuple; /*!< The tuple Expression */
int index; /*!< The index of the tuple with -1 meaning arbitrary */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TupleGetItemPatternNode>()
.def_ro("tuple", &TupleGetItemPatternNode::tuple)
.def_ro("index", &TupleGetItemPatternNode::index);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TupleGetItemPattern", TupleGetItemPatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to TupleGetItemPatternNode.
* \sa TupleGetItemPatternNode
*/
class TupleGetItemPattern : public DFPattern {
public:
TVM_DLL TupleGetItemPattern(DFPattern tuple, int index);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TupleGetItemPattern, DFPattern,
TupleGetItemPatternNode);
};
/*!
* \brief Match a conjunction of other patterns.
* \sa AndPattern
*/
class AndPatternNode : public DFPatternNode {
public:
DFPattern left; /*!< The left hand side of the conjunction */
DFPattern right; /*!< The right hand side of the conjunction */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AndPatternNode>()
.def_ro("left", &AndPatternNode::left)
.def_ro("right", &AndPatternNode::right);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AndPattern", AndPatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to AndPatternNode.
* \sa AndPatternNode
*/
class AndPattern : public DFPattern {
public:
TVM_DLL AndPattern(DFPattern lhs, DFPattern rhs);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AndPattern, DFPattern, AndPatternNode);
};
/*!
* \brief Match a disjunction of other patterns.
* \sa OrPattern
*/
class OrPatternNode : public DFPatternNode {
public:
DFPattern left; /*!< The left hand side of the disjunction */
DFPattern right; /*!< The right hand side of the disjunction */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<OrPatternNode>()
.def_ro("left", &OrPatternNode::left)
.def_ro("right", &OrPatternNode::right);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.OrPattern", OrPatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to OrPatternNode.
* \sa OrPatternNode
*/
class OrPattern : public DFPattern {
public:
TVM_DLL OrPattern(DFPattern left, DFPattern right);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(OrPattern, DFPattern, OrPatternNode);
};
/*!
* \brief Pattern for rejecting a certain pattern.
* \sa NotPattern
*/
class NotPatternNode : public DFPatternNode {
public:
DFPattern reject; /*!< The pattern to reject */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<NotPatternNode>().def_ro("reject", &NotPatternNode::reject);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.NotPattern", NotPatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to NotPatternNode.
* \sa NotPatternNode
*/
class NotPattern : public DFPattern {
public:
TVM_DLL NotPattern(DFPattern reject);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(NotPattern, DFPattern, NotPatternNode);
};
/*!
* \brief Wildcard Pattern is a pattern that can match anything.
* \sa WildcardPattern
*/
class WildcardPatternNode : public DFPatternNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<WildcardPatternNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.WildcardPattern", WildcardPatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to WildcardPatternNode.
* \sa WildcardPatternNode
*/
class WildcardPattern : public DFPattern {
public:
WildcardPattern();
explicit WildcardPattern(ffi::ObjectPtr<WildcardPatternNode> data)
: DFPattern(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
// Declaring WildcardPattern declared as non-nullable avoids the
// default zero-parameter constructor for ffi::ObjectRef with `data_ =
// nullptr`. This allows a zero-parameter constructor to be
// declared here, to create a valid wildcard instance.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WildcardPattern, DFPattern, WildcardPatternNode);
};
/*!
* \brief Pattern for matching a certain type.
* \sa TypePattern
*/
class TypePatternNode : public DFPatternNode {
public:
DFPattern pattern; /*!< The pattern to match */
Type ty = Type::Missing(); /*!< The type to match */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TypePatternNode>()
.def_ro("pattern", &TypePatternNode::pattern)
.def_ro("ty", &TypePatternNode::ty);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TypePattern", TypePatternNode, DFPatternNode);
};
class TypePattern : public DFPattern {
public:
TVM_DLL TypePattern(DFPattern pattern, Type ty);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TypePattern, DFPattern, TypePatternNode);
};
/*!
* \brief A pattern that asserting a root pattern has a certain shape.
* \sa ShapePattern
*/
class ShapePatternNode : public DFPatternNode {
public:
DFPattern pattern; /*!< The root pattern to match */
ffi::Array<PrimExpr> shape; /*!< The shape to match */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ShapePatternNode>()
.def_ro("pattern", &ShapePatternNode::pattern)
.def_ro("shape", &ShapePatternNode::shape);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ShapePattern", ShapePatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to ShapePatternNode.
* \sa ShapePatternNode
*/
class ShapePattern : public DFPattern {
public:
TVM_DLL ShapePattern(DFPattern pattern, ffi::Array<PrimExpr> type);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ShapePattern, DFPattern, ShapePatternNode);
};
/*!
* \brief A pattern that asserting multiple root patterns have the same shape
* \sa SameShapePattern
*/
class SameShapeConstraintNode : public DFConstraintNode {
public:
ffi::Array<DFPattern> args; /*!< The patterns with matching shapes */
ffi::Array<DFPattern> GetDependentPatterns() const override { return args; }
std::tuple<PrimExpr, bool> AsCondition(
std::function<ffi::Optional<Var>(const DFPatternNode*)> match_state) const override;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SameShapeConstraintNode>().def_ro("args", &SameShapeConstraintNode::args);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.SameShapeConstraint", SameShapeConstraintNode,
DFConstraintNode);
};
/*!
* \brief Managed reference to SameShapePatternNode.
* \sa SameShapePatternNode
*/
class SameShapeConstraint : public DFConstraint {
public:
TVM_DLL SameShapeConstraint(ffi::Array<DFPattern> args);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SameShapeConstraint, DFConstraint,
SameShapeConstraintNode);
};
/*!
* \brief A pattern that asserting a root pattern has a certain data type.
* \sa DataTypePattern
*/
class DataTypePatternNode : public DFPatternNode {
public:
DFPattern pattern; /*!< The root pattern to match */
DLDataType dtype; /*!< The data type to match */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DataTypePatternNode>()
.def_ro("pattern", &DataTypePatternNode::pattern)
.def_ro("dtype", &DataTypePatternNode::dtype);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataTypePattern", DataTypePatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to DataTypePatternNode.
* \sa DataTypePatternNode
*/
class DataTypePattern : public DFPattern {
public:
TVM_DLL DataTypePattern(DFPattern pattern, DLDataType dtype);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataTypePattern, DFPattern, DataTypePatternNode);
};
/*!
* \brief A pattern that asserting a root pattern has certain attributes.
* \sa AttrPattern
*/
class AttrPatternNode : public DFPatternNode {
public:
DFPattern pattern; /*!< The root pattern to match */
DictAttrs attrs; /*!< The attributes (a map/dictionary) to match */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AttrPatternNode>()
.def_ro("pattern", &AttrPatternNode::pattern)
.def_ro("attrs", &AttrPatternNode::attrs);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AttrPattern", AttrPatternNode, DFPatternNode);
};
/*!
* \brief Managed reference to AttrPatternNode.
* \sa AttrPatternNode
*/
class AttrPattern : public DFPattern {
public:
TVM_DLL AttrPattern(DFPattern pattern, DictAttrs attrs);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AttrPattern, DFPattern, AttrPatternNode);
};
/*!
* \brief A pattern of external function.
* \sa ExternFunc
* \sa ExternFuncPattern
*/
class ExternFuncPatternNode : public DFPatternNode {
public:
ffi::String global_symbol_; /*!< The global symbol name of the external function */
/*! \brief The external function name */
const ffi::String& global_symbol() const { return global_symbol_; }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExternFuncPatternNode>().def_ro("global_symbol",
&ExternFuncPatternNode::global_symbol_);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExternFuncPattern", ExternFuncPatternNode,
DFPatternNode);
};
/*!
* \brief Managed reference to ExternFuncPatternNode.
* \sa ExternFuncPatternNode
*/
class ExternFuncPattern : public DFPattern {
public:
TVM_DLL ExternFuncPattern(ffi::String global_symbol);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExternFuncPattern, DFPattern, ExternFuncPatternNode);
};
/*! \brief Syntatic Sugar for creating a VarPattern with a name */
VarPattern IsVar(const ffi::String& name);
/*! \brief Syntatic Sugar for creating a ConstantPattern */
ConstantPattern IsConst();
/*! \brief Syntatic Sugar for creating a WildcardPattern */
WildcardPattern Wildcard();
/*! \brief Syntatic Sugar for creating a ExprPattern */
ExprPattern IsExpr(const Expr& expr);
/*! \brief Syntatic Sugar for creating a ExprPattern base on an Op */
ExprPattern IsOp(const ffi::String& op_name);
/*! \brief Syntatic Sugar for call_tir (return a tensor) */
// Todo(relax-team): Dataflow pattern for Type, and match out_ty
CallPattern IsCallTIR(const ffi::String& name, ffi::Optional<TuplePattern> args = std::nullopt);
/*! \brief Syntatic Sugar for call_tir (return a tuple of tensor) */
CallPattern IsCallTIR(const ffi::String& name, TuplePattern var_args);
/*! \brief Syntatic Sugar for call_dps_packed (return a tensor) */
CallPattern IsCallDPSPacked(const ffi::String& name,
ffi::Optional<TuplePattern> args = std::nullopt);
/*! \brief Syntatic Sugar for call_dps_packed (return a tuple of tensor) */
CallPattern IsCallDPSPacked(const ffi::String& name, TuplePattern var_args);
/*! \brief Syntatic Sugar for creating TuplePattern or UnorderedTuplePattern (unordered=true) */
DFPattern IsTuple(const ffi::Array<DFPattern>& fields, bool unordered = false);
/*! \brief Syntatic Sugar for creating a TupleGetItemPattern */
TupleGetItemPattern IsTupleGetItem(const DFPattern tuple, int index = -1);
/*! \brief Implementation of the templated CallPattern syntax sugar */
template <typename... Args>
CallPattern DFPattern::operator()(Args&&... args) const {
return CallPattern(ffi::GetRef<DFPattern>(this->get()),
ffi::Array<DFPattern>({std::forward<Args>(args)...}));
}
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_DATAFLOW_PATTERN_H_
include/tvm/relax/dataflow_pattern_functor.h
+185
View File
@@ -0,0 +1,185 @@
/*
* 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/relax/dataflow_pattern_functor.h
* \brief Functors and visitors for dataflow patterns.
*/
#ifndef TVM_RELAX_DATAFLOW_PATTERN_FUNCTOR_H_
#define TVM_RELAX_DATAFLOW_PATTERN_FUNCTOR_H_
#include <tvm/relax/dataflow_pattern.h>
#include <unordered_set>
#include <utility>
namespace tvm {
namespace relax {
/*!
* \brief A dynamical functor that dispatches on in the first DFPattern argument.
*
* \tparam FType function signature
* This type is only defined for FType with function signature R(const DFPattern&,
* Args...)
*/
template <typename FType>
class DFPatternFunctor;
// functions to be overriden.
#define DFPATTERN_FUNCTOR_DEFAULT \
{ \
return VisitDFPatternDefault_(op, std::forward<Args>(args)...); \
}
#define RELAX_DFPATTERN_FUNCTOR_DISPATCH(OP) \
vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \
return self->VisitDFPattern_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
});
template <typename R, typename... Args>
class DFPatternFunctor<R(const DFPattern& n, Args...)> {
private:
using TSelf = DFPatternFunctor<R(const DFPattern& n, Args...)>;
using FType = tvm::NodeFunctor<R(const ffi::ObjectRef& n, TSelf* self, Args...)>;
public:
/*! \brief virtual destructor */
virtual ~DFPatternFunctor() {}
/*!
* \brief Same as call.
* \param n The expression node.
* \param args Additional arguments.
* \return The result of the call
*/
R operator()(const DFPattern& n, Args... args) {
return VisitDFPattern(n, std::forward<Args>(args)...);
}
/*!
* \brief The functor call.
* \param n The expression node.
* \param args Additional arguments.
* \return The result of the call
*/
virtual R VisitDFPattern(const DFPattern& n, Args... args) {
TVM_FFI_ICHECK(n.defined());
static FType vtable = InitVTable();
return vtable(n, this, std::forward<Args>(args)...);
}
// Functions that can be overriden by subclass
virtual R VisitDFPattern_(const OrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const AndPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const NotPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const AttrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const CallPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const ConstantPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const DataTypePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const ExprPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const FunctionPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const ShapePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const TupleGetItemPatternNode* op,
Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const TuplePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const TypePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const WildcardPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const VarPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const DataflowVarPatternNode* op,
Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const GlobalVarPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const ExternFuncPatternNode* op,
Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const PrimArrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const UnorderedTuplePatternNode* op,
Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPatternDefault_(const ffi::Object* op, Args...) {
TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
throw;
}
private:
// initialize the vtable.
static FType InitVTable() {
FType vtable;
// Set dispatch
RELAX_DFPATTERN_FUNCTOR_DISPATCH(OrPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(AndPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(NotPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(AttrPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(CallPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(ConstantPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataTypePatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExprPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(FunctionPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(ShapePatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(TupleGetItemPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(TuplePatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(TypePatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(WildcardPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(VarPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataflowVarPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(GlobalVarPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExternFuncPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(PrimArrPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(UnorderedTuplePatternNode);
vtable.Finalize();
return vtable;
}
};
/*!
* \brief A simple visitor wrapper around DFPatternFunctor.
* Recursively visit the content.
*
* DFPatternVisitor treats the Pattern as dataflow graph,and only visit each Expr node once.
*/
class DFPatternVisitor : public DFPatternFunctor<void(const DFPattern&)> {
public:
void VisitDFPattern(const DFPattern& pattern) override;
void VisitDFPattern_(const OrPatternNode* op) override;
void VisitDFPattern_(const AndPatternNode* op) override;
void VisitDFPattern_(const NotPatternNode* op) override;
void VisitDFPattern_(const AttrPatternNode* op) override;
void VisitDFPattern_(const CallPatternNode* op) override;
void VisitDFPattern_(const ConstantPatternNode* op) override;
void VisitDFPattern_(const DataTypePatternNode* op) override;
void VisitDFPattern_(const ExprPatternNode* op) override;
void VisitDFPattern_(const FunctionPatternNode* op) override;
void VisitDFPattern_(const ShapePatternNode* op) override;
void VisitDFPattern_(const TupleGetItemPatternNode* op) override;
void VisitDFPattern_(const TuplePatternNode* op) override;
void VisitDFPattern_(const TypePatternNode* op) override;
void VisitDFPattern_(const WildcardPatternNode* op) override;
void VisitDFPattern_(const VarPatternNode* op) override;
void VisitDFPattern_(const DataflowVarPatternNode* op) override;
void VisitDFPattern_(const GlobalVarPatternNode* op) override;
void VisitDFPattern_(const ExternFuncPatternNode* op) override;
void VisitDFPattern_(const PrimArrPatternNode* op) override;
void VisitDFPattern_(const UnorderedTuplePatternNode* op) override;
protected:
// set of already-visited nodes
std::unordered_set<const ffi::Object*> visited_;
};
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_DATAFLOW_PATTERN_FUNCTOR_H_
include/tvm/relax/distributed/axis_group_graph.h
+480
View File
@@ -0,0 +1,480 @@
/*
* 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.
*/
#ifndef TVM_RELAX_DISTRIBUTED_AXIS_GROUP_GRAPH_H_
#define TVM_RELAX_DISTRIBUTED_AXIS_GROUP_GRAPH_H_
#include <tvm/arith/iter_affine_map.h>
#include <tvm/relax/distributed/type.h>
#include <tvm/relax/expr.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
#include <limits>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
namespace tvm {
namespace tirx {
// (var, axis)
using TIRVarAxis = std::pair<Var, int>;
// (buffer, axis)
using BufferAxis = std::pair<Buffer, int>;
class BufferAxisHash {
public:
size_t operator()(const BufferAxis& buffer_axis) const {
size_t const h1(ffi::ObjectPtrHash()(buffer_axis.first));
size_t const h2(std::hash<int>()(buffer_axis.second));
return h1 ^ (h2 << 1);
}
};
/*!
* \brief Suppose we want to shard a buffer along a specific dimension, we need to know how
* to rewrite the access index of the buffer. To make it simple, we only support the case that
* the access can be rewritten by changing the extent of an iter var.
* \param index The access index
* \param var_range The range of each iter var
* \param analyzer The analyzer
* \return The iter var whose extent to be changed
*/
Var GetShardingVarFromIndex(PrimExpr index, ffi::Map<Var, Range> var_range,
const arith::Analyzer& analyzer);
/*!
* \brief Construct an axis group graph from a PrimFunc. Two buffer axis are connected if they
* are accessed by the same index.
*/
class BufferAxisGraphExtractor : public StmtExprVisitor {
public:
static std::vector<std::vector<TIRVarAxis>> GetTIRVarAxisGraph(const PrimFunc& prim_func) {
BufferAxisGraphExtractor extractor;
extractor(prim_func->body);
ffi::Map<Buffer, Var> inverse_buffer_map;
for (const auto& pr : prim_func->buffer_map) {
inverse_buffer_map.Set(pr.second, pr.first);
}
std::vector<std::vector<TIRVarAxis>> tir_var_axis_group_list;
std::unordered_set<BufferAxis, BufferAxisHash> visited;
for (const auto& pr : prim_func->buffer_map) {
Var param = pr.first;
Buffer buffer = pr.second;
for (int i = 0; i < static_cast<int>(buffer->shape.size()); i++) {
if (extractor.buffer_axis_graph_.count({buffer, i})) {
std::vector<BufferAxis> buffer_axis_group;
extractor.DFSGraph({buffer, i}, &visited, &buffer_axis_group);
if (buffer_axis_group.size() <= 1) {
continue;
}
std::vector<TIRVarAxis> tir_var_axis_group;
for (const auto& buffer_axis : buffer_axis_group) {
if (!inverse_buffer_map.count(buffer_axis.first)) {
continue;
}
tir_var_axis_group.push_back(
{inverse_buffer_map[buffer_axis.first], buffer_axis.second});
}
tir_var_axis_group_list.push_back(tir_var_axis_group);
}
}
}
return tir_var_axis_group_list;
}
void DFSGraph(BufferAxis cur, std::unordered_set<BufferAxis, BufferAxisHash>* visited,
std::vector<BufferAxis>* buffer_axis_group) {
if (visited->count(cur)) {
return;
}
visited->insert(cur);
buffer_axis_group->push_back(cur);
for (const auto& next : buffer_axis_graph_[cur]) {
DFSGraph(next, visited, buffer_axis_group);
}
}
private:
void VisitStmt_(const BufferStoreNode* op) final {
StmtExprVisitor::VisitStmt_(op);
buffer_access_indices_.push_back({op->buffer, op->indices});
}
void VisitExpr_(const BufferLoadNode* op) final {
StmtExprVisitor::VisitExpr_(op);
buffer_access_indices_.push_back({op->buffer, op->indices});
}
bool Match(PrimExpr a, PrimExpr buffer_shape_a, PrimExpr b, PrimExpr buffer_shape_b,
const arith::Analyzer& analyzer) {
if (b.as<VarNode>()) {
std::swap(a, b);
std::swap(buffer_shape_a, buffer_shape_b);
}
if (!a.as<VarNode>()) {
return false;
}
Var var = a.as_or_throw<Var>();
analyzer->Bind(iter_var_range_);
b = analyzer->Simplify(b);
// index var `a` must access whole range of a specific buffer dimension
arith::IntSet intset_b = arith::EvalSet(b, arith::AsIntSet(iter_var_range_));
if (!analyzer->CanProveEqual(buffer_shape_a, iter_var_range_[var]->extent) ||
!intset_b.MatchRange(Range::FromMinExtent(0, buffer_shape_b))) {
return false;
}
Var matched_var = GetShardingVarFromIndex(b, iter_var_range_, analyzer);
if (!matched_var.same_as(var)) {
return false;
}
return true;
}
void VisitStmt_(const SBlockNode* op) final {
if (op->name_hint == "root") {
StmtExprVisitor::VisitStmt_(op);
return;
}
buffer_access_indices_.clear();
StmtExprVisitor::VisitStmt_(op);
iter_var_range_.clear();
for (const auto& iter_var : op->iter_vars) {
iter_var_range_.Set(iter_var->var, iter_var->dom);
}
arith::Analyzer analyzer;
for (const auto& access_pr : buffer_access_indices_) {
Buffer buffer = access_pr.first;
ffi::Array<PrimExpr> indices = access_pr.second;
for (int i = 0; i < static_cast<int>(indices.size()); i++) {
for (const auto& another_access_pr : buffer_access_indices_) {
if (another_access_pr.first.same_as(buffer)) {
continue;
}
Buffer another_buffer = another_access_pr.first;
ffi::Array<PrimExpr> another_indices = another_access_pr.second;
for (int j = 0; j < static_cast<int>(another_indices.size()); j++) {
if (Match(indices[i], buffer->shape[i], another_indices[j], another_buffer->shape[j],
analyzer)) {
JoinBufferAxis({buffer, i}, {another_buffer, j});
}
}
}
}
}
}
void JoinBufferAxis(BufferAxis axis1, BufferAxis axis2) {
if (!buffer_axis_graph_.count(axis1)) {
buffer_axis_graph_[axis1] = {};
}
if (!buffer_axis_graph_.count(axis2)) {
buffer_axis_graph_[axis2] = {};
}
buffer_axis_graph_[axis1].push_back(axis2);
buffer_axis_graph_[axis2].push_back(axis1);
}
std::vector<std::pair<Buffer, ffi::Array<PrimExpr>>> buffer_access_indices_;
std::unordered_map<BufferAxis, std::vector<BufferAxis>, BufferAxisHash> buffer_axis_graph_;
ffi::Map<Var, Range> iter_var_range_;
std::string func_name;
};
} // namespace tirx
} // namespace tvm
namespace tvm {
namespace relax {
namespace distributed {
/*! \brief tensor axis*/
struct Axis {
const ExprNode* tensor;
int dim = 0;
int tuple_index = 0;
Axis(const ExprNode* tensor, int dim, int tuple_index = 0)
: tensor(tensor), dim(dim), tuple_index(tuple_index) {
TVM_FFI_ICHECK(tensor->IsInstance<ConstantNode>() || tensor->IsInstance<VarNode>());
}
bool operator==(const Axis& other) const {
return tensor == other.tensor && dim == other.dim && tuple_index == other.tuple_index;
}
};
class AxisHash {
public:
size_t operator()(const Axis& axis) const {
size_t const h1(std::hash<const ExprNode*>()(axis.tensor));
size_t const h2(std::hash<int>()(axis.dim));
size_t const h3(std::hash<int>()(axis.tuple_index));
return h1 ^ (h2 << 1) ^ (h3 << 2);
}
};
using AxisGroup = std::unordered_set<Axis, AxisHash>;
class AxisGroupHash {
public:
size_t operator()(const AxisGroup& axis_group) const {
size_t seed = 0;
for (auto axis : axis_group) {
seed ^= AxisHash()(axis) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
}
return seed;
}
};
using ShardingSpec = std::pair<DeviceMesh, Placement>;
// device mesh and the device mesh axis that the tensor axis maps to
using AxisShardingSpec = std::pair<DeviceMesh, int>;
class AxisShardingSpecEqual {
public:
bool operator()(const AxisShardingSpec& lhs, const AxisShardingSpec& rhs) const {
return ffi::StructuralEqual()(lhs.first, rhs.first) && lhs.second == rhs.second;
}
};
class AxisShardingSpecHash {
public:
size_t operator()(const AxisShardingSpec& sharding_spec) const {
size_t seed = 0;
seed ^= ffi::StructuralHash()(sharding_spec.first);
seed ^= std::hash<int>()(sharding_spec.second) << 1;
return seed;
}
};
/*!
* \brief A graph whose nodes are tensor axes, and the edge means some information can be propagated
* through the two axes. Although it only does sharding propagation, this data structure can be
* extended to perform all kinds of propagation that happens on tensor axes.
*/
class AxisGroupGraph {
public:
enum class EdgeType { kAscend, kDescend, kSimbling };
private:
static EdgeType ReverseEdgeType(EdgeType type) {
switch (type) {
case EdgeType::kAscend:
return EdgeType::kDescend;
case EdgeType::kDescend:
return EdgeType::kAscend;
case EdgeType::kSimbling:
return EdgeType::kSimbling;
}
TVM_FFI_THROW(InternalError) << "Unreachable code";
throw;
}
static int GetEdgePriority(EdgeType type) {
switch (type) {
case EdgeType::kAscend:
return 0;
case EdgeType::kDescend:
return 2;
case EdgeType::kSimbling:
return 1;
}
TVM_FFI_THROW(InternalError) << "Unreachable code";
throw;
}
struct AxisGraphEdge {
Axis src;
Axis dst;
// the producer-consumer relationship between src tensor and dst tensor
// kAscend means consumer->producer
// kDescend means producer->consumer
// kSimbling means other cases
EdgeType type;
bool operator==(const AxisGraphEdge& other) const {
return src == other.src && dst == other.dst && type == other.type;
}
};
struct Path {
int direction = 0;
Path AddEdge(EdgeType type) { return {direction |= (1 << GetEdgePriority(type))}; }
int GetPriority() const {
switch (direction) {
case 1: // ascend only
return 0;
case 4: // descend only
return 2;
case 0: // empty path (source node)
return 3; // source node must have max priority
default: // mixed path
return 1;
}
}
};
public:
AxisGroupGraph() = default;
/*!
* \brief add edge between two axes
* \param axis1 The src axis
* \param axis2 The dst axis
* \param type The producer-consumer relationship between src tensor and dst tensor
* kAscend means consumer->producer
* kDescend means producer->consumer
* kSimbling means other cases
*/
void JoinAxis(Axis axis1, Axis axis2, EdgeType type) {
AddEdge(axis1, axis2, type);
AddEdge(axis2, axis1, ReverseEdgeType(type));
}
/*!
* \brief add a source shardingspec to propagate
* \param axis The source axis
* \param spec The axis's sharding spec
*/
void AddSrcShardingPoint(Axis axis, AxisShardingSpec spec) {
src_axis_sharding_spec_[axis] = spec;
}
/*!
* \brief propagate sharding specs from source axes
*/
void PropagateShardingSpec() {
axis_sharding_specs_priority_.clear();
for (const auto& pr : src_axis_sharding_spec_) {
std::unordered_set<Axis, AxisHash> visited;
PropagateShardingSpec(pr.first, pr.second, Path(), &visited);
}
ChooseAxisShardingSpec();
}
/*!
* \brief add a cut point that stops the propagation of a certain sharding spec
*
* \param axis The cut point
* \param spec The spec to stop propagation
*/
void AddPropagationCutPoint(Axis axis, AxisShardingSpec spec) {
cutpoint_axis_sharding_spec_[axis] = spec;
}
/*!
* \brief Get the Sharding Spec of an axis after propagation
*
* \param axis the specified axis
* \return if a sharding spec is found, return (axis_sharding_spec, true)
* otherwise, return (null axis_sharding_spec, false)
*/
std::tuple<AxisShardingSpec, bool> GetAxisShardingSpec(Axis axis) {
if (axis_sharding_specs_priority_.count(axis)) {
return {axis_sharding_specs_priority_[axis].begin()->first, true};
} else {
return {{DeviceMesh(), -1}, false};
}
}
private:
void AddEdge(Axis src, Axis dst, EdgeType type) {
if (!graph_.count(src)) {
graph_[src] = {};
}
graph_[src].push_back({src, dst, type});
}
void PropagateShardingSpec(Axis axis, AxisShardingSpec spec, Path path,
std::unordered_set<Axis, AxisHash>* visited) {
if (cutpoint_axis_sharding_spec_.count(axis) ||
(src_axis_sharding_spec_.count(axis) &&
!AxisShardingSpecEqual()(src_axis_sharding_spec_[axis], spec)) ||
visited->count(axis)) {
return;
}
visited->insert(axis);
if (!axis_sharding_specs_priority_.count(axis)) {
axis_sharding_specs_priority_[axis] = {};
}
axis_sharding_specs_priority_[axis][spec] = path.GetPriority();
for (auto edge : graph_[axis]) {
PropagateShardingSpec(edge.dst, spec, path.AddEdge(edge.type), visited);
}
}
void ChooseAxisShardingSpec() {
for (auto& pr : axis_sharding_specs_priority_) {
auto& axis = pr.first;
auto& specs = pr.second;
int max_priority = std::numeric_limits<int>::min();
for (auto& pr2 : specs) {
max_priority = std::max(max_priority, pr2.second);
}
for (auto it = specs.begin(); it != specs.end();) {
if (it->second != max_priority) {
it = specs.erase(it);
} else {
it++;
}
}
TVM_FFI_ICHECK(specs.size() == 1)
<< "multiple possible sharding for axis: (" << ffi::GetRef<Expr>(axis.tensor) << ", "
<< axis.dim << ")";
}
}
// union set
std::unordered_map<Axis, std::vector<AxisGraphEdge>, AxisHash> graph_;
std::unordered_map<Axis, AxisShardingSpec, AxisHash> src_axis_sharding_spec_;
std::unordered_map<Axis, AxisShardingSpec, AxisHash> cutpoint_axis_sharding_spec_;
std::unordered_map<
Axis, std::unordered_map<AxisShardingSpec, int, AxisShardingSpecHash, AxisShardingSpecEqual>,
AxisHash>
axis_sharding_specs_priority_;
};
using FBuildAxisGraph = std::function<void(const Var& output_var, const Call& call,
distributed::AxisGroupGraph* axis_group_graph)>;
void BuildAxisGraphUnary(const Var& output_var, const Call& call,
distributed::AxisGroupGraph* axis_group_graph);
void BuildAxisGraphBinary(const Var& output_var, const Call& call,
distributed::AxisGroupGraph* axis_group_graph);
void BuildAxisGraphReduce(const Var& output_var, const Call& call,
distributed::AxisGroupGraph* axis_group_graph);
void BuildAxisGraphMatmul(const Var& output_var, const Call& call,
distributed::AxisGroupGraph* axis_group_graph);
void BuildAxisGraphPermuteDims(const Var& output_var, const Call& call,
distributed::AxisGroupGraph* axis_group_graph);
void BuildAxisGraphReshape(const Var& output_var, const Call& call,
distributed::AxisGroupGraph* axis_group_graph);
void BuildAxisGraphCallTIR(const Var& output_var, const Call& call, const tirx::PrimFunc& func,
distributed::AxisGroupGraph* axis_group_graph);
} // namespace distributed
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_DISTRIBUTED_AXIS_GROUP_GRAPH_H_
include/tvm/relax/distributed/global_info.h
+74
View File
@@ -0,0 +1,74 @@
/*
* 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/relax/distributed/global_info.h
* \brief Data structure for distributed inference
*/
#ifndef TVM_RELAX_DISTRIBUTED_GLOBAL_INFO_H_
#define TVM_RELAX_DISTRIBUTED_GLOBAL_INFO_H_
#include <tvm/ffi/container/shape.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
namespace tvm {
namespace relax {
namespace distributed {
/*
* \brief Device mesh express a view of topology of devices, represented by an n-d matrix of
* device ids
*/
class DeviceMeshNode : public GlobalInfoNode {
public:
/*! \brief logical shape of the mesh*/
ffi::Shape shape;
/*! \brief device ids in the mesh*/
ffi::Array<int64_t> device_ids;
/*! \brief Optionally use range to represent device_ids*/
ffi::Optional<Range> device_range;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DeviceMeshNode>()
.def_ro("shape", &DeviceMeshNode::shape)
.def_ro("device_ids", &DeviceMeshNode::device_ids)
.def_ro("device_range", &DeviceMeshNode::device_range);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.DeviceMesh", DeviceMeshNode, GlobalInfoNode);
};
/*!
* \brief Managed reference to a DeviceMesh.
* \sa DeviceMeshNode
*/
class DeviceMesh : public GlobalInfo {
public:
TVM_DLL DeviceMesh(ffi::Shape shape, ffi::Array<int64_t> device_ids);
TVM_DLL DeviceMesh(ffi::Shape shape, Range device_range);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DeviceMesh, GlobalInfo, DeviceMeshNode);
};
} // namespace distributed
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_DISTRIBUTED_GLOBAL_INFO_H_
include/tvm/relax/distributed/transform.h
+78
View File
@@ -0,0 +1,78 @@
/*
* 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/relax/distributed/transform.h
* \brief Relax distributed specific transformation passes.
*/
#ifndef TVM_RELAX_DISTRIBUTED_TRANSFORM_H_
#define TVM_RELAX_DISTRIBUTED_TRANSFORM_H_
#include <tvm/ir/transform.h>
#include <tvm/relax/dataflow_pattern.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/transform.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/index_map.h>
namespace tvm {
namespace relax {
namespace distributed {
namespace transform {
using Pass = tvm::transform::Pass;
using PassInfo = tvm::transform::PassInfo;
using PassContext = tvm::transform::PassContext;
using Function = tvm::relax::Function;
using DataflowBlock = tvm::relax::DataflowBlock;
using tvm::transform::CreateModulePass;
/*!
* \brief Propagate sharding information.
*
* \return The Pass.
*/
TVM_DLL Pass PropagateSharding();
/*!
* \brief Lower global view TensorIR into local view.
*
* \return The Pass.
*/
TVM_DLL Pass LowerGlobalViewToLocalView();
/*!
* \brief Legalize redistribute op to ccl op.
*
* \return The Pass.
*/
TVM_DLL Pass LegalizeRedistribute();
/*!
* \brief Lower DistIR to Relax
*
* \return The Pass.
*/
TVM_DLL Pass LowerDistIR();
} // namespace transform
} // namespace distributed
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_DISTRIBUTED_TRANSFORM_H_
include/tvm/relax/distributed/type.h
+175
View File
@@ -0,0 +1,175 @@
/*
* 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/relax/distributed/type.h
* \brief Type definitions for DTensor (Distributed Tensor)
*/
#ifndef TVM_RELAX_DISTRIBUTED_TYPE_H_
#define TVM_RELAX_DISTRIBUTED_TYPE_H_
#include <tvm/relax/distributed/global_info.h>
#include <tvm/relax/type.h>
#include <utility>
namespace tvm {
namespace relax {
namespace distributed {
enum class PlacementSpecKind : int { kSharding = 0, kReplica = 1 };
/*! \brief Describes how data is distributed in one dimension of the device mesh*/
class PlacementSpecNode : public ffi::Object {
public:
/*! \brief If the kind is sharding, this value represents the tensor dimension to shard.
* otherwise, axis is -1.
*/
int axis;
/*! \brief The kind of placement spec. Possible values: kSharding and kReplica. */
PlacementSpecKind kind;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PlacementSpecNode>()
.def_ro("axis", &PlacementSpecNode::axis)
.def_ro("kind", &PlacementSpecNode::kind);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("relax.distributed.PlacementSpec", PlacementSpecNode, ffi::Object);
};
/*!
* \brief Managed reference to PlacementSpecNode.
* \sa PlacementSpecNode
*/
class PlacementSpec : public ffi::ObjectRef {
public:
TVM_DLL static PlacementSpec Sharding(int axis);
TVM_DLL static PlacementSpec Replica();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PlacementSpec, ffi::ObjectRef, PlacementSpecNode);
};
class ShardingNode : public PlacementSpecNode {
public:
/*! \brief The dimension of tensor we shard*/
int64_t sharding_dim;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ShardingNode>().def_ro("sharding_dim", &ShardingNode::sharding_dim);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.Sharding", ShardingNode, PlacementSpecNode);
};
/*! \brief Describes how data is distributed in each dimension of the device mesh*/
class PlacementNode : public ffi::Object {
public:
/*! \brief specs for each dim of device mesh.*/
ffi::Array<PlacementSpec> dim_specs;
ffi::String ToString() const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PlacementNode>().def_ro("dim_specs", &PlacementNode::dim_specs);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.Placement", PlacementNode, ffi::Object);
};
/*!
* \brief Managed reference to a Placement.
* \sa PlacementNode
*/
class Placement : public ffi::ObjectRef {
public:
TVM_DLL explicit Placement(ffi::Array<PlacementSpec> dim_specs);
/*! \brief replica dim is printed as "R" and sharding dim is printed as "S[i]".]*/
static Placement FromText(ffi::String text_repr);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Placement, ffi::ObjectRef, PlacementNode);
};
/*!
* \brief Type of DTensor (Distributed Tensor).
*/
class DTensorTypeNode : public TypeNode {
public:
explicit DTensorTypeNode(ffi::UnsafeInit)
: tensor_ty(ffi::UnsafeInit{}), device_mesh(), placement() {}
DTensorTypeNode(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement)
: tensor_ty(std::move(tensor_ty)),
device_mesh(std::move(device_mesh)),
placement(std::move(placement)) {}
/*!
* \brief The tensor type carried by the DTensor type.
*/
TensorType tensor_ty;
/*!
* \brief The device mesh of the tensor.
*/
DeviceMesh device_mesh;
/*!
* \brief The placement of the tensor among the device mesh.
*/
Placement placement;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DTensorTypeNode>()
.def_ro("device_mesh", &DTensorTypeNode::device_mesh)
.def_ro("placement", &DTensorTypeNode::placement)
.def_ro("tensor_ty", &DTensorTypeNode::tensor_ty);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DTensorType", DTensorTypeNode, TypeNode);
};
/*!
* \brief Managed reference to DTensorTypeNode.
* \sa DTensorTypeNode
*/
class DTensorType : public Type {
public:
/*!
* \brief Construction with device mesh and placement.
* \param tensor_ty The tensor type carried by the DTensor type.
* \param device_mesh The device mesh of the tensor.
* \param placement The placement of the tensor among the device mesh.
* \param span The span of the AST.
*/
TVM_DLL DTensorType(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement,
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DTensorType, Type, DTensorTypeNode);
};
} // namespace distributed
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_DISTRIBUTED_TYPE_H_
include/tvm/relax/exec_builder.h
+194
View File
@@ -0,0 +1,194 @@
/*
* 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/relax/exec_builder.h
*/
#ifndef TVM_RELAX_EXEC_BUILDER_H_
#define TVM_RELAX_EXEC_BUILDER_H_
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/runtime/vm/bytecode.h>
#include <tvm/runtime/vm/executable.h>
#include <string>
#include <unordered_map>
#include <vector>
namespace tvm {
namespace relax {
namespace vm = tvm::runtime::vm;
class ExecBuilder;
/*!
* \brief A builder provides api to build VM executable with instructions.
*/
class ExecBuilderNode : public ffi::Object {
public:
/*!
* \brief Declare a function, it is OK to have multiple declarations.
* \param func The function name.
* \param kind The kind of the function.
*/
void DeclareFunction(const std::string& func, vm::VMFuncInfo::FuncKind kind);
/*!
* \brief To annotate the start of a vm function.
* \param func The function name.
* \param num_inputs The number of inputs.
* \param param_names The function parameter names.
* \param kind The kind of the function.
* \param init_register_size Initial setting of register file size.
*/
void EmitFunction(const std::string& func, int64_t num_inputs,
ffi::Optional<ffi::Array<ffi::String>> param_names,
vm::VMFuncInfo::FuncKind kind = vm::VMFuncInfo::FuncKind::kVMFunc,
int64_t init_register_size = 0);
/*!
* \brief Annotate the end of a vm function.
* \param func The function name.
*/
void EndFunction(const std::string& func);
/*!
* \brief Emit a call instruction for a packed function.
* \param func The packed function name.
* \param args The arguments of the function.
* \param ret The return register.
*/
void EmitCall(const std::string& func, std::vector<vm::Instruction::Arg> args, vm::RegName ret);
/*!
* \brief Emit a call instruction with func as argument.
* \param func The packed function index.
* \param args The arguments of the function.
* \param ret The return register.
*/
void EmitCall(vm::Instruction::Arg func, std::vector<vm::Instruction::Arg> args, vm::RegName ret);
/*!
* \brief Emit a ret instruction.
* \param result The return result.
* \note result must be a register.
*/
void EmitRet(vm::Instruction::Arg result);
/*!
* \brief Emit a goto instruction.
* \param pc_offset The program counter offset as the jump offset.
*/
void EmitGoto(vm::Index pc_offset);
/*!
* \brief Emit an If instruction.
* \param cond The register containing the cond value.
* \param false_offset The program counter offset for the false branch.
* \note result must be a register.
*/
void EmitIf(vm::Instruction::Arg cond, vm::Index false_offset);
/*!
* \brief Get function index by its name.
* \param name The name of the function.
* \return The argument corresponding to the function index.
*/
vm::Instruction::Arg GetFunction(const std::string& name);
/*!
* \brief Convert a constant value something that exec builder can understand.
*
* This function may update the constant pool to include the obj value.
*
* \param value The input constant value
* \return An Arg that represents the result of constant argument.
*/
template <typename T>
vm::Instruction::Arg ConvertConstant(T value) {
ffi::Any rv;
rv = value;
return ConvertConstant_(rv);
}
/*!
* \brief update memory scopes.
*
* This function builds the memory scopes for constants.
*
* \param idx Index of the constant
* \param scope The memory scope.
*/
void SaveMemoryScope(vm::Instruction::Arg idx, ffi::String scope);
/*!
* \brief Raw access to underlying executable build in progress.
*/
vm::VMExecutable* exec() const;
/*!
* \brief Finalize the build, run formalize and get the final result.
* \note This function should not be called during construction.
*/
ffi::ObjectPtr<vm::VMExecutable> Get();
/*!
* \brief Create an ExecBuilder.
* \return The ExecBuilder.
*/
TVM_DLL static ExecBuilder Create();
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExecBuilderNode>();
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ExecBuilder", ExecBuilderNode, ffi::Object);
private:
/*!
* \brief Convert a constant value something that exec builder can understand.
*
* This function may update the constant pool to include the obj value.
*
* \param obj The constant value to be emitted
* \return An Arg that represents the result of constant argument.
*/
vm::Instruction::Arg ConvertConstant_(ffi::Any obj);
/*!
* \brief A helper function to check if an executable is legal by checking if registers are used
* properly
*/
void CheckExecutable();
/*!
* \brief Formalize the executable.
*/
void Formalize();
/*! \brief The mutable internal executable. */
ffi::ObjectPtr<vm::VMExecutable> exec_; // mutable
/*! \brief internal dedup map when creating index for a new constant */
std::unordered_map<ffi::Any, vm::Index, ffi::StructuralHash, ffi::StructuralEqual>
const_dedup_map_;
};
class ExecBuilder : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExecBuilder, ffi::ObjectRef, ExecBuilderNode);
};
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_EXEC_BUILDER_H_
include/tvm/relax/expr.h
+701
View File
@@ -0,0 +1,701 @@
/*
* 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.
*/
#ifndef TVM_RELAX_EXPR_H_
#define TVM_RELAX_EXPR_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
#include <tvm/ir/source_map.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/tensor.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <functional>
namespace tvm {
namespace relax {
/*! \brief Tuple container */
class TupleNode : public ExprNode {
public:
/*! \brief the fields of the tuple */
tvm::ffi::Array<Expr> fields;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TupleNode>().def_ro("fields", &TupleNode::fields);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.Tuple", TupleNode, ExprNode);
};
class Tuple : public Expr {
public:
/*!
* \brief The constructor
* \param fields The fields of a tuple.
* \param span The source span of the expression.
*/
TVM_DLL explicit Tuple(tvm::ffi::Array<Expr> fields, Span span = Span());
/*!
* \brief Utility constructor to handle conversion to relax::Expr
*
* If the calling scope already has an array of a specific type of
* relax expression (e.g. `ffi::Array<relax::Var>`), it must be converted
* into an array of base type. This constructor handles the
* conversion to the base `ffi::Array<relax::Expr>`.
*
* \tparam ExprType The type of relax expression passed in as an argument.
*
* \param fields The fields of a tuple.
*
* \param span The source span of the expression.
*/
template <typename ExprType, typename = std::enable_if_t<std::is_base_of_v<Expr, ExprType>>>
TVM_DLL explicit Tuple(tvm::ffi::Array<ExprType> fields, Span span = Span())
: Tuple(fields.Map([](const ExprType& expr) -> Expr { return expr; }), span) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tuple, Expr, TupleNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(TupleNode);
};
/*! \brief Get index-th field out of a tuple. */
class TupleGetItemNode : public ExprNode {
public:
/*! \brief The tuple Expression */
Expr tuple;
/*! \brief which value to get */
int index;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TupleGetItemNode>()
.def_ro("tuple_value", &TupleGetItemNode::tuple)
.def_ro("index", &TupleGetItemNode::index);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.TupleGetItem", TupleGetItemNode, ExprNode);
};
class TupleGetItem : public Expr {
public:
/*!
* \brief The constructor
* \param tuple The tuple to get an element from.
* \param index The index for extracting a value in the tuple.
* \param span The source span of the expression.
*/
TVM_DLL TupleGetItem(Expr tuple, int index, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TupleGetItem, Expr, TupleGetItemNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(TupleGetItemNode);
};
/*! \brief A shape expression which allows users to construct a shape containing PrimExpr.
*/
class ShapeExprNode : public ExprNode {
public:
/*! The values of the shape expression. */
ffi::Array<PrimExpr> values;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ShapeExprNode>().def_ro("values", &ShapeExprNode::values);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ShapeExpr", ShapeExprNode, ExprNode);
};
class ShapeExpr : public Expr {
public:
TVM_DLL explicit ShapeExpr(ffi::Array<PrimExpr> values, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ShapeExpr, Expr, ShapeExprNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ShapeExprNode);
};
/*! \brief The variable class for all Relax bindings. */
class VarNode : public ExprNode {
public:
/*!
* \brief The hint to the variable name.
* \note Each variable is uniquely identified by its address.
*/
ffi::String name_hint;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<VarNode>().def_ro("name_hint", &VarNode::name_hint,
refl::AttachFieldFlag::SEqHashIgnore());
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
static constexpr const uint32_t _type_child_slots = 1;
TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.Var", VarNode, ExprNode);
};
class Var : public Expr {
public:
TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional<Type> ty_annotation,
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, Expr, VarNode);
};
/*! \brief A sub-type of the variable node used to mark dataflow variables from
* normal visible "function local" bindings.
*/
class DataflowVarNode : public VarNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DataflowVarNode>();
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowVar", DataflowVarNode, VarNode);
};
class DataflowVar : public Var {
public:
TVM_DLL explicit DataflowVar(ffi::String name_hint, ffi::Optional<Type> ty_annotation,
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowVar, Var, DataflowVarNode);
};
/*!
* \brief Constant tensor.
*
* \note Scalar constants are represented by ndim-0 constant tensors.
*/
class ConstantNode : public ExprNode {
public:
/*! \brief The data of the tensor */
runtime::Tensor data;
/*! \return The corresponding tensor type of the data */
TensorType tensor_type() const;
/*! \return Whether it is scalar(ndim-0 tensor) */
bool is_scalar() const { return data->ndim == 0; }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ConstantNode>().def_ro("data", &ConstantNode::data);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.Constant", ConstantNode, ExprNode);
};
class Constant : public Expr {
public:
/*!
* \brief The constructor
* \param data The data of the constant tensor.
* \param ty_annotation The type of the constant tensor.
* If not specified, infer it from data.
* \param span The source span of the expression.
*/
TVM_DLL explicit Constant(runtime::Tensor data, ffi::Optional<Type> ty_annotation = std::nullopt,
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Constant, Expr, ConstantNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ConstantNode);
};
/*!
* \brief Represent a string literal constant.
*/
class StringImmNode : public ExprNode {
public:
/*! \brief The data value. */
ffi::String value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<StringImmNode>().def_ro("value", &StringImmNode::value);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.StringImm", StringImmNode, ExprNode);
};
/*!
* \brief Managed reference to StringImm
* \sa StringImmNode
*/
class StringImm : public Expr {
public:
/*!
* \brief The constructor
* \param value The value input.
* \param span The source span of the expression.
*/
TVM_DLL explicit StringImm(ffi::String value, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StringImm, Expr, StringImmNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(StringImmNode);
};
/*!
* \brief Represent a data type constant.
*/
class DataTypeImmNode : public ExprNode {
public:
/*! \brief The data value. */
DLDataType value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DataTypeImmNode>().def_ro("value", &DataTypeImmNode::value);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataTypeImm", DataTypeImmNode, ExprNode);
};
/*!
* \brief Managed reference to DataTypeImm
* \sa DataTypeImmNode
*/
class DataTypeImm : public Expr {
public:
/*!
* \brief The constructor
* \param value The value input.
* \param span The source span of the expression.
*/
TVM_DLL explicit DataTypeImm(DLDataType value, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataTypeImm, Expr, DataTypeImmNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(DataTypeImmNode);
};
/*! \brief The base class of a variable binding in Relax. */
class BindingNode : public ffi::Object {
public:
mutable Span span;
/*! \brief The return variable to bound to. */
Var var;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BindingNode>()
.def_ro("span", &BindingNode::span, refl::AttachFieldFlag::SEqHashIgnore())
// TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
.def_ro("var", &BindingNode::var, refl::AttachFieldFlag::SEqHashDefRecursive());
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.Binding", BindingNode, ffi::Object);
};
class Binding : public ffi::ObjectRef {
protected:
Binding() = default;
public:
explicit Binding(ffi::ObjectPtr<BindingNode> n) : ffi::ObjectRef(n) {}
explicit Binding(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
Binding(const Binding&) = default;
Binding(Binding&&) = default;
Binding& operator=(const Binding&) = default;
Binding& operator=(Binding&&) = default;
const BindingNode* operator->() const { return static_cast<const BindingNode*>(data_.get()); }
const BindingNode* get() const { return operator->(); }
using ContainerType = BindingNode;
};
/*!
* \brief Runtime-match the value to the type.
*
* This operation does runtime check, populates the un-defined symbolic shape vars
* and vars in ty in first occurance, and insert equality assertions in
* other cases.
*/
class MatchCastNode : public BindingNode {
public:
/*! \brief The input value to match cast. */
Expr value;
/*! \brief The type pattern to match to. */
Type ty = Type::Missing();
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MatchCastNode>()
.def_ro("value", &MatchCastNode::value)
// TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
.def_ro("ty", &MatchCastNode::ty, refl::AttachFieldFlag::SEqHashDefRecursive());
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.MatchCast", MatchCastNode, BindingNode);
};
/*!
* \brief Managed reference to MatchCastNode.
* \sa MatchCastNode
*/
class MatchCast : public Binding {
public:
TVM_DLL explicit MatchCast(Var var, Expr value, Type ty, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchCast, Binding, MatchCastNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchCastNode);
};
class VarBindingNode : public BindingNode {
public:
/*! \brief The binding value. */
Expr value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<VarBindingNode>().def_ro("value", &VarBindingNode::value);
// customize the SEqual and SHash methods for better error messages
refl::TypeAttrDef<VarBindingNode>()
.def("__s_equal__", &VarBindingNode::SEqual)
.def("__s_hash__", &VarBindingNode::SHash);
}
bool SEqual(const VarBindingNode* other,
ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.VarBinding", VarBindingNode, BindingNode);
};
class VarBinding : public Binding {
public:
TVM_DLL explicit VarBinding(Var var, Expr value, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VarBinding, Binding, VarBindingNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(VarBindingNode);
};
class BindingBlockNode : public ffi::Object {
public:
ffi::Array<Binding> bindings;
mutable Span span;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BindingBlockNode>()
.def_ro("bindings", &BindingBlockNode::bindings)
.def_ro("span", &BindingBlockNode::span, refl::AttachFieldFlag::SEqHashIgnore(),
refl::DefaultValue(Span()));
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.BindingBlock", BindingBlockNode, ffi::Object);
};
class BindingBlock : public ffi::ObjectRef {
public:
TVM_DLL explicit BindingBlock(ffi::Array<Binding> bindings, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BindingBlock, ffi::ObjectRef, BindingBlockNode);
BindingBlockNode* CopyOnWrite();
};
class DataflowBlockNode : public BindingBlockNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DataflowBlockNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowBlock", DataflowBlockNode,
BindingBlockNode);
};
class DataflowBlock : public BindingBlock {
public:
TVM_DLL explicit DataflowBlock(ffi::Array<Binding> bindings, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowBlock, BindingBlock, DataflowBlockNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(DataflowBlockNode);
};
/*! \brief A sequence of blocks followed by an expression.
*
* The order of blocks enforces scoping and ordering.
*/
class SeqExprNode : public ExprNode {
public:
ffi::Array<BindingBlock> blocks;
Expr body;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SeqExprNode>()
.def_ro("blocks", &SeqExprNode::blocks)
.def_ro("body", &SeqExprNode::body);
refl::TypeAttrDef<SeqExprNode>()
.def("__s_equal__", &SeqExprNode::SEqual)
.def("__s_hash__", &SeqExprNode::SHash);
}
bool SEqual(const SeqExprNode* other,
ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const {
// Establish mappings for symbolic variables defined by bindings before
// comparing their uses in the SeqExpr result type and body.
return equal(blocks, other->blocks, false, "blocks") && equal(ty, other->ty, false, "ty") &&
equal(body, other->body, false, "body");
}
int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const {
int64_t hash_value = init_hash;
hash_value = hash(blocks, hash_value, false);
hash_value = hash(ty, hash_value, false);
hash_value = hash(body, hash_value, false);
return hash_value;
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.SeqExpr", SeqExprNode, ExprNode);
};
class SeqExpr : public Expr {
public:
/* \brief Implicit conversion constructor
*
* Relax nodes that introduce a new scope (e.g. `relax::Function`)
* are required to be held as SeqExpr. This implicit conversion
* provides allows callsites to use these member variables when the
* C++ compile-time type is a `relax::Expr`. For example,
* a transform may use `func.CopyOnWrite()->body = expr;`.
*
* If the expression is already a `relax::SeqExpr`, the same
* underlying `relax::SeqExprNode` is used, and no copies are made.
*/
TVM_DLL SeqExpr(Expr body); // NOLINT(*)
TVM_DLL explicit SeqExpr(ffi::Array<BindingBlock> blocks, Expr body, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SeqExpr, Expr, SeqExprNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(SeqExprNode);
};
/*!
* \brief Condition expression
*
* Unlike traditional statement `if`s, the if evalutes
* to the result of the branch taken.
*
* x = if (true) { 1 } else { 0 }; // x is 1
* y = if (false) { 1 } else { 0 }; // y is 0
*
* \note This is similar to C's ternary operator.
*/
class IfNode : public ExprNode {
public:
/*! \brief The condition. */
Expr cond;
/*! \brief The expression evaluated when condition is true. */
SeqExpr true_branch;
/*! \brief The expression evaluated when condition is false */
SeqExpr false_branch;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IfNode>()
.def_ro("cond", &IfNode::cond)
.def_ro("true_branch", &IfNode::true_branch)
.def_ro("false_branch", &IfNode::false_branch);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.If", IfNode, ExprNode);
};
class If : public Expr {
public:
/*!
* \brief The constructor
*
* \param cond The condition of a if node.
*
* \param true_branch The fall through branch. If this is not a
* SeqExpr, it will be wrapped in a SeqExpr, to satisfy the
* Relax IR requirement that all scopes be contained in a
* SeqExpr.
*
* \param false_branch The branch for execution when condition is
* false. If this is not a SeqExpr, it will be wrapped in a
* SeqExpr, to satisfy the Relax IR requirement that all scopes
* be contained in a SeqExpr.
*
* \param span The source span of the expression.
*/
TVM_DLL If(Expr cond, Expr true_branch, Expr false_branch, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(If, Expr, IfNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IfNode);
};
/*! \brief A Relax function. */
class FunctionNode : public BaseFuncNode {
public:
/*! \brief The parameters to the function. */
ffi::Array<Var> params;
/*! \brief The body of the function. */
SeqExpr body;
/*! \brief The return type of the function. */
Type ret_ty = Type::Missing();
/*! \brief Whether the function is annotated as pure or not. */
bool is_pure;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FunctionNode>()
.def_ro("params", &FunctionNode::params, refl::AttachFieldFlag::SEqHashDefRecursive())
.def_ro("body", &FunctionNode::body)
.def_ro("ret_ty", &FunctionNode::ret_ty)
.def_ro("is_pure", &FunctionNode::is_pure);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.Function", FunctionNode, BaseFuncNode);
};
class Function : public BaseFunc {
public:
/*!
* \brief Construct a Relax Function
*
* \param params The parameters accepted by the function
*
* \param body The body of the function. If this is not a
* SeqExpr, it will be wrapped in a SeqExpr, to satisfy the
* Relax IR requirement that all scopes be contained in a
* SeqExpr.
*
* \param ret_ty The Type returned by the function.
* If std::nullopt, will be inferred from the Type of the
* function's body.
*
* \param is_pure The purity of the function.
*
* \param attrs Any attributes associated with the function.
* Defaults to an empty dictionary.
*
* \param span The source span of the expression.
*/
TVM_DLL explicit Function(ffi::Array<Var> params, Expr body, ffi::Optional<Type> ret_ty,
bool is_pure = true, DictAttrs attrs = DictAttrs(), Span span = Span());
/*!
* \brief Mimics the constructor but without body Expr.
* \note ret_ty is required, since it can not deduced by the body.
*/
TVM_DLL static Function CreateEmpty(ffi::Array<Var> params, Type ret_ty, bool is_pure = true,
DictAttrs attrs = DictAttrs(), Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Function, BaseFunc, FunctionNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode);
};
// TODO(@sunggg): Investigate the exact usage of kComposite, kPartitionedFromPattern, and
// kPrimitive.
namespace attr {
/*! \brief Mark the function as a primitive function. */
constexpr const char* kPrimitive = "Primitive";
/*!
* \brief Indicate the codegen that should be used for building this function.
* When this is unset or set to "default", the default compilation pipeline will be used.
*/
constexpr const char* kCodegen = "Codegen";
/*! \brief Treat the function as a composite operator. */
constexpr const char* kComposite = "Composite";
/*! \brief Indicate the function was created by the Pattern Partitioning Pass. */
constexpr const char* kPartitionedFromPattern = "PartitionedFromPattern";
/*! \brief The required workspace for an external function. */
constexpr const char* kWorkspaceSize = "WorkspaceSize";
// Note: in the future, we prefer snake_case instead of CamelCase for attributes.
// Past ones will be kept for backwards compatibility.
/*! \brief Override checking purity for this function and treat as pure
* (is_pure must be set to true) */
constexpr const char* kForcePure = "relax.force_pure";
/*!
* \brief The number of inputs of a function.
* If a function has the num_input attribute, the last func->params.size() - num_inputs
* arguments are assumed to be weights that are fixed across invocations.
*/
constexpr const char* kNumInput = "num_input";
} // namespace attr
/*! \brief The extern function, which can represent packed function. */
class ExternFuncNode : public BaseFuncNode {
public:
/*! \brief The name of global symbol. */
ffi::String global_symbol;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExternFuncNode>().def_ro("global_symbol", &ExternFuncNode::global_symbol);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ExternFunc", ExternFuncNode, BaseFuncNode);
};
class ExternFunc : public BaseFunc {
public:
TVM_DLL ExternFunc(ffi::String global_symbol, Span span = Span());
TVM_DLL ExternFunc(ffi::String global_symbol, Type ty, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExternFunc, BaseFunc, ExternFuncNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ExternFuncNode);
};
/*!
* \brief Get the shape of Expr.
* \param expr The input expr.
* \return The corresonding shape.
*
* \note This function requires expr to be normalized.
* The function will report an error if expr's Type is not TensorType.
* It will try to return symbolic function when possible. If the tensor do not
* have a compile-time symbolic shape, the function will then choose to return
* Call(relax.op.shape_of, [expr]).
*/
TVM_DLL Expr GetShapeOf(const Expr& expr);
} // namespace relax
} // namespace tvm
/* \brief Allow relax.Var as key in STL tables
*
* For most Relax expressions, it would be ambiguous whether the
* expression should follow reference equality or structural equality.
* This is not the case for variables, which do not contain nested
* internal structure, and are frequently used as keys in lookup
* tables.
*
* Providing `std::hash` and `std::equal_to` specializations for
* `relax::Var` allows it to be used as a key in STL tables. For
* `relax::Expr`, the user must specify the type of equality used
* (e.g. `std::unordered_set<T, StructuralHash, StructuralEqual>` or
* `std::unordered_set<T, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>`).
*/
template <>
struct std::hash<tvm::relax::Var> {
std::size_t operator()(const tvm::relax::Var& var) const {
return tvm::ffi::ObjectPtrHash()(var);
}
};
template <>
struct std::equal_to<tvm::relax::Var> {
bool operator()(const tvm::relax::Var& var_a, const tvm::relax::Var& var_b) const {
return tvm::ffi::ObjectPtrEqual()(var_a, var_b);
}
};
#endif // TVM_RELAX_EXPR_H_
include/tvm/relax/expr_functor.h
+574
View File
@@ -0,0 +1,574 @@
/*
* 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/relax/expr_functor.h
* \brief A more powerful visitor which enables defining arbitrary function
* signatures with type based dispatch on first argument.
*/
#ifndef TVM_RELAX_EXPR_FUNCTOR_H_
#define TVM_RELAX_EXPR_FUNCTOR_H_
#include <tvm/ir/node_functor.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
#include <tvm/relax/type_functor.h>
#include <tvm/tirx/function.h>
#include <unordered_map>
#include <utility>
namespace tvm {
namespace relax {
/*!
* \brief A dynamical functor that dispatches on in the first Expr argument.
* You can use this as a more powerful Visitor, since it allows you to
* define function signatures of Visit Function.
*
* \sa tvm/ir_functor.h
*
* \tparam FType function signature
* This type is only defined for FType with function signature R(const Expr&,
* Args...)
*/
template <typename FType>
class ExprFunctor;
// functions to be overriden.
#define EXPR_FUNCTOR_DEFAULT \
{ \
return VisitExprDefault_(op, std::forward<Args>(args)...); \
}
#define RELAX_EXPR_FUNCTOR_DISPATCH(OP) \
vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \
return self->VisitExpr_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
});
#define PY_EXPR_VISITOR_DEFAULT(N, PY_FUNC, DEFAULT_FUNC) \
{ \
if (PY_FUNC != nullptr) \
PY_FUNC(N); \
else \
DEFAULT_FUNC; \
}
#define PY_EXPR_MUTATOR_DEFAULT(N, PY_FUNC, DEFAULT_FUNC, RET_TYPE) \
{ \
if (PY_FUNC != nullptr) { \
RET_TYPE ret = PY_FUNC(N).cast<RET_TYPE>(); \
return ret; \
} else { \
return DEFAULT_FUNC; \
} \
}
#define PY_EXPR_VISITOR_DISPATCH(OP, PY_FUNC) \
vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self) { \
if (self->PY_FUNC != nullptr) \
self->PY_FUNC(n); \
else \
self->VisitExpr_(static_cast<const OP*>(n.get())); \
});
#define PY_EXPR_MUTATOR_DISPATCH(OP, PY_FUNC) \
vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self) { \
if (self->PY_FUNC != nullptr) { \
Expr expr = self->PY_FUNC(n).cast<Expr>(); \
return expr; \
} else { \
return self->VisitExpr_(static_cast<const OP*>(n.get())); \
} \
});
#define PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(OP) \
post_order_vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self) { \
return self->VisitExprPostOrder_(static_cast<const OP*>(n.get())); \
});
template <typename R, typename... Args>
class ExprFunctor<R(const Expr& n, Args...)> {
private:
using TSelf = ExprFunctor<R(const Expr& n, Args...)>;
using FType = tvm::NodeFunctor<R(const ffi::ObjectRef& n, TSelf* self, Args...)>;
public:
/*! \brief the result type of this functor */
using result_type = R;
/*! \brief virtual destructor */
virtual ~ExprFunctor() {}
/*!
* \brief Same as call.
* \param n The expression node.
* \param args Additional arguments.
* \return The result of the call
*/
R operator()(const Expr& n, Args... args) { return VisitExpr(n, std::forward<Args>(args)...); }
/*!
* \brief The functor call.
* \param n The expression node.
* \param args Additional arguments.
* \return The result of the call
*/
virtual R VisitExpr(const Expr& n, Args... args) {
TVM_FFI_ICHECK(n.defined())
<< "Found null pointer node while traversing AST. The previous pass may "
"have generated invalid data.";
static FType vtable = InitVTable();
if (vtable.can_dispatch(n)) {
return vtable(n, this, std::forward<Args>(args)...);
}
return VisitExprFallback_(n.get(), std::forward<Args>(args)...);
}
// Functions that can be overriden by subclass
// NOTE: cross dialect calls are invoked through global var
// We do not expect inline PrimFunc to appear in relax IR.
virtual R VisitExpr_(const ConstantNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const DataflowVarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const ShapeExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const ExternFuncNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const GlobalVarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const FunctionNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const CallNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const SeqExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const IfNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const OpNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TupleGetItemNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExprFallback_(const ExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const StringImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const DataTypeImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExprDefault_(const ffi::Object* op, Args...) {
TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
throw;
}
private:
// initialize the vtable.
static FType InitVTable() {
FType vtable;
// Set dispatch
RELAX_EXPR_FUNCTOR_DISPATCH(ConstantNode);
RELAX_EXPR_FUNCTOR_DISPATCH(TupleNode);
RELAX_EXPR_FUNCTOR_DISPATCH(VarNode);
RELAX_EXPR_FUNCTOR_DISPATCH(DataflowVarNode);
RELAX_EXPR_FUNCTOR_DISPATCH(ShapeExprNode);
RELAX_EXPR_FUNCTOR_DISPATCH(ExternFuncNode);
RELAX_EXPR_FUNCTOR_DISPATCH(GlobalVarNode);
RELAX_EXPR_FUNCTOR_DISPATCH(FunctionNode);
RELAX_EXPR_FUNCTOR_DISPATCH(CallNode);
RELAX_EXPR_FUNCTOR_DISPATCH(SeqExprNode);
RELAX_EXPR_FUNCTOR_DISPATCH(IfNode);
RELAX_EXPR_FUNCTOR_DISPATCH(OpNode);
RELAX_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode);
RELAX_EXPR_FUNCTOR_DISPATCH(StringImmNode);
RELAX_EXPR_FUNCTOR_DISPATCH(DataTypeImmNode);
vtable.Finalize();
return vtable;
}
};
/*!
* \brief A simple visitor wrapper around ExprFunctor.
* Recursively visit the content.
*/
class ExprVisitor : public ExprFunctor<void(const Expr&)> {
public:
/*!
* \brief Generic dispatcher for Expr.
* \param expr The expr to be visited.
*/
void VisitExpr(const Expr& expr) override;
// specific leaf level visitor functions
void VisitExpr_(const ConstantNode* op) override;
void VisitExpr_(const TupleNode* op) override;
void VisitExpr_(const VarNode* op) override;
void VisitExpr_(const DataflowVarNode* op) override;
void VisitExpr_(const ShapeExprNode* op) override;
void VisitExpr_(const ExternFuncNode* op) override;
void VisitExpr_(const GlobalVarNode* op) override;
void VisitExpr_(const FunctionNode* op) override;
void VisitExpr_(const CallNode* op) override;
void VisitExpr_(const SeqExprNode* op) override;
void VisitExpr_(const IfNode* op) override;
void VisitExpr_(const OpNode* op) override;
void VisitExpr_(const TupleGetItemNode* op) override;
void VisitExprFallback_(const ExprNode* op) override;
void VisitExpr_(const StringImmNode* op) override;
void VisitExpr_(const DataTypeImmNode* op) override;
/*!
* \brief Generic dispatcher for bindings.
* \param binding The binding to be visited.
*/
virtual void VisitBinding(const Binding& binding);
// specific leaf level visitor functions
virtual void VisitBinding_(const VarBindingNode* binding);
virtual void VisitBinding_(const MatchCastNode* binding);
// second level dispatching based on binding value type.
// these dispatching functions get called from first-level dispatch on VarBinding
virtual void VisitBinding_(const VarBindingNode* binding, const ConstantNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const TupleNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const VarNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ShapeExprNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ExternFuncNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const GlobalVarNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const FunctionNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const CallNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const SeqExprNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const IfNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const OpNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const StringImmNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const DataTypeImmNode* val);
/*!
* \brief Generic dispatcher for binding blocks.
* \param block The binding block to be visited.
*/
virtual void VisitBindingBlock(const BindingBlock& block);
// specific leaf level visitor functions
virtual void VisitBindingBlock_(const BindingBlockNode* block);
virtual void VisitBindingBlock_(const DataflowBlockNode* block);
/*!
* \brief Generic dispatcher for visiting the var definition site.
* \param var The var to be visited.
* \note VisitExpr_(const VarNode*) will only visit the usage site of an Var
*/
virtual void VisitVarDef(const Var& var);
/*!
* \brief Visit ty may recursively contain Expr/PrimExpr.
*
* By default, this function recurse into type such as
* TensorType and ShapeType and call VisitExpr/VisitTypePrimExprField
* accordingly. It does not recurse into FunctionType as it does
* not contain Expr defined in the current scope.
*
* Pass writers can overload this function to change to other behaviors.
* For example, if we are not interested in Expr in Type, we can
* override this function by a no-op.
*
* \param ty Input type field.
*/
virtual void VisitExprDepTypeField(const Type& ty);
// specific leaf level visitor functions
virtual void VisitVarDef_(const VarNode* var);
virtual void VisitVarDef_(const DataflowVarNode* var);
virtual void VisitSpan(const Span& span);
virtual void VisitTypePrimExprField(const PrimExpr& expr);
private:
using TSelf = ExprVisitor;
using VisitBindingVTable = tvm::NodeFunctor<void(const ffi::ObjectRef& n, ExprVisitor* self,
const VarBindingNode* binding)>;
// initialize the vtable.
static VisitBindingVTable InitVisitBindingVTable();
/*!
* \brief Private internal type field visitor.
*
* Support default visiting of type field and recursive into
* their Expr fields.
*
* We use component instead of sub-classing so there can be other
* joint inheritance between ExprVisitor and TypeVisitor.
*/
class DefaultTypeFieldVisitor : public TypeVisitor {
public:
explicit DefaultTypeFieldVisitor(ExprVisitor* parent);
// Override defaults in type visitor.
void VisitTypeExprField(const Expr& expr) final;
void VisitTypeExprField(const PrimExpr& expr) final;
void VisitType_(const FuncTypeNode* op) final;
private:
ExprVisitor* parent_;
};
// This visitor is not visible to child classes and only
// used to supported default visiting behavior.
DefaultTypeFieldVisitor default_tyfield_visitor_{this};
};
void PostOrderVisit(const Expr& node, std::function<void(const Expr&)> fvisit);
/*!
* \brief A mutator works in unnormalized form.
*
* ExprMutatorBase expects input AST to be in the unnormalized form, i.e., ty
* of expressions can be nullptr, and the expressions may nest(and as a result the AST is not in
* ANF).
*/
class ExprMutatorBase : public ExprFunctor<Expr(const Expr&)> {
public:
Expr VisitExpr(const Expr& expr) override;
Expr VisitExpr_(const ConstantNode* op) override;
Expr VisitExpr_(const TupleNode* op) override;
Expr VisitExpr_(const VarNode* op) override;
Expr VisitExpr_(const DataflowVarNode* op) override;
Expr VisitExpr_(const ShapeExprNode* op) override;
Expr VisitExpr_(const ExternFuncNode* op) override;
Expr VisitExpr_(const GlobalVarNode* op) override;
Expr VisitExpr_(const FunctionNode* op) override;
Expr VisitExpr_(const CallNode* op) override;
Expr VisitExpr_(const SeqExprNode* op) override;
Expr VisitExpr_(const IfNode* op) override;
Expr VisitExpr_(const OpNode* op) override;
Expr VisitExpr_(const TupleGetItemNode* op) override;
Expr VisitExprFallback_(const ExprNode* op) override;
Expr VisitExpr_(const StringImmNode* op) override;
Expr VisitExpr_(const DataTypeImmNode* op) override;
/*!
* \brief Mutate BindingBlock.
* \param block The binding block to be visited.
* \return The binding block after transformation.
*/
virtual BindingBlock VisitBindingBlock(const BindingBlock& block);
/*!
* \brief Used to visit the PrimExpr inside of expressions.
*
* Can be overloaded to transform the shape expressions.
*/
virtual PrimExpr VisitTypePrimExprField(const PrimExpr& expr);
/*!
* \brief Visit ty that may recursively contain Expr/PrimExpr.
*
* By default, this function recurse into type such as
* TensorType and ShapeType and call VisitExpr/VisitTypePrimExprField
* accordingly. It does not recurse into FunctionType as it does
* not contain Expr defined in the current scope.
*
* Pass writers can overload this function to change to other behaviors.
* For example, if in Expr in Type won't change, we can
* override this function by an identity function.
*
* \param ty Input type field.
* \return The updated type.
*/
virtual Type VisitExprDepTypeField(const Type& ty);
protected:
/*!
* \brief Check whether VisitExprDepTypeField change ty.
* \return Whether type changed.
* \note This function is used by mutator implementations to check if
* previous Expr update will trigger a change in ty.
* If change is detected, the implementation can generate a fresh
* node without ty, and trigger normalizer to re-derive.
*/
bool VisitAndCheckTypeFieldUnchanged(const ffi::ObjectRef& ty) {
if (const TypeNode* ty_node = ty.as<TypeNode>()) {
Type type = ffi::GetRef<Type>(ty_node);
return type.IsMissing() || this->VisitExprDepTypeField(type).same_as(ty);
} else {
return true;
}
}
private:
/*!
* \brief Private internal type field visitor to support
* Default visiting of type field and recursive into their Expr fields.
*
* We use component instead of sub-classing so there can be other
* joint inheritance between ExprMutator and TypeMutator.
*/
class DefaultTypeFieldMutator : public TypeMutator {
public:
explicit DefaultTypeFieldMutator(ExprMutatorBase* parent);
// Override defaults in type visitor.
Expr VisitTypeExprField(const Expr& expr) final;
PrimExpr VisitTypeExprField(const PrimExpr& expr) final;
Type VisitType_(const FuncTypeNode* op) final;
private:
ExprMutatorBase* parent_;
};
// This visitor is not visible to child classes and only
// used to supported default visiting behavior.
DefaultTypeFieldMutator default_tyfield_mutator_{this};
};
/*!
* \brief A mutator works in normal form.
*
* ExprMutator expects input AST to be in the normal form, i.e., the expressions are normalized(no
* nesting and hence the AST is in ANF), and all ty of expressions are
* available.
*/
class ExprMutator : public ExprMutatorBase {
public:
using ExprMutatorBase::VisitExpr_;
ExprMutator(ffi::Optional<IRModule> mod = std::nullopt) { builder_ = BlockBuilder::Create(mod); }
Expr VisitExpr(const Expr& expr) override;
Expr VisitExpr_(const VarNode* op) override;
Expr VisitExpr_(const DataflowVarNode* op) override;
Expr VisitExpr_(const FunctionNode* op) override;
Expr VisitExpr_(const SeqExprNode* op) override;
Expr VisitExpr_(const IfNode* op) override;
/*!
* \brief Generic dispatcher for bindings.
* \param binding The binding to be visited.
*/
virtual void VisitBinding(const Binding& binding);
// specific leaf level visitor functions
virtual void VisitBinding_(const VarBindingNode* binding);
virtual void VisitBinding_(const MatchCastNode* binding);
// second level dispatching based on binding value type.
// these dispatching functions get called from first-level dispatch on VarBinding
virtual void VisitBinding_(const VarBindingNode* binding, const ConstantNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const TupleNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const VarNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ShapeExprNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ExternFuncNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const GlobalVarNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const FunctionNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const CallNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const SeqExprNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const IfNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const OpNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const StringImmNode* val);
virtual void VisitBinding_(const VarBindingNode* binding, const DataTypeImmNode* val);
/*!
* \brief Generic dispatcher for binding blocks.
* \param block The binding block to be visited.
* \return The binding block after transformation.
*/
virtual BindingBlock VisitBindingBlock(const BindingBlock& block) override; // NOLINT(*)
// specific leaf level visitor functions
virtual BindingBlock VisitBindingBlock_(const BindingBlockNode* block);
virtual BindingBlock VisitBindingBlock_(const DataflowBlockNode* block);
/*!
* \brief Generic dispatcher for rewriting the var definition site.
* \param var The var to be visited.
* \return The var after post-order rewritten.
* \note VisitExpr_(const VarNode*) will only visit the usage site of an Var
*/
virtual Var VisitVarDef(const Var& var);
// specific leaf level visitor functions
virtual Var VisitVarDef_(const VarNode* var);
virtual Var VisitVarDef_(const DataflowVarNode* var);
protected:
/*!
* \brief Try to remit binding and bind it to a new_value
*
* This function is called after VisitExpr(binding->value) in
* VisitBinding_(const VarBinding*).
* It will try to reuse the current binding when the new value's shape/type
* matches the original binding and no changes in var is needed.
*
* Otherwise, a new binding will be emitted to replace the var specified in
* the current binding.
*/
void ReEmitBinding(const VarBindingNode* binding, Expr new_value);
/*!
* \brief Rewrite the expr with a new scope, used in a Function's body.
*
* Visit an expression that may neither access variables from the
* current scope, nor may export definitions into the current scope.
*
* \param body_expr The body to be visited.
* \param params Optional parameters that are visible within the scope.
* \return The expr after visiting.
*
* \note The body_expr must be an SeqExpr in the normal form.
*/
Expr VisitWithNewScope(const Expr& body_expr,
ffi::Optional<ffi::Array<Var>> params = std::nullopt);
/*!
* \brief Rewrite the expr with a new scope, used in the branches of If.
*
* Visit an expression that may access variables from the current
* scope, but may not export definitions into the current scope.
*
* \param body_expr The body to be visited.
*
* \return The expr after visiting.
*
* \sa VisitWithNewScope
*
* \note The body_expr must be an SeqExpr in the normal form.
*/
Expr VisitWithInnerScope(const Expr& body_expr);
/*!
* \brief Look up the value bound to a variable.
* \param var The var to be looked up.
* \return The value bound to the input \p var.
* \note For function parameters, this function returns std::nullopt.
*/
ffi::Optional<Expr> LookupBinding(const Var& var);
/*!
* \brief Post-order rewrite a node and normalize.
* \tparam T The node type to be rewritten.
* \param op The node to be rewritten.
* \return The node after post rewritten.
*/
template <typename T>
Expr VisitExprPostOrder_(const T* op) {
return builder_->Normalize(ExprMutator::VisitExpr_(op));
}
/*!
* \brief Create a new var with specified type if the original var's shape or type does not
* match with the specified ones.
* \param var The var to be updated.
* \param ty The type to be updated.
* \return The var filled with type information.
*/
Var WithType(Var var, Type ty);
/*! \brief Internal block builder to emit bindings during rewriting. */
BlockBuilder builder_;
/*! \brief Remap a var to a new var in use-site. */
std::unordered_map<Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> var_remap_;
private:
using TSelf = ExprMutator;
using VisitBindingVTable = tvm::NodeFunctor<void(const ffi::ObjectRef& n, ExprMutator* self,
const VarBindingNode* binding)>;
// initialize the vtable.
static VisitBindingVTable InitVisitBindingVTable();
};
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_EXPR_FUNCTOR_H_
include/tvm/relax/nested_msg.h
+687
View File
@@ -0,0 +1,687 @@
/*
* 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/relax/nested_msg.h
* \brief Helper container to store nested message for robust tuple-aware analysis.
*
* Please see NestedMsg for description of usage.
*
* \sa NestedMsg
*/
#ifndef TVM_RELAX_NESTED_MSG_H_
#define TVM_RELAX_NESTED_MSG_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/optional.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
#include <string>
#include <utility>
#include <vector>
namespace tvm {
namespace relax {
/*!
* \brief Container that stores possibly nested message with leaf message type T.
*
* NestedMsg is a helper structure to store intermediate
* message state in pass analysis so we can robustly handle message
* passing with the presence of nested tuple types.
*
* Under the hood, NestedMsg[T] = Union[T, std::nullopt, Array[NestedMsg[T]]].
* Each nested message corresponds to the same nesting structure as
* the nested tuple types when we encounter them in analysis.
*
* Relax support nested tuple structures in the IR. Nested tuple structure
* is important to support advanced groupings in cases such as gradient calculation
* and other scenarios.
*
* The possible presence of nested tuple does mean that we need to
* to robustly handle analysis that contains nested tuple structures
* in a dataflow graph.
*
* \code
*
* v1 = relu(v0)
* v2 = exp(v0)
* t = ((v0, v1), (v2,), v0)
* t1 = t[0]
* v3 = concat(t1)
* v4 = t[2]
* v5 = add(v4, v3)
*
* \endcode
*
* Consider the above code sequence that contains a mixture of tuple
* nesting and normal operations. A common message-passing-based analysis
* will track messages attached to each intermediate variable.
*
* Because the intermediate value can contain nested-tuples, we need to have
* abilities to nest messages according to tuple structure and propagate them
* along the way. In python, this simply corresponds to using a tuple to hold
* nested messages. This class provides a helper wrapper in C++ to present such
* possibly nested message for a given leaf message.
*
* This design pattern is necessary to handle tuple values regardless of
* the normal form design of the IR to enable different messages for each
* tuple component without enforcing all tuple elements to have the same message.
*
* Please consider the following patterns in our pass:
*
* On a forward propagation message passing analysis:
* - Create map [leafnode=>NestedMsg<T>], scan forward
* - input_msg = [MapToNestedMsg<T>(x, lookup_map) for x in call->args]
* - output_msg = ForwardProp[call->op](input_msg, call)
* - map[binding->var] = output_msg
* - Use MapToNestedMsg to remap the remaining body.
*
* On a backward propagation message passing analysis:
* - Create map [leafnode=>NestedMsg<T>], scan backward
* - output_msg = lookup map(binding->var)
* - handle case when output_msg is null
* - input_msg = BackProp[call->op](out_msg, call)
* - for arg, msg in zip(call->args, input_msg),
* DecomposeNestedMessage(arg, msg, lambda node, m: update_map(node, m))
* - update_map(node, m) => CombineNestedMessage(map[node], m)
*
* Here leafnode is a node that you would like to propagate messages to
* such as constant, var and should not include tuple.
*
* We also recommend writing unit-test cases that involve nested tuple composition
* and decomposition.
*
* \sa MapToNestedMsg, DecomposeNestedMsg, CombineNestedMsg, ForEachLeaf, Equal
*
* \note If you want to write robust message passing-based analysis for
* programs that can contain nested tuples, you likely need to
* use this class or logic of a similar kind.
*/
template <typename T>
class NestedMsg {
public:
// default constructors.
NestedMsg() = default;
NestedMsg(const NestedMsg<T>&) = default;
NestedMsg(NestedMsg<T>&&) = default;
NestedMsg<T>& operator=(const NestedMsg<T>&) = default;
NestedMsg<T>& operator=(NestedMsg<T>&&) = default;
/*! \brief Nullopt handling */
NestedMsg(std::nullopt_t) {} // NOLINT(*)
// nullptr handling.
// disallow implicit conversion as 0 can be implicitly converted to nullptr_t
explicit NestedMsg(std::nullptr_t) {}
NestedMsg<T>& operator=(std::nullptr_t) {
data_ = nullptr;
return *this;
}
// normal value handling.
NestedMsg(T other) // NOLINT(*)
: data_(std::move(other)) {}
NestedMsg<T>& operator=(T other) {
data_ = std::move(other);
return *this;
}
// ffi::Array<NestedMsg<T>> handling
NestedMsg(ffi::Array<NestedMsg<T>, void> other) // NOLINT(*)
: data_(other) {}
NestedMsg<T>& operator=(ffi::Array<NestedMsg<T>, void> other) {
data_ = std::move(other);
return *this;
}
// initializer list handling
NestedMsg(std::initializer_list<NestedMsg<T>> other) // NOLINT(*)
: NestedMsg(ffi::Array<NestedMsg<T>, void>(other)) {}
NestedMsg<T>& operator=(std::initializer_list<NestedMsg<T>> other) {
return operator=(ffi::Array<NestedMsg<T>, void>(other));
}
// delete the int constructor
// since NestedMsg<IntImm>(0) is ambiguous
// 0 can be implicitly casted to nullptr_t
explicit NestedMsg(int val) = delete;
NestedMsg<T>& operator=(int val) = delete;
// operator overloadings
bool operator==(std::nullptr_t) const { return data_ == nullptr; }
bool operator!=(std::nullptr_t) const { return data_ != nullptr; }
/*! \return Whether the nested message is not-null leaf value */
bool IsLeaf() const {
return data_.type_index() != ffi::TypeIndex::kTVMFFINone &&
data_.type_index() != ffi::TypeIndex::kTVMFFIArray;
}
/*! \return Whether the nested message is null */
bool IsNull() const { return data_.type_index() == ffi::TypeIndex::kTVMFFINone; }
/*! \return Whether the nested message is nested */
bool IsNested() const { return data_.type_index() == ffi::TypeIndex::kTVMFFIArray; }
/*!
* \return The underlying leaf value.
* \note This function checks if the msg is leaf.
*/
T LeafValue() const {
TVM_FFI_ICHECK(IsLeaf());
return ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck<T>(data_);
}
/*!
* \return a corresponding nested array.
* \note This checks if the underlying data type is array.
*/
ffi::Array<NestedMsg<T>, void> NestedArray() const {
return ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck<ffi::Array<NestedMsg<T>, void>>(
data_);
}
private:
ffi::Any data_;
// private constructor
explicit NestedMsg(ffi::Any data) : data_(data) {}
template <typename, typename>
friend struct ffi::TypeTraits;
};
/*!
* \brief Apply fvisit for each leaf elements in the nested message.
* \param fvisit The visit callback.
* \param msg The input nested message.
* \tparam T the content type of nested msg
* \tparam FType the visitor type with signature void fvisit(T)
*/
template <typename T, typename FType>
void ForEachLeaf(const NestedMsg<T>& msg, FType fvisit) {
if (msg == nullptr) return;
if (msg.IsLeaf()) {
fvisit(msg.LeafValue());
} else {
for (NestedMsg<T> x : msg.NestedArray()) {
ForEachLeaf(x, fvisit);
}
}
}
/*!
* \brief Recursively compare two nested messages.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \param fequal The equal functor with signature bool fequal(T, T)
* \tparam T the content type of nested msg
* \tparam FType the equal comparator type
*/
template <typename T, typename FType>
bool Equal(const NestedMsg<T>& lhs, const NestedMsg<T>& rhs, FType fequal) {
if (lhs.IsNull()) return rhs.IsNull();
if (rhs.IsNull()) return lhs.IsNull();
if (lhs.IsLeaf()) {
return rhs.IsLeaf() && fequal(lhs.LeafValue(), rhs.LeafValue());
} else {
if (!rhs.IsNested()) return false;
ffi::Array<NestedMsg<T>> arr_lhs = lhs.NestedArray();
ffi::Array<NestedMsg<T>> arr_rhs = rhs.NestedArray();
if (arr_lhs.size() != arr_rhs.size()) return false;
for (size_t i = 0; i < arr_lhs.size(); ++i) {
if (!Equal(arr_lhs[i], arr_rhs[i], fequal)) return false;
}
return true;
}
}
/*!
* \brief Map expr with possible nested-tuple to nested message.
*
* This function will unpack recursive tuples and run fmapleaf for each leaf,
* then recursively combines the results together into a NestedMsg.
*
* The nesting structure will corresponds to the tuple structure.
*
* \param expr The input expression.
* \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmap(Expr)`
* \tparam T the content type of nested msg
* \tparam FType The mapping function type
*/
template <typename T, typename FType>
NestedMsg<T> MapToNestedMsg(Expr expr, FType fmapleaf) {
if (auto* tuple = expr.as<TupleNode>()) {
ffi::Array<NestedMsg<T>> res;
res.reserve(tuple->fields.size());
for (Expr x : tuple->fields) {
res.push_back(MapToNestedMsg<T, FType>(x, fmapleaf));
}
return res;
} else {
return fmapleaf(expr);
}
}
/*!
* \brief Map structinfo with possible nested-ty to nested message.
*
* This function will unpack recursive ty and run fmapleaf for each leaf,
* then recursively combines the results together into a NestedMsg.
*
* The nesting structure will corresponds to the tuple structure.
*
* \param ty The input type.
* \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmap(Type)`
* \tparam T the content type of nested msg
* \tparam FType The mapping function type
*/
template <typename T, typename FType>
NestedMsg<T> MapToNestedMsg(Type ty, FType fmapleaf) {
if (auto* tuple = ty.as<TupleTypeNode>()) {
ffi::Array<NestedMsg<T>> res;
res.reserve(tuple->fields.size());
for (Type x : tuple->fields) {
res.push_back(MapToNestedMsg<T, FType>(x, fmapleaf));
}
return res;
} else {
return fmapleaf(ty);
}
}
/*!
* \brief Map expr with possible nested-tuple to nested message.
*
* This function will unpack recursive expr by its type and
* run fmapleaf for each leaf, then recursively combines the results
* together into a NestedMsg.
*
* The nesting structure will corresponds to the type of expr.
*
* \param expr The input expression which should have type.
* \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmapleaf(Expr)`
* \tparam T the content type of nested msg
* \tparam FType The mapping function type
*/
template <typename T, typename FType>
NestedMsg<T> MapToNestedMsgByType(Expr expr, FType fmapleaf) {
auto ty = GetType(expr);
if (auto* tuple = ty.as<TupleTypeNode>()) {
ffi::Array<NestedMsg<T>> res;
res.reserve(tuple->fields.size());
for (size_t i = 0; i < tuple->fields.size(); ++i) {
Expr field;
if (const auto* expr_tuple = expr.as<TupleNode>()) {
field = expr_tuple->fields[i];
} else {
field = TupleGetItem(expr, i);
}
res.push_back(MapToNestedMsgByType<T, FType>(field, fmapleaf));
}
return res;
} else {
return fmapleaf(expr);
}
}
/*!
* \brief Map nested message back to TargetType.
*
* This function will decompose the nested message and
* run fmapleaf for each leaf message and get the leaf value,
* then recursively combines the results by fcombine.
*
* \param msg The input nested message.
* \param fmapleaf The mapping function for each leaf with signature
* `TargetType fmapleaf(ffi::Optional<T>)`.
* \param fcombine The function for combining all childs of a node into TargetType with signature
* `TargetType fmapleaf(ffi::Array<TargetType>)`.
* \tparam TargetType the target type to map nested msg to.
* \tparam T the content type of nested msg.
* \tparam FMapLeaf The leaf mapping function type.
* \tparam FCombine The combining function type.
*/
template <typename TargetType, typename T, typename FMapLeaf, typename FCombine>
TargetType NestedMsgTo(NestedMsg<T> msg, FMapLeaf fmapleaf, FCombine fcombine) {
if (msg.IsNull()) {
return fmapleaf(std::nullopt);
} else if (msg.IsLeaf()) {
return fmapleaf(msg.LeafValue());
} else {
TVM_FFI_ICHECK(msg.IsNested());
ffi::Array<NestedMsg<T>> arr = msg.NestedArray();
ffi::Array<TargetType> subexpr;
subexpr.reserve(arr.size());
for (size_t i = 0; i < arr.size(); ++i) {
subexpr.push_back(NestedMsgTo<TargetType>(arr[i], fmapleaf, fcombine));
}
return fcombine(subexpr);
}
}
/*!
* \brief Map nested message back to the expr.
*
* This function will decompose the nested message and
* run fmapleaf for each leaf message and get the leaf expr,
* then recursively combines the results as tuple expr.
*
* \param msg The input nested message.
* \param fmapleaf The mapping function for each leaf with signature `Expr
* fmapleaf(ffi::Optional<T>)`. \tparam T the content type of nested msg. \tparam FType The mapping
* function type.
*/
template <typename T, typename FType>
Expr NestedMsgToExpr(NestedMsg<T> msg, FType fmapleaf) {
return NestedMsgTo<Expr>(msg, fmapleaf, [](ffi::Array<Expr> arr) {
ffi::Optional<Expr> simplified_tuple;
bool simplified_flag = false;
if (arr.size() >= 1) {
simplified_flag = true;
for (size_t i = 0; i < arr.size() && simplified_flag; ++i) {
auto* node = arr[i].as<TupleGetItemNode>();
if (node == nullptr || node->index != static_cast<int>(i)) {
simplified_flag = false;
} else {
if (simplified_tuple.has_value()) {
simplified_flag &= simplified_tuple.value().same_as(node->tuple);
} else {
simplified_tuple = node->tuple;
TVM_FFI_ICHECK(simplified_tuple.has_value());
}
}
}
}
return simplified_flag ? simplified_tuple.value() : Tuple(arr);
});
}
/*!
* \brief Recursively combine two nested message into one.
*
* This function requires the two messages to be compatible with each other.
* The combination rule is as follows:
* - combine(null, msg) => msg
* - combine(leaf1, leaf2) => fcombine(leaf1, leaf2)
* - combine(array1, array2) => [combine(x, y) for x, y in zip(array1, array2)]
* - This function will throw an error if array have different size
*
* \param lhs The left operand.
* \param rhs The right operand.
* \param fcombine with signature T fcombine(T lhs, T rhs)
* \tparam T the content type of nested msg
* \tparam FType combine function type.
*/
template <typename T, typename FType>
NestedMsg<T> CombineNestedMsg(NestedMsg<T> lhs, NestedMsg<T> rhs, FType fcombine) {
if (lhs.IsNull()) return rhs;
if (rhs.IsNull()) return lhs;
if (lhs.IsLeaf()) {
TVM_FFI_ICHECK(rhs.IsLeaf()) << "Cannot combine leaf with nested";
return NestedMsg<T>(fcombine(lhs.LeafValue(), rhs.LeafValue()));
} else {
TVM_FFI_ICHECK(lhs.IsNested());
TVM_FFI_ICHECK(rhs.IsNested()) << "Cannot combine leaf with nested";
ffi::Array<NestedMsg<T>> arr_lhs = lhs.NestedArray();
ffi::Array<NestedMsg<T>> arr_rhs = rhs.NestedArray();
TVM_FFI_ICHECK_EQ(arr_lhs.size(), arr_rhs.size())
<< "Cannot combine two nested array with different sizes";
ffi::Array<NestedMsg<T>> res;
res.reserve(arr_lhs.size());
for (size_t i = 0; i < arr_lhs.size(); ++i) {
res.push_back(CombineNestedMsg<T, FType>(arr_lhs[i], arr_rhs[i], fcombine));
}
return NestedMsg<T>(res);
}
}
/*!
* \brief Recursively map a nested message to another one, with leaf mapped by the input fmapleaf.
* \param msg The nested message to be mapped.
* \param fmapleaf The leaf map function, with signature NestedMsg<T> fmapleaf(T msg)
* \tparam T The content type of nested message.
* \tparam FType The leaf map function type.
* \return The new nested message.
*/
template <typename T, typename FType>
NestedMsg<T> MapNestedMsg(NestedMsg<T> msg, FType fmapleaf) {
if (msg.IsNull()) {
return msg;
} else if (msg.IsLeaf()) {
return fmapleaf(msg.LeafValue());
} else {
TVM_FFI_ICHECK(msg.IsNested());
ffi::Array<NestedMsg<T>> arr = msg.NestedArray();
ffi::Array<NestedMsg<T>> res;
res.reserve(arr.size());
for (int i = 0; i < static_cast<int>(arr.size()); ++i) {
res.push_back(MapNestedMsg(arr[i], fmapleaf));
}
return NestedMsg<T>(res);
}
}
/*!
* \brief Recursively decompose the tuple structure in expr and msg along with it.
*
* This function will call fvisitleaf for each leaf expression in expr.
* This function will throw an error if the nesting structure in msg does not
* match the tuple nesting structure in expr.
*
* \param expr The input expression to be decomposed.
* \param msg The input nested message.
* \param fvisitleaf with signature fvisitleaf(Expr expr, NestedMsg<T> msg)
* \tparam T the content type of nested msg
* \tparam FType The visit function type.
*/
template <typename T, typename FType>
void DecomposeNestedMsg(Expr expr, NestedMsg<T> msg, FType fvisitleaf) {
if (auto* tuple = expr.as<TupleNode>()) {
TVM_FFI_ICHECK(msg.IsNested()) << "Expected nested to match tuple";
ffi::Array<NestedMsg<T>> arr = msg.NestedArray();
TVM_FFI_ICHECK_EQ(arr.size(), tuple->fields.size())
<< "Expected nested array size to match tuple size";
for (size_t i = 0; i < arr.size(); ++i) {
DecomposeNestedMsg(tuple->fields[i], arr[i], fvisitleaf);
}
} else {
fvisitleaf(expr, msg);
}
}
/*!
* \brief Recursively transform the tuple structure in expr and msgs along with it.
*
* This function will call ftransleaf for each leaf expression in expr.
* This function will throw an error if the nesting structure in msg does not
* match the tuple nesting structure in expr.
*
* \param expr The input expression to be transform. 
* \param msgs The input messages to guide the transformation.
* \param ftransleaf with signature ftransleaf(Expr, ffi::Array<NestedMsg<T>>)->Expr
* \tparam T the content type of nested msg
* \tparam N the number of messages
* \tparam FType The visit function type.
*/
template <typename T, std::size_t N, typename FType>
Expr TransformTupleLeaf(Expr expr, std::array<NestedMsg<T>, N> msgs, FType ftransleaf) {
Type ty = GetType(expr);
if (const auto* tuple = ty.as<TupleTypeNode>()) {
std::array<ffi::Array<NestedMsg<T>>, N> msg_arrays;
for (size_t i = 0; i < N; ++i) {
TVM_FFI_ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple";
msg_arrays[i] = msgs[i].NestedArray();
}
bool same = true;
ffi::Array<Expr> fields;
fields.reserve(tuple->fields.size());
for (size_t i = 0; i < tuple->fields.size(); ++i) {
Expr field;
if (const auto* expr_tuple = expr.as<TupleNode>()) {
field = expr_tuple->fields[i];
} else {
field = TupleGetItem(expr, i);
}
std::array<NestedMsg<T>, N> sub_msgs;
for (size_t j = 0; j < N; ++j) {
sub_msgs[j] = msg_arrays[j][i];
}
fields.push_back(TransformTupleLeaf(field, std::move(sub_msgs), ftransleaf));
same &= (fields.back().same_as(field));
}
return same ? expr : Tuple(fields);
} else {
for (const auto& msg : msgs) {
TVM_FFI_ICHECK(msg.IsLeaf()) << "Expected leaf to match non-tuple";
}
return ftransleaf(expr, msgs);
}
}
/*!
* \brief Recursively transform the tuple structure in ty and msgs along with it.
*
* This function will call ftransleaf for each leaf ty in ty.
* This function will throw an error if the nesting structure in msg does not
* match the tuple nesting structure in ty.
*
* \param ty The input ty to be transform. 
* \param msgs The input messages to guide the transformation.
* \param ftransleaf with signature ftransleaf(Type, ffi::Array<NestedMsg<T>>)->Type
* \tparam T the content type of nested msg
* \tparam N the number of messages
* \tparam FType The visit function type.
*/
template <typename T, std::size_t N, typename FType>
Type TransformTupleLeaf(Type ty, std::array<NestedMsg<T>, N> msgs, FType ftransleaf) {
if (const auto* tuple = ty.as<TupleTypeNode>()) {
std::array<ffi::Array<NestedMsg<T>>, N> msg_arrays;
for (size_t i = 0; i < N; ++i) {
TVM_FFI_ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple";
msg_arrays[i] = msgs[i].NestedArray();
}
bool same = true;
ffi::Array<Type> fields;
fields.reserve(tuple->fields.size());
for (size_t i = 0; i < tuple->fields.size(); ++i) {
Type field = tuple->fields[i];
std::array<NestedMsg<T>, N> sub_msgs;
for (size_t j = 0; j < N; ++j) {
sub_msgs[j] = msg_arrays[j][i];
}
fields.push_back(TransformTupleLeaf(field, std::move(sub_msgs), ftransleaf));
same &= (fields.back().same_as(field));
}
return same ? ty : TupleType(fields);
} else {
for (const auto& msg : msgs) {
TVM_FFI_ICHECK(msg.IsLeaf()) << "Expected leaf to match non-tuple";
}
return ftransleaf(ty, msgs);
}
}
} // namespace relax
namespace ffi {
template <typename T>
inline constexpr bool use_default_type_traits_v<relax::NestedMsg<T>> = false;
template <typename T>
struct TypeTraits<relax::NestedMsg<T>> : public TypeTraitsBase {
TVM_FFI_INLINE static void CopyToAnyView(const relax::NestedMsg<T>& src, TVMFFIAny* result) {
*result = ffi::AnyView(src.data_).CopyToTVMFFIAny();
}
TVM_FFI_INLINE static void MoveToAny(relax::NestedMsg<T> src, TVMFFIAny* result) {
*result = details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(src.data_));
}
TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
return TypeTraitsBase::GetMismatchTypeInfo(src);
}
static bool CheckAnyStrict(const TVMFFIAny* src) {
if (src->type_index == TypeIndex::kTVMFFINone) {
return true;
}
if (TypeTraits<T>::CheckAnyStrict(src)) {
return true;
}
if (src->type_index == TypeIndex::kTVMFFIArray) {
const ffi::ArrayObj* array = reinterpret_cast<const ffi::ArrayObj*>(src->v_obj);
for (size_t i = 0; i < array->size(); ++i) {
const Any& any_v = (*array)[i];
if (!details::AnyUnsafe::CheckAnyStrict<relax::NestedMsg<T>>(any_v)) return false;
}
}
return true;
}
TVM_FFI_INLINE static relax::NestedMsg<T> CopyFromAnyViewAfterCheck(const TVMFFIAny* src) {
return relax::NestedMsg<T>(Any(AnyView::CopyFromTVMFFIAny(*src)));
}
TVM_FFI_INLINE static relax::NestedMsg<T> MoveFromAnyAfterCheck(TVMFFIAny* src) {
return relax::NestedMsg<T>(details::AnyUnsafe::MoveTVMFFIAnyToAny(src));
}
static std::optional<relax::NestedMsg<T>> TryCastFromAnyView(const TVMFFIAny* src) {
if (CheckAnyStrict(src)) {
return CopyFromAnyViewAfterCheck(src);
}
// slow path run conversion
if (src->type_index == TypeIndex::kTVMFFINone) {
return relax::NestedMsg<T>(std::nullopt);
}
if (auto opt_value = TypeTraits<T>::TryCastFromAnyView(src)) {
return relax::NestedMsg<T>(*std::move(opt_value));
}
if (src->type_index == TypeIndex::kTVMFFIArray) {
const ArrayObj* n = reinterpret_cast<const ArrayObj*>(src->v_obj);
ffi::Array<relax::NestedMsg<T>> result;
result.reserve(n->size());
for (size_t i = 0; i < n->size(); i++) {
const Any& any_v = (*n)[i];
if (auto opt_v = any_v.try_cast<relax::NestedMsg<T>>()) {
result.push_back(*std::move(opt_v));
} else {
return std::nullopt;
}
}
return relax::NestedMsg<T>(result);
}
return std::nullopt;
}
TVM_FFI_INLINE static std::string TypeStr() {
return "NestedMsg<" + details::Type2Str<T>::v() + ">";
}
TVM_FFI_INLINE static std::string TypeSchema() {
std::ostringstream oss;
oss << R"({"type":"NestedMsg","args":[)";
oss << details::TypeSchema<T>::v();
oss << "]}";
return oss.str();
}
};
} // namespace ffi
} // namespace tvm
#endif // TVM_RELAX_NESTED_MSG_H_
include/tvm/relax/op_attr_types.h
+151
View File
@@ -0,0 +1,151 @@
/*
* 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/relax/op_attr_types.h
* \brief Data structures that can appear in operator attributes.
*/
#ifndef TVM_RELAX_OP_ATTR_TYPES_H_
#define TVM_RELAX_OP_ATTR_TYPES_H_
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
#include <tvm/te/tensor.h>
namespace tvm {
namespace relax {
enum OpPatternKind {
// Elementwise operation
kElemWise = 0,
// Broadcasting operator, can always map output axis to the input in order.
// for example :code:`out[i, ax1, j, ax2] = input[i, j]`.
// Note that the axis need to be in order so transpose is not a bcast operator.
kBroadcast = 1,
// Injective operator, can always injectively map output axis to a single input axis.
// All injective operator can still be safely fused to injective and reduction.
kInjective = 2,
// Communicative reduction operator.
kCommReduce = 3,
// Complex operation, can still fuse elemwise operations into its output.
// but cannot chain another complex op
kOutEWiseFusable = 4,
// The pattern for tuple nodes. Can fuse into subsequent injective ops,
// but treated specially
kTuple = 7,
// Opaque operation, cannot fuse anything.
kOpaque = 8
};
/*!
* \brief Packed function implementation for operators. The relax operator will be lowered to
* this packed function call during codegen.
*/
using FCallPacked = ffi::String;
/*!
* \brief Infer output type given the call
*
* \param call The call expression to be derived.
* \param ctx The builder context.
*/
using FInferType = ffi::TypedFunction<Type(const Call& call, const BlockBuilder& ctx)>;
/*!
* \brief The function type of a normalization function.
*
* A normalization function is used when a `relax::Call` may be
* expressed in multiple syntactically valid and semantically
* equivalent forms, to normalize to a single representation.
*
* Note: `FNormalize` is applied for each expression as part of the
* `relax::BlockBuilder`. While operator-specific validation may
* be performed within the `FNormalize` implementation, ensuring
* that errors are caught as early as possible, this should only be
* used when validation is fast to apply. If the validation logic
* may be slow, it should instead be implemented in `FValidate`,
* which is only run as part of the well-formed checker.
*
* \param bb The BlockBuilder context.
*
* \param call The call to be normalized. It is provided by-value, to
* avoid copies for the common case where the call is already normalized.
*/
using FNormalize = ffi::TypedFunction<Expr(const BlockBuilder& bb, Call call)>;
/*!
* \brief The function type of a validation function.
*
* A validation function is used to define constraints that should be
* verified for an operator as part of the well-formed checker.
*
* Note: `FValidate` is only applied as part of the well-formed
* checker. While this minimizes overhead while compiling Relax,
* this delay between generating an ill-formed `relax::Call` and
* identifying the ill-formed call may complicate debugging. If
* the validation logic is very fast to check, and doing so would
* not introduce a significant overhead, consider validating as part
* of `FNormalize`, which is applied by the block builder for each
* `relax::Call`.
*
* \param call The call to be validated.
*/
using FValidate = ffi::TypedFunction<void(const Call& call)>;
/*! \brief The function type of a legalization function.
*
* A legalization function is used to replace a `relax::Call` with
* more concrete implementations. For example, the operation
* `relax.op.add` may be replaced with a call to a TIR function
* implementing addition of two tensors.
*
* The purpose of `FLegalize` is to remove calls to the operator while
* lowering. Therefore, unlike `FNormalize`, the resulting expression
* may *not* contain the original operator.
*
* \param bb The BlockBuilder context.
* \param call The call to be legalized.
*/
using FLegalize = ffi::TypedFunction<Expr(const BlockBuilder& bb, const Call& call)>;
/*! \brief The function type of a function to lower the runtime builtin.
*
* A builtin function may be lowered to a lowered form in `LowerRuntimeBuiltin`.
*
* \param bb The BlockBuilder context.
* \param call The call to be lowered.
*/
using FLowerBuiltin = ffi::TypedFunction<Expr(const BlockBuilder& bb, const Call& call)>;
/*!
* \brief Gradient for a specific op.
*
* \param orig_var the original var corresponding to orig_call.
* \param orig_call the original Call(op) expr.
* \param output_grad the gradient of the Expr.
* \param ctx the current block builder context.
* \return the gradient for each parameter.
*/
using FPrimalGradient = ffi::TypedFunction<tvm::ffi::Array<Expr>(
const Var& orig_var, const Call& orig_call, const Var& output_grad, const BlockBuilder& ctx)>;
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_OP_ATTR_TYPES_H_
include/tvm/relax/script/builder/frame.h
+332
View File
@@ -0,0 +1,332 @@
/*
* 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.
*/
#ifndef TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
#define TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <tvm/script/ir_builder/ir/ir.h>
#include <utility>
namespace tvm {
namespace script {
namespace ir_builder {
namespace relax {
/*! \brief The base ir_builder frame for the relax dialect. */
class RelaxFrameNode : public IRBuilderFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RelaxFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.RelaxFrame", RelaxFrameNode,
IRBuilderFrameNode);
};
class RelaxFrame : public IRBuilderFrame {
public:
explicit RelaxFrame(ffi::ObjectPtr<RelaxFrameNode> data) : IRBuilderFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RelaxFrame, IRBuilderFrame, RelaxFrameNode);
protected:
RelaxFrame() = default;
};
/*! \brief The base ir_builder frame for frames with SeqExpr
i.e. Functions, If branches
*/
class SeqExprFrameNode : public RelaxFrameNode {
public:
/*! \brief The binding blocks inside the frame. */
ffi::Array<tvm::relax::BindingBlock> binding_blocks;
/*! \brief The frame output expr. `std::nullopt` when undefined. */
ffi::Optional<tvm::relax::Expr> output;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SeqExprFrameNode>()
.def_ro("binding_blocks", &SeqExprFrameNode::binding_blocks)
.def_ro("output", &SeqExprFrameNode::output);
}
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.SeqExprFrame", SeqExprFrameNode,
RelaxFrameNode);
public:
void EnterWithScope() override;
void ExitWithScope() override;
};
class SeqExprFrame : public RelaxFrame {
public:
explicit SeqExprFrame(ffi::ObjectPtr<SeqExprFrameNode> data) : RelaxFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqExprFrame, RelaxFrame, SeqExprFrameNode);
};
/*! \brief The ir_builder frame for the relax function. */
class FunctionFrameNode : public SeqExprFrameNode {
public:
/*!
* \brief The function name.
* \note The name will not be specified in constructor, so it is "Optional",
* However, we must specify the name by `R.func_name` before exit this frame.
*/
ffi::Optional<ffi::String> name;
/*! \brief The function params. */
ffi::Array<tvm::relax::Var> params;
/*!
* \brief The function return type.
* \note Usually the function return type can be deduced by the function body.
* But we can use this field to specify a more "accurate" return type.
* i.e. If the `ret_ty` is None, try to use the deduced type from body
* If the `ret_ty` is not None, we can still take body.ty
* if we ret_ty is base of body.ty. If not, we will
* take the specified `ret_ty`.
*/
ffi::Optional<tvm::Type> ret_ty;
/*! \brief Whether the function is annotated as pure */
ffi::Optional<bool> is_pure;
/*! \brief Whether the function is annotated as private */
ffi::Optional<bool> is_private;
/*! \brief The function attributes. */
ffi::Map<ffi::String, Any> attrs;
/*! \brief The block builder to create Relax function. */
tvm::relax::BlockBuilder block_builder;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FunctionFrameNode>()
.def_ro("name", &FunctionFrameNode::name)
.def_ro("params", &FunctionFrameNode::params)
.def_ro("ret_ty", &FunctionFrameNode::ret_ty)
.def_ro("is_pure", &FunctionFrameNode::is_pure)
.def_ro("attrs", &FunctionFrameNode::attrs);
// `binding_blocks` and `output` are inherited from SeqExprFrameNode.
// `block_builder` is not registered as it's not visited.
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.FunctionFrame", FunctionFrameNode,
SeqExprFrameNode);
public:
void EnterWithScope() final;
void ExitWithScope() final;
};
class FunctionFrame : public SeqExprFrame {
public:
explicit FunctionFrame(ffi::ObjectPtr<FunctionFrameNode> data) : SeqExprFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionFrame, SeqExprFrame, FunctionFrameNode);
};
/*! \brief The ir_builder frame for relax binding blocks. */
class BindingBlockFrameNode : public RelaxFrameNode {
public:
/*! \brief The flag that indicates whether the block is a dataflow block. */
bool is_dataflow;
/*! \brief The variables emitted in this block. */
ffi::Array<tvm::relax::Var> emitted_vars;
/*!
* \brief A boolean indicating if the dataflow block is ended of construction.
* If it is true, any new binding trying to be emitted into this block will cause an error.
* \note Only used for a dataflow block.
*/
bool block_ended;
/*!
* \brief The output vars of the dataflow block.
* \note Only used for a dataflow block.
*/
ffi::Array<tvm::relax::Var> output_vars;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BindingBlockFrameNode>()
.def_ro("is_dataflow", &BindingBlockFrameNode::is_dataflow)
.def_ro("emitted_vars", &BindingBlockFrameNode::emitted_vars)
.def_ro("output_vars", &BindingBlockFrameNode::output_vars);
// `block_ended` is not registered as it's not visited.
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.BindingBlockFrame",
BindingBlockFrameNode, RelaxFrameNode);
public:
void EnterWithScope() final;
void ExitWithScope() final;
};
class BindingBlockFrame : public RelaxFrame {
public:
explicit BindingBlockFrame(ffi::ObjectPtr<BindingBlockFrameNode> data) : RelaxFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BindingBlockFrame, RelaxFrame,
BindingBlockFrameNode);
};
/*!
* \brief A frame that represents if statement.
*
* \sa IfFrame
*/
class IfFrameNode : public RelaxFrameNode {
public:
/*! \brief The condition of the if statement. */
tvm::relax::Expr condition;
/*! \brief The Bindings in the true branch. */
ffi::Optional<tvm::relax::Expr> then_expr;
/*! \brief The Bindings in the false branch. */
ffi::Optional<tvm::relax::Expr> else_expr;
/*! \brief The Binding var. */
tvm::relax::Var var;
/*! \brief The binding var name. */
ffi::String var_name;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IfFrameNode>()
.def_ro("condition", &IfFrameNode::condition)
.def_ro("then_expr", &IfFrameNode::then_expr)
.def_ro("else_expr", &IfFrameNode::else_expr)
.def_ro("var", &IfFrameNode::var)
.def_ro("var_name", &IfFrameNode::var_name);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.IfFrame", IfFrameNode, RelaxFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to IfFrameNode.
*
* \sa IfFrameNode
*/
class IfFrame : public RelaxFrame {
public:
explicit IfFrame(ffi::ObjectPtr<IfFrameNode> data) : RelaxFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfFrame, RelaxFrame, IfFrameNode);
};
/*!
* \brief A frame that represents then.
*
* \sa ThenFrame
*/
class ThenFrameNode : public SeqExprFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ThenFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ThenFrame", ThenFrameNode,
SeqExprFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ThenFrameNode.
*
* \sa ThenFrameNode
*/
class ThenFrame : public SeqExprFrame {
public:
explicit ThenFrame(ffi::ObjectPtr<ThenFrameNode> data) : SeqExprFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ThenFrame, SeqExprFrame, ThenFrameNode);
};
/*!
* \brief A frame that represents else.
*
* \sa ElseFrame
*/
class ElseFrameNode : public SeqExprFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ElseFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ElseFrame", ElseFrameNode,
SeqExprFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ElseFrameNode.
*
* \sa ElseFrameNode
*/
class ElseFrame : public SeqExprFrame {
public:
explicit ElseFrame(ffi::ObjectPtr<ElseFrameNode> data) : SeqExprFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, SeqExprFrame, ElseFrameNode);
};
} // namespace relax
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
include/tvm/relax/script/builder/ir.h
+144
View File
@@ -0,0 +1,144 @@
/*
* 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.
*/
#ifndef TVM_RELAX_SCRIPT_BUILDER_IR_H_
#define TVM_RELAX_SCRIPT_BUILDER_IR_H_
#include <tvm/relax/expr.h>
#include <tvm/relax/script/builder/frame.h>
#include <tvm/relax/type.h>
#include <tvm/script/ir_builder/base.h>
namespace tvm {
namespace script {
namespace ir_builder {
namespace relax {
/////////////////////////////// Function ////////////////////////////////
/*!
* \brief Start a function frame.
* \param is_pure Whether the function is annotated as pure.
* \param is_private Whether the function is annotated as private.
* \return The created ir_builder Function frame.
*/
TVM_DLL FunctionFrame Function(bool is_pure, bool is_private);
/*!
* \brief Add a parameter to the last function frame.
* \param name The name of the parameter.
* \param ty The ty of the parameter.
* \return The created function parameter var.
*/
TVM_DLL tvm::relax::Var Arg(const ffi::String& name, const tvm::Type& ty);
/*!
* \brief Specify the name of the last function frame.
* \param name The function name.
*/
TVM_DLL void FuncName(const ffi::String& name);
/*!
* \brief Specify the attrs of the last function frame.
* \param attrs The function attrs.
*/
TVM_DLL void FuncAttrs(ffi::Map<ffi::String, Any> attrs);
/*!
* \brief Specify the return type of the last function frame.
* \param ret_ty The return type.
*/
TVM_DLL void FuncRetType(const tvm::Type& ret_ty);
/*!
* \brief Specify the return value of the last function frame.
* \param value The return value.
*/
TVM_DLL void FuncRetValue(const tvm::relax::Expr& value);
///////////////////////////// BindingBlock //////////////////////////////
/*!
* \brief Start a binding block frame.
* \return The created ir_builder Block frame.
*/
TVM_DLL BindingBlockFrame BindingBlock();
/*!
* \brief Start a dataflow binding block frame.
* \return The created ir_builder Block frame.
*/
TVM_DLL BindingBlockFrame Dataflow();
/*!
* \brief Expose the dataflow block output variables as global ones
* \param vars The output variables of a dataflow block
*/
TVM_DLL void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars);
////////////////////////////// Bindings ////////////////////////////////
/*!
* \brief Emit a binding to the last binding block frame.
* \param value The right side value of the bindings to be emitted.
* \param annotate_ty The optional type annotation for the emitted value.
* \return The left side var of the emitted binding.
*/
TVM_DLL tvm::relax::Var Emit(const tvm::relax::Expr& value,
const ffi::Optional<tvm::Type>& annotate_ty = std::nullopt);
/*!
* \brief Emit a match_cast binding to the last binding block frame.
* \param value The value of the MatchCast to be emitted.
* \param ty The type of the MatchCast to be emitted.
* \return The left side var of the emitted binding.
*/
TVM_DLL tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& ty);
/*!
* \brief Emit a binding to the last binding block frame.
* \param binding The binding to be emitted.
* \return The left side var of the emitted binding.
*/
TVM_DLL tvm::relax::Var EmitVarBinding(const tvm::relax::VarBinding& binding);
///////////////////////////// If Then Else /////////////////////////////
/*!
* \brief Create an if statement.
* \param condition The condition of if statement.
* \return The result IfFrame.
*/
IfFrame If(tvm::relax::Expr condition);
/*!
* \brief Create a then.
* \return The result ThenFrame.
*/
ThenFrame Then();
/*!
* \brief Create an else.
* \return The result ElseFrame.
*/
ElseFrame Else();
} // namespace relax
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_RELAX_SCRIPT_BUILDER_IR_H_
include/tvm/relax/tir_pattern.h
+78
View File
@@ -0,0 +1,78 @@
/*
* 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 tir_pattern.h
* \brief Data Structure of TIR Pattern used for matching.
*/
#ifndef TVM_RELAX_TIR_PATTERN_H_
#define TVM_RELAX_TIR_PATTERN_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/function.h>
namespace tvm {
namespace relax {
using TIRPattern = tirx::PrimFunc;
/*
* \brief The match result of a TIR pattern.
*/
class MatchResultNode : public ffi::Object {
public:
/*! The matched tirx pattern*/
TIRPattern pattern;
/*! \brief The evaluated values of symbolic vars. */
ffi::Array<PrimExpr> symbol_values;
/*! \brief The matched buffers of input and output. */
ffi::Array<tirx::Buffer> matched_buffers;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MatchResultNode>()
.def_ro("pattern", &MatchResultNode::pattern)
.def_ro("symbol_values", &MatchResultNode::symbol_values)
.def_ro("matched_buffers", &MatchResultNode::matched_buffers);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.MatchResult", MatchResultNode, ffi::Object);
};
/*!
* \brief Managed reference to MatchResultNode.
*/
class MatchResult : public ffi::ObjectRef {
public:
/*!
* \brief Constructor
* \param pattern The matched tirx pattern.
* \param symbol_values The evaluated values of symbolic vars.
* \param matched_buffers The matched buffers of input and output.
*/
TVM_DLL explicit MatchResult(TIRPattern pattern, ffi::Array<PrimExpr> symbol_values,
ffi::Array<tirx::Buffer> matched_buffers);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchResult, ffi::ObjectRef, MatchResultNode);
};
using FCodegen = ffi::TypedFunction<ffi::Array<ffi::Any>(ffi::Array<MatchResult> match_results)>;
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_TIR_PATTERN_H_
include/tvm/relax/transform.h
+687
View File
@@ -0,0 +1,687 @@
/*
* 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/relax/transform.h
* \brief Relax specific transformation passes.
*/
#ifndef TVM_RELAX_TRANSFORM_H_
#define TVM_RELAX_TRANSFORM_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/transform.h>
#include <tvm/relax/dataflow_pattern.h>
#include <tvm/relax/expr.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/index_map.h>
namespace tvm {
namespace relax {
namespace transform {
using Pass = tvm::transform::Pass;
using PassInfo = tvm::transform::PassInfo;
using PassContext = tvm::transform::PassContext;
using Function = tvm::relax::Function;
using DataflowBlock = tvm::relax::DataflowBlock;
using tvm::transform::CreateModulePass;
using LayoutCb = ffi::TypedFunction<ffi::Map<ffi::String, ffi::Array<ffi::String>>(Call)>;
/*!
* \brief Create a function pass.
*
* \param pass_func The packed function that contains the optimization.
* \param opt_level The optimization level of the function pass.
* \param name The name of the function pass.
* \param required The list of the passes that the function pass is dependent on.
* \param traceable Boolean variable whether the dataflowblock pass is traceable.
*
* \return The created function pass.
*/
TVM_DLL Pass CreateFunctionPass(std::function<Function(Function, IRModule, PassContext)> pass_func,
int opt_level, ffi::String name,
tvm::ffi::Array<ffi::String> required, bool traceable = false);
/*!
* \brief Create a dataflowblock pass.
*
* \param pass_func The packed function that contains the optimization.
* \param opt_level The optimization level of the dataflowblock pass.
* \param name The name of the dataflowblock pass.
* \param required The list of the passes that the dataflowblock pass is dependent on.
* \param traceable Boolean variable whether the dataflowblock pass is traceable.
*
* \return The created dataflowblock pass.
*/
TVM_DLL Pass CreateDataflowBlockPass(
std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func, int opt_level,
ffi::String name, tvm::ffi::Array<ffi::String> required, bool traceable = false);
/*!
* \brief Perform lambda lifting to lift functions from nested into global.
*
* \return The Pass.
*/
TVM_DLL Pass LambdaLift();
/*!
* \brief Transform all dataflow structure to non-dataflow version.
*
* \return The Pass.
*/
TVM_DLL Pass ToNonDataflow();
/*!
* \brief Activate force_pure on all pure functions in the module
* and unwrap all pure override ops into the normal versions.
*
* This effectively means that there will be no more purity tracking,
* useful for low-level code generation.
*
* \return The Pass.
*
* \note Should be used after ToNonDataflow()
*/
TVM_DLL Pass RemovePurityChecking();
/*!
* \brief Perform explicit tensor allocation for call_tir and call_dps_packed.
*
* \return The Pass.
*/
TVM_DLL Pass CallTIRRewrite();
/*!
* \brief Convert all reshape-like call_tir whose corresponding binding
* vars are DataflowVars to relax.reshape operator calls. The relax.reshape
* calls will be lowered an external builtin function call in a subsequent
* pass, where the external builtin function does a CreateView operation
* at runtime, instead of doing real data copy.
* Here "reshape-like" includes reshape, expand_dims, flatten, etc.
*
* \return The Pass.
* \note The pass is applied at the first stage of Relax VM build, before
* rewriting call_tir, as this pass requires dataflow information.
*/
TVM_DLL Pass RewriteDataflowReshape();
/*!
* \brief The static memory planning pass on BindingBlock level.
* The pass will reuse allocated memory to its best effort, in order to
* reduce the total amount of allocated memory size.
*
* The pass "supports" dynamic shape in the way of TIR variable bound
* annotations. We can optionally annotate the attributes "tir_var_upper_bound"
* and "tir_var_lower_bound" to Relax functions. The attribute values are dicts
* from strings to integers, denoting the name of TIR variables to the bound
* values of the TIR vars.
* Note: The annotated bound attributes only apply to TIR vars in the
* function signature for clarity.
*
* For example, we can annotate a Relax function with
* `R.func_attr({"tir_var_lower_bound": {"n": 1}, "tir_var_upper_bound": {"n": 1024}})`.
* It means the variable that names "n" in the function signature will have
* range [1, 1024]. And we will use these bounds during memory planning.
* If lower bound is not specified, it defaults to 0.
*
* \return The pass.
*/
TVM_DLL Pass StaticPlanBlockMemory();
/*!
* \brief Attach global_symbol to Relax functions and TIR Primfuncs for codegen.
*
* \return The Pass.
*/
TVM_DLL Pass AttachGlobalSymbol();
/*!
* \brief Transform Relax IR to normal form: transform AST to A-normal form, and fill the
* ty of expressions.
*
* \return The Pass.
*/
TVM_DLL Pass Normalize();
/*!
* \brief Possibly rename the GlobalVar in an IRModule to ensure these properties:
* 1. (Invariant) First ensure every public function has the same name as its "global_symbol"
* attribute;
* 2. To ensure 1., we may need to rename private functions with conflicting names;
* 3. Finally, the name of every GlobalVar is unique in the IRModule.
*/
TVM_DLL Pass NormalizeGlobalVar();
/*!
* \brief Simplify a Relax module by folding var bindings and match shape nodes,
* as well as tuple indices.
* Best used alongside constant folding and eliminating unused bindings.
*
* \note If a dataflow var is used only in a binding to the dataflow block
* output var (i.e., a non-dataflow var), this pass will also remove the dataflow var
* and replaces the output var's binding with the dataflow var's direct definition.
*
* \return The Pass.
*/
TVM_DLL Pass CanonicalizeBindings();
/*!
* Eliminate common subexpressions within functions.
* \return The pass that eliminates common subexpressions.
*
* \note For nested functions, this pass performs CSE *within* those functions.
* \param call_only If true, enable eliminating only call nodes.
*/
TVM_DLL Pass EliminateCommonSubexpr(bool call_only = false);
/*!
* \brief Bind params of function of the module to constant tensors.
*
* \param func_name The name of the function to bind parameters.
* \param params The parameters to bind.
*
* \return The Pass.
*/
TVM_DLL Pass BindParams(ffi::String func_name, ffi::Map<Any, ffi::ObjectRef> params);
/*!
* \brief Bind symbolic vars to constant shape values.
*
* \param binding_map The dictionary of symbolic variables and their
* constant shape values. Dictionary keys may be either a
* `tirx.Var` or a string name of the variable. If the variables
* are referred to by name, the name must uniquely identify a
* symbolic variable in each function where it is used.
*
* \param func_name The name of the function in which to bind shape
* values. If std::nullopt, all functions in the module will be
* updated.
*
* \return The Pass.
*/
TVM_DLL Pass BindSymbolicVars(ffi::Map<ffi::Variant<tirx::Var, ffi::String>, PrimExpr> binding_map,
ffi::Optional<ffi::String> func_name = std::nullopt);
/*!
* \brief Fold constant expressions within dataflow blocks.
*
* \note ConvertToDataflow may need to be called first to provide dataflow blocks.
*
* \return The Pass.
*/
TVM_DLL Pass FoldConstant();
/*!
* \brief Legalize high-level operator calls in Relax functions to call_tir
* with corresponding low-level TIR PrimFuncs.
*
* For each high-level operator, we register the way of legalizing it as a
* function, which takes a context BlockBuilder and the Call being legalized
* as input, and returns the legalized call. Here the input BlockBuilder is
* mainly used for adding the PrimFunc created by call_te into the context
* IRModule.
*
* The legalization function for each operator is registered as an attribute (with
* attribute key `FLegalize`) of the operator.
*
* For customizability, the user can pass their own legalization by an optional customized map,
* with the key to be the operator name and value to be the legalization function.
* The default legalization function will be overridden by the customized one.
*
* \param cmap The customized operator legalization function map. The customized function
* will override the default one.
* \param skip_ops The list operator names which need to be skipped from legalization
* \param enable_warning A boolean value indicating if to print warnings for TIR functions not
* showing up in the database.
* \return The Pass.
*/
TVM_DLL Pass LegalizeOps(ffi::Optional<ffi::Map<ffi::String, ffi::Function>> cmap,
ffi::Optional<ffi::Array<ffi::String>> skip_ops,
bool enable_warning = false);
/*!
* \brief Propagate virtual device information.
* \return The Pass.
*/
TVM_DLL Pass RealizeVDevice();
/*!
* \brief Attach layout free buffers to the tirx::PrimFunc.
*
* This pass is used to attach layout free buffers to the tirx::PrimFunc according to
* the function usage in the relax function. Currently, the layout free buffers are the model
* weights and relax constants.
*
* \note We recommend applying CanonicalizeBindings before this pass.
* \return The Pass.
*/
TVM_DLL Pass AttachAttrLayoutFreeBuffers();
/*!
* \brief Split the layout rewrite preproc block to a separate tirx::PrimFunc.
*
* This pass is used in the prepack weight after meta_schedule tuning.
*
* \return The Pass.
*/
TVM_DLL Pass SplitLayoutRewritePreproc();
/*!
* \brief Lift transformation of the parameters of a function.
*
* When some inputs of the function is marked as 'parameters' (the model weights), this pass
* identifies the transformation of the parameters and lifts them to a separate function called
* `transform_params`. `transform_params` takes a tuple of the original parameters as input and
* returns a tuple of the transformed parameters. The original function will be rewritten to accept
* a tuple of transformed parameters as input.
*
* Users are expected to invoke the `transform_params` function in runtime and pass the transformed
* parameters to the original function as input.
*
* \param shared_transform Indicates how the parameter transformation function will be produced.
* - `False` (default): A separate parameter transformation function will be produced for each
* function with the `"num_input"` attribute.
*
* - `True`: A single parameter transformation function will be produced, containing the
* preprocessing steps common across all functions with the `"num_input"` attribute.
*
* - List[str]: A single parameter transformation function will be produced, containing the
* preprocessing steps common across each function whose name is in the list. Passing a list of
* all functions with the `"num_input"` attribute or an empty list is equivalent to passing
* `True`.
*
* \return The Pass.
*/
TVM_DLL Pass
LiftTransformParams(ffi::Variant<bool, ffi::Array<ffi::String>> shared_transform = false);
/*!
* \brief Update virtual device.
* \param new_vdevice The new virtual device.
* \param index The device index indicates the device on which the update will be performed.
* \return The Pass.
*/
TVM_DLL Pass UpdateVDevice(VDevice new_vdevice, int64_t index);
/*! \brief Expand tuple arguments to internal functions
*
* \return The Pass
*/
TVM_DLL Pass ExpandTupleArguments();
/*! \brief Remove unused parameters to internal functions
*
* \return The Pass
*/
TVM_DLL Pass RemoveUnusedParameters();
/*! \brief Remove unused outputs from internal functions
*
* \return The Pass
*/
TVM_DLL Pass RemoveUnusedOutputs();
/*!
* \brief Annotate Op Pattern Kind for TIR functions, which is used in FuseOps.
* \note It is an auto-detect pass for "unscheduled prim_funcs", the op_pattern will be
* "opaque" of we can't detect it. Users can manually annotate the attr `op_pattern`
* to prim_func.
* \return The Pass.
*/
TVM_DLL Pass AnnotateTIROpPattern();
/*!
* \brief This pass groups bindings in a dataflow block of Relax functions and generates a new
* grouped Relax function for each group, according to the fusion algorithm described in the pass
* implementation. By grouping bindings into new Relax functions, we substitute the bindings in the
* function being manipulated into function calls to the new grouped function.
*
* A follow-up pass named "FuseTIR" will generate a TIR PrimFunc for each grouped function.
* \param fuse_opt_level The level of fuse optimization.
* -1 indicates that the level will be inferred from pass context.
* \return The Pass.
*/
TVM_DLL Pass FuseOps(int fuse_opt_level = -1);
/*!
* \brief The pattern object used as the input of FuseOpsByPattern. For bindings to be
* fused, it needs to be matched with `pattern` and the `check` function needs to return
* true.
*/
class FusionPatternNode : public ffi::Object {
public:
/*!
* \brief The name of pattern. It becomes the value of the kComposite attribute
* of a fused function after successful matching
*/
ffi::String name;
/*!
* \brief The dataflow pattern that will be used to match expression in the DataflowBlock.
* All the call nodes covered by the pattern will be extracted into the fused function.
*/
DFPattern pattern;
/*!
* \brief The map which is used to extract important expressions from the pattern match
* result. All DFPattern in this map should be part of the `pattern`.
*/
ffi::Map<ffi::String, DFPattern> annotation_patterns;
/*!
* \brief The function to determine whether the match result is accepted. This can be
* std::nullopt if check function is not necessary for this pattern.
*
* It should have signature
* bool(const PatternCheckContext& context)
*/
ffi::Optional<ffi::Function> check;
/*!
* \brief The function to get attributes for fused function
*
* It should have signature
* ffi::Map<ffi::String, Any>(const ffi::Map<ffi::String, Expr>& context)
*/
ffi::Optional<ffi::Function> attrs_getter;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FusionPatternNode>()
.def_ro("name", &FusionPatternNode::name)
.def_ro("pattern", &FusionPatternNode::pattern)
.def_ro("annotation_patterns", &FusionPatternNode::annotation_patterns)
.def_ro("check", &FusionPatternNode::check)
.def_ro("attrs_getter", &FusionPatternNode::attrs_getter);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.transform.FusionPattern", FusionPatternNode,
ffi::Object);
};
class FusionPattern : public ffi::ObjectRef {
public:
FusionPattern(ffi::String name, DFPattern pattern,
ffi::Map<ffi::String, DFPattern> annotation_patterns,
ffi::Optional<ffi::Function> check, ffi::Optional<ffi::Function> attrs_getter);
FusionPattern(ffi::String name, DFPattern pattern)
: FusionPattern(name, pattern, {}, std::nullopt, std::nullopt) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FusionPattern, ffi::ObjectRef, FusionPatternNode);
};
/*!
* \brief The input of FusionPattern::check.
*/
class PatternCheckContextNode : public ffi::Object {
public:
/*!
* \brief The expression that's matched with the FusionPattern::pattern.
*/
Expr matched_expr;
/*!
* \brief A map which contains all expressions matched by the sub patterns in
* FusionPattern::annotation_patterns.
*/
ffi::Map<ffi::String, Expr> annotated_expr;
/*!
* \brief Map from variable to its value. It contains variables from bindings that
* is being fused by FuseOpsByPattern.
*/
ffi::Map<Var, Expr> matched_bindings;
/*!
* \brief A map mapping variable definitions to a set of uses. It has all variables
* used in the function.
*/
ffi::Map<Var, ffi::Array<Var>> var_usages;
/*!
* \brief Map from value to its bound variable. It doesn't have variables after the
* matched expression.
*/
ffi::Map<Expr, Var> value_to_bound_var;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PatternCheckContextNode>()
.def_ro("matched_expr", &PatternCheckContextNode::matched_expr)
.def_ro("annotated_expr", &PatternCheckContextNode::annotated_expr)
.def_ro("matched_bindings", &PatternCheckContextNode::matched_bindings)
.def_ro("var_usages", &PatternCheckContextNode::var_usages)
.def_ro("value_to_bound_var", &PatternCheckContextNode::value_to_bound_var);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.transform.PatternCheckContext", PatternCheckContextNode,
ffi::Object);
};
class PatternCheckContext : public ffi::ObjectRef {
public:
PatternCheckContext(Expr matched_expr, ffi::Map<ffi::String, Expr> annotated_expr,
ffi::Map<Var, Expr> matched_bindings,
ffi::Map<Var, ffi::Array<Var>> var_usages,
ffi::Map<Expr, Var> value_to_bound_var);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PatternCheckContext, ffi::ObjectRef,
PatternCheckContextNode);
};
/*!
* \brief Reverse-mode automatic differentiation.
*
* This pass will differentiate one function in the IRModule. Now the input function must have only
* one dataflow block.
*
* For a given function specified by `func_name`, it generates a new function with the name
* `func_name + "_adjoint"`. The new function computes the gradient of the **differentiation
* target** with respect to the arguments specified by `require_grads` of the original function.
*
* If the function has only one return value, the return value will be specified as target. If the
* function has more than one return values, the target will be specified as the target_index-th
* return value. The target must be a scalar (0-dim tensor).
*
* \param func_name The name of the specified function.
* \param require_grads The relax variables whose adjoints is needed. Must be parameters of the
* given function and should not be duplicate. If it is not specified, adjoints of all parameters
* would be computed.
* \param target_index If the specified function has more than one return values, specify the index
* of the return value as the target. If it is not specified, the first return value will be the
* target.
* \return The Pass.
*
* \note ConvertToDataflow may need to be called first to provide dataflow blocks.
*/
TVM_DLL Pass Gradient(ffi::String func_name,
ffi::Optional<ffi::Array<Var>> require_grads = std::nullopt,
int target_index = 0);
/*!
* \brief Apply pattern matching to each function in the given module, and group matched
* expressions into a new function. The end result is similar to FuseOps, but fusion is driven
* completely by the provided patterns.
*
* \param patterns The patterns to detect. The order of the patterns determines the order
* of priority in which they are matched. Higher-priority patterns should come earlier in the list.
* \param bind_constants Whether or not to keep bound constants of the grouped function.
* \param annotate_codegen If true, wrap each created composite function with another function,
* whose body consists only of a call to the composite function, and annotate the outer function
* with kCodegen and kGlobalSymbol attributes. The kCodegen attribute is set as the prefix of the
* corresponding pattern name. For example, "dnnl" if the pattern name is "dnnl.conv2d_relu".
* This must be True if the created composite functions are intended to be offloaded to
* an external backend without using the MergeCompositeFunctions pass.
* \param entry_function_names The names of functions that should be considered as entry points. If
* not specified, all externally exposed functions will be considered as entry points.
* \return The Pass.
*
* \note Only operates within dataflow blocks. ConvertToDataflow may need to be called first.
*/
TVM_DLL Pass FuseOpsByPattern(const tvm::ffi::Array<FusionPattern>& patterns,
bool bind_constants = true, bool annotate_codegen = false,
const tvm::ffi::Array<ffi::String>& entry_function_names = {});
/*!
* \brief Group one or multiple composite functions created by FuseOpsByPattern into a new
* function. The new function will be annotated with kCodegen and GlobalSymbol attributes,
* and it is intented to be offloaded to an external backend.
*
* \return The Pass.
*/
TVM_DLL Pass MergeCompositeFunctions();
/*!
* \brief Fuse relax sub-function into a larger TIR function if possible.
this pass works together with FuseOps to perform operator fusion.
* \return The Pass.
*/
TVM_DLL Pass FuseTIR();
/*!
* \brief Run codegen.
* \param target_options pairs of target name and compilation options
* \param entry_functions list of entry functions
* \return The Pass.
*/
TVM_DLL Pass
RunCodegen(ffi::Optional<ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::Any>>> target_options,
ffi::Array<ffi::String> entry_functions);
/*!
* \brief Decompose composite operators during inference. For example, The result of batch norm (a
* triple) will be simplified. Operators like Attention, Erf, etc. can be also simplified into
* several operators as well.
*
* \param func_name The name of the specified function. If not specified, the pass will run in
* all functions.
*/
TVM_DLL Pass DecomposeOpsForInference(ffi::Optional<ffi::String> func_name);
/*!
* \brief Decompose composite operators during training. For example, The result of batch norm (a
* triple) will be simplified. Operators like Attention, Erf, etc. can be also simplified into
* several operators as well.
*
* \param func_name The name of the specified function. If not specified, the pass will run in
* all functions.
*/
TVM_DLL Pass DecomposeOpsForTraining(ffi::Optional<ffi::String> func_name);
/*!
* \brief Returns a pass which replaces PrimFuncs which have matching kOperatorName attribute in \p
* op_impl_map, with replacement PrimFunc that could possibly have different layouts on i/o
* buffers. The layout transformations on i/o buffers is present in the \p op_buffer_transforms. The
* pass inserts the layout transformations in the call sites of PrimFuncs being replaced to
* transform i/o buffers into expected layout.
*
* \param op_impl_map Map from kOperatorName attr (e.g., relax.conv2d) to replacement PrimFunc
* \param op_buffer_transforms Map from kOperatorName attr to layout transformations on each of the
* PrimFunc i/o buffers.
* \param axis_separators Map from kOperatorName attr to axis_separators of each buffer_transforms
* \param input_axis_separators Map from kOperatorName attr to axis_separator for input buffer
* \return The Pass.
*/
TVM_DLL Pass AlterOpImpl(
const ffi::Map<ffi::String, tirx::PrimFunc>& op_impl_map,
const ffi::Map<ffi::String, ffi::Array<tirx::IndexMap>>& op_buffer_transforms,
const ffi::Map<ffi::String, ffi::Optional<ffi::Array<ffi::Array<IntImm>>>>& axis_separators,
const ffi::Map<ffi::String, ffi::Optional<ffi::Array<ffi::Array<IntImm>>>>&
input_axis_separators);
/*!
* \brief Layout conversion pass.
* \param desired_layouts The desired layouts for some operators.
* \param layout_cb custom call back to define layouts dynamically.
* \return The Pass.
* \note Operates only on dataflow blocks. ConvertToDataflow may need to be called first.
*/
TVM_DLL Pass ConvertLayout(ffi::Map<ffi::String, ffi::Array<ffi::String>> desired_layouts,
LayoutCb layout_cb);
/*!
* \brief A pass that converts consecutive dataflow operations
* inside binding blocks into dataflow blocks.
* \param min_size The minimum number of consecutive dataflow bindings
* required for the pass to create a new dataflow block
* \return The Pass.
*/
TVM_DLL Pass ConvertToDataflow(int min_size = 2);
/*!
* \brief Dead code elimination.
* \sa RemoveAllUnused
* Currently it removes:
* 1. Unused local VarBindings
* (those where the bound var is unused and no impure operation is used).
* 2. Unused Relax functions in the module.
* We detect the call chain from the entry function, and remove all unused functions.
*
* Any binding blocks that are left empty will be removed by the normalizer.
*
* \param entry_functions Names of functions that should be considered
* as entry points, in addition to any externally exposed functions.
*
* \return The Pass.
*/
TVM_DLL Pass DeadCodeElimination(ffi::Array<ffi::String> entry_functions = {});
/*!
* \brief Pass that changes calls to operators that can be done in-place
* (generally, these are elementwise operations) in dataflow blocks into in-place implementations.
* Supported operators will be replaced by calls to `call_tir_inplace` that invoke in-place
* PrimFunc implementations of those operators (which are based on the legalizations of those
* operators).
* \note ConvertToDataflow may need to be called first to provide dataflow blocks.
* \return The pass.
*/
TVM_DLL Pass DataflowUseInplaceCalls();
/*!
* \brief Automatic mixed precision pass. Currently the pass assumes the input module to be fp32
* only, and will automatically cast fp32 to fp16 for certain ops.
* \param out_dtype The output data type of gemm/conv, which is the data type of the accumulator.
* \param fp16_input_names The names of function parameters whose dtype should become fp16. The
* function signature would change accordingly.
* \return The Pass.
*
* \note Mainly operates within dataflow blocks. ConvertToDataflow may need to be called first.
*/
TVM_DLL Pass ToMixedPrecision(
DLDataType out_dtype, ffi::Optional<ffi::Array<ffi::String>> fp16_input_names = std::nullopt);
/*!
* \brief Rewrite a Relax module for executing with CUDA graph. This pass identifies
* the regions that can be executed with CUDA graph and lifts them into new functions for runtime
* graph capturing.
*/
TVM_DLL Pass RewriteCUDAGraph();
/*!
* \brief This pass updates the var_buffer mapping of PrimFunctions from the call_tir info.
* Primarily used to update the VDevice information if any changes occurred from the caller.
* This pass recreates the buffers and updates the map.
*/
TVM_DLL Pass SpecializePrimFuncBasedOnCallSite();
} // namespace transform
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_TRANSFORM_H_
include/tvm/relax/type.h
+429
View File
@@ -0,0 +1,429 @@
/*
* 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/relax/type.h
* \brief Relax types, including the richer dependent Relax type nodes.
*/
#ifndef TVM_RELAX_TYPE_H_
#define TVM_RELAX_TYPE_H_
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/global_info.h>
#include <tvm/ir/type.h>
#include <tvm/tirx/expr.h>
#include <string>
#include <utility>
namespace tvm {
namespace relax {
using Expr = tvm::Expr;
using ExprNode = tvm::ExprNode;
class BlockBuilder;
/*! \brief Indicates the number of dimensions of a tensor is unknown at compile time. */
static constexpr int kUnknownNDim = -1;
using tvm::TupleType;
using tvm::TupleTypeNode;
class PackedFuncTypeNode : public TypeNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PackedFuncTypeNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PackedFuncType", PackedFuncTypeNode, TypeNode);
};
class PackedFuncType : public Type {
public:
TVM_DLL PackedFuncType(Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PackedFuncType, Type, PackedFuncTypeNode);
};
/*!
* \brief Base type of all Relax type information.
*
* Type stores possible type information deduced during compile-time.
* It encapsulates both static type and runtime information such as shape.
*
* Type of each non-primitive Expr can be deduced during compilation in a
* "best-effort" manner.
*
* When ty appears in function parameter and return signatures, it
* implies a runtime check that matches the type information with the value.
*
* When it appears in Expr, it follows "assume-semantics", which means the
* compiler will take the deduced information as it is and only do best effort
* proofs and checks.
*
* Each type can be uniquely erased to a static-type. The compiler will
* still compile the code, with less information, when we erase to the static
* type.
*
* If a Type contains an Expr field, then that field must already be
* normalized through NormalizeArg. This invariant is checked in constructors
* and simplifies assumptions during type deduction.
*/
/*!
* \brief Any Relax value.
*/
class AnyTypeNode : public TypeNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AnyTypeNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.AnyType", AnyTypeNode, TypeNode);
};
/*!
* \brief Managed reference to AnyTypeNode.
* \sa AnyTypeNode
*/
class AnyType : public Type {
public:
TVM_DLL AnyType(Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AnyType, Type, AnyTypeNode);
};
// Compatibility aliases for existing C++ callers. New code should use AnyType.
using ObjectTypeNode = AnyTypeNode;
using ObjectType = AnyType;
/*!
* \brief Type of shape value.
*/
class ShapeTypeNode : public TypeNode {
public:
/*! \brief optionally stores the symbolic value patterns of the shape */
ffi::Optional<ffi::Array<PrimExpr>> values;
/*!
* \brief The number of dimension of the shape, can be unknown.
* \sa kUnknownNDim
*/
int ndim{kUnknownNDim};
/*! \return Whether the type contains unknown ndim. */
bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ShapeTypeNode>()
.def_ro("values", &ShapeTypeNode::values)
.def_ro("ndim", &ShapeTypeNode::ndim);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeType", ShapeTypeNode, TypeNode);
};
/*!
* \brief Managed reference to ShapeTypeNode.
* \sa ShapeTypeNode
*/
class ShapeType : public Type {
public:
/*!
* \brief Construction with known symbolic shape patterns
* \param values The symbolic shape values
* \param span The span of the AST.
*/
TVM_DLL ShapeType(ffi::Array<PrimExpr> values, Span span = Span());
/*!
* \brief Construction with known unknown symbolic shape patterns.
* \param ndim Number of dimensions -- can be kUnknownNDim
* \param span The span of the AST.
*/
TVM_DLL ShapeType(int ndim, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeType, Type, ShapeTypeNode);
};
/*!
* \brief Type of Tensor.
*/
class TensorTypeNode : public TypeNode {
public:
/*!
* \brief optionally store the shape expression of the tensor.
* \note shape must be normalized: it can only be std::nullopt or ShapeExpr or Var.
*/
ffi::Optional<Expr> shape;
/*! \brief The virtual device, indicates where the tensor
* is expected to be executed.
*/
ffi::Optional<VDevice> vdevice;
/*! \brief The content dtype, or nullopt if the dtype is unknown. */
ffi::Optional<tvm::PrimType> dtype{std::nullopt};
/*!
* \brief The number of dimension of the tensor, can be unknown.
* \sa kUnknownNDim
*/
int ndim{kUnknownNDim};
/*! \return Whether the type contains unknown ndim. */
bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
/*! \return Whether the type contains unknown dtype. */
bool IsUnknownDtype() const { return !dtype.has_value(); }
/*! \return Shape if it is known. */
ffi::Optional<ffi::Array<PrimExpr>> GetShape() const {
if (!shape.has_value()) return {};
const Expr& shape_expr = this->shape.value();
if (shape_expr->ty.IsMissing()) return {};
if (const auto* shape_ty = shape_expr->ty.as<ShapeTypeNode>()) {
return shape_ty->values;
}
return {};
}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TensorTypeNode>()
.def_ro("shape", &TensorTypeNode::shape)
.def_ro("dtype", &TensorTypeNode::dtype)
.def_ro("vdevice", &TensorTypeNode::vdevice)
.def_ro("ndim", &TensorTypeNode::ndim);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TensorType", TensorTypeNode, TypeNode);
};
/*!
* \brief Managed reference to TensorTypeNode.
* \sa TensorTypeNode
*/
class TensorType : public Type {
public:
explicit TensorType(ffi::ObjectPtr<TensorTypeNode> data) : Type(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
/*!
* \brief Construction with a known shape expression.
* \param shape The shape of the tensor.
* \param dtype The data type of tensor's elements.
* \param vdevice The virtual device.
* \param span The span of the AST.
*
* \note shape must already be normalized.
*/
TVM_DLL TensorType(Expr shape, ffi::Optional<tvm::PrimType> dtype = std::nullopt,
ffi::Optional<VDevice> vdevice = std::nullopt, Span span = Span());
/*!
* \brief Construction with an unknown shape expression.
* \param dtype The data type of tensor's elements.
* \param ndim The number of dimensions
* \param vdevice The virtual device.
* \param span The span of the AST.
*/
TVM_DLL TensorType(ffi::Optional<tvm::PrimType> dtype, int ndim,
ffi::Optional<VDevice> vdevice = std::nullopt, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeNode);
};
/*!
* \brief custom-defined Type derivation function.
* \param call The call expression to be derived.
* \param ctx The builder context.
* \return The derived type of the call.
*/
using TypeDeriveFunc = TypedEnvFunc<Type(const Call& call, const BlockBuilder& ctx)>;
/*!
* \brief Function type information.
*
* This data structure contains enough information for us to do best-effort
* type deduction.
*/
class FuncTypeNode : public TypeNode {
public:
/*!
* \brief The parameter type of the function.
* \note When params is std::nullopt means the function can take arbitrary number of arguments.
* We define such functions as Opaque function.
*/
ffi::Optional<ffi::Array<Type>> params;
/*!
* \brief The type of the function's return value.
*/
Type ret = Type::Missing();
/*!
* \brief Derivation function of opaque functions that may take any number of parameters.
* \note When derive_func is not empty, then params should be std::nullopt,
* ret should be AnyType()
*/
ffi::Optional<TypeDeriveFunc> derive_func;
/*!
* \brief Whether the function is pure.
* \note This parameter should be set to true only if the function is pure on all inputs.
* If the function _may_ have visible side effects, set it to false.
*/
bool purity;
/*!
* \return Whether the func type is opaque.
* \note We define a function as opaque we have no constraints on params.
*/
bool IsOpaque() const { return !params.has_value(); }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FuncTypeNode>()
.def_ro("params", &FuncTypeNode::params, refl::AttachFieldFlag::SEqHashDefRecursive())
.def_ro("ret", &FuncTypeNode::ret)
.def_ro("derive_func", &FuncTypeNode::derive_func)
.def_ro("purity", &FuncTypeNode::purity);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FuncType", FuncTypeNode, TypeNode);
};
/*!
* \brief Managed reference to FuncTypeNode.
* \sa FuncTypeNode
*/
class FuncType : public Type {
public:
explicit FuncType(ffi::ObjectPtr<FuncTypeNode> data) : Type(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
/*!
* \brief Constructor from parameter type and return value type.
* \param params The type of function parameters.
* \param ret The return value type.
* \param purity The purity of the function (true by default).
* \param span The span of the AST.
*
* \note If the ret contains variables(tirx::Var and relax::Var), they must be deducible from
* params. If you are unsure, you can always erase ret to static.
*/
TVM_DLL FuncType(ffi::Array<Type> params, Type ret, bool purity = true, Span span = Span());
/*!
* \brief Constructing an opaque function type using derive_func.
*
* \param derive_func Derivation function.
* \param purity The purity of the function
* (false by default: most external functions are not pure).
* \param span The span of the AST.
*
* \return The FuncType for opaque packedfunc.
* \note Defaults to an derive func that always return AnyType if not specified.
*/
TVM_DLL static FuncType OpaqueFunc(TypeDeriveFunc derive_func, bool purity = false,
Span span = Span());
/*!
* \brief Construct an opaque function using from return type.
*
* \param ret The type of the return value.
* \param purity The purity of the function
* (false by default: most external functions are not pure).
* \param span The span of the AST.
*
* \return The FuncType for opaque packedfunc.
* \note Defaults to an derive func that always return AnyType if not specified.
*/
TVM_DLL static FuncType OpaqueFunc(Type ret = AnyType(), bool purity = false, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
};
/*!
* \brief Match and check if expr has Relax type T and return it.
*
* \param expr The input expression.
* \return The result of match.
* \tparam T the underlying Relax type
*/
template <typename T>
inline ffi::Optional<T> MatchType(const Expr& expr) {
if (!expr.defined()) {
return std::nullopt;
}
using TNode = typename T::ContainerType;
if (const TNode* ptr = expr->ty.as<TNode>()) {
return ffi::GetRef<T>(ptr);
} else {
return std::nullopt;
}
}
/*!
* \brief Get the type of a given expr and try to cast it as const T*.
*
* \param expr The input expression.
* \return The pointer. Returns nullptr if the type does not match.
* \tparam T the underlying Relax type node
*/
template <typename T>
inline const T* GetTypeAs(const Expr& expr) {
TVM_FFI_ICHECK(!expr->ty.IsMissing())
<< "The type is not populated, check if you have normalized the expr";
return expr->ty.as<T>();
}
/*!
* \brief Get the underlying Relax type of expr.
*
* \param expr The input expression.
* \return underlying Relax type.
*/
inline Type GetType(const Expr& expr) {
TVM_FFI_ICHECK(!expr->ty.IsMissing())
<< "The type is not populated, check if you have normalized the expr";
return expr->ty;
}
/*!
* \brief Whether the expr has void type.
*
* \param expr The input expression.
* \return Whether the expr has void type.
*/
inline bool HasVoidType(const Expr& expr) {
auto* ptr = expr->ty.as<TupleTypeNode>();
return ptr != nullptr && ptr->fields.size() == 0;
}
/*!
* \brief Update the type of an Expr.
* \param expr The Expr whose type to be updated.
* \param ty The type assigned.
* \note We ensure idempotence, that is we can only update the type of an Expr only
* if the original one is nullptr.
*/
TVM_DLL void UpdateType(Expr expr, Type ty);
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_TYPE_H_
include/tvm/relax/type_functor.h
+154
View File
@@ -0,0 +1,154 @@
/*
* 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/relax/type_functor.h
* \brief Functors and visitors for Relax type nodes.
*/
#ifndef TVM_RELAX_TYPE_FUNCTOR_H_
#define TVM_RELAX_TYPE_FUNCTOR_H_
#include <tvm/ir/node_functor.h>
#include <tvm/relax/distributed/type.h>
#include <tvm/relax/type.h>
#include <utility>
namespace tvm {
namespace relax {
template <typename FType>
class TypeFunctor;
// functions to be overriden.
#define RELAX_TYPE_FUNCTOR_DEFAULT \
{ \
return VisitTypeDefault_(op, std::forward<Args>(args)...); \
}
#define TVM_RELAX_TYPE_FUNCTOR_DISPATCH(OP) \
vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \
return self->VisitType_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
});
template <typename R, typename... Args>
class TypeFunctor<R(const Type& n, Args...)> {
private:
using TSelf = TypeFunctor<R(const Type& n, Args...)>;
using FType = tvm::NodeFunctor<R(const ffi::ObjectRef& n, TSelf* self, Args...)>;
public:
/*! \brief the result type of this functor */
using result_type = R;
/*! \brief virtual destructor */
virtual ~TypeFunctor() {}
/*!
* \brief Same as call.
* \param n The type node.
* \param args Additional arguments.
* \return The result of the call
*/
R operator()(const Type& n, Args... args) { return VisitType(n, std::forward<Args>(args)...); }
/*!
* \brief The functor call.
* \param n The type node.
* \param args Additional arguments.
* \return The result of the call
*/
virtual R VisitType(const Type& n, Args... args) {
TVM_FFI_ICHECK(n.defined());
TVM_FFI_ICHECK_NE(n->type_index(), TypeNode::RuntimeTypeIndex())
<< "TypeFunctor cannot visit Type::Missing()";
static FType vtable = InitVTable();
return vtable(n, this, std::forward<Args>(args)...);
}
// Functions that can be overriden by subclass
virtual R VisitType_(const AnyTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitType_(const PrimTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitType_(const ShapeTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitType_(const TensorTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitType_(const distributed::DTensorTypeNode* op,
Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitType_(const TupleTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitType_(const FuncTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitTypeDefault_(const ffi::Object* op, Args...) {
TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
throw; // unreachable, written to stop compiler warning
}
private:
// initialize the vtable.
static FType InitVTable() {
FType vtable;
// Set dispatch
TVM_RELAX_TYPE_FUNCTOR_DISPATCH(AnyTypeNode);
TVM_RELAX_TYPE_FUNCTOR_DISPATCH(PrimTypeNode);
TVM_RELAX_TYPE_FUNCTOR_DISPATCH(ShapeTypeNode);
TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TensorTypeNode);
TVM_RELAX_TYPE_FUNCTOR_DISPATCH(distributed::DTensorTypeNode);
TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TupleTypeNode);
TVM_RELAX_TYPE_FUNCTOR_DISPATCH(FuncTypeNode);
vtable.Finalize();
return vtable;
}
};
#undef TVM_RELAX_TYPE_FUNCTOR_DISPATCH
/*!
* \brief A type visitor.
*/
class TVM_DLL TypeVisitor : public TypeFunctor<void(const Type& n)> {
public:
void VisitType_(const AnyTypeNode* op) override;
void VisitType_(const PrimTypeNode* op) override;
void VisitType_(const ShapeTypeNode* op) override;
void VisitType_(const TensorTypeNode* op) override;
void VisitType_(const distributed::DTensorTypeNode* op) override;
void VisitType_(const TupleTypeNode* op) override;
void VisitType_(const FuncTypeNode* op) override;
protected:
// two functions to override when visit expr fields in type nodes.
virtual void VisitTypeExprField(const Expr& expr) {}
virtual void VisitTypeExprField(const PrimExpr& expr) {}
};
/*!
* \brief TypeMutator that mutates Relax type nodes.
*/
class TVM_DLL TypeMutator : public TypeFunctor<Type(const Type& n)> {
public:
Type VisitType_(const AnyTypeNode* op) override;
Type VisitType_(const PrimTypeNode* op) override;
Type VisitType_(const ShapeTypeNode* op) override;
Type VisitType_(const TensorTypeNode* op) override;
Type VisitType_(const distributed::DTensorTypeNode* op) override;
Type VisitType_(const TupleTypeNode* op) override;
Type VisitType_(const FuncTypeNode* op) override;
protected:
// two functions to override when visit expr fields in type nodes.
virtual Expr VisitTypeExprField(const Expr& expr) { return expr; }
virtual PrimExpr VisitTypeExprField(const PrimExpr& expr) { return expr; }
};
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_TYPE_FUNCTOR_H_
include/tvm/relax/utils.h
+149
View File
@@ -0,0 +1,149 @@
/*
* 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/relax/utils.h
* \brief Utility classes and functions for working with the Relax IR.
*/
#ifndef TVM_RELAX_UTILS_H_
#define TVM_RELAX_UTILS_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/error.h>
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
namespace tvm {
namespace relax {
/*!
* \brief Bind the variables to a Relax expression. This is a helper
* function usually called by other pass functions to help optimizations.
* If any free variables are introduced into a function, those are added
* to the function parameters.
* Additionally this may change the order of parameters if you map a variable
* to a variable.
*
* \param expr The input expression.
* \param binds The variable to expression map that will be used to help the
* binding.
* \param symbolic_var_map The map from symbolic var to the expr it binds to.
*
* \return The updated expression.
*/
TVM_DLL Expr Bind(const Expr& expr, const tvm::ffi::Map<Var, Expr>& binds,
const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map = {});
/*!
* \brief Bind the symbolic variables to a Type. This is a helper function usually called by
* other pass functions to help optimizations.
*/
TVM_DLL Type Bind(const Type& ty, const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map);
/*!
* \brief Infer a binding map for symbolic variables
*
* If a set of relax variables are replaced within an expression, this
* may result in removal of the definition site of a symbolic
* variable. This utility function determines the symbolic variable
* replacements that can be inferred based on the replaced relax
* variables, and can be used alongside the `Bind` utility function to
* replace both the relax variables and the implied symbolic
* variables.
*
* \param binds A map of relax variables to relax expressions
*
* \param analyzer The analyzer to use for simplifications
*
* \return A map of TIR variables to TIR expressions
*/
TVM_DLL tvm::ffi::Map<tirx::Var, PrimExpr> InferSymbolicVarMap(
const tvm::ffi::Map<relax::Var, relax::Expr>& binds, const arith::Analyzer& analyzer);
/*!
* \brief Check if the given Type is for a boolean scalar (tensor of rank 0 with a boolean
* dtype).
*
* \param ty The input Type.
* \param permit_unknown_rank If true, it will permit the input type to have unknown rank
* (ndim of -1), which will require a dynamic check.
* \param permit_unknown_dtype If true, it will permit the input type to have an unknown dtype
* (namely, void), which will require a dynamic check.
*
* \return True iff the input type is a boolean scalar type (or, depending on options, has unknown
* rank or dtype)
*/
TVM_DLL bool IsBoolType(const Type& ty, bool permit_unknown_rank = true,
bool permit_unknown_dtype = true);
/*!
* \brief Check if the given expression is a "leaf" node or tuple node for normalization purposes.
*
* The following expressions are defined as leaf nodes: Var, Constant, ShapeExpr,
* GlobalVar, Op, ExternFunc.
*
* Tuples are included in this list mainly for convenience in grouping operator arguments.
* *Note*: Since tuples can contain nested expressions, it is necessary to ensure that
* values nested inside them are also leaves.
*
* \param expr The input expression
*
* \return True iff the input expression is a "leaf" node (a value allowed to appear
* inline without being bound to a var during normalization).
*/
TVM_DLL bool IsLeafOrTuple(const Expr& expr);
/*!
* \brief Check if the given Call node is an impure operation. If the callee is a general
* expression, this simply requires checking the purity field of the FuncType. If it is an Op,
* then this checks the `fPurity` field.
*
* \param call The input call
*
* \return True iff the call is impure (definitely or possibly results in a visible side effect).
* That is, a call is considered pure only if definitely does not result in a visible side effect.
*/
TVM_DLL bool IsImpureCall(const Call& call);
/*!
* \brief Copy the given function. All variables that are bound inside the original function
* would be copied to satisfy the restriction in the well-formed check: Variables in
* Relax must be bound exactly once. This also ensures that both the function and its copy
* can be inserted into the same IRModule, and be asserted on the structural equality
* agaisnt IRModule created by TVMScript.
*
* \param func The relax function to copy.
* \return The copied function.
*/
TVM_DLL Function CopyWithNewVars(Function func);
/*!
* \brief Transform all dataflow structure to non-dataflow version.
*/
Expr ToNonDataflow(const Expr& e);
/*!
* \brief Get the value bound in the binding.
*/
Expr GetBoundValue(const Binding& b);
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_UTILS_H_
include/tvm/runtime/base.h
+96
View File
@@ -0,0 +1,96 @@
/*
* 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/runtime/base.h
* \brief base macros
*/
#ifndef TVM_RUNTIME_BASE_H_
#define TVM_RUNTIME_BASE_H_
// TVM runtime fully relies on TVM FFI C API
// we will avoid defining extra C APIs here
#include <tvm/ffi/c_api.h>
// TVM version. Overridable at build time via -DTVM_VERSION="..." (scikit-build-core
// passes the setuptools_scm-resolved version through CMake). The literal below is the
// fallback for a bare build with no override.
#ifndef TVM_VERSION
#define TVM_VERSION "0.26.dev0"
#endif
// TVM ships two shared libraries: libtvm_compiler and libtvm_runtime.
// Each exposes its own DLL macro pair. The two families are defined
// independently so that each can be overridden separately by downstream
// embedders who need custom visibility on only one of the two libraries.
//
// TVM_DLL / TVM_DLL_EXPORT: symbols in libtvm_compiler.
// - TVM_DLL is dllexport when TVM_EXPORTS is defined (compiler build),
// dllimport otherwise (downstream consumers, runtime TUs).
// - TVM_DLL_EXPORT is always dllexport.
//
// TVM_RUNTIME_DLL / TVM_RUNTIME_DLL_EXPORT: symbols in libtvm_runtime.
// - TVM_RUNTIME_DLL is dllexport when TVM_RUNTIME_EXPORTS is defined
// (runtime build), dllimport otherwise.
// - TVM_RUNTIME_DLL_EXPORT is always dllexport.
//
// On non-MSVC platforms the import/export decision is made by the dynamic
// loader, so all four macros expand to visibility("default"). Under
// Emscripten they expand to EMSCRIPTEN_KEEPALIVE.
#ifdef __EMSCRIPTEN__
#include <emscripten/emscripten.h>
#endif
// --- TVM_DLL family (libtvm_compiler) ---
#if !defined(TVM_DLL) && defined(__EMSCRIPTEN__)
#define TVM_DLL EMSCRIPTEN_KEEPALIVE
#define TVM_DLL_EXPORT EMSCRIPTEN_KEEPALIVE
#endif
#if !defined(TVM_DLL) && defined(_MSC_VER)
#ifdef TVM_EXPORTS
#define TVM_DLL __declspec(dllexport)
#else
#define TVM_DLL __declspec(dllimport)
#endif
#define TVM_DLL_EXPORT __declspec(dllexport)
#endif
#ifndef TVM_DLL
#define TVM_DLL __attribute__((visibility("default")))
#define TVM_DLL_EXPORT __attribute__((visibility("default")))
#endif
// --- TVM_RUNTIME_DLL family (libtvm_runtime) ---
#if !defined(TVM_RUNTIME_DLL) && defined(__EMSCRIPTEN__)
#define TVM_RUNTIME_DLL EMSCRIPTEN_KEEPALIVE
#define TVM_RUNTIME_DLL_EXPORT EMSCRIPTEN_KEEPALIVE
#endif
#if !defined(TVM_RUNTIME_DLL) && defined(_MSC_VER)
#ifdef TVM_RUNTIME_EXPORTS
#define TVM_RUNTIME_DLL __declspec(dllexport)
#else
#define TVM_RUNTIME_DLL __declspec(dllimport)
#endif
#define TVM_RUNTIME_DLL_EXPORT __declspec(dllexport)
#endif
#ifndef TVM_RUNTIME_DLL
#define TVM_RUNTIME_DLL __attribute__((visibility("default")))
#define TVM_RUNTIME_DLL_EXPORT __attribute__((visibility("default")))
#endif
#endif // TVM_RUNTIME_BASE_H_
include/tvm/runtime/c_backend_api.h
+109
View File
@@ -0,0 +1,109 @@
/*
* 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/runtime/c_backend_api.h
* \brief TVM runtime backend API.
*
* The functions defined in this header are intended to be
* used by compiled tvm operators, usually user do not need to use these
* function directly.
*/
#ifndef TVM_RUNTIME_C_BACKEND_API_H_
#define TVM_RUNTIME_C_BACKEND_API_H_
#include <tvm/runtime/base.h>
#ifdef __cplusplus
extern "C" {
#endif
/*!
* \brief Backend function to allocate temporal workspace.
*
* \note The result allocated space is ensured to be aligned to kTempAllocaAlignment.
*
* \param nbytes The size of the space requested.
* \param device_type The device type which the space will be allocated.
* \param device_id The device id which the space will be allocated.
* \param dtype_code_hint The type code of the array elements. Only used in
* certain backends such as OpenGL.
* \param dtype_bits_hint The type bits of the array elements. Only used in
* certain backends such as OpenGL.
* \return nullptr when error is thrown, a valid ptr if success
*/
TVM_RUNTIME_DLL void* TVMBackendAllocWorkspace(int device_type, int device_id, uint64_t nbytes,
int dtype_code_hint, int dtype_bits_hint);
/*!
* \brief Backend function to free temporal workspace.
*
* \param ptr The result allocated space pointer.
* \param device_type The device type which the space will be allocated.
* \param device_id The device id which the space will be allocated.
* \return 0 when no error is thrown, -1 when failure happens
*
* \sa TVMBackendAllocWorkspace
*/
TVM_RUNTIME_DLL int TVMBackendFreeWorkspace(int device_type, int device_id, void* ptr);
/*!
* \brief Environment for TVM parallel task.
*/
typedef struct {
/*!
* \brief Auxiliary used for synchronization
*/
void* sync_handle;
/*! \brief total amount of task */
int32_t num_task;
} TVMParallelGroupEnv;
/*!
* \brief The callback function to execute a parallel lambda
* \param task_id the task id of the function.
* \param penv The parallel environment backs the execution.
* \param cdata The supporting closure data.
*/
typedef int (*FTVMParallelLambda)(int task_id, TVMParallelGroupEnv* penv, void* cdata);
/*!
* \brief Backend function for running parallel jobs.
*
* \param flambda The parallel function to be launched.
* \param cdata The closure data.
* \param num_task Number of tasks to launch, can be 0, means launch
* with all available threads.
*
* \return 0 when no error is thrown, -1 when failure happens
*/
TVM_RUNTIME_DLL int TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* cdata, int num_task);
/*!
* \brief BSP barrrier between parallel threads
* \param task_id the task id of the function.
* \param penv The parallel environment backs the execution.
* \return 0 when no error is thrown, -1 when failure happens
*/
TVM_RUNTIME_DLL int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv);
#ifdef __cplusplus
} // TVM_EXTERN_C
#endif
#endif // TVM_RUNTIME_C_BACKEND_API_H_
include/tvm/runtime/device_api.h
+427
View File
@@ -0,0 +1,427 @@
/*
* 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/runtime/device_api.h
* \brief Abstract device memory management API
*/
#ifndef TVM_RUNTIME_DEVICE_API_H_
#define TVM_RUNTIME_DEVICE_API_H_
#include <tvm/ffi/any.h>
#include <tvm/ffi/error.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/base.h>
#include <string>
/*!
* \brief The stream that is specific to device
* can be NULL, which indicates the default one.
*/
typedef void* TVMStreamHandle;
namespace tvm {
// alias DLDevice
using Device = DLDevice;
namespace runtime {
/*! \brief Extension device types in TVM
*
* Additional enumerators to supplement those provided by
* DLPack's `DLDeviceType` enumeration.
*
* MAINTAINERS NOTE #1: We need to ensure that the two devices
* are identified by the same integer.
* Currently this requires manual verification.
* Discussed here: https://github.com/dmlc/dlpack/issues/111
* As of DLPack v0.7, the highest-valued enumerator in
* `DLDeviceType` is kDLHexagon = 16.
*
* MAINTAINERS NOTE #2: As of DLPack v0.7, the definition for
* `DLDeviceType` specifies an underlying storage type of
* `int32_t`. That guarantees a variable of type
* `DLDeviceType` is capable of holding any integers provided
* by *either* of these enumerations.
*
* However, the `int32_t` specification only applies when the
* header file is compiled as C++, and this header file is also
* meant to work as C code. So the unspecified storage type
* could be a latent bug when compiled as C.
*/
#ifdef __cplusplus
typedef enum : int32_t {
#else
typedef enum {
#endif
// To help avoid accidental conflicts between `DLDeviceType`
// and this enumeration, start numbering the new enumerators
// a little higher than (currently) seems necessary.
TVMDeviceExtType_End = 36, // sentinel value
} TVMDeviceExtType;
/*!
* \brief the query type into GetAttr
*/
enum DeviceAttrKind : int {
kExist = 0,
kMaxThreadsPerBlock = 1,
kWarpSize = 2,
kMaxSharedMemoryPerBlock = 3,
kComputeVersion = 4,
kDeviceName = 5,
kMaxClockRate = 6,
kMultiProcessorCount = 7,
kMaxThreadDimensions = 8,
kMaxRegistersPerBlock = 9,
kGcnArch = 10,
kApiVersion = 11,
kDriverVersion = 12,
kL2CacheSizeBytes = 13,
kTotalGlobalMemory = 14,
kAvailableGlobalMemory = 15,
kImagePitchAlignment = 16,
};
#ifdef TVM_KALLOC_ALIGNMENT
/*! \brief Number of bytes each allocation must align to */
constexpr int kAllocAlignment = TVM_KALLOC_ALIGNMENT;
/*! \brief Number of bytes each allocation must align to in temporary allocation */
constexpr int kTempAllocaAlignment = TVM_KALLOC_ALIGNMENT;
#else
/*! \brief Number of bytes each allocation must align to */
constexpr int kAllocAlignment = 64;
/*! \brief Number of bytes each allocation must align to in temporary allocation */
constexpr int kTempAllocaAlignment = 64;
#endif // TVM_KALLOC_ALIGNMENT
/*! \brief Maximum size that can be allocated on stack */
constexpr int kMaxStackAlloca = 1024;
/*! \brief Number of bytes each allocation must align to by default in the workspace buffer to
* service intermediate tensors */
constexpr int kDefaultWorkspaceAlignment = 1;
/*!
* \brief TVM Runtime Device API, abstracts the device
* specific interface for memory management.
*/
class TVM_RUNTIME_DLL DeviceAPI {
public:
/*! \brief virtual destructor */
virtual ~DeviceAPI() {}
/*!
* \brief Set the environment device id to device
* \param dev The device to be set.
*/
virtual void SetDevice(Device dev) = 0;
/*!
* \brief Get attribute of specified device.
* \param dev The device device
* \param kind The result kind
* \param rv The return value.
* \sa DeviceAttrKind
*/
virtual void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) = 0;
/*!
* \brief Get the physical memory size required.
* \param arr the tensor object.
* \param mem_scope the memory scope if any
* \return the memory size.
*/
virtual size_t GetDataSize(const DLTensor& arr,
ffi::Optional<ffi::String> mem_scope = std::nullopt);
/*!
* \brief Query the device for specified properties.
*
* This is used to expand "-from_device=N" in the target string to
* all properties that can be determined from that device.
*/
virtual void GetTargetProperty(Device dev, const std::string& property, ffi::Any* rv) {}
/*!
* \brief Allocate a data space on device.
* \param dev The device device to perform operation.
* \param nbytes The number of bytes in memory.
* \param alignment The alignment of the memory.
* \param type_hint The type of elements. Only needed by certain backends such
* as OpenGL, as nbytes & alignment are sufficient for most backends.
* \return The allocated device pointer.
*/
virtual void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment,
DLDataType type_hint) = 0;
/*!
* \brief Allocate a data space on device with memory scope support.
* \param dev The device device to perform operation.
* \param ndim The number of dimension of allocated tensor.
* \param shape The shape of allocated tensor.
* \param dtype The type of elements.
* \param mem_scope The memory scope of allocated tensor.
* \return The allocated device pointer.
*/
virtual void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype,
ffi::Optional<ffi::String> mem_scope = std::nullopt);
/*!
* \brief Free a data space on device.
* \param dev The device device to perform operation.
* \param ptr The data space.
*/
virtual void FreeDataSpace(Device dev, void* ptr) = 0;
/*!
* \brief copy data from one place to another
* \note This API is designed to support special memory with shape dependent layout.
* We pass in DLTensor* with shape information to support these cases.
* \param from The source array.
* \param to The target array.
* \param stream Optional stream object.
* \note The copy may happen asynchronously if it involves a GPU context.
* Call StreamSync to ensure the copy completes from host's pov.
*/
virtual void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream);
/*!
* \brief Create a new stream of execution.
*
* \param dev The device of allocation.
*/
virtual TVMStreamHandle CreateStream(Device dev);
/*!
* \brief Free a stream of execution
*
* \param dev The device of the stream
* \param stream The pointer to be freed.
*/
virtual void FreeStream(Device dev, TVMStreamHandle stream);
/*!
* \brief Synchronize the stream
* \param dev The device to perform operation.
* \param stream The stream to be sync.
*/
virtual void StreamSync(Device dev, TVMStreamHandle stream) = 0;
/*!
* \brief Set the stream
* \param dev The device to set stream.
* \param stream The stream to be set.
*/
virtual void SetStream(Device dev, TVMStreamHandle stream);
/*!
* \brief Get the current stream
* \param dev The device to get stream.
* \return The current stream of the device.
*/
virtual TVMStreamHandle GetCurrentStream(Device dev);
/*!
* \brief Synchronize 2 streams of execution.
*
* An event is created in event_src stream that the second then
* stream waits on. Neither event_src or event_dst need to be of
* the same device ID as the device, but they must be of the same
* device type.
*
* \param dev The device of the streams.
* \param event_src The source stream to synchronize.
* \param event_dst The destination stream to synchronize.
*/
virtual void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst);
/*!
* \brief Allocate temporal workspace for backend execution.
*
* \note We have the following assumption about backend temporal
* workspace allocation, and backend will optimize for such assumption:
*
* - Only a few allocation will happen, and space will be released after use.
* - The release order is usually in reverse order of allocate (stack style).
* - Repeative pattern of same allocations over different runs.
* - Workspace should not overlap between different threads(i.e. be threadlocal)
*
* \param dev The device of allocation.
* \param nbytes The size to be allocated.
* \param type_hint The type of elements. Only needed by certain backends such
* as OpenGL, as nbytes is sufficient for most backends.
*/
virtual void* AllocWorkspace(Device dev, size_t nbytes, DLDataType type_hint = {});
/*!
* \brief Free temporal workspace in backend execution.
*
* \param dev The device of allocation.
* \param ptr The pointer to be freed.
*/
virtual void FreeWorkspace(Device dev, void* ptr);
/*!
* \brief Get device API based on device.
* \param dev The device
* \param allow_missing Whether allow missing
* \return The corresponding device API.
*/
static DeviceAPI* Get(Device dev, bool allow_missing = false);
/*!
* \brief Whether a certian device type requires set device device
* before launching the kernel function.
* \param device_type The device type.
*/
static bool NeedSetDevice(int device_type) { return device_type != kDLCPU; }
/*!
* \brief Whether pointer arithmetics on a device owned pointer may be performed on the host.
*/
virtual bool SupportsDevicePointerArithmeticsOnHost() { return false; }
protected:
/*!
* \brief copy data from one place to another
* \param from The source array.
* \param from_offset The byte offeset in the from.
* \param to The target array.
* \param to_offset The byte offset in the to.
* \param num_bytes The size of the memory in bytes
* \param dev_from The source device
* \param dev_to The target device
* \param type_hint The type of elements, only neded by certain backends.
* can be useful for cross device endian converison.
* \param stream Optional stream object.
*/
virtual void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset,
size_t num_bytes, Device dev_from, Device dev_to,
DLDataType type_hint, TVMStreamHandle stream);
};
/*!
* \brief The name of DLDeviceType.
* \param type The device type.
* \return the device name.
*/
inline const char* DLDeviceType2Str(int type) {
switch (type) {
case kDLCPU:
return "cpu";
case kDLCUDA:
return "cuda";
case kDLCUDAHost:
return "cuda_host";
case kDLCUDAManaged:
return "cuda_managed";
case kDLOpenCL:
return "opencl";
case kDLVulkan:
return "vulkan";
case kDLMetal:
return "metal";
case kDLVPI:
return "vpi";
case kDLROCM:
return "rocm";
case kDLROCMHost:
return "rocm_host";
case kDLExtDev:
return "ext_dev";
case kDLOneAPI:
return "oneapi";
case kDLWebGPU:
return "webgpu";
case kDLHexagon:
return "hexagon";
case kDLTrn:
return "trn";
default:
TVM_FFI_THROW(InternalError) << "unknown type = " << type;
}
throw;
}
/*! \brief The device type bigger than this is RPC device */
constexpr int kRPCSessMask = 128;
static_assert(kRPCSessMask >= TVMDeviceExtType_End);
/*!
* \brief Return true if a Device is owned by an RPC session.
*/
inline bool IsRPCSessionDevice(Device dev) { return (dev.device_type / kRPCSessMask) > 0; }
/*!
* \brief Return the RPCSessTable index of the RPC Session that owns this device.
* \return the table index.
*/
inline int GetRPCSessionIndex(Device dev) {
TVM_FFI_ICHECK(IsRPCSessionDevice(dev)) << "GetRPCSessionIndex: dev has no RPC session";
return dev.device_type / kRPCSessMask - 1;
}
/*!
* \brief Remove the RPC session mask from a Device.
* RPC clients typically do this when encoding a Device for transmission to an RPC remote.
* On the wire, RPCdevice are expected to be valid on the server without interpretation.
* \param dev A Device with non-zero RPC Session mask, valid on the RPC client.
* \return A Device without any RPC Session mask, valid on the RPC server.
*/
inline Device RemoveRPCSessionMask(Device dev) {
dev.device_type = static_cast<DLDeviceType>(dev.device_type % kRPCSessMask);
return dev;
}
inline std::ostream& operator<<(std::ostream& os, DLDevice dev) { // NOLINT(*)
if (tvm::runtime::IsRPCSessionDevice(dev)) {
os << "remote[" << tvm::runtime::GetRPCSessionIndex(dev) << "]-";
dev = tvm::runtime::RemoveRPCSessionMask(dev);
}
os << tvm::runtime::DLDeviceType2Str(static_cast<int>(dev.device_type)) << ":" << dev.device_id;
return os;
}
/*!
* \brief Add a RPC session mask to a Device.
* RPC clients typically do this when decoding a Device received from a RPC remote.
* \param dev A Device without any RPC Session mask, valid on the RPC server.
* \param session_table_index Numeric index of the RPC session in the session table.
* \return A Device with RPC session mask added, valid on the RPC client.
*/
inline Device AddRPCSessionMask(Device dev, int session_table_index) {
TVM_FFI_ICHECK(!IsRPCSessionDevice(dev))
<< "AddRPCSessionMask: dev already non-zero RPCSessionIndex: " << dev;
dev.device_type =
static_cast<DLDeviceType>(dev.device_type | (kRPCSessMask * (session_table_index + 1)));
return dev;
}
/*!
* \brief Check if runtime module is enabled for target.
* \param target The target module name.
* \return Whether runtime is enabled.
*/
TVM_RUNTIME_DLL bool RuntimeEnabled(const ffi::String& target);
/*! \brief namespace for constant symbols */
namespace symbol {
constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
/*! \brief global function to set device */
constexpr const char* tvm_set_device = "__tvm_set_device";
} // namespace symbol
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DEVICE_API_H_
include/tvm/runtime/disco/builtin.h
+154
View File
@@ -0,0 +1,154 @@
/*
* 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.
*/
#ifndef TVM_RUNTIME_DISCO_BUILTIN_H_
#define TVM_RUNTIME_DISCO_BUILTIN_H_
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/runtime/tensor.h>
#include <string>
namespace tvm {
namespace runtime {
/*!
* \brief Possible kinds of reduction operations.
*/
enum class ReduceKind : int32_t {
kSum = 0,
kProd = 1,
kMin = 2,
kMax = 3,
kAvg = 4,
};
/*! \brief Converts `ReduceKind` to string */
inline std::string ReduceKind2String(ReduceKind kind) {
switch (kind) {
case ReduceKind::kSum:
return "kSum";
case ReduceKind::kProd:
return "kProd";
case ReduceKind::kMin:
return "kMin";
case ReduceKind::kMax:
return "kMax";
case ReduceKind::kAvg:
return "kAvg";
}
TVM_FFI_THROW(ValueError) << "Unknown ReduceKind: " << static_cast<int>(kind);
}
/*!
* \brief Load a runtime Module, then create and initialize a RelaxVM
* \param path The path to the runtime Module (a DSO file) to be loaded
* \param device The default device used to initialize the RelaxVM
* \return The RelaxVM as a runtime Module
*/
TVM_RUNTIME_DLL ffi::Module LoadVMModule(std::string path, ffi::Optional<Device> device);
/*!
* \brief Create an uninitialized empty Tensor
* \param shape The shape of the Tensor
* \param dtype The dtype of the Tensor
* \param device The device the Tensor is created on. If None, use the thread local default device
* \return The Tensor created
*/
TVM_RUNTIME_DLL Tensor DiscoEmptyTensor(ffi::Shape shape, DLDataType dtype,
ffi::Optional<Device> device);
/*!
* \brief Perform an allreduce operation using the underlying communication library
* \param send The array send to perform allreduce on
* \param reduce_kind The kind of reduction operation (e.g. sum, avg, min, max)
* \param in_group Whether the allreduce operation performs globally or in group as default.
* \param recv The array receives the outcome of allreduce
*/
TVM_RUNTIME_DLL void AllReduce(Tensor send, ReduceKind reduce_kind, bool in_group, Tensor recv);
/*!
* \brief Perform an allgather operation using the underlying communication library
* \param send The array send to perform allgather on
* \param in_group Whether the allgather operation performs globally or in group as default.
* \param recv The array receives the outcome of allgather
*/
TVM_RUNTIME_DLL void AllGather(Tensor send, bool in_group, Tensor recv);
/*!
* \brief Perform a broadcast operation from worker-0
* \param send The buffer to be broadcasted
* \param in_group Whether the broadcast operation performs globally or in group as default.
* \param recv The buffer receives the broadcasted array
*/
TVM_RUNTIME_DLL void BroadcastFromWorker0(Tensor send, bool in_group, Tensor recv);
/*!
* \brief Perform a scatter operation from worker-0, chunking the given buffer into equal parts.
* \param send For worker-0, it must be provided, and otherwise, the buffer must be None.
* The buffer will be divided into equal parts and sent to each worker accordingly.
* \param in_group Whether the scatter operation performs globally or in group as default.
* \param recv The receiving buffer, which must not be None.
*/
TVM_RUNTIME_DLL void ScatterFromWorker0(ffi::Optional<Tensor> send, bool in_group, Tensor recv);
/*!
* \brief Perform a gather operation to worker-0.
* \param send The sending buffer, which must not be None.
* \param in_group Whether the gather operation performs globally or in group as default.
* \param recv For worker-0, it must be provided, and otherwise, the buffer must be None. The
* receiving buffer will be divided into equal parts and receive from each worker accordingly.
*/
TVM_RUNTIME_DLL void GatherToWorker0(Tensor send, bool in_group, ffi::Optional<Tensor> recv);
/*!
* \brief Receive a buffer from worker-0. No-op if the current worker is worker-0.
* \param buffer The buffer to be received
*/
TVM_RUNTIME_DLL void RecvFromWorker0(Tensor buffer);
/*!
* \brief Send a buffer to the corresponding worker in the next group.
* An error is thrown if the worker is already in the last group.
* \param buffer The sending buffer.
*/
TVM_RUNTIME_DLL void SendToNextGroup(Tensor buffer);
/*!
* \brief Receive a buffer from the corresponding worker in the previous group.
* An error is thrown if the worker is already in the first group.
* \param buffer The receiving buffer.
*/
TVM_RUNTIME_DLL void RecvFromPrevGroup(Tensor buffer);
/*!
* \brief Send a buffer to the target receiver worker (globally across all groups).
* \param buffer The sending buffer.
* \param receiver_id The global receiver worker id.
*/
TVM_RUNTIME_DLL void SendToWorker(Tensor buffer, int receiver_id);
/*!
* \brief Receive a buffer from the target sender worker (globally across all groups).
* \param buffer The receiving buffer.
* \param sender_id The global sender worker id.
*/
TVM_RUNTIME_DLL void RecvFromWorker(Tensor buffer, int sender_id);
/*! \brief Get the local worker id */
TVM_RUNTIME_DLL int WorkerId();
/*!
* \brief Called by the worker thread. Waiting until the worker completes all its tasks.
* As a specific example, on a CUDA worker, it blocks until all kernels are launched and
* cudaStreamSynchronize is complete.
*/
TVM_RUNTIME_DLL void SyncWorker();
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_BUILTIN_H_
include/tvm/runtime/disco/cuda_ipc_memory.h
+99
View File
@@ -0,0 +1,99 @@
/*
* 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.
*/
#ifndef TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
#define TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
#include <tvm/runtime/base.h>
#include <tvm/runtime/memory/memory_manager.h>
#include <vector>
namespace tvm {
namespace runtime {
namespace cuda_ipc {
/*!
* \brief The CUDA IPC (interprocess communication) memory object,
* which internally contains data pointers to CUDA IPC memory.
* It is be useful for efficient all-reduce implementation.
* \note Right now the class members are closely tied with customized
* all-reduce kernel. They may also be extended for other uses in
* the future.
*/
class CUDAIPCMemoryObj : public ffi::Object {
public:
/*! \brief The number of GPU workers. */
int num_workers;
/*! \brief The worker id corresponding to this IPC memory object. */
int worker_id;
/*!
* \brief The data pointers of all all-reduce inputs.
* It has "num_workers" pointers. The i-th pointer is the data pointer on worker i.
* If "i != worker_id", the pointer is an IPC data pointer.
* Otherwise, the pointer is a local CUDA data pointer.
*/
std::vector<void*> remote_data;
// We introduce the barrier helper data below per CUDAIPCMemory object
// so that they can be used by custom collective operations and allow
// fine-grained synchronization on each buffer. These barriers have
// low overhead, and can potentially enable concurrent execution of
// kernels in future.
/*!
* \brief The pointers to input barrier signals of all workers for all-reduce.
* It has "num_workers" pointers, and the pointer arrangement is the same as "remote_data".
*/
std::vector<void*> barrier_in;
/*!
* \brief The pointers to output barrier signals of all workers for all-reduce.
* It has "num_workers" pointers, and the pointer arrangement is the same as "remote_data".
*/
std::vector<void*> barrier_out;
/*! \brief The integer buffer flag for all-reduce. */
int barrier_flag;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("tvm.runtime.disco.cuda_ipc_memory", CUDAIPCMemoryObj, ffi::Object);
};
/*!
* \brief Managed reference to CUDAIPCMemoryObj.
* \sa CUDAIPCMemory
*/
class CUDAIPCMemory : public ffi::ObjectRef {
public:
/*! \brief Get the global singleton CUDAIPCMemory allocator. */
TVM_RUNTIME_DLL static memory::Allocator* GlobalAllocator();
/*!
* \brief Given a local CUDA data pointer, return the CUDAIPCMemory object of the pointer.
* \note The pointer's CUDAIPCMemory is expected to have been allocated
* through global function "cuda_ipc.alloc_storage". Or otherwise this
* function will raise exception.
*/
TVM_RUNTIME_DLL static CUDAIPCMemory GetIPCMemoryFromDevicePtr(void* ptr);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CUDAIPCMemory, ffi::ObjectRef, CUDAIPCMemoryObj);
};
} // namespace cuda_ipc
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
include/tvm/runtime/disco/disco_worker.h
+118
View File
@@ -0,0 +1,118 @@
/*
* 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 disco_worker.h
* \brief This file defines a worker in Disco. A worker can be launched in a separate thread or
* process as long as the channel supports bi-directional communication in-between the worker and
* the controler.
*/
#ifndef TVM_RUNTIME_DISCO_DISCO_WORKER_H_
#define TVM_RUNTIME_DISCO_DISCO_WORKER_H_
#include <tvm/ffi/function.h>
#include <tvm/runtime/disco/session.h>
#include <vector>
namespace tvm {
namespace runtime {
/*!
* \brief A worker in Disco. It takes a channel to communication with the controler.
* The worker can be run in a separate thread or process as long as the channel supports
* bi-directional communication in-between.
*/
class DiscoWorker {
public:
/*!
* \brief Construct a worker.
* \param worker_id The id of the worker.
* \param num_workers The number of the workers.
* \param num_groups The number of the worker groups.
* \param worker_zero_data The data shared between worker-0 and the controler. It's a nullptr if
* the worker is not worker-0.
* \param channel The communication channel between the worker and the controler.
*/
explicit DiscoWorker(int worker_id, int num_workers, int num_groups,
WorkerZeroData* worker_zero_data, DiscoChannel* channel)
: worker_id(worker_id),
local_worker_id(worker_id),
num_workers(num_workers),
num_groups(num_groups),
default_device(Device{DLDeviceType::kDLCPU, 0}),
worker_zero_data(worker_zero_data),
channel(channel),
register_file{} {}
/*! \brief Main loop of the worker */
void MainLoop();
/*! \brief Get the worker instance on the current thread */
TVM_RUNTIME_DLL static DiscoWorker* ThreadLocal();
/*! \brief Set the specific register to a specific value */
void SetRegister(int reg_id, ffi::AnyView value);
/*! \brief The id of the worker.*/
int worker_id;
/*! \brief The local id of the worker. This can be different from worker_id if the session is
* consisted with multiple sub-sessions. */
int local_worker_id;
/*! \brief Total number of workers */
int num_workers;
/*! \brief Total number of workers */
int num_groups;
/*! \brief The default device to allocate data if not specified */
Device default_device;
/*! \brief The name of the underlying collective communication library. */
ffi::String ccl;
/*!
* \brief The data shared between worker-0 and the controler. It's a nullptr if
* the worker is not worker-0.
* \note This data structure is owned by the controler.
*/
WorkerZeroData* worker_zero_data;
/*!
* \brief The communication channel between the worker and the controler.
* \note This data structure is owned by the controler.
*/
DiscoChannel* channel;
/*! \brief The registers in the worker */
std::vector<ffi::Any> register_file;
struct Impl;
friend struct DiscoWorker::Impl;
};
/*!
* \brief A threadlocal wrapper of DiscoWorker.
*/
struct ThreadLocalDiscoWorker {
/*! \brief The Disco worker */
DiscoWorker* worker;
/*!
* \brief Get the threadlocal Disco worker.
*/
static ThreadLocalDiscoWorker* Get() {
thread_local static ThreadLocalDiscoWorker worker;
return &worker;
}
};
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_DISCO_WORKER_H_
include/tvm/runtime/disco/session.h
+385
View File
@@ -0,0 +1,385 @@
/*
* 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 session.h
* \brief This file serves as the entry point of Disco and defines key data structures and
* interfaces.
*
* Disco is a distributed runtime that consists of a controler and a cluster of workers. The
* controler is responsible for managing the workers by broadcasting commands to all the workers
* together, and the workers are responsible for executing the commands and. The controler and
* workers communicate with each other through a bi-directional channel.
*
* Different from a generic system, Disco is designed to as "single-program-multiple-data" (SPMD)
* runtime, which means that all the workers execute the same instruction at the same time, but the
* data they are working on may be different. For example, in data parallelism, each worker may
* work on a different batches of the data, but they all execute the same set of instructions.
* Therefore, imagine there is a virtual machine that executes the program, the structures of
* workers' register files could be considered as "identical" (single program) although the values
* may differ (multiple data).
*
*
* **DRef.** Following the design above, consider the program in SPMD in a virtual ISA, then each
* worker is a virtual machine instance to execute the ISA maintaining its own register file.
* The controler denotes each of their register files with a unique integer "register id",
* and the workers use this id to refer to the register file that resides on itself.
* DRef is a control-side object backed by such a register id. The data it contains is not assumed
* to be directly accessible by the controler, with an exception for worker-0, which is a special
* worker that is always co-located with the controler.
*
* **Worker-0.** Worker-0 is a special worker that is always co-located with the controler.
* It is assumed that the controler can synchronize with and access the registers of worker-0.
* The Disco session provides multiple APIs to interact specifically with the worker-0.
* To shared data with other workers, a common paradigm in Disco is to copy data from the
* controler-side Tensor to the worker-0, and then copy it to other workers using primitives on
* the data plane, for example, `broadcast` and `send`.
*
* **Control plane.** The controler broadcasts commands to all the workers as control signals.
* For example, the control may ask all workers to load a library or call a function respectively.
* Common control signals include: shutdown, retrievel a global ffi::Function, call packed function,
* etc. The controler is assumed to keep a message channel to each worker to implement the broadcast
* behavior, and the message channel may vary depends on usecases.
*
* **Data plane.** The data channel is usually used to exchange data between workers, especially for
* tensor data which is usually large. For example, performing an allreduce operator for sharded
* matrix multiplication, or broadcasting for an input tensor. For efficiency, the data channel is
* usually backed by NCCL on NVIDIA GPUs, RCCL on AMD GPUs, or MPI on CPUs.
*
* **Session.** A Disco session is a primary interface to interact with the Disco runtime, serving
* as a global context that manages the control and workers. It could be implemented as a
* multi-threaded with a pool of workers for single-node multi-gpu scenarios, or TCP sockets for
* workloads that span over a cluster of nodes.
*
* **Channel.** Disco channel is a bi-directional communication channel between the controler and
* workers for exchanging control signals. It is no different from a generic RPC channel, but
* adopts TVM's ffi::Function calling convention to support polymorphic and variadic arguments.
*/
#ifndef TVM_RUNTIME_DISCO_SESSION_H_
#define TVM_RUNTIME_DISCO_SESSION_H_
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/tensor.h>
#include <mutex>
#include <queue>
#include <string>
#include <utility>
namespace tvm {
namespace runtime {
/*!
* \brief Static FFI type index for `runtime::disco::DRef`.
*
* Allocated within the [kTVMFFIDynObjectBegin - 16, kTVMFFIDynObjectBegin)
* custom-static slot range. The sibling constant `kRuntimeRPCObjectRef`
* lives in `src/runtime/rpc/rpc_session.h` and uses `... - 13`; values must
* remain disjoint across this small reserved block.
*/
constexpr int32_t kRuntimeDiscoDRef = TVMFFITypeIndex::kTVMFFIDynObjectBegin - 14;
static_assert(kRuntimeDiscoDRef >= TVMFFITypeIndex::kTVMFFIStaticObjectEnd &&
kRuntimeDiscoDRef < TVMFFITypeIndex::kTVMFFIDynObjectBegin,
"kRuntimeDiscoDRef must live in the static custom-index slot range");
/*!
* \brief All possible kinds of Disco commands.
*/
enum class DiscoAction : int32_t {
kShutDown = 0,
kKillReg = 1,
kGetGlobalFunc = 2,
kCallPacked = 3,
kSyncWorker = 4,
kCopyFromWorker0 = 5,
kCopyToWorker0 = 6,
kDebugGetFromRemote = 7,
kDebugSetRegister = 8,
};
/*! \brief Converts the enum class `DiscoAction` to string */
inline std::string DiscoAction2String(DiscoAction action) {
switch (action) {
case DiscoAction::kShutDown:
return "kShutDown";
case DiscoAction::kKillReg:
return "kKillReg";
case DiscoAction::kGetGlobalFunc:
return "kGetGlobalFunc";
case DiscoAction::kCallPacked:
return "kCallPacked";
case DiscoAction::kSyncWorker:
return "kSyncWorker";
case DiscoAction::kCopyFromWorker0:
return "kCopyFromWorker0";
case DiscoAction::kCopyToWorker0:
return "kCopyToWorker0";
case DiscoAction::kDebugGetFromRemote:
return "kDebugGetFromRemote";
case DiscoAction::kDebugSetRegister:
return "kDebugSetRegister";
}
TVM_FFI_THROW(ValueError) << "Unknown DiscoAction: " << static_cast<int>(action);
}
class SessionObj;
/*!
* \brief An object that exists on all workers.
*
* The controler assigns a unique "register id" to each object, and the worker uses this id to
* refer to the object residing on itself.
*/
class DRefObj : public ffi::Object {
public:
/*!\ brief Send dellocation command for `reg_id` */
inline ~DRefObj();
/*!
* \brief Get the value of a DRef from a remote worker.
* \param worker_id The id of the worker to be fetched from.
* \return The value of the register.
*/
inline ffi::Any DebugGetFromRemote(int worker_id);
/*!
* \brief Copy from the Tensor provided to a remote worker.
* \param worker_id The id of the worker to be copied to.
* \param source The Tensor to be copied.
*/
inline void DebugCopyFrom(int worker_id, ffi::AnyView source);
static constexpr const uint32_t _type_index = kRuntimeDiscoDRef;
static const constexpr bool _type_final = true;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_STATIC("runtime.disco.DRef", DRefObj, ffi::Object);
/*! \brief The id of the register */
int64_t reg_id;
/*! \brief Back-pointer to the host controler session */
ffi::ObjectRef session{nullptr};
private:
inline SessionObj* GetSession();
};
/*!
* \brief Managed reference to DRefObj.
* \sa DRefObj
* \note No public constructor is provided as it is not supposed to be directly created by users.
*/
class DRef : public ffi::ObjectRef {
public:
explicit DRef(ffi::ObjectPtr<DRefObj> data) : ffi::ObjectRef(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DRef, ffi::ObjectRef, DRefObj);
};
/*!
* \brief A Disco interactive session. It allows users to interact with the Disco command queue with
* various ffi::Function calling convention.
*/
class SessionObj : public ffi::Object {
public:
virtual ~SessionObj() = default;
/*!
* \brief Call a ffi::Function on workers providing variadic arguments.
* \tparam Args In the variadic arguments, the supported types include:
* - integers and floating point numbers;
* - DataType;
* - Device;
* - std::string;
* - DRef.
* Examples of unsupported types:
* - Tensor, DLTensor;
* - TVM Objects, including ffi::Function, Module and String;
* \param func The function to be called.
* \param args The variadic arguments.
* \return The return value of function call
*/
template <typename... Args>
TVM_FFI_INLINE DRef CallPacked(const DRef& func, Args&&... args);
/*!
* \brief Call packed function on each worker using a packed sequence. The calling convention:
* The first element must be DiscoAction::kCallPacked,
* The second element must be 0, which will later be updated by the session to return reg_id
* The thirtd element is the function to be called.
*/
TVM_RUNTIME_DLL virtual DRef CallWithPacked(const ffi::PackedArgs& args) = 0;
/*! \brief Get the number of workers in the session. */
TVM_RUNTIME_DLL virtual int64_t GetNumWorkers() = 0;
/*! \brief Get a global functions on workers. */
TVM_RUNTIME_DLL virtual DRef GetGlobalFunc(const std::string& name) = 0;
/*!
* \brief Copy an Tensor from worker-0 to the controler-side Tensor
* \param host_array The array to be copied to worker-0
* \param remote_array The Tensor on worker-0
*/
TVM_RUNTIME_DLL virtual void CopyFromWorker0(const Tensor& host_array,
const DRef& remote_array) = 0;
/*!
* \brief Copy the controler-side Tensor to worker-0
* \param host_array The array to be copied to worker-0
* \param remote_array The Tensor on worker-0
*/
TVM_RUNTIME_DLL virtual void CopyToWorker0(const Tensor& host_array,
const DRef& remote_array) = 0;
/*!
* \brief Synchrnoize the controler with a worker, and it will wait until worker finishes
* executing this instruction.
* \param worker_id The id of the worker to be synced with.
* \note This function is usually used for worker-0, because it is the only worker that is
* assumed to collocate with the controler. Syncing with other workers may not be supported.
*/
TVM_RUNTIME_DLL virtual void SyncWorker(int worker_id) = 0;
/*! \brief Signal all the workers to shutdown */
TVM_RUNTIME_DLL virtual void Shutdown() = 0;
/*!
* \brief Initialize the data plane between workers.
* \param ccl The name of the communication backend, e.g., nccl, rccl, mpi.
* \param device_ids The device ids of the workers.
*/
TVM_RUNTIME_DLL virtual void InitCCL(ffi::String ccl, ffi::Shape device_ids) = 0;
/*!
* \brief Get the value of a register from a remote worker.
* \param reg_id The id of the register to be fetched.
* \param worker_id The id of the worker to be fetched from.
* \return The value of the register.
*/
TVM_RUNTIME_DLL virtual ffi::Any DebugGetFromRemote(int64_t reg_id, int worker_id) = 0;
/*!
* \brief Set the value of a register on a remote worker.
* \param reg_id The id of the register to be set.
* \param value The value to be set.
* \param worker_id The id of the worker to be set.
*/
TVM_RUNTIME_DLL virtual void DebugSetRegister(int64_t reg_id, ffi::AnyView value,
int worker_id) = 0;
struct FFI;
friend struct SessionObj::FFI;
friend class DRefObj;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("runtime.disco.Session", SessionObj, ffi::Object);
protected:
/*! \brief Deallocate a register id, kill it on all workers, and append it to `free_regs_`. */
virtual void DeallocReg(int reg_id) = 0;
};
/*!
* \brief Managed reference to SessionObj
* \sa SessionObj
*/
class Session : public ffi::ObjectRef {
public:
/*!
* \brief Create a session backed by a thread pool of workers
* \param num_workers The number of workers.
* \param num_groups The number of worker groups.
*/
TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups);
/*!
* \brief Create a session backed by pipe-based multiprocessing
* \param num_workers The number of workers.
* \param num_groups The number of worker groups.
* \param process_pool_creator The name of a global function that takes `num_workers` as an input,
* and returns a ffi::Function, which takes an integer `worker_id` as the input and returns None.
* When `worker-id` is 0, it shuts down the process pool; Otherwise, it retursn a tuple
* (read_fd, writefd) used to communicate with the corresponding worker.
* \param entrypoint The entrypoint of DiscoWorker main worker function.
* \note Worker-0 is always co-located with the controler as a separate thread, and therefore
* worker-0 does not exist in the process pool.
*/
TVM_RUNTIME_DLL static Session ProcessSession(int num_workers, int num_groups,
ffi::String process_pool_creator,
ffi::String entrypoint);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Session, ffi::ObjectRef, SessionObj);
};
/*!
* \brief A bi-directional channel for controler-worker communication.
* This channel is primarily used to transfer control messages but not data.
*/
class DiscoChannel {
public:
virtual ~DiscoChannel() = default;
/*! \brief Send a packed sequence to the receiver */
virtual void Send(const ffi::PackedArgs& args) = 0;
/*! \brief Receive a packed sequence from worker */
virtual ffi::PackedArgs Recv() = 0;
/*! \brief Reply a packed sequence to the sender */
virtual void Reply(const ffi::PackedArgs& args) = 0;
/*! \brief Receive a reply from the worker */
virtual ffi::PackedArgs RecvReply() = 0;
};
/*!
* \brief A special communication channel between controler and worker-0,
* assuming they are always collocated in the same process.
*/
class WorkerZeroData {
public:
/*!
* \brief The host-side arrays to passed to worker-0 for special uses, for example,
* copy-to-worker0 and copy-from-worker0
*/
std::queue<Tensor> host_arrays;
/*! \brief The mutex that guards `host_arrays` */
std::mutex queue_mutex_;
};
// Implementation details
inline SessionObj* DRefObj::GetSession() {
return const_cast<SessionObj*>(static_cast<const SessionObj*>(session.get()));
}
DRefObj::~DRefObj() {
if (this->session.defined()) {
GetSession()->DeallocReg(reg_id);
}
}
ffi::Any DRefObj::DebugGetFromRemote(int worker_id) {
return GetSession()->DebugGetFromRemote(this->reg_id, worker_id);
}
void DRefObj::DebugCopyFrom(int worker_id, ffi::AnyView value) {
return GetSession()->DebugSetRegister(this->reg_id, value, worker_id);
}
template <typename... Args>
DRef SessionObj::CallPacked(const DRef& func, Args&&... args) {
constexpr int offset = 3;
constexpr int kNumArgs = offset + sizeof...(Args);
ffi::AnyView packed_args[kNumArgs];
ffi::PackedArgs::Fill(packed_args,
/*.0=*/static_cast<int>(DiscoAction::kCallPacked), // action
/*.1=*/0, // reg_id, which will be updated by this->CallWithPacked
/*.2=*/func, // the function to be called
std::forward<Args>(args)...);
return this->CallWithPacked(ffi::PackedArgs(packed_args, kNumArgs));
}
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_SESSION_H_
include/tvm/runtime/logging.h
+402
View File
@@ -0,0 +1,402 @@
/*
* 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/runtime/logging.h
* \brief logging utilities
*
* We use the following facilities from tvm/ffi/error.h for
* error handling and checking:
*
* - TVM_FFI_THROW(ErrorKind) << "msg";
* - TVM_FFI_CHECK(cond, ErrorKind) << "msg";
* - TVM_FFI_CHECK_EQ(x, y, ErrorKind) << "msg";
* - TVM_FFI_ICHECK(x) << "msg"; // InternalError
* - TVM_FFI_ICHECK_EQ(x, y) << "msg";
* - TVM_FFI_DCHECK(x) << "msg"; // Debug-only InternalError
*
* LOG(INFO), LOG(WARNING), LOG(ERROR) are kept for logging.
* LOG(FATAL) is kept for completeness, it throws InternalError.
*/
#ifndef TVM_RUNTIME_LOGGING_H_
#define TVM_RUNTIME_LOGGING_H_
#include <tvm/ffi/error.h>
#include <tvm/runtime/base.h>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
/*!
* \brief Whether or not customize the logging output.
* If log customize is enabled, the user must implement
* tvm::runtime::detail::LogFatalImpl and tvm::runtime::detail::LogMessageImpl.
*/
#ifndef TVM_LOG_CUSTOMIZE
#define TVM_LOG_CUSTOMIZE 0
#endif
namespace tvm {
namespace runtime {
/*! \brief Internal implementation */
namespace detail {
// Provide support for customized logging.
#if TVM_LOG_CUSTOMIZE
/*!
* \brief Custom implementations of LogFatal.
*
* \sa TVM_LOG_CUSTOMIZE
*/
[[noreturn]] TVM_RUNTIME_DLL void LogFatalImpl(const std::string& file, int lineno,
const std::string& message);
/*!
* \brief Custom implementations of LogMessage.
*
* \sa TVM_LOG_CUSTOMIZE
*/
TVM_RUNTIME_DLL void LogMessageImpl(const std::string& file, int lineno, int level,
const std::string& message);
/*!
* \brief Class to accumulate an error message and throw it. Do not use
* directly, instead use LOG(FATAL).
*/
class LogFatal {
public:
LogFatal(const std::string& file, int lineno) : file_(file), lineno_(lineno) {}
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4722)
#endif
[[noreturn]] ~LogFatal() noexcept(false) { LogFatalImpl(file_, lineno_, stream_.str()); }
#ifdef _MSC_VER
#pragma warning(pop)
#endif
std::ostringstream& stream() { return stream_; }
private:
std::ostringstream stream_;
std::string file_;
int lineno_;
};
/*!
* \brief Class to accumulate an log message. Do not use directly, instead use
* LOG(INFO), LOG(WARNING), LOG(ERROR).
*/
class LogMessage {
public:
LogMessage(const std::string& file, int lineno, int level)
: file_(file), lineno_(lineno), level_(level) {}
~LogMessage() { LogMessageImpl(file_, lineno_, level_, stream_.str()); }
std::ostringstream& stream() { return stream_; }
private:
std::string file_;
int lineno_;
int level_;
std::ostringstream stream_;
};
#else
/*!
* \brief Class to accumulate an error message and throw it. Do not use
* directly, instead use LOG(FATAL).
* \note The `LogFatal` class is designed to be an empty class to reduce stack size usage.
* To play this trick, we use the thread-local storage to store its internal data.
*/
class LogFatal {
public:
TVM_FFI_NO_INLINE LogFatal(const char* file, int lineno) { GetEntry().Init(file, lineno); }
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4722)
#endif
[[noreturn]] ~LogFatal() noexcept(false) {
GetEntry().Finalize();
throw;
}
#ifdef _MSC_VER
#pragma warning(pop)
#endif
std::ostringstream& stream() { return GetEntry().stream_; }
private:
struct Entry {
void Init(const char* file, int lineno) {
this->stream_.str("");
this->file_ = file;
this->lineno_ = lineno;
}
[[noreturn]] TVM_FFI_NO_INLINE ffi::Error Finalize() noexcept(false) {
ffi::Error error("InternalError", stream_.str(),
TVMFFIBacktrace(file_.c_str(), lineno_, "", 0));
throw error;
}
std::ostringstream stream_;
std::string file_;
int lineno_;
};
TVM_FFI_NO_INLINE TVM_RUNTIME_DLL static Entry& GetEntry();
};
/*!
* \brief Class to accumulate an log message. Do not use directly, instead use
* LOG(INFO), LOG(WARNING), LOG(ERROR).
*/
class LogMessage {
public:
LogMessage(const std::string& file, int lineno, int level) {
// Use inline constexpr to avoid ODR-issues with hidden static members
// when libtvm_runtime.so is built with -fvisibility=hidden.
static constexpr const char* kLevelStrings[] = {
": Debug: ", // TVM_LOG_LEVEL_DEBUG
": ", // TVM_LOG_LEVEL_INFO
": Warning: ", // TVM_LOG_LEVEL_WARNING
": Error: ", // TVM_LOG_LEVEL_ERROR
};
std::time_t t = std::time(nullptr);
stream_ << "[" << std::put_time(std::localtime(&t), "%H:%M:%S") << "] " << file << ":" << lineno
<< kLevelStrings[level];
}
TVM_FFI_NO_INLINE ~LogMessage() { std::cerr << stream_.str() << std::endl; }
std::ostringstream& stream() { return stream_; }
private:
std::ostringstream stream_;
};
#endif
// This class is used to explicitly ignore values in the conditional
// logging macros. This avoids compiler warnings like "value computed
// is not used" and "statement has no effect".
class LogMessageVoidify {
public:
LogMessageVoidify() {}
// This has to be an operator with a precedence lower than << but
// higher than "?:". See its usage.
void operator&(std::ostream&) {}
};
/*! \brief Captures the state of the \p TVM_LOG_DEBUG environment flag. */
class TVM_RUNTIME_DLL TvmLogDebugSettings {
public:
/*!
* \brief Parses the \p TVM_LOG_DEBUG environment flag as per the specification given by
* \p DebugLoggingEnabled and \p VerboseLoggingEnabled, and caches the result.
*/
inline static const TvmLogDebugSettings& FromFlag() {
// Parse and cache the verbosity level map.
static const auto* settings =
new TvmLogDebugSettings(TvmLogDebugSettings::ParseSpec(std::getenv("TVM_LOG_DEBUG")));
return *settings;
}
/*!
* \brief Parses \p opt_spec as per specification for \p TVM_LOG_DEBUG given by
* \p DebugLoggingEnabled and \p VerboseLoggingEnabled. Throws if specification is ill-formed.
*/
static TvmLogDebugSettings ParseSpec(const char* opt_spec);
/*!
* \brief Implements \p VerboseLoggingEnabled below w.r.t. the already parsed \p TVM_LOG_DEBUG
* environment variable.
*/
inline bool VerboseEnabled(const char* opt_filename, int level) const {
if (opt_filename == nullptr || level < 0 || vlog_level_map_.empty()) {
return false;
}
return VerboseEnabledImpl(opt_filename, level);
}
/*! \brief Returns true if \p DLOG statements should be executed. */
bool dlog_enabled() const { return dlog_enabled_; }
private:
// Slow path for VerboseEnabled.
bool VerboseEnabledImpl(const std::string& filename, int level) const;
/*! \brief If true, DLOG statements are enabled. */
bool dlog_enabled_ = false;
/*!
* \brief A map from canonicalized filenames to the maximum VLOG verbosity level for that file.
* May also contain the 'wildcard' entry \p "DEFAULT" representing the level for all other files.
*/
std::unordered_map<std::string, int> vlog_level_map_;
};
/*!
* \brief Returns true if a DLOG statement is enabled by the \p TVM_LOG_DEBUG environment
* variable. Requires:
* \code
* TVM_LOG_DEBUG=1
* \endcode
* or a valid setting as described by \p VerboseLoggingEnabled below.
*/
inline bool DebugLoggingEnabled() {
static int state = 0;
if (state == 0) {
state = TvmLogDebugSettings::FromFlag().dlog_enabled() ? 1 : -1;
}
return state == 1;
}
/*!
* \brief Returns true if a VLOG statement in \p filename is enabled by the \p TVM_LOG_DEBUG
* environment variable for logging at verbosity \p level. Levels should be non-negative.
*
* Filenames are canonicalized to be w.r.t. the src/ dir of the TVM tree. (VLOG's should not
* appear under include/).
*
* To enable file \p ir/bar.cc for level 0 only set:
* \code
* TVM_LOG_DEBUG="ir/bar.cc=0"
* \endcode
*
* To enable all files up to level 3 but disable \p ir/bar.cc set:
* \code
* TVM_LOG_DEBUG="DEFAULT=2,ir/bar.cc=-1"
* \endcode
*
* Any of these settings will also enable DLOG statements.
*/
inline bool VerboseLoggingEnabled(const char* opt_filename, int level) {
return TvmLogDebugSettings::FromFlag().VerboseEnabled(opt_filename, level);
}
/*!
* \brief A stack of VLOG context messages.
*
* For use by \p VLOG_CONTEXT macro only.
*/
class VLogContext {
public:
void Push(std::stringstream* stream) { context_stack_.push_back(stream); }
void Pop() {
if (!context_stack_.empty()) {
context_stack_.pop_back();
}
}
std::string str() const;
private:
std::vector<std::stringstream*> context_stack_;
};
/*! \brief Get thread local \p VLogContext for tracking a stack of VLOG context messages. */
inline VLogContext* ThreadLocalVLogContext() {
static thread_local VLogContext inst;
return &inst;
}
/*!
* \brief A RAII class to push/pos a VLOG context message onto the thread-local stack.
*
* For use by \p VLOG_CONTEXT macro only.
*/
class VLogContextEntry {
public:
VLogContextEntry() { ThreadLocalVLogContext()->Push(&sstream_); }
~VLogContextEntry() { ThreadLocalVLogContext()->Pop(); }
std::ostream& stream() { return sstream_; }
private:
std::stringstream sstream_;
};
} // namespace detail
#define TVM_LOG_LEVEL_DEBUG 0
#define TVM_LOG_LEVEL_INFO 1
#define TVM_LOG_LEVEL_WARNING 2
#define TVM_LOG_LEVEL_ERROR 3
#define TVM_LOG_LEVEL_FATAL 4
#define LOG(level) LOG_##level
#define LOG_DEBUG \
::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_DEBUG).stream()
#define LOG_FATAL ::tvm::runtime::detail::LogFatal(__FILE__, __LINE__).stream()
#define LOG_INFO ::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_INFO).stream()
#define LOG_ERROR \
::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_ERROR).stream()
#define LOG_WARNING \
::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_WARNING).stream()
#define LOG_IF(severity, condition) \
!(condition) ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(severity)
#if TVM_LOG_DEBUG
#define LOG_DFATAL LOG_FATAL
#define DFATAL FATAL
#define DLOG(severity) LOG_IF(severity, ::tvm::runtime::detail::DebugLoggingEnabled())
#define DLOG_IF(severity, condition) \
LOG_IF(severity, ::tvm::runtime::detail::DebugLoggingEnabled() && (condition))
/*!
* \brief If the \p TVM_LOG_DEBUG build flag is enabled, push a context message onto an internal
* stack. All VLOG messages will include this stack in their prefix to help with debugging. E.g.:
* \code
* VLOG_CONTEXT << "my context";
* VLOG(1) << "my log message";
* \endcode
* Thread safe. No-op with no execution overhead if the \p TVM_LOG_DEBUG build flag is not enabled.
*/
#define VLOG_CONTEXT \
::tvm::runtime::detail::VLogContextEntry vlog_entry_; \
vlog_entry_.stream()
#else
#define LOG_DFATAL LOG_ERROR
#define DFATAL ERROR
#define DLOG(severity) true ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(severity)
#define DLOG_IF(severity, condition) \
(true || !(condition)) ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(severity)
#define VLOG_CONTEXT true ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(INFO)
#endif
/*!
* \brief If the \p TVM_LOG_DEBUG build flag is enabled, and the containing file has been enabled
* at \p level or greater in the \p TVM_LOG_DEBUG environment variable, then log a message at
* \p INFO severity.
*
* See \p VerboseLoggingEnabled for the format of the \p TVM_LOG_DEBUG environment variable.
* Thread safe. No-op with no execution overhead if the \p TVM_LOG_DEBUG build flag is not enabled.
* No-op with some execution overhead if the \p TVM_LOG_DEBUG build flag is enabled but the
* containing file is not enabled.
*/
#define VLOG(level) \
DLOG_IF(INFO, ::tvm::runtime::detail::VerboseLoggingEnabled(__FILE__, (level))) \
<< ::tvm::runtime::detail::ThreadLocalVLogContext()->str()
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_LOGGING_H_
include/tvm/runtime/memory/memory_manager.h
+199
View File
@@ -0,0 +1,199 @@
/*
* 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/runtime/memory/memory_manager.h
* \brief Abstract device memory management API
*/
#ifndef TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
#define TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
#include <tvm/runtime/base.h>
#include <tvm/runtime/tensor.h>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
namespace tvm {
namespace runtime {
namespace memory {
enum AllocatorType {
kNaive = 1,
kPooled,
};
struct Buffer {
/*! \brief The pointer to the allocated block of memory. */
void* data{nullptr};
/*! \brief The size of the block. */
size_t size{0};
/*! \brief The context of the allocated buffers. */
Device device;
/*! \brief The allocator that created this buffer. */
AllocatorType alloc_type;
};
class TVM_RUNTIME_DLL Allocator {
public:
explicit Allocator(AllocatorType type) : type_(type) {}
virtual ~Allocator() = default;
/*! \brief Allocate an empty Tensor using from the allocator.
* \param shape The shape of the Tensor.
* \param dtype The datatype of the Tensor.
* \param dev The device where the array is allocated.
* \param mem_scope The device memory scope hint.
* \return The empty Tensor.
*/
Tensor Empty(ffi::Shape shape, DLDataType dtype, Device dev,
ffi::Optional<ffi::String> mem_scope = std::nullopt);
/*! \brief Return the allocator type. */
inline AllocatorType type() const { return type_; }
/*! \brief Allocate a buffer given a size, alignment and type.
* \param dev The device where the array is allocated.
* \param nbytes The size of the buffer.
* \param alignment The alignment of the buffer.
* \param type_hint A type hint to the allocator.
* \return A sized allocation in the form of a buffer.
*/
virtual Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) = 0;
/*! \brief Allocate a buffer given a shape and type.
* \param dev The device where the array is allocated.
* \param shape The shape of the tensor.
* \param type_hint A type hint to the allocator.
* \param mem_scope A memory scope of the buffer.
* \return A sized allocation in the form of a buffer.
*/
virtual Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
const std::string& mem_scope = "");
/*! \brief Create a view for the buffer given a shape, type and scope.
* \param buffer The existing buffer upon which we need to create a view.
* \param shape The shape of the view.
* \param type_hint A type hint to the view.
* \param mem_scope A memory scope of the view.
* \return A device pointer to the created view.
*/
virtual void* CreateView(const Buffer& buffer, ffi::Shape shape, DLDataType type_hint,
const std::string& mem_scope = "global") {
return buffer.data;
}
/*! \brief Release the view .
* \param dev is the device where this view is created
* \param data The view pointer to be freed.
*/
virtual void FreeView(Device dev, void* data) {}
/*! \brief Free a buffer allocated by the allocator.
* \param buffer The buffer to free.
*/
virtual void Free(const Buffer& buffer) = 0;
/*! \brief Clear the allocated memory. */
virtual void Clear();
/*! \brief The amount of memory currently allocated.
* \return The amount of memory currently allocated.
*/
virtual size_t UsedMemory() const = 0;
protected:
/*! \brief Check if the given memory scope is allowed to allocate by the allocator. */
virtual bool AllowMemoryScope(const std::string& mem_scope) const;
private:
AllocatorType type_;
};
class MemoryManager {
public:
TVM_RUNTIME_DLL static MemoryManager* Global();
/*!
* \brief Get or create an allocator given the context and allocator type.
* \param dev The TVM device
* \param type The allocator type
* \return The memory allocator.
*/
TVM_RUNTIME_DLL static Allocator* GetOrCreateAllocator(Device dev, AllocatorType type);
/*!
* \brief Get an allocator given the context.
* \param dev The TVM device
* \param type The allocator type
* \return The memory allocator.
*/
TVM_RUNTIME_DLL static Allocator* GetAllocator(Device dev, AllocatorType type);
/*! \brief Clear the allocators. */
static void Clear();
private:
MemoryManager() {}
protected:
std::mutex mu_;
std::unordered_map<Device, std::unordered_map<AllocatorType, std::unique_ptr<Allocator>>>
allocators_;
};
/*! \brief An object representing a storage allocation. */
class StorageObj : public ffi::Object {
public:
/*! \brief The index into the VM function table. */
Buffer buffer;
/*! \brief The allocator where the storage buffer is allocated from. */
Allocator* allocator = nullptr;
/*! \brief Allocate an Tensor from a given piece of storage. */
TVM_RUNTIME_DLL Tensor AllocTensor(int64_t offset, ffi::Shape shape, DLDataType dtype);
/*! \brief Allocate an Tensor with memory scope from a given piece of storage. */
TVM_RUNTIME_DLL Tensor AllocTensorScoped(int64_t offset, ffi::Shape shape, DLDataType dtype,
ffi::String scope = "global");
~StorageObj() {
if (allocator) {
allocator->Free(buffer);
}
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("vm.Storage", StorageObj, ffi::Object);
};
/*! \brief reference to storage. */
class Storage : public ffi::ObjectRef {
public:
TVM_RUNTIME_DLL explicit Storage(Buffer buffer, Allocator* allocator);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Storage, ffi::ObjectRef, StorageObj);
};
} // namespace memory
using memory::Allocator;
using memory::AllocatorType;
using memory::MemoryManager;
using memory::StorageObj;
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
include/tvm/runtime/tensor.h
+360
View File
@@ -0,0 +1,360 @@
/*
* 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/runtime/tensor.h
* \brief A device-independent managed Tensor abstraction.
*/
#ifndef TVM_RUNTIME_TENSOR_H_
#define TVM_RUNTIME_TENSOR_H_
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/base.h>
#include <tvm/runtime/device_api.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
#include <atomic>
#include <functional>
#include <utility>
#include <vector>
namespace tvm {
namespace runtime {
/*!
* \brief Managed Tensor.
* The array is backed by reference counted blocks.
*/
class Tensor : public tvm::ffi::Tensor {
public:
Tensor() = default;
/*!
* \brief constructor.
* \param data ffi::ObjectPtr to the data container.
*/
explicit Tensor(ffi::ObjectPtr<ffi::TensorObj> data) : tvm::ffi::Tensor(data) {}
explicit Tensor(ffi::UnsafeInit tag) : tvm::ffi::Tensor(tag) {}
Tensor(ffi::Tensor&& other) : tvm::ffi::Tensor(std::move(other)) {} // NOLINT(*)
Tensor(const ffi::Tensor& other) : tvm::ffi::Tensor(other) {} // NOLINT(*)
ffi::ShapeView Shape() const { return this->shape(); }
DLDataType DataType() const { return this->dtype(); }
// DLPack handling
static Tensor FromDLPack(DLManagedTensor* tensor) {
return tvm::ffi::Tensor::FromDLPack(tensor, kAllocAlignment, true);
}
static Tensor FromDLPackVersioned(DLManagedTensorVersioned* tensor) {
return tvm::ffi::Tensor::FromDLPackVersioned(tensor, kAllocAlignment, true);
}
inline const DLTensor* operator->() const { return this->get(); }
/*!
* \brief Copy data content from another array.
* \param other The source array to be copied from.
* \note The copy may happen asynchronously if it involves a GPU context.
* TVMSynchronize is necessary.
*/
inline void CopyFrom(const DLTensor* other);
inline void CopyFrom(const Tensor& other);
/*!
* \brief Copy data content from a byte buffer.
* \param data The source bytes to be copied from.
* \param nbytes The size of the buffer in bytes
* Must be equal to the size of the Tensor.
* \note The copy always triggers a TVMSynchronize.
*/
TVM_RUNTIME_DLL void CopyFromBytes(const void* data, size_t nbytes);
/*!
* \brief Copy data content into another array.
* \param other The source array to be copied from.
* \note The copy may happen asynchronously if it involves a GPU context.
* TVMSynchronize is necessary.
*/
inline void CopyTo(DLTensor* other) const;
inline void CopyTo(const Tensor& other) const;
/*!
* \brief Copy data content into another array.
* \param data The source bytes to be copied from.
* \param nbytes The size of the data buffer.
* Must be equal to the size of the Tensor.
* \note The copy always triggers a TVMSynchronize.
*/
TVM_RUNTIME_DLL void CopyToBytes(void* data, size_t nbytes) const;
/*!
* \brief Copy the data to another device.
* \param dev The target device.
* \param mem_scope The memory scope of the target array.
* \return The array under another device.
* \note The copy always triggers a TVMSynchronize.
*/
TVM_RUNTIME_DLL Tensor CopyTo(const Device& dev,
ffi::Optional<ffi::String> mem_scope = std::nullopt) const;
/*!
* \brief Load Tensor from stream
* \param stream The input data stream
* \return Whether load is successful
*/
inline bool Load(support::Stream* stream);
/*!
* \brief Save Tensor to stream
* \param stream The output data stream
*/
inline void Save(support::Stream* stream) const;
/*!
* \brief Create a Tensor that shares the data memory with the current one.
*
* \param shape The shape of the new array.
*
* \param dtype The data type of the new array.
*
* \param relative_byte_offset The offset of the output Tensor,
* relative to the current byte offset.
*
* By default, the offset of the view is the same as the offset
* of the current array.
*
* \note The new array must not allow access of addresses which
* would be out of bounds in the current array. If the new
* array is larger than the current array, or if the
* `relative_byte_offset` would place the end of the new array
* outside the bounds of the current array, this function will
* raise an exception.
*/
TVM_RUNTIME_DLL Tensor CreateView(ffi::Shape shape, DLDataType dtype,
uint64_t relative_byte_offset = 0) const;
/*!
* \brief Create an empty Tensor.
* \param shape The shape of the new array.
* \param dtype The data type of the new array.
* \param dev The device of the array.
* \param mem_scope The memory scope of the array.
* \return The created Array
*/
TVM_RUNTIME_DLL static Tensor Empty(ffi::Shape shape, DLDataType dtype, Device dev,
ffi::Optional<ffi::String> mem_scope = std::nullopt);
/*!
* \brief Function to copy data from one array to another.
* \param from The source array.
* \param to The target array.
* \param stream The stream used in copy.
*/
TVM_RUNTIME_DLL static void CopyFromTo(const DLTensor* from, DLTensor* to,
TVMStreamHandle stream = nullptr);
/*!
* \brief Function to copy data from one array to a byte buffer.
* \param from The source array.
* \param to The target byte buffer.
* \param nbytes The size of the data buffer.
* \param stream The stream used in copy.
*/
TVM_RUNTIME_DLL static void CopyToBytes(const DLTensor* from, void* to, size_t nbytes,
TVMStreamHandle stream = nullptr);
/*!
* \brief Function to copy data from one array to a byte buffer.
* \param from The source array.
* \param to The target byte buffer.
* \param nbytes The size of the data buffer.
* \param stream The stream used in copy.
*/
TVM_RUNTIME_DLL static void CopyFromBytes(const DLTensor* to, void* from, size_t nbytes,
TVMStreamHandle stream = nullptr);
/*!
* \brief Check if two tensors share the same underlying storage.
*
* This detects runtime storage aliasing (e.g. views from CreateView, etc.) but does
* not imply either tensor was created by CreateView.
*
* \param a The first tensor.
* \param b The second tensor.
* \return True if the tensors share the same storage.
*/
TVM_RUNTIME_DLL static bool IsStorageShared(const DLTensor* a, const DLTensor* b);
/*!
* \brief Tensor overload of IsStorageShared.
* \param a The first tensor.
* \param b The second tensor.
* \return True if the tensors share the same storage.
*/
static bool IsStorageShared(const Tensor& a, const Tensor& b);
};
/*!
* \brief Save a DLTensor to stream
* \param strm The output stream
* \param tensor The tensor to be saved.
*/
inline bool SaveDLTensor(support::Stream* strm, const DLTensor* tensor);
inline void Tensor::CopyFrom(const DLTensor* other) {
TVM_FFI_ICHECK(data_ != nullptr);
CopyFromTo(other, get_mutable());
}
inline void Tensor::CopyFrom(const Tensor& other) {
TVM_FFI_ICHECK(data_ != nullptr);
TVM_FFI_ICHECK(other.data_ != nullptr);
CopyFromTo(other.get_mutable(), get_mutable());
}
inline void Tensor::CopyTo(DLTensor* other) const {
TVM_FFI_ICHECK(data_ != nullptr);
CopyFromTo(get_mutable(), other);
}
inline void Tensor::CopyTo(const Tensor& other) const {
TVM_FFI_ICHECK(data_ != nullptr);
TVM_FFI_ICHECK(other.data_ != nullptr);
CopyFromTo(get_mutable(), other.get_mutable());
}
/*! \brief Magic number for Tensor file */
constexpr uint64_t kTVMTensorMagic = 0xDD5E40F096B4A13F;
inline bool SaveDLTensor(support::Stream* strm, const DLTensor* tensor) {
uint64_t header = kTVMTensorMagic, reserved = 0;
strm->Write(header);
strm->Write(reserved);
// Always save data as CPU context
//
// Parameters that get serialized should be in CPU by default.
// So even the array's context is GPU, it will be stored as CPU array.
// This is used to prevent case when another user loads the parameters
// back on machine that do not have GPU or related context.
//
// We can always do array.CopyTo(target_dev) to get a corresponding
// array in the target context.
Device cpu_dev;
cpu_dev.device_type = kDLCPU;
cpu_dev.device_id = 0;
strm->Write(cpu_dev);
strm->Write(tensor->ndim);
strm->Write(tensor->dtype);
int ndim = tensor->ndim;
strm->WriteArray(tensor->shape, ndim);
int type_bytes = (tensor->dtype.bits + 7) / 8;
int64_t num_elems = 1;
for (int i = 0; i < ndim; ++i) {
num_elems *= tensor->shape[i];
}
int64_t data_byte_size = type_bytes * num_elems;
strm->Write(data_byte_size);
if (TVM_FFI_IO_NO_ENDIAN_SWAP && tensor->device.device_type == kDLCPU &&
ffi::IsContiguous(*tensor) && tensor->byte_offset == 0) {
// quick path
strm->Write(tensor->data, data_byte_size);
} else {
std::vector<uint8_t> bytes(data_byte_size);
Tensor::CopyToBytes(const_cast<DLTensor*>(tensor), bytes.data(), data_byte_size);
if (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
ffi::ByteSwap(bytes.data(), type_bytes, num_elems);
}
strm->Write(bytes.data(), data_byte_size);
}
return true;
}
inline void Tensor::Save(support::Stream* strm) const { SaveDLTensor(strm, operator->()); }
inline bool Tensor::Load(support::Stream* strm) {
uint64_t header, reserved;
TVM_FFI_ICHECK(strm->Read(&header)) << "Invalid DLTensor file format";
TVM_FFI_ICHECK(strm->Read(&reserved)) << "Invalid DLTensor file format";
TVM_FFI_ICHECK(header == kTVMTensorMagic) << "Invalid DLTensor file format";
Device dev;
int ndim;
DLDataType dtype;
TVM_FFI_ICHECK(strm->Read(&dev)) << "Invalid DLTensor file format";
TVM_FFI_ICHECK(strm->Read(&ndim)) << "Invalid DLTensor file format";
TVM_FFI_ICHECK(strm->Read(&dtype)) << "Invalid DLTensor file format";
TVM_FFI_ICHECK_EQ(dev.device_type, kDLCPU)
<< "Invalid DLTensor device: can only save as CPU tensor";
std::vector<int64_t> shape(ndim);
if (ndim != 0) {
TVM_FFI_ICHECK(strm->ReadArray(&shape[0], ndim)) << "Invalid DLTensor file format";
}
Tensor ret = Tensor::Empty(ffi::Shape(shape), dtype, dev);
int64_t num_elems = 1;
int elem_bytes = (ret->dtype.bits + 7) / 8;
for (int i = 0; i < ret->ndim; ++i) {
num_elems *= ret->shape[i];
}
int64_t data_byte_size;
TVM_FFI_ICHECK(strm->Read(&data_byte_size)) << "Invalid DLTensor file format";
TVM_FFI_ICHECK(data_byte_size == num_elems * elem_bytes) << "Invalid DLTensor file format";
auto read_ret = strm->Read(ret->data, data_byte_size);
// Only check non-empty data
if (ndim > 0 && shape[0] != 0) {
TVM_FFI_ICHECK(read_ret) << "Invalid DLTensor file format";
}
if (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
ffi::ByteSwap(ret->data, elem_bytes, num_elems);
}
*this = ret;
return true;
}
/*!
* \brief Get the preferred host device from the input device.
* - For CUDA and ROCm, CUDAHost and ROCMHost will be returned for pinned memory,
* since pinned memory reduces copy overhead.
* - For other devices, CPU is returned as a fallback.
*/
inline Device GetPreferredHostDevice(Device device) {
if (device.device_type == DLDeviceType::kDLCUDA) {
return Device{DLDeviceType::kDLCUDAHost, 0};
} else if (device.device_type == DLDeviceType::kDLROCM) {
return Device{DLDeviceType::kDLROCMHost, 0};
} else {
// Fallback to CPU.
return Device{DLDeviceType::kDLCPU, 0};
}
}
} // namespace runtime
} // namespace tvm
namespace std {
template <>
struct hash<tvm::Device> {
std::size_t operator()(const tvm::Device& dev) const {
return ((dev.device_id << 8) | dev.device_type);
}
};
template <>
struct equal_to<tvm::Device> {
bool operator()(const tvm::Device& lhs, const tvm::Device& rhs) const {
return (lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id);
}
};
} // namespace std
#endif // TVM_RUNTIME_TENSOR_H_
include/tvm/runtime/timer.h
+194
View File
@@ -0,0 +1,194 @@
/*
* 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 include/tvm/runtime/timer.h
* \brief Runtime timer primitives: Timer, TimerNode, WrapTimeEvaluator.
*/
#ifndef TVM_RUNTIME_TIMER_H_
#define TVM_RUNTIME_TIMER_H_
#include <tvm/ffi/function.h>
#include <tvm/runtime/base.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/tensor.h>
namespace tvm {
namespace runtime {
/*! \brief Base class for all timer implementations.
*
* New implementations of this interface should make sure that `Start` and `Stop`
* are as lightweight as possible. Expensive state synchronization should be
* done in `SyncAndGetElapsedNanos`.
*/
class TimerNode : public ffi::Object {
public:
/*! \brief Start the timer.
*
* Note: this function should only be called once per object.
*/
virtual void Start() = 0;
/*! \brief Stop the timer.
*
* Note: this function should only be called once per object.
*/
virtual void Stop() = 0;
/*! \brief Synchronize timer state and return elapsed time between `Start` and `Stop`.
* \return The time in nanoseconds between `Start` and `Stop`.
*
* This function is necessary because we want to avoid timing the overhead of
* doing timing. When using multiple timers, it is recommended to stop all of
* them before calling `SyncAndGetElapsedNanos` on any of them.
*
* Note: this function should be only called once per object. It may incur
* a large synchronization overhead (for example, with GPUs).
*/
virtual int64_t SyncAndGetElapsedNanos() = 0;
virtual ~TimerNode() {}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("runtime.TimerNode", TimerNode, ffi::Object);
};
/*! \brief Timer for a specific device.
*
* This is a managed reference to a TimerNode.
*
* \sa TimerNode
*/
class Timer : public ffi::ObjectRef {
public:
/*!
* \brief Get a device specific timer.
* \param dev The device to time.
* \return A `Timer` that has already been started.
*
* Use this function to time runtime of arbitrary regions of code on a specific
* device. The code that you want to time should be running on the device
* otherwise the timer will not return correct results. This is a lower level
* interface than TimeEvaluator and only runs the timed code once
* (TimeEvaluator runs the code multiple times).
*
* A default timer is used if a device specific one does not exist. This
* timer performs synchronization between the device and CPU, which can lead
* to overhead in the reported results.
*
* Example usage:
* \code{.cpp}
* Timer t = Timer::Start(Device::cpu());
* my_long_running_function();
* t->Stop();
* ... // some more computation
* int64_t nanosecs = t->SyncAndGetElapsedNanos() // elapsed time in nanoseconds
* \endcode
*
* To add a new device-specific timer, register a new function
* "runtime.timer.my_device" (where `my_device` is the `DeviceName` of your
* device). This function should accept a `Device` and return a new `Timer`
* that has already been started.
*
* For example, this is how the CPU timer is implemented:
* \code{.cpp}
* class CPUTimerNode : public TimerNode {
* public:
* virtual void Start() { start_ = std::chrono::high_resolution_clock::now(); }
* virtual void Stop() { duration_ = std::chrono::high_resolution_clock::now() - start_; }
* virtual int64_t SyncAndGetElapsedNanos() { return duration_.count(); }
* virtual ~CPUTimerNode() {}
*
* static constexpr const char* _type_key = "runtime.CPUTimerNode";
* TVM_FFI_DECLARE_OBJECT_INFO_FINAL(CPUTimerNode, TimerNode);
*
* private:
* std::chrono::high_resolution_clock::time_point start_;
* std::chrono::duration<int64_t, std::nano> duration_;
* };
*
*
* TVM_FFI_STATIC_INIT_BLOCK() {
* namespace refl = tvm::ffi::reflection;
* refl::GlobalDef().def("runtime.timer.cpu", [](Device dev) {
* return Timer(ffi::make_object<CPUTimerNode>());
* });
* }
* \endcode
*/
static TVM_RUNTIME_DLL Timer Start(Device dev);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Timer, ffi::ObjectRef, TimerNode);
};
/*!
* \brief Wrap a timer function to measure the time cost of a given packed function.
*
* Approximate implementation:
* \code{.py}
* f() // warmup
* for i in range(repeat)
* f_preproc()
* while True:
* start = time()
* for j in range(number):
* f()
* duration_ms = time() - start
* if duration_ms >= min_repeat_ms:
* break
* else:
* number = (min_repeat_ms / (duration_ms / number) + 1
* if cooldown_interval_ms and i % repeats_to_cooldown == 0:
* sleep(cooldown_interval_ms)
* \endcode
*
* \param f The function argument.
* \param dev The device.
* \param number The number of times to run this function for taking average.
* We call these runs as one `repeat` of measurement.
* \param repeat The number of times to repeat the measurement.
* In total, the function will be invoked (1 + number x repeat) times,
* where the first one is warm up and will be discarded.
* The returned result contains `repeat` costs,
* each of which is an average of `number` costs.
* \param min_repeat_ms The minimum duration of one `repeat` in milliseconds.
* By default, one `repeat` contains `number` runs. If this parameter is set,
* the parameters `number` will be dynamically adjusted to meet the
* minimum duration requirement of one `repeat`.
* i.e., When the run time of one `repeat` falls below this time,
* the `number` parameter will be automatically increased.
* \param limit_zero_time_iterations The maximum number of repeats when
* measured time is equal to 0. It helps to avoid hanging during measurements.
* \param cooldown_interval_ms The cooldown interval in milliseconds between the number of repeats
* defined by `repeats_to_cooldown`.
* \param repeats_to_cooldown The number of repeats before the
* cooldown is activated.
* \param cache_flush_bytes The number of bytes to flush from cache before
* \param f_preproc The function to be executed before we execute time
* evaluator.
* \return f_timer A timer function.
*/
ffi::Function WrapTimeEvaluator(ffi::Function f, Device dev, int number, int repeat,
int min_repeat_ms, int limit_zero_time_iterations,
int cooldown_interval_ms, int repeats_to_cooldown,
int cache_flush_bytes = 0, ffi::Function f_preproc = nullptr);
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_TIMER_H_
include/tvm/runtime/vm/builtin.h
+89
View File
@@ -0,0 +1,89 @@
/*
* 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/runtime/vm/builtin.h
* \brief Builtin runtime APIs.
*/
#ifndef TVM_RUNTIME_VM_BUILTIN_H_
#define TVM_RUNTIME_VM_BUILTIN_H_
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief Op code used in built-in match-shape function.
*
* The function takes the following signature:
* MatchShape(input_shape, shape_heap, n, c[0], r[0], c[1], r[1], ... c[n], r[n], err_ctx)
*
* This function provides runtime shape population and checking support for match-cast.
* When a shape variable appears in the first time, we should load the shape and
* populate the variable. When a shape variable already appears, we should
* assert that it already equals an existing shape value.
*
* NOTE: It is OK to pass nullptr shape_heap if all code are AssertEqualToImm.
*/
enum class MatchShapeCode : int {
/*!
* \brief Perform an assertion that shape equals immediate.
*
* assert input_shape[i] == r[i]
*/
kAssertEqualToImm = 0,
/*!
* \brief This is the first time we see a symbolic shape variable, store to heap.
*
* shape_heap[r[i]] = input_shape[i]
*/
kStoreToHeap = 1,
/*!
* \brief skip and do not do anything.
*/
kNoOp = 2,
/*!
* \brief Peform an assertion that the shape equals a loaded value.
*
* assert input_shape[i] == shape_heap[r[i]]
*/
kAssertEqualToLoad = 3,
};
/*!
* \brief Op code used in builtin function MakeShape.
*
* MakeShape(shape_heap, n, c[0], r[0], c[1], r[1], ... c[n], r[n]).
*
* \note It is OK to pass nullptr to shape_heap if all code are UseImm.
*/
enum class MakeShapeCode : int {
/*! \brief Use the following r[i] as immediate shape value. */
kUseImm = 0,
/*!
* \brief Load shape value from the shape_heap[[r[i]].
*/
kLoadShape = 1,
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_BUILTIN_H_
include/tvm/runtime/vm/bytecode.h
+246
View File
@@ -0,0 +1,246 @@
/*
* 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/runtime/vm/bytecode.h
* \brief The bytecode for the virtual machine.
*/
#ifndef TVM_RUNTIME_VM_BYTECODE_H_
#define TVM_RUNTIME_VM_BYTECODE_H_
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/error.h>
#include <iostream>
#include <vector>
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief The storage type for the bytecode in the VM.
*/
using ExecWord = int64_t;
/*! \brief A register name. */
using RegName = ExecWord;
/*!
* \brief An alias for the integer type used ubiquitously in the VM.
*/
using Index = ExecWord;
/*!
* \brief An enumeration of Relax's opcodes.
*
* The opcode is used to implement instruction
* as a tagged union.
*/
enum class Opcode {
Call = 1U,
Ret = 2U,
Goto = 3U,
If = 4U,
};
/*! \brief A single virtual machine instruction.
*
* The representation of the instruction is as
* a tagged union.
*
* The first field represents which instruction,
* and by extension which field of the union
* is active.
*/
struct Instruction {
/*! \brief The number of bit for storing value. */
static constexpr ExecWord kKindBit = 8;
/*! \brief The number of bit for storing value. */
static constexpr ExecWord kValueBit = sizeof(ExecWord) * 8 - kKindBit;
/*! \brief The bit mask of the value part. */
static constexpr ExecWord kValueMask = (static_cast<ExecWord>(1) << kValueBit) - 1;
/*! \brief Maximum possible value, use 1 bit for sign. */
static constexpr ExecWord kValueMaxLimit = (static_cast<ExecWord>(1) << (kValueBit - 1)) - 1;
/*! \brief Minimum possible value, remove 1 slot to keep things symmetric. */
static constexpr ExecWord kValueMinLimit = -kValueMaxLimit;
/*! \brief Beginning of special register section. */
static constexpr RegName kBeginSpecialReg = static_cast<ExecWord>(1) << 54;
/*! \brief Random magic number that represents void argument, indicate null value */
static constexpr RegName kVoidRegister = kBeginSpecialReg + 0;
/*! \brief Random magic number that represents the VM context */
static constexpr RegName kVMRegister = kBeginSpecialReg + 1;
/*!
* \brief The kind of instruction's argument.
*/
enum class ArgKind : int { kRegister = 0, kImmediate = 1, kConstIdx = 2, kFuncIdx = 3 };
friend std::ostream& operator<<(std::ostream& os, const ArgKind& kind) {
switch (kind) {
case ArgKind::kRegister:
os << "kRegister";
break;
case ArgKind::kImmediate:
os << "kImmediate";
break;
case ArgKind::kConstIdx:
os << "kConstIdx";
break;
case ArgKind::kFuncIdx:
os << "kFuncIdx";
break;
default:
TVM_FFI_THROW(InternalError)
<< "Internal error: "
<< "Invalid ArgKind with integer value " << static_cast<int>(kind);
}
return os;
}
/*!
* \brief The auxiliary data structure for instruction argument.
*/
class Arg {
public:
/*! \brief Construct a void argument. */
Arg() : data_(Instruction::kVoidRegister) {}
/*!
* \brief construct Arg from data.
* \param data The data repr.
*/
static Arg FromData(ExecWord data) { return Arg(data); }
/*!
* \brief construct a register Arg.
* \param reg The register number.
* \return The constructed arg.
*/
static Arg Register(RegName reg) { return Arg(ArgKind::kRegister, reg); }
/*!
* \brief construct a ConstIdx arg.
* \param index The constant index.
* \return The constructed arg.
*/
static Arg ConstIdx(Index index) { return Arg(ArgKind::kConstIdx, index); }
/*!
* \brief construct a immediate arg.
* \param imm_value The immediate value.
* \return The constructed arg.
*/
static Arg Immediate(int64_t imm_value) { return Arg(ArgKind::kImmediate, imm_value); }
/*!
* \brief construct a FuncIdx arg.
* \param index The func index in the function table.
* \return The constructed arg.
*/
static Arg FuncIdx(Index index) { return Arg(ArgKind::kFuncIdx, index); }
/*!
* \brief Get the kind of argument.
* \return The kind of argument.
*/
ArgKind kind() const {
uint8_t kind = (data_ >> kValueBit) & 0xFF;
return Instruction::ArgKind(kind);
}
/*!
* \brief Get the value of argument.
* \return The value of argument.
* \note We store both positive and negative values by sign extension.
*/
ExecWord value() const { return ((data_ & kValueMask) << kKindBit) >> kKindBit; }
/*!
* \brief Get the raw data repr of the arg.
* \return The raw data.
*/
ExecWord data() const { return data_; }
private:
/*! \brief Construct from the data. */
explicit Arg(ExecWord data) : data_(data) {}
/*! \brief Construct from the kind and value. */
Arg(ArgKind kind, Index value) {
TVM_FFI_ICHECK_LE(value, kValueMaxLimit);
TVM_FFI_ICHECK_GE(value, kValueMinLimit);
data_ = (static_cast<ExecWord>(kind) << kValueBit) | (value & kValueMask);
}
/*! \brief The underlying stored data. */
ExecWord data_;
};
/*! \brief The instruction opcode. */
Opcode op;
union {
struct /* Call */ {
/*! \brief The destination register. */
RegName dst;
/*! \brief The index into the packed function table. */
Index func_idx;
/*! \brief The number of arguments to the packed function. */
Index num_args;
/*! \brief The arguments of the packed function. */
Arg* args;
};
struct /* Ret */ {
/*! \brief The return result. */
RegName result;
};
struct /* Goto */ {
/*! \brief The jump offset. */
Index pc_offset;
};
struct /* If */ {
/*! \brief The register containing the cond value. */
RegName cond;
/*! \brief The program counter offset for the false branch. */
Index false_offset;
};
};
/*!
* \brief Construct a Call instruction.
* \param func_idx The index of the function to call.
* \param num_args The number of arguments.
* \param args The input arguments.
* \param dst The destination register.
* \return The call instruction.
*/
static Instruction Call(Index func_idx, Index num_args, Arg* args, RegName dst);
/*!
* \brief Construct a return instruction.
* \param result The register containing the return value.
* \return The return instruction.
*/
static Instruction Ret(RegName result);
/*!
* \brief Construct a goto instruction.
* \param pc_offset The register containing the jump offset.
* \return The goto instruction.
*/
static Instruction Goto(RegName pc_offset);
/*!
* \brief Construct an If instruction.
* \param cond The register containing the cond value.
* \param false_offset The program counter offset for the false branch.
* \return The If instruction.
*/
static Instruction If(RegName cond, Index false_offset);
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_BYTECODE_H_
include/tvm/runtime/vm/executable.h
+242
View File
@@ -0,0 +1,242 @@
/*
* 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/runtime/vm/executable.h
*/
#ifndef TVM_RUNTIME_VM_EXECUTABLE_H_
#define TVM_RUNTIME_VM_EXECUTABLE_H_
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
#include <string>
#include <unordered_map>
#include <vector>
#include "./bytecode.h"
// Convention: this version should set to minimum TVM version it support
// NOTE: this file only changes if we change relax vm format
// for example if relax vm format do not change in 0.15, this should remain as 0.14
// if it changes in 0.16, we will change it to 0.16
#define VM_VERSION "0.14"
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief Information entry in executable function table.
*
* Contains metadata about the compiled function, as
* well as the compiled VM instructions.
*/
struct VMFuncInfo {
/*! \brief kind of the function. */
enum class FuncKind : int {
/*! \brief system level packed function */
kPackedFunc = 0,
/*! \brief VM function. */
kVMFunc = 1,
/*! \brief VMTIR function. */
kVMTIRFunc = 2,
};
/*! \brief The kind of function. */
FuncKind kind;
/*! \brief The function's name, global symbol */
std::string name;
/*! \brief The start instruction index of the function. */
Index start_instr = 0;
/*! \brief The end instruction index of the function. */
Index end_instr = 0;
/*! \brief The number of arguments of the function. */
Index num_args = 0;
/*! \brief The register file size of the function. */
Index register_file_size = 0;
/*! \brief The function parameter names.*/
std::vector<std::string> param_names;
// defined customized loader save
void Save(support::Stream* writer) const;
bool Load(support::Stream* reader);
};
/*!
* \brief The virtual machine executable emitted by the VM compiler.
*
* The executable contains information (e.g. data in different memory regions)
* to run in a virtual machine.
*/
class TVM_RUNTIME_DLL VMExecutable : public ffi::ModuleObj {
public:
/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ffi::Module::kBinarySerializable; };
/*!
* \brief Print the detailed statistics of the given code, i.e. number of
* globals and constants, etc.
* \return The statistics represented by a string.
*/
std::string Stats() const;
/*!
* \brief Get the i-th instruction from the executable.
* \param i The index of the instruction to be fetched.
* \return The instruction.
*/
Instruction GetInstruction(Index i) const;
/*!
* \brief Set j-th byte data of i-th instruction to val.
* \param i The index of the instruction to be updated.
* \param j The index of the byte data of the instruction to be updated.
* \param val The value to be set
*/
void SetInstructionData(Index i, Index j, ExecWord val);
/*!
* \brief Print the instructions as text format.
* \return The text format of the instructions.
*/
ffi::String AsText() const;
/*!
* \brief Print the instructions as python program.
* \return The python program of the instructions, represented by a string.
*/
ffi::String AsPython() const;
/*!
* \brief Write the VMExecutable to the binary stream in serialized form.
* \return The binary bytes that save the executable to.
*/
ffi::Bytes SaveToBytes() const final;
/*!
* \brief Load VMExecutable from the binary stream in serialized form.
* \param bytes The binary bytes that load the executable from.
* \return The loaded executable, in the form of a `runtime::Module`.
*/
static ffi::Module LoadFromBytes(const ffi::Bytes& bytes);
/*!
* \brief Write the VMExecutable to the provided path as a file containing its serialized content.
* \param file_name The name of the file to write the serialized data to.
* \param format The target format of the saved file.
*/
void WriteToFile(const ffi::String& file_name, const ffi::String& format) const final;
/*! \brief Create a Relax virtual machine and load `this` as the executable. */
ffi::Module VMLoadExecutable() const;
/*! \brief Check if the VMExecutable contains a specific function. */
bool HasFunction(const ffi::String& name) const;
/*!
* \brief Load VMExecutable from the file.
* \param file_name The path of the file that load the executable from.
* \return The loaded executable, in the form of a `runtime::Module`.
*/
static ffi::Module LoadFromFile(const ffi::String& file_name);
/*! \brief The virtual machine's function table. */
std::vector<VMFuncInfo> func_table;
/*! \brief A map from globals (as strings) to their index in the function map. */
std::unordered_map<std::string, Index> func_map;
/*! \brief The global constant pool. */
std::vector<ffi::Any> constants;
/*! \brief The VDevice memory scopes */
std::unordered_map<Index, std::string> memory_scopes;
/*! \brief The offset of instruction. */
std::vector<Index> instr_offset;
/*! \brief The byte data of instruction. */
std::vector<ExecWord> instr_data;
virtual ~VMExecutable() {}
/*! \brief Module type key. */
const char* kind() const final;
/*!
* \brief Look up an exported function by name.
* \param name The function name.
* \return The function if found, otherwise std::nullopt.
*/
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) override;
private:
/*!
* \brief Save the globals.
* \param strm The input stream.
*/
void SaveGlobalSection(support::Stream* strm) const;
/*!
* \brief Save the memory scopes.
* \param strm The output stream.
*/
void SaveMemoryScopeSection(support::Stream* strm) const;
/*!
* \brief Save the constant pool.
* \param strm The input stream.
*/
void SaveConstantSection(support::Stream* strm) const;
/*!
* \brief Save the instructions.
* \param strm The input stream.
*/
void SaveCodeSection(support::Stream* strm) const;
/*!
* \brief Save the packed functions.
* \param strm The input stream.
*/
void SavePackedFuncNames(support::Stream* strm) const;
/*!
* \brief Load the globals.
* \param strm The input stream.
*/
void LoadGlobalSection(support::Stream* strm);
/*!
* \brief Load the memory scopes.
* \param strm The input stream.
*/
void LoadMemoryScopeSection(support::Stream* strm);
/*!
* \brief Load the constant pool.
* \param strm The input stream.
*/
void LoadConstantSection(support::Stream* strm);
/*!
* \brief Load the instructions.
* \param strm The input stream.
*/
void LoadCodeSection(support::Stream* strm);
/*!
* \brief Save the packed functions.
* \param strm The input stream.
*/
void LoadPackedFuncNames(support::Stream* strm);
};
} // namespace vm
} // namespace runtime
} // namespace tvm
namespace tvm {
namespace support {
template <>
struct Serializer<::tvm::runtime::vm::VMFuncInfo> {
static constexpr bool enabled = true;
static void Write(Stream* strm, const ::tvm::runtime::vm::VMFuncInfo& data) { data.Save(strm); }
static bool Read(Stream* strm, ::tvm::runtime::vm::VMFuncInfo* data) { return data->Load(strm); }
};
} // namespace support
} // namespace tvm
#endif // TVM_RUNTIME_VM_EXECUTABLE_H_
include/tvm/runtime/vm/tensor_cache_support.h
+97
View File
@@ -0,0 +1,97 @@
/*
* 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.
*/
#ifndef TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
#define TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/tensor.h>
#include <string>
#include <unordered_map>
#include <vector>
namespace tvm {
namespace runtime {
namespace vm {
/*!
* \brief Metadata for Tensor cache, which by default, is named as "tensor-cache.json".
*/
struct TensorCacheMetadata {
/*! \brief Each shard of Tensor cache, which by default, is named as "params_shard_x.bin". */
struct FileRecord {
/*! \brief Metadata of each parameter */
struct ParamRecord {
/*!
* \brief Load the parameter from raw data.
* \param device The device to load the parameter onto.
* \param raw_data The raw data stream
* \param staging_buffer The buffer to be used to avoid extra OpenCL copies. Pass in a nullptr
* in other cases
*/
TVM_RUNTIME_DLL Tensor Load(Device device, const std::string* raw_data,
ffi::Optional<Tensor>* staging_buffer = nullptr) const;
/*! \brief Name of the parameter */
std::string name;
/*! \brief Shape of the parameter */
ffi::Shape shape;
/*! \brief Data type of the parameter */
DLDataType dtype;
/*! \brief Format of the parameter */
std::string format;
/*! \brief Number of bytes */
int64_t nbytes;
/*! \brief Offset from the raw stream */
int64_t byte_offset;
};
/*! \brief Load a FileRecord into memory */
TVM_RUNTIME_DLL ffi::Array<Tensor> Load(Device device, //
const std::string& path_prefix, //
std::string* raw_data_buffer, //
ffi::Optional<Tensor>* staging_buffer = nullptr) const;
/*! \brief Relative path to the bin file */
std::string data_path;
/*! \brief Format of the file */
std::string format;
/*! \brief Size of the file */
int64_t nbytes;
/*! \brief The parameters in the file */
std::vector<ParamRecord> records;
};
/*! \brief The files in the Tensor cache */
std::vector<FileRecord> records;
/*! \brief The path to the `tensor-cache.json` file */
std::string path;
/*! \brief Load the metadata from a specific directory */
TVM_RUNTIME_DLL static TensorCacheMetadata Load(const std::string& path);
/*! \brief Load the metadata from a given JSON string */
TVM_RUNTIME_DLL static TensorCacheMetadata LoadFromStr(const std::string& json_str,
const std::string& path);
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
include/tvm/runtime/vm/vm.h
+226
View File
@@ -0,0 +1,226 @@
/*
* 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/runtime/vm/vm.h
*/
#ifndef TVM_RUNTIME_VM_VM_H_
#define TVM_RUNTIME_VM_VM_H_
#include <tvm/ffi/extra/module.h>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "../memory/memory_manager.h"
#include "./bytecode.h"
#include "./executable.h"
namespace tvm {
namespace runtime {
using memory::Allocator;
using memory::AllocatorType;
using memory::MemoryManager;
using memory::Storage;
using memory::StorageObj;
namespace vm {
/*!
* \brief Possible instrument actions.
*/
enum class VMInstrumentReturnKind : int {
/*! \brief Running as normal. */
kNoOp = 0,
/*! \brief Skip the following run, only valid in before. */
kSkipRun = 1,
};
/*!
* \brief An object representing a vm closure.
*/
class VMClosureObj : public ffi::Object {
public:
/*!
* \brief The function name. The function could be any
* function object that is compatible to the VM runtime.
*/
ffi::String func_name;
/*!
* \brief The implementation of the Closure.
* \note This function takes context pointer(VirtualMachine*)
* as the first argument. The rest of arguments follows
* the same arguments as the normal function call.
*/
ffi::Function impl;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.vm.Closure", VMClosureObj, ffi::Object);
};
/*! \brief reference to closure. */
class VMClosure : public ffi::ObjectRef {
public:
VMClosure(ffi::String func_name, ffi::Function impl);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMClosure, ffi::ObjectRef, VMClosureObj);
/*!
* \brief Create another ffi::Function with last arguments already bound to last_args.
*
* This is a helper function to create captured closures.
* \param func The input func, can be a VMClosure or ffi::Function.
* \param last_args The arguments to bound to in the end of the function.
* \note The new function takes in arguments and append the last_args in the end.
*/
static ffi::Function BindLastArgs(ffi::Function func, std::vector<ffi::Any> last_args);
};
/*!
* \brief Represent a VM extension.
* A VM extension allows the user to extend the VM with target specific functionalities.
* The VM holds the reference of the extensions to ensure the extensions have the same lifetime
* as the VM.
*
* This is the base class for all VM extensions and should not be used directly.
*/
class VMExtensionNode : public ffi::Object {
protected:
TVM_FFI_DECLARE_OBJECT_INFO("runtime.VMExtension", VMExtensionNode, ffi::Object);
};
/*! \brief Managed reference to VM extension. */
class VMExtension : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMExtension, ffi::ObjectRef, VMExtensionNode);
};
/*!
* \brief The virtual machine.
*
* The virtual machine contains all the current execution state,
* as well as the executable.
*
* The goal is to have a single self-contained object,
* enabling one to easily pass around VMs, execute them on
* multiple threads, or serialize them to disk or over the
* wire.
*/
class VirtualMachine : public ffi::ModuleObj {
public:
/*!
* \brief Initialize the virtual machine for a set of devices.
* \param devices The set of TVM devices.
* \param alloc_types The allocator types for each device.
*/
virtual void Init(const std::vector<Device>& devices,
const std::vector<AllocatorType>& alloc_types) = 0;
/*!
* \brief Load the executable for the virtual machine.
* \param exec The executable.
*/
virtual void LoadExecutable(ffi::ObjectPtr<VMExecutable> exec) = 0;
/*!
* \brief Get global function in the VM.
* \param func_name The name of the function.
* \return The closure
*/
virtual VMClosure GetClosure(const ffi::String& func_name) = 0;
/*!
* \brief Invoke closure or packed function using ffi::Function convention.
* \param closure_or_packedfunc A VM closure or a packed_func.
* \param args The input arguments.
* \param rv The return value.
*/
virtual void InvokeClosurePacked(const ffi::ObjectRef& closure_or_packedfunc,
ffi::PackedArgs args, ffi::Any* rv) = 0;
/*!
* \brief Set an instrumentation function.
*
* If instrument is present, the function will be called
* before/after each Call instruction.
*
* bool instrument(func, func_symbol, before_run, args...)
*
* - func: Union[VMClosure, ffi::Function], the function object.
* - func_symbol: string, the symbol of the function.
* - before_run: bool, whether it is before or after call.
* - ret_value: Only valid in after run, otherwise it is null.
* - args: the arguments being passed to call.
*
* instrument can return an int which corresponds to the action value.
* \sa VMInstrumentAction
*
* \param instrument The instrument function.
*/
virtual void SetInstrument(ffi::Function instrument) = 0;
/*!
* \brief Get or create a VM extension. Once created, the extension will be stored in the VM
* and held until the VM is destructed.
*
* \tparam T The type of the extension
* \return The extension instance
*/
template <typename T, typename = std::enable_if_t<std::is_base_of<VMExtension, T>::value>>
T GetOrCreateExtension() {
using ContainerType = typename T::ContainerType;
uint32_t key = ContainerType::RuntimeTypeIndex();
if (auto it = extensions.find(key); it != extensions.end()) {
ffi::Any value = (*it).second;
return value.cast<T>();
}
auto [it, _] = extensions.emplace(key, T::Create());
ffi::Any value = (*it).second;
return value.cast<T>();
}
/*!
* \brief Create a specific instance of VM.
* \return Created VM
*/
static ffi::ObjectPtr<VirtualMachine> Create();
/*!
* \brief Helper function for vm closure functions to get the context ptr
* \param arg The argument value.
*/
static VirtualMachine* GetContextPtr(ffi::AnyView arg) {
return static_cast<VirtualMachine*>(arg.cast<void*>());
}
~VirtualMachine() {}
//--------------------------------------------------------------------------
// The following section contains states that other builtin can depend on
//--------------------------------------------------------------------------
/*! \brief The memory allocators. */
std::vector<Allocator*> allocators;
/*! \brief Runtime physical device list. */
std::vector<Device> devices;
/*! \brief The VM extensions. Mapping from the type index of the extension to the extension
* instance. */
std::unordered_map<uint32_t, Any> extensions;
};
} // namespace vm
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_VM_VM_H_
include/tvm/s_tir/analysis.h
+211
View File
@@ -0,0 +1,211 @@
/*
* 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/s_tir/analysis.h
* \brief Analysis utilities for Schedulable TensorIR (S-TIR).
*/
#ifndef TVM_S_TIR_ANALYSIS_H_
#define TVM_S_TIR_ANALYSIS_H_
#include <tvm/ir/module.h>
#include <tvm/ir/transform.h>
#include <tvm/target/target.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
#include <optional>
namespace tvm {
namespace tirx {
/*!
* \brief Auto detect the block access region according to its body stmt
* It will detect the access region as an array in order of appearance in AST
* \param block The block to be detected
* \param buffer_var_map The outside buffers which may be accessed the block.
* It is a map from buffer var to the buffer.
* \return Array of access regions.
* There are three arrays of BufferRegion:
* - first: read regions
* - second: write regions
* - third: opaque regions
*/
TVM_DLL ffi::Array<ffi::Array<BufferRegion>> GetSBlockAccessRegion(
const SBlock& block, const ffi::Map<Var, Buffer>& buffer_var_map);
/*!
* \brief Auto detect the block read/write region according to its body stmt. An opaque access will
* be counted as both a read and a write access
* \param block The block to be detected
* \param buffer_var_map The outside buffers which may be accessed the block.
* It is a map from buffer var to the buffer
* \return An array only consisting of the read regions and write regions of the input block
*/
TVM_DLL ffi::Array<ffi::Array<BufferRegion>> GetSBlockReadWriteRegion(
const SBlock& block, const ffi::Map<Var, Buffer>& buffer_var_map);
/*!
* \brief Detect the lowest common ancestor(LCA) of buffer access, including both high-level
* access(BufferLoad, BufferStore) and low-level access(Load, Store and opaque access).
* The LCA may be a For loop or a Block.
* \param func The PrimFunc to be detected.
* \return The Map from buffer to the LCA of all access to it. The lca is function root if the
* return stmt is std::nullopt.
*/
TVM_DLL ffi::Map<Buffer, ffi::Optional<Stmt>> DetectBufferAccessLCA(const PrimFunc& func);
/*!
* \brief Find the "anchor block" of the given module.
* We define the anchor block to be the block with (1) an init statement and (2) having
* the biggest flops count. The latter condition is only used when there are multiple blocks
* with an init statement.
* For example, if the input module is conv2d + fused spatial blocks, conv2d is the anchor block.
* The input module may not contain more than one such block. For example, a module having
* two conv2d is not allowed as an input.
* However, a module created from winograd convolution has multiple blocks with an init statement
* (input transform, batched GEMM, and output transform). We use the second condition, the flops
* count, to determine that the batched GEMM block is the anchor block.
* \param mod The input TIR module.
* \return The anchor block if found, nullptr otherwise.
*/
const tirx::SBlockNode* FindAnchorBlock(const IRModule& mod);
} // namespace tirx
namespace arith {
class AnalyzerObj;
class Analyzer;
} // namespace arith
namespace s_tir {
using namespace tvm::tirx;
/*!
* \brief Estimate the FLOPs of a TIR fragment.
* \param stmt The TIR fragment to be estimated.
* \return The estimated FLOPs.
*/
TVM_DLL double EstimateTIRFlops(const Stmt& stmt);
/*!
* \brief Estimate the FLOPs of TIRs in an IRModule.
* \param mod The IRModule to be estimated.
* \return The estimated FLOPs.
*/
TVM_DLL double EstimateTIRFlops(const IRModule& mod);
/*!
* \brief Analyze the side effect of a function
* \param func The function to be checked.
* \param assert_on_error If true, an error will be thrown for an impure function.
* \return The purity of the function.
*/
TVM_DLL bool IsPureFunction(const PrimFunc& func, bool assert_on_error = false);
/*!
* \brief Verify the correctness of a GPU code
* \param func The function to be checked.
* \param constraints The dict to specify constraints to check.
* \return valid Whether it is a valid GPU code.
*/
TVM_DLL bool VerifyGPUCode(const PrimFunc& func, ffi::Map<ffi::String, PrimExpr> constraints);
/*! \brief Helper struct for return value of IdentifyMemCpy */
struct MemCpyDetails {
BufferRegion source;
BufferRegion dest;
};
/*! \brief Identify whether a For loop is semantically equivalent to MemCpy
* \param loop The loop to be checked
* \param analyzer The analyzer with which to check any algebraic expressions
* \returns The source and destination regions being copied, if the loop is equivalent to memcpy.
*/
TVM_DLL std::optional<MemCpyDetails> IdentifyMemCpy(const For& loop,
const arith::Analyzer& analyzer);
/*!
* \brief Calculate the allocated memory per scope in bytes needed inside the TIR PrimFunc
* \param func The TIR PrimFunc for which the allocated memory size to be calculated
* \return Allocated memory size per scope in bytes.
*/
TVM_DLL ffi::Map<ffi::String, ffi::Map<ffi::String, int64_t>> CalculateAllocatedBytes(
const PrimFunc& func);
/*!
* \brief Calculate the allocated memory per scope in bytes for each function inside the module
* \param mod The IRModule for which the allocated memory size has to be calculated
* \return Allocated memory size per scope in bytes for each function.
*/
TVM_DLL ffi::Map<ffi::String, ffi::Map<ffi::String, int64_t>> CalculateAllocatedBytes(
const IRModule& mod);
/**
* \brief Get the list of lowering passes to calculate the compacted VTCM allocation size.
* \return The list of passes.
*/
TVM_DLL ffi::Array<tvm::transform::Pass> GetVTCMCompactionPasses();
/*!
* \brief Verifies that the VTCM usage for all prim_funcs in the given IRModule.
* \param mod The module to be checked.
* \param limit The limit to check.
* \return true if the VTCM usage is within the provided limit.
*/
TVM_DLL bool VerifyVTCMLimit(const IRModule& mod, int64_t limit);
/*!
* \brief Verifies that the VTCM usage of the given prim_func is within the provided limit.
* \param func The function to be checked.
* \param limit The limit to check.
* \return true if the VTCM usage is within the provided limit.
*/
TVM_DLL bool VerifyVTCMLimit(const PrimFunc& func, int64_t limit);
namespace transform {
using tvm::transform::Pass;
using tvm::transform::PassContext;
/*!
* \brief Pass to verify GPU code constraints.
* \param constraints The dict to specify constraints.
* \return The pass.
*/
TVM_DLL Pass VerifyGPUCode(ffi::Map<ffi::String, PrimExpr> constraints);
/*!
* \brief Pass to check if VTCM usage is within limit.
* \param default_target The default target for functions without target attribute.
* \return The pass.
*/
TVM_DLL Pass VerifyVTCMLimit(ffi::Optional<Target> default_target = std::nullopt);
/*!
* \brief Statically check TIR code for out of bounds array access.
* \return The pass.
*/
TVM_DLL Pass OOBChecker();
} // namespace transform
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_ANALYSIS_H_
include/tvm/s_tir/backend/adreno/transform.h
+64
View File
@@ -0,0 +1,64 @@
/*
* 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/s_tir/backend/adreno/transform.h
* \brief S-TIR specific Adreno GPU transformation passes.
*/
#ifndef TVM_S_TIR_BACKEND_ADRENO_TRANSFORM_H_
#define TVM_S_TIR_BACKEND_ADRENO_TRANSFORM_H_
#include <tvm/ir/transform.h>
#include <tvm/s_tir/transform.h>
#include <tvm/target/target.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/function.h>
#include <string>
#include <vector>
namespace tvm {
namespace s_tir {
namespace backend {
namespace adreno {
namespace transform {
using tirx::transform::CreatePrimFuncPass;
using tvm::transform::Pass;
using tvm::transform::PassContext;
/*!
* \brief Texture flattening pass.
* \return The pass.
*/
TVM_DLL Pass TextureFlatten();
/*!
* \brief Inject Texture Allocation intrinsic.
* \return The pass.
*/
TVM_DLL Pass InjectTextureAlloc();
} // namespace transform
} // namespace adreno
} // namespace backend
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_BACKEND_ADRENO_TRANSFORM_H_
include/tvm/s_tir/data_layout.h
+411
View File
@@ -0,0 +1,411 @@
/*
* 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/s_tir/data_layout.h
* \brief SLayout expression to describe the data organization of a tensor.
* And SBijectiveLayout to mapping two data layouts between each other.
*/
#ifndef TVM_S_TIR_DATA_LAYOUT_H_
#define TVM_S_TIR_DATA_LAYOUT_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <algorithm>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
#include "tvm/tirx/var.h"
namespace tvm {
namespace tirx {
class SLayout;
class SLayoutAxis {
public:
static const SLayoutAxis& Get(const char name);
// Get the singleton SLayoutAxis using itvar->var->name_hint
static const SLayoutAxis& Get(const tirx::IterVar& itvar);
// Get the singleton SLayoutAxis using name[0] (size of name must be 1).
static const SLayoutAxis& Get(const std::string& name);
inline bool IsPrimal() const { return name_ >= 'A' && name_ <= 'Z'; }
inline std::string name() const { return std::string(1, name_); }
// if current axis is primal, switch the axis to its subordinate one,
// else switch to the primal.
inline const SLayoutAxis& ToDual() const {
if (name_ >= 'A' && name_ <= 'Z') {
return SLayoutAxis::Get(name_ - 'A' + 'a');
} else {
return SLayoutAxis::Get(name_ - 'a' + 'A');
}
}
// return the primal axis. If it is already primal, return itself.
const SLayoutAxis& ToPrimal() const { return IsPrimal() ? *this : ToDual(); }
// return the subordinate axis. If it is already subordinate, return itself.
const SLayoutAxis& ToSubordinate() const { return IsPrimal() ? ToDual() : *this; }
inline bool operator==(const SLayoutAxis& rhs) const { return name_ == rhs.name_; }
friend std::ostream& operator<<(std::ostream& os, const SLayoutAxis& l) {
os << l.name();
return os;
}
private:
static const SLayoutAxis UPPER_CASE[];
static const SLayoutAxis LOWER_CASE[];
SLayoutAxis(const SLayoutAxis&);
SLayoutAxis& operator=(const SLayoutAxis&);
explicit SLayoutAxis(const char name) : name_(name) {}
const char name_;
};
/*!
* \brief SLayout is to describe how data is organized within an N-dimention tensor.
* It is composed of upper cases, lower cases and numbers,
* where upper case indicates a primal axis and
* the corresponding lower case with factor size indicates the subordinate axis.
* For example, NCHW16c can describe a 5-D tensor of
* [batch_size, channel, height, width, channel_block].
* Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel).
* SLayout for scalar is defined, while both its name and axes have size 0.
*/
class SLayoutNode : public ffi::Object {
public:
/*! \brief string representation of layout, "" for scalar. */
ffi::String name;
/*! \brief specify each axis of the layout,
* in which the variable name is the name of the axis.
* The IterVar's extent indicates the size of the axis,
* it is a variable for a primal axis, but a constant for a subordinate axis.
* Empty for scalar's layout.
*/
ffi::Array<tirx::IterVar> axes;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SLayoutNode>()
.def_ro("name", &SLayoutNode::name)
.def_ro("axes", &SLayoutNode::axes);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SLayout", SLayoutNode, ffi::Object);
};
/*!
* \brief Managed reference to SLayoutNode
* \sa SLayoutNode
*/
class SLayout : public ffi::ObjectRef {
public:
explicit SLayout(const ffi::Array<tirx::IterVar>& axes);
/*! \brief construct from a string */
SLayout(const tvm::ffi::String& name) : SLayout(name.operator std::string()) {} // NOLINT(*)
/*! \brief construct from a string */
SLayout(const char* name) : SLayout(std::string(name)) {} // NOLINT(*)
/*!
* \brief construct from a string.
* \param name input in layout convention:
* upper case indicates a dimension and
* the corresponding lower case with factor size
* indicates the split dimension.
* return undefined layout if "__undef__" is passed.
* \param index_ty The type of generated axes vars in the returned layout.
* It is required to be integer type.
*/
TVM_DLL SLayout(const std::string& name, PrimType index_ty = PrimType::Int(32)); // NOLINT(*)
/*!
* \brief access the internal node container
* \return the pointer to the internal node container
*/
SLayoutNode* operator->() { return static_cast<SLayoutNode*>(get_mutable()); }
/*!
* \brief Return an undefined layout.
* \return a (global) undefined layout.
*/
static const SLayout& Undef() {
static SLayout undef;
return undef;
}
/*!
* \brief Packs the Given Array of IterVars into a Single IterVar. Each IterVar in the Array
* should represent either a single primal axis or one or more subordinate axis
* \param iters Array of iter vars to be packed
* \return A packed iter var
*/
static IterVar PackIterVar(ffi::Array<IterVar> iters);
/*!
* \brief Unpacks a Packed IterVar into its constituents
* \param packed_iter A Packed IterVar containing a single primal axis or one or more subordinate
* axis
* \return Constituent IterVars
*/
static ffi::Array<IterVar> UnpackIterVar(IterVar packed_iter);
/*!
* \brief Returns a sub-layout which is the portion of the object
* that starts at dimension \p pos and spans \p len dimensions
* (or until the end of the layout, whichever comes first).
* \param pos The start position.
* \param len The length of the sub-layout. if 0, return layout of scalar
* \return A newly constructed SLayout object.
*/
SLayout SubLayout(size_t pos, size_t len) const;
/*!
* \brief Split \p axis by \p size and put the sub-axis to position \p target_pos.
* \param axis The source axis to be split. It must be a primal-axis;
* \param target_pos The target position of the newly split subordinate-axis.
* \param factor size of the sub-dimension.
* \return A newly constructed SLayout object.
*/
SLayout Split(const SLayoutAxis& axis, size_t target_pos, int32_t factor) const;
/*! \return number of dimensions */
inline size_t ndim() const {
if (!defined()) return 0;
return operator->()->axes.size();
}
/*! \return number of super dimensions */
inline size_t ndim_primal() const {
if (!defined()) return 0;
size_t ct = 0;
for (auto px : operator->()->axes) {
auto iter_vars = UnpackIterVar(px);
for (auto x : iter_vars) {
if (SLayoutAxis::Get(x).IsPrimal()) {
ct++;
}
}
}
return ct;
}
/*!
* \brief Returns a new layout where the dims have been expanded to match the primal dimensions.
* \param dst_layout The dst layout to which current layout has to be expanded.
* \return The expanded SLayout.
*/
inline SLayout ExpandPrimal(const SLayout& dst_layout) {
SLayout new_src_layout;
// 1) Find the axis which are missing in the current layout. Make them the prefix.
std::string new_src_layout_str = "";
for (auto packed_axis : dst_layout->axes) {
auto iter_vars = UnpackIterVar(packed_axis);
for (auto dst_axis : iter_vars) {
if (SLayoutAxis::Get(dst_axis).IsPrimal()) {
if (!this->Contains(SLayoutAxis::Get(dst_axis))) {
new_src_layout_str += dst_axis->var->name_hint;
}
}
}
}
// 2) Now, add the primal axis of the current layout.
new_src_layout_str += this->name();
new_src_layout = SLayout(new_src_layout_str);
return new_src_layout;
}
/*!
* \brief return the index of the input axis.
* If it is not found in the layout or the layout is undefined,
* return -1.
* \param axis The input axis either a layout axis, or a packed axis
* \return the index or -1 if not found.
*/
inline int32_t IndexOf(const std::string& axis) const {
if (!this->defined()) return -1;
const auto axes = operator->()->axes;
for (size_t i = 0; i < axes.size(); ++i) {
if (axes[i]->var->name_hint == axis) return static_cast<int32_t>(i);
}
return -1;
}
/*!
* \brief return the index of the input axis.
* If it is not found in the layout or the layout is undefined,
* return -1.
* \param axis the input layout axis.
* \return the index or -1 if not found.
*/
inline int32_t IndexOf(const SLayoutAxis& axis) const { return IndexOf(axis.name()); }
/*!
* \brief return the index of the input axis.
* If it is not found in the layout or the layout is undefined,
* return -1.
* \param iter the input iter var.
* \return the index or -1 if not found.
*/
inline int32_t IndexOf(const tirx::IterVar& iter) const { return IndexOf(iter->var->name_hint); }
/*!
* \brief Get the factor size of the subordinate axis.
* \param axis the input primal-axis or subordinate-axis.
* \return the size of the subordinate-axis of \p axis (if \p axis is a primal-axis),
* or the size of \p axis itself (if \p axis is a subordinate-axis).
* Return -1 if \p axis is not in the layout the layout is undefined.
*/
int32_t FactorOf(const SLayoutAxis& axis) const;
/*!
* \brief Whether the layout contains an axis.
* \param axis axis to be checked.
* \return Whether the layout contains the axis.
*/
bool Contains(const SLayoutAxis& axis) const {
if (!defined()) return false;
for (const tirx::IterVar packed_var : operator->()->axes) {
auto iter_vars = UnpackIterVar(packed_var);
for (auto var : iter_vars) {
if (var->var->name_hint == axis.name()) {
return true;
}
}
}
return false;
}
const SLayoutAxis& operator[](int32_t i) const {
TVM_FFI_ICHECK(defined()) << "Try to access axis from an undefined layout.";
int32_t index = i < 0 ? static_cast<int32_t>(ndim() + i) : i;
TVM_FFI_ICHECK(index >= 0 && static_cast<size_t>(index) < ndim()) << "Invalid index " << i;
const tirx::IterVar axis = operator->()->axes[index];
return SLayoutAxis::Get(axis);
}
IterVar PackedAxisAt(int32_t i) const {
TVM_FFI_ICHECK(defined()) << "Try to access axis from an undefined layout.";
int32_t index = i < 0 ? static_cast<int32_t>(ndim() + i) : i;
TVM_FFI_ICHECK(index >= 0 && static_cast<size_t>(index) < ndim()) << "Invalid index " << i;
const tirx::IterVar axis = operator->()->axes[index];
return axis;
}
/*! \return the string description of the layout */
inline std::string name() const {
if (!defined()) return "__undef__";
return operator->()->name;
}
/*!
* \brief Whether the two layouts are equal.
* \param rhs Another layout.
* \return whether the two layouts are equal.
*/
inline bool Equals(const SLayout& rhs) const { return name() == rhs.name(); }
/*!
* \brief allow output string of layout to ostream
* \param os the output stream
* \param l the layout
* \return the ostream
*/
friend std::ostream& operator<<(std::ostream& os, const SLayout& l) {
os << l.name();
return os;
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SLayout, ffi::ObjectRef, SLayoutNode);
};
// Internal node container SBijectiveLayout
class SBijectiveLayoutNode : public ffi::Object {
public:
/*! \brief Describes how source axes can be mapped to the destination axes,
* e.g., [i0 / 16, i1, i0 % 16] can describe NC -> NC16n
*/
ffi::Array<PrimExpr> index_forward_rule;
/*! \brief Describes how destination axes can be mapped to the source axes */
ffi::Array<PrimExpr> index_backward_rule;
/*! \brief Describes how source shapes can be mapped to the destination shapes */
ffi::Array<PrimExpr> shape_forward_rule;
/*! \brief Describes how destination shapes can be mapped to the source shapes */
ffi::Array<PrimExpr> shape_backward_rule;
/*! \brief The source layout */
SLayout src_layout;
/*! \brief The destination layout */
SLayout dst_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SBijectiveLayoutNode>()
.def_ro("src_layout", &SBijectiveLayoutNode::src_layout)
.def_ro("dst_layout", &SBijectiveLayoutNode::dst_layout)
.def_ro("index_forward_rule", &SBijectiveLayoutNode::index_forward_rule)
.def_ro("index_backward_rule", &SBijectiveLayoutNode::index_backward_rule)
.def_ro("shape_forward_rule", &SBijectiveLayoutNode::shape_forward_rule)
.def_ro("shape_backward_rule", &SBijectiveLayoutNode::shape_backward_rule);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBijectiveLayout", SBijectiveLayoutNode, ffi::Object);
};
/*!
* \brief Bijective function mapping for data layout transformation.
* Given two SLayout, SBijectiveLayout build and store the mapping rules,
* provides API to transform N-dimention tensor from the source indices (i0, i1, .., im)
* to the destination indices (j0, j1, .., jm).
*/
class SBijectiveLayout : public ffi::ObjectRef {
public:
/*!
* \brief The constructor
* \param src_layout The source layout
* \param dst_layout The destination layout
*/
TVM_DLL SBijectiveLayout(SLayout src_layout, SLayout dst_layout);
// Given the source shape, infer the destination shape.
TVM_DLL ffi::Array<PrimExpr> ForwardShape(const ffi::Array<PrimExpr>& shape) const;
// Given the destination shape, recover the source shape.
TVM_DLL ffi::Array<PrimExpr> BackwardShape(const ffi::Array<PrimExpr>& dst_shape) const;
// Given the destination indices, infer the destination indices.
TVM_DLL ffi::Array<PrimExpr> ForwardIndex(const ffi::Array<PrimExpr>& index) const;
// Given the destination indices, recover the source indices.
TVM_DLL ffi::Array<PrimExpr> BackwardIndex(const ffi::Array<PrimExpr>& dst_index) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBijectiveLayout, ffi::ObjectRef,
SBijectiveLayoutNode);
};
} // namespace tirx
} // namespace tvm
#endif // TVM_S_TIR_DATA_LAYOUT_H_
include/tvm/s_tir/meta_schedule/arg_info.h
+121
View File
@@ -0,0 +1,121 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_ARG_INFO_H_
#define TVM_S_TIR_META_SCHEDULE_ARG_INFO_H_
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/module.h>
#include <tvm/tirx/function.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
/*! \brief The argument information. */
class ArgInfoNode : public ffi::Object {
public:
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.ArgInfo", ArgInfoNode, ffi::Object);
public:
/*! \brief Default destructor. */
virtual ~ArgInfoNode() = default;
/*! \brief Converts the ArgInfo to its corresponding JSON representation. */
virtual ffi::ObjectRef AsJSON() const = 0;
};
/*!
* \brief Managed reference to ArgInfoNode
* \sa ArgInfoNode
*/
class ArgInfo : public ffi::ObjectRef {
public:
/*!
* \brief Parse the argument information from a JSON object.
* \param json_obj The json object to parse.
* \return The argument information parsed.
*/
TVM_DLL static ArgInfo FromJSON(const ffi::ObjectRef& json_obj);
/*!
* \brief Extract a list of the argument information from PrimFunc.
* \param func The PrimFunc to get argument information from.
* \return An array of the argument information derived.
*/
TVM_DLL static ffi::Array<ArgInfo, void> FromPrimFunc(const tirx::PrimFunc& func);
/*!
* \brief Extract a list of the argument information from the entry func of an IRModule
* \param mod The IRModule to extract argument information from.
* \param remove_preproc Whether to remove the preprocessing blocks.
* \return An array of the argument information derived.
*/
TVM_DLL static ffi::Array<ArgInfo, void> FromEntryFunc(const IRModule& mod, bool remove_preproc);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ArgInfo, ffi::ObjectRef, ArgInfoNode);
protected:
ArgInfo() = default;
};
/*! \brief The tensor argument information. */
class TensorInfoNode : public ArgInfoNode {
public:
/*! \brief The data type of the tensor. */
DLDataType dtype;
/*! \brief The shape of the tensor. */
ffi::Shape shape;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TensorInfoNode>()
.def_ro("dtype", &TensorInfoNode::dtype)
.def_ro("shape", &TensorInfoNode::shape);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.TensorInfo", TensorInfoNode, ArgInfoNode);
public:
ffi::ObjectRef AsJSON() const;
};
/*!
* \brief Managed reference to TensorInfoNode
* \sa TensorInfoNode
*/
class TensorInfo : public ArgInfo {
public:
/*!
* \brief Constructor of TensorInfo.
* \param dtype The data type of the tensor argument.
* \param shape The shape tuple of the tensor argument.
*/
TVM_DLL explicit TensorInfo(DLDataType dtype, ffi::Shape shape);
/*!
* \brief Parse the argument information from a JSON object.
* \param json_obj The json object to parse.
* \return The argument information parsed.
*/
TVM_DLL static TensorInfo FromJSON(const ffi::ObjectRef& json_obj);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorInfo, ArgInfo, TensorInfoNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_ARG_INFO_H_
include/tvm/s_tir/meta_schedule/builder.h
+175
View File
@@ -0,0 +1,175 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_BUILDER_H_
#define TVM_S_TIR_META_SCHEDULE_BUILDER_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/module.h>
#include <tvm/runtime/tensor.h>
#include <tvm/target/target.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
/*! \brief The builder's input, containing an IRModule and the target. */
class BuilderInputNode : public ffi::Object {
public:
/*! \brief The IRModule to be built. */
IRModule mod;
/*! \brief The target to be built for. */
Target target;
/*! \brief Parameters for Relax build module. */
ffi::Optional<ffi::Map<ffi::String, runtime::Tensor>> params;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BuilderInputNode>()
.def_ro("mod", &BuilderInputNode::mod)
.def_ro("target", &BuilderInputNode::target)
.def_ro("params", &BuilderInputNode::params);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.BuilderInput", BuilderInputNode,
ffi::Object);
};
/*!
* \brief Managed reference to BuilderInputNode
* \sa BuilderInputNode
*/
class BuilderInput : public ffi::ObjectRef {
public:
/*!
* \brief Constructor of BuilderInput.
* \param mod The IRModule to be built.
* \param target The target to be built for.
* \param params Parameters for Relax build module.
*/
TVM_DLL explicit BuilderInput(
IRModule mod, Target target,
ffi::Optional<ffi::Map<ffi::String, runtime::Tensor>> params = std::nullopt);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BuilderInput, ffi::ObjectRef, BuilderInputNode);
};
/*! \brief The builder's output, containing the artifact path or error message if any. */
class BuilderResultNode : public ffi::Object {
public:
/*! \brief The path to the built artifact. */
ffi::Optional<ffi::String> artifact_path;
/*! \brief The error message if any. */
ffi::Optional<ffi::String> error_msg;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BuilderResultNode>()
.def_ro("artifact_path", &BuilderResultNode::artifact_path)
.def_ro("error_msg", &BuilderResultNode::error_msg);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.BuilderResult", BuilderResultNode,
ffi::Object);
};
/*!
* \brief Managed reference to BuilderResultNode
* \sa BuilderResultNode
*/
class BuilderResult : public ffi::ObjectRef {
public:
/*!
* \brief Constructor of BuilderResult.
* \param artifact_path The path to the built artifact.
* \param error_msg The error message if any.
*/
TVM_DLL explicit BuilderResult(ffi::Optional<ffi::String> artifact_path,
ffi::Optional<ffi::String> error_msg);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BuilderResult, ffi::ObjectRef, BuilderResultNode);
};
/*! \brief The abstract builder interface. */
class BuilderNode : public ffi::Object {
public:
/*! \brief Default destructor */
virtual ~BuilderNode() = default;
/*!
* \brief Generate the build results from build inputs.
* \param build_inputs The inputs to be built.
* \return The build results.
*/
virtual ffi::Array<BuilderResult> Build(const ffi::Array<BuilderInput>& build_inputs) = 0;
/*!
* \brief The function type of `Build` method.
* \param build_inputs The inputs to be built.
* \return The build results.
*/
using FBuild = ffi::TypedFunction<ffi::Array<BuilderResult>(const ffi::Array<BuilderInput>&)>;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.Builder", BuilderNode, ffi::Object);
};
/*!
* \brief Managed reference to BuilderNode
* \sa BuilderNode
*/
class Builder : public ffi::ObjectRef {
public:
/*!
* \brief Constructor from ffi::ObjectPtr<BuilderNode>.
* \param data The object pointer.
*/
explicit Builder(ffi::ObjectPtr<BuilderNode> data) : ffi::ObjectRef(data) {
TVM_FFI_ICHECK(data != nullptr);
}
/*!
* \brief Create a builder with customized build method on the python-side.
* \param f_build The packed function to the `Build` function..
* \return The Builder created.
*/
static Builder PyBuilder(BuilderNode::FBuild f_build);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Builder, ffi::ObjectRef, BuilderNode);
};
/*! \brief An abstract builder with customized build method on the python-side. */
class PyBuilderNode : public BuilderNode {
public:
/*! \brief The packed function to the `Build` function. */
FBuild f_build;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PyBuilderNode>().def_ro("f_build", &PyBuilderNode::f_build);
}
ffi::Array<BuilderResult> Build(const ffi::Array<BuilderInput>& build_inputs) final {
TVM_FFI_ICHECK(f_build != nullptr) << "PyBuilder's Build method not implemented!";
return f_build(build_inputs);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyBuilder", PyBuilderNode, BuilderNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_BUILDER_H_
include/tvm/s_tir/meta_schedule/cost_model.h
+155
View File
@@ -0,0 +1,155 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_COST_MODEL_H_
#define TVM_S_TIR_META_SCHEDULE_COST_MODEL_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/base.h>
#include <tvm/s_tir/meta_schedule/arg_info.h>
#include <tvm/s_tir/meta_schedule/measure_candidate.h>
#include <tvm/s_tir/meta_schedule/runner.h>
#include <tvm/s_tir/schedule/schedule.h>
#include <vector>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
class TuneContext;
/*! \brief Cost model. */
class CostModelNode : public ffi::Object {
public:
/*! \brief Virtual destructor. */
virtual ~CostModelNode() = default;
/*!
* \brief Load the cost model from given file location.
* \param path The file path.
*/
virtual void Load(const ffi::String& path) = 0;
/*!
* \brief Save the cost model to given file location.
* \param path The file path.
*/
virtual void Save(const ffi::String& path) = 0;
/*!
* \brief Update the cost model given running results.
* \param context The tuning context.
* \param candidates The measure candidates.
* \param results The running results of the measure candidates.
*/
virtual void Update(const TuneContext& context, const ffi::Array<MeasureCandidate>& candidates,
const ffi::Array<RunnerResult>& results) = 0;
/*!
* \brief Predict the normalized score (the larger the better) of given measure candidates.
* \param context The tuning context.
* \param candidates The measure candidates.
* \return The predicted normalized score.
*/
virtual std::vector<double> Predict(const TuneContext& context,
const ffi::Array<MeasureCandidate>& candidates) = 0;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.CostModel", CostModelNode, ffi::Object);
};
/*! \brief The cost model with customized methods on the python-side. */
class PyCostModelNode : public CostModelNode {
public:
/*!
* \brief Load the cost model from given file location.
* \param path The file path.
*/
using FLoad = ffi::TypedFunction<void(ffi::String)>;
/*!
* \brief Save the cost model to given file location.
* \param path The file path.
*/
using FSave = ffi::TypedFunction<void(ffi::String)>;
/*!
* \brief Update the cost model given running results.
* \param context The tuning context.
* \param candidates The measure candidates.
* \param results The running results of the measure candidates.
* \return Whether cost model was updated successfully.
*/
using FUpdate = ffi::TypedFunction<void(const TuneContext&, const ffi::Array<MeasureCandidate>&,
const ffi::Array<RunnerResult>&)>;
/*!
* \brief Predict the running results of given measure candidates.
* \param context The tuning context.
* \param candidates The measure candidates.
* \param p_addr The address to save the estimated running results.
*/
using FPredict = ffi::TypedFunction<void(const TuneContext&, const ffi::Array<MeasureCandidate>&,
void* p_addr)>;
/*! \brief The packed function to the `Load` function. */
FLoad f_load;
/*! \brief The packed function to the `Save` function. */
FSave f_save;
/*! \brief The packed function to the `Update` function. */
FUpdate f_update;
/*! \brief The packed function to the `Predict` function. */
FPredict f_predict;
void Load(const ffi::String& path);
void Save(const ffi::String& path);
void Update(const TuneContext& context, const ffi::Array<MeasureCandidate>& candidates,
const ffi::Array<RunnerResult>& results);
std::vector<double> Predict(const TuneContext& context,
const ffi::Array<MeasureCandidate>& candidates);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyCostModel", PyCostModelNode,
CostModelNode);
};
/*!
* \brief Managed reference to CostModelNode
* \sa CostModelNode
*/
class CostModel : public ffi::ObjectRef {
public:
/*!
* \brief Create a cost model with customized methods on the python-side.
* \param f_load The packed function of `Load`.
* \param f_save The packed function of `Save`.
* \param f_update The packed function of `Update`.
* \param f_predict The packed function of `Predict`.
* \return The cost model created.
*/
TVM_DLL static CostModel PyCostModel(PyCostModelNode::FLoad f_load, //
PyCostModelNode::FSave f_save, //
PyCostModelNode::FUpdate f_update, //
PyCostModelNode::FPredict f_predict);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CostModel, ffi::ObjectRef, CostModelNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_COST_MODEL_H_
include/tvm/s_tir/meta_schedule/database.h
+552
View File
@@ -0,0 +1,552 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_DATABASE_H_
#define TVM_S_TIR_META_SCHEDULE_DATABASE_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
#include <tvm/s_tir/meta_schedule/arg_info.h>
#include <tvm/s_tir/schedule/schedule.h>
#include <tvm/s_tir/schedule/trace.h>
#include <tvm/target/target.h>
#include <filesystem>
#include <memory>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
class ModuleEquality;
/*! \brief A workload, i.e. an IRModule and its structural hash. */
class WorkloadNode : public ffi::Object {
public:
/*! \brief The type of structural hash */
using THashCode = size_t;
/*! \brief The workload's IRModule. */
IRModule mod;
/*! \brief The workload's structural hash. */
THashCode shash;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<WorkloadNode>().def_ro("mod", &WorkloadNode::mod);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.Workload", WorkloadNode, ffi::Object);
/*!
* \brief Export the workload to a JSON string.
* \return An array containing the structural hash and the base64 json string.
*/
ffi::ObjectRef AsJSON() const;
};
/*!
* \brief Managed reference to WorkloadNode.
* \sa WorkloadNode
*/
class Workload : public ffi::ObjectRef {
public:
using THashCode = WorkloadNode::THashCode;
explicit Workload(ffi::ObjectPtr<WorkloadNode> data) : ffi::ObjectRef(data) {}
/*!
* \brief Constructor of Workload.
* \param mod The workload's IRModule.
*/
TVM_DLL explicit Workload(IRModule mod);
/*!
* \brief Constructor of Workload.
* \param mod The workload's IRModule.
* \param shash The workload's structural hash.
*/
TVM_DLL explicit Workload(IRModule mod, THashCode shash);
/*!
* \brief Create a workload from a json object.
* \param json_obj The json object.
* \return The created workload.
*/
TVM_DLL static Workload FromJSON(const ffi::ObjectRef& json_obj);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Workload, ffi::ObjectRef, WorkloadNode);
};
/*! \brief The hash method for Workload */
struct WorkloadHash {
size_t operator()(const Workload& a) const { return a->shash; }
};
/*! \brief The equality check for Workload */
struct WorkloadEqual {
explicit WorkloadEqual(const ModuleEquality& mod_eq) : mod_eq_(mod_eq) {}
bool operator()(const Workload& a, const Workload& b) const;
private:
/*! \brief The module equality testing and hashing method */
const ModuleEquality& mod_eq_;
};
/*! \brief The class of measure candidates. */
class MeasureCandidate;
/*! \brief The class of tuning records. */
class TuningRecordNode : public ffi::Object {
public:
/*! \brief The trace tuned. */
s_tir::Trace trace;
/*! \brief The workload. */
Workload workload{ffi::UnsafeInit()};
/*! \brief The profiling result in seconds. */
ffi::Optional<ffi::Array<FloatImm>> run_secs;
/*! \brief The target for tuning. */
ffi::Optional<Target> target;
/*! \brief The argument information. */
ffi::Optional<ffi::Array<ArgInfo>> args_info;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TuningRecordNode>()
.def_ro("trace", &TuningRecordNode::trace)
.def_ro("workload", &TuningRecordNode::workload)
.def_ro("run_secs", &TuningRecordNode::run_secs)
.def_ro("target", &TuningRecordNode::target)
.def_ro("args_info", &TuningRecordNode::args_info);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.TuningRecord", TuningRecordNode,
ffi::Object);
/*! \brief Construct the measure candidate given the initial IR module and trace
* stored in the tuning record. */
MeasureCandidate AsMeasureCandidate() const;
/*!
* \brief Export the tuning record to a JSON string.
* \return An array containing the trace, running secs, serialized target, and
* argument information.
*/
ffi::ObjectRef AsJSON() const;
/*!
* \brief Check if this tuning record has valid trace instructions and successful run results.
* \return The check result.
*/
bool IsValid() const;
};
/*!
* \brief The managed reference of TuningRecordNode.
* \sa TuningRecordNode
*/
class TuningRecord : public ffi::ObjectRef {
public:
/*!
\brief Constructor of a tuning record.
\param trace The trace of the tuning record.
\param workload The workload of the tuning record.
\param run_secs The running time of the tuning record.
\param target The target of the tuning record.
\param args_info The argument information of the tuning record.
*/
TVM_DLL explicit TuningRecord(s_tir::Trace trace, Workload workload,
ffi::Optional<ffi::Array<FloatImm>> run_secs,
ffi::Optional<Target> target,
ffi::Optional<ffi::Array<ArgInfo>> args_info);
/*!
* \brief Create a tuning record from a json object.
* \param json_obj The json object.
* \param workload The workload.
* \return The tuning record created.
*/
TVM_DLL static TuningRecord FromJSON(const ffi::ObjectRef& json_obj, const Workload& workload);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TuningRecord, ffi::ObjectRef, TuningRecordNode);
};
class Database;
/* \brief The abstract interface of database. */
class DatabaseNode : public ffi::Object {
public:
/*!
* \brief Constructor
* \param mod_eq_name A string to specify the module equality testing and hashing method.
* It must be one of the followings:
* - "structural": Use StructuralEqual/Hash
* - "ignore-tensor": Same as "structural", but ignore tensor raw data during
* equality testing and hashing.
* - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
* given module. The "ignore-tensor" varint is used for the extracted blocks
* or in case no anchor block is found.
* For the definition of the anchor block, see tvm/tirx/analysis.h.
*/
explicit DatabaseNode(ffi::String mod_eq_name = "structural");
/*! \brief Default destructor */
virtual ~DatabaseNode();
/*!
* \brief Check if the database has the given workload.
* \param mod The IRModule to be searched for.
* \return Whether the database has the given workload.
*/
virtual bool HasWorkload(const IRModule& mod) = 0;
/*!
* \brief Look up or add workload to the database if missing.
* \param mod The IRModule to be searched for or added.
* \return The workload corresponding to the given IRModule.
*/
virtual Workload CommitWorkload(const IRModule& mod) = 0;
/*!
* \brief Add a tuning record to the database.
* \param record The tuning record to be added.
*/
virtual void CommitTuningRecord(const TuningRecord& record) = 0;
/*!
* \brief Get the top K valid tuning records of given workload from the database.
* \param workload The workload to be searched for.
* \param top_k The number of top records to be returned.
* \return An array of top K tuning records for the given workload.
*/
virtual ffi::Array<TuningRecord> GetTopK(const Workload& workload, int top_k) = 0;
/*!
* \brief Get all tuning records from the database.
* \return An Array of all the tuning records in the database.
*/
virtual ffi::Array<TuningRecord> GetAllTuningRecords() = 0;
/*!
* \brief Get the size of the database.
* \return The size of the database.
*/
virtual int64_t Size() = 0;
/*!
* \brief Query the best record of the given workload from the database.
* \param mod The IRModule to be searched for.
* \param target The target to be searched for.
* \param workload_name The name of the workload to be searched for.
* \return The best record of the given workload; std::nullopt if not found.
*/
virtual ffi::Optional<TuningRecord> QueryTuningRecord(const IRModule& mod, const Target& target,
const ffi::String& workload_name);
/*!
* \brief Query the best schedule of the given workload from the database.
* \param mod The IRModule to be searched for.
* \param target The target to be searched for.
* \param workload_name The name of the workload to be searched for.
* \return The schedule in the best schedule of the given workload; std::nullopt if not found.
*/
virtual ffi::Optional<s_tir::Schedule> QuerySchedule(const IRModule& mod, const Target& target,
const ffi::String& workload_name);
/*!
* \brief Query the best IRModule of the given workload from the database.
* \param mod The IRModule to be searched for.
* \param target The target to be searched for.
* \param workload_name The name of the workload to be searched for.
* \return The IRModule in the best IRModule of the given workload; std::nullopt if not found.
*/
virtual ffi::Optional<IRModule> QueryIRModule(const IRModule& mod, const Target& target,
const ffi::String& workload_name);
/*!
* \brief Prune the database and dump it a given database.
* \param destination The destination database to be dumped to.
*/
void DumpPruned(Database destination);
/*! \brief Return a reference to the owned module equality method instance. */
const ModuleEquality& GetModuleEquality() const {
TVM_FFI_ICHECK(mod_eq_);
return *mod_eq_;
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.Database", DatabaseNode, ffi::Object);
private:
/*! \brief The module equality testing and hashing method */
std::unique_ptr<ModuleEquality> mod_eq_;
};
/*! \brief The database with customized methods on the python-side. */
class PyDatabaseNode : public DatabaseNode {
public:
/*!
* \brief Constructor
* \param mod_eq_name A string to specify the module equality testing and hashing method.
* It must be one of the followings:
* - "structural": Use StructuralEqual/Hash
* - "ignore-tensor": Same as "structural", but ignore tensor raw data during
* equality testing and hashing.
* - "anchor-block": Apply equality testing and hashing on the anchor block extracted from a
* given module. The "ignore-tensor" varint is used for the extracted blocks
* or in case no anchor block is found.
* For the definition of the anchor block, see tvm/tirx/analysis.h.
*/
explicit PyDatabaseNode(ffi::String mod_eq_name = "structural");
/*!
* \brief The function type of `HasWorkload` method.
* \param mod The IRModule to be searched for.
* \return Whether the database has the given workload.
*/
using FHasWorkload = ffi::TypedFunction<bool(const IRModule&)>;
/*!
* \brief The function type of `CommitWorkload` method.
* \param mod The IRModule to be searched for or added.
* \return The workload corresponding to the given IRModule.
*/
using FCommitWorkload = ffi::TypedFunction<Workload(const IRModule&)>;
/*!
* \brief The function type of `CommitTuningRecord` method.
* \param record The tuning record to be added.
*/
using FCommitTuningRecord = ffi::TypedFunction<void(const TuningRecord&)>;
/*!
* \brief The function type of `GetTopK` method.
* \param workload The workload to be searched for.
* \param top_k The number of top records to be returned.
* \return An array of top K tuning records for the given workload.
*/
using FGetTopK = ffi::TypedFunction<ffi::Array<TuningRecord>(const Workload&, int)>;
/*!
* \brief The function type of `GetAllTuningRecords` method.
* \return An Array of all the tuning records in the database.
*/
using FGetAllTuningRecords = ffi::TypedFunction<ffi::Array<TuningRecord>()>;
/*!
* \brief The function type of `QueryTuningRecord` method.
* \param mod The IRModule to be searched for.
* \param target The target to be searched for.
* \param workload_name The name of the workload to be searched for.
* \return The best record of the given workload; std::nullopt if not found.
*/
using FQueryTuningRecord = ffi::TypedFunction<ffi::Optional<TuningRecord>(
const IRModule&, const Target&, const ffi::String&)>;
/*!
* \brief The function type of `QuerySchedule` method.
* \param mod The IRModule to be searched for.
* \param target The target to be searched for.
* \param workload_name The name of the workload to be searched for.
* \return The schedule in the best schedule of the given workload; std::nullopt if not found.
*/
using FQuerySchedule = ffi::TypedFunction<ffi::Optional<s_tir::Schedule>(
const IRModule&, const Target&, const ffi::String&)>;
/*!
* \brief The function type of `QueryIRModule` method.
* \param mod The IRModule to be searched for.
* \param target The target to be searched for.
* \param workload_name The name of the workload to be searched for.
* \return The IRModule in the best IRModule of the given workload; std::nullopt if not found.
*/
using FQueryIRModule = ffi::TypedFunction<ffi::Optional<IRModule>(const IRModule&, const Target&,
const ffi::String&)>;
/*!
* \brief The function type of `Size` method.
* \return The size of the database.
*/
using FSize = ffi::TypedFunction<int64_t()>;
/*! \brief The packed function to the `HasWorkload` function. */
FHasWorkload f_has_workload;
/*! \brief The packed function to the `CommitWorkload` function. */
FCommitWorkload f_commit_workload;
/*! \brief The packed function to the `CommitTuningRecord` function. */
FCommitTuningRecord f_commit_tuning_record;
/*! \brief The packed function to the `GetTopK` function. */
FGetTopK f_get_top_k;
/*! \brief The packed function to the `GetAllTuningRecords` function. */
FGetAllTuningRecords f_get_all_tuning_records;
/*! \brief The packed function to the `QueryTuningRecord` function. */
FQueryTuningRecord f_query_tuning_record;
/*! \brief The packed function to the `QuerySchedule` function. */
FQuerySchedule f_query_schedule;
/*! \brief The packed function to the `QueryIRModule` function. */
FQueryIRModule f_query_ir_module;
/*! \brief The packed function to the `Size` function. */
FSize f_size;
static void RegisterReflection() {
// ffi::Functions are all not registered, because the reflection system doesn't take care of
// them, so it cannot be accessible on the python side. If there is such need from the future,
// we can then add corresponding accessor methods to help access on python.
// `f_has_workload` is not registered
// `f_commit_workload` is not registered
// `f_commit_tuning_record` is not registered
// `f_get_top_k` is not registered
// `f_get_all_tuning_records` is not registered
// `f_query_tuning_record` is not registered
// `f_query_schedule` is not registered
// `f_query_ir_module` is not registered
// `f_size` is not registered
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PyDatabaseNode>();
}
bool HasWorkload(const IRModule& mod) final {
TVM_FFI_ICHECK(f_has_workload != nullptr) << "PyDatabase's HasWorkload method not implemented!";
return f_has_workload(mod);
}
Workload CommitWorkload(const IRModule& mod) final {
TVM_FFI_ICHECK(f_commit_workload != nullptr)
<< "PyDatabase's CommitWorkload method not implemented!";
return f_commit_workload(mod);
}
void CommitTuningRecord(const TuningRecord& record) final {
TVM_FFI_ICHECK(f_commit_tuning_record != nullptr)
<< "PyDatabase's CommitTuningRecord method not implemented!";
f_commit_tuning_record(record);
}
ffi::Array<TuningRecord> GetTopK(const Workload& workload, int top_k) final {
TVM_FFI_ICHECK(f_get_top_k != nullptr) << "PyDatabase's GetTopK method not implemented!";
return f_get_top_k(workload, top_k);
}
ffi::Array<TuningRecord> GetAllTuningRecords() final {
TVM_FFI_ICHECK(f_get_all_tuning_records != nullptr)
<< "PyDatabase's GetAllTuningRecords method not implemented!";
return f_get_all_tuning_records();
}
ffi::Optional<TuningRecord> QueryTuningRecord(const IRModule& mod, const Target& target,
const ffi::String& workload_name) final {
if (f_query_tuning_record == nullptr) {
return DatabaseNode::QueryTuningRecord(mod, target, workload_name);
} else {
return f_query_tuning_record(mod, target, workload_name);
}
}
ffi::Optional<s_tir::Schedule> QuerySchedule(const IRModule& mod, const Target& target,
const ffi::String& workload_name) final {
if (f_query_schedule == nullptr) {
return DatabaseNode::QuerySchedule(mod, target, workload_name);
} else {
return f_query_schedule(mod, target, workload_name);
}
}
ffi::Optional<IRModule> QueryIRModule(const IRModule& mod, const Target& target,
const ffi::String& workload_name) final {
if (f_query_ir_module == nullptr) {
return DatabaseNode::QueryIRModule(mod, target, workload_name);
} else {
return f_query_ir_module(mod, target, workload_name);
}
}
int64_t Size() final {
TVM_FFI_ICHECK(f_size != nullptr) << "PyDatabase's Size method not implemented!";
return f_size();
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyDatabase", PyDatabaseNode, DatabaseNode);
};
/*!
* \brief Managed reference to DatabaseNode.
* \sa DatabaseNode
*/
class Database : public ffi::ObjectRef {
public:
/*!
* \brief Constructor from ffi::ObjectPtr<DatabaseNode>.
* \param data The object pointer.
*/
explicit Database(ffi::ObjectPtr<DatabaseNode> data) : ffi::ObjectRef(data) {
TVM_FFI_ICHECK(data != nullptr);
}
/*!
* \brief An in-memory database.
* \param mod_eq_name A string to specify the module equality testing and hashing method.
*/
TVM_DLL static Database MemoryDatabase(ffi::String mod_eq_name = "structural");
/*!
* \brief A database for injecting handcrafted schedule functions.
* \param schedule_fn The function to do scheduling, which takes a TIR schedule,
* and returns a boolean indicating if the schedule is successful.
* \param mod_eq_name A string to specify the module equality testing and hashing method.
*/
TVM_DLL static Database ScheduleFnDatabase(ffi::TypedFunction<bool(s_tir::Schedule)> schedule_fn,
ffi::String mod_eq_name = "structural");
/*!
* \brief Create a default database that uses JSON file for tuning records.
* \param path_workload The path to the workload table.
* \param path_tuning_record The path to the database table.
* \param allow_missing Whether to create new file when the given path is not found.
* \param mod_eq_name A string to specify the module equality testing and hashing method.
*/
TVM_DLL static Database JSONDatabase(ffi::String path_workload, ffi::String path_tuning_record,
bool allow_missing, ffi::String mod_eq_name = "structural");
/*!
* \brief A database composed of multiple databases, allowing users to guide IR rewriting using
* combined knowledge of those databases. To each query, it returns the best record among all the
* databases given.
* \param databases The list of databases to be combined.
* \return The combined database.
*/
TVM_DLL static Database UnionDatabase(ffi::Array<Database, void> databases);
/*!
* \brief A database composed of multiple databases, allowing users to guide IR rewriting using
* combined knowledge of those databases. To each query, it returns the record from the first
* database that responds to the query.
* \param databases The database to be subsetted.
* \return The subsetted database.
*/
TVM_DLL static Database OrderedUnionDatabase(ffi::Array<Database, void> databases);
/*!
* \brief Create a database with customized methods on the python-side.
* \param f_has_workload The packed function of `HasWorkload`.
* \param f_commit_workload The packed function of `CommitWorkload`.
* \param f_commit_tuning_record The packed function of `CommitTuningRecord`.
* \param f_get_top_k The packed function of `GetTopK`.
* \param f_get_all_tuning_records The packed function of `GetAllTuningRecords`.
* \param f_query_tuning_record The packed function of `QueryTuningRecord`.
* \param f_query_schedule The packed function of `QuerySchedule`.
* \param f_query_ir_module The packed function of `QueryIRModule`.
* \param f_size The packed function of `Size`.
* \param mod_eq_name A string to specify the module equality testing and hashing method.
* \return The created database.
*/
TVM_DLL static Database PyDatabase(PyDatabaseNode::FHasWorkload f_has_workload,
PyDatabaseNode::FCommitWorkload f_commit_workload,
PyDatabaseNode::FCommitTuningRecord f_commit_tuning_record,
PyDatabaseNode::FGetTopK f_get_top_k,
PyDatabaseNode::FGetAllTuningRecords f_get_all_tuning_records,
PyDatabaseNode::FQueryTuningRecord f_query_tuning_record,
PyDatabaseNode::FQuerySchedule f_query_schedule,
PyDatabaseNode::FQueryIRModule f_query_ir_module,
PyDatabaseNode::FSize f_size,
ffi::String mod_eq_name = "structural");
/*! \return The current Database in the scope. */
static ffi::Optional<Database> Current();
/*! \brief Entering the scope of the context manager */
void EnterWithScope();
/*! \brief Exiting the scope of the context manager */
void ExitWithScope();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Database, ffi::ObjectRef, DatabaseNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_DATABASE_H_
include/tvm/s_tir/meta_schedule/extracted_task.h
+85
View File
@@ -0,0 +1,85 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_EXTRACTED_TASK_H_
#define TVM_S_TIR_META_SCHEDULE_EXTRACTED_TASK_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/module.h>
#include <tvm/target/target.h>
namespace tvm {
namespace tirx {
class PrimFunc;
} // namespace tirx
namespace te {
class Tensor;
} // namespace te
} // namespace tvm
namespace tvm {
namespace s_tir {
namespace meta_schedule {
/*! \brief A tuning task extracted from the high-level IR */
class ExtractedTaskNode : public ffi::Object {
public:
/*! \brief The name of the task extracted */
ffi::String task_name;
/*! \brief The high-level IR */
IRModule mod;
/*! \brief Target */
Target target;
/*! \brief A list of low-level IRs that the high-level IR could potentially dispatch to */
ffi::Array<IRModule> dispatched;
/*! \brief Weight of the task */
int weight;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExtractedTaskNode>()
.def_ro("task_name", &ExtractedTaskNode::task_name)
.def_ro("mod", &ExtractedTaskNode::mod)
.def_ro("target", &ExtractedTaskNode::target)
.def_ro("dispatched", &ExtractedTaskNode::dispatched)
.def_ro("weight", &ExtractedTaskNode::weight);
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.ExtractedTask", ExtractedTaskNode,
ffi::Object);
};
/*!
* \brief Managed reference to ExtractedTaskNode
* \sa ExtractedTaskNode
*/
class ExtractedTask : public ffi::ObjectRef {
public:
explicit ExtractedTask(ffi::String task_name, IRModule mod, Target target,
ffi::Array<IRModule> dispatched, int weight);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExtractedTask, ffi::ObjectRef, ExtractedTaskNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_EXTRACTED_TASK_H_
include/tvm/s_tir/meta_schedule/feature_extractor.h
+120
View File
@@ -0,0 +1,120 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_FEATURE_EXTRACTOR_H_
#define TVM_S_TIR_META_SCHEDULE_FEATURE_EXTRACTOR_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/tensor.h>
#include <tvm/s_tir/meta_schedule/measure_candidate.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
class TuneContext;
/*! \brief Extractor for features from measure candidates for use in cost model. */
class FeatureExtractorNode : public ffi::Object {
public:
/*! \brief Virtual destructor. */
virtual ~FeatureExtractorNode() = default;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FeatureExtractorNode>();
}
/*!
* \brief Extract features from the given measure candidate.
* \param context The tuning context for feature extraction.
* \param candidates The measure candidates to extract features from.
* \return The feature tensor extracted.
*/
virtual ffi::Array<tvm::runtime::Tensor> ExtractFrom(
const TuneContext& context, const ffi::Array<MeasureCandidate>& candidates) = 0;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.FeatureExtractor", FeatureExtractorNode,
ffi::Object);
};
/*! \brief The feature extractor with customized methods on the python-side. */
class PyFeatureExtractorNode : public FeatureExtractorNode {
public:
/*!
* \brief Extract features from the given measure candidate.
* \param context The tuning context for feature extraction.
* \param candidates The measure candidates to extract features from.
* \return The feature tensor extracted.
*/
using FExtractFrom = ffi::TypedFunction<ffi::Array<tvm::runtime::Tensor>(
const TuneContext& context, const ffi::Array<MeasureCandidate>& candidates)>;
/*! \brief The packed function to the `ExtractFrom` function. */
FExtractFrom f_extract_from;
static void RegisterReflection() {
// `f_extract_from` is not registered
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PyFeatureExtractorNode>();
}
ffi::Array<tvm::runtime::Tensor> ExtractFrom(
const TuneContext& context, const ffi::Array<MeasureCandidate>& candidates) final;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyFeatureExtractor",
PyFeatureExtractorNode, FeatureExtractorNode);
};
/*!
* \brief Managed reference to FeatureExtractorNode
* \sa FeatureExtractorNode
*/
class FeatureExtractor : public ffi::ObjectRef {
public:
/*!
* \brief Create a feature extractor that extracts features from each BufferStore
* \param buffers_per_store The number of buffers in each BufferStore; Pad or truncate if
* necessary.
* \param arith_intensity_curve_num_samples The number of samples used in the arithmetic intensity
* curve.
* \param cache_line_bytes The number of bytes in a cache line.
* \param extract_workload Whether to extract features in the workload in tuning context or not.
* \return The feature extractor created.
*/
TVM_DLL static FeatureExtractor PerStoreFeature(int buffers_per_store = 5,
int arith_intensity_curve_num_samples = 10,
int cache_line_bytes = 64,
bool extract_workload = false);
/*!
* \brief Create a feature extractor with customized methods on the python-side.
* \param f_extract_from The packed function of `ExtractFrom`.
* \return The feature extractor created.
*/
TVM_DLL static FeatureExtractor PyFeatureExtractor(
PyFeatureExtractorNode::FExtractFrom f_extract_from);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FeatureExtractor, ffi::ObjectRef,
FeatureExtractorNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_FEATURE_EXTRACTOR_H_
include/tvm/s_tir/meta_schedule/measure_callback.h
+141
View File
@@ -0,0 +1,141 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_MEASURE_CALLBACK_H_
#define TVM_S_TIR_META_SCHEDULE_MEASURE_CALLBACK_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/base.h>
#include <tvm/s_tir/meta_schedule/builder.h>
#include <tvm/s_tir/meta_schedule/measure_candidate.h>
#include <tvm/s_tir/meta_schedule/runner.h>
#include <tvm/s_tir/meta_schedule/search_strategy.h>
#include <tvm/s_tir/meta_schedule/tune_context.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
class TaskScheduler;
/*! \brief Rules to apply after measure results is available. */
class MeasureCallbackNode : public ffi::Object {
public:
/*! \brief Virtual destructor. */
virtual ~MeasureCallbackNode() = default;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MeasureCallbackNode>();
}
/*!
* \brief Apply a measure callback rule with given arguments.
* \param task_scheduler The task scheduler.
* \param task_id The id of the task (tune context) to apply measure callbacks.
* \param measure_candidates The measure candidates.
* \param builder_results The builder results by building the measure candidates.
* \param runner_results The runner results by running the built measure candidates.
*/
virtual void Apply(const TaskScheduler& task_scheduler, //
int task_id, //
const ffi::Array<MeasureCandidate>& measure_candidates, //
const ffi::Array<BuilderResult>& builder_results, //
const ffi::Array<RunnerResult>& runner_results) = 0;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.MeasureCallback", MeasureCallbackNode,
ffi::Object);
};
/*! \brief The measure callback with customized methods on the python-side. */
class PyMeasureCallbackNode : public MeasureCallbackNode {
public:
/*!
* \brief Apply a measure callback to the given schedule.
* \param task_scheduler The task scheduler.
* \param tasks The list of tune context to process.
* \param measure_candidates The measure candidates.
* \param builds The builder results by building the measure candidates.
* \param results The runner results by running the built measure candidates.
* \return Whether the measure callback was successfully applied.
*/
using FApply = ffi::TypedFunction<void(const TaskScheduler& task_scheduler, //
int task_id, //
const ffi::Array<MeasureCandidate>& measure_candidates, //
const ffi::Array<BuilderResult>& builds, //
const ffi::Array<RunnerResult>& results)>;
/*! \brief The packed function to the `Apply` function. */
FApply f_apply;
static void RegisterReflection() {
// `f_apply` is not registered
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PyMeasureCallbackNode>();
}
void Apply(const TaskScheduler& task_scheduler, //
int task_id, //
const ffi::Array<MeasureCandidate>& measure_candidates, //
const ffi::Array<BuilderResult>& builds, //
const ffi::Array<RunnerResult>& results);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyMeasureCallback", PyMeasureCallbackNode,
MeasureCallbackNode);
};
/*!
* \brief Managed reference to MeasureCallbackNode
* \sa MeasureCallbackNode
*/
class MeasureCallback : public ffi::ObjectRef {
public:
/*!
* \brief Create a measure callback that adds the measurement results into the database
* \return The measure callback created.
*/
TVM_DLL static MeasureCallback AddToDatabase();
/*!
* \brief Create a measure callback that removes the build artifacts from the disk
* \return The measure callback created.
*/
TVM_DLL static MeasureCallback RemoveBuildArtifact();
/*!
* \brief Create a measure callback that updates the cost model with measurement result.
* \return The measure callback created.
*/
TVM_DLL static MeasureCallback UpdateCostModel();
/*!
* \brief Create a measure callback with customized methods on the python-side.
* \param f_apply The packed function of `Apply`.
* \return The measure callback created.
*/
TVM_DLL static MeasureCallback PyMeasureCallback(PyMeasureCallbackNode::FApply f_apply);
/*! \brief The default list of measure callbacks. */
TVM_DLL static ffi::Array<MeasureCallback, void> Default();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MeasureCallback, ffi::ObjectRef, MeasureCallbackNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_MEASURE_CALLBACK_H_
include/tvm/s_tir/meta_schedule/measure_candidate.h
+71
View File
@@ -0,0 +1,71 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_MEASURE_CANDIDATE_H_
#define TVM_S_TIR_META_SCHEDULE_MEASURE_CANDIDATE_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/base.h>
#include <tvm/s_tir/meta_schedule/arg_info.h>
#include <tvm/s_tir/schedule/schedule.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
/*! \brief The schedule (with input shapes) to be measured. */
class MeasureCandidateNode : public ffi::Object {
public:
/*! \brief The schedule for measurement. */
s_tir::Schedule sch;
/*! \brief The argument information, e.g., (shape, dtype) for tensors. */
ffi::Array<ArgInfo> args_info;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MeasureCandidateNode>()
.def_ro("sch", &MeasureCandidateNode::sch)
.def_ro("args_info", &MeasureCandidateNode::args_info);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.MeasureCandidate", MeasureCandidateNode,
ffi::Object);
};
/*!
* \brief Managed reference to MeasureCandidateNode.
* \sa MeasureCandidateNode
*/
class MeasureCandidate : public ffi::ObjectRef {
public:
/*!
* \brief Constructor of MeasureCandidate.
* \param sch The schedule for measurement.
* \param args_info The argument information, e.g., (shape, dtype) for tensors.
*/
TVM_DLL MeasureCandidate(s_tir::Schedule sch, ffi::Array<ArgInfo> args_info);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MeasureCandidate, ffi::ObjectRef,
MeasureCandidateNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_MEASURE_CANDIDATE_H_
include/tvm/s_tir/meta_schedule/mutator.h
+174
View File
@@ -0,0 +1,174 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_MUTATOR_H_
#define TVM_S_TIR_META_SCHEDULE_MUTATOR_H_
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/base.h>
#include <tvm/s_tir/random_engine.h>
#include <tvm/s_tir/schedule/schedule.h>
#include <tvm/s_tir/schedule/trace.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
class TuneContext;
class Mutator;
/*! \brief Mutator is designed to mutate the trace to explore the design space. */
class MutatorNode : public ffi::Object {
public:
/*! \brief Virtual destructor. */
virtual ~MutatorNode() = default;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MutatorNode>();
}
/*!
* \brief Initialize the design space generator with tuning context.
* \param context The tuning context for initialization.
* \note This method is supposed to be called only once before every other method.
*/
virtual void InitializeWithTuneContext(const TuneContext& context) = 0;
/*!
* \brief Apply the mutator function to the given trace.
* \param trace The given trace for mutation.
* \param rand_state The random state for mutation.
* \return None if mutator failed, otherwise return the mutated trace.
*/
virtual ffi::Optional<s_tir::Trace> Apply(const s_tir::Trace& trace,
LinearCongruentialEngine::TRandState* rand_state) = 0;
/*!
* \brief Clone the mutator.
* \return The cloned mutator.
*/
virtual Mutator Clone() const = 0;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.Mutator", MutatorNode, ffi::Object);
};
/*!
* \brief Managed reference to MutatorNode
* \sa MutatorNode
*/
class Mutator : public ffi::ObjectRef {
public:
/*!
* \brief The function type of `InitializeWithTuneContext` method.
* \param context The tuning context for initialization.
*/
using FInitializeWithTuneContext = ffi::TypedFunction<void(const TuneContext&)>;
/*!
* \brief Apply the mutator function to the given trace.
* \param trace The given trace for mutation.
* \return None if mutator failed, otherwise return the mutated trace.
*/
using FApply = ffi::TypedFunction<ffi::Optional<s_tir::Trace>(
const s_tir::Trace&, LinearCongruentialEngine::TRandState rand_state)>;
/*!
* \brief Clone the mutator.
* \return The cloned mutator.
*/
using FClone = ffi::TypedFunction<Mutator()>;
/*! \brief Create a Mutator that mutates the decision of instruction Sample-Perfect-Tile */
TVM_DLL static Mutator MutateTileSize();
/*!
* \brief Create a Mutator that mutates the parallel extent
* \param max_jobs_per_core The maximum number of parallel jobs per core.
* \return The created mutator.
*/
TVM_DLL static Mutator MutateParallel(int64_t max_jobs_per_core);
/*!
* \brief Create a Mutator that mutates auto unroll step
* \return The mutator created
*/
TVM_DLL static Mutator MutateUnroll();
/*!
* \brief Create a Mutator that mutates the outcome of SampleComputeLocation
* \return The mutator created
*/
TVM_DLL static Mutator MutateComputeLocation();
/*!
* \brief Create a Mutator that mutates auto thread binding.
* \return The mutator created
*/
TVM_DLL static Mutator MutateThreadBinding();
/*!
* \brief Create a mutator with customized methods on the python-side.
* \param f_initialize_with_tune_context The packed function of `InitializeWithTuneContext`.
* \param f_apply The packed function of `Apply`.
* \param f_clone The packed function of `Clone`.
* \return The mutator created.
*/
TVM_DLL static Mutator PyMutator(FInitializeWithTuneContext f_initialize_with_tune_context,
FApply f_apply, FClone f_clone);
/*! \brief Create default mutators for LLVM */
TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultLLVM();
/*! \brief Create default mutators for CUDA */
TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultCUDA();
/*! \brief Create default mutators for CUDA with TensorCore */
TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultCUDATensorCore();
/*! \brief Create default mutators for Hexagon */
TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultHexagon();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Mutator, ffi::ObjectRef, MutatorNode);
};
/*! \brief The mutator with customized methods on the python-side. */
class PyMutatorNode : public MutatorNode {
public:
using FInitializeWithTuneContext = Mutator::FInitializeWithTuneContext;
using FApply = Mutator::FApply;
using FClone = Mutator::FClone;
/*! \brief The packed function to the `InitializeWithTuneContext` function. */
FInitializeWithTuneContext f_initialize_with_tune_context;
/*! \brief The packed function to the `Apply` function. */
FApply f_apply;
/*! \brief The packed function to the `Clone` function. */
FClone f_clone;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PyMutatorNode>();
// `f_initialize_with_tune_context` is not registered
// `f_apply` is not registered
// `f_clone` is not registered
}
void InitializeWithTuneContext(const TuneContext& context) final;
ffi::Optional<s_tir::Trace> Apply(const s_tir::Trace& trace,
LinearCongruentialEngine::TRandState* rand_state) final;
Mutator Clone() const final;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyMutator", PyMutatorNode, MutatorNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_MUTATOR_H_
include/tvm/s_tir/meta_schedule/postproc.h
+207
View File
@@ -0,0 +1,207 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_POSTPROC_H_
#define TVM_S_TIR_META_SCHEDULE_POSTPROC_H_
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/base.h>
#include <tvm/s_tir/schedule/schedule.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
class TuneContext;
class Postproc;
/*!
* \brief Rules to apply a postprocessor to a schedule.
*/
class PostprocNode : public ffi::Object {
public:
/*! \brief Virtual destructor. */
virtual ~PostprocNode() = default;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PostprocNode>();
}
/*!
* \brief Initialize the design space generator with tuning context.
* \param context The tuning context for initialization.
* \note This method is supposed to be called only once before every other method.
*/
virtual void InitializeWithTuneContext(const TuneContext& context) = 0;
/*!
* \brief Apply a postprocessor to the given schedule.
* \param sch The schedule to be post processed.
* \return Whether the postprocessor was successfully applied.
*/
virtual bool Apply(const s_tir::Schedule& sch) = 0;
/*!
* \brief Clone the postprocessor.
* \return The cloned postprocessor.
*/
virtual Postproc Clone() const = 0;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.Postproc", PostprocNode, ffi::Object);
};
/*!
* \brief Managed reference to PostprocNode
* \sa PostprocNode
*/
class Postproc : public ffi::ObjectRef {
public:
/*!
* \brief The function type of `InitializeWithTuneContext` method.
* \param context The tuning context for initialization.
*/
using FInitializeWithTuneContext = ffi::TypedFunction<void(const TuneContext&)>;
/*!
* \brief Apply a postprocessor to the given schedule.
* \param sch The schedule to be post processed.
* \return Whether the postprocessor was successfully applied.
*/
using FApply = ffi::TypedFunction<bool(const s_tir::Schedule&)>;
/*!
* \brief Clone the postprocessor.
* \return The cloned postprocessor.
*/
using FClone = ffi::TypedFunction<Postproc()>;
/*!
* \brief Create a postprocessor with customized methods on the python-side.
* \param f_initialize_with_tune_context The packed function of `InitializeWithTuneContext`.
* \param f_apply The packed function of `Apply`.
* \param f_clone The packed function of `Clone`.
* \return The postprocessor created.
*/
TVM_DLL static Postproc PyPostproc(FInitializeWithTuneContext f_initialize_with_tune_context, //
FApply f_apply, //
FClone f_clone);
/*!
* \brief Create a postprocessor that checks if all loops are static
* \return The postprocessor created
*/
TVM_DLL static Postproc DisallowDynamicLoop();
/*!
* \brief Create a postprocessor that checks if all async mem copies are not strided.
* \return The postprocessor created
*/
TVM_DLL static Postproc DisallowAsyncStridedMemCopy();
/*!
* \brief Create a postprocessor that rewrites the cooperative fetch annotation to
* actual vectorized cooperative fetching in loop bindings.
* \return The postprocessor created.
*/
TVM_DLL static Postproc RewriteCooperativeFetch();
/*!
* \brief Creates a postprocessor that applies parallelization, vectorization and auto unrolling
* according to the annotation of each block
* \return The postprocessor created
*/
TVM_DLL static Postproc RewriteParallelVectorizeUnroll();
/*!
* \brief Create a postprocessor that rewrites reduction block by moving the init block out.
* \return The postprocessor created.
*/
TVM_DLL static Postproc RewriteReductionBlock();
/*!
* \brief Create a postprocessor that adds thread binding to unbound blocks
* \param max_threadblocks The max number of threadblocks in the CUDA device.
* \return The postprocessor created.
*/
TVM_DLL static Postproc RewriteUnboundBlock(int max_threadblocks);
/*!
* \brief Create a postprocessor that applies tensorization to annotated blocks
* \param vectorize_init_loop Whether or not vectorize the initialization loop produced by
* DecomposeReduction
* \return The postprocessor created.
*/
TVM_DLL static Postproc RewriteTensorize(bool vectorize_init_loop = false);
/*!
* \brief Creates a postprocessor that verifies if the GPU code is correct
* \return The postprocessor created
*/
TVM_DLL static Postproc VerifyGPUCode();
/*!
* \brief Verifies that the VTCM usage of a given schedule is within the provided limit.
* \return The postprocessor created
*/
TVM_DLL static Postproc VerifyVTCMLimit();
/*!
* \brief Creates a postprocessor that rewrites the layout of input tensor
* \note Weight layout rewrite is supported so far, activation layout rewrite will be added.
* \return The postprocessor created
*/
TVM_DLL static Postproc RewriteLayout();
/*! \brief Create default postprocessors for LLVM */
TVM_DLL static ffi::Array<Postproc, void> DefaultLLVM();
/*! \brief Create default postprocessors for x86 (AVX512 and VNNI) */
TVM_DLL static ffi::Array<Postproc, void> DefaultCPUTensorization();
/*! \brief Create default postprocessors for RISCV */
TVM_DLL static ffi::Array<Postproc, void> DefaultRISCV();
/*! \brief Create default postprocessors for CUDA */
TVM_DLL static ffi::Array<Postproc, void> DefaultCUDA();
/*! \brief Create default postprocessors for CUDA with TensorCore */
TVM_DLL static ffi::Array<Postproc, void> DefaultCUDATensorCore();
/*! \brief Create default postprocessors for Hexagon */
TVM_DLL static ffi::Array<Postproc, void> DefaultHexagon();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Postproc, ffi::ObjectRef, PostprocNode);
};
/*! \brief The postprocessor with customized methods on the python-side. */
class PyPostprocNode : public PostprocNode {
public:
using FInitializeWithTuneContext = Postproc::FInitializeWithTuneContext;
using FApply = Postproc::FApply;
using FClone = Postproc::FClone;
/*! \brief The packed function to the `InitializeWithTuneContext` function. */
FInitializeWithTuneContext f_initialize_with_tune_context;
/*! \brief The packed function to the `Apply` function. */
FApply f_apply;
/*! \brief The packed function to the `Clone` function. */
FClone f_clone;
static void RegisterReflection() {
// `f_initialize_with_tune_context` is not registered
// `f_apply` is not registered
// `f_clone` is not registered
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PyPostprocNode>();
}
void InitializeWithTuneContext(const TuneContext& context) final;
bool Apply(const s_tir::Schedule& sch) final;
Postproc Clone() const final;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyPostproc", PyPostprocNode, PostprocNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_POSTPROC_H_
include/tvm/s_tir/meta_schedule/profiler.h
+104
View File
@@ -0,0 +1,104 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_PROFILER_H_
#define TVM_S_TIR_META_SCHEDULE_PROFILER_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/module.h>
#include <tvm/target/target.h>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
class ScopedTimer {
public:
~ScopedTimer() {
if (deferred_ != nullptr) {
deferred_();
}
}
private:
friend class Profiler;
explicit ScopedTimer(ffi::TypedFunction<void()> deferred) : deferred_(deferred) {}
ffi::TypedFunction<void()> deferred_;
};
/*! \brief A generic profiler */
class ProfilerNode : public ffi::Object {
public:
/*! \brief The segments that are already profiled */
std::unordered_map<std::string, double> stats_sec;
/*! \brief Counter for the total time used */
ffi::Function total_timer;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ProfilerNode>();
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.Profiler", ProfilerNode, ffi::Object);
public:
/*! \brief Get the internal stats of the running time */
ffi::Map<ffi::String, FloatImm> Get() const;
/*! \brief Return a summary of profiling results as table format */
ffi::String Table() const;
};
/*!
* \brief Managed reference to ProfilerNode
* \sa ProfilerNode
*/
class Profiler : public ffi::ObjectRef {
public:
Profiler();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Profiler, ffi::ObjectRef, ProfilerNode);
/*! \brief Entering the scope of the context manager */
void EnterWithScope();
/*! \brief Exiting the scope of the context manager */
void ExitWithScope();
/*! \brief Returns the current profiler */
static ffi::Optional<Profiler> Current();
/*!
* \brief Profile the time usage in the given scope in the given name.
* \param name Name for the scope.
* \return A scope timer for time profiling.
*/
static ScopedTimer TimedScope(ffi::String name);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_PROFILER_H_
include/tvm/s_tir/meta_schedule/runner.h
+250
View File
@@ -0,0 +1,250 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_RUNNER_H_
#define TVM_S_TIR_META_SCHEDULE_RUNNER_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
#include <tvm/s_tir/meta_schedule/arg_info.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
/*! \brief Runner's input containing path of artifact, type of device and argument info. */
class RunnerInputNode : public ffi::Object {
public:
/*! \brief The path to the built artifact. */
ffi::String artifact_path;
/*! \brief The type of device. */
ffi::String device_type;
/*! \brief The argument information. */
ffi::Array<ArgInfo> args_info;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RunnerInputNode>()
.def_ro("artifact_path", &RunnerInputNode::artifact_path)
.def_ro("device_type", &RunnerInputNode::device_type)
.def_ro("args_info", &RunnerInputNode::args_info);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.RunnerInput", RunnerInputNode,
ffi::Object);
};
/*!
* \brief Managed reference to RunnerInputNode
* \sa RunnerInputNode
*/
class RunnerInput : public ffi::ObjectRef {
public:
/*!
* \brief Constructor of RunnerInput
* \param artifact_path The path to the built artifact.
* \param device_type The type of device.
* \param args_info The argument information.
*/
TVM_DLL explicit RunnerInput(ffi::String artifact_path, ffi::String device_type,
ffi::Array<ArgInfo> args_info);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RunnerInput, ffi::ObjectRef, RunnerInputNode);
};
/*! \brief Runner's output containing measurement result of MeasureCandidate or error msg if any. */
class RunnerResultNode : public ffi::Object {
public:
/*! \brief The run time in seconds.*/
ffi::Optional<ffi::Array<FloatImm>> run_secs;
/*! \brief The error message, if any. */
ffi::Optional<ffi::String> error_msg;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RunnerResultNode>()
.def_ro("run_secs", &RunnerResultNode::run_secs)
.def_ro("error_msg", &RunnerResultNode::error_msg);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.RunnerResult", RunnerResultNode,
ffi::Object);
};
/*!
* \brief Managed reference to RunnerResultNode
* \sa RunnerResultNode
*/
class RunnerResult : public ffi::ObjectRef {
public:
/*!
* \brief Constructor
* \brief The run time in seconds.
* \brief The error message, if any.
*/
TVM_DLL explicit RunnerResult(ffi::Optional<ffi::Array<FloatImm>> run_secs,
ffi::Optional<ffi::String> error_msg);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RunnerResult, ffi::ObjectRef, RunnerResultNode);
};
/*!
* \brief A class to asynchronously fetch runner's output.
* \note The API design is consistent with python's concurrent.futures.Future:
* https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.Future
*/
class RunnerFutureNode : public ffi::Object {
public:
/*!
* \brief The function type to check whether the runner has finished.
* \return Whether the runner's output is ready.
*/
using FDone = ffi::TypedFunction<bool()>;
/*!
* \brief The function type to fetch runner output if it is ready.
* \return The runner's output.
*/
using FResult = ffi::TypedFunction<RunnerResult()>;
/*! \brief The packed function to check whether the runner has finished. */
FDone f_done;
/*! \brief The packed function to fetch runner output if it is ready. */
FResult f_result;
static void RegisterReflection() {
// `f_done` is not registered
// `f_result` is not registered
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RunnerFutureNode>();
}
/*!
* \brief Check whether the runner has finished.
* \return A boolean indicating whether the runner has finished.
*/
bool Done() const {
TVM_FFI_ICHECK(f_done != nullptr) << "PyRunnerFuture's Done method not implemented!";
return f_done();
}
/*!
* \brief Fetch the runner's output if it is ready.
* \return The runner's output.
*/
RunnerResult Result() const {
TVM_FFI_ICHECK(f_result != nullptr) << "PyRunnerFuture's Result method not implemented!";
return f_result();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.RunnerFuture", RunnerFutureNode,
ffi::Object);
};
/*!
* \brief Managed reference to RunnerFutureNode
* \sa RunnerFutureNode
*/
class RunnerFuture : public ffi::ObjectRef {
public:
using FDone = RunnerFutureNode::FDone;
using FResult = RunnerFutureNode::FResult;
/*!
* \brief Constructor of RunnerFuture
* \param f_done The packed function to check whether the runner has finished.
* \param f_result The packed function to fetch runner output if it is ready.
*/
TVM_DLL explicit RunnerFuture(FDone f_done, FResult f_result);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RunnerFuture, ffi::ObjectRef, RunnerFutureNode);
};
/*! \brief The abstract runner interface. */
class RunnerNode : public ffi::Object {
public:
/*!
* \brief The function type to run the built artifacts and get runner futures.
* \param input The runner's inputs.
* \return The runner futures.
* \sa RunnerFuture
*/
using FRun = ffi::TypedFunction<ffi::Array<RunnerFuture>(ffi::Array<RunnerInput>)>;
/*! \brief Default destructor */
virtual ~RunnerNode() = default;
/*!
* \brief Run the built artifact and get runner futures.
* \param runner_inputs The runner's inputs.
* \return The runner futures.
*/
virtual ffi::Array<RunnerFuture> Run(ffi::Array<RunnerInput> runner_inputs) = 0;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RunnerNode>();
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.Runner", RunnerNode, ffi::Object);
};
/*!
* \brief Managed reference to RunnerNode
* \sa RunnerNode
*/
class Runner : public ffi::ObjectRef {
public:
using FRun = RunnerNode::FRun;
/*!
* \brief Constructor from ffi::ObjectPtr<RunnerNode>.
* \param data The object pointer.
*/
explicit Runner(ffi::ObjectPtr<RunnerNode> data) : ffi::ObjectRef(data) {
TVM_FFI_ICHECK(data != nullptr);
}
/*!
* \brief Create a runner with customized build method on the python-side.
* \param f_run The packed function to run the built artifacts and get runner futures.
* \return The runner created.
*/
TVM_DLL static Runner PyRunner(FRun f_run);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Runner, ffi::ObjectRef, RunnerNode);
};
/*! \brief An abstract runner with customized build method on the python-side. */
class PyRunnerNode : public RunnerNode {
public:
/*! \brief The packed function to run the built artifacts and get runner futures. */
FRun f_run;
static void RegisterReflection() {
// `f_run` is not registered
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PyRunnerNode>();
}
ffi::Array<RunnerFuture> Run(ffi::Array<RunnerInput> runner_inputs) final {
TVM_FFI_ICHECK(f_run != nullptr) << "PyRunner's Run method not implemented!";
return f_run(runner_inputs);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyRunner", PyRunnerNode, RunnerNode);
};
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_RUNNER_H_
include/tvm/s_tir/meta_schedule/schedule/cpu/.gitignore
View File
include/tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h
+72
View File
@@ -0,0 +1,72 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_SCHEDULE_CUDA_THREAD_BIND_H_
#define TVM_S_TIR_META_SCHEDULE_SCHEDULE_CUDA_THREAD_BIND_H_
#include <tvm/s_tir/schedule/schedule.h>
#include <algorithm>
#include <limits>
#include <utility>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
/*!
* \brief Given candidates of thread_extents, make a sampler that use `sch->SampleCategorical`
* to return a random thread extent.
* \param sch The schedule
* \param thread_extents The candidate thread extents.
* \return A sampler that returns a random thread extent.
*/
std::function<s_tir::ExprRV(int64_t)> MakeFactorSampler(s_tir::Schedule sch,
ffi::Array<int64_t> thread_extents);
/*!
* \brief Bind blockIdx.x and threadIdx.x to the given loop
* \param sch The schedule.
* \param loop The loop to be bound.
* \param max_threadblocks The maximum number of threadblocks allowed.
* \param max_threads_per_block The maximum number of threads allowed.
* \param get_factor A function that returns the tiling factor.
* \return The binded loops in the order of blockIdx.x, threadIdx.x, and the rest.
*/
ffi::Array<s_tir::LoopRV> BindSpatialLoop(
s_tir::Schedule sch, s_tir::LoopRV loop, //
int64_t max_threadblocks, int64_t max_threads_per_block,
std::function<s_tir::ExprRV(int64_t)> get_factor = nullptr);
/*!
* \brief Bind the given block if it is not bound to blockIdx or threadIdx.
* \param sch The schedule.
* \param block The block to be bound.
* \param max_threadblocks The maximum number of threadblocks allowed.
* \param max_threads_per_block The maximum number of threads allowed.
* \param get_factor A function that returns the tiling factor.
*/
void BindBlockThreadIdx(s_tir::Schedule sch, s_tir::SBlockRV block, //
int64_t max_threadblocks, int64_t max_threads_per_block,
std::function<s_tir::ExprRV(int64_t max_extent)> get_factor = nullptr);
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_SCHEDULE_CUDA_THREAD_BIND_H_
include/tvm/s_tir/meta_schedule/schedule/generic/winograd.h
+39
View File
@@ -0,0 +1,39 @@
/*
* 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.
*/
#ifndef TVM_S_TIR_META_SCHEDULE_SCHEDULE_GENERIC_WINOGRAD_H_
#define TVM_S_TIR_META_SCHEDULE_SCHEDULE_GENERIC_WINOGRAD_H_
#include <tvm/s_tir/schedule/schedule.h>
namespace tvm {
namespace s_tir {
namespace meta_schedule {
/*!
* \brief Get the producer block of a given block.
* If there is a constant winograd transform matrix, inline it.
* \return The only producer block.
*/
s_tir::SBlockRV GetWinogradProducerAndInlineConst(s_tir::Schedule sch, s_tir::SBlockRV block);
} // namespace meta_schedule
} // namespace s_tir
} // namespace tvm
#endif // TVM_S_TIR_META_SCHEDULE_SCHEDULE_GENERIC_WINOGRAD_H_
include/tvm/s_tir/meta_schedule/schedule/x86/.gitignore
View File

Some files were not shown because too many files have changed in this diff Show More