24 #ifndef TVM_TE_OPERATION_H_
25 #define TVM_TE_OPERATION_H_
28 #include <tvm/ffi/reflection/registry.h>
35 #include <unordered_map>
50 std::vector<std::vector<IntSet>>
data;
63 ffi::Map<ffi::String, ffi::Any>
attrs;
88 refl::ObjectDef<OperationNode>()
113 refl::ObjectDef<PlaceholderOpNode>()
146 refl::ObjectDef<BaseComputeOpNode>()
165 ffi::Array<
Tensor> InputTensors() const final;
167 static
void RegisterReflection() {
180 TVM_DLL
ComputeOp(std::string name, std::string tag, ffi::Map<ffi::String, ffi::Any> attrs,
181 ffi::Array<IterVar> axis, ffi::Array<PrimExpr> body);
225 refl::ObjectDef<ScanOpNode>()
242 TVM_DLL
ScanOp(std::string name, std::string tag,
243 ffi::Optional<ffi::Map<ffi::String, ffi::Any>> attrs,
IterVar axis,
244 ffi::Array<Tensor> init, ffi::Array<Tensor> update,
245 ffi::Array<Tensor> state_placeholder, ffi::Array<Tensor> input);
274 refl::ObjectDef<ExternOpNode>()
289 TVM_DLL
ExternOp(std::string name, std::string tag, ffi::Map<ffi::String, ffi::Any> attrs,
290 ffi::Array<Tensor> inputs, ffi::Array<Buffer> input_placeholders,
291 ffi::Array<Buffer> output_placeholders,
Stmt body);
323 using FBatchCompute = std::function<ffi::Array<PrimExpr>(
const ffi::Array<Var>& i)>;
332 std::string name =
"placeholder");
344 std::string tag =
"", ffi::Map<ffi::String, ffi::Any> attrs = {});
356 std::string name =
"tensor", std::string tag =
"",
357 ffi::Map<ffi::String, ffi::Any> attrs = {});
371 TVM_DLL ffi::Array<Tensor>
scan(ffi::Array<Tensor> init, ffi::Array<Tensor> update,
372 ffi::Array<Tensor> state_placeholder,
373 ffi::Array<Tensor> inputs = ffi::Array<Tensor>(),
374 std::string name =
"scan", std::string tag =
"",
375 ffi::Map<ffi::String, ffi::Any> attrs = {});
379 std::string name =
"tensor", std::string tag =
"",
380 ffi::Map<ffi::String, ffi::Any> attrs = {}) {
381 FCompute fc = [f](
const ffi::Array<Var>& i) {
return f(i[0]); };
385 std::string name =
"tensor", std::string tag =
"",
386 ffi::Map<ffi::String, ffi::Any> attrs = {}) {
387 FCompute fc = [f](
const ffi::Array<Var>& i) {
return f(i[0], i[1]); };
391 std::string name =
"tensor", std::string tag =
"",
392 ffi::Map<ffi::String, ffi::Any> attrs = {}) {
393 FCompute fc = [f](
const ffi::Array<Var>& i) {
return f(i[0], i[1], i[2]); };
397 std::string name =
"tensor", std::string tag =
"",
398 ffi::Map<ffi::String, ffi::Any> attrs = {}) {
399 FCompute fc = [f](
const ffi::Array<Var>& i) {
return f(i[0], i[1], i[2], i[3]); };
Algebra expression simplifications.
Symbolic n-dimensional array, to represent a memory buffer.
Reference to PrimExprNode.
Definition: expr.h:124
Range container
Definition: expr.h:689
Runtime primitive data type.
Definition: data_type.h:47
static DataType Float(int bits, int lanes=1)
Construct an float type.
Definition: data_type.h:291
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:274
A Compute op that compute a tensor on certain domain. This is the base class for ComputeOp (operating...
Definition: operation.h:135
ffi::Array< IterVar > reduce_axis
IterVar on each reduction axis, if the body is a Reduce.
Definition: operation.h:140
static void RegisterReflection()
Definition: operation.h:144
ffi::Array< PrimExpr > output_shape(size_t idx) const final
Get shape of i-th output tensor.
TVM_FFI_DECLARE_OBJECT_INFO("te.BaseComputeOp", BaseComputeOpNode, OperationNode)
ffi::Array< IterVar > axis
IterVar on each axis.
Definition: operation.h:138
A Compute op that compute a tensor on certain domain.
Definition: operation.h:156
ComputeOpNode()
constructor
Definition: operation.h:161
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ComputeOp", ComputeOpNode, BaseComputeOpNode)
int num_outputs() const final
ffi::Array< PrimExpr > body
the compute expression
Definition: operation.h:159
Managed reference to ComputeOpNode.
Definition: operation.h:178
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ComputeOp, Operation, ComputeOpNode)
ComputeOp(std::string name, std::string tag, ffi::Map< ffi::String, ffi::Any > attrs, ffi::Array< IterVar > axis, ffi::Array< PrimExpr > body)
TVM_DEFINE_OBJECT_REF_COW_METHOD(ComputeOpNode)
External computation that cannot be splitted.
Definition: operation.h:253
int num_outputs() const final
ffi::Array< Tensor > inputs
The input tensors.
Definition: operation.h:256
static void RegisterReflection()
Definition: operation.h:272
ffi::Array< PrimExpr > output_shape(size_t i) const final
Get shape of i-th output tensor.
ffi::Array< Buffer > input_placeholders
Symbolic placeholder representation of inputs.
Definition: operation.h:258
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ExternOp", ExternOpNode, OperationNode)
ffi::Array< Tensor > InputTensors() const final
List all the input Tensors.
Stmt body
the statement that generates the computation.
Definition: operation.h:262
ExternOpNode()
constructor
Definition: operation.h:265
ffi::Array< Buffer > output_placeholders
Symbolic placeholder representation of outputs.
Definition: operation.h:260
DataType output_dtype(size_t i) const final
Get data type. i-th output tensor.
Managed reference to ExternOpNode.
Definition: operation.h:287
ExternOp(std::string name, std::string tag, ffi::Map< ffi::String, ffi::Any > attrs, ffi::Array< Tensor > inputs, ffi::Array< Buffer > input_placeholders, ffi::Array< Buffer > output_placeholders, Stmt body)
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExternOp, Operation, ExternOpNode)
Base class of all operation nodes.
Definition: operation.h:56
virtual ~OperationNode()
Definition: operation.h:65
TVM_FFI_DECLARE_OBJECT_INFO("te.Operation", OperationNode, Object)
virtual ffi::Array< PrimExpr > output_shape(size_t i) const =0
Get shape of i-th output tensor.
static void RegisterReflection()
Definition: operation.h:86
virtual ffi::Array< Tensor > InputTensors() const =0
List all the input Tensors.
virtual DataType output_dtype(size_t i) const =0
Get data type. i-th output tensor.
virtual int num_outputs() const =0
ffi::Map< ffi::String, ffi::Any > attrs
additional attributes of the operation
Definition: operation.h:63
std::string name
optional name of the operation
Definition: operation.h:59
std::string tag
optional tag of the operation
Definition: operation.h:61
Operation that produces tensors.
Definition: tensor.h:48
const OperationNode * operator->() const
access the internal node container
Definition: operation.h:404
A placeholder op represents an input placeholder.
Definition: operation.h:99
TVM_FFI_DECLARE_OBJECT_INFO("te.PlaceholderOp", PlaceholderOpNode, OperationNode)
DataType dtype
The data type of the input.
Definition: operation.h:104
static void RegisterReflection()
Definition: operation.h:111
ffi::Array< PrimExpr > shape
The shape of the input.
Definition: operation.h:102
int num_outputs() const final
DataType output_dtype(size_t i) const final
Get data type. i-th output tensor.
ffi::Array< Tensor > InputTensors() const final
List all the input Tensors.
ffi::Array< PrimExpr > output_shape(size_t i) const final
Get shape of i-th output tensor.
Managed reference to PlaceholderOpNode.
Definition: operation.h:124
PlaceholderOp(std::string name, ffi::Array< PrimExpr > shape, DataType dtype)
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PlaceholderOp, Operation, PlaceholderOpNode)
Symbolic scan.
Definition: operation.h:190
ffi::Array< Tensor > state_placeholder
The placeholder to refer as states in update.
Definition: operation.h:199
ScanOpNode()
constructor
Definition: operation.h:216
int num_outputs() const final
ffi::Array< Tensor > init
the initialization tensors
Definition: operation.h:195
DataType output_dtype(size_t i) const final
Get data type. i-th output tensor.
ffi::Array< Tensor > update
the update function represented by tensor
Definition: operation.h:197
ffi::Array< PrimExpr > output_shape(size_t i) const final
Get shape of i-th output tensor.
ffi::Array< Tensor > inputs
the inputs to the scan, these are optionally provided But they can be helpful to provide hints to spe...
Definition: operation.h:204
ffi::Array< Tensor > InputTensors() const final
List all the input Tensors.
static void RegisterReflection()
Definition: operation.h:223
IterVar scan_axis
IterVar to scan over.
Definition: operation.h:193
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ScanOp", ScanOpNode, OperationNode)
ffi::Array< IterVar > spatial_axis_
Spatial axis to indicate spatial dimension of each output. They corresponds to flattened spatial axis...
Definition: operation.h:214
Managed reference to ScanOpNode.
Definition: operation.h:240
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ScanOp, Operation, ScanOpNode)
ScanOp(std::string name, std::string tag, ffi::Optional< ffi::Map< ffi::String, ffi::Any >> attrs, IterVar axis, ffi::Array< Tensor > init, ffi::Array< Tensor > update, ffi::Array< Tensor > state_placeholder, ffi::Array< Tensor > input)
Tensor structure representing a possible input, or intermediate computation result.
Definition: tensor.h:100
Iteration Variable, represents an iteration over an integer interval.
Definition: var.h:297
Container of all statements.
Definition: stmt.h:63
a named variable in TIR
Definition: var.h:77
Definition: repr_printer.h:91
std::function< PrimExpr(const ffi::Array< Var > &i)> FCompute
The compute function to specify the input source of a Tensor.
Definition: operation.h:320
ffi::Array< Tensor > scan(ffi::Array< Tensor > init, ffi::Array< Tensor > update, ffi::Array< Tensor > state_placeholder, ffi::Array< Tensor > inputs=ffi::Array< Tensor >(), std::string name="scan", std::string tag="", ffi::Map< ffi::String, ffi::Any > attrs={})
Construct new tensors by scan.
std::function< ffi::Array< PrimExpr >(const ffi::Array< Var > &i)> FBatchCompute
The compute function to specify the inputs source of Tensors.
Definition: operation.h:323
IterVar thread_axis(Range dom, std::string tag)
Create a new IterVar that represents an axis in thread.
IterVar reduce_axis(Range dom, std::string name="rv")
Create a new IterVar for reduction operations.
Tensor compute(ffi::Array< PrimExpr > shape, FCompute fcompute, std::string name="tensor", std::string tag="", ffi::Map< ffi::String, ffi::Any > attrs={})
Construct a new tensor by computing over shape, using the computation rule: result_tensor[axis] = fco...
Var var(std::string name_hint, DataType t=DataType::Int(32))
Construct a new Var expression.
Tensor placeholder(ffi::Array< PrimExpr > shape, DataType dtype=DataType::Float(32), std::string name="placeholder")
create a place holder tensor.
Tensor shape(const Tensor &src, DataType dtype, const std::string name="T_shape", const std::string tag=kInjective)
Get the shape of input tensor.
Definition: transform.h:1960
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Temporary data structure to store union of bounds of each axis of Tensor.
Definition: operation.h:46
TensorDom(int ndim)
Definition: operation.h:48
std::vector< std::vector< IntSet > > data
The domain data.
Definition: operation.h:50
Common operators defined for Expr.