!31724 Refactor gpu factory and coexist with old one for now.

Merge pull request !31724 from TronZhang/refactor_gpu_factory
This commit is contained in:
i-robot 2022-03-26 09:01:44 +00:00 committed by Gitee
commit ff8d2e612b
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
15 changed files with 610 additions and 285 deletions

View File

@ -56,15 +56,28 @@ void CreateGPUKernel(const std::vector<CNodePtr> &kernels) {
}
akg_nodes.push_back(kernel);
} else if (!common::AnfAlgo::IsControlOpExecInBackend(kernel)) {
auto gpu_kernel_ptr = kernel::NativeGpuKernelModFactory::GetInstance().Create(kernel_name, kernel);
if (!gpu_kernel_ptr) {
std::shared_ptr<kernel::NativeGpuKernelMod> gpu_kernel_mod = nullptr;
bool new_factory = true;
// TODO(tronzhang): When old kernel has been rectified, remove the condition and keep the true branch.
if (kernel::Factory<kernel::NativeGpuKernelMod>::Instance().IsRegistered(kernel_name)) {
gpu_kernel_mod = kernel::Factory<kernel::NativeGpuKernelMod>::Instance().Create(kernel_name);
} else {
gpu_kernel_mod =
(std::shared_ptr<kernel::NativeGpuKernelMod>)(kernel::NativeGpuKernelModFactory::GetInstance().Create(
kernel_name, kernel));
new_factory = false;
}
if (!gpu_kernel_mod) {
MS_LOG(EXCEPTION) << "Build gpu kernel op[" << kernel->fullname_with_scope() << "] failed";
}
MS_EXCEPTION_IF_NULL(kernel);
if (!gpu_kernel_ptr->Init(kernel)) {
if (new_factory) {
gpu_kernel_mod->SetGpuRefMapToKernelInfo(kernel);
}
if (!gpu_kernel_mod->Init(kernel)) {
MS_LOG(EXCEPTION) << "Initialize gpu kernel op[" << kernel->fullname_with_scope() << "] failed.";
}
session::AnfRuntimeAlgorithm::SetKernelMod((kernel::KernelModPtr)gpu_kernel_ptr, kernel.get());
session::AnfRuntimeAlgorithm::SetKernelMod(gpu_kernel_mod, kernel.get());
}
}

View File

@ -19,7 +19,8 @@
#include <memory>
#include <string>
#include "kernel/common_utils.h"
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
#include "plugin/factory/ms_factory.h"
#include "plugin/device/gpu/kernel/gpu_kernel.h"
#include "kernel/kernel.h"
#include "kernel/kernel_build_info.h"
#include "kernel/oplib/opinfo.h"
@ -82,16 +83,41 @@ bool CheckKernelInfo(const std::shared_ptr<KernelBuildInfo> &alternative_kernel_
return true;
}
std::string SupportedTypeList(const CNodePtr &kernel_node, KernelType kernel_type) {
std::string GetSupportedTypesStr(const CNodePtr &kernel_node, KernelType kernel_type) {
std::string supported_type_lists;
// Custom op gets reg info from OpLib instead of NativeGpuKernelModFactory.
// Custom op gets reg info from OpLib instead of NativeGpuKernelMod.
if (!IsPrimitiveCNode(kernel_node, prim::kPrimCustom)) {
supported_type_lists =
kernel::NativeGpuKernelModFactory::GetInstance().SupportedTypeList(common::AnfAlgo::GetCNodeName(kernel_node));
if (!supported_type_lists.empty()) {
return supported_type_lists;
auto kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
// TODO(tronzhang): When old kernel has been rectified, remove the condition and keep the true branch.
if (kernel::Factory<kernel::NativeGpuKernelMod>::Instance().IsRegistered(kernel_name)) {
auto kernel_attr_list = kernel::NativeGpuKernelMod::GetGpuSupportedList(kernel_name);
if (!kernel_attr_list.empty()) {
for (size_t attr_index = 0; attr_index < kernel_attr_list.size(); ++attr_index) {
std::string type_list = "input[";
auto attr = kernel_attr_list[attr_index];
for (size_t input_index = 0; input_index < attr.GetInputSize(); ++input_index) {
type_list = type_list + TypeIdToString(attr.GetInputAttr(input_index).first) +
((input_index == (attr.GetInputSize() - 1)) ? "" : " ");
}
type_list = type_list + "], output[";
for (size_t input_index = 0; input_index < attr.GetOutputSize(); ++input_index) {
type_list = type_list + TypeIdToString(attr.GetOutputAttr(input_index).first) +
((input_index == (attr.GetOutputSize() - 1)) ? "" : " ");
}
supported_type_lists = supported_type_lists + type_list + "]; ";
}
return supported_type_lists;
}
} else {
supported_type_lists =
kernel::NativeGpuKernelModFactory::GetInstance().SupportedTypeList(common::AnfAlgo::GetCNodeName(kernel_node));
if (!supported_type_lists.empty()) {
return supported_type_lists;
}
}
}
std::vector<std::shared_ptr<KernelBuildInfo>> kernel_info_list;
std::string op_name = common::AnfAlgo::GetCNodeName(kernel_node);
kernel::OpImplyType imply_type = GetImplyType(kernel_type);
@ -116,6 +142,7 @@ std::string SupportedTypeList(const CNodePtr &kernel_node, KernelType kernel_typ
}
supported_type_lists = supported_type_lists + supported_akg_type_list + "]; ";
}
return supported_type_lists;
}
@ -164,9 +191,9 @@ bool SelectCustomKernel(const CNodePtr &kernel_node, const std::shared_ptr<Kerne
auto func_type = common::AnfAlgo::GetNodeAttr<std::string>(kernel_node, kAttrFuncType);
if (func_type == kCustomTypeAOT) {
*kernel_type = KernelType::GPU_KERNEL;
if (!kernel::NativeGpuKernelModFactory::GetInstance().SearchRegistered(op_name, selected_kernel_info)) {
kernel::GpuKernelRegister(op_name, mindspore::kernel::KernelAttr(),
[]() { return new kernel::CustomAOTGpuKernelMod(); });
if (!kernel::Factory<kernel::NativeGpuKernelMod>::Instance().IsRegistered(op_name)) {
kernel::Factory<kernel::NativeGpuKernelMod>::Instance().Register(
op_name, []() { return std::make_shared<kernel::CustomAOTGpuKernelMod>(); });
}
} else if (kCustomTypeAkg.find(func_type) != kCustomTypeAkg.end()) {
*kernel_type = KernelType::AKG_KERNEL;
@ -200,7 +227,8 @@ bool SelectCustomKernel(const CNodePtr &kernel_node, const std::shared_ptr<Kerne
return true;
}
void SetTensorDeviceInfo(const kernel::KernelBuildInfo &selected_kernel_info, const CNodePtr &kernel_node) {
void SetTensorDeviceInfo(const kernel::KernelBuildInfo &selected_kernel_info, const CNodePtr &kernel_node,
const std::vector<std::tuple<size_t, TypeId, TypeId>> &input_reduce_detail) {
MS_EXCEPTION_IF_NULL(kernel_node);
size_t input_num = common::AnfAlgo::GetInputTensorNum(kernel_node);
for (size_t input_index = 0; input_index < input_num; ++input_index) {
@ -229,17 +257,27 @@ void SetTensorDeviceInfo(const kernel::KernelBuildInfo &selected_kernel_info, co
(common::AnfAlgo::GetCNodeName(kernel_node) == "ApplyMomentum")) {
std::vector<std::string> output_format = {selected_kernel_info.GetInputFormat(input_index)};
builder->SetOutputsFormat(output_format);
auto reduce_flag = kernel::NativeGpuKernelModFactory::GetInstance().reduce_flag_;
std::vector<TypeId> output_type;
// TODO(tronzhang): When old kernel has been rectified, remove the condition and keep the false branch.
auto reduce_flag = kernel::NativeGpuKernelModFactory::GetInstance().reduce_flag_;
if (std::find(reduce_flag.first.begin(), reduce_flag.first.end(), input_index) != reduce_flag.first.end()) {
output_type = {reduce_flag.second};
} else {
output_type = {selected_kernel_info.GetInputDeviceType(input_index)};
auto iter = std::find_if(input_reduce_detail.begin(), input_reduce_detail.end(),
[input_index](const std::tuple<size_t, TypeId, TypeId> &reduce_detail) {
return std::get<0>(reduce_detail) == input_index;
});
if (iter != input_reduce_detail.end()) {
output_type = {std::get<1>(*iter)};
} else {
output_type = {selected_kernel_info.GetInputDeviceType(input_index)};
}
}
builder->SetOutputsDeviceType(output_type);
AnfAlgo::SetSelectKernelBuildInfo(builder->Build(), real_input_node.get());
}
}
// TODO(tronzhang): When old kernel has been rectified, remove the call of NativeGpuKernelModFactory.
kernel::NativeGpuKernelModFactory::GetInstance().reduce_flag_.first.clear();
}
@ -397,7 +435,7 @@ void SetGraphKernelInfo(const CNodePtr &kernel_node, const FuncGraphPtr &func_gr
auto graph_selected_info = graph_info_builder.Build();
MS_EXCEPTION_IF_NULL(graph_selected_info);
AnfAlgo::SetSelectKernelBuildInfo(graph_selected_info, kernel_node.get());
SetTensorDeviceInfo(*graph_selected_info, kernel_node);
SetTensorDeviceInfo(*graph_selected_info, kernel_node, {});
}
void PrintUnsupportedTypeException(const CNodePtr &kernel_node, const std::vector<TypeId> &inputs_type,
@ -410,7 +448,7 @@ void PrintUnsupportedTypeException(const CNodePtr &kernel_node, const std::vecto
std::for_each(std::begin(outputs_type), std::end(outputs_type),
[&build_type](auto i) { build_type += TypeIdToString(i) + " "; });
build_type += "]";
auto supported_type_lists = SupportedTypeList(kernel_node, kernel_type);
auto supported_type_lists = GetSupportedTypesStr(kernel_node, kernel_type);
MS_EXCEPTION(TypeError) << "Select GPU kernel op[" << kernel_name
<< "] fail! Incompatible data type!\nThe supported data types are " << supported_type_lists
<< ", but get " << build_type;
@ -487,16 +525,42 @@ void SetKernelInfo(const CNodePtr &kernel_node, KernelType kernel_type) {
builder->SetOutputsFormat(outputs_format);
builder->SetOutputsDeviceType(outputs_type);
bool result = false;
std::vector<std::tuple<size_t, TypeId, TypeId>> input_reduce_index;
std::vector<std::tuple<size_t, TypeId, TypeId>> output_reduce_index;
if (IsPrimitiveCNode(kernel_node, prim::kPrimCustom)) {
// Custom op select kernel from OpLib
result = SelectCustomKernel(kernel_node, builder->Build(), &kernel_type);
} else if (kernel_type == UNKNOWN_KERNEL_TYPE) {
result = kernel::NativeGpuKernelModFactory::GetInstance().SearchRegistered(
common::AnfAlgo::GetCNodeName(kernel_node), builder->Build());
if (!result) {
result = kernel::NativeGpuKernelModFactory::GetInstance().ReducePrecision(
common::AnfAlgo::GetCNodeName(kernel_node), builder);
auto kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
// TODO(tronzhang): When old kernel has been rectified, remove the condition and keep the true branch.
if (kernel::Factory<kernel::NativeGpuKernelMod>::Instance().IsRegistered(kernel_name)) {
result = kernel::NativeGpuKernelMod::GpuCheckSupport(kernel_name, GetKernelAttrFromBuildInfo(builder->Build()));
if (!result) {
std::tie(result, input_reduce_index, output_reduce_index) = kernel::NativeGpuKernelMod::GpuReducePrecisionCheck(
kernel_name, GetKernelAttrFromBuildInfo(builder->Build()));
if (result) {
const size_t kReduceToTypeIdx = 2;
for (const auto &item : input_reduce_index) {
auto idx = std::get<0>(item);
auto to_type_id = std::get<kReduceToTypeIdx>(item);
builder->SetInputDeviceType(to_type_id, idx);
}
for (const auto &item : output_reduce_index) {
auto idx = std::get<0>(item);
auto to_type_id = std::get<kReduceToTypeIdx>(item);
builder->SetOutputDeviceType(to_type_id, idx);
}
}
}
} else {
result = kernel::NativeGpuKernelModFactory::GetInstance().SearchRegistered(
common::AnfAlgo::GetCNodeName(kernel_node), builder->Build());
if (!result) {
result = kernel::NativeGpuKernelModFactory::GetInstance().ReducePrecision(
common::AnfAlgo::GetCNodeName(kernel_node), builder);
}
}
if (!result && (!common::AnfAlgo::IsControlOpExecInBackend(kernel_node))) {
result = SelectAkgKernel(kernel_node, builder->Build());
kernel_type = AKG_KERNEL;
@ -511,7 +575,7 @@ void SetKernelInfo(const CNodePtr &kernel_node, KernelType kernel_type) {
builder->SetKernelType(kernel_type);
builder->SetProcessor(kernel::Processor::CUDA);
AnfAlgo::SetSelectKernelBuildInfo(builder->Build(), kernel_node.get());
SetTensorDeviceInfo(*(builder->Build()), kernel_node);
SetTensorDeviceInfo(*(builder->Build()), kernel_node, input_reduce_index);
}
} // namespace gpu
} // namespace device

View File

@ -15,28 +15,68 @@
*/
#include "plugin/device/gpu/kernel/arrays/batchtospace_gpu_kernel.h"
#include <algorithm>
#include <utility>
namespace mindspore {
namespace kernel {
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
BatchToSpaceGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
BatchToSpaceGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
BatchToSpaceGpuKernelMod, int)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64),
BatchToSpaceGpuKernelMod, int64_t)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeInt16).AddOutputAttr(kNumberTypeInt16),
BatchToSpaceGpuKernelMod, int16_t)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeInt8).AddOutputAttr(kNumberTypeInt8),
BatchToSpaceGpuKernelMod, int8_t)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeUInt8).AddOutputAttr(kNumberTypeUInt8),
BatchToSpaceGpuKernelMod, uint8_t)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeUInt16).AddOutputAttr(kNumberTypeUInt16),
BatchToSpaceGpuKernelMod, uint16_t)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeUInt32).AddOutputAttr(kNumberTypeUInt32),
BatchToSpaceGpuKernelMod, uint32_t)
MS_REG_GPU_KERNEL_ONE(BatchToSpace, KernelAttr().AddInputAttr(kNumberTypeUInt64).AddOutputAttr(kNumberTypeUInt64),
BatchToSpaceGpuKernelMod, uint64_t)
namespace {
template <typename T>
std::unique_ptr<cukernel::GpuKernelHelperBase> CreateBTSKernelPtr(const std::string &kernel_name) {
return std::make_unique<cukernel::BatchToSpaceHelperGpuKernel<T>>(kernel_name);
}
using BTSPtrCreatorFunc = std::function<std::unique_ptr<cukernel::GpuKernelHelperBase>(const std::string &)>;
const std::vector<std::pair<KernelAttr, BTSPtrCreatorFunc>> kernel_attr = {
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateBTSKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateBTSKernelPtr<half>},
{KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32), CreateBTSKernelPtr<int>},
{KernelAttr().AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64), CreateBTSKernelPtr<int64_t>},
{KernelAttr().AddInputAttr(kNumberTypeInt16).AddOutputAttr(kNumberTypeInt16), CreateBTSKernelPtr<int16_t>},
{KernelAttr().AddInputAttr(kNumberTypeInt8).AddOutputAttr(kNumberTypeInt8), CreateBTSKernelPtr<int8_t>},
{KernelAttr().AddInputAttr(kNumberTypeUInt8).AddOutputAttr(kNumberTypeUInt8), CreateBTSKernelPtr<uint8_t>},
{KernelAttr().AddInputAttr(kNumberTypeUInt16).AddOutputAttr(kNumberTypeUInt16), CreateBTSKernelPtr<uint16_t>},
{KernelAttr().AddInputAttr(kNumberTypeUInt32).AddOutputAttr(kNumberTypeUInt32), CreateBTSKernelPtr<uint32_t>},
{KernelAttr().AddInputAttr(kNumberTypeUInt64).AddOutputAttr(kNumberTypeUInt64), CreateBTSKernelPtr<uint64_t>}};
} // namespace
bool BatchToSpaceGpuKernelMod::Init(const CNodePtr &kernel_node) {
kernel_node_ = kernel_node;
kernel_name_ = common::AnfAlgo::GetCNodeName(kernel_node);
auto index = GetMatchKernelAttrIdxWithException(kernel_node, GetOpSupport());
helper_ptr_ = std::move(kernel_attr[index].second(kernel_name_));
helper_ptr_->ResetResource();
std::vector<std::vector<size_t>> input_shapes;
std::vector<std::vector<size_t>> output_shapes;
auto input_shape = AnfAlgo::GetInputDeviceShapeAdaptively(kernel_node, 0);
auto output_shape = AnfAlgo::GetOutputDeviceShapeAdaptively(kernel_node, 0);
input_shapes.emplace_back(input_shape);
output_shapes.emplace_back(output_shape);
attr_.block_size = GetAttr<int64_t>(kernel_node, "block_size");
attr_.crops = GetAttr<std::vector<std::vector<int64_t>>>(kernel_node, "crops");
attr_.input_shape = input_shape;
int flag = helper_ptr_->CheckKernelParam(&attr_);
if (flag != 0) {
return false;
}
flag = helper_ptr_->CalMemSize(input_shapes, output_shapes);
if (flag != 0) {
return false;
}
InitSizeLists();
return true;
}
std::vector<KernelAttr> BatchToSpaceGpuKernelMod::GetOpSupport() {
std::vector<KernelAttr> support_list;
(void)std::transform(kernel_attr.begin(), kernel_attr.end(), std::back_inserter(support_list),
[](const std::pair<KernelAttr, BTSPtrCreatorFunc> &item) { return item.first; });
return support_list;
}
MS_KERNEL_FACTORY_REG(NativeGpuKernelMod, BatchToSpace, BatchToSpaceGpuKernelMod);
} // namespace kernel
} // namespace mindspore

