!6728 [Ascend][DynamicShape] Dynamic shape feature
Merge pull request !6728 from caifubi/dynamic_shape_share_2pull/6728/MERGE
mindspore-ci-bot
4 years ago
committed by
Gitee
commit
c951d42c2c
@ -0,0 +1,52 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
#include "backend/kernel_compiler/host/dynamic_shape_kernel.h"
|
||||||
|
#include "backend/session/anf_runtime_algorithm.h"
|
||||||
|
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
void DynamicShapeKernel::Execute() {
|
||||||
|
MS_LOG(INFO) << "Execute DynamicShapeKernel Start";
|
||||||
|
auto input_num = AnfAlgo::GetInputTensorNum(cnode_ptr_);
|
||||||
|
if (input_num != 1) {
|
||||||
|
MS_LOG(EXCEPTION) << "Invalid Input Num:" << input_num;
|
||||||
|
}
|
||||||
|
|
||||||
|
auto prev_output_shape = AnfAlgo::GetPrevNodeOutputInferShape(cnode_ptr_, 0);
|
||||||
|
auto output_shape = std::vector<int>(SizeToInt(prev_output_shape.size()));
|
||||||
|
|
||||||
|
auto output_type = TypeId::kNumberTypeInt32;
|
||||||
|
|
||||||
|
auto output_tensor_for_sync = std::make_shared<tensor::Tensor>(output_type, output_shape);
|
||||||
|
auto data_ptr = static_cast<int32_t *>(output_tensor_for_sync->data_c());
|
||||||
|
for (size_t i = 0; i < prev_output_shape.size(); ++i) {
|
||||||
|
MS_LOG(INFO) << "DEBUG prev_output_shape[" << i << "]:" << prev_output_shape[i];
|
||||||
|
*(data_ptr + i) = prev_output_shape[i];
|
||||||
|
}
|
||||||
|
|
||||||
|
auto output_addr = AnfAlgo::GetOutputAddr(cnode_ptr_, 0);
|
||||||
|
MS_EXCEPTION_IF_NULL(output_addr);
|
||||||
|
output_addr->SyncHostToDevice(output_shape, LongToSize(output_tensor_for_sync->data().nbytes()),
|
||||||
|
output_tensor_for_sync->data_type(), output_tensor_for_sync->data_c());
|
||||||
|
MS_LOG(INFO) << "Execute DynamicShapeKernel End";
|
||||||
|
}
|
||||||
|
|
||||||
|
device::DynamicKernelPtr DynamicShapeKernelMod::GenDynamicKernel(const CNodePtr &cnode_ptr, void *stream_ptr) {
|
||||||
|
return std::make_shared<DynamicShapeKernel>(stream_ptr, cnode_ptr);
|
||||||
|
}
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
@ -0,0 +1,43 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_DYNAMIC_SHAPE_KERNEL_H_
|
||||||
|
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_DYNAMIC_SHAPE_KERNEL_H_
|
||||||
|
#include <vector>
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include "runtime/device/ascend/executor/host_dynamic_kernel.h"
|
||||||
|
#include "backend/kernel_compiler/host/host_kernel_mod.h"
|
||||||
|
using HostDynamicKernel = mindspore::device::ascend::HostDynamicKernel;
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
class DynamicShapeKernel : public HostDynamicKernel {
|
||||||
|
public:
|
||||||
|
DynamicShapeKernel(void *stream, const CNodePtr &cnode_ptr) : HostDynamicKernel(stream, cnode_ptr) {}
|
||||||
|
~DynamicShapeKernel() override = default;
|
||||||
|
void Execute() override;
|
||||||
|
};
|
||||||
|
|
||||||
|
class DynamicShapeKernelMod : public HostKernelMod {
|
||||||
|
public:
|
||||||
|
DynamicShapeKernelMod() = default;
|
||||||
|
~DynamicShapeKernelMod() override = default;
|
||||||
|
device::DynamicKernelPtr GenDynamicKernel(const CNodePtr &cnode_ptr, void *stream_ptr) override;
|
||||||
|
};
|
||||||
|
MS_HOST_REG_KERNEL(DynamicShape, DynamicShapeKernelMod);
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
|
|
||||||
|
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_DYNAMIC_SHAPE_KERNEL_H_
|
||||||
@ -0,0 +1,42 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
#include "backend/kernel_compiler/host/host_kernel_build.h"
|
||||||
|
#include <string>
|
||||||
|
#include "runtime/device/kernel_runtime.h"
|
||||||
|
#include "backend/kernel_compiler/host/host_kernel_mod.h"
|
||||||
|
#include "backend/session/anf_runtime_algorithm.h"
|
||||||
|
#include "backend/session/kernel_graph.h"
|
||||||
|
#include "backend/kernel_compiler/common_utils.h"
|
||||||
|
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
KernelModPtr HostOpBuild(const std::shared_ptr<AnfNode> &anf_node) {
|
||||||
|
MS_EXCEPTION_IF_NULL(anf_node);
|
||||||
|
std::string opname = AnfAlgo::GetCNodeName(anf_node);
|
||||||
|
MS_LOG(INFO) << "Host op [" << opname << "]";
|
||||||
|
auto kerPtr = HostKernelFactory::Get(opname);
|
||||||
|
if (kerPtr == nullptr) {
|
||||||
|
MS_LOG(ERROR) << "Host can't find Kernel[" << opname << "]";
|
||||||
|
return nullptr;
|
||||||
|
}
|
||||||
|
if (!kerPtr->Init(anf_node)) {
|
||||||
|
MS_LOG(ERROR) << "Host Kernel initialize failed!";
|
||||||
|
return nullptr;
|
||||||
|
}
|
||||||
|
return kerPtr;
|
||||||
|
}
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
@ -0,0 +1,27 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_BUILD_H_
|
||||||
|
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_BUILD_H_
|
||||||
|
#include <memory>
|
||||||
|
#include "backend/kernel_compiler/kernel.h"
|
||||||
|
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
KernelModPtr HostOpBuild(const std::shared_ptr<AnfNode> &anf_node);
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
|
|
||||||
|
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_BUILD_H_
|
||||||
@ -0,0 +1,59 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
#include "backend/kernel_compiler/host/host_kernel_metadata.h"
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include "backend/kernel_compiler/oplib/oplib.h"
|
||||||
|
#include "backend/kernel_compiler/common_utils.h"
|
||||||
|
#include "backend/session/anf_runtime_algorithm.h"
|
||||||
|
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
constexpr auto kDynamicShape = "DynamicShape";
|
||||||
|
|
||||||
|
void HostMetadataInfo(const CNodePtr &kernel_node, std::vector<std::shared_ptr<KernelBuildInfo>> *kernel_info_list) {
|
||||||
|
MS_LOG(INFO) << "HostMetadataInfo.";
|
||||||
|
MS_EXCEPTION_IF_NULL(kernel_node);
|
||||||
|
MS_EXCEPTION_IF_NULL(kernel_info_list);
|
||||||
|
std::string op_name = AnfAlgo::GetCNodeName(kernel_node);
|
||||||
|
if (op_name != kDynamicShape) {
|
||||||
|
MS_LOG(DEBUG) << "Host does not have op [" << op_name << "]";
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<std::string> inputs_format{};
|
||||||
|
std::vector<TypeId> inputs_type{};
|
||||||
|
for (size_t input_index = 0; input_index < AnfAlgo::GetInputTensorNum(kernel_node); ++input_index) {
|
||||||
|
inputs_format.emplace_back(kOpFormat_DEFAULT);
|
||||||
|
inputs_type.push_back(AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, input_index));
|
||||||
|
}
|
||||||
|
std::vector<std::string> outputs_format;
|
||||||
|
std::vector<TypeId> outputs_type;
|
||||||
|
for (size_t output_index = 0; output_index < AnfAlgo::GetOutputTensorNum(kernel_node); ++output_index) {
|
||||||
|
outputs_format.emplace_back(kOpFormat_DEFAULT);
|
||||||
|
outputs_type.push_back(AnfAlgo::GetOutputInferDataType(kernel_node, output_index));
|
||||||
|
}
|
||||||
|
auto builder = KernelBuildInfo::KernelBuildInfoBuilder();
|
||||||
|
builder.SetInputsFormat(inputs_format);
|
||||||
|
builder.SetInputsDeviceType(inputs_type);
|
||||||
|
builder.SetOutputsFormat(outputs_format);
|
||||||
|
builder.SetOutputsDeviceType(outputs_type);
|
||||||
|
builder.SetKernelType(HOST_KERNEL);
|
||||||
|
kernel_info_list->push_back(builder.Build());
|
||||||
|
}
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
@ -0,0 +1,30 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_META_DATA_H_
|
||||||
|
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_META_DATA_H_
|
||||||
|
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
#include <memory>
|
||||||
|
#include "backend/kernel_compiler/kernel_build_info.h"
|
||||||
|
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
void HostMetadataInfo(const CNodePtr &kernel_node, std::vector<std::shared_ptr<KernelBuildInfo>> *kernel_info_list);
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
|
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_META_DATA_H_
|
||||||
@ -0,0 +1,98 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
#include "backend/kernel_compiler/host/host_kernel_mod.h"
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <vector>
|
||||||
|
#include <string>
|
||||||
|
#include <utility>
|
||||||
|
#include "runtime/mem.h"
|
||||||
|
#include "utils/ms_context.h"
|
||||||
|
#include "runtime/device/kernel_runtime.h"
|
||||||
|
#include "runtime/device/ascend/executor/host_dynamic_kernel.h"
|
||||||
|
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
void HostKernelFactory::Registe(const std::string &name, HostKernelCreater &&fun) {
|
||||||
|
hostKernelMap_.emplace(name, std::move(fun));
|
||||||
|
}
|
||||||
|
|
||||||
|
std::shared_ptr<HostKernelMod> HostKernelFactory::Get(const std::string &name) {
|
||||||
|
const auto &map = Get().hostKernelMap_;
|
||||||
|
auto it = map.find(name);
|
||||||
|
if (it != map.end() && it->second) {
|
||||||
|
return (it->second)();
|
||||||
|
}
|
||||||
|
return nullptr;
|
||||||
|
}
|
||||||
|
|
||||||
|
HostKernelFactory &HostKernelFactory::Get() {
|
||||||
|
static HostKernelFactory instance;
|
||||||
|
return instance;
|
||||||
|
}
|
||||||
|
|
||||||
|
const std::vector<size_t> &HostKernelMod::GetInputSizeList() const { return input_size_list_; }
|
||||||
|
const std::vector<size_t> &HostKernelMod::GetOutputSizeList() const { return output_size_list_; }
|
||||||
|
const std::vector<size_t> &HostKernelMod::GetWorkspaceSizeList() const { return workspace_size_list_; }
|
||||||
|
bool HostKernelMod::Init(const AnfNodePtr &anf_node) {
|
||||||
|
MS_EXCEPTION_IF_NULL(anf_node);
|
||||||
|
size_t input_num = AnfAlgo::GetInputTensorNum(anf_node);
|
||||||
|
size_t output_num = AnfAlgo::GetOutputTensorNum(anf_node);
|
||||||
|
|
||||||
|
for (size_t i = 0; i < input_num; i++) {
|
||||||
|
std::vector<size_t> shape_i = AnfAlgo::GetInputDeviceShape(anf_node, i);
|
||||||
|
TypePtr type_ptr = TypeIdToType(AnfAlgo::GetInputDeviceDataType(anf_node, i));
|
||||||
|
MS_EXCEPTION_IF_NULL(type_ptr);
|
||||||
|
int64_t size_i = 1;
|
||||||
|
for (size_t j = 0; j < shape_i.size(); j++) {
|
||||||
|
size_i = LongMulWithOverflowCheck(size_i, static_cast<int>(shape_i[j]));
|
||||||
|
}
|
||||||
|
size_t type_byte = GetTypeByte(type_ptr);
|
||||||
|
if (type_byte == 0) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
size_i = LongMulWithOverflowCheck(size_i, SizeToInt(type_byte));
|
||||||
|
input_size_list_.push_back(LongToSize(size_i));
|
||||||
|
}
|
||||||
|
|
||||||
|
for (size_t i = 0; i < output_num; i++) {
|
||||||
|
std::vector<size_t> shape_i = AnfAlgo::GetOutputDeviceShape(anf_node, i);
|
||||||
|
TypePtr type_ptr = TypeIdToType(AnfAlgo::GetOutputDeviceDataType(anf_node, i));
|
||||||
|
MS_EXCEPTION_IF_NULL(type_ptr);
|
||||||
|
int64_t size_i = 1;
|
||||||
|
for (size_t j = 0; j < shape_i.size(); j++) {
|
||||||
|
size_i = LongMulWithOverflowCheck(size_i, static_cast<int>(shape_i[j]));
|
||||||
|
}
|
||||||
|
size_t type_byte = GetTypeByte(type_ptr);
|
||||||
|
if (type_byte == 0) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
size_i = LongMulWithOverflowCheck(size_i, SizeToInt(type_byte));
|
||||||
|
output_size_list_.push_back(LongToSize(size_i));
|
||||||
|
}
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
bool HostKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||||
|
const std::vector<AddressPtr> &outputs, void *stream_ptr) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
std::vector<TaskInfoPtr> HostKernelMod::GenTask(const std::vector<AddressPtr> &, const std::vector<AddressPtr> &,
|
||||||
|
const std::vector<AddressPtr> &, uint32_t) {
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
@ -0,0 +1,86 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_MOD_H_
|
||||||
|
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_MOD_H_
|
||||||
|
#include <vector>
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include <map>
|
||||||
|
#include <utility>
|
||||||
|
#include "backend/kernel_compiler/ascend_kernel_mod.h"
|
||||||
|
namespace mindspore {
|
||||||
|
namespace kernel {
|
||||||
|
class HostKernelMod : public AscendKernelMod {
|
||||||
|
public:
|
||||||
|
HostKernelMod() = default;
|
||||||
|
~HostKernelMod() override = default;
|
||||||
|
const std::vector<size_t> &GetInputSizeList() const override;
|
||||||
|
const std::vector<size_t> &GetOutputSizeList() const override;
|
||||||
|
const std::vector<size_t> &GetWorkspaceSizeList() const override;
|
||||||
|
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||||
|
const std::vector<AddressPtr> &outputs, void *stream_ptr) override;
|
||||||
|
std::vector<TaskInfoPtr> GenTask(const std::vector<AddressPtr> &, const std::vector<AddressPtr> &,
|
||||||
|
const std::vector<AddressPtr> &, uint32_t) override;
|
||||||
|
device::DynamicKernelPtr GenDynamicKernel(const CNodePtr &cnode_ptr, void *stream_ptr) override = 0;
|
||||||
|
bool Init(const AnfNodePtr &anf_node);
|
||||||
|
|
||||||
|
protected:
|
||||||
|
AnfNodePtr anf_node_;
|
||||||
|
std::string op_name_;
|
||||||
|
std::vector<size_t> input_size_list_;
|
||||||
|
std::vector<size_t> output_size_list_;
|
||||||
|
std::vector<size_t> workspace_size_list_;
|
||||||
|
};
|
||||||
|
|
||||||
|
using HostKernelModPtr = std::shared_ptr<HostKernelMod>;
|
||||||
|
using HostKernelModPtrList = std::vector<HostKernelModPtr>;
|
||||||
|
using HostKernelCreater = std::function<std::shared_ptr<HostKernelMod>()>;
|
||||||
|
|
||||||
|
class HostKernelFactory {
|
||||||
|
HostKernelFactory() = default;
|
||||||
|
~HostKernelFactory() = default;
|
||||||
|
|
||||||
|
public:
|
||||||
|
static HostKernelFactory &Get();
|
||||||
|
void Registe(const string &name, HostKernelCreater &&fun);
|
||||||
|
static std::shared_ptr<HostKernelMod> Get(const string &name);
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::map<string, HostKernelCreater> hostKernelMap_;
|
||||||
|
};
|
||||||
|
|
||||||
|
class _HostKernelRegister {
|
||||||
|
public:
|
||||||
|
_HostKernelRegister(const string &name, HostKernelCreater &&fun) {
|
||||||
|
HostKernelFactory::Get().Registe(name, std::move(fun));
|
||||||
|
}
|
||||||
|
~_HostKernelRegister() = default;
|
||||||
|
};
|
||||||
|
|
||||||
|
#define _MS_HOST_REG_KERNEL_REG(KNAME, clazz) \
|
||||||
|
static_assert(std::is_base_of<HostKernelMod, clazz>::value, " must be base of HostKernelMod"); \
|
||||||
|
static const _HostKernelRegister g_##KNAME##_##_kernel_reg(#KNAME, []() { \
|
||||||
|
std::shared_ptr<clazz> ptr = nullptr; \
|
||||||
|
ptr = std::make_shared<clazz>(); \
|
||||||
|
MS_EXCEPTION_IF_NULL(ptr); \
|
||||||
|
return ptr; \
|
||||||
|
});
|
||||||
|
|
||||||
|
#define MS_HOST_REG_KERNEL(KNAME, clazz) _MS_HOST_REG_KERNEL_REG(KNAME, clazz)
|
||||||
|
} // namespace kernel
|
||||||
|
} // namespace mindspore
|
||||||
|
|
||||||
|
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_HOST_HOST_KERNEL_MOD_H_
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue