Browse Source

expend broadcast_gpu_kernel with truncate_div and truncate_mod

tags/v1.5.0-rc1
zong_shuai 4 years ago
parent
commit
f1eb2fe6bf
1 changed files with 16 additions and 8 deletions
  1. +16
    -8
      mindspore/ccsrc/backend/kernel_compiler/gpu/math/broadcast_gpu_kernel.h

+ 16
- 8
mindspore/ccsrc/backend/kernel_compiler/gpu/math/broadcast_gpu_kernel.h View File

@@ -165,14 +165,22 @@ class BroadcastOpGpuKernel : public GpuKernel {
}

static const std::map<std::string, BroadcastOpType> kBroadcastArithmetricTypeMap = {
{"Maximum", BROADCAST_TYPE_MAXIMUM}, {"Minimum", BROADCAST_TYPE_MINIMUM},
{"Pow", BROADCAST_TYPE_POWER}, {"RealDiv", BROADCAST_TYPE_REALDIV},
{"Mul", BROADCAST_TYPE_MUL}, {"Sub", BROADCAST_TYPE_SUB},
{"Add", BROADCAST_TYPE_ADD}, {"FloorDiv", BROADCAST_TYPE_FLOORDIV},
{"AbsGrad", BROADCAST_TYPE_ABSGRAD}, {"Div", BROADCAST_TYPE_DIV},
{"DivNoNan", BROADCAST_TYPE_DIVNONAN}, {"Mod", BROADCAST_TYPE_MOD},
{"FloorMod", BROADCAST_TYPE_FLOORMOD}, {"Atan2", BROADCAST_TYPE_ATAN2},
{"TruncateDiv", BROADCAST_TYPE_TRUNCATEDIV}, {"TruncateMod", BROADCAST_TYPE_TRUNCATEMOD},
{"Maximum", BROADCAST_TYPE_MAXIMUM},
{"Minimum", BROADCAST_TYPE_MINIMUM},
{"Pow", BROADCAST_TYPE_POWER},
{"RealDiv", BROADCAST_TYPE_REALDIV},
{"Mul", BROADCAST_TYPE_MUL},
{"Sub", BROADCAST_TYPE_SUB},
{"Add", BROADCAST_TYPE_ADD},
{"FloorDiv", BROADCAST_TYPE_FLOORDIV},
{"AbsGrad", BROADCAST_TYPE_ABSGRAD},
{"Div", BROADCAST_TYPE_DIV},
{"DivNoNan", BROADCAST_TYPE_DIVNONAN},
{"Mod", BROADCAST_TYPE_MOD},
{"FloorMod", BROADCAST_TYPE_FLOORMOD},
{"Atan2", BROADCAST_TYPE_ATAN2},
{"TruncateDiv", BROADCAST_TYPE_TRUNCATEDIV},
{"TruncateMod", BROADCAST_TYPE_TRUNCATEMOD},
};

iter = kBroadcastArithmetricTypeMap.find(kernel_name);


Loading…
Cancel
Save