diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.cc index 39f535a2af..881abe6d7d 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.cc @@ -1,5 +1,5 @@ /** - * Copyright 2019 Huawei Technologies Co., Ltd + * 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. diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.h index cbcf0c8f48..18b6f231f2 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.h @@ -1,5 +1,5 @@ /** - * Copyright 2019 Huawei Technologies Co., Ltd + * 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. @@ -14,8 +14,8 @@ * limitations under the License. */ -#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_ARGMAXGPUKERNEL_H_ -#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_ARGMAXGPUKERNEL_H_ +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_ARRAYS_ARGMAX_GPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_ARRAYS_ARGMAX_GPU_KERNEL_H_ #include #include "backend/kernel_compiler/gpu/gpu_kernel.h" @@ -60,13 +60,13 @@ class ArgmaxGpuKernel : public GpuKernel { } auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0); if (input_shape.size() > ARGMAX_MAX_DIMENSION) { - MS_LOG(EXCEPTION) << "Input is " << input_shape.size() << "-D, but argmax supports max " << ARGMAX_MAX_DIMENSION + MS_LOG(EXCEPTION) << "Input is " << input_shape.size() << "-D, but Argmax supports max " << ARGMAX_MAX_DIMENSION << "-D inputs."; } - axis_ = static_cast(GetAttr(kernel_node, "axis")); + axis_ = GetAttr(kernel_node, "axis"); if (axis_ < 0) { - axis_ += SizeToInt(input_shape.size()); + axis_ += static_cast(input_shape.size()); } if (input_shape.size() == 1) { batch_size_ = 0; @@ -98,9 +98,9 @@ class ArgmaxGpuKernel : public GpuKernel { std::vector workspace_size_list_; size_t batch_size_; size_t channel_size_; - int axis_; + int64_t axis_; }; } // namespace kernel } // namespace mindspore -#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_ARGMAXGPUKERNEL_H_ +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_ARRAYS_ARGMAX_GPU_KERNEL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cu index a4f1f6680b..116bd9f42d 100755 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cu @@ -1,5 +1,5 @@ /** - * Copyright 2019 Huawei Technologies Co., Ltd + * 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. @@ -18,7 +18,7 @@ #include "runtime/device/gpu/cuda_common.h" #include "include/cuda_fp16.h" template -__global__ void Argmax1D(const T* input, const int channel_size, int* output) { +__global__ void Argmax1D(const T *input, const int channel_size, int *output) { int max_index = 0; T max = input[0]; for (int pos = 1; pos < channel_size; pos++) { @@ -31,7 +31,7 @@ __global__ void Argmax1D(const T* input, const int channel_size, int* output) { return; } template -__global__ void ArgmaxDefault2D(const T* input, const int batch_size, const int channel_size, int* output) { +__global__ void ArgmaxDefault2D(const T *input, const int batch_size, const int channel_size, int *output) { int pos; int max_index; T max; @@ -51,7 +51,7 @@ __global__ void ArgmaxDefault2D(const T* input, const int batch_size, const int return; } template -__global__ void ArgmaxAxis2D(const T* input, const int batch_size, const int channel_size, int* output) { +__global__ void ArgmaxAxis2D(const T *input, const int batch_size, const int channel_size, int *output) { int pos; int max_index; T max; @@ -70,7 +70,7 @@ __global__ void ArgmaxAxis2D(const T* input, const int batch_size, const int cha return; } template -void CalArgmax(const T* input, const int batch_size, const int channel_size, const int axis, int* output, +void CalArgmax(const T *input, const int batch_size, const int channel_size, const int64_t axis, int *output, cudaStream_t cuda_stream) { if (batch_size == 0) { Argmax1D<<<1, 1, 0, cuda_stream>>>(input, channel_size, output); @@ -82,7 +82,7 @@ void CalArgmax(const T* input, const int batch_size, const int channel_size, con return; } -template void CalArgmax(const float* input, const int batch_size, const int channel_size, const int axis, - int* output, cudaStream_t cuda_stream); -template void CalArgmax(const half* input, const int batch_size, const int channel_size, const int axis, - int* output, cudaStream_t cuda_stream); +template void CalArgmax(const float *input, const int batch_size, const int channel_size, const int64_t axis, + int *output, cudaStream_t cuda_stream); +template void CalArgmax(const half *input, const int batch_size, const int channel_size, const int64_t axis, + int *output, cudaStream_t cuda_stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cuh index 1e90f133ff..ddebaca7e1 100755 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cuh @@ -1,5 +1,5 @@ /** - * Copyright 2019 Huawei Technologies Co., Ltd + * 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. @@ -14,10 +14,10 @@ * limitations under the License. */ -#ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMP_ARGMAX_H_ -#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMP_ARGMAX_H_ +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_ARGMAX_IMPL_CUH_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_ARGMAX_IMPL_CUH_ template -void CalArgmax(const T* input, const int batch_size, const int channel_size, const int axis, int* output, +void CalArgmax(const T *input, const int batch_size, const int channel_size, const int64_t axis, int *output, cudaStream_t cuda_stream); -#endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMP_ARGMAX_H_ +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_ARGMAX_IMPL_CUH_ diff --git a/tests/st/ops/gpu/test_argmax_op.py b/tests/st/ops/gpu/test_argmax_op.py index 7b90bc4eaa..ccf492cfb4 100644 --- a/tests/st/ops/gpu/test_argmax_op.py +++ b/tests/st/ops/gpu/test_argmax_op.py @@ -1,4 +1,4 @@ -# Copyright 2019 Huawei Technologies Co., Ltd +# 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.