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...
#include <operation.h>
|
Array< IterVar > | root_iter_vars () const final |
|
Array< PrimExpr > | output_shape (size_t idx) const final |
| Get shape of i-th output tensor. More...
|
|
void | GatherBound (const Operation &self, const std::unordered_map< Tensor, TensorDom > &tensor_dom, std::unordered_map< IterVar, Range > *out_dom_map) const final |
| Gather the bound from output tensor. Set the range of each root_iter_vars in the op to out_dom_map. More...
|
|
Stmt | BuildRealize (const Stage &stage, const std::unordered_map< IterVar, Range > &realize_map, const Stmt &body, String storage_scope="") const final |
| Build the Realize statement that realizes the op's output tensors. More...
|
|
virtual size_t | num_schedulable_dims () const =0 |
|
| TVM_DECLARE_BASE_OBJECT_INFO (BaseComputeOpNode, OperationNode) |
|
virtual | ~OperationNode () |
|
virtual int | num_outputs () const =0 |
|
virtual DataType | output_dtype (size_t i) const =0 |
| Get data type. i-th output tensor. More...
|
|
virtual Array< Tensor > | InputTensors () const =0 |
| List all the input Tensors. More...
|
|
virtual Operation | ReplaceInputs (const Operation &self, const std::unordered_map< Tensor, Tensor > &rmap) const =0 |
| Replace the input of the operation by pattern specified by rmap. More...
|
|
virtual void | PropBoundToInputs (const Operation &self, arith::Analyzer *analyzer, const std::unordered_map< const VarNode *, IntSet > &dom_map, std::unordered_map< Tensor, TensorDom > *out_dom_map) const =0 |
| Propagate the bounds to inputs. More...
|
|
virtual Stmt | BuildProvide (const Stage &stage, const std::unordered_map< IterVar, Range > &dom_map, bool debug_keep_trivial_loop) const =0 |
| Build the statement that provide the output tensors. More...
|
|
| TVM_DECLARE_BASE_OBJECT_INFO (OperationNode, Object) |
|
uint32_t | type_index () const |
|
std::string | GetTypeKey () const |
|
size_t | GetTypeKeyHash () const |
|
template<typename TargetType > |
bool | IsInstance () const |
|
bool | unique () const |
|
| Object () |
|
| Object (const Object &other) |
|
| Object (Object &&other) |
|
Object & | operator= (const Object &other) |
|
Object & | operator= (Object &&other) |
|
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)
◆ BuildRealize()
Stmt tvm::te::BaseComputeOpNode::BuildRealize |
( |
const Stage & |
stage, |
|
|
const std::unordered_map< IterVar, Range > & |
realize_map, |
|
|
const Stmt & |
body, |
|
|
String |
storage_scope = "" |
|
) |
| const |
|
finalvirtual |
Build the Realize statement that realizes the op's output tensors.
- Parameters
-
stage | the op's stage. |
realize_map | The realization domain map of the operators. |
body | The body that is going to get |
storage_scope | The storage scope associated with this realization |
- Returns
- A realization statement that wraps body.
Implements tvm::te::OperationNode.
◆ GatherBound()
void tvm::te::BaseComputeOpNode::GatherBound |
( |
const Operation & |
self, |
|
|
const std::unordered_map< Tensor, TensorDom > & |
tensor_dom, |
|
|
std::unordered_map< IterVar, Range > * |
out_dom_map |
|
) |
| const |
|
finalvirtual |
Gather the bound from output tensor. Set the range of each root_iter_vars in the op to out_dom_map.
- Parameters
-
self | The reference to self. |
tensor_dom | Domain map of Tensor->access set of each dimension. |
out_dom_map | The output domain map of each IterVar to be setted. |
Implements tvm::te::OperationNode.
◆ num_schedulable_dims()
virtual size_t tvm::te::BaseComputeOpNode::num_schedulable_dims |
( |
| ) |
const |
|
pure virtual |
◆ output_shape()
Array<PrimExpr> tvm::te::BaseComputeOpNode::output_shape |
( |
size_t |
i | ) |
const |
|
finalvirtual |
◆ root_iter_vars()
Array<IterVar> tvm::te::BaseComputeOpNode::root_iter_vars |
( |
| ) |
const |
|
finalvirtual |
- Returns
- The list of iteration variable at root
- Note
- root_iter_vars decides the shape of the outputs.
Implements tvm::te::OperationNode.
◆ TVM_DECLARE_BASE_OBJECT_INFO()
◆ _type_key
constexpr const char* tvm::te::BaseComputeOpNode::_type_key = "BaseComputeOp" |
|
staticconstexpr |
◆ axis
◆ reduce_axis
IterVar on each reduction axis, if the body is a Reduce.
The documentation for this class was generated from the following file: