Add MKL and CUFFT helpers

ssnl · ssnl · commit fb3a0256739f · 2018-03-19T14:03:39.000-04:00
diff --git a/aten/src/ATen/mkl/Descriptors.h b/aten/src/ATen/mkl/Descriptors.h
@@ -0,0 +1,44 @@
+#pragma once
+
+#include "Exceptions.h"
+#include <mkl_dfti.h>
+#include <ATen/Tensor.h>
+
+namespace at { namespace native {
+
+struct DftiDescriptorDeleter {
+  void operator()(DFTI_DESCRIPTOR* desc) {
+    if (desc != nullptr) {
+      MKL_DFTI_CHECK(DftiFreeDescriptor(&desc));
+    }
+  }
+};
+
+class DftiDescriptor {
+public:
+  void init(DFTI_CONFIG_VALUE precision, DFTI_CONFIG_VALUE signal_type, MKL_LONG signal_ndim, MKL_LONG* sizes) {
+    if (desc_ != nullptr) {
+      throw std::runtime_error("DFTI DESCRIPTOR can only be initialized once");
+    }
+    DFTI_DESCRIPTOR *raw_desc;
+    if (signal_ndim == 1) {
+      MKL_DFTI_CHECK(DftiCreateDescriptor(&raw_desc, precision, signal_type, 1, sizes[0]));
+    } else {
+      MKL_DFTI_CHECK(DftiCreateDescriptor(&raw_desc, precision, signal_type, signal_ndim, sizes));
+    }
+    desc_.reset(raw_desc);
+  }
+
+  DFTI_DESCRIPTOR *get() const {
+    if (desc_ == nullptr) {
+      throw std::runtime_error("DFTI DESCRIPTOR has not been initialized");
+    }
+    return desc_.get();
+  }
+
+private:
+  std::unique_ptr<DFTI_DESCRIPTOR, DftiDescriptorDeleter> desc_;
+};
+
+
+}}  // at::native
diff --git a/aten/src/ATen/mkl/Exceptions.h b/aten/src/ATen/mkl/Exceptions.h
@@ -0,0 +1,19 @@
+#pragma once
+
+#include <string>
+#include <stdexcept>
+#include <sstream>
+#include <mkl_dfti.h>
+
+namespace at { namespace native {
+
+static inline void MKL_DFTI_CHECK(MKL_INT status)
+{
+  if (status && !DftiErrorClass(status, DFTI_NO_ERROR)) {
+    std::ostringstream ss;
+    ss << "MKL FFT error: " << DftiErrorMessage(status);
+    throw std::runtime_error(ss.str());
+  }
+}
+
+}}  // namespace at::native
diff --git a/aten/src/ATen/mkl/Limits.h b/aten/src/ATen/mkl/Limits.h
@@ -0,0 +1,11 @@
+#pragma once
+
+#include <mkl_types.h>
+
+namespace at { namespace native {
+
+  // Since size of MKL_LONG varies on different platforms (linux 64 bit, windows
+  // 32 bit), we need to programmatically calculate the max.
+  static int64_t MKL_LONG_MAX = ((1LL << (sizeof(MKL_LONG) * 8 - 2)) - 1) * 2 + 1;
+
+}}  // namespace
diff --git a/aten/src/ATen/mkl/README.md b/aten/src/ATen/mkl/README.md
@@ -0,0 +1,4 @@
+All files living in this directory are written with the assumption that MKL is available,
+which means that these code are not guarded by `#if AT_MKL_ENABLED()`. Therefore, whenever
+you need to use definitions from here, please guard the `#include<ATen/mkl/*.h>` and
+definition usages with `#if AT_MKL_ENABLED()` macro, e.g. [SpectralOps.cpp](native/mkl/SpectralOps.cpp).
diff --git a/aten/src/ATen/native/cuda/CuFFTUtils.h b/aten/src/ATen/native/cuda/CuFFTUtils.h
@@ -0,0 +1,84 @@
+#pragma once
+
+#include "ATen/ATen.h"
+#include "ATen/Config.h"
+
+#include <string>
+#include <stdexcept>
+#include <sstream>
+#include <cufft.h>
+#include <cufftXt.h>
+
+
+namespace at { namespace native {
+
+static inline std::string _cudaGetErrorEnum(cufftResult error)
+{
+  switch (error)
+  {
+    case CUFFT_SUCCESS:
+      return "CUFFT_SUCCESS";
+    case CUFFT_INVALID_PLAN:
+      return "CUFFT_INVALID_PLAN";
+    case CUFFT_ALLOC_FAILED:
+      return "CUFFT_ALLOC_FAILED";
+    case CUFFT_INVALID_TYPE:
+      return "CUFFT_INVALID_TYPE";
+    case CUFFT_INVALID_VALUE:
+      return "CUFFT_INVALID_VALUE";
+    case CUFFT_INTERNAL_ERROR:
+      return "CUFFT_INTERNAL_ERROR";
+    case CUFFT_EXEC_FAILED:
+      return "CUFFT_EXEC_FAILED";
+    case CUFFT_SETUP_FAILED:
+      return "CUFFT_SETUP_FAILED";
+    case CUFFT_INVALID_SIZE:
+      return "CUFFT_INVALID_SIZE";
+    case CUFFT_UNALIGNED_DATA:
+      return "CUFFT_UNALIGNED_DATA";
+    case CUFFT_INCOMPLETE_PARAMETER_LIST:
+      return "CUFFT_INCOMPLETE_PARAMETER_LIST";
+    case CUFFT_INVALID_DEVICE:
+      return "CUFFT_INVALID_DEVICE";
+    case CUFFT_PARSE_ERROR:
+      return "CUFFT_PARSE_ERROR";
+    case CUFFT_NO_WORKSPACE:
+      return "CUFFT_NO_WORKSPACE";
+    case CUFFT_NOT_IMPLEMENTED:
+      return "CUFFT_NOT_IMPLEMENTED";
+    case CUFFT_LICENSE_ERROR:
+      return "CUFFT_LICENSE_ERROR";
+    case CUFFT_NOT_SUPPORTED:
+      return "CUFFT_NOT_SUPPORTED";
+    default:
+      std::ostringstream ss;
+      ss << "unknown error " << error;
+      return ss.str();
+  }
+}
+
+static inline void CUFFT_CHECK(cufftResult error)
+{
+  if (error != CUFFT_SUCCESS) {
+    std::ostringstream ss;
+    ss << "cuFFT error: " << _cudaGetErrorEnum(error);
+    throw std::runtime_error(ss.str());
+  }
+}
+
+class CufftHandle {
+public:
+  explicit CufftHandle() {
+    CUFFT_CHECK(cufftCreate(&raw_plan));
+  }
+
+  const cufftHandle &get() const { return raw_plan; }
+
+  ~CufftHandle() {
+    CUFFT_CHECK(cufftDestroy(raw_plan));
+  }
+private:
+  cufftHandle raw_plan;
+};
+
+}} // at::native
