| @@ -1,5 +1,5 @@ | |||
| /** | |||
| * Copyright 2019 Huawei Technologies Co., Ltd | |||
| * Copyright 2019-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,17 +14,17 @@ | |||
| * limitations under the License. | |||
| */ | |||
| #include "backend/kernel_compiler/gpu/arrays/gather_gpu_kernel.h" | |||
| #include "backend/kernel_compiler/gpu/arrays/gatherv2_gpu_kernel.h" | |||
| namespace mindspore { | |||
| namespace kernel { | |||
| MS_REG_GPU_KERNEL_TWO( | |||
| GatherV2, | |||
| KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat32), | |||
| GatherGpuFwdKernel, float, int) | |||
| GatherV2GpuFwdKernel, float, int) | |||
| MS_REG_GPU_KERNEL_TWO( | |||
| GatherV2, | |||
| KernelAttr().AddInputAttr(kNumberTypeFloat16).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat16), | |||
| GatherGpuFwdKernel, half, int) | |||
| GatherV2GpuFwdKernel, half, int) | |||
| } // namespace kernel | |||
| } // namespace mindspore | |||
| @@ -1,5 +1,5 @@ | |||
| /** | |||
| * Copyright 2019 Huawei Technologies Co., Ltd | |||
| * Copyright 2019-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. | |||
| @@ -20,15 +20,15 @@ | |||
| #include <vector> | |||
| #include "backend/kernel_compiler/gpu/gpu_kernel.h" | |||
| #include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" | |||
| #include "backend/kernel_compiler/gpu/cuda_impl/gather.cuh" | |||
| #include "backend/kernel_compiler/gpu/cuda_impl/gatherv2.cuh" | |||
| namespace mindspore { | |||
| namespace kernel { | |||
| template <typename T, typename S> | |||
| class GatherGpuFwdKernel : public GpuKernel { | |||
| class GatherV2GpuFwdKernel : public GpuKernel { | |||
| public: | |||
| GatherGpuFwdKernel() : axis_(0), handle_(nullptr) {} | |||
| ~GatherGpuFwdKernel() = default; | |||
| GatherV2GpuFwdKernel() : axis_(0), handle_(nullptr) {} | |||
| ~GatherV2GpuFwdKernel() = 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_; } | |||
| @@ -42,15 +42,15 @@ class GatherGpuFwdKernel : public GpuKernel { | |||
| T *output_addr = GetDeviceAddress<T>(outputs, 0); | |||
| auto input_dim1 = input_shapes_[IntToSize(axis_)]; | |||
| Gather(input_addr, indices_addr, output_addr, dims_[0], dims_[1], dims_[2], input_dim1, | |||
| reinterpret_cast<cudaStream_t>(stream_ptr)); | |||
| GatherV2(input_addr, indices_addr, output_addr, dims_[0], dims_[1], dims_[2], input_dim1, | |||
| reinterpret_cast<cudaStream_t>(stream_ptr)); | |||
| return true; | |||
| } | |||
| bool Init(const CNodePtr &kernel_node) override { | |||
| InitResource(); | |||
| size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node); | |||
| if (input_num != 2) { | |||
| MS_LOG(EXCEPTION) << "Argument number is " << input_num << ", but GatherGpuFwdKernel needs 2."; | |||
| MS_LOG(EXCEPTION) << "Argument number is " << input_num << ", but GatherGpuV2FwdKernel needs 2."; | |||
| } | |||
| input_shapes_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0); | |||
| indices_shapes_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1); | |||
| @@ -1,5 +1,5 @@ | |||
| /** | |||
| * Copyright 2019 Huawei Technologies Co., Ltd | |||
| * Copyright 2019-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. | |||
| @@ -15,10 +15,10 @@ | |||
| */ | |||
| #include <iostream> | |||
| #include "backend/kernel_compiler/gpu/cuda_impl/gather.cuh" | |||
| #include "backend/kernel_compiler/gpu/cuda_impl/gatherv2.cuh" | |||
| #include "runtime/device/gpu/cuda_common.h" | |||
| template <typename T, typename S> | |||
| __global__ void GatherKernel(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, | |||
| __global__ void GatherV2Kernel(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, | |||
| size_t output_dim2, size_t input_dim1) { | |||
| int num = output_dim0 * output_dim1 * output_dim2; | |||
| int i, j, k; | |||
| @@ -39,16 +39,16 @@ __global__ void GatherKernel(T *input, S *indices, T *output, size_t output_dim0 | |||
| return; | |||
| } | |||
| template <typename T, typename S> | |||
| void Gather(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, size_t output_dim2, | |||
| void GatherV2(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, size_t output_dim2, | |||
| size_t input_dim1, cudaStream_t stream) { | |||
| int size = output_dim0 * output_dim1 * output_dim2; | |||
| GatherKernel<<<GET_BLOCKS(size), GET_THREADS, 0, stream>>>(input, indices, output, output_dim0, output_dim1, | |||
| GatherV2Kernel<<<GET_BLOCKS(size), GET_THREADS, 0, stream>>>(input, indices, output, output_dim0, output_dim1, | |||
| output_dim2, input_dim1); | |||
| return; | |||
| } | |||
| template void Gather<float, int>(float *input, int *indices, float *output, size_t output_dim0, size_t output_dim1, | |||
| template void GatherV2<float, int>(float *input, int *indices, float *output, size_t output_dim0, size_t output_dim1, | |||
| size_t output_dim2, size_t input_dim1, cudaStream_t stream); | |||
| template void Gather<half, int>(half *input, int *indices, half *output, size_t output_dim0, size_t output_dim1, | |||
| template void GatherV2<half, int>(half *input, int *indices, half *output, size_t output_dim0, size_t output_dim1, | |||
| size_t output_dim2, size_t input_dim1, cudaStream_t stream); | |||
| @@ -1,5 +1,5 @@ | |||
| /** | |||
| * Copyright 2019 Huawei Technologies Co., Ltd | |||
| * Copyright 2019-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. | |||
| @@ -17,7 +17,7 @@ | |||
| #ifndef MINDSPORE_GATHER_GPU_CU_H | |||
| #define MINDSPORE_GATHER_GPU_CU_H | |||
| template <typename T, typename S> | |||
| void Gather(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, size_t output_dim2, | |||
| size_t input_dim1, cudaStream_t stream); | |||
| void GatherV2(T *input, S *indices, T *output, size_t output_dim0, size_t output_dim1, | |||
| size_t output_dim2, size_t input_dim1, cudaStream_t stream); | |||
| #endif | |||