diff --git a/mindspore/ccsrc/include/common/utils/utils.h b/mindspore/ccsrc/include/common/utils/utils.h index 2c78e008502..0cdd5657713 100644 --- a/mindspore/ccsrc/include/common/utils/utils.h +++ b/mindspore/ccsrc/include/common/utils/utils.h @@ -350,6 +350,10 @@ constexpr auto kEnvironSetOpName = "EnvironSet"; constexpr auto kEnvironGetOpName = "EnvironGet"; constexpr auto kEnvironDestroyAllOpName = "EnvironDestroyAll"; constexpr auto kUpdateStateOpName = "UpdateState"; +constexpr auto kPriorityReplayBufferCreate = "PriorityReplayBufferCreate"; +constexpr auto kPriorityReplayBufferPush = "PriorityReplayBufferPush"; +constexpr auto kPriorityReplayBufferSample = "PriorityReplayBufferSample"; +constexpr auto kPriorityReplayBufferUpdate = "PriorityReplayBufferUpdate"; // Communication world group constexpr auto kNcclWorldGroup = "nccl_world_group"; diff --git a/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_util.h b/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_util.h index 71554609482..aff71090b7d 100644 --- a/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_util.h +++ b/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_util.h @@ -75,14 +75,26 @@ constexpr auto kEnvironCreate = "EnvironCreate"; constexpr auto kEnvironSet = "EnvironSet"; constexpr auto kEnvironGet = "EnvironGet"; constexpr auto kEnvironDestroyAll = "EnvironDestroyAll"; +constexpr auto kPriorityReplayBufferCreate = "PriorityReplayBufferCreate"; +constexpr auto kPriorityReplayBufferPush = "PriorityReplayBufferPush"; +constexpr auto kPriorityReplayBufferSample = "PriorityReplayBufferSample"; +constexpr auto kPriorityReplayBufferUpdate = "PriorityReplayBufferUpdate"; const std::set kCpuKernelOps{kIdentity, kMaskedSelect, kMaskedSelectGrad, kDynamicStitch, kSearchSorted, kResizeBilinear, kResizeBilinearGrad, kScatterElements}; const std::set kCacheKernelOps{kUpdateCache, kCacheSwapTable, kSubAndFilter, kPadAndShift, kDropout3D, kDropout2D, kNonMaxSuppressionV3, kGetNext, kInitData, kPrint}; -const std::set kCpuKernelBaseOps{kRandomChoiceWithMask, kEnvironCreate, kEnvironSet, kEnvironGet, - kEnvironDestroyAll}; +const std::set kCpuKernelBaseOps{kRandomChoiceWithMask, + kEnvironCreate, + kEnvironSet, + kEnvironGet, + kEnvironDestroyAll, + kPriorityReplayBufferCreate, + kPriorityReplayBufferPush, + kPriorityReplayBufferSample, + kPriorityReplayBufferUpdate}; const std::set kDynamicInputOps{ - kPrint, kPack, kMeshgrid, kStackInitOpName, kStackDestroyOpName, kStackPushOpName, kStackPopOpName, kDynamicStitch}; + kPrint, kPack, kMeshgrid, kStackInitOpName, kStackDestroyOpName, + kStackPushOpName, kStackPopOpName, kDynamicStitch, kPriorityReplayBufferPush, kPriorityReplayBufferSample}; struct AicpuParamHead { uint32_t length; // Total length: include cunstom message uint32_t ioAddrNum; // Input and output address number diff --git a/mindspore/ccsrc/plugin/device/ascend/optimizer/ascend_backend_optimization.cc b/mindspore/ccsrc/plugin/device/ascend/optimizer/ascend_backend_optimization.cc index c0f782538c0..d0b66e47a53 100644 --- a/mindspore/ccsrc/plugin/device/ascend/optimizer/ascend_backend_optimization.cc +++ b/mindspore/ccsrc/plugin/device/ascend/optimizer/ascend_backend_optimization.cc @@ -139,7 +139,7 @@ #include "plugin/device/ascend/optimizer/enhancer/add_attr_for_3d_graph.h" #include "plugin/device/ascend/optimizer/enhancer/split_n_optimizer.h" #include "plugin/device/ascend/optimizer/mindir/space_batch_nd_attr_update.h" -#include "plugin/device/ascend/optimizer/mindir/env_op_attr_update.h" +#include "plugin/device/ascend/optimizer/mindir/aicpu_lib_select.h" #include "plugin/device/ascend/optimizer/mindir/dropout_unify_mindir.h" #include "plugin/device/ascend/optimizer/mindir/maxpool_to_maxpool_with_argmax.h" #include "plugin/device/ascend/optimizer/mindir/maxpool_with_argmax_unify_mindir.h" @@ -575,7 +575,7 @@ void AscendUnifyMindIR(const std::shared_ptr &graph) { auto unify_mindir_pm = std::make_shared("unify_mindir_pm"); unify_mindir_pm->AddPass(std::make_shared()); unify_mindir_pm->AddPass(std::make_shared()); - unify_mindir_pm->AddPass(std::make_shared()); + unify_mindir_pm->AddPass(std::make_shared()); unify_mindir_pm->AddPass(std::make_shared()); unify_mindir_pm->AddPass(std::make_shared()); unify_mindir_pm->AddPass(std::make_shared()); diff --git a/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/env_op_attr_update.cc b/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.cc similarity index 67% rename from mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/env_op_attr_update.cc rename to mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.cc index fc1ab6e24ae..8f01660a137 100644 --- a/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/env_op_attr_update.cc +++ b/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.cc @@ -14,7 +14,7 @@ * limitations under the License. */ -#include "plugin/device/ascend/optimizer/mindir/env_op_attr_update.h" +#include "plugin/device/ascend/optimizer/mindir/aicpu_lib_select.h" #include #include #include "include/common/utils/utils.h" @@ -23,21 +23,22 @@ namespace mindspore { namespace opt { -const AnfNodePtr EnvOpAttrUpdate::Process(const FuncGraphPtr &graph, const AnfNodePtr &node, - const EquivPtr &equiv) const { +const AnfNodePtr AICpuLibSelectPass::Process(const FuncGraphPtr &graph, const AnfNodePtr &node, + const EquivPtr &equiv) const { MS_EXCEPTION_IF_NULL(graph); MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(equiv); - static const std::set kEnvOpNames = {kEnvironCreateOpName, kEnvironSetOpName, kEnvironGetOpName, - kEnvironDestroyAllOpName}; + static const std::set kAICpuOpNames = { + kEnvironCreateOpName, kEnvironSetOpName, kEnvironGetOpName, kEnvironDestroyAllOpName, + kPriorityReplayBufferCreate, kPriorityReplayBufferPush, kPriorityReplayBufferSample, kPriorityReplayBufferUpdate}; static const std::string kEnvOpSoNames = "mindspore_aicpu_kernels"; if (!node->isa()) { return node; } auto kernel_name = common::AnfAlgo::GetCNodeName(node); - if (kEnvOpNames.find(kernel_name) != kEnvOpNames.end()) { + if (kAICpuOpNames.find(kernel_name) != kAICpuOpNames.end()) { common::AnfAlgo::SetNodeAttr(kAttrCustAicpu, MakeValue(kEnvOpSoNames), node); } diff --git a/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/env_op_attr_update.h b/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.h similarity index 66% rename from mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/env_op_attr_update.h rename to mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.h index b9c84e8b957..67620444ba2 100644 --- a/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/env_op_attr_update.h +++ b/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.h @@ -13,19 +13,19 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#ifndef MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_ENV_OP_ATTR_UPDATE_H_ -#define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_ENV_OP_ATTR_UPDATE_H_ +#ifndef MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_AICPU_LIB_SELECT_H_ +#define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_AICPU_LIB_SELECT_H_ #include "backend/common/optimizer/optimizer.h" namespace mindspore { namespace opt { -class EnvOpAttrUpdate : public PatternProcessPass { +class AICpuLibSelectPass : public PatternProcessPass { public: - explicit EnvOpAttrUpdate(bool multigraph = true) : PatternProcessPass("env_op_attr_update", multigraph) {} - ~EnvOpAttrUpdate() override = default; + explicit AICpuLibSelectPass(bool multigraph = true) : PatternProcessPass("env_op_attr_update", multigraph) {} + ~AICpuLibSelectPass() override = default; const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; }; } // namespace opt } // namespace mindspore -#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_ENV_OP_ATTR_UPDATE_H_ +#endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_ASCEND_MINDIR_AICPU_LIB_SELECT_H_