View File

@ -26,7 +26,6 @@
namespace mindspore {
namespace kernel {
template <typename T>
class BatchToSpaceGpuKernelMod : public NativeGpuKernelMod {
public:
BatchToSpaceGpuKernelMod() { ResetResource(); }
@ -44,33 +43,11 @@ class BatchToSpaceGpuKernelMod : public NativeGpuKernelMod {
return true;
}
bool Init(const CNodePtr &kernel_node) override {
kernel_node_ = kernel_node;
kernel_name_ = common::AnfAlgo::GetCNodeName(kernel_node);
helper_ptr_ = std::make_unique<cukernel::BatchToSpaceHelperGpuKernel<T>>(kernel_name_);
helper_ptr_->ResetResource();
std::vector<std::vector<size_t>> input_shapes;
std::vector<std::vector<size_t>> output_shapes;
auto input_shape = AnfAlgo::GetInputDeviceShapeAdaptively(kernel_node, 0);
auto output_shape = AnfAlgo::GetOutputDeviceShapeAdaptively(kernel_node, 0);
input_shapes.emplace_back(input_shape);
output_shapes.emplace_back(output_shape);
attr_.block_size = GetAttr<int64_t>(kernel_node, "block_size");
attr_.crops = GetAttr<std::vector<std::vector<int64_t>>>(kernel_node, "crops");
attr_.input_shape = input_shape;
int flag = helper_ptr_->CheckKernelParam(&attr_);
if (flag != 0) {
return false;
bool Init(const CNodePtr &kernel_node) override;
void ResetResource() noexcept override {
if (helper_ptr_) {
helper_ptr_->ResetResource();
}
flag = helper_ptr_->CalMemSize(input_shapes, output_shapes);
if (flag != 0) {
return false;
}
InitSizeLists();
return true;
}
protected:
@ -78,10 +55,11 @@ class BatchToSpaceGpuKernelMod : public NativeGpuKernelMod {
input_size_list_ = helper_ptr_->GetInputSizeList();
output_size_list_ = helper_ptr_->GetOutputSizeList();
}
std::vector<KernelAttr> GetOpSupport() override;
private:
std::string kernel_name_;
std::unique_ptr<cukernel::BatchToSpaceHelperGpuKernel<T>> helper_ptr_ = nullptr;
std::unique_ptr<cukernel::GpuKernelHelperBase> helper_ptr_ = nullptr;
cukernel::BatchToSpaceAttr attr_;
};
} // namespace kernel

View File

@ -15,22 +15,59 @@
*/
#include "plugin/device/gpu/kernel/arrays/unique_gpu_kernel.h"
#include <functional>
#include <utility>
#include <string>
#include <algorithm>
namespace mindspore {
namespace kernel {
MS_REG_GPU_KERNEL_TWO(
Unique,
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt32),
UniqueGpuKernelMod, float, int)
MS_REG_GPU_KERNEL_TWO(
Unique,
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeInt32),
UniqueGpuKernelMod, half, int)
MS_REG_GPU_KERNEL_TWO(
Unique, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
UniqueGpuKernelMod, int, int)
MS_REG_GPU_KERNEL_TWO(
Unique, KernelAttr().AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64),
UniqueGpuKernelMod, int64_t, int64_t)
namespace {
template <typename T, typename S>
std::unique_ptr<cukernel::GpuKernelHelperBase> CreateUniqueKernelPtr(const std::string &kernel_name) {
return std::make_unique<cukernel::UniqueHelperGpuKernel<T, S>>(kernel_name);
}
using UniquePtrCreatorFunc = std::function<std::unique_ptr<cukernel::GpuKernelHelperBase>(const std::string &)>;
const std::vector<std::pair<KernelAttr, UniquePtrCreatorFunc>> kernel_attr = {
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt32),
CreateUniqueKernelPtr<float, int>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeInt32),
CreateUniqueKernelPtr<half, int>},
{KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
CreateUniqueKernelPtr<int, int>},
{KernelAttr().AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64),
CreateUniqueKernelPtr<int64_t, int64_t>}};
} // namespace
bool UniqueGpuKernelMod::Init(const CNodePtr &kernel_node) {
kernel_node_ = kernel_node;
auto kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
auto index = GetMatchKernelAttrIdxWithException(kernel_node, GetOpSupport());
helper_ptr_ = std::move(kernel_attr[index].second(kernel_name));
helper_ptr_->ResetResource();
std::vector<std::vector<size_t>> input_shapes;
std::vector<std::vector<size_t>> output_shapes;
std::vector<size_t> shape = AnfAlgo::GetInputDeviceShapeAdaptively(kernel_node, 0);
is_null_input_ = CHECK_SHAPE_NULL(shape, kernel_name, "input");
if (is_null_input_) {
InitSizeLists();
return true;
}
input_shapes.emplace_back(shape);
helper_ptr_->CalMemSize(input_shapes, output_shapes);
InitSizeLists();
is_need_updateop_ = true;
return true;
}
std::vector<KernelAttr> UniqueGpuKernelMod::GetOpSupport() {
std::vector<KernelAttr> support_list;
(void)std::transform(kernel_attr.begin(), kernel_attr.end(), std::back_inserter(support_list),
[](const std::pair<KernelAttr, UniquePtrCreatorFunc> &item) { return item.first; });
return support_list;
}
MS_KERNEL_FACTORY_REG(NativeGpuKernelMod, Unique, UniqueGpuKernelMod);
} // namespace kernel
} // namespace mindspore

