libtrilinos-tpetra-dev 12.10.1-3 / usr / include / trilinos / Kokkos_MV_GEMM.hpp

/*
//@HEADER
// ************************************************************************
//
//          Kokkos: Node API and Parallel Node Kernels
//              Copyright (2008) Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Michael A. Heroux (maherou@sandia.gov)
//
// ************************************************************************
//@HEADER
*/

#ifndef KOKKOS_MV_GEMM_HPP
#define KOKKOS_MV_GEMM_HPP

// Note this code lives only temporarily in TpetraCore.  As soon as
// GEMM kernels exist in the TpetraKernels subpackage, and thus a
// dependency on Teuchos can be eliminated, the code will move to
// TpetraKernels.

#include <Teuchos_BLAS.hpp>
#include <Kokkos_Blas2_MV.hpp>

#ifdef KOKKOS_HAVE_CUDA
#include <cublas.h>
#endif

namespace Teuchos {

  // mfh 11 Nov 2014: The DeviceGEMM specializations below need to be
  // able to use Teuchos::BLAS::{GEMM, GEMV}.  We provide just enough
  // of a specialization for Kokkos::complex<{float, double}> to make
  // DeviceGEMM work.  They just defer to BLAS<int,
  // std::complex<{float, double}> > via reinterpret_cast.  Please
  // feel free to expand these specializations if you need to.

  template<>
  class BLAS<int, ::Kokkos::complex<float> > {
  public:
    typedef float mag_type;
    typedef ::Kokkos::complex<float> val_type;
    typedef std::complex<float> impl_type;

    BLAS () {}
    BLAS (const BLAS<int, val_type>&) {}
    virtual ~BLAS () {}

    // void ROTG (val_type* da, val_type* db, mag_type* c, val_type* s) const;
    // void ROT (const int n, val_type* dx, const int incx, val_type* dy, const int incy, RealType* c, val_type* s) const;
    // RealType ASUM (const int n, const val_type* x, const int incx) const;
    //void AXPY (const int n, const val_type alpha, const val_type* x, const int incx, val_type* y, const int incy) const;
    //void COPY (const int n, const val_type* x, const int incx, val_type* y, const int incy) const;
    //val_type DOT(const int n, const val_type* x, const int incx, const val_type* y, const int incy) const;
    //RealType NRM2(const int n, const val_type* x, const int incx) const;
    //void SCAL(const int n, const val_type alpha, val_type* x, const int incx) const;
    //int IAMAX(const int n, const val_type* x, const int incx) const;

    void
    GEMV (ETransp trans, const int m, const int n, const val_type alpha,
          const val_type* A, const int lda, const val_type* x, const int incx,
          const val_type beta, val_type* y, const int incy) const
    {
      BLAS<int, impl_type> blas;
      blas.GEMV (trans, m, n, static_cast<impl_type> (alpha),
                 reinterpret_cast<const impl_type*> (A), lda,
                 reinterpret_cast<const impl_type*> (x), incx,
                 static_cast<impl_type> (beta),
                 reinterpret_cast<impl_type*> (y), incy);
    }

    //void TRMV(EUplo uplo, ETransp trans, EDiag diag, const int n, const val_type* A, const int lda, val_type* x, const int incx) const;
    //void GER(const int m, const int n, const val_type alpha, const val_type* x, const int incx, const val_type* y, const int incy, val_type* A, const int lda) const;

    void
    GEMM (ETransp transa, ETransp transb, const int m, const int n, const int k,
          const val_type alpha, const val_type* A, const int lda,
          const val_type* B, const int ldb, const val_type beta, val_type* C,
          const int ldc) const
    {
      BLAS<int, impl_type> blas;
      blas.GEMM (transa, transb, m, n, k,
                 static_cast<impl_type> (alpha),
                 reinterpret_cast<const impl_type*> (A), lda,
                 reinterpret_cast<const impl_type*> (B), ldb,
                 static_cast<impl_type> (beta),
                 reinterpret_cast<impl_type*> (C), ldc);
    }

    //void SYMM(ESide side, EUplo uplo, const int m, const int n, const val_type alpha, const val_type* A, const int lda, const val_type *B, const int ldb, const val_type beta, val_type *C, const int ldc) const;
    //void SYRK(EUplo uplo, ETransp trans, const int n, const int k, const val_type alpha, const val_type* A, const int lda, const val_type beta, val_type* C, const int ldc) const;
    //void TRMM(ESide side, EUplo uplo, ETransp transa, EDiag diag, const int m, const int n, const val_type alpha, const val_type* A, const int lda, val_type* B, const int ldb) const;
    //void TRSM(ESide side, EUplo uplo, ETransp transa, EDiag diag, const int m, const int n, const val_type alpha, const val_type* A, const int lda, val_type* B, const int ldb) const;
  };

  template<>
  class BLAS<int, ::Kokkos::complex<double> > {
  public:
    typedef double mag_type;
    typedef ::Kokkos::complex<double> val_type;
    typedef std::complex<double> impl_type;

    BLAS () {}
    BLAS (const BLAS<int, val_type>&) {}
    virtual ~BLAS () {}

    // void ROTG (val_type* da, val_type* db, mag_type* c, val_type* s) const;
    // void ROT (const int n, val_type* dx, const int incx, val_type* dy, const int incy, RealType* c, val_type* s) const;
    // RealType ASUM (const int n, const val_type* x, const int incx) const;
    //void AXPY (const int n, const val_type alpha, const val_type* x, const int incx, val_type* y, const int incy) const;
    //void COPY (const int n, const val_type* x, const int incx, val_type* y, const int incy) const;
    //val_type DOT(const int n, const val_type* x, const int incx, const val_type* y, const int incy) const;
    //RealType NRM2(const int n, const val_type* x, const int incx) const;
    //void SCAL(const int n, const val_type alpha, val_type* x, const int incx) const;
    //int IAMAX(const int n, const val_type* x, const int incx) const;

    void
    GEMV (ETransp trans, const int m, const int n, const val_type alpha,
          const val_type* A, const int lda, const val_type* x, const int incx,
          const val_type beta, val_type* y, const int incy) const
    {
      BLAS<int, impl_type> blas;
      blas.GEMV (trans, m, n, static_cast<impl_type> (alpha),
                 reinterpret_cast<const impl_type*> (A), lda,
                 reinterpret_cast<const impl_type*> (x), incx,
                 static_cast<impl_type> (beta),
                 reinterpret_cast<impl_type*> (y), incy);
    }

    //void TRMV(EUplo uplo, ETransp trans, EDiag diag, const int n, const val_type* A, const int lda, val_type* x, const int incx) const;
    //void GER(const int m, const int n, const val_type alpha, const val_type* x, const int incx, const val_type* y, const int incy, val_type* A, const int lda) const;

    void
    GEMM (ETransp transa, ETransp transb, const int m, const int n, const int k,
          const val_type alpha, const val_type* A, const int lda,
          const val_type* B, const int ldb, const val_type beta, val_type* C,
          const int ldc) const
    {
      BLAS<int, impl_type> blas;
      blas.GEMM (transa, transb, m, n, k,
                 static_cast<impl_type> (alpha),
                 reinterpret_cast<const impl_type*> (A), lda,
                 reinterpret_cast<const impl_type*> (B), ldb,
                 static_cast<impl_type> (beta),
                 reinterpret_cast<impl_type*> (C), ldc);
    }

    //void SYMM(ESide side, EUplo uplo, const int m, const int n, const val_type alpha, const val_type* A, const int lda, const val_type *B, const int ldb, const val_type beta, val_type *C, const int ldc) const;
    //void SYRK(EUplo uplo, ETransp trans, const int n, const int k, const val_type alpha, const val_type* A, const int lda, const val_type beta, val_type* C, const int ldc) const;
    //void TRMM(ESide side, EUplo uplo, ETransp transa, EDiag diag, const int m, const int n, const val_type alpha, const val_type* A, const int lda, val_type* B, const int ldb) const;
    //void TRSM(ESide side, EUplo uplo, ETransp transa, EDiag diag, const int m, const int n, const val_type alpha, const val_type* A, const int lda, val_type* B, const int ldb) const;
  };

} // namespace Teuchos


namespace Kokkos {
  namespace Impl {

