wehub-resource-sync
414 lines
15 KiB
C++
414 lines
15 KiB
C++
/*
|
|
* Licensed to the Apache Software Foundation (ASF) under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you under the Apache License, Version 2.0 (the
|
|
* "License"); you may not use this file except in compliance
|
|
* with the License. You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing,
|
|
* software distributed under the License is distributed on an
|
|
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
* KIND, either express or implied. See the License for the
|
|
* specific language governing permissions and limitations
|
|
* under the License.
|
|
*/
|
|
|
|
/*!
|
|
* \file tvm/te/operation.h
|
|
* \brief Operation node can generate one or multiple Tensors
|
|
*/
|
|
#ifndef TVM_TE_OPERATION_H_
|
|
#define TVM_TE_OPERATION_H_
|
|
|
|
#include <tvm/arith/analyzer.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
#include <tvm/ir/cow.h>
|
|
#include <tvm/te/tensor.h>
|
|
#include <tvm/tirx/buffer.h>
|
|
#include <tvm/tirx/expr.h>
|
|
#include <tvm/tirx/op.h>
|
|
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace tvm {
|
|
/*! \brief Tensor expression language DSL. */
|
|
namespace te {
|
|
|
|
/*!
|
|
* \brief Temporary data structure to store union
|
|
* of bounds of each axis of Tensor.
|
|
*/
|
|
struct TensorDom {
|
|
// constructor
|
|
explicit TensorDom(int ndim) : data(ndim) {}
|
|
/*! \brief The domain data */
|
|
std::vector<std::vector<IntSet>> data;
|
|
};
|
|
|
|
/*!
|
|
* \brief Base class of all operation nodes
|
|
*/
|
|
class TVM_DLL OperationNode : public ffi::Object {
|
|
public:
|
|
/*! \brief optional name of the operation */
|
|
std::string name;
|
|
/*! \brief optional tag of the operation */
|
|
std::string tag;
|
|
/*! \brief additional attributes of the operation*/
|
|
ffi::Map<ffi::String, ffi::Any> attrs;
|
|
// virtual destructor.
|
|
virtual ~OperationNode() {}
|
|
/*! \return number of outputs */
|
|
virtual int num_outputs() const = 0;
|
|
/*!
|
|
* \brief Get the primitive element type of the i-th output tensor.
|
|
* \param i The output index.
|
|
* \return primitive element type of i-th output.
|
|
*/
|
|
virtual PrimType output_dtype(size_t i) const = 0;
|
|
/*!
|
|
* \brief Get shape of i-th output tensor.
|
|
* \param i The output index.
|
|
* \return shape of i-th output.
|
|
*/
|
|
virtual ffi::Array<PrimExpr> output_shape(size_t i) const = 0;
|
|
/*!
|
|
* \brief List all the input Tensors.
|
|
* \return List of input tensors.
|
|
*/
|
|
virtual ffi::Array<Tensor> InputTensors() const = 0;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<OperationNode>()
|
|
.def_ro("name", &OperationNode::name)
|
|
.def_ro("tag", &OperationNode::tag)
|
|
.def_ro("attrs", &OperationNode::attrs);
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO("te.Operation", OperationNode, ffi::Object);
|
|
};
|
|
|
|
/*!
|
|
* \brief A placeholder op represents an input placeholder.
|
|
*/
|
|
class PlaceholderOpNode : public OperationNode {
|
|
public:
|
|
/*! \brief The shape of the input */
|
|
ffi::Array<PrimExpr> shape;
|
|
/*! \brief The dtype of the input. */
|
|
PrimType dtype = PrimType::Void();
|
|
// override behavior.
|
|
int num_outputs() const final;
|
|
PrimType output_dtype(size_t i) const final;
|
|
ffi::Array<PrimExpr> output_shape(size_t i) const final;
|
|
ffi::Array<Tensor> InputTensors() const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<PlaceholderOpNode>()
|
|
.def_ro("shape", &PlaceholderOpNode::shape)
|
|
.def_ro("dtype", &PlaceholderOpNode::dtype);
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO("te.PlaceholderOp", PlaceholderOpNode, OperationNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to PlaceholderOpNode
|
|
* \sa PlaceholderOpNode
|
|
*/
|
|
class PlaceholderOp : public Operation {
|
|
public:
|
|
TVM_DLL PlaceholderOp(std::string name, ffi::Array<PrimExpr> shape, PrimType dtype);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PlaceholderOp, Operation, PlaceholderOpNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief A Compute op that compute a tensor on certain domain.
|
|
* This is the base class for ComputeOp (operating on a scalar at a time)
|
|
*/
|
|
class TVM_DLL BaseComputeOpNode : public OperationNode {
|
|
public:
|
|
/*! \brief IterVar on each axis */
|
|
ffi::Array<IterVar> axis;
|
|
/*! \brief IterVar on each reduction axis, if the body is a Reduce */
|
|
ffi::Array<IterVar> reduce_axis;
|
|
// override functions
|
|
ffi::Array<PrimExpr> output_shape(size_t idx) const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<BaseComputeOpNode>()
|
|
.def_ro("axis", &BaseComputeOpNode::axis)
|
|
.def_ro("reduce_axis", &BaseComputeOpNode::reduce_axis);
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO("te.BaseComputeOp", BaseComputeOpNode, OperationNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief A Compute op that compute a tensor on certain domain.
|
|
*/
|
|
class TVM_DLL ComputeOpNode : public BaseComputeOpNode {
|
|
public:
|
|
/*! \brief the compute expression */
|
|
ffi::Array<PrimExpr> body;
|
|
/*! \brief constructor */
|
|
ComputeOpNode() {}
|
|
// override functions
|
|
int num_outputs() const final;
|
|
PrimType output_dtype(size_t i) const final;
|
|
ffi::Array<Tensor> InputTensors() const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<ComputeOpNode>().def_ro("body", &ComputeOpNode::body);
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ComputeOp", ComputeOpNode, BaseComputeOpNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to ComputeOpNode
|
|
* \sa ComputeOpNode
|
|
*/
|
|
class ComputeOp : public Operation {
|
|
public:
|
|
TVM_DLL ComputeOp(std::string name, std::string tag, ffi::Map<ffi::String, ffi::Any> attrs,
|
|
ffi::Array<IterVar> axis, ffi::Array<PrimExpr> body);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ComputeOp, Operation, ComputeOpNode);
|
|
TVM_DEFINE_OBJECT_REF_COW_METHOD(ComputeOpNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief Symbolic scan.
|
|
*/
|
|
class ScanOpNode : public OperationNode {
|
|
public:
|
|
/*! \brief IterVar to scan over */
|
|
IterVar scan_axis;
|
|
/*! \brief the initialization tensors */
|
|
ffi::Array<Tensor> init;
|
|
/*! \brief the update function represented by tensor */
|
|
ffi::Array<Tensor> update;
|
|
/*! \brief The placeholder to refer as states in update. */
|
|
ffi::Array<Tensor> state_placeholder;
|
|
/*!
|
|
* \brief the inputs to the scan, these are optionally provided
|
|
* But they can be helpful to provide hints to speedup get of scan body.
|
|
*/
|
|
ffi::Array<Tensor> inputs;
|
|
/*!
|
|
* \brief Spatial axis to indicate spatial dimension of each output.
|
|
* They corresponds to flattened spatial axis of the outputs.
|
|
*
|
|
* [output[0].axis[1], output[0].axis[2]... output[k].axis[j]...]
|
|
* These are auxiliary data structure for storing result of bound inference.
|
|
* They do not corresponds to splittable iterations, thus the name comes
|
|
* with underscore.
|
|
*/
|
|
ffi::Array<IterVar> spatial_axis_;
|
|
/*! \brief constructor */
|
|
ScanOpNode() {}
|
|
// override behavior.
|
|
int num_outputs() const final;
|
|
PrimType output_dtype(size_t i) const final;
|
|
ffi::Array<PrimExpr> output_shape(size_t i) const final;
|
|
ffi::Array<Tensor> InputTensors() const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<ScanOpNode>()
|
|
.def_ro("scan_axis", &ScanOpNode::scan_axis)
|
|
.def_ro("init", &ScanOpNode::init)
|
|
.def_ro("update", &ScanOpNode::update)
|
|
.def_ro("state_placeholder", &ScanOpNode::state_placeholder)
|
|
.def_ro("inputs", &ScanOpNode::inputs)
|
|
.def_ro("spatial_axis_", &ScanOpNode::spatial_axis_);
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ScanOp", ScanOpNode, OperationNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to ScanOpNode
|
|
* \sa ScanOpNode
|
|
*/
|
|
class ScanOp : public Operation {
|
|
public:
|
|
TVM_DLL 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);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ScanOp, Operation, ScanOpNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief External computation that cannot be splitted.
|
|
*/
|
|
class ExternOpNode : public OperationNode {
|
|
public:
|
|
/*! \brief The input tensors */
|
|
ffi::Array<Tensor> inputs;
|
|
/*! \brief Symbolic placeholder representation of inputs */
|
|
ffi::Array<Buffer> input_placeholders;
|
|
/*! \brief Symbolic placeholder representation of outputs */
|
|
ffi::Array<Buffer> output_placeholders;
|
|
/*! \brief the statement that generates the computation. */
|
|
Stmt body;
|
|
|
|
/*! \brief constructor */
|
|
ExternOpNode() {}
|
|
// override functions
|
|
int num_outputs() const final;
|
|
PrimType output_dtype(size_t i) const final;
|
|
ffi::Array<PrimExpr> output_shape(size_t i) const final;
|
|
ffi::Array<Tensor> InputTensors() const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<ExternOpNode>()
|
|
.def_ro("inputs", &ExternOpNode::inputs)
|
|
.def_ro("input_placeholders", &ExternOpNode::input_placeholders)
|
|
.def_ro("output_placeholders", &ExternOpNode::output_placeholders)
|
|
.def_ro("body", &ExternOpNode::body);
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ExternOp", ExternOpNode, OperationNode);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to ExternOpNode
|
|
* \sa ExternOpNode
|
|
*/
|
|
class ExternOp : public Operation {
|
|
public:
|
|
TVM_DLL 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);
|
|
};
|
|
|
|
/*!
|
|
* \brief Construct a new Var expression
|
|
* \param name_hint The name hint for the expression
|
|
* \param t The type of the expression
|
|
*/
|
|
TVM_DLL PrimVar var(std::string name_hint, PrimType t = PrimType::Int(32));
|
|
|
|
/*!
|
|
* \brief Create a new IterVar that represents an axis in thread.
|
|
*
|
|
* \param dom Optional, domain of the thread axis.
|
|
* \param tag The thread tag of the axis.
|
|
*/
|
|
TVM_DLL IterVar thread_axis(Range dom, std::string tag);
|
|
|
|
/*!
|
|
* \brief Create a new IterVar for reduction operations.
|
|
*
|
|
* \param dom The domain of the reduction axis.
|
|
* \param name The name of the reduction axis.
|
|
*/
|
|
TVM_DLL IterVar reduce_axis(Range dom, std::string name = "rv");
|
|
|
|
/*! \brief The compute function to specify the input source of a Tensor */
|
|
using FCompute = std::function<PrimExpr(const ffi::Array<PrimVar>& i)>;
|
|
|
|
/*! \brief The compute function to specify the inputs source of Tensors */
|
|
using FBatchCompute = std::function<ffi::Array<PrimExpr>(const ffi::Array<PrimVar>& i)>;
|
|
|
|
/*!
|
|
* \brief create a place holder tensor.
|
|
* \param shape The shape of the tensor.
|
|
* \param dtype the data type of the tensor.
|
|
* \param name The name of the Tensor.
|
|
*/
|
|
TVM_DLL Tensor placeholder(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
|
|
std::string name = "placeholder");
|
|
|
|
/*!
|
|
* \brief Construct a new tensor by computing over shape,
|
|
* using the computation rule: result_tensor[axis] = fcompute(axis)
|
|
* \param shape Shape of the tensor.
|
|
* \param fcompute The compute function to create the tensor.
|
|
* \param name The optional name of the tensor.
|
|
* \param tag The optional tag of the tensor.
|
|
* \param attrs Optional additional attributes of the compute.
|
|
*/
|
|
TVM_DLL Tensor compute(ffi::Array<PrimExpr> shape, FCompute fcompute, std::string name = "tensor",
|
|
std::string tag = "", ffi::Map<ffi::String, ffi::Any> attrs = {});
|
|
|
|
/*!
|
|
* \brief Construct a new tensor by computing over shape,
|
|
* using the computation rule: result_tensor[axis] = fcompute(axis)
|
|
* \param shape Shape of the tensor.
|
|
* \param fcompute The compute function to create the tensors.
|
|
* \param name The optional name of the tensor.
|
|
* \param tag The optional tag of the tensor.
|
|
* \param attrs Optional additional attributes of the compute.
|
|
*/
|
|
TVM_DLL ffi::Array<Tensor> compute(ffi::Array<PrimExpr> shape, FBatchCompute fcompute,
|
|
std::string name = "tensor", std::string tag = "",
|
|
ffi::Map<ffi::String, ffi::Any> attrs = {});
|
|
|
|
/*!
|
|
* \brief Construct new tensors by scan.
|
|
*
|
|
* \param init The intialize tensor of first K steps.
|
|
* \param update The update tensor indicated the updated result after each timestamp.
|
|
* \param state_placeholder The placeholder for the states.
|
|
* \param inputs The inputs to the scan body, this is optional,
|
|
* but recommended to provide concrete information about scan body.
|
|
* \param name The optional name of the tensor.
|
|
* \param tag The optional tag of the tensor.
|
|
* \param attrs Optional additional attributes of the compute.
|
|
*/
|
|
TVM_DLL 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 = {});
|
|
|
|
// same as compute, specialized for different fcompute function
|
|
inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(PrimVar)> f,
|
|
std::string name = "tensor", std::string tag = "",
|
|
ffi::Map<ffi::String, ffi::Any> attrs = {}) {
|
|
FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0]); };
|
|
return compute(shape, fc, name, tag, attrs);
|
|
}
|
|
inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(PrimVar, PrimVar)> f,
|
|
std::string name = "tensor", std::string tag = "",
|
|
ffi::Map<ffi::String, ffi::Any> attrs = {}) {
|
|
FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0], i[1]); };
|
|
return compute(shape, fc, name, tag, attrs);
|
|
}
|
|
inline Tensor compute(ffi::Array<PrimExpr> shape,
|
|
std::function<PrimExpr(PrimVar, PrimVar, PrimVar)> f,
|
|
std::string name = "tensor", std::string tag = "",
|
|
ffi::Map<ffi::String, ffi::Any> attrs = {}) {
|
|
FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0], i[1], i[2]); };
|
|
return compute(shape, fc, name, tag, attrs);
|
|
}
|
|
inline Tensor compute(ffi::Array<PrimExpr> shape,
|
|
std::function<PrimExpr(PrimVar, PrimVar, PrimVar, PrimVar)> f,
|
|
std::string name = "tensor", std::string tag = "",
|
|
ffi::Map<ffi::String, ffi::Any> attrs = {}) {
|
|
FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0], i[1], i[2], i[3]); };
|
|
return compute(shape, fc, name, tag, attrs);
|
|
}
|
|
|
|
// inline function.
|
|
inline const OperationNode* Operation::operator->() const {
|
|
return static_cast<const OperationNode*>(get());
|
|
}
|
|
} // namespace te
|
|
} // namespace tvm
|
|
#endif // TVM_TE_OPERATION_H_
|