View File

@ -20,11 +20,10 @@
#include <vector>
#include <memory>
#include "plugin/device/gpu/kernel/gpu_kernel.h"
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
#include "plugin/factory/ms_factory.h"
#include "plugin/device/gpu/kernel/cuda_impl/cuda_class/unique_helper.h"
namespace mindspore {
namespace kernel {
template <typename T, typename S>
class UniqueGpuKernelMod : public NativeGpuKernelMod {
public:
UniqueGpuKernelMod() { ResetResource(); }
@ -45,25 +44,7 @@ class UniqueGpuKernelMod : public NativeGpuKernelMod {
return true;
}
bool Init(const CNodePtr &kernel_node) override {
kernel_node_ = kernel_node;
auto kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
helper_ptr_ = std::make_unique<cukernel::UniqueHelperGpuKernel<T, S>>(kernel_name);
helper_ptr_->ResetResource();
std::vector<std::vector<size_t>> input_shapes;
std::vector<std::vector<size_t>> output_shapes;
std::vector<size_t> shape = AnfAlgo::GetInputDeviceShapeAdaptively(kernel_node, 0);
is_null_input_ = CHECK_SHAPE_NULL(shape, kernel_name, "input");
if (is_null_input_) {
InitSizeLists();
return true;
}
input_shapes.emplace_back(shape);
helper_ptr_->CalMemSize(input_shapes, output_shapes);
InitSizeLists();
is_need_updateop_ = true;
return true;
}
bool Init(const CNodePtr &kernel_node) override;
void UpdateOp() override {
CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_, cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream_ptr_)),
@ -74,7 +55,10 @@ class UniqueGpuKernelMod : public NativeGpuKernelMod {
for (size_t i = 0; i < output_num; ++i) {
std::vector<size_t> shape = common::AnfAlgo::GetOutputInferShape(kernel_node_.lock(), i);
if (i == 0) {
shape[0] = helper_ptr_->GetOutSize();
auto dyn_out = helper_ptr_->GetDynOutInfo();
MS_EXCEPTION_IF_CHECK_FAIL(dyn_out.shapes.size() == 1 && dyn_out.shapes[0].size() == 1,
"Unique output info error.");
shape[0] = dyn_out.shapes[0][0];
}
TypeId type_id = common::AnfAlgo::GetOutputInferDataType(kernel_node_.lock(), i);
type_ids.emplace_back(type_id);
@ -98,10 +82,12 @@ class UniqueGpuKernelMod : public NativeGpuKernelMod {
workspace_size_list_ = helper_ptr_->GetWorkSizeList();
}
std::vector<KernelAttr> GetOpSupport() override;
private:
void *stream_ptr_;
bool is_null_input_;
std::unique_ptr<cukernel::UniqueHelperGpuKernel<T, S>> helper_ptr_ = nullptr;
std::unique_ptr<cukernel::GpuKernelHelperBase> helper_ptr_ = nullptr;
};
} // namespace kernel
} // namespace mindspore

View File

@ -38,7 +38,7 @@ struct BatchToSpaceAttr : public GpuKernelAttrBase {
template <typename T>
class BatchToSpaceHelperGpuKernel : public GpuKernelHelperBase {
public:
explicit BatchToSpaceHelperGpuKernel(std::string &kernel_name) : GpuKernelHelperBase(kernel_name) {}
explicit BatchToSpaceHelperGpuKernel(const std::string &kernel_name) : GpuKernelHelperBase(kernel_name) {}
virtual ~BatchToSpaceHelperGpuKernel() = default;
int CalMemSize(const std::vector<std::vector<size_t>> &input_shapes,
const std::vector<std::vector<size_t>> &output_shapes) override {

View File

@ -21,15 +21,23 @@
#include <vector>
#include "mindspore/core/utils/log_adapter.h"
#include "plugin/device/gpu/kernel/cuda_impl/cuda_class/cuda_class_common.h"
#include "ir/dtype/type_id.h"
#include "include/api/format.h"
namespace mindspore {
namespace cukernel {
struct GpuKernelAttrBase {
virtual ~GpuKernelAttrBase() = default;
};
struct DynamicOutInfo {
std::vector<std::vector<int>> shapes;
std::vector<std::vector<TypeId>> types;
std::vector<std::vector<Format>> formats;
};
class GpuKernelHelperBase {
public:
explicit GpuKernelHelperBase(std::string &kernel_name) : kernel_name_(kernel_name) {}
explicit GpuKernelHelperBase(const std::string &kernel_name) : kernel_name_(kernel_name) {}
virtual ~GpuKernelHelperBase() {
input_size_list_.clear();
output_size_list_.clear();
@ -52,6 +60,9 @@ class GpuKernelHelperBase {
virtual int CheckKernelParam(GpuKernelAttrBase *kernel_attr) { return 0; }
// Dynamic kernel can pass output information by this interface.
virtual DynamicOutInfo GetDynOutInfo() { return DynamicOutInfo(); }
protected:
std::vector<size_t> input_size_list_;
std::vector<size_t> output_size_list_;

View File

@ -72,7 +72,7 @@ static const std::map<std::string, UnaryOptype> kUnaryOpTypeMap = {
template <typename T>
class UnaryHelperGpuKernel : public GpuKernelHelperBase {
public:
explicit UnaryHelperGpuKernel(std::string &kernel_name) : GpuKernelHelperBase(kernel_name) {}
explicit UnaryHelperGpuKernel(const std::string &kernel_name) : GpuKernelHelperBase(kernel_name) {}
virtual ~UnaryHelperGpuKernel() = default;
int CalMemSize(const std::vector<std::vector<size_t>> &input_shapes,
const std::vector<std::vector<size_t>> &output_shapes) override {

View File

@ -33,7 +33,7 @@ constexpr size_t CROPS_SHAPE_1 = 2;
template <typename T, typename S>
class UniqueHelperGpuKernel : public GpuKernelHelperBase {
public:
explicit UniqueHelperGpuKernel(std::string &kernel_name) : GpuKernelHelperBase(kernel_name) {}
explicit UniqueHelperGpuKernel(const std::string &kernel_name) : GpuKernelHelperBase(kernel_name) {}
virtual ~UniqueHelperGpuKernel() = default;
int CalMemSize(const std::vector<std::vector<size_t>> &input_shapes,
const std::vector<std::vector<size_t>> &output_shapes) override {
@ -94,7 +94,11 @@ class UniqueHelperGpuKernel : public GpuKernelHelperBase {
work_size_list_.clear();
}
int GetOutSize() { return post_output_size_; }
DynamicOutInfo GetDynOutInfo() override {
DynamicOutInfo dyn_out;
dyn_out.shapes.push_back({{post_output_size_}});
return dyn_out;
}
private:
int num_elements_;

View File

@ -15,9 +15,34 @@
*/
#include "plugin/device/gpu/kernel/gpu_kernel.h"
#include <tuple>
#include <set>
namespace mindspore {
namespace kernel {
namespace {
void CheckDeviceSm(const KernelAttr &kernel_attr) {
const int major_sm = GET_MAJOR_SM;
if (!mindspore::device::gpu::CudaCommon::GetInstance().check_sm() || major_sm >= RECOMMEND_SM) {
return;
}
for (size_t i = 0; i < kernel_attr.GetInputSize(); ++i) {
if (kernel_attr.GetInputAttr(i).first != kNumberTypeFloat16) {
continue;
}
if (major_sm < MINIUM_SM) {
MS_LOG(EXCEPTION) << "Half precision ops can be used on Devices which computing capacity is >= " << MINIUM_SM
<< ", but the current device's computing capacity is " << major_sm;
}
MS_LOG(WARNING) << "It is recommended to use devices with a computing capacity >= " << RECOMMEND_SM
<< ", but the current device's computing capacity is " << major_sm;
mindspore::device::gpu::CudaCommon::GetInstance().set_check_sm(false);
return;
}
}
} // namespace
void NativeGpuKernelMod::InferOp() {
anf_node_ = kernel_node_.lock();
if (common::AnfAlgo::IsDynamicShape(kernel_node_.lock())) {
@ -46,5 +71,85 @@ void NativeGpuKernelMod::InitOp() {
ResetResource();
Init(cnode);
}
void NativeGpuKernelMod::SetGpuRefMapToKernelInfo(const CNodePtr &apply_kernel) {
MS_EXCEPTION_IF_NULL(apply_kernel);
auto kernel_attrs = GetOpSupport();
if (kernel_attrs.empty()) {
return;
}
auto index = GetMatchKernelAttrIdxWithException(apply_kernel, kernel_attrs);
auto kernel_info = dynamic_cast<device::KernelInfo *>(apply_kernel->kernel_info());
MS_EXCEPTION_IF_NULL(kernel_info);
const KernelBuildInfo *kernel_build_Info = kernel_info->select_kernel_build_info();
MS_EXCEPTION_IF_NULL(kernel_build_Info);
const auto &matched_kernel_attr = kernel_attrs[index];
if (!matched_kernel_attr.GetOutInRefMap().empty()) {
kernel_info->set_ref_map(matched_kernel_attr.GetOutInRefMap());
}
}
bool NativeGpuKernelMod::GpuCheckSupport(const std::string &kernel_name, const KernelAttr &kernel_attr) {
return kernel::Factory<NativeGpuKernelMod>::Instance().Create(kernel_name)->CheckSupport(kernel_name, kernel_attr);
}
std::vector<KernelAttr> NativeGpuKernelMod::GetAllSupportedList(const std::string &kernel_name) {
if (initialize_.count(kernel_name) == 0) {
auto kernel_support = GetOpSupport();
(void)support_map_.emplace(kernel_name, kernel_support);
(void)initialize_.insert(kernel_name);
}
return support_map_[kernel_name];
}
bool NativeGpuKernelMod::CheckSupport(const std::string &kernel_name, const KernelAttr &kernel_attr_to_check) {
CheckDeviceSm(kernel_attr_to_check);
auto kernel_attrs = GetAllSupportedList(kernel_name);
bool is_match;
std::tie(is_match, std::ignore) = MatchKernelAttr(kernel_attr_to_check, kernel_attrs);
return is_match;
}
NativeGpuKernelMod::ReducePrecisonRes NativeGpuKernelMod::ReducePrecisionCheck(const std::string &kernel_name,
const KernelAttr &kernel_attr_to_check) {
KernelAttr reduce_kernel_attr;
std::vector<ReduceDetail> input_reduce_index;
std::vector<ReduceDetail> output_reduce_index;
std::vector<KernelAttr> kernel_attr_list = this->GetOpSupport();
const TypeId from_precision = kNumberTypeInt64;
const TypeId to_precision = kNumberTypeInt32;
for (size_t attr_index = 0; attr_index < kernel_attr_list.size(); ++attr_index) {
auto &cur_kernel_attr = kernel_attr_list[attr_index];
auto attr_size = cur_kernel_attr.GetInputSize();
MS_EXCEPTION_IF_ZERO("kernel attr input size", attr_size);
for (size_t iidx = 0; iidx < kernel_attr_to_check.GetInputSize(); iidx++) {
auto [type_id, format] = kernel_attr_to_check.GetInputAttr(iidx);
if (type_id == from_precision && cur_kernel_attr.GetInputAttr(iidx % attr_size).first == to_precision) {
(void)input_reduce_index.emplace_back(iidx, from_precision, to_precision);
type_id = to_precision;
MS_LOG(WARNING) << "Kernel [" << kernel_name << "] does not support int64, cast input " << iidx << " to int32.";
reduce_kernel_attr.AddInputAttr(type_id, format);
}
}
for (size_t oidx = 0; oidx < kernel_attr_to_check.GetOutputSize(); oidx++) {
auto [type_id, format] = kernel_attr_to_check.GetOutputAttr(oidx);
if (type_id == from_precision && cur_kernel_attr.GetOutputAttr(oidx % attr_size).first == to_precision) {
(void)output_reduce_index.emplace_back(oidx, from_precision, to_precision);
type_id = to_precision;
MS_LOG(WARNING) << "Kernel [" << kernel_name << "] does not support int64, cast output " << oidx
<< " to int32.";
reduce_kernel_attr.AddOutputAttr(type_id, format);
}
}
}
return std::make_tuple(CheckSupport(kernel_name, reduce_kernel_attr), input_reduce_index, output_reduce_index);
}
std::map<std::string, std::vector<KernelAttr>> NativeGpuKernelMod::support_map_{};
std::set<std::string> NativeGpuKernelMod::initialize_{};
} // namespace kernel
} // namespace mindspore

View File

@ -14,8 +14,8 @@
* limitations under the License.
*/
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_GPU_KERNEL_H_
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_GPU_KERNEL_H_
#ifndef MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_GPU_KERNEL_H_
#define MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_GPU_KERNEL_H_
#include <cuda.h>
#include <cudnn.h>
@ -28,14 +28,21 @@
#include <numeric>
#include <functional>
#include <algorithm>
#include <tuple>
#include <set>
#include "kernel/kernel.h"
#include "plugin/device/gpu/kernel/gpu_kernel_mod.h"
#include "plugin/factory/ms_factory.h"
#include "plugin/device/gpu/kernel/kernel_constants.h"
#include "plugin/device/gpu/hal/device/gpu_device_manager.h"
#include "plugin/device/gpu/hal/device/gpu_common.h"
#include "backend/common/session/anf_runtime_algorithm.h"
#include "include/common/utils/anfalgo.h"
#include "runtime/device/executor/dynamic_kernel.h"
#include "kernel/kernel_build_info.h"
#include "kernel/common_utils.h"
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
using AnfAlgo = mindspore::session::AnfRuntimeAlgorithm;
// The max_limit of tensor shape size: 2 Giga-elements(2^31, the largest number in 32 bits).
@ -82,15 +89,47 @@ class NativeGpuKernelMod : public GpuKernelMod {
}
virtual void DestroyResource() noexcept {}
bool IsDynamicShape() { return common::AnfAlgo::IsDynamicShape(kernel_node_.lock()); }
static std::vector<KernelAttr> GetGpuSupportedList(const std::string &kernel_name) {
if (!Factory<NativeGpuKernelMod>::Instance().IsRegistered(kernel_name)) {
return {};
}
return Factory<NativeGpuKernelMod>::Instance().Create(kernel_name)->GetAllSupportedList(kernel_name);
}
static bool GpuCheckSupport(const std::string &kernel_name, const KernelAttr &kernel_attr);
using ReduceDetail = std::tuple<size_t, TypeId, TypeId>;
using ReducePrecisonRes = std::tuple<bool, std::vector<ReduceDetail>, std::vector<ReduceDetail>>;
static ReducePrecisonRes GpuReducePrecisionCheck(const std::string &kernel_name, const KernelAttr &kernel_attr) {
return Factory<NativeGpuKernelMod>::Instance().Create(kernel_name)->ReducePrecisionCheck(kernel_name, kernel_attr);
}
void SetGpuRefMapToKernelInfo(const CNodePtr &apply_kernel);
bool IsDynamicShape() { return common::AnfAlgo::IsDynamicShape(kernel_node_.lock()); }
void InferOp() override;
void InitOp() override;
protected:
virtual void InitResource() {}
virtual void InitSizeLists() = 0;
virtual std::vector<KernelAttr> GetOpSupport() { return {}; }
bool CheckSupport(const std::string &kernel_name, const KernelAttr &kernel_attr);
std::vector<KernelAttr> GetAllSupportedList(const std::string &kernel_name);
ReducePrecisonRes ReducePrecisionCheck(const std::string &kernel_name, const KernelAttr &kernel_attr);
std::weak_ptr<CNode> kernel_node_;
static std::map<std::string, std::vector<KernelAttr>> support_map_;
static std::set<std::string> initialize_;
size_t GetMatchKernelAttrIdxWithException(const AnfNodePtr &node, const std::vector<KernelAttr> &kernel_attrs) {
auto kernel_attr = GetKernelAttrFromNode(node);
auto [is_match, index] = MatchKernelAttr(kernel_attr, kernel_attrs);
if (!is_match) {
MS_LOG(EXCEPTION) << common::AnfAlgo::GetCNodeName(node)
<< " does not support this kernel data type: " << kernel_attr;
}
return index;
}
inline void ResetSizeLists() {
input_size_list_.clear();
@ -401,7 +440,15 @@ class NativeGpuKernelMod : public GpuKernelMod {
return true;
}
};
// This is necessary for gpu kernels to support uint8 data type. In cuda, an unsigned,
// 8 bit integral type is represented by an unsigned char, but the MS_REG_GPU_KERNEL
// macros defined below will create compilation errors when datatype T contains a space,
// because the variable created by the macro will also contain a space. So, we solve this
// problem by writing uchar when calling these macros, and expanding uchar after the
// variable has been created.
using uchar = unsigned char;
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_GPU_KERNEL_H_
#endif // MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_GPU_KERNEL_H_

View File

@ -77,14 +77,6 @@ class GpuKernelRegister {
~GpuKernelRegister() = default;
};
// This is necessary for gpu kernels to support uint8 data type. In cuda, an unsigned,
// 8 bit integral type is represented by an unsigned char, but the MS_REG_GPU_KERNEL
// macros defined below will create compilation errors when datatype T contains a space,
// because the variable created by the macro will also contain a space. So, we solve this
// problem by writing uchar when calling these macros, and expanding uchar after the
// variable has been created.
using uchar = unsigned char;
#define UNIQUE_KERNEL_NAME(kernel) KERNEL_NAME(g_##kernel##_gpu_kernel_reg, __COUNTER__)
#define KERNEL_NAME(kernel, cnt) MERGE(kernel, cnt)
#define MERGE(kernel, cnt) kernel##cnt

View File

@ -15,128 +15,206 @@
*/
#include "plugin/device/gpu/kernel/math/unary_op_gpu_kernel.h"
#include <map>
#include <functional>
#include <utility>
#include <algorithm>
namespace mindspore {
namespace kernel {
MS_REG_GPU_KERNEL_ONE(Exp, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Exp, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Expm1, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Expm1, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Log, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Log, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Log1p, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Log1p, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Erf, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Erf, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Erfc, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Erfc, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Neg, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Neg, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Neg, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Neg, KernelAttr().AddInputAttr(kNumberTypeInt8).AddOutputAttr(kNumberTypeInt8),
UnaryOpGpuKernelMod, char)
MS_REG_GPU_KERNEL_ONE(Neg, KernelAttr().AddInputAttr(kNumberTypeUInt8).AddOutputAttr(kNumberTypeUInt8),
UnaryOpGpuKernelMod, uchar)
MS_REG_GPU_KERNEL_ONE(Neg, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
UnaryOpGpuKernelMod, int)
MS_REG_GPU_KERNEL_ONE(Reciprocal, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Reciprocal, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Square, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Square, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Sqrt, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Sqrt, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Sqrt, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Rsqrt, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Rsqrt, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Sin, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Sin, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Sin, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Asin, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Asin, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Asinh, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Asinh, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Cos, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Cos, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Cos, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(ACos, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(ACos, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Acosh, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Acosh, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Atan, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Atan, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Abs, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Abs, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Abs, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Abs, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
UnaryOpGpuKernelMod, int)
MS_REG_GPU_KERNEL_ONE(Floor, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Floor, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Rint, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Rint, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Rint, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Round, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
UnaryOpGpuKernelMod, int)
MS_REG_GPU_KERNEL_ONE(Round, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Round, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Round, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
MS_REG_GPU_KERNEL_ONE(Sign, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
UnaryOpGpuKernelMod, int)
MS_REG_GPU_KERNEL_ONE(Sign, KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64),
UnaryOpGpuKernelMod, double)
MS_REG_GPU_KERNEL_ONE(Sign, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
UnaryOpGpuKernelMod, float)
MS_REG_GPU_KERNEL_ONE(Sign, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
UnaryOpGpuKernelMod, half)
namespace {
constexpr auto kAbs = "Abs";
constexpr auto kACos = "ACos";
constexpr auto kAcosh = "Acosh";
constexpr auto kAsin = "Asin";
constexpr auto kAsinh = "Asinh";
constexpr auto kAtan = "Atan";
constexpr auto kCos = "Cos";
constexpr auto kErf = "Erf";
constexpr auto kErfc = "Erfc";
constexpr auto kExp = "Exp";
constexpr auto kExpm1 = "Expm1";
constexpr auto kFloor = "Floor";
constexpr auto kLog = "Log";
constexpr auto kLog1p = "Log1p";
constexpr auto kNeg = "Neg";
constexpr auto kReciprocal = "Reciprocal";
constexpr auto kRint = "Rint";
constexpr auto kRound = "Round";
constexpr auto kRsqrt = "Rsqrt";
constexpr auto kSign = "Sign";
constexpr auto kSin = "Sin";
constexpr auto kSqrt = "Sqrt";
constexpr auto kSquare = "Square";
template <typename T>
std::unique_ptr<cukernel::GpuKernelHelperBase> CreateUnaryKernelPtr(const std::string &kernel_name) {
return std::make_unique<cukernel::UnaryHelperGpuKernel<T>>(kernel_name);
}
using UnaryPtrCreatorFunc = std::function<std::unique_ptr<cukernel::GpuKernelHelperBase>(const std::string &)>;
const std::map<std::string, std::vector<std::pair<KernelAttr, UnaryPtrCreatorFunc>>> kernel_attr_map = {
{kExp,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kExpm1,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kLog,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kLog1p,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kErf,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kErfc,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kNeg,
{{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>},
{KernelAttr().AddInputAttr(kNumberTypeInt8).AddOutputAttr(kNumberTypeInt8), CreateUnaryKernelPtr<char>},
{KernelAttr().AddInputAttr(kNumberTypeUInt8).AddOutputAttr(kNumberTypeUInt8), CreateUnaryKernelPtr<uchar>},
{KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32), CreateUnaryKernelPtr<int>}}},
{kReciprocal,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kSquare,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kSqrt,
{{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kRsqrt,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kSin,
{{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kAsin,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kAsinh,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kCos,
{{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kACos,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kAcosh,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kAtan,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kAbs,
{{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>},
{KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32), CreateUnaryKernelPtr<int>}}},
{kFloor,
{{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kRint,
{{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kRound,
{{KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32), CreateUnaryKernelPtr<int>},
{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}},
{kSign,
{{KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32), CreateUnaryKernelPtr<int>},
{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), CreateUnaryKernelPtr<double>},
{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), CreateUnaryKernelPtr<float>},
{KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16), CreateUnaryKernelPtr<half>}}}};
} // namespace
bool UnaryOpGpuKernelMod::Init(const CNodePtr &kernel_node) {
kernel_node_ = kernel_node;
std::string kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
auto index = GetMatchKernelAttrIdxWithException(kernel_node, GetOpSupport());
helper_ptr_ = std::move(kernel_attr_map.at(kernel_type_)[index].second(kernel_name));
helper_ptr_->ResetResource();
std::vector<std::vector<size_t>> input_shapes;
std::vector<std::vector<size_t>> output_shapes;
auto input_shape = AnfAlgo::GetInputDeviceShapeAdaptively(kernel_node, 0);
auto output_shape = AnfAlgo::GetOutputDeviceShapeAdaptively(kernel_node, 0);
is_null_input_ = CHECK_SHAPE_NULL(input_shape, kernel_name, "input");
if (is_null_input_) {
input_size_list_.emplace_back(0);
output_size_list_.emplace_back(0);
return true;
}
input_shapes.emplace_back(input_shape);
output_shapes.emplace_back(output_shape);
int flag = helper_ptr_->CalMemSize(input_shapes, output_shapes);
if (flag != 0) {
return false;
}
InitSizeLists();
return true;
}
std::vector<KernelAttr> UnaryOpGpuKernelMod::GetOpSupport() {
auto iter = kernel_attr_map.find(kernel_type_);
if (iter == kernel_attr_map.end()) {
MS_LOG(EXCEPTION) << "Unary gpu do not support " << kernel_type_;
}
std::vector<KernelAttr> support_list;
(void)std::transform(iter->second.begin(), iter->second.end(), std::back_inserter(support_list),
[](const std::pair<KernelAttr, UnaryPtrCreatorFunc> &item) { return item.first; });
return support_list;
}
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Abs, []() { return std::make_shared<UnaryOpGpuKernelMod>(kAbs); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, ACos,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kACos); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Acosh,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kAcosh); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Asin,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kAsin); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Asinh,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kAsinh); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Atan,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kAtan); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Cos, []() { return std::make_shared<UnaryOpGpuKernelMod>(kCos); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Erf, []() { return std::make_shared<UnaryOpGpuKernelMod>(kErf); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Erfc,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kErfc); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Exp, []() { return std::make_shared<UnaryOpGpuKernelMod>(kExp); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Expm1,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kExpm1); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Floor,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kFloor); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Log, []() { return std::make_shared<UnaryOpGpuKernelMod>(kLog); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Log1p,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kLog1p); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Neg, []() { return std::make_shared<UnaryOpGpuKernelMod>(kNeg); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Reciprocal,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kReciprocal); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Rint,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kRint); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Round,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kRound); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Rsqrt,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kRsqrt); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Sign,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kSign); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Sin, []() { return std::make_shared<UnaryOpGpuKernelMod>(kSin); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Sqrt,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kSqrt); });
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeGpuKernelMod, Square,
[]() { return std::make_shared<UnaryOpGpuKernelMod>(kSquare); });
} // namespace kernel
} // namespace mindspore

View File

@ -23,15 +23,14 @@
#include <string>
#include <memory>
#include "plugin/device/gpu/kernel/gpu_kernel.h"
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
#include "plugin/factory/ms_factory.h"
#include "plugin/device/gpu/kernel/cuda_impl/cuda_class/unary_helper.h"
namespace mindspore {
namespace kernel {
template <typename T>
class UnaryOpGpuKernelMod : public NativeGpuKernelMod {
public:
UnaryOpGpuKernelMod() { ResetResource(); }
explicit UnaryOpGpuKernelMod(const std::string &kernel_type) : kernel_type_(kernel_type) { ResetResource(); }
~UnaryOpGpuKernelMod() override = default;
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &,
@ -39,17 +38,9 @@ class UnaryOpGpuKernelMod : public NativeGpuKernelMod {
if (is_null_input_) {
return true;
}
std::vector<void *> input_addrs;
std::vector<void *> output_addrs;
std::vector<void *> input_addrs = ConvertPtrs(inputs);
std::vector<void *> output_addrs = ConvertPtrs(outputs);
std::vector<void *> work_addrs;
for (size_t idx = 0; idx < inputs.size(); ++idx) {
void *cur_ptr = reinterpret_cast<void *>(GetDeviceAddress<T>(inputs, idx));
input_addrs.emplace_back(cur_ptr);
}
for (size_t idx = 0; idx < outputs.size(); ++idx) {
void *cur_ptr = reinterpret_cast<void *>(GetDeviceAddress<T>(outputs, idx));
output_addrs.emplace_back(cur_ptr);
}
int flag = helper_ptr_->Process(input_addrs, output_addrs, work_addrs, stream_ptr);
if (flag != 0) {
return false;
@ -57,30 +48,7 @@ class UnaryOpGpuKernelMod : public NativeGpuKernelMod {
return true;
}
bool Init(const CNodePtr &kernel_node) override {
kernel_node_ = kernel_node;
std::string kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
helper_ptr_ = std::make_unique<cukernel::UnaryHelperGpuKernel<T>>(kernel_name);
helper_ptr_->ResetResource();
std::vector<std::vector<size_t>> input_shapes;
std::vector<std::vector<size_t>> output_shapes;
auto input_shape = AnfAlgo::GetInputDeviceShapeAdaptively(kernel_node, 0);
auto output_shape = AnfAlgo::GetOutputDeviceShapeAdaptively(kernel_node, 0);
is_null_input_ = CHECK_SHAPE_NULL(input_shape, kernel_name, "input");
if (is_null_input_) {
input_size_list_.emplace_back(0);
output_size_list_.emplace_back(0);
return true;
}
input_shapes.emplace_back(input_shape);
output_shapes.emplace_back(output_shape);
int flag = helper_ptr_->CalMemSize(input_shapes, output_shapes);
if (flag != 0) {
return false;
}
InitSizeLists();
return true;
}
bool Init(const CNodePtr &kernel_node) override;
void ResetResource() noexcept override {
input_size_list_.clear();
@ -94,10 +62,12 @@ class UnaryOpGpuKernelMod : public NativeGpuKernelMod {
output_size_list_ = helper_ptr_->GetOutputSizeList();
workspace_size_list_ = helper_ptr_->GetWorkSizeList();
}
std::vector<KernelAttr> GetOpSupport() override;
private:
std::unique_ptr<cukernel::UnaryHelperGpuKernel<T>> helper_ptr_ = nullptr;
std::unique_ptr<cukernel::GpuKernelHelperBase> helper_ptr_ = nullptr;
bool is_null_input_;
std::string kernel_type_{"Unknown"};
};
} // namespace kernel
} // namespace mindspore