    template<class ViewType>
    size_t getStride2DView (ViewType A) {
      size_t stride[8];
      A.stride (stride);
      return A.dimension_1 () > 1 ? stride[1] : A.dimension_0 ();
    }
  }

  /// \struct DeviceGEMM
  /// \brief Class that provides GEMM for a particular Kokkos Device.
  /// \tparam Scalar The type of the entries in the matrices.
  /// \tparam DeviceType The Kokkos Device type.
  ///
  /// GEMM refers to the BLAS' dense matrix-matrix multiply routine.
  template <typename Scalar, typename DeviceType>
  struct DeviceGEMM {
  public:
    static void
    GEMM (const Teuchos::ETransp transA,
          const Teuchos::ETransp transB,
          const Scalar& alpha,
          const View<const Scalar**, LayoutLeft, DeviceType>& A,
          const View<const Scalar**, LayoutLeft, DeviceType>& B,
          const Scalar& beta,
          const View<Scalar**, LayoutLeft, DeviceType>& C)
    {
      const int n = static_cast<int> (C.dimension_1 ());
      const int lda = static_cast<int> (Impl::getStride2DView (A));
      Teuchos::BLAS<int,Scalar> blas;

      // For some BLAS implementations (e.g., MKL), GEMM when B has
      // one column may be signficantly less efficient than GEMV.
      if (n == 1 && transB == Teuchos::NO_TRANS) {
        blas.GEMV (transA, A.dimension_0 (), A.dimension_1 (),
                   alpha, A.ptr_on_device (), lda,
                   B.ptr_on_device (), static_cast<int> (1),
                   beta, C.ptr_on_device (), static_cast<int> (1));
      }
      else {
        const int m = static_cast<int> (C.dimension_0 ());
        const int k = static_cast<int> (transA == Teuchos::NO_TRANS ?
                                        A.dimension_1 () : A.dimension_0 ());
        const int ldb = static_cast<int> (Impl::getStride2DView (B));
        const int ldc = static_cast<int> (Impl::getStride2DView (C));

        blas.GEMM (transA, transB, m, n, k, alpha,
                   A.ptr_on_device(), lda,
                   B.ptr_on_device(), ldb,
                   beta, C.ptr_on_device(), ldc);
      }
    }
  };

  // FIXME (mfh 10 May 2016) Temporary work-around for #243.
  // Don't call MKL for this case.
#ifdef HAVE_TPETRAKERNELS_MKL
  template <typename DeviceType>
  struct DeviceGEMM<double, DeviceType> {
  public:
    static void
    GEMM (const Teuchos::ETransp transA,
          const Teuchos::ETransp transB,
          const double& alpha,
          const View<const double**, LayoutLeft, DeviceType>& A,
          const View<const double**, LayoutLeft, DeviceType>& B,
          const double& beta,
          const View<double**, LayoutLeft, DeviceType>& C)
    {
      const int n = static_cast<int> (C.dimension_1 ());

      // For some BLAS implementations (e.g., MKL), GEMM when B has
      // one column may be signficantly less efficient than GEMV.
      if (n == 1 && transB == Teuchos::NO_TRANS) {
        char trans = 'N';
        if (transA == Teuchos::TRANS) {
          trans = 'T';
        }
        else if (transA == Teuchos::CONJ_TRANS) {
          trans = 'C';
        }
        auto B_0 = Kokkos::subview (B, Kokkos::ALL (), 0);
        auto C_0 = Kokkos::subview (C, Kokkos::ALL (), 0);
        KokkosBlas::gemv (&trans, alpha, A, B_0, beta, C_0);
      }
      else {
        const int m = static_cast<int> (C.dimension_0 ());
        const int k = static_cast<int> (transA == Teuchos::NO_TRANS ? A.dimension_1 () : A.dimension_0 ());
        const int lda = static_cast<int> (Impl::getStride2DView (A));
        const int ldb = static_cast<int> (Impl::getStride2DView (B));
        const int ldc = static_cast<int> (Impl::getStride2DView (C));

        Teuchos::BLAS<int,double> blas;
        blas.GEMM (transA, transB, m, n, k, alpha,
                   A.ptr_on_device(), lda,
                   B.ptr_on_device(), ldb,
                   beta, C.ptr_on_device(), ldc);
      }
    }
  };
#endif // HAVE_TPETRAKERNELS_MKL

#ifdef KOKKOS_HAVE_CUDA
  template <typename Scalar>
  struct DeviceGEMM<Scalar,Cuda> {
    public:
      static void
      GEMM (const Teuchos::ETransp transA,
            const Teuchos::ETransp transB,
            const Scalar& alpha,
            const View<const Scalar**, LayoutLeft, Cuda>& A,
            const View<const Scalar**,LayoutLeft,Cuda>& B,
            const Scalar& beta,
            const View<Scalar**,LayoutLeft,Cuda>& C)
    {
      TEUCHOS_TEST_FOR_EXCEPTION
        (true, std::logic_error, "DeviceGEMM: Kokkos::Cuda has no support "
         "for GEMM operations over Scalar=" << Teuchos::typeName(alpha) << ".");
      }
  };

