Add blas_nrm2

Author: Cstandardlib

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

@@ -3,6 +3,7 @@
 namespace container {
 namespace kernels {
 
+
 template <typename T>
 struct blas_copy<T, DEVICE_CPU> {
     void operator()(
@@ -16,6 +17,18 @@ struct blas_copy<T, DEVICE_CPU> {
     }
 };
 
+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()(
@@ -194,6 +207,11 @@ 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>;

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

@@ -19,6 +19,14 @@ struct blas_copy {
         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 {

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

@@ -22,6 +22,20 @@ 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()(
@@ -209,11 +223,18 @@ 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>;

source/source_base/module_container/ATen/kernels/rocm/blas.hip.cu

@@ -23,6 +23,19 @@ void destroyGpuBlasHandle() {
 }
 
 
+template <typename T>
+struct blas_nrm2<T, DEVICE_GPU> {
+    T operator()(
+        const int n,
+        const T *x,
+        const int incx)
+    {
+        T result;
+        hipBlasConnector::nrm2(hipblas_handle, n, x, incx, &result);
+        return result;
+    }
+};
+
 template <typename T>
 struct blas_dot<T, DEVICE_GPU> {
     void operator()(
@@ -196,6 +209,11 @@ struct blas_gemm_batched_strided<T, DEVICE_GPU> {
 };
 
 // Explicitly instantiate functors for the types of functor registered.
+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>;

source/source_base/module_container/ATen/kernels/test/blas_test.cpp

@@ -36,6 +36,22 @@ TYPED_TEST(BlasTest, Copy) {
     EXPECT_EQ(y, expected);
 }
 
+TYPED_TEST(BlasTest, Nrm2) {
+    using Type = typename std::tuple_element<0, decltype(TypeParam())>::type;
+    using Device = typename std::tuple_element<1, decltype(TypeParam())>::type;
+
+    blas_nrm2<Type, Device> nrm2Calculator;
+
+    const int n = 3;
+    const Tensor x = std::move(Tensor({static_cast<Type>(3.0), static_cast<Type>(4.0), static_cast<Type>(0.0)}).to_device<Device>());
+
+    Type result = {};
+    nrm2Calculator(n, x.data<Type>(), 1, &result);
+    const Type expected = static_cast<Type>(5.0);
+
+    EXPECT_NEAR(result, expected, static_cast<Type>(1e-6));
+}
+
 TYPED_TEST(BlasTest, Dot) {
     using Type = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using Device = typename std::tuple_element<1, decltype(TypeParam())>::type;

source/source_base/module_container/base/third_party/blas.h

@@ -43,6 +43,7 @@ double ddot_(const int *N, const double *x, const int *incx, const double *y, co
 // Peize Lin add ?nrm2 2018-06-12, to compute out = ||x||_2 = \sqrt{ \sum_i x_i**2 }
 float snrm2_( const int *n, const float *x, const int *incx );
 double dnrm2_( const int *n, const double *x, const int *incx );
+double scnrm2_( const int *n, const std::complex<float> *x, const int *incx );
 double dznrm2_( const int *n, const std::complex<double> *x, const int *incx );
 
 // level 2: matrix-std::vector operations, O(n^2) data and O(n^2) work.
@@ -334,6 +335,11 @@ double nrm2( const int n, const double *x, const int incx )
     return dnrm2_( &n, x, &incx );
 }
 static inline
+double nrm2( const int n, const std::complex<float> *x, const int incx )
+{
+    return scnrm2_( &n, x, &incx );
+}
+static inline
 double nrm2( const int n, const std::complex<double> *x, const int incx )
 {
     return dznrm2_( &n, x, &incx );

source/source_base/module_container/base/third_party/cublas.h

@@ -29,6 +29,27 @@ void copy(cublasHandle_t& handle, const int& n, const std::complex<double> *x, c
     cublasErrcheck(cublasZcopy(handle, n, reinterpret_cast<const cuDoubleComplex*>(x), incx, reinterpret_cast<cuDoubleComplex*>(y), incy));
 }
 
+static inline
+void nrm2(cublasHandle_t& handle, const int& n, const float *x, const int& incx, float* result)
+{
+    cublasErrcheck(cublasSnrm2(handle, n, x, incx, result));
+}
+static inline
+void nrm2(cublasHandle_t& handle, const int& n, const double *x, const int& incx, double* result)
+{
+    cublasErrcheck(cublasDnrm2(handle, n, x, incx, result));
+}
+static inline
+void nrm2(cublasHandle_t& handle, const int& n, const std::complex<float> *x, const int& incx, float* result)
+{
+    cublasErrcheck(cublasScnrm2(handle, n, reinterpret_cast<const cuComplex*>(x), incx, result));
+}
+static inline
+void nrm2(cublasHandle_t& handle, const int& n, const std::complex<double> *x, const int& incx, double* result)
+{
+    cublasErrcheck(cublasDznrm2(handle, n, reinterpret_cast<const cuDoubleComplex*>(x), incx, result));
+}
+
 static inline
 void dot(cublasHandle_t& handle, const int& n, const float *x, const int& incx, const float *y, const int& incy, float* result)
 {