Add new blas/lapack interfaces

source/source_base/module_container/ATen/kernels/blas.cpp

@@ -3,6 +3,32 @@
 namespace container {
 namespace kernels {
 
+
+template <typename T>
+struct blas_copy<T, DEVICE_CPU> {
+    void operator()(
+        const int n,
+        const T *x,
+        const int incx,
+        T *y,
+        const int incy)
+    {
+        BlasConnector::copy(n, x, incx, y, incy);
+    }
+};
+
+template <typename T>
+struct blas_nrm2<T, DEVICE_CPU> {
+    using Real = typename GetTypeReal<T>::type;
+    Real operator()(
+        const int n,
+        const T *x,
+        const int incx)
+    {
+        return BlasConnector::nrm2(n, x, incx);
+    }
+};
+
 template <typename T>
 struct blas_dot<T, DEVICE_CPU> {
     void operator()(
@@ -175,6 +201,17 @@ struct blas_gemm_batched_strided<T, DEVICE_CPU> {
 };
 
 // Explicitly instantiate functors for the types of functor registered.
+
+template struct blas_copy<float , DEVICE_CPU>;
+template struct blas_copy<double, DEVICE_CPU>;
+template struct blas_copy<std::complex<float >, DEVICE_CPU>;
+template struct blas_copy<std::complex<double>, DEVICE_CPU>;
+
+template struct blas_nrm2<float , DEVICE_CPU>;
+template struct blas_nrm2<double, DEVICE_CPU>;
+template struct blas_nrm2<std::complex<float >, DEVICE_CPU>;
+template struct blas_nrm2<std::complex<double>, DEVICE_CPU>;
+
 template struct blas_dot<float , DEVICE_CPU>;
 template struct blas_dot<double, DEVICE_CPU>;
 template struct blas_dot<std::complex<float >, DEVICE_CPU>;
@@ -221,4 +258,4 @@ template struct blas_gemm_batched_strided<std::complex<float >, DEVICE_CPU>;
 template struct blas_gemm_batched_strided<std::complex<double>, DEVICE_CPU>;
 
 } // namespace kernels
-} // namespace container
+} // namespace container

source/source_base/module_container/ATen/kernels/blas.h

@@ -9,6 +9,26 @@
 namespace container {
 namespace kernels {
 
+template <typename T, typename Device>
+struct blas_copy {
+    // DCOPY copies a vector, x, to a vector, y.
+    void operator()(
+        const int n,
+        const T *x,
+        const int incx,
+        T *y,
+        const int incy);
+};
+
+template <typename T, typename Device>
+struct blas_nrm2 {
+    using Real = typename GetTypeReal<T>::type;
+    Real operator()(
+        const int n,
+        const T *x,
+        const int incx);
+};
+
 template <typename T, typename Device>
 struct blas_dot {
     void operator()(
@@ -168,4 +188,4 @@ void destroyGpuBlasHandle(); // destory blas handle
 } // namespace kernels
 } // namespace container
 
-#endif // ATEN_KERNELS_BLAS_H_
+#endif // ATEN_KERNELS_BLAS_H_

source/source_base/module_container/ATen/kernels/cuda/blas.cu

@@ -22,6 +22,32 @@ void destroyGpuBlasHandle() {
     }
 }
 
+template <typename T>
+struct blas_nrm2<T, DEVICE_GPU> {
+    using Real = typename GetTypeReal<T>::type;
+    Real operator()(
+        const int n,
+        const T *x,
+        const int incx)
+    {
+        Real result;
+        cuBlasConnector::nrm2(cublas_handle, n, x, incx, &result);
+        return result;
+    }
+};
+
+template <typename T>
+struct blas_copy<T, DEVICE_GPU> {
+    void operator()(
+        const int n,
+        const T * x,
+        const int incx,
+        T *y,
+        const int incy)
+    {
+        cuBlasConnector::copy(cublas_handle, n, x, incx, y, incy);
+    }
+};
 
 template <typename T>
 struct blas_dot<T, DEVICE_GPU> {
@@ -76,7 +102,7 @@ struct blas_gemv<T, DEVICE_GPU> {
         const int& incx,
         const T* beta,
         T* y,
-        const int& incy) 
+        const int& incy)
     {
         cuBlasConnector::gemv(cublas_handle, trans, m, n, *alpha, A, lda, x, incx, *beta, y, incy);
     }
@@ -196,6 +222,19 @@ struct blas_gemm_batched_strided<T, DEVICE_GPU> {
 };
 
 // Explicitly instantiate functors for the types of functor registered.
+
+
+
+template struct blas_copy<float , DEVICE_GPU>;
+template struct blas_copy<double, DEVICE_GPU>;
+template struct blas_copy<std::complex<float> , DEVICE_GPU>;
+template struct blas_copy<std::complex<double>, DEVICE_GPU>;
+
+template struct blas_nrm2<float , DEVICE_GPU>;
+template struct blas_nrm2<double, DEVICE_GPU>;
+template struct blas_nrm2<std::complex<float> , DEVICE_GPU>;
+template struct blas_nrm2<std::complex<double>, DEVICE_GPU>;
+
 template struct blas_dot<float , DEVICE_GPU>;
 template struct blas_dot<double, DEVICE_GPU>;
 template struct blas_dot<std::complex<float> , DEVICE_GPU>;
@@ -242,4 +281,4 @@ template struct blas_gemm_batched_strided<std::complex<float >, DEVICE_GPU>;
 template struct blas_gemm_batched_strided<std::complex<double>, DEVICE_GPU>;
 
 } // namespace kernels
-} // namespace container
+} // namespace container