Merge pull request #37 from iotamudelta/rocfft_improvements

Enable FFT operations on the GPU for ROCm

Author: whchung

tensorflow/stream_executor/rocm/rocm_fft.cc

@@ -70,14 +70,12 @@ namespace wrap {
   __macro(hipfftMakePlan3d)          \
   __macro(hipfftGetSizeMany)         \
   __macro(hipfftMakePlanMany)        \
-
-  // ROCM TODO disable before rocFFT uses proper HIP complex types
-  //__macro(hipfftExecD2Z)             \
-  //__macro(hipfftExecZ2D)             \
-  //__macro(hipfftExecC2C)             \
-  //__macro(hipfftExecC2R)             \
-  //__macro(hipfftExecZ2Z)             \
-  //__macro(hipfftExecR2C)             \
+  __macro(hipfftExecD2Z)             \
+  __macro(hipfftExecZ2D)             \
+  __macro(hipfftExecC2C)             \
+  __macro(hipfftExecC2R)             \
+  __macro(hipfftExecZ2Z)             \
+  __macro(hipfftExecR2C)             \
 
 HIPFFT_ROUTINE_EACH(PERFTOOLS_GPUTOOLS_HIPFFT_WRAP)
 
@@ -536,32 +534,20 @@ bool ROCMFft::DoFftWithDirectionInternal(Stream *stream, fft::Plan *plan,
   bool ROCMFft::DoFft(Stream *stream, fft::Plan *plan,                       \
                       const DeviceMemory<std::complex<__type>> &input,       \
                       DeviceMemory<std::complex<__type>> *output) {          \
-    LOG(ERROR) << "rocFFT does not current support complex<float> "          \
-	       << "/ complex<double> datatypes" ;                            \
-    /* ROCM TODO disable for now until rocFFT properly honors HIP complex types */ \
-    /* return DoFftWithDirectionInternal(                                 */ \
-    /*     stream, plan, wrap::hipfftExec##__fft_type1, input, output);   */ \
-    return false;                                                            \
+    return DoFftWithDirectionInternal(                                       \
+         stream, plan, wrap::hipfftExec##__fft_type1, input, output);        \
   }                                                                          \
   bool ROCMFft::DoFft(Stream *stream, fft::Plan *plan,                       \
                       const DeviceMemory<__type> &input,                     \
                       DeviceMemory<std::complex<__type>> *output) {          \
-    LOG(ERROR) << "rocFFT does not current support complex<float> "          \
-	       << "/ complex<double> datatypes" ;                            \
-    /* ROCM TODO disable for now until rocFFT properly honors HIP complex types */ \
-    /* return DoFftInternal(stream, plan, wrap::hipfftExec##__fft_type2, input, */  \
-    /*                     output);                                       */ \
-    return false;                                                            \
+    return DoFftInternal(stream, plan, wrap::hipfftExec##__fft_type2, input, \
+                         output);                                            \
   }                                                                          \
   bool ROCMFft::DoFft(Stream *stream, fft::Plan *plan,                       \
                       const DeviceMemory<std::complex<__type>> &input,       \
                       DeviceMemory<__type> *output) {                        \
-    LOG(ERROR) << "rocFFT does not current support complex<float> "          \
-	       << "/ complex<double> datatypes" ;                            \
-    /* ROCM TODO disable for now until rocFFT properly honors HIP complex types */ \
-    /* return DoFftInternal(stream, plan, wrap::hipfftExec##__fft_type3, input, */ \
-    /*                     output);                                       */ \
-    return false;                                                            \
+    return DoFftInternal(stream, plan, wrap::hipfftExec##__fft_type3, input, \
+                         output);                                            \
   }
 
 PERFTOOLS_GPUTOOLS_ROCM_DEFINE_FFT(float, C2C, R2C, C2R)

tensorflow/stream_executor/rocm/rocm_fft.h

@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// ROCM-specific support for FFT functionality -- this wraps the hipFFT library
+// ROCM-specific support for FFT functionality -- this wraps the rocFFT library
 // capabilities, and is only included into ROCM implementation code -- it will
 // not introduce rocm headers into other code.
 
@@ -86,7 +86,7 @@ class ROCMFftPlan : public fft::Plan {
   bool is_initialized_;
 };
 
-// FFT support for ROCM platform via hipFFT library.
+// FFT support for ROCM platform via rocFFT library.
 //
 // This satisfies the platform-agnostic FftSupport interface.
 //