parent
025bb7b125
commit
a52cd685fc
@ -1,28 +0,0 @@
|
||||
/**
|
||||
* 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 <cstdint>
|
||||
|
||||
#include "backend/kernel_compiler/gpu/arrays/repeat_elements_gpu_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(RepeatElements, KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
|
||||
RepeatElementsGpuKernel, half)
|
||||
|
||||
MS_REG_GPU_KERNEL_ONE(RepeatElements, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
|
||||
RepeatElementsGpuKernel, int32_t)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
@ -1,161 +0,0 @@
|
||||
/**
|
||||
* 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_GPU_REPEAT_ELEMENTS_GPU_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_REPEAT_ELEMENTS_GPU_KERNEL_H_
|
||||
|
||||
#include "backend/kernel_compiler/gpu/cuda_impl/repeat_elements_impl.cuh"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <vector>
|
||||
|
||||
#include "backend/kernel_compiler/gpu/gpu_kernel.h"
|
||||
#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
class RepeatElementsGpuKernel : public GpuKernel {
|
||||
public:
|
||||
RepeatElementsGpuKernel() : rep_(1), axis_(0), input_size_(1), output_size_(0) {}
|
||||
~RepeatElementsGpuKernel() = default;
|
||||
|
||||
const std::vector<size_t> &GetInputSizeList() const override { return input_size_list_; }
|
||||
const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; }
|
||||
const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; }
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
T *input_device_address = GetDeviceAddress<T>(inputs, 0);
|
||||
T *output_device_address = GetDeviceAddress<T>(outputs, 0);
|
||||
|
||||
switch (input_dim_) {
|
||||
case 1:
|
||||
CalRepeatElements1d(input_device_address, rep_, axis_, output_device_address, output_size_,
|
||||
reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
break;
|
||||
case 2:
|
||||
CalRepeatElements2d(input_device_address, input_shape_[1], rep_, axis_, output_device_address, output_shape_[1],
|
||||
output_size_, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
break;
|
||||
case 3:
|
||||
CalRepeatElements3d(input_device_address, input_shape_[1], input_shape_[2], rep_, axis_, output_device_address,
|
||||
output_shape_[1], output_shape_[2], output_size_,
|
||||
reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
break;
|
||||
case 4:
|
||||
CalRepeatElements4d(input_device_address, input_shape_[1], input_shape_[2], input_shape_[3], rep_, axis_,
|
||||
output_device_address, output_shape_[1], output_shape_[2], output_shape_[3], output_size_,
|
||||
reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
break;
|
||||
case 5:
|
||||
CalRepeatElements5d(input_device_address, input_shape_[1], input_shape_[2], input_shape_[3], input_shape_[4],
|
||||
rep_, axis_, output_device_address, output_shape_[1], output_shape_[2], output_shape_[3],
|
||||
output_shape_[4], output_size_, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
break;
|
||||
default:
|
||||
int *input_shape_device_address = GetDeviceAddress<int>(workspace, 0);
|
||||
int *output_shape_device_address = GetDeviceAddress<int>(workspace, 1);
|
||||
int *input_shape_cumulative_product_device_address = GetDeviceAddress<int>(workspace, 2);
|
||||
CHECK_CUDA_RET_WITH_EXCEPT(
|
||||
cudaMemcpyAsync(input_shape_device_address, input_shape_.data(), workspace_size_list_[0],
|
||||
cudaMemcpyHostToDevice, reinterpret_cast<cudaStream_t>(stream_ptr)),
|
||||
"cudaMemcpyAsync input_shape failed");
|
||||
CHECK_CUDA_RET_WITH_EXCEPT(
|
||||
cudaMemcpyAsync(output_shape_device_address, output_shape_.data(), workspace_size_list_[1],
|
||||
cudaMemcpyHostToDevice, reinterpret_cast<cudaStream_t>(stream_ptr)),
|
||||
"cudaMemcpyAsync output_shape failed");
|
||||
CHECK_CUDA_RET_WITH_EXCEPT(
|
||||
cudaMemcpyAsync(input_shape_cumulative_product_device_address, input_shape_cumulative_product_.data(),
|
||||
workspace_size_list_[2], cudaMemcpyHostToDevice, reinterpret_cast<cudaStream_t>(stream_ptr)),
|
||||
"cudaMemcpyAsync input_shape_cumulative_product_device_address failed");
|
||||
|
||||
CalRepeatElements(input_device_address, input_dim_, input_shape_device_address,
|
||||
input_shape_cumulative_product_device_address, rep_, axis_, output_device_address,
|
||||
output_shape_device_address, output_size_, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
break;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
size_t input_count = AnfAlgo::GetInputTensorNum(kernel_node);
|
||||
if (input_count != 1) {
|
||||
MS_LOG(EXCEPTION) << input_count << " arguments were provided, but RepeatElementsGpuKernel expects 1.";
|
||||
}
|
||||
|
||||
std::vector<size_t> temp_input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
|
||||
input_dim_ = temp_input_shape.size();
|
||||
for (size_t e : temp_input_shape) {
|
||||
input_size_ *= e;
|
||||
input_shape_.push_back(e);
|
||||
}
|
||||
|
||||
int cumulative_product = 1;
|
||||
for (size_t i = input_dim_ - 1; i > 0; i--) {
|
||||
cumulative_product *= input_shape_[i];
|
||||
input_shape_cumulative_product_.push_back(cumulative_product);
|
||||
}
|
||||
std::reverse(input_shape_cumulative_product_.begin(), input_shape_cumulative_product_.end());
|
||||
|
||||
axis_ = static_cast<int>(GetAttr<int64_t>(kernel_node, "axis"));
|
||||
if (axis_ < 0) {
|
||||
axis_ += input_dim_;
|
||||
}
|
||||
|
||||
rep_ = static_cast<int>(GetAttr<int64_t>(kernel_node, "rep"));
|
||||
output_size_ = input_size_ * rep_;
|
||||
output_shape_ = input_shape_;
|
||||
output_shape_[axis_] *= rep_;
|
||||
|
||||
InitSizeLists();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(input_size_ * sizeof(T));
|
||||
output_size_list_.push_back(output_size_ * sizeof(T));
|
||||
|
||||
// workspaces for input shape, output shape and cumulative sum
|
||||
workspace_size_list_.push_back(input_dim_ * sizeof(int));
|
||||
workspace_size_list_.push_back(input_dim_ * sizeof(int));
|
||||
workspace_size_list_.push_back((input_dim_ - 1) * sizeof(int));
|
||||
}
|
||||
|
||||
private:
|
||||
int rep_;
|
||||
int axis_;
|
||||
int input_dim_;
|
||||
std::vector<int> input_shape_;
|
||||
std::vector<int> input_shape_cumulative_product_;
|
||||
std::vector<int> output_shape_;
|
||||
|
||||
size_t input_size_;
|
||||
size_t output_size_;
|
||||
|
||||
std::vector<size_t> input_size_list_;
|
||||
std::vector<size_t> output_size_list_;
|
||||
std::vector<size_t> workspace_size_list_;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_REPEAT_ELEMENTS_GPU_KERNEL_H_
|
@ -1,29 +0,0 @@
|
||||
/**
|
||||
* 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 <cstdint>
|
||||
|
||||
#include "backend/kernel_compiler/gpu/arrays/repeat_elements_grad_gpu_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(RepeatElementsGrad,
|
||||
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeFloat16),
|
||||
RepeatElementsGradGpuKernel, half)
|
||||
|
||||
MS_REG_GPU_KERNEL_ONE(RepeatElementsGrad, KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
|
||||
RepeatElementsGradGpuKernel, int32_t)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
@ -1,119 +0,0 @@
|
||||
/**
|
||||
* 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_GPU_REPEAT_ELEMENTS_GRAD_GPU_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_REPEAT_ELEMENTS_GRAD_GPU_KERNEL_H_
|
||||
|
||||
#include "backend/kernel_compiler/gpu/cuda_impl/repeat_elements_grad_impl.cuh"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <vector>
|
||||
|
||||
#include "backend/kernel_compiler/gpu/gpu_kernel.h"
|
||||
#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
class RepeatElementsGradGpuKernel : public GpuKernel {
|
||||
public:
|
||||
RepeatElementsGradGpuKernel()
|
||||
: rep_(1), axis_(0), input_size_(1), output_size_(0), outer_size_(1), repeat_dim_size_(1), inner_size_(1) {}
|
||||
~RepeatElementsGradGpuKernel() = default;
|
||||
|
||||
const std::vector<size_t> &GetInputSizeList() const override { return input_size_list_; }
|
||||
const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; }
|
||||
const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; }
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
T *dy = GetDeviceAddress<T>(inputs, 0);
|
||||
T *dx = GetDeviceAddress<T>(outputs, 0);
|
||||
|
||||
CalRepeatElementsGrad(dy, rep_, dx, outer_size_, repeat_dim_size_, inner_size_,
|
||||
reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
size_t input_count = AnfAlgo::GetInputTensorNum(kernel_node);
|
||||
if (input_count != 1) {
|
||||
MS_LOG(EXCEPTION) << input_count << " arguments were provided, but RepeatElementGradGpuKernel expects 1.";
|
||||
}
|
||||
|
||||
std::vector<size_t> dy_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
|
||||
int dy_dim = dy_shape.size();
|
||||
|
||||
axis_ = static_cast<int>(GetAttr<int64_t>(kernel_node, "axis"));
|
||||
if (axis_ < 0) {
|
||||
axis_ += dy_dim;
|
||||
}
|
||||
rep_ = static_cast<int>(GetAttr<int64_t>(kernel_node, "rep"));
|
||||
if (axis_ >= dy_dim) {
|
||||
axis_ = dy_dim - 1;
|
||||
rep_ = 1;
|
||||
}
|
||||
|
||||
for (int i = 0; i < dy_dim; i++) {
|
||||
auto e = dy_shape[i];
|
||||
input_size_ *= e;
|
||||
input_shape_.push_back(e);
|
||||
if (i < axis_) {
|
||||
outer_size_ *= e;
|
||||
} else if (i > axis_) {
|
||||
inner_size_ *= e;
|
||||
} else {
|
||||
repeat_dim_size_ = e / rep_;
|
||||
}
|
||||
}
|
||||
|
||||
output_size_ = input_size_ / rep_;
|
||||
output_shape_ = input_shape_;
|
||||
output_shape_[axis_] /= rep_;
|
||||
|
||||
InitSizeLists();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(input_size_ * sizeof(T));
|
||||
output_size_list_.push_back(output_size_ * sizeof(T));
|
||||
}
|
||||
|
||||
private:
|
||||
int rep_;
|
||||
int axis_;
|
||||
size_t input_size_;
|
||||
size_t output_size_;
|
||||
int outer_size_;
|
||||
int repeat_dim_size_;
|
||||
int inner_size_;
|
||||
std::vector<int> input_shape_;
|
||||
std::vector<int> output_shape_;
|
||||
|
||||
std::vector<size_t> input_size_list_;
|
||||
std::vector<size_t> output_size_list_;
|
||||
std::vector<size_t> workspace_size_list_;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_REPEAT_ELEMENTS_GRAD_GPU_KERNEL_H_
|
@ -1,48 +0,0 @@
|
||||
/**
|
||||
* 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 <cuda_runtime.h>
|
||||
|
||||
#include "repeat_elements_grad_impl.cuh"
|
||||
#include "runtime/device/gpu/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
__global__ void RepeatElementsGrad(const int dx_size, const T *dy, const int rep, T *dx, const int outer_size,
|
||||
const int repeat_dim_size, const int inner_size) {
|
||||
for (size_t t_id = blockIdx.x * blockDim.x + threadIdx.x; t_id < dx_size; t_id += gridDim.x * blockDim.x) {
|
||||
int inner_id = t_id % inner_size;
|
||||
int repeat_dim_id = t_id / inner_size % repeat_dim_size;
|
||||
int outer_id = t_id / inner_size / repeat_dim_size;
|
||||
T dx_i = static_cast<T>(0);
|
||||
for (int i = 0; i < rep; i++) {
|
||||
dx_i += dy[(outer_id * rep * repeat_dim_size * inner_size) + (repeat_dim_id * rep * inner_size) +
|
||||
(i * inner_size) + inner_id];
|
||||
}
|
||||
dx[t_id] = dx_i;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElementsGrad(const T *dy, const int rep, T *dx, const int outer_size, const int repeat_dim_size,
|
||||
const int inner_size, cudaStream_t cuda_stream) {
|
||||
const int dx_size = outer_size * repeat_dim_size * inner_size;
|
||||
RepeatElementsGrad<<<GET_BLOCKS(dx_size), GET_THREADS, 0, cuda_stream>>>(dx_size, dy, rep, dx, outer_size,
|
||||
repeat_dim_size, inner_size);
|
||||
}
|
||||
|
||||
template void CalRepeatElementsGrad<int>(const int *dy, const int rep, int *dx, const int outer_size,
|
||||
const int repeat_dim_size, const int inner_size, cudaStream_t cuda_stream);
|
||||
template void CalRepeatElementsGrad<half>(const half *dy, const int rep, half *dx, const int outer_size,
|
||||
const int repeat_dim_size, const int inner_size, cudaStream_t cuda_stream);
|
@ -1,26 +0,0 @@
|
||||
/**
|
||||
* 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_GPU_CUDA_IMPL_REPEAT_ELEMENTS_GRAD_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_REPEAT_ELEMENTS_GRAD_H_
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElementsGrad(const T *dy, const int rep, T *dx, const int outer_size, const int repeat_dim_size,
|
||||
const int inner_size, cudaStream_t cuda_stream);
|
||||
|
||||
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_REPEAT_ELEMENTS_GRAD_H_
|
File diff suppressed because it is too large
Load Diff
@ -1,52 +0,0 @@
|
||||
/**
|
||||
* 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_GPU_CUDA_IMPL_REPEAT_ELEMENTS_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_REPEAT_ELEMENTS_H_
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define REPEAT_ELEMENTS_MAX_INPUT_DIM 100
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElements1d(
|
||||
const T *input, const int rep, const int axis, T *output, const int output_size, cudaStream_t cuda_stream);
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElements2d(const T *input, const int input_d1, const int rep, const int axis, T *output,
|
||||
const int output_d1, const int output_size, cudaStream_t cuda_stream);
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElements3d(const T *input, const int input_d1, const int input_d2, const int rep, const int axis,
|
||||
T *output, const int output_d1, const int output_d2, const int output_size,
|
||||
cudaStream_t cuda_stream);
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElements4d(const T *input, const int input_d1, const int input_d2, const int input_d3, const int rep,
|
||||
const int axis, T *output, const int output_d1, const int output_d2, const int output_d3,
|
||||
const int output_size, cudaStream_t cuda_stream);
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElements5d(const T *input, const int input_d1, const int input_d2, const int input_d3, const int input_d4,
|
||||
const int rep, const int axis, T *output, const int output_d1, const int output_d2,
|
||||
const int output_d3, const int output_d4, const int output_size, cudaStream_t cuda_stream);
|
||||
|
||||
template <typename T>
|
||||
void CalRepeatElements(const T *input, const int input_dim, const int* const input_shape,
|
||||
const int* const input_shape_cumulative_product, const int rep, const int axis, T *output,
|
||||
const int* const output_shape, const int output_size, cudaStream_t cuda_stream);
|
||||
|
||||
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_REPEAT_ELEMENTS_H_
|
@ -0,0 +1,100 @@
|
||||
# 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.
|
||||
# ============================================================================
|
||||
"""math Operations."""
|
||||
from mindspore.ops.composite.multitype_ops import _constexpr_utils as const_utils
|
||||
from mindspore.common import dtype as mstype
|
||||
from mindspore._checkparam import Validator as validator
|
||||
from mindspore._checkparam import Rel
|
||||
from mindspore.ops.primitive import constexpr
|
||||
from mindspore.ops import functional as F
|
||||
from .. import operations as P
|
||||
|
||||
|
||||
@constexpr
|
||||
def _check_is_int(arg_value, arg_name, op_name):
|
||||
arg_value = validator.check_is_int(arg_value, arg_name, op_name)
|
||||
return arg_value
|
||||
|
||||
|
||||
@constexpr
|
||||
def _check_positive_int(arg_value, arg_name, op_name):
|
||||
arg_value = validator.check_positive_int(arg_value, arg_name, op_name)
|
||||
return arg_value
|
||||
|
||||
|
||||
@constexpr
|
||||
def _check_axis_range(arg_value, limit, arg_name, op_name):
|
||||
arg_value = validator.check_int_range(arg_value, -limit, limit, Rel.INC_LEFT, arg_name, op_name)
|
||||
return arg_value
|
||||
|
||||
|
||||
@constexpr
|
||||
def _cal_repeat_dims(x_rank, rep, expand_axis):
|
||||
rep_dims = [1] * (x_rank + 1)
|
||||
rep_dims[expand_axis] = rep
|
||||
return tuple(rep_dims)
|
||||
|
||||
|
||||
@constexpr
|
||||
def _cal_reshape(x_shape, rep, axis):
|
||||
x_reshape = list(x_shape)
|
||||
x_reshape[axis] *= rep
|
||||
return tuple(x_reshape)
|
||||
|
||||
|
||||
def repeat_elements(x, rep, axis=0):
|
||||
"""
|
||||
Repeat elements of a tensor along an axis, like np.repeat.
|
||||
|
||||
Args:
|
||||
- **x** (Tensor) - The tensor to repeat values for.
|
||||
- **rep** (int) - The number of times to repeat, must be positive, required.
|
||||
- **axis** (int) - The axis along which to repeat, default 0.
|
||||
|
||||
Outputs:
|
||||
One tensor with values repeated along the specified axis. If x has shape
|
||||
(s1, s2, ..., sn) and axis is i, the output will have shape (s1, s2, ..., si * rep, ..., sn)
|
||||
|
||||
Examples:
|
||||
>>> x = Tensor(np.array([[0, 1, 2], [3, 4, 5]]), mindspore.int32)
|
||||
>>> output = C.RepeatElements(x, rep = 2, axis = 0)
|
||||
>>> print(output)
|
||||
[[0, 1, 2],
|
||||
[0, 1, 2],
|
||||
[3, 4, 5],
|
||||
[3, 4, 5]],
|
||||
"""
|
||||
const_utils.check_valid_type(F.dtype(x), mstype.number_type, 'input x')
|
||||
rep = _check_positive_int(rep, "rep", "repeat_elements")
|
||||
axis = _check_is_int(axis, "axis", "repeat_elements")
|
||||
|
||||
shape_op = P.Shape()
|
||||
rank_op = P.Rank()
|
||||
tile_op = P.Tile()
|
||||
expand_dims_op = P.ExpandDims()
|
||||
reshape_op = P.Reshape()
|
||||
|
||||
x_rank = rank_op(x)
|
||||
axis = _check_axis_range(axis, x_rank, "axis", "repeat_elements")
|
||||
|
||||
expand_axis = axis + 1
|
||||
x_expand = expand_dims_op(x, expand_axis)
|
||||
rep_dims = _cal_repeat_dims(x_rank, rep, expand_axis)
|
||||
x_expand = tile_op(x_expand, rep_dims)
|
||||
x_shape = shape_op(x)
|
||||
x_reshape = _cal_reshape(x_shape, rep, axis)
|
||||
x_rep = reshape_op(x_expand, x_reshape)
|
||||
|
||||
return x_rep
|
File diff suppressed because it is too large
Load Diff
@ -1,86 +0,0 @@
|
||||
# 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.
|
||||
|
||||
import numpy as np
|
||||
|
||||
import mindspore as ms
|
||||
from mindspore import context, Tensor, Parameter
|
||||
from mindspore.common.api import _executor
|
||||
from mindspore.nn import Cell, TrainOneStepCell, Momentum
|
||||
from mindspore.ops import operations as P
|
||||
|
||||
|
||||
class Net(Cell):
|
||||
def __init__(self, mul_weight, strategy1=None, strategy2=None):
|
||||
super().__init__()
|
||||
self.mul = P.Mul().shard(strategy1)
|
||||
self.repeat = P.RepeatElements(rep=2, axis=1).shard(strategy2)
|
||||
self.mul_weight = Parameter(mul_weight, "w1")
|
||||
|
||||
def construct(self, x, b):
|
||||
out = self.mul(x, self.mul_weight)
|
||||
out = self.repeat(out)
|
||||
return out
|
||||
|
||||
|
||||
_x = Tensor(np.ones([128, 64, 32]), dtype=ms.float32)
|
||||
_w1 = Tensor(np.ones([128, 64, 32]), dtype=ms.float32)
|
||||
_b = Tensor(np.ones([128, 64, 32]), dtype=ms.float32)
|
||||
|
||||
|
||||
def compile_net(net):
|
||||
optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
|
||||
train_net = TrainOneStepCell(net, optimizer)
|
||||
train_net.set_auto_parallel()
|
||||
train_net.set_train()
|
||||
_executor.compile(train_net, _x, _b)
|
||||
context.reset_auto_parallel_context()
|
||||
|
||||
|
||||
def test_repeat_elements_data_parallel():
|
||||
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
|
||||
strategy1 = ((16, 1, 1), (16, 1, 1))
|
||||
strategy2 = ((16, 1, 1),)
|
||||
net = Net(_w1, strategy1, strategy2)
|
||||
compile_net(net)
|
||||
|
||||
|
||||
def test_repeat_elements_model_parallel():
|
||||
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
|
||||
strategy1 = ((1, 1, 16), (1, 1, 16))
|
||||
strategy2 = ((1, 1, 16),)
|
||||
net = Net(_w1, strategy1, strategy2)
|
||||
compile_net(net)
|
||||
|
||||
|
||||
def test_repeat_elements_hybrid_parallel():
|
||||
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
|
||||
strategy1 = ((2, 2, 4), (2, 2, 4))
|
||||
strategy2 = ((2, 2, 4),)
|
||||
net = Net(_w1, strategy1, strategy2)
|
||||
compile_net(net)
|
||||
|
||||
|
||||
def test_repeat_elements_auto_parallel():
|
||||
context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=16, global_rank=0)
|
||||
net = Net(_w1)
|
||||
compile_net(net)
|
||||
|
||||
|
||||
def test_repeat_elements_repeat_calc():
|
||||
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
|
||||
strategy1 = ((2, 2, 4), (2, 2, 4))
|
||||
strategy2 = ((1, 2, 2),)
|
||||
net = Net(_w1, strategy1, strategy2)
|
||||
compile_net(net)
|
Loading…
Reference in new issue