!37766 PMEFFT1D, 2D； PMEBatchedFFT2D；PMEIFFT1D, 2D；PMERFFT2D；PMEIRFFT2D算子前后端实现
Merge pull request !37766 from 李佳兴/lijiaxing1
This commit is contained in:
Gitee
commit
180d0ed6a5
|
|
@ -0,0 +1,28 @@
|
|||
/**
|
||||
* Copyright 2021 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 "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_batched_fft_2d_impl.cuh"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_common.cuh"
|
||||
|
||||
template <typename T>
|
||||
void PMEBatchedFFT2D(Complex<T> *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, int direction, cudaStream_t stream) {
|
||||
cufftExecC2C(FFT_plan_c2c, reinterpret_cast<cufftComplex *>(input_tensor),
|
||||
reinterpret_cast<cufftComplex *>(output_tensor), direction);
|
||||
return;
|
||||
}
|
||||
|
||||
template void PMEBatchedFFT2D<float>(Complex<float> *input_tensor, Complex<float> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, int direction, cudaStream_t stream);
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_BATCHED_FFT_2D_IMPL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_BATCHED_FFT_2D_IMPL_H_
|
||||
|
||||
#include <cufft.h>
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/complex.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
void PMEBatchedFFT2D(Complex<T> *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, int direction, cudaStream_t stream);
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
/**
|
||||
* Copyright 2021 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 "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_fft_1d_impl.cuh"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_common.cuh"
|
||||
|
||||
template <typename T>
|
||||
void PMEFFT1D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor, const cufftHandle &FFT_plan_c2c,
|
||||
cudaStream_t stream) {
|
||||
cufftExecC2C(FFT_plan_c2c, reinterpret_cast<cufftComplex *>(input_tensor),
|
||||
reinterpret_cast<cufftComplex *>(output_tensor), CUFFT_FORWARD);
|
||||
return;
|
||||
}
|
||||
|
||||
template void PMEFFT1D<float>(int Nfft, Complex<float> *input_tensor, Complex<float> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_FFT_1D_IMPL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_FFT_1D_IMPL_H_
|
||||
|
||||
#include <cufft.h>
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/complex.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
void PMEFFT1D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
/**
|
||||
* Copyright 2021 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 "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_fft_2d_impl.cuh"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_common.cuh"
|
||||
|
||||
template <typename T>
|
||||
void PMEFFT2D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream) {
|
||||
cufftExecC2C(FFT_plan_c2c, reinterpret_cast<cufftComplex *>(input_tensor),
|
||||
reinterpret_cast<cufftComplex *>(output_tensor), CUFFT_FORWARD);
|
||||
return;
|
||||
}
|
||||
|
||||
template void PMEFFT2D<float>(int Nfft, Complex<float> *input_tensor, Complex<float> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_FFT_2D_IMPL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_FFT_2D_IMPL_H_
|
||||
|
||||
#include <cufft.h>
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/complex.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
void PMEFFT2D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
/**
|
||||
* Copyright 2021 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 "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_ifft_1d_impl.cuh"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_common.cuh"
|
||||
|
||||
template <typename T>
|
||||
void PMEIFFT1D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor, const cufftHandle &FFT_plan_c2c,
|
||||
cudaStream_t stream) {
|
||||
cufftExecC2C(FFT_plan_c2c, reinterpret_cast<cufftComplex *>(input_tensor),
|
||||
reinterpret_cast<cufftComplex *>(output_tensor), CUFFT_INVERSE);
|
||||
return;
|
||||
}
|
||||
|
||||
template void PMEIFFT1D<float>(int Nfft, Complex<float> *input_tensor, Complex<float> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_IFFT_1D_IMPL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_IFFT_1D_IMPL_H_
|
||||
|
||||
#include <cufft.h>
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/complex.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
void PMEIFFT1D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
/**
|
||||
* Copyright 2021 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 "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_ifft_2d_impl.cuh"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_common.cuh"
|
||||
|
||||
template <typename T>
|
||||
void PMEIFFT2D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor, const cufftHandle &FFT_plan_c2c,
|
||||
cudaStream_t stream) {
|
||||
cufftExecC2C(FFT_plan_c2c, reinterpret_cast<cufftComplex *>(input_tensor),
|
||||
reinterpret_cast<cufftComplex *>(output_tensor), CUFFT_INVERSE);
|
||||
return;
|
||||
}
|
||||
|
||||
template void PMEIFFT2D<float>(int Nfft, Complex<float> *input_tensor, Complex<float> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_IFFT_2D_IMPL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_IFFT_2D_IMPL_H_
|
||||
|
||||
#include <cufft.h>
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/complex.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
void PMEIFFT2D(int Nfft, Complex<T> *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2c, cudaStream_t stream);
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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 "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_irfft_2d_impl.cuh"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_common.cuh"
|
||||
|
||||
template <typename T>
|
||||
void PMEIRFFT2D(int Nfft, Complex<T> *input_tensor, T *output_tensor, const cufftHandle &FFT_plan_c2r,
|
||||
cudaStream_t stream) {
|
||||
cufftExecC2R(FFT_plan_c2r, reinterpret_cast<cufftComplex *>(input_tensor), output_tensor);
|
||||
return;
|
||||
}
|
||||
|
||||
template void PMEIRFFT2D<float>(int Nfft, Complex<float> *input_tensor, float *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2r, cudaStream_t stream);
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_IRFFT_2D_IMPL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_IRFFT_2D_IMPL_H_
|
||||
|
||||
#include <cufft.h>
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/complex.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
void PMEIRFFT2D(int Nfft, Complex<T> *input_tensor, T *output_tensor,
|
||||
const cufftHandle &FFT_plan_c2r, cudaStream_t stream);
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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 "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_rfft_2d_impl.cuh"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_common.cuh"
|
||||
|
||||
template <typename T>
|
||||
void PMERFFT2D(int Nfft, T *input_tensor, Complex<T> *output_tensor,
|
||||
const cufftHandle &FFT_plan_r2c, cudaStream_t stream) {
|
||||
cufftExecR2C(FFT_plan_r2c, input_tensor, reinterpret_cast<cufftComplex *>(output_tensor));
|
||||
return;
|
||||
}
|
||||
|
||||
template void PMERFFT2D<float>(int Nfft, float *input_tensor, Complex<float> *output_tensor,
|
||||
const cufftHandle &FFT_plan_r2c, cudaStream_t stream);
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/**
|
||||
* Copyright 2021 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_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_RFFT_2D_IMPL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_RFFT_2D_IMPL_H_
|
||||
|
||||
#include <cufft.h>
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/complex.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
|
||||
template <typename T>
|
||||
void PMERFFT2D(int Nfft, T *input_tensor, Complex<T> *output_tensor, const cufftHandle &FFT_plan_r2c,
|
||||
cudaStream_t stream);
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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 "plugin/device/gpu/kernel/sponge/pme/pme_batched_fft_2d_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(PMEBatchedFFT2D,
|
||||
KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeComplex64),
|
||||
PMEBatchedFFT2DGpuKernelMod, float)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,82 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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_SPONGE_PME_PME_BATCHED_FFT_2D_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONGE_PME_PME_BATCHED_FFT_2D_KERNEL_H_
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cufft.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel.h"
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_batched_fft_2d_impl.cuh"
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
using Complex = mindspore::utils::Complex<T>;
|
||||
template <typename T>
|
||||
class PMEBatchedFFT2DGpuKernelMod : public DeprecatedNativeGpuKernelMod {
|
||||
public:
|
||||
PMEBatchedFFT2DGpuKernelMod() = default;
|
||||
~PMEBatchedFFT2DGpuKernelMod() override = default;
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
// direction: forward 1, inverse -1
|
||||
direction = static_cast<int>(GetAttr<int64_t>(kernel_node, "direction"));
|
||||
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
|
||||
ffty = static_cast<int>(GetAttr<int64_t>(kernel_node, "ffty"));
|
||||
fftz = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftz"));
|
||||
Nall = fftx * ffty * fftz;
|
||||
|
||||
// config plan
|
||||
int batch = fftx;
|
||||
int dimSize[] = {ffty, fftz};
|
||||
int rank = sizeof(dimSize) / sizeof(dimSize[0]);
|
||||
cufftPlanMany(&FFT_plan_c2c, rank, dimSize, NULL, 0, 0, NULL, 0, 0, CUFFT_C2C, batch);
|
||||
|
||||
InitSizeLists();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
auto input_tensor = GetDeviceAddress<Complex<T>>(inputs, 0);
|
||||
auto output_tensor = GetDeviceAddress<Complex<T>>(outputs, 0);
|
||||
|
||||
cufftSetStream(FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
PMEBatchedFFT2D<T>(input_tensor, output_tensor, FFT_plan_c2c, direction,
|
||||
reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(Nall * sizeof(Complex<T>));
|
||||
output_size_list_.push_back(Nall * sizeof(Complex<T>));
|
||||
}
|
||||
|
||||
private:
|
||||
int direction;
|
||||
int fftx;
|
||||
int ffty;
|
||||
int fftz;
|
||||
int Nall;
|
||||
cufftHandle FFT_plan_c2c;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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 "plugin/device/gpu/kernel/sponge/pme/pme_fft_1d_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(PMEFFT1D, KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeComplex64),
|
||||
PMEFFT1DGpuKernelMod, float)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,73 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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_SPONGE_PME_PME_FFT_1D_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONGE_PME_PME_FFT_1D_KERNEL_H_
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cufft.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel.h"
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_fft_1d_impl.cuh"
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
using Complex = mindspore::utils::Complex<T>;
|
||||
template <typename T>
|
||||
class PMEFFT1DGpuKernelMod : public DeprecatedNativeGpuKernelMod {
|
||||
public:
|
||||
PMEFFT1DGpuKernelMod() = default;
|
||||
~PMEFFT1DGpuKernelMod() override = default;
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
|
||||
Nall = fftx;
|
||||
Nfft = fftx;
|
||||
|
||||
cufftPlan1d(&FFT_plan_c2c, fftx, CUFFT_C2C, batch);
|
||||
|
||||
InitSizeLists();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
auto input_tensor = GetDeviceAddress<Complex<T>>(inputs, 0);
|
||||
auto output_tensor = GetDeviceAddress<Complex<T>>(outputs, 0);
|
||||
|
||||
cufftSetStream(FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
PMEFFT1D<T>(Nfft, input_tensor, output_tensor, FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(Nall * sizeof(Complex<T>));
|
||||
output_size_list_.push_back(Nfft * sizeof(Complex<T>));
|
||||
}
|
||||
|
||||
private:
|
||||
int batch = 1;
|
||||
int fftx;
|
||||
int Nall;
|
||||
int Nfft;
|
||||
cufftHandle FFT_plan_c2c;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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 "plugin/device/gpu/kernel/sponge/pme/pme_fft_2d_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(PMEFFT2D, KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeComplex64),
|
||||
PMEFFT2DGpuKernelMod, float)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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_SPONGE_PME_PME_FFT_2D_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONGE_PME_PME_FFT_2D_KERNEL_H_
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cufft.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel.h"
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_fft_2d_impl.cuh"
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
using Complex = mindspore::utils::Complex<T>;
|
||||
template <typename T>
|
||||
class PMEFFT2DGpuKernelMod : public DeprecatedNativeGpuKernelMod {
|
||||
public:
|
||||
PMEFFT2DGpuKernelMod() = default;
|
||||
~PMEFFT2DGpuKernelMod() override = default;
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
|
||||
ffty = static_cast<int>(GetAttr<int64_t>(kernel_node, "ffty"));
|
||||
Nall = fftx * ffty;
|
||||
Nfft = fftx * ffty;
|
||||
|
||||
cufftPlan2d(&FFT_plan_c2c, fftx, ffty, CUFFT_C2C);
|
||||
|
||||
InitSizeLists();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
auto input_tensor = GetDeviceAddress<Complex<T>>(inputs, 0);
|
||||
auto output_tensor = GetDeviceAddress<Complex<T>>(outputs, 0);
|
||||
|
||||
cufftSetStream(FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
PMEFFT2D<T>(Nfft, input_tensor, output_tensor, FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(Nall * sizeof(Complex<T>));
|
||||
output_size_list_.push_back(Nfft * sizeof(Complex<T>));
|
||||
}
|
||||
|
||||
private:
|
||||
int fftx;
|
||||
int ffty;
|
||||
int Nall;
|
||||
int Nfft;
|
||||
cufftHandle FFT_plan_c2c;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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 "plugin/device/gpu/kernel/sponge/pme/pme_ifft_1d_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(PMEIFFT1D, KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeComplex64),
|
||||
PMEIFFT1DGpuKernelMod, float)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,73 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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_SPONGE_PME_PME_IFFT_1D_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONGE_PME_PME_IFFT_1D_KERNEL_H_
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cufft.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel.h"
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_ifft_1d_impl.cuh"
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
using Complex = mindspore::utils::Complex<T>;
|
||||
template <typename T>
|
||||
class PMEIFFT1DGpuKernelMod : public DeprecatedNativeGpuKernelMod {
|
||||
public:
|
||||
PMEIFFT1DGpuKernelMod() = default;
|
||||
~PMEIFFT1DGpuKernelMod() override = default;
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
|
||||
Nfft = fftx;
|
||||
Nall = fftx;
|
||||
|
||||
cufftPlan1d(&FFT_plan_c2c, fftx, CUFFT_C2C, batch);
|
||||
|
||||
InitSizeLists();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
auto input_tensor = GetDeviceAddress<Complex<T>>(inputs, 0);
|
||||
auto output_tensor = GetDeviceAddress<Complex<T>>(outputs, 0);
|
||||
|
||||
cufftSetStream(FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
PMEIFFT1D<T>(Nfft, input_tensor, output_tensor, FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(Nfft * sizeof(Complex<T>));
|
||||
output_size_list_.push_back(Nall * sizeof(Complex<T>));
|
||||
}
|
||||
|
||||
private:
|
||||
int batch = 1;
|
||||
int fftx;
|
||||
int Nall;
|
||||
int Nfft;
|
||||
cufftHandle FFT_plan_c2c;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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 "plugin/device/gpu/kernel/sponge/pme/pme_ifft_2d_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(PMEIFFT2D, KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeComplex64),
|
||||
PMEIFFT2DGpuKernelMod, float)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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_SPONGE_PME_PME_IFFT_2D_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONGE_PME_PME_IFFT_2D_KERNEL_H_
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cufft.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel.h"
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_ifft_2d_impl.cuh"
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
using Complex = mindspore::utils::Complex<T>;
|
||||
template <typename T>
|
||||
class PMEIFFT2DGpuKernelMod : public DeprecatedNativeGpuKernelMod {
|
||||
public:
|
||||
PMEIFFT2DGpuKernelMod() = default;
|
||||
~PMEIFFT2DGpuKernelMod() override = default;
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
|
||||
ffty = static_cast<int>(GetAttr<int64_t>(kernel_node, "ffty"));
|
||||
Nfft = fftx * ffty;
|
||||
Nall = fftx * ffty;
|
||||
|
||||
cufftPlan2d(&FFT_plan_c2c, fftx, ffty, CUFFT_C2C);
|
||||
|
||||
InitSizeLists();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
auto input_tensor = GetDeviceAddress<Complex<T>>(inputs, 0);
|
||||
auto output_tensor = GetDeviceAddress<Complex<T>>(outputs, 0);
|
||||
|
||||
cufftSetStream(FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
PMEIFFT2D<T>(Nfft, input_tensor, output_tensor, FFT_plan_c2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(Nfft * sizeof(Complex<T>));
|
||||
output_size_list_.push_back(Nall * sizeof(Complex<T>));
|
||||
}
|
||||
|
||||
private:
|
||||
int fftx;
|
||||
int ffty;
|
||||
int Nall;
|
||||
int Nfft;
|
||||
cufftHandle FFT_plan_c2c;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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 "plugin/device/gpu/kernel/sponge/pme/pme_irfft_2d_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(PMEIRFFT2D, KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeFloat32),
|
||||
PMEIRFFT2DGpuKernelMod, float)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,75 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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_SPONGE_PME_PME_IRFFT_2D_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONGE_PME_PME_IRFFT_2D_KERNEL_H_
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cufft.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel.h"
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_irfft_2d_impl.cuh"
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
using Complex = mindspore::utils::Complex<T>;
|
||||
template <typename T>
|
||||
class PMEIRFFT2DGpuKernelMod : public DeprecatedNativeGpuKernelMod {
|
||||
public:
|
||||
PMEIRFFT2DGpuKernelMod() = default;
|
||||
~PMEIRFFT2DGpuKernelMod() override = default;
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
constexpr int NUM_2 = 2;
|
||||
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
|
||||
ffty = static_cast<int>(GetAttr<int64_t>(kernel_node, "ffty"));
|
||||
Nfft = fftx * ffty;
|
||||
Nall = fftx * (ffty - 1) * NUM_2;
|
||||
|
||||
cufftPlan2d(&FFT_plan_c2r, fftx, (ffty - 1) * NUM_2, CUFFT_C2R);
|
||||
|
||||
InitSizeLists();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
auto input_tensor = GetDeviceAddress<Complex<T>>(inputs, 0);
|
||||
auto output_tensor = GetDeviceAddress<T>(outputs, 0);
|
||||
|
||||
cufftSetStream(FFT_plan_c2r, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
PMEIRFFT2D<T>(Nfft, input_tensor, output_tensor, FFT_plan_c2r, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(Nfft * sizeof(Complex<T>));
|
||||
output_size_list_.push_back(Nall * sizeof(T));
|
||||
}
|
||||
|
||||
private:
|
||||
int fftx;
|
||||
int ffty;
|
||||
int Nall;
|
||||
int Nfft;
|
||||
cufftHandle FFT_plan_c2r;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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 "plugin/device/gpu/kernel/sponge/pme/pme_rfft_2d_kernel.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
MS_REG_GPU_KERNEL_ONE(PMERFFT2D, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeComplex64),
|
||||
PMERFFT2DGpuKernelMod, float)
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,75 @@
|
|||
/**
|
||||
* Copyright 2021-2022 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_SPONGE_PME_PME_RFFT_2D_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONGE_PME_PME_RFFT_2D_KERNEL_H_
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cufft.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel.h"
|
||||
#include "plugin/device/gpu/kernel/gpu_kernel_factory.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/cuda_ops/cuda_common.h"
|
||||
#include "plugin/device/gpu/kernel/cuda_impl/sponge/pme/pme_rfft_2d_impl.cuh"
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
template <typename T>
|
||||
using Complex = mindspore::utils::Complex<T>;
|
||||
template <typename T>
|
||||
class PMERFFT2DGpuKernelMod : public DeprecatedNativeGpuKernelMod {
|
||||
public:
|
||||
PMERFFT2DGpuKernelMod() = default;
|
||||
~PMERFFT2DGpuKernelMod() override = default;
|
||||
|
||||
bool Init(const CNodePtr &kernel_node) override {
|
||||
constexpr int NUM_2 = 2;
|
||||
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
|
||||
ffty = static_cast<int>(GetAttr<int64_t>(kernel_node, "ffty"));
|
||||
Nall = fftx * ffty;
|
||||
Nfft = fftx * (ffty / NUM_2 + 1);
|
||||
|
||||
cufftPlan2d(&FFT_plan_r2c, fftx, ffty, CUFFT_R2C);
|
||||
|
||||
InitSizeLists();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
|
||||
auto input_tensor = GetDeviceAddress<T>(inputs, 0);
|
||||
auto output_tensor = GetDeviceAddress<Complex<T>>(outputs, 0);
|
||||
|
||||
cufftSetStream(FFT_plan_r2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
PMERFFT2D<T>(Nfft, input_tensor, output_tensor, FFT_plan_r2c, reinterpret_cast<cudaStream_t>(stream_ptr));
|
||||
return true;
|
||||
}
|
||||
|
||||
protected:
|
||||
void InitSizeLists() override {
|
||||
input_size_list_.push_back(Nall * sizeof(T));
|
||||
output_size_list_.push_back(Nfft * sizeof(Complex<T>));
|
||||
}
|
||||
|
||||
private:
|
||||
int fftx;
|
||||
int ffty;
|
||||
int Nall;
|
||||
int Nfft;
|
||||
cufftHandle FFT_plan_r2c;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif
|
||||
|
|
@ -2989,6 +2989,254 @@ class TransferCrd(PrimitiveWithInfer):
|
|||
return mstype.float32, mstype.float32, mstype.float32, mstype.int32
|
||||
|
||||
|
||||
class PMEBatchedFFT2D(PrimitiveWithInfer):
|
||||
"""
|
||||
Forward FFT with N-Dimensional Input. currently this is only batched ifft2.
|
||||
|
||||
.. warning::
|
||||
This is an experimental prototype that is subject to change and/or deletion.
|
||||
|
||||
Inputs:
|
||||
- **input_tensor** (Tensor) - Three dimentsional tensor, supported
|
||||
data type is complex64.
|
||||
|
||||
Outputs:
|
||||
- **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
|
||||
transform, the data type is complex64.
|
||||
|
||||
Supported Platforms:
|
||||
``GPU``
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self, direction):
|
||||
self.init_prim_io_names(
|
||||
inputs=['input_tensor'],
|
||||
outputs=['output_tensor'])
|
||||
|
||||
validator.check_value_type('direction', direction, int, self.name)
|
||||
self.direction = direction
|
||||
self.add_prim_attr('direction', self.direction)
|
||||
|
||||
def infer_shape(self, input_shape):
|
||||
self.add_prim_attr("fftx", input_shape[0])
|
||||
self.add_prim_attr("ffty", input_shape[1])
|
||||
self.add_prim_attr("fftz", input_shape[2])
|
||||
return [input_shape[0], input_shape[1], input_shape[2]]
|
||||
|
||||
def infer_dtype(self, input_dtype):
|
||||
validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
|
||||
return mstype.complex64
|
||||
|
||||
|
||||
class PMEFFT1D(PrimitiveWithInfer):
|
||||
"""
|
||||
Forward FFT with One-Dimensional Input.
|
||||
|
||||
.. warning::
|
||||
This is an experimental prototype that is subject to change and/or deletion.
|
||||
|
||||
Inputs:
|
||||
- **input_tensor** (Tensor) - Three dimentsional tensor, supported
|
||||
data type is complex64.
|
||||
|
||||
Outputs:
|
||||
- **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
|
||||
transform, the data type is complex64.
|
||||
|
||||
Supported Platforms:
|
||||
``GPU``
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self):
|
||||
self.init_prim_io_names(
|
||||
inputs=['input_tensor'],
|
||||
outputs=['output_tensor'])
|
||||
|
||||
def infer_shape(self, input_shape):
|
||||
self.add_prim_attr('fftx', input_shape[0])
|
||||
return [input_shape[0]]
|
||||
|
||||
def infer_dtype(self, input_dtype):
|
||||
validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
|
||||
return mstype.complex64
|
||||
|
||||
|
||||
class PMEIFFT1D(PrimitiveWithInfer):
|
||||
"""
|
||||
Inverse FFT with One-Dimensional Input.
|
||||
|
||||
.. warning::
|
||||
This is an experimental prototype that is subject to change and/or deletion.
|
||||
|
||||
Inputs:
|
||||
- **input_tensor** (Tensor) - Three dimentsional input tensor, supported data
|
||||
type is complex64.
|
||||
|
||||
Outputs:
|
||||
- **output_tensor** (Tensor) - Returns the tensor after undergoing
|
||||
inverse Fourier transform, the data type is complex64.
|
||||
|
||||
Supported Platforms:
|
||||
``GPU``
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self):
|
||||
self.init_prim_io_names(
|
||||
inputs=['input_tensor'],
|
||||
outputs=['output_tensor'])
|
||||
|
||||
def infer_shape(self, input_shape):
|
||||
self.add_prim_attr('fftx', input_shape[0])
|
||||
return [input_shape[0]]
|
||||
|
||||
def infer_dtype(self, input_dtype):
|
||||
validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
|
||||
return mstype.complex64
|
||||
|
||||
|
||||
class PMEFFT2D(PrimitiveWithInfer):
|
||||
"""
|
||||
Forward FFT with Two-Dimensional Input.
|
||||
|
||||
.. warning::
|
||||
This is an experimental prototype that is subject to change and/or deletion.
|
||||
|
||||
Inputs:
|
||||
- **input_tensor** (Tensor) - Three dimentsional tensor, supported
|
||||
data type is complex64.
|
||||
|
||||
Outputs:
|
||||
- **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
|
||||
transform, the data type is complex64.
|
||||
|
||||
Supported Platforms:
|
||||
``GPU``
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self):
|
||||
self.init_prim_io_names(
|
||||
inputs=['input_tensor'],
|
||||
outputs=['output_tensor'])
|
||||
|
||||
def infer_shape(self, input_shape):
|
||||
self.add_prim_attr('fftx', input_shape[0])
|
||||
self.add_prim_attr('ffty', input_shape[1])
|
||||
return [input_shape[0], input_shape[1]]
|
||||
|
||||
def infer_dtype(self, input_dtype):
|
||||
validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
|
||||
return mstype.complex64
|
||||
|
||||
|
||||
class PMEIFFT2D(PrimitiveWithInfer):
|
||||
"""
|
||||
Inverse FFT with Two-Dimensional Input.
|
||||
|
||||
.. warning::
|
||||
This is an experimental prototype that is subject to change and/or deletion.
|
||||
|
||||
Inputs:
|
||||
- **input_tensor** (Tensor) - Three dimentsional input tensor, supported data
|
||||
type is complex64.
|
||||
|
||||
Outputs:
|
||||
- **output_tensor** (Tensor) - Return the tensor after undergoing
|
||||
inverse Fourier transform, the data type is complex64.
|
||||
|
||||
Supported Platforms:
|
||||
``GPU``
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self):
|
||||
self.init_prim_io_names(
|
||||
inputs=['input_tensor'],
|
||||
outputs=['output_tensor'])
|
||||
|
||||
def infer_shape(self, input_shape):
|
||||
self.add_prim_attr('fftx', input_shape[0])
|
||||
self.add_prim_attr('ffty', input_shape[1])
|
||||
return [input_shape[0], input_shape[1]]
|
||||
|
||||
def infer_dtype(self, input_dtype):
|
||||
validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
|
||||
return mstype.complex64
|
||||
|
||||
|
||||
class PMERFFT2D(PrimitiveWithInfer):
|
||||
"""
|
||||
Forward FFT with Two-Dimensional Input for real -> complex.
|
||||
|
||||
.. warning::
|
||||
This is an experimental prototype that is subject to change and/or deletion.
|
||||
|
||||
Inputs:
|
||||
- **input_tensor** (Tensor) - Three dimentsional tensor, supported
|
||||
data type is float32.
|
||||
|
||||
Outputs:
|
||||
- **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
|
||||
transform, the data type is complex64.
|
||||
|
||||
Supported Platforms:
|
||||
``GPU``
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self):
|
||||
self.init_prim_io_names(
|
||||
inputs=['input_tensor'],
|
||||
outputs=['output_tensor'])
|
||||
|
||||
def infer_shape(self, input_shape):
|
||||
self.add_prim_attr('fftx', input_shape[0])
|
||||
self.add_prim_attr('ffty', input_shape[1])
|
||||
return [input_shape[0], int(input_shape[1]/2)+1]
|
||||
|
||||
def infer_dtype(self, input_dtype):
|
||||
validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.float32], self.name)
|
||||
return mstype.complex64
|
||||
|
||||
|
||||
class PMEIRFFT2D(PrimitiveWithInfer):
|
||||
"""
|
||||
Inverse RFFT with Two-Dimensional Input.
|
||||
|
||||
.. warning::
|
||||
This is an experimental prototype that is subject to change and/or deletion.
|
||||
|
||||
Inputs:
|
||||
- **input_tensor** (Tensor) - Three dimentsional input tensor, supported data
|
||||
type is complex64.
|
||||
|
||||
Outputs:
|
||||
- **output_tensor** (Tensor) - Return the tensor after undergoing
|
||||
inverse Fourier transform, the data type is float32.
|
||||
|
||||
Supported Platforms:
|
||||
``GPU``
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self):
|
||||
self.init_prim_io_names(
|
||||
inputs=['input_tensor'],
|
||||
outputs=['output_tensor'])
|
||||
|
||||
def infer_shape(self, input_shape):
|
||||
self.add_prim_attr('fftx', input_shape[0])
|
||||
self.add_prim_attr('ffty', input_shape[1])
|
||||
return [input_shape[0], input_shape[1]]
|
||||
|
||||
def infer_dtype(self, input_dtype):
|
||||
validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
|
||||
return mstype.float32
|
||||
|
||||
|
||||
class FFT3D(PrimitiveWithInfer):
|
||||
"""
|
||||
Forward FFT with Three-Dimensional Input.
|
||||
|
|
|
|||
Loading…
Reference in New Issue