p85971406
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!39544 fix unstack with big shape 

Merge pull request !39544 from 范吉斌/fix_unstack
This commit is contained in:
i-robot 2022-08-05 07:37:30 +00:00 committed by Gitee
parent 68af199203 3371aaea97
commit 7b5345bf95
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
2 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mindspore/ccsrc/plugin/device/cpu/kernel/unpack_cpu_kernel.cc
View File

@ -24,6 +24,7 @@ namespace {
constexpr size_t kUnpackInputsNum = 1;
constexpr size_t kUnpackOutputsMinNum = 1;
constexpr size_t kUnpackWorkspaceMinNum = 1;
constexpr size_t kMaxDataSize = 2147483648; // 2GB
} // namespace
void UnpackCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
@ -40,6 +41,7 @@ void UnpackCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
unstack_param_.pre_dims_ = 1;
unstack_param_.axis_dim_ = 1;
unstack_param_.after_dims_ = 1;
input_size_ = 1;
for (size_t i = 0; i < input_shape.size(); i++) {
if (i < IntToSize(unstack_param_.axis_)) {
@ -49,6 +51,7 @@ void UnpackCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
} else {
unstack_param_.axis_dim_ = LongToInt(input_shape[i]);
}
input_size_ *= LongToSize(input_shape[i]);
}
auto kernel_attr = GetKernelAttrFromNode(kernel_node);
@ -88,6 +91,12 @@ bool UnpackCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &inputs,
for (size_t i = 0; i < outputs.size(); i++) {
outputs_host[i] = reinterpret_cast<T *>(outputs[i]->addr);
}
size_t total_size = input_size_ * sizeof(T);
if (total_size >= kMaxDataSize) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << ", the input data size cannot larger than 2GB, but got "
<< total_size << " bytes";
}
int data_size = SizeToInt(sizeof(T));
Unstack(input, outputs_host, &unstack_param_, data_size);
return true;

1
mindspore/ccsrc/plugin/device/cpu/kernel/unpack_cpu_kernel.h
View File

@ -63,6 +63,7 @@ class UnpackCpuKernelMod : public DeprecatedNativeCpuKernelMod {
UnstackParameter unstack_param_{};
size_t output_num_{0};
size_t input_size_{1};
};
} // namespace kernel
} // namespace mindspore