61 lines
2.5 KiB
C++
61 lines
2.5 KiB
C++
|
#include "register/op_def_registry.h"
|
||
|
|
|
||
|
|
namespace ops {
|
||
|
|
class NotifyDispatch : public OpDef {
|
||
|
|
public:
|
||
|
|
explicit NotifyDispatch(const char *name) : OpDef(name)
|
||
|
|
{
|
||
|
|
this->Input("sendData")
|
||
|
|
.ParamType(REQUIRED)
|
||
|
|
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
|
||
|
|
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||
|
|
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||
|
|
this->Input("tokenPerExpertData")
|
||
|
|
.ParamType(REQUIRED)
|
||
|
|
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
|
||
|
|
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||
|
|
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||
|
|
this->Output("sendDataOffset")
|
||
|
|
.ParamType(REQUIRED)
|
||
|
|
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
|
||
|
|
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||
|
|
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||
|
|
this->Output("recvData")
|
||
|
|
.ParamType(REQUIRED)
|
||
|
|
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
|
||
|
|
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||
|
|
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||
|
|
|
||
|
|
this->Attr("sendCount").Int();
|
||
|
|
this->Attr("num_tokens").Int();
|
||
|
|
this->Attr("comm_group").String();
|
||
|
|
this->Attr("rank_size").Int();
|
||
|
|
this->Attr("rank_id").Int();
|
||
|
|
this->Attr("local_rank_size").Int();
|
||
|
|
this->Attr("local_rank_id").Int();
|
||
|
|
|
||
|
|
OpAICoreConfig aicore_config_base;
|
||
|
|
aicore_config_base.DynamicCompileStaticFlag(true)
|
||
|
|
.DynamicFormatFlag(true)
|
||
|
|
.DynamicRankSupportFlag(true)
|
||
|
|
.DynamicShapeSupportFlag(true)
|
||
|
|
.NeedCheckSupportFlag(false)
|
||
|
|
.PrecisionReduceFlag(true)
|
||
|
|
.ExtendCfgInfo("aclnnSupport.value", "support_aclnn")
|
||
|
|
.ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel");
|
||
|
|
|
||
|
|
OpAICoreConfig aicore_config_A2 = aicore_config_base;
|
||
|
|
aicore_config_A2.ExtendCfgInfo("jitCompile.flag", "static_false");
|
||
|
|
|
||
|
|
OpAICoreConfig aicore_config = aicore_config_base;
|
||
|
|
aicore_config.ExtendCfgInfo("jitCompile.flag", "static_true");
|
||
|
|
|
||
|
|
this->AICore().AddConfig("ascend910_93", aicore_config);
|
||
|
|
this->AICore().AddConfig("ascend910b", aicore_config_A2);
|
||
|
|
this->MC2().HcclGroup("comm_group");
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
OP_ADD(NotifyDispatch);
|
||
|
|
} // namespace ops
|