tvm
operation.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
24 #ifndef TVM_TE_OPERATION_H_
25 #define TVM_TE_OPERATION_H_
26 
27 #include <tvm/arith/analyzer.h>
28 #include <tvm/ffi/reflection/registry.h>
29 #include <tvm/te/tensor.h>
30 #include <tvm/tir/buffer.h>
31 #include <tvm/tir/expr.h>
32 #include <tvm/tir/op.h>
33 
34 #include <string>
35 #include <unordered_map>
36 #include <vector>
37 
38 namespace tvm {
40 namespace te {
41 
46 struct TensorDom {
47  // constructor
48  explicit TensorDom(int ndim) : data(ndim) {}
50  std::vector<std::vector<IntSet>> data;
51 };
52 
56 class TVM_DLL OperationNode : public Object {
57  public:
59  std::string name;
61  std::string tag;
63  ffi::Map<ffi::String, ffi::Any> attrs;
64  // virtual destructor.
65  virtual ~OperationNode() {}
67  virtual int num_outputs() const = 0;
73  virtual DataType output_dtype(size_t i) const = 0;
79  virtual ffi::Array<PrimExpr> output_shape(size_t i) const = 0;
84  virtual ffi::Array<Tensor> InputTensors() const = 0;
85 
86  static void RegisterReflection() {
87  namespace refl = tvm::ffi::reflection;
88  refl::ObjectDef<OperationNode>()
89  .def_ro("name", &OperationNode::name)
90  .def_ro("tag", &OperationNode::tag)
91  .def_ro("attrs", &OperationNode::attrs);
92  }
93  TVM_FFI_DECLARE_OBJECT_INFO("te.Operation", OperationNode, Object);
94 };
95 
100  public:
102  ffi::Array<PrimExpr> shape;
105  // override behavior.
106  int num_outputs() const final;
107  DataType output_dtype(size_t i) const final;
108  ffi::Array<PrimExpr> output_shape(size_t i) const final;
109  ffi::Array<Tensor> InputTensors() const final;
110 
111  static void RegisterReflection() {
112  namespace refl = tvm::ffi::reflection;
113  refl::ObjectDef<PlaceholderOpNode>()
114  .def_ro("shape", &PlaceholderOpNode::shape)
115  .def_ro("dtype", &PlaceholderOpNode::dtype);
116  }
118 };
119 
124 class PlaceholderOp : public Operation {
125  public:
126  TVM_DLL PlaceholderOp(std::string name, ffi::Array<PrimExpr> shape, DataType dtype);
127 
129 };
130 
135 class TVM_DLL BaseComputeOpNode : public OperationNode {
136  public:
138  ffi::Array<IterVar> axis;
140  ffi::Array<IterVar> reduce_axis;
141  // override functions
142  ffi::Array<PrimExpr> output_shape(size_t idx) const final;
143 
144  static void RegisterReflection() {
145  namespace refl = tvm::ffi::reflection;
146  refl::ObjectDef<BaseComputeOpNode>()
147  .def_ro("axis", &BaseComputeOpNode::axis)
148  .def_ro("reduce_axis", &BaseComputeOpNode::reduce_axis);
149  }
151 };
152 
156 class TVM_DLL ComputeOpNode : public BaseComputeOpNode {
157  public:
159  ffi::Array<PrimExpr> body;
162  // override functions
163  int num_outputs() const final;
164  DataType output_dtype(size_t i) const final;
165  ffi::Array<Tensor> InputTensors() const final;
166 
167  static void RegisterReflection() {
168  namespace refl = tvm::ffi::reflection;
169  refl::ObjectDef<ComputeOpNode>().def_ro("body", &ComputeOpNode::body);
170  }
172 };
173 
178 class ComputeOp : public Operation {
179  public:
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);
182 
185 };
186 
190 class ScanOpNode : public OperationNode {
191  public:
195  ffi::Array<Tensor> init;
197  ffi::Array<Tensor> update;
199  ffi::Array<Tensor> state_placeholder;
204  ffi::Array<Tensor> inputs;
214  ffi::Array<IterVar> spatial_axis_;
217  // override behavior.
218  int num_outputs() const final;
219  DataType output_dtype(size_t i) const final;
220  ffi::Array<PrimExpr> output_shape(size_t i) const final;
221  ffi::Array<Tensor> InputTensors() const final;
222 
223  static void RegisterReflection() {
224  namespace refl = tvm::ffi::reflection;
225  refl::ObjectDef<ScanOpNode>()
226  .def_ro("scan_axis", &ScanOpNode::scan_axis)
227  .def_ro("init", &ScanOpNode::init)
228  .def_ro("update", &ScanOpNode::update)
229  .def_ro("state_placeholder", &ScanOpNode::state_placeholder)
230  .def_ro("inputs", &ScanOpNode::inputs)
231  .def_ro("spatial_axis_", &ScanOpNode::spatial_axis_);
232  }
234 };
235 
240 class ScanOp : public Operation {
241  public:
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);
246 
248 };
249 
253 class ExternOpNode : public OperationNode {
254  public:
256  ffi::Array<Tensor> inputs;
258  ffi::Array<Buffer> input_placeholders;
260  ffi::Array<Buffer> output_placeholders;
263 
266  // override functions
267  int num_outputs() const final;
268  DataType output_dtype(size_t i) const final;
269  ffi::Array<PrimExpr> output_shape(size_t i) const final;
270  ffi::Array<Tensor> InputTensors() const final;
271 
272  static void RegisterReflection() {
273  namespace refl = tvm::ffi::reflection;
274  refl::ObjectDef<ExternOpNode>()
275  .def_ro("inputs", &ExternOpNode::inputs)
276  .def_ro("input_placeholders", &ExternOpNode::input_placeholders)
277  .def_ro("output_placeholders", &ExternOpNode::output_placeholders)
278  .def_ro("body", &ExternOpNode::body);
279  }
281 };
282 
287 class ExternOp : public Operation {
288  public:
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);
292 
294 };
295 
301 TVM_DLL Var var(std::string name_hint, DataType t = DataType::Int(32));
302 
309 TVM_DLL IterVar thread_axis(Range dom, std::string tag);
310 
317 TVM_DLL IterVar reduce_axis(Range dom, std::string name = "rv");
318 
320 using FCompute = std::function<PrimExpr(const ffi::Array<Var>& i)>;
321 
323 using FBatchCompute = std::function<ffi::Array<PrimExpr>(const ffi::Array<Var>& i)>;
324 
331 TVM_DLL Tensor placeholder(ffi::Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
332  std::string name = "placeholder");
333 
343 TVM_DLL Tensor compute(ffi::Array<PrimExpr> shape, FCompute fcompute, std::string name = "tensor",
344  std::string tag = "", ffi::Map<ffi::String, ffi::Any> attrs = {});
345 
355 TVM_DLL ffi::Array<Tensor> compute(ffi::Array<PrimExpr> shape, FBatchCompute fcompute,
356  std::string name = "tensor", std::string tag = "",
357  ffi::Map<ffi::String, ffi::Any> attrs = {});
358 
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 = {});
376 
377 // same as compute, specialized for different fcompute function
378 inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(Var)> f,
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]); };
382  return compute(shape, fc, name, tag, attrs);
383 }
384 inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(Var, Var)> f,
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]); };
388  return compute(shape, fc, name, tag, attrs);
389 }
390 inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(Var, Var, Var)> f,
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]); };
394  return compute(shape, fc, name, tag, attrs);
395 }
396 inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(Var, Var, Var, Var)> f,
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]); };
400  return compute(shape, fc, name, tag, attrs);
401 }
402 
403 // inline function.
404 inline const OperationNode* Operation::operator->() const {
405  return static_cast<const OperationNode*>(get());
406 }
407 } // namespace te
408 } // namespace tvm
409 #endif // TVM_TE_OPERATION_H_
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
Dataflow tensor object.
TIR expressions.
Common operators defined for Expr.