onnxruntime-tvm/include/tvm/tensor_intrin.h

94 строки
2.8 KiB
C++

/*!
* Copyright (c) 2017 by Contributors
* \file tensor_intrin.h
* \brief Tensor intrinsic operations.
*/
#ifndef TVM_TENSOR_INTRIN_H_
#define TVM_TENSOR_INTRIN_H_
#include <string>
#include "./tensor.h"
#include "./buffer.h"
namespace tvm {
// Internal node container of tensor intrinsics.
class TensorIntrinNode;
/*! \brief Tensor intrinsic node. */
class TensorIntrin : public NodeRef {
public:
TensorIntrin() {}
explicit TensorIntrin(std::shared_ptr<Node> n) : NodeRef(n) {}
/*!
* \brief access the internal node container
* \return the pointer to the internal node container
*/
inline const TensorIntrinNode* operator->() const;
/*! \brief specify container node */
using ContainerType = TensorIntrinNode;
};
/*! \brief Node to represent a Tensor intrinsic operator */
class TensorIntrinNode : public Node {
public:
/*! \brief The name of the intrinsic */
std::string name;
/*! \brief The operation this intrinsics is carrying out */
Operation op;
/*! \brief List of inputs of operator, placeholder in postdfs order */
Array<Tensor> inputs;
/*!
* \brief Symbolic buffers of each output/input tensor
* buffers[0:len(inputs)] are buffers of the inputs.
* buffers[len(inputs):] are buffers of each output.
*
* \note When a field in Buffer is Var, it means we can be flexible
* wrt that field and Var can occur in body.
* When it is a constant, it means we can only take data in that shape.
*/
Array<Buffer> buffers;
/*! \brief The normal statement to execute the intrinsic */
Stmt body;
/*!
* \brief Special statement for reduction op, can be None
* reset the value of output buffer to identity value.
*/
Stmt reduce_init;
/*!
* \brief Special statement for reduction op, can be None
* Reduce: do a reduction of current output buffer with the result.
*/
Stmt reduce_update;
/*! \brief constructor */
TensorIntrinNode() {}
void VisitAttrs(AttrVisitor* v) final {
v->Visit("name", &name);
v->Visit("op", &op);
v->Visit("inputs", &inputs);
v->Visit("buffers", &buffers);
v->Visit("body", &body);
v->Visit("reduce_init", &reduce_init);
v->Visit("reduce_update", &reduce_update);
}
static TensorIntrin make(std::string name,
Operation op,
Array<Tensor> inputs,
Array<Buffer> buffers,
Stmt body,
Stmt reduce_init,
Stmt reduce_update);
static constexpr const char* _type_key = "TensorIntrin";
TVM_DECLARE_NODE_TYPE_INFO(TensorIntrinNode, Node);
};
inline const TensorIntrinNode* TensorIntrin::operator->() const {
return static_cast<const TensorIntrinNode*>(node_.get());
}
} // namespace tvm
#endif // TVM_TENSOR_INTRIN_H_