change tensor shape from int to size_t for GPU Argmax

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.

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 <vector>
 #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<int>(GetAttr<int64_t>(kernel_node, "axis"));
+    axis_ = GetAttr<int64_t>(kernel_node, "axis");
     if (axis_ < 0) {
-      axis_ += SizeToInt(input_shape.size());
+      axis_ += static_cast<int64_t>(input_shape.size());
     }
     if (input_shape.size() == 1) {
       batch_size_ = 0;
@@ -98,9 +98,9 @@ class ArgmaxGpuKernel : public GpuKernel {
   std::vector<size_t> 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_

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 <typename T>
-__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 <typename T>
-__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 <typename T>
-__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 <typename T>
-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<float>(const float* input, const int batch_size, const int channel_size, const int axis,
-                               int* output, cudaStream_t cuda_stream);
-template void CalArgmax<half>(const half* input, const int batch_size, const int channel_size, const int axis,
-                              int* output, cudaStream_t cuda_stream);
+template void CalArgmax<float>(const float *input, const int batch_size, const int channel_size, const int64_t axis,
+                               int *output, cudaStream_t cuda_stream);
+template void CalArgmax<half>(const half *input, const int batch_size, const int channel_size, const int64_t axis,
+                              int *output, cudaStream_t cuda_stream);

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 <typename T>
-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_

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.