GitOrigin-RevId: 485290aa09
tags/v0.4.0
@@ -469,9 +469,7 @@ using Split = SplitForward; | |||||
* large number of inputs and can handle alignment requirements. Axis is also | * large number of inputs and can handle alignment requirements. Axis is also | ||||
* not supported. | * not supported. | ||||
* | * | ||||
* The offsets can be generated by gen_offsets(). The \p srcs in ParamPackSplit and | |||||
* \p dsts in ParamPackConcat must be on CPU, and must remain valid until the | |||||
* execution stream is synchronized. | |||||
* The offsets can be generated by gen_offsets(). | |||||
*/ | */ | ||||
class ParamPackConcatSplitBase : public OperatorBase { | class ParamPackConcatSplitBase : public OperatorBase { | ||||
protected: | protected: | ||||
@@ -502,7 +500,7 @@ public: | |||||
* address of i-th Tensor. | * address of i-th Tensor. | ||||
* \param[in] offsets: with size `2 * srcs.shape[0]`. | * \param[in] offsets: with size `2 * srcs.shape[0]`. | ||||
* offsets[i * 2] and offsets[i * 2 + 1] means | * offsets[i * 2] and offsets[i * 2 + 1] means | ||||
* the begin and the end of offset in | |||||
* the begin and the end of srcs[i]'s offsets in dst | |||||
* \param[out] dst: output TensorND, live on cpu or gpu | * \param[out] dst: output TensorND, live on cpu or gpu | ||||
*/ | */ | ||||
virtual void exec(_megdnn_tensor_in srcs, _megdnn_tensor_in offsets, | virtual void exec(_megdnn_tensor_in srcs, _megdnn_tensor_in offsets, | ||||
@@ -535,6 +535,10 @@ MGB_DEFINE_OPR_CLASS(Concat, cg::SingleCNOutshapePureByInshapeOprBase) // { | |||||
/*! | /*! | ||||
* \brief Opr used to pack parameter, all input node must in same device, dtype | * \brief Opr used to pack parameter, all input node must in same device, dtype | ||||
* and shape is not needed to be same | * and shape is not needed to be same | ||||
* \param offsets: size of 2 * inputs.size() | |||||
* offsets[i * 2] and offsets[i * 2 + 1] means | |||||
* the begin and the end of inputs[i]'s offsets in output | |||||
* \param offsets_val: offsets value on cpu | |||||
*/ | */ | ||||
MGB_DEFINE_OPR_CLASS(ParamPackConcat, cg::SingleCNOperatorNodeBase) // { | MGB_DEFINE_OPR_CLASS(ParamPackConcat, cg::SingleCNOperatorNodeBase) // { | ||||
//! input pointer buffer | //! input pointer buffer | ||||
@@ -577,6 +581,11 @@ public: | |||||
/*! | /*! | ||||
* \brief Opr used to split parameter | * \brief Opr used to split parameter | ||||
* \param offsets: size of 2 * outputs.size() | |||||
* offsets[i * 2] and offsets[i * 2 + 1] means | |||||
* the begin and the end of output[i]'s offsets in input | |||||
* \param offsets_val: offsets value on cpu | |||||
* \param shapes: shape of each output | |||||
*/ | */ | ||||
MGB_DEFINE_OPR_CLASS(ParamPackSplit, cg::SingleCNOperatorNodeBase) // { | MGB_DEFINE_OPR_CLASS(ParamPackSplit, cg::SingleCNOperatorNodeBase) // { | ||||
TensorShapeArray m_shapes; | TensorShapeArray m_shapes; | ||||