diff --git a/atvc/docs/02_developer_guide.md b/atvc/docs/02_developer_guide.md index 455a522e92b532c0f10c745b682ef6a8f2269ef4..c24e416bcf26965c4914e6b598ff423e78b2597d 100644 --- a/atvc/docs/02_developer_guide.md +++ b/atvc/docs/02_developer_guide.md @@ -1382,6 +1382,225 @@ int32_t main(int32_t argc, char* argv[]) } ``` +## 2.4 组合算子开发 +ATVC框架支持Broadcast与Elementwise组合的算子通过扩展BroadcastOpTemplate的模板参数对用户提供接口，开发者可以根据算子实际需求来定制组合，框架支持以下组合：Elementwise + Broadcast、Broadcast + Elementwise、Elementwise + Broadcast + Elementwise。下面以Broadcast与Elementwise组合为例进行详细讲解。 + +组合算子模板类的模块之间的交互如下（ATVC框架提供的模板及接口用黄色表示；开发自定义的模块用蓝色表示）： + +![](images/broadcast_fusion_dataflow.png) + +组合算子模板内部根据计算的数据大小，shape完成了不同计算调度代码的实现，考虑到Broadcast的tiling复杂度，组合算子的计算调度场景复用Broadcast的调度策略`ATVC::BroadcastPolicy`。在算子调用阶段，分派策略API可根据Tiling API计算出的`ATVC::BroadcastPolicy`转化为编译态参数，结合定制的Elementwise和Broadcast计算模板来实例化`ATVC::Kernel::BroadcastOpTemplate`算子模板类。 + +### 2.4.1 Components +根据组合算子在框架内部的交互场景，ATVC提供如下的接口及模板类帮助开发搭建自定义Broadcast与Elementwise组合算子： + +需按照以下顺序完成模块之间的组装： + +1.自定义前置或后置Elementwise计算模板 + +2.自定义Broadcast计算模板/使用框架内置Broadcast计算模板，并组合Elementwise计算模板。 + +3.将计算模板传入`Kernel`层模板算子完成核函数功能实现； + +4.定义Kernel层算子入口API， 内部实例化计算模板类； + +下面将以Add算子（Broadcast + Elementwise, 为区别Add单算字，命名为AddWithBroadcast算子）搭建为样例，按照组装顺序介绍组合算子类的开发流程。 + +### 2.4.2 计算模板 +组合计算模板复用已有的Elementwise计算模板（详见[2.1.2章节](#212-计算模板)）和Broadcst计算模板（参见[2.3.2章节](#232-计算模板)），具体使用方法和约束参考对应章节。 + +根据实际的算子诉求，构建1个或2个Elementwise计算模板，与1个Broadcast计算模板，作为模板参数传入`ATVC：：BroadcastOpTemplate`中，并在数据计算以及同步阶段被调用。 + +存在Broadcast的组合计算模板设计多核之间的数据结果同步以及核内分块的对其计算，用户自定义难度较高，因此ATVC框架提供了Broadcast的内置计算模板，并实现了Broadcast在单核内与多核间的计算与同步等函数接口。 +#### 2.4.2.1 Elementwise计算模板 +前置或后置Elementwise计算模板除了需要满足基本计算模板的要求外，还需要定义两个额外接口SetArgs和SetParam，分别用来接受组合算子的向量参数和标量参数。 + +以PostCompute为例， Elementwise计算模板定义如下： + +```cpp +template +struct PostComputeAddOfBroadcast { + /* ! + * \brief set scaler param for compute fuction + * \param [in] args, args are mutable parameters, and are passed transparently from the parameters of + * global kernel functions, which are the parameters after broadcastParam + */ + template + __aicore__ inline void SetParam(Args... args) {} + + /* ! + * \brief set tensor param for compute fuction + * \param [in] args, args are mutable parameters, and are passed transparently from the parameters of + * global kernel functions, which are the parameters before broadcastParam, the num of args is + * decided by Traits + */ + template + __aicore__ inline void SetArgs(Args... args) {} + + /* ! + * \brief process function of compute struct + * \param [in] y, local tensor of y + * \param [in] z, local tensor of z + * \param [in] x, local tensor of x, x is the output of broadcast, must be the last local tensor + * \param [in] copyOutOffset, copy out offset for DataCopy + * \param [in] copyOutParams, copy out params for DataCopy + */ + template + __aicore__ inline void operator()(AscendC::LocalTensor y, AscendC::LocalTensor z, AscendC::LocalTensor x, + uint32_t copyOutOffset, AscendC::DataCopyExtParams ©OutParams) {} +}; +``` + +Elementwise计算模板需要定义三个接口， 分别为： + +1、SetParam函数，入参为可变参数，用户调用Kernel函数传递的scaler参数经过Broadcast模板的分拣，会通过SetParam函数传递给每一个计算函数模板。 + +2、SetArgs函数，入参为可变参数，用户调用Kernel函数传递的向量参数经过Broadcast模板的分拣，会根据OpTrats计算每个计算模板需要的参数个数，并通过SetArgs函数按顺序传递给对应的计算模板。 + +3、()函数，入参为可变参数。用户调用Kernel函数后，Broadcast模板会判断用户是否有定义PreCompute或PostCompute函数，并用PreCompute替换CopyIn操作，用PostCompute替换CopyOut操作。该函数的参数按顺序主要分为3部分： + +* 计算模板所需要的LocalTensor，包括输入输出和临时Buffer，内存由Broadcast模板申请，用户可以直接使用，参数的个数模板的所有参数个数-1（剩下一个参数必为Broadcast模板的输入或者输出，不需要额外申请） +* 单个LocalTensor，复用Broadcast模板的输入或者输出。 +* CopyIn或CopyOut的参数，包括offset和DataCopyParam，为原来CopyIn或CopyOut操作时， 调用AscendC::DataCopy接口所传递的参数， 用户可以在完成自定义计算后，直接使用该参数做DataCopy操作。 + + +完整实现见[post_compute_add_of_broadcast.h](../examples/add_with_broadcast/post_compute_add_of_broadcast.h)。 + + + + +#### 2.4.2.2 Broadcast计算模板 +Broadcast计算模板在组合算子中与Broadcast单算子无任何区别，可直接复用单算子计算模板或使用内置计算模板。 + +### 2.4.3 BroadcastOpTemplate模板 +Broadcast与Elementwise组合的算子模板以BroadcastOpTemplate为基础进行扩展,`BroadcastOpTemplate`的介绍可以参考章节[2.3.3](#233-内置broadcast算子模板)。下面为组合算子场景`ATVC::Kernel::BroadcastOpTemplate`新引入的接口或定义，以及调用计算模板函数的示意代码，完整模板定义请参考[`atvc/broadcast/broadcast_op_template.h`](../include/broadcast/broadcast_op_template.h)： + +```cpp +template +class BroadcastOpTemplate { + // 萃取PreCompute和PostCompute中的信息 + static constexpr bool HAS_PRE_COMPUTE = !AscendC::Std::is_same_v; + static constexpr bool HAS_POST_COMPUTE = !AscendC::Std::is_same_v; + using PreComputeTraits = AscendC::Std::conditional_t::ComputeTraits, VoidComputeTraits>; + using PostComputeTraits = AscendC::Std::conditional_t::ComputeTraits, VoidComputeTraits>; + using PreInputs = typename PreComputeTraits::In::types; + using PreOutputs = typename PreComputeTraits::Out::types; + using PreTemp = typename PreComputeTraits::Temp::types; + using PostInputs = typename PostComputeTraits::In::types; + using PostOutputs = typename PostComputeTraits::Out::types; + using PostTemp = typename PostComputeTraits::Temp::types; + using DataType = typename BroadcastCompute::DataType; + + + template + __aicore__ inline void SetParam(BroadcastParam *param, Args... args) + { + if constexpr (HAS_PRE_COMPUTE) { + preCompute_.SetParam(args...); + } + if constexpr (HAS_POST_COMPUTE) { + postCompute_.SetParam(args...); + } + } + // Run接口参数改为可变参数,用来传递PreCompute和PostCompute的参数。 + template + __aicore__ inline void Run(Args&&... args) + { + // 分拣出tensor参数并按使用个数传递给计算函数 + preCompute_.SetArgs(arg1...); + broadcastCompute_.SetArgs(arg2...); + postCompute_.SetArgs(args3...); + + // 分拣出scaler参数并传给PreCompute和PostCompute + SetParam(scalerArgs); + } + + __aicore__ inline void CopyIn(AscendC::LocalTensor &input, uint32_t copyInOffset, BroadcastDataView &view) + { + if constexpr(HAS_PRE_COMPUTE) { + ProcessPreCompute(input, copyInOffset, copyInParams); + return; + } + } + + __aicore__ inline void CopyOut(AscendC::LocalTensor &output, + uint32_t copyOutOffset, BroadcastDataView &view) + { + if constexpr(HAS_POST_COMPUTE) { + ProcessPostCompute(output, copyOutOffset, copyOutParams); + return; + } + } +}; +``` + +### 2.4.4 核函数定义 +在ATVC提供了Broadcast内部实现后，用户需要定义封装核函数接口。核函数功能可依赖`ATVC::Kernel::BroadcastOpTemplate`的相关接口实现。 + +基于`2.4.2`和`2.4.3`的样例代码，Kernel层的自定义API样例如下： + +```cpp +// AddWithBroadcast算子的描述：两个输入，一个输出，类型均为float +using BroadcastOpTraits = ATVC::OpTraits, ATVC::OpOutputs, ATVC::OpTemps>; + +/* ! + * \brief z = x + y, the shape of x must be able to be broadcasted to the shape of y + * \param [in] x, input global memory of x + * \param [in] y, input global memory of y + * \param [out] z, output global memory + * \param [in] broadcastParam, params of broadcast + */ +template +__global__ __aicore__ void AddWithBroadcastCustom(GM_ADDR x, + GM_ADDR y, + GM_ADDR z, + ATVC::BroadcastParam broadcastParam) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); + + // 1. get input and output for kernel op from host Traits + using KernelOpIn = typename Traits::In::types; + using KernelOpOut = typename Traits::Out::types; + using KernelOpTemp = typename Traits::Temp::types; + + // 2. define input and output for broadcast + using BroadcastOpInput = ATVC::OpInputs::Type>; + using BroadcastOpOutput = ATVC::OpOutputs::Type>; + using BroadcastOpTraits = ATVC::OpTraits; + + // 3. define input and output for post compute + using AddOpInput = ATVC::OpInputs::Type, typename ATVC::TypeListGet::Type>; + using AddOpOutput = ATVC::OpOutputs::Type>; + using AddOpTraits = ATVC::OpTraits; + using PostCompute = PostComputeAddOfBroadcast; + + // 4. call op run + auto op = ATVC::Kernel::BroadcastOpTemplate, Policy, void, PostCompute>(); + ATVC::BroadcastParam *param = &broadcastParam; + op.Run(x, y, z, param); +} +``` + +在组合算子中，用户定义的OpTraits为组合算子的整体定义，核函数内部根据算子的组合形式，用组合算子的OpTraits定义来表达Elementwise和Broadcast算子的Optraits定义。例如，Broadcast计算函数的OpTraits定义就表示：broadcast的输入是组合算子的第一个输入，broadcast算子的输出组合算子的临时资源。 + + +### 2.4.5 Host层API +#### 2.4.5.1 CalcBroadcastTiling +组合类算子的TilingAPI可以复用`ATVC::Host::CalcBroacastTiling`功能，在此框架上引入PreCompute和PostCompute对应的OpTraits。具体信息参考[2.2.4](#224-核函数定义)章节。在计算基本快所需UB空间时，从PreCompute和PostCompute的OpTraits萃取对于UB空间的需求，确保不会出现溢出的情况。简单的，如果用户不用对Tiling做扩展，可以直接用组合算子的OpTraits来计算Tiling，不需要分别传单算子的OpTraits。 + +```cpp +// AddWithBroadcast算子的描述：两个输入，一个输出，类型均为float +using BroadcastOpTraits = ATVC::OpTraits, ATVC::OpOutputs, ATVC::OpTemps>; +ATVC::Host::CalcBroadcastTiling(shapeIn, shapeOut, &policy, ¶m); +``` + +#### 2.4.5.2 BroadcastAdapter +`BroadcastAdapter`的介绍可参考章节2.3.5.2。 + +### 2.4.6 完整样例 +参考sample [add_with_broadcast](../examples/add_with_broadcast/README.md) + + # 3 ATVC的调试调优功能 为了用户在使用ATVC进行算子开发时能快速进行精度调试和性能调优，ATVC支持多种调试调优能力。 ## 3.1 OpTraits校验接口 @@ -1394,6 +1613,7 @@ bool DebugCheck() } } ``` + 其中，模板参数`OpTraits`是用户定义的待校验的输入输出描述信息, 模板参数`templateType`是校验规则分类的标识， 定义如下: ```cpp enum class TemplateType {