  template <>
  struct DeviceGEMM<float,Cuda> {
    public:
      static void
      GEMM (const Teuchos::ETransp transA,
            const Teuchos::ETransp transB,
            const float alpha,
            const View<const float**,LayoutLeft,Cuda>& A,
            const View<const float**,LayoutLeft,Cuda>& B,
            const float beta,
            const View<float**,LayoutLeft,Cuda>& C)
    {
      const int m = static_cast<int>(C.dimension_0()),
        n = static_cast<int>(C.dimension_1()),
        k = (transA == Teuchos::NO_TRANS ? A.dimension_1() : A.dimension_0()),
        lda = static_cast<int>(Impl::getStride2DView(A)),
        ldb = static_cast<int>(Impl::getStride2DView(B)),
        ldc = static_cast<int>(Impl::getStride2DView(C));
      const char char_transA = (transA == Teuchos::NO_TRANS ? 'N' : 'T'),
        char_transB = (transB == Teuchos::NO_TRANS ? 'N' : 'T');
      cublasSgemm (char_transA, char_transB, m, n, k, alpha,
                   A.ptr_on_device(), lda, B.ptr_on_device(),
                   ldb, beta, C.ptr_on_device(), ldc);

#ifdef HAVE_KOKKOS_DEBUG
      const cublasStatus info = cublasGetError ();
      TEUCHOS_TEST_FOR_EXCEPTION
        (info != CUBLAS_STATUS_SUCCESS, std::runtime_error,
         "cublasSgemm failed with status " << info << "." );
#endif // HAVE_KOKKOS_DEBUG
      }
  };

  template <>
  struct DeviceGEMM<double,Cuda> {
    public:
      static void GEMM(Teuchos::ETransp transA, Teuchos::ETransp transB, double alpha,
          View<const double**,LayoutLeft,Cuda> A, View<const double**,LayoutLeft,Cuda> B,
          double beta, View<double**,LayoutLeft,Cuda> C) {
        const int m = static_cast<int>(C.dimension_0()),
                  n = static_cast<int>(C.dimension_1()),
                  k = (transA == Teuchos::NO_TRANS ? A.dimension_1() : A.dimension_0()),
                  lda = static_cast<int>(Impl::getStride2DView(A)),
                  ldb = static_cast<int>(Impl::getStride2DView(B)),
                  ldc = static_cast<int>(Impl::getStride2DView(C));
        const char char_transA = (transA == Teuchos::NO_TRANS ? 'N' : 'T'),
                   char_transB = (transB == Teuchos::NO_TRANS ? 'N' : 'T');
        cublasDgemm(char_transA, char_transB, m, n, k, alpha, A.ptr_on_device(), lda, B.ptr_on_device(), ldb, beta, C.ptr_on_device(), ldc);
#ifdef HAVE_KOKKOS_DEBUG
        cublasStatus info = cublasGetError();
        TEUCHOS_TEST_FOR_EXCEPTION( info != CUBLAS_STATUS_SUCCESS, std::runtime_error, "cublasDgemm failed with status " << info << "." );
#endif
      }
  };


#endif
}
#endif // KOKKOS_MV_GEMM_HPP