tvm
|
Tensor expression language DSL. More...
Classes | |
class | BaseComputeOpNode |
A Compute op that compute a tensor on certain domain. This is the base class for ComputeOp (operating on a scalar at a time) and TensorComputeOp (operating on a TensorSlice at a time) More... | |
class | ComputeOp |
Managed reference to ComputeOpNode. More... | |
class | ComputeOpNode |
A Compute op that compute a tensor on certain domain. More... | |
class | ExternOp |
Managed reference to ExternOpNode. More... | |
class | ExternOpNode |
External computation that cannot be splitted. More... | |
class | Fuse |
Managed reference to FuseNode. More... | |
class | FuseNode |
Fuse two domains into one domain. More... | |
class | HybridOp |
Managed reference to HybridOpNode. More... | |
class | HybridOpNode |
A computation operator that generated by hybrid script. More... | |
class | IterVarAttr |
Additional scheduable attributes about IterVar. More... | |
class | IterVarAttrNode |
node container for IterVar attr More... | |
class | IterVarRelation |
The schedule relation between IterVars can be Split, Fuse. More... | |
class | IterVarRelationNode |
base node of iteration var More... | |
class | Operation |
Operation that produces tensors. More... | |
class | OperationNode |
Base class of all operation nodes. More... | |
class | PlaceholderOp |
Managed reference to PlaceholderOpNode. More... | |
class | PlaceholderOpNode |
A placeholder op represents an input placeholder. More... | |
class | Rebase |
Managed reference to RebaseNode. More... | |
class | RebaseNode |
Rebase the iteration to make min to be 0. This is useful to normalize the Schedule to make every leaf variable's min to be 0. More... | |
class | ScanOp |
Managed reference to ScanOpNode. More... | |
class | ScanOpNode |
Symbolic scan. More... | |
class | Schedule |
Global schedule container For operations and all the operations they depend on. The schedule per Operation is named as stage. More... | |
class | ScheduleContext |
Context helper to collect debug information of Schedule. More... | |
class | ScheduleNode |
node container for schedule More... | |
class | Singleton |
Managed reference to SingletonNode. More... | |
class | SingletonNode |
Singleton iterator [0, 1) More... | |
class | SpecializedCondition |
Specialized condition to enable op specialization. More... | |
class | SpecializedConditionNode |
Container for specialization conditions. More... | |
class | Split |
Managed reference to SplitNode. More... | |
class | SplitNode |
Split the parent domain into product of outer and iter. More... | |
class | Stage |
Stage, contains scheduling for a stage of computation. More... | |
class | StageNode |
represents a stage. More... | |
class | Tensor |
Tensor structure representing a possible input, or intermediate computation result. More... | |
class | TensorComputeOp |
Managed reference to TensorComputeOpNode. More... | |
class | TensorComputeOpNode |
A TenorCompute op that compute a tensor with an tensor intrinsic. More... | |
struct | TensorDom |
Temporary data structure to store union of bounds of each axis of Tensor. More... | |
class | TensorIntrin |
Managed reference to TensorIntrinNode. More... | |
class | TensorIntrinCall |
Managed reference to TensorIntrinCallNode. More... | |
class | TensorIntrinCallNode |
class | TensorIntrinNode |
Node to represent a Tensor intrinsic operator. More... | |
class | TensorNode |
Node to represent a tensor. More... | |
class | Transform |
class | TransformNode |
Transform iterator according to some arbitrary expression. More... | |
Typedefs | |
using | FCompute = std::function< PrimExpr(const Array< Var > &i)> |
The compute function to specify the input source of a Tensor. More... | |
using | FBatchCompute = std::function< Array< PrimExpr >(const Array< Var > &i)> |
The compute function to specify the inputs source of Tensors. More... | |
Enumerations | |
enum | AttachType : int { kGroupRoot = 1, kInline = 2, kInlinedAlready = 3, kScope = 4, kScanUpdate = 5 } |
the attachment type More... | |
Functions | |
PrimExpr | Derivative (const PrimExpr &expr, const Var &var) |
Take the derivative of the expression with respect to the given variable. More... | |
Tensor | Jacobian (const Tensor &output, const Tensor &input) |
Get the tensor representing the Jacobian of the output with respect to the input. More... | |
Tensor | VectorJacobianProduct (const Tensor &output, const Tensor &input, const Tensor &head) |
The building block for reverse-mode AD. More... | |
Array< Tensor > | Gradient (const Tensor &output, const Array< Tensor > &inputs, const Tensor &head=Tensor()) |
Perform reverse mode automatic differentiation. More... | |
Var | var (std::string name_hint, DataType t=DataType::Int(32)) |
Construct a new Var expression. More... | |
IterVar | thread_axis (Range dom, std::string tag) |
Create a new IterVar that represents an axis in thread. More... | |
IterVar | reduce_axis (Range dom, std::string name="rv") |
Create a new IterVar for reduction operations. More... | |
Tensor | placeholder (Array< PrimExpr > shape, DataType dtype=DataType::Float(32), std::string name="placeholder") |
create a place holder tensor. More... | |
Tensor | compute (Array< PrimExpr > shape, FCompute fcompute, std::string name="tensor", std::string tag="", Map< String, ObjectRef > attrs={}) |
Construct a new tensor by computing over shape, using the computation rule: result_tensor[axis] = fcompute(axis) More... | |
Array< Tensor > | compute (Array< PrimExpr > shape, FBatchCompute fcompute, std::string name="tensor", std::string tag="", Map< String, ObjectRef > attrs={}) |
Construct a new tensor by computing over shape, using the computation rule: result_tensor[axis] = fcompute(axis) More... | |
Array< Tensor > | scan (Array< Tensor > init, Array< Tensor > update, Array< Tensor > state_placeholder, Array< Tensor > inputs=Array< Tensor >(), std::string name="scan", std::string tag="", Map< String, ObjectRef > attrs={}) |
Construct new tensors by scan. More... | |
Tensor | compute (Array< PrimExpr > shape, std::function< PrimExpr(Var)> f, std::string name="tensor", std::string tag="", Map< String, ObjectRef > attrs={}) |
Tensor | compute (Array< PrimExpr > shape, std::function< PrimExpr(Var, Var)> f, std::string name="tensor", std::string tag="", Map< String, ObjectRef > attrs={}) |
Tensor | compute (Array< PrimExpr > shape, std::function< PrimExpr(Var, Var, Var)> f, std::string name="tensor", std::string tag="", Map< String, ObjectRef > attrs={}) |
Tensor | compute (Array< PrimExpr > shape, std::function< PrimExpr(Var, Var, Var, Var)> f, std::string name="tensor", std::string tag="", Map< String, ObjectRef > attrs={}) |
Schedule | create_schedule (Array< Operation > ops) |
Create a schedule for array of ops(and their dependencies). More... | |
void | AutoInlineElemWise (Schedule sch) |
To automatically inline the element-wise operations. More... | |
void | AutoInlineBroarcast (Schedule sch) |
To automatically inline the broadcast operations. More... | |
void | AutoInlineInjective (Schedule sch) |
To automatically inline operations with injective writes (i.e. writes without reduction or sequential loops). Note that in this case, guarantees about contiguity, transpose, stride, alignemnt and memory footprint in general do not hold. More... | |
Map< IterVar, Range > | InferBound (const Schedule &sch) |
Infer the bound of all iteration variables relates to the schedule. More... | |
bool | VerifyCompactBuffer (const Stmt &stmt) |
Verify if there is any argument bound to compact buffer. More... | |
Stmt | ScheduleOps (Schedule s, Map< IterVar, Range > dom_map, bool debug_keep_trivial_loop) |
Schedule s' dependent operations. More... | |
PrimFunc | SchedulePostProcToPrimFunc (Array< ObjectRef > arg_list, Stmt body, Optional< Map< Tensor, Buffer >> bindings) |
Postprocessing the Stmt generated by ScheduleOps to create a PrimFunc that can then be used for further TIR optimizations. More... | |
PrimExpr | operator! (const Tensor::Slice &a) |
PrimExpr | operator- (const Tensor::Slice &a) |
template<typename T > | |
PrimExpr | operator+ (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator+ (const T &a, const Tensor::Slice &b) |
PrimExpr | operator+ (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator- (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator- (const T &a, const Tensor::Slice &b) |
PrimExpr | operator- (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator* (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator* (const T &a, const Tensor::Slice &b) |
PrimExpr | operator* (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator== (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator== (const T &a, const Tensor::Slice &b) |
PrimExpr | operator== (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator<= (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator<= (const T &a, const Tensor::Slice &b) |
PrimExpr | operator<= (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator>= (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator>= (const T &a, const Tensor::Slice &b) |
PrimExpr | operator>= (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator!= (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator!= (const T &a, const Tensor::Slice &b) |
PrimExpr | operator!= (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator&& (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator&& (const T &a, const Tensor::Slice &b) |
PrimExpr | operator&& (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator|| (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator|| (const T &a, const Tensor::Slice &b) |
PrimExpr | operator|| (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator>> (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator>> (const T &a, const Tensor::Slice &b) |
PrimExpr | operator>> (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator<< (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator<< (const T &a, const Tensor::Slice &b) |
PrimExpr | operator<< (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator> (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator> (const T &a, const Tensor::Slice &b) |
PrimExpr | operator> (const Tensor::Slice &a, const Tensor::Slice &b) |
template<typename T > | |
PrimExpr | operator< (const Tensor::Slice &a, const T &b) |
template<typename T > | |
PrimExpr | operator< (const T &a, const Tensor::Slice &b) |
PrimExpr | operator< (const Tensor::Slice &a, const Tensor::Slice &b) |
Tensor expression language DSL.
using tvm::te::FBatchCompute = typedef std::function<Array<PrimExpr>(const Array<Var>& i)> |
The compute function to specify the inputs source of Tensors.
using tvm::te::FCompute = typedef std::function<PrimExpr(const Array<Var>& i)> |
The compute function to specify the input source of a Tensor.
enum tvm::te::AttachType : int |
void tvm::te::AutoInlineBroarcast | ( | Schedule | sch | ) |
To automatically inline the broadcast operations.
sch | The schedule to be inlined. |
void tvm::te::AutoInlineElemWise | ( | Schedule | sch | ) |
To automatically inline the element-wise operations.
sch | The schedule to be inlined. |
void tvm::te::AutoInlineInjective | ( | Schedule | sch | ) |
To automatically inline operations with injective writes (i.e. writes without reduction or sequential loops). Note that in this case, guarantees about contiguity, transpose, stride, alignemnt and memory footprint in general do not hold.
sch | The schedule to be inlined. |
Tensor tvm::te::compute | ( | Array< PrimExpr > | shape, |
FCompute | fcompute, | ||
std::string | name = "tensor" , |
||
std::string | tag = "" , |
||
Map< String, ObjectRef > | attrs = {} |
||
) |
Construct a new tensor by computing over shape, using the computation rule: result_tensor[axis] = fcompute(axis)
shape | Shape of the tensor. |
fcompute | The compute function to create the tensor. |
name | The optional name of the tensor. |
tag | The optional tag of the tensor. |
attrs | Optional additional attributes of the compute. |
Array<Tensor> tvm::te::compute | ( | Array< PrimExpr > | shape, |
FBatchCompute | fcompute, | ||
std::string | name = "tensor" , |
||
std::string | tag = "" , |
||
Map< String, ObjectRef > | attrs = {} |
||
) |
Construct a new tensor by computing over shape, using the computation rule: result_tensor[axis] = fcompute(axis)
shape | Shape of the tensor. |
fcompute | The compute function to create the tensors. |
name | The optional name of the tensor. |
tag | The optional tag of the tensor. |
attrs | Optional additional attributes of the compute. |
|
inline |
|
inline |
|
inline |
|
inline |
Create a schedule for array of ops(and their dependencies).
ops | The ops to be scheduled. |
Take the derivative of the expression with respect to the given variable.
expr | The expression to differentiate. |
var | The variable to differentiate with respect to. |
Array<Tensor> tvm::te::Gradient | ( | const Tensor & | output, |
const Array< Tensor > & | inputs, | ||
const Tensor & | head = Tensor() |
||
) |
Perform reverse mode automatic differentiation.
Each item of the result
field of the result is an adjoint for the corresponding item of inputs
, i.e. head
multiplied by the Jacobian of output
with respect to the corresponding item of inputs
.
output | The tensor to differentiate. |
inputs | The array of input tensors. When the array is empty, will perform differentiation wrt all tensors the output depends on. |
head | The adjoint of the output, in other words, some tensor, by which the Jacobians will be multiplied (using tensordot axes=output.shape ). Its shape must be of the form prefix + output.shape . If the null pointer is provided, the identity tensor of shape output.shape + output.shape will be used. |
inputs
. Infer the bound of all iteration variables relates to the schedule.
sch | The root schedule to infer all the bounds. |
Get the tensor representing the Jacobian of the output with respect to the input.
Note that if output
depends on input
indirectly (by using some other tensor depending on input
), this dependency won't contribute to the resulting Jacobian. For such cases use the function Gradient.
output | The tensor to differentiate. |
input | The input tensor, which output should directly use. |
output.shape + input.shape
.
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
|
inline |
Tensor tvm::te::placeholder | ( | Array< PrimExpr > | shape, |
DataType | dtype = DataType::Float(32) , |
||
std::string | name = "placeholder" |
||
) |
create a place holder tensor.
shape | The shape of the tensor. |
dtype | the data type of the tensor. |
name | The name of the Tensor. |
Create a new IterVar for reduction operations.
dom | The domain of the reduction axis. |
name | The name of the reduction axis. |
Array<Tensor> tvm::te::scan | ( | Array< Tensor > | init, |
Array< Tensor > | update, | ||
Array< Tensor > | state_placeholder, | ||
Array< Tensor > | inputs = Array< Tensor >() , |
||
std::string | name = "scan" , |
||
std::string | tag = "" , |
||
Map< String, ObjectRef > | attrs = {} |
||
) |
Construct new tensors by scan.
init | The intialize tensor of first K steps. |
update | The update tensor indicated the updated result after each timestamp. |
state_placeholder | The placeholder for the states. |
inputs | The inputs to the scan body, this is optional, but recommended to provide concrete information about scan body. |
name | The optional name of the tensor. |
tag | The optional tag of the tensor. |
attrs | Optional additional attributes of the compute. |
Stmt tvm::te::ScheduleOps | ( | Schedule | s, |
Map< IterVar, Range > | dom_map, | ||
bool | debug_keep_trivial_loop | ||
) |
Schedule s' dependent operations.
s | The schedule to be realized |
dom_map | The domain of each iter vars. |
debug_keep_trivial_loop | Whether keep trivial loops with extent of 1 during lowering. This is a debug feature for dataflow/axis analysis. Note: If this is true, The lowered IR may be incorrect, because we will also delete the init part of reduction |
PrimFunc tvm::te::SchedulePostProcToPrimFunc | ( | Array< ObjectRef > | arg_list, |
Stmt | body, | ||
Optional< Map< Tensor, Buffer >> | bindings | ||
) |
Postprocessing the Stmt generated by ScheduleOps to create a PrimFunc that can then be used for further TIR optimizations.
Perform this translation before running any TIR optimizations.
List of actions taken by the function:
arg_list | Array of Tensor/Var/Buffer arguments to the function. |
body | The body of the function. |
bindings | potential Tensor to Buffer bindings for the Tensors in the body. |
Create a new IterVar that represents an axis in thread.
dom | Optional, domain of the thread axis. |
tag | The thread tag of the axis. |
Var tvm::te::var | ( | std::string | name_hint, |
DataType | t = DataType::Int(32) |
||
) |
Construct a new Var expression.
name_hint | The name hint for the expression |
t | The type of the expression |
Tensor tvm::te::VectorJacobianProduct | ( | const Tensor & | output, |
const Tensor & | input, | ||
const Tensor & | head | ||
) |
The building block for reverse-mode AD.
Differentiate output
wrt input
and multiply the result by head
on the left using tensor dot product. input
must be an immediate dependency of output
(must be called from within the body of output
). That is, the function will compute one summand of the adjoint for input
given the adjoint for output
(which is called head
here).
output | The tensor to differentiate. |
input | The input tensor, which output should directly use. |
head | The adjoint of output . Must be of shape prefix + output.shape |
prefix + input.shape
representing the partial adjoint of input
wrt one of its consumers (output) bool tvm::te::VerifyCompactBuffer | ( | const Stmt & | stmt | ) |
Verify if there is any argument bound to compact buffer.
stmt | The stmt to be verified. |