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

change tensor shape from int to size_t for GPU Argmax

This commit is contained in:
TFbunny 2020-12-09 17:45:10 -05:00
parent a15356c779
commit d274f425c7
5 changed files with 24 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.cc
View File

@ -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.

16
mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/argmax_gpu_kernel.h
View File

@ -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_

18
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cu
View File

@ -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);

10
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/argmax_impl.cuh
View File

@ -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_

2
tests/st/ops/gpu/test_argmax_op.py
View File

@ -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.