Tpetra_BlockView.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
 
#ifndef TPETRA_BLOCKVIEW_HPP
#define TPETRA_BLOCKVIEW_HPP

 
#include "TpetraCore_config.h"
#include "KokkosKernels_ArithTraits.hpp"
#include "Kokkos_Complex.hpp"
 
namespace Tpetra {

 
namespace Impl {

template <class ViewType1,
          class ViewType2,
          const int rank1 = ViewType1::rank>
struct AbsMax {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType2& Y, const ViewType1& X);
};

template <class ViewType1,
          class ViewType2>
struct AbsMax<ViewType1, ViewType2, 2> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType2& Y, const ViewType1& X) {
    static_assert(static_cast<int>(ViewType1::rank) == static_cast<int>(ViewType2::rank),
                  "AbsMax: ViewType1 and ViewType2 must have the same rank.");
    typedef typename std::remove_reference<decltype(Y(0, 0))>::type STY;
    static_assert(!std::is_const<STY>::value,
                  "AbsMax: The type of each entry of Y must be nonconst.");
    typedef typename std::decay<decltype(X(0, 0))>::type STX;
    static_assert(std::is_same<STX, STY>::value,
                  "AbsMax: The type of each entry of X and Y must be the same.");
    typedef KokkosKernels::ArithTraits<STY> KAT;
 
    const int numCols = Y.extent(1);
    const int numRows = Y.extent(0);
    for (int j = 0; j < numCols; ++j) {
      for (int i = 0; i < numRows; ++i) {
        STY& Y_ij      = Y(i, j);  // use ref here to avoid 2nd op() call on Y
        const STX X_ij = X(i, j);
        // NOTE: no std::max (not a CUDA __device__ function); must
        // cast back up to complex.
        const auto Y_ij_abs = KAT::abs(Y_ij);
        const auto X_ij_abs = KAT::abs(X_ij);
        Y_ij                = (Y_ij_abs >= X_ij_abs) ? static_cast<STY>(Y_ij_abs) : static_cast<STY>(X_ij_abs);
      }
    }
  }
};

template <class ViewType1,
          class ViewType2>
struct AbsMax<ViewType1, ViewType2, 1> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType2& Y, const ViewType1& X) {
    static_assert(static_cast<int>(ViewType1::rank) == static_cast<int>(ViewType2::rank),
                  "AbsMax: ViewType1 and ViewType2 must have the same rank.");
 
    typedef typename std::remove_reference<decltype(Y(0))>::type STY;
    static_assert(!std::is_const<STY>::value,
                  "AbsMax: The type of each entry of Y must be nonconst.");
    typedef typename std::remove_const<typename std::remove_reference<decltype(X(0))>::type>::type STX;
    static_assert(std::is_same<STX, STY>::value,
                  "AbsMax: The type of each entry of X and Y must be the same.");
    typedef KokkosKernels::ArithTraits<STY> KAT;
 
    const int numRows = Y.extent(0);
    for (int i = 0; i < numRows; ++i) {
      STY& Y_i      = Y(i);  // use ref here to avoid 2nd op() call on Y
      const STX X_i = X(i);
      // NOTE: no std::max (not a CUDA __device__ function); must
      // cast back up to complex.
      const auto Y_i_abs = KAT::abs(Y_i);
      const auto X_i_abs = KAT::abs(X_i);
      Y_i                = (Y_i_abs >= X_i_abs) ? static_cast<STY>(Y_i_abs) : static_cast<STY>(X_i_abs);
    }
  }
};

template <class ViewType1, class ViewType2, const int rank = ViewType1::rank>
KOKKOS_INLINE_FUNCTION void
absMax(const ViewType2& Y, const ViewType1& X) {
  static_assert(static_cast<int>(ViewType1::rank) == static_cast<int>(ViewType2::rank),
                "absMax: ViewType1 and ViewType2 must have the same rank.");
  AbsMax<ViewType1, ViewType2, rank>::run(Y, X);
}

template <class ViewType,
          class CoefficientType,
          class IndexType          = int,
          const bool is_contiguous = false,
          const int rank           = ViewType::rank>
struct SCAL {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha, const ViewType& x);
};

template <class ViewType,
          class CoefficientType,
          class IndexType>
struct SCAL<ViewType, CoefficientType, IndexType, false, 1> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha, const ViewType& x) {
    const IndexType numRows = static_cast<IndexType>(x.extent(0));

#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
    for (IndexType i = 0; i < numRows; ++i)
      x(i) = alpha * x(i);
  }
};

template <class ViewType,
          class CoefficientType,
          class IndexType>
struct SCAL<ViewType, CoefficientType, IndexType, false, 2> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha, const ViewType& A) {
    const IndexType numRows = static_cast<IndexType>(A.extent(0));
    const IndexType numCols = static_cast<IndexType>(A.extent(1));
 
    for (IndexType j = 0; j < numCols; ++j)
      for (IndexType i = 0; i < numRows; ++i)
        A(i, j) = alpha * A(i, j);
  }
};
template <class ViewType,
          class CoefficientType,
          class IndexType,
          const int rank>
struct SCAL<ViewType, CoefficientType, IndexType, true, rank> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha, const ViewType& x) {
    using x_value_type   = typename std::decay<decltype(*x.data())>::type;
    const IndexType span = static_cast<IndexType>(x.span());
    x_value_type* __restrict x_ptr(x.data());
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
    for (IndexType i = 0; i < span; ++i)
      x_ptr[i] = alpha * x_ptr[i];
  }
};

template <class ViewType,
          class InputType,
          class IndexType          = int,
          const bool is_contiguous = false,
          const int rank           = ViewType::rank>
struct FILL {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType& x, const InputType& val);
};

template <class ViewType,
          class InputType,
          class IndexType>
struct FILL<ViewType, InputType, IndexType, false, 1> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType& x, const InputType& val) {
    const IndexType numRows = static_cast<IndexType>(x.extent(0));
    for (IndexType i = 0; i < numRows; ++i)
      x(i) = val;
  }
};

template <class ViewType,
          class InputType,
          class IndexType>
struct FILL<ViewType, InputType, IndexType, false, 2> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType& X, const InputType& val) {
    const IndexType numRows = static_cast<IndexType>(X.extent(0));
    const IndexType numCols = static_cast<IndexType>(X.extent(1));
    for (IndexType j = 0; j < numCols; ++j)
      for (IndexType i = 0; i < numRows; ++i)
        X(i, j) = val;
  }
};
template <class ViewType,
          class InputType,
          class IndexType,
          const int rank>
struct FILL<ViewType, InputType, IndexType, true, rank> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType& x, const InputType& val) {
    const IndexType span = static_cast<IndexType>(x.span());
    auto x_ptr           = x.data();
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
    for (IndexType i = 0; i < span; ++i)
      x_ptr[i] = val;
  }
};

template <class CoefficientType,
          class ViewType1,
          class ViewType2,
          class IndexType          = int,
          const bool is_contiguous = false,
          const int rank           = ViewType1::rank>
struct AXPY {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha,
      const ViewType1& x,
      const ViewType2& y);
};

template <class CoefficientType,
          class ViewType1,
          class ViewType2,
          class IndexType>
struct AXPY<CoefficientType, ViewType1, ViewType2, IndexType, false, 1> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha,
      const ViewType1& x,
      const ViewType2& y) {
    using KokkosKernels::ArithTraits;
    static_assert(static_cast<int>(ViewType1::rank) == static_cast<int>(ViewType2::rank),
                  "AXPY: x and y must have the same rank.");
 
    const IndexType numRows = static_cast<IndexType>(y.extent(0));
    if (alpha != ArithTraits<CoefficientType>::zero()) {
      for (IndexType i = 0; i < numRows; ++i)
        y(i) += alpha * x(i);
    }
  }
};

template <class CoefficientType,
          class ViewType1,
          class ViewType2,
          class IndexType>
struct AXPY<CoefficientType, ViewType1, ViewType2, IndexType, false, 2> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha,
      const ViewType1& X,
      const ViewType2& Y) {
    using KokkosKernels::ArithTraits;
    static_assert(static_cast<int>(ViewType1::rank) == static_cast<int>(ViewType2::rank),
                  "AXPY: X and Y must have the same rank.");
    const IndexType numRows = static_cast<IndexType>(Y.extent(0));
    const IndexType numCols = static_cast<IndexType>(Y.extent(1));
 
    if (alpha != ArithTraits<CoefficientType>::zero()) {
      for (IndexType j = 0; j < numCols; ++j)
        for (IndexType i = 0; i < numRows; ++i)
          Y(i, j) += alpha * X(i, j);
    }
  }
};
 
template <class CoefficientType,
          class ViewType1,
          class ViewType2,
          class IndexType,
          const int rank>
struct AXPY<CoefficientType, ViewType1, ViewType2, IndexType, true, rank> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoefficientType& alpha,
      const ViewType1& x,
      const ViewType2& y) {
    using KokkosKernels::ArithTraits;
    static_assert(static_cast<int>(ViewType1::rank) == static_cast<int>(ViewType2::rank),
                  "AXPY: x and y must have the same rank.");
 
    if (alpha != ArithTraits<CoefficientType>::zero()) {
      using x_value_type   = typename std::decay<decltype(*x.data())>::type;
      using y_value_type   = typename std::decay<decltype(*y.data())>::type;
      const IndexType span = static_cast<IndexType>(y.span());
      const x_value_type* __restrict x_ptr(x.data());
      y_value_type* __restrict y_ptr(y.data());
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
      for (IndexType i = 0; i < span; ++i)
        y_ptr[i] += alpha * x_ptr[i];
    }
  }
};

template <class ViewType1,
          class ViewType2,
          class IndexType          = int,
          const bool is_contiguous = false,
          const int rank           = ViewType1::rank>
struct COPY {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType1& x, const ViewType2& y);
};

template <class ViewType1,
          class ViewType2,
          class IndexType>
struct COPY<ViewType1, ViewType2, IndexType, false, 1> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType1& x, const ViewType2& y) {
    const IndexType numRows = static_cast<IndexType>(x.extent(0));
    for (IndexType i = 0; i < numRows; ++i)
      y(i) = x(i);
  }
};

template <class ViewType1,
          class ViewType2,
          class IndexType>
struct COPY<ViewType1, ViewType2, IndexType, false, 2> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType1& X, const ViewType2& Y) {
    const IndexType numRows = static_cast<IndexType>(Y.extent(0));
    const IndexType numCols = static_cast<IndexType>(Y.extent(1));
    for (IndexType j = 0; j < numCols; ++j)
      for (IndexType i = 0; i < numRows; ++i)
        Y(i, j) = X(i, j);
  }
};
 
template <class ViewType1,
          class ViewType2,
          class IndexType,
          const int rank>
struct COPY<ViewType1, ViewType2, IndexType, true, rank> {
  static KOKKOS_INLINE_FUNCTION void
  run(const ViewType1& x, const ViewType2& y) {
    const IndexType span = static_cast<IndexType>(x.span());
    using x_value_type   = typename std::decay<decltype(*x.data())>::type;
    using y_value_type   = typename std::decay<decltype(*y.data())>::type;
    const x_value_type* __restrict x_ptr(x.data());
    y_value_type* __restrict y_ptr(y.data());
 
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
    for (IndexType i = 0; i < span; ++i)
      y_ptr[i] = x_ptr[i];
  }
};
 
template <class VecType1,
          class BlkType,
          class VecType2,
          class CoeffType,
          class IndexType     = int,
          bool is_contiguous  = false,
          class BlkLayoutType = typename BlkType::array_layout>
struct GEMV {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoeffType& alpha,
      const BlkType& A,
      const VecType1& x,
      const VecType2& y) {
    static_assert(VecType1::rank == 1, "GEMV: VecType1 must have rank 1.");
    static_assert(BlkType::rank == 2, "GEMV: BlkType must have rank 2.");
    static_assert(VecType2::rank == 1, "GEMV: VecType2 must have rank 1.");
 
    const IndexType numRows = static_cast<IndexType>(A.extent(0));
    const IndexType numCols = static_cast<IndexType>(A.extent(1));

    for (IndexType i = 0; i < numRows; ++i)
      for (IndexType j = 0; j < numCols; ++j)
        y(i) += alpha * A(i, j) * x(j);
  }
};
 
template <class VecType1,
          class BlkType,
          class VecType2,
          class CoeffType,
          class IndexType>
struct GEMV<VecType1, BlkType, VecType2, CoeffType, IndexType, true, Kokkos::LayoutLeft> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoeffType& alpha,
      const BlkType& A,
      const VecType1& x,
      const VecType2& y) {
    static_assert(VecType1::rank == 1, "GEMV: VecType1 must have rank 1.");
    static_assert(BlkType::rank == 2, "GEMV: BlkType must have rank 2.");
    static_assert(VecType2::rank == 1, "GEMV: VecType2 must have rank 1.");
 
    using A_value_type = typename std::decay<decltype(A(0, 0))>::type;
    using x_value_type = typename std::decay<decltype(x(0))>::type;
    using y_value_type = typename std::decay<decltype(y(0))>::type;
 
    const IndexType numRows = static_cast<IndexType>(A.extent(0));
    const IndexType numCols = static_cast<IndexType>(A.extent(1));
 
    const A_value_type* __restrict A_ptr(A.data());
    const IndexType as1(A.stride(1));
    const x_value_type* __restrict x_ptr(x.data());
    y_value_type* __restrict y_ptr(y.data());
 
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
    for (IndexType j = 0; j < numCols; ++j) {
      const x_value_type x_at_j             = alpha * x_ptr[j];
      const A_value_type* __restrict A_at_j = A_ptr + j * as1;
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
      for (IndexType i = 0; i < numRows; ++i)
        y_ptr[i] += A_at_j[i] * x_at_j;
    }
  }
};
 
template <class VecType1,
          class BlkType,
          class VecType2,
          class CoeffType,
          class IndexType>
struct GEMV<VecType1, BlkType, VecType2, CoeffType, IndexType, true, Kokkos::LayoutRight> {
  static KOKKOS_INLINE_FUNCTION void
  run(const CoeffType& alpha,
      const BlkType& A,
      const VecType1& x,
      const VecType2& y) {
    static_assert(VecType1::rank == 1, "GEMV: VecType1 must have rank 1.");
    static_assert(BlkType::rank == 2, "GEMV: BlkType must have rank 2.");
    static_assert(VecType2::rank == 1, "GEMV: VecType2 must have rank 1.");
 
    using A_value_type = typename std::decay<decltype(A(0, 0))>::type;
    using x_value_type = typename std::decay<decltype(x(0))>::type;
    using y_value_type = typename std::decay<decltype(y(0))>::type;
 
    const IndexType numRows = static_cast<IndexType>(A.extent(0));
    const IndexType numCols = static_cast<IndexType>(A.extent(1));
 
    const A_value_type* __restrict A_ptr(A.data());
    const IndexType as0(A.stride(0));
    const x_value_type* __restrict x_ptr(x.data());
    y_value_type* __restrict y_ptr(y.data());
 
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
    for (IndexType i = 0; i < numRows; ++i) {
      y_value_type y_at_i(0);
      const auto A_at_i = A_ptr + i * as0;
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
      for (IndexType j = 0; j < numCols; ++j)
        y_at_i += A_at_i[j] * x_ptr[j];
      y_ptr[i] += alpha * y_at_i;
    }
  }
};
 
}  // namespace Impl

template <class ViewType,
          class CoefficientType,
          class IndexType = int,
          const int rank  = ViewType::rank>
KOKKOS_INLINE_FUNCTION void
SCAL(const CoefficientType& alpha, const ViewType& x) {
  using LayoutType             = typename ViewType::array_layout;
  constexpr bool is_contiguous = (std::is_same<LayoutType, Kokkos::LayoutLeft>::value ||
                                  std::is_same<LayoutType, Kokkos::LayoutRight>::value);
  Impl::SCAL<ViewType, CoefficientType, IndexType, is_contiguous, rank>::run(alpha, x);
}

template <class ViewType,
          class InputType,
          class IndexType = int,
          const int rank  = ViewType::rank>
KOKKOS_INLINE_FUNCTION void
FILL(const ViewType& x, const InputType& val) {
  using LayoutType             = typename ViewType::array_layout;
  constexpr bool is_contiguous = (std::is_same<LayoutType, Kokkos::LayoutLeft>::value ||
                                  std::is_same<LayoutType, Kokkos::LayoutRight>::value);
  Impl::FILL<ViewType, InputType, IndexType, is_contiguous, rank>::run(x, val);
}

template <class CoefficientType,
          class ViewType1,
          class ViewType2,
          class IndexType = int,
          const int rank  = ViewType1::rank>
KOKKOS_INLINE_FUNCTION void
AXPY(const CoefficientType& alpha,
     const ViewType1& x,
     const ViewType2& y) {
  static_assert(static_cast<int>(ViewType1::rank) ==
                    static_cast<int>(ViewType2::rank),
                "AXPY: x and y must have the same rank.");
  using LayoutType1              = typename ViewType1::array_layout;
  using LayoutType2              = typename ViewType2::array_layout;
  constexpr bool is_layout_same  = std::is_same<LayoutType1, LayoutType2>::value;
  constexpr bool is_x_contiguous = (std::is_same<LayoutType1, Kokkos::LayoutLeft>::value ||
                                    std::is_same<LayoutType1, Kokkos::LayoutRight>::value);
  constexpr bool is_contiguous   = is_layout_same && is_x_contiguous;
  Impl::AXPY<CoefficientType, ViewType1, ViewType2, IndexType, is_contiguous, rank>::run(alpha, x, y);
}

template <class ViewType1,
          class ViewType2,
          class IndexType = int,
          const int rank  = ViewType1::rank>
KOKKOS_INLINE_FUNCTION void
COPY(const ViewType1& x, const ViewType2& y) {
  static_assert(static_cast<int>(ViewType1::rank) ==
                    static_cast<int>(ViewType2::rank),
                "COPY: x and y must have the same rank.");
  using LayoutType1              = typename ViewType1::array_layout;
  using LayoutType2              = typename ViewType2::array_layout;
  constexpr bool is_layout_same  = std::is_same<LayoutType1, LayoutType2>::value;
  constexpr bool is_x_contiguous = (std::is_same<LayoutType1, Kokkos::LayoutLeft>::value ||
                                    std::is_same<LayoutType1, Kokkos::LayoutRight>::value);
  constexpr bool is_contiguous   = is_layout_same && is_x_contiguous;
  Impl::COPY<ViewType1, ViewType2, IndexType, is_contiguous, rank>::run(x, y);
}

template <class VecType1,
          class BlkType,
          class VecType2,
          class CoeffType,
          class IndexType = int>
KOKKOS_INLINE_FUNCTION void
GEMV(const CoeffType& alpha,
     const BlkType& A,
     const VecType1& x,
     const VecType2& y) {
  constexpr bool is_A_contiguous = (std::is_same<typename BlkType::array_layout, Kokkos::LayoutLeft>::value ||
                                    std::is_same<typename BlkType::array_layout, Kokkos::LayoutRight>::value);
  constexpr bool is_x_contiguous = (std::is_same<typename VecType1::array_layout, Kokkos::LayoutLeft>::value ||
                                    std::is_same<typename VecType1::array_layout, Kokkos::LayoutRight>::value);
  constexpr bool is_y_contiguous = (std::is_same<typename VecType2::array_layout, Kokkos::LayoutLeft>::value ||
                                    std::is_same<typename VecType2::array_layout, Kokkos::LayoutRight>::value);
  constexpr bool is_contiguous   = is_A_contiguous && is_x_contiguous && is_y_contiguous;
  Impl::GEMV<VecType1, BlkType, VecType2, CoeffType, IndexType, is_contiguous>::run(alpha, A, x, y);
}

template <class ViewType1,
          class ViewType2,
          class ViewType3,
          class CoefficientType,
          class IndexType = int>
KOKKOS_INLINE_FUNCTION void
GEMM(const char transA[],
     const char transB[],
     const CoefficientType& alpha,
     const ViewType1& A,
     const ViewType2& B,
     const CoefficientType& beta,
     const ViewType3& C) {
  // Assert that A, B, and C are in fact matrices
  static_assert(ViewType1::rank == 2, "GEMM: A must have rank 2 (be a matrix).");
  static_assert(ViewType2::rank == 2, "GEMM: B must have rank 2 (be a matrix).");
  static_assert(ViewType3::rank == 2, "GEMM: C must have rank 2 (be a matrix).");
 
  typedef typename std::remove_reference<decltype(A(0, 0))>::type Scalar;
  typedef KokkosKernels::ArithTraits<Scalar> STS;
  const Scalar ZERO = STS::zero();
  const Scalar ONE  = STS::one();
 
  // Get the dimensions
  IndexType m, n, k;
  if (transA[0] == 'N' || transA[0] == 'n') {
    m = static_cast<IndexType>(A.extent(0));
    n = static_cast<IndexType>(A.extent(1));
  } else {
    m = static_cast<IndexType>(A.extent(1));
    n = static_cast<IndexType>(A.extent(0));
  }
  k = static_cast<IndexType>(C.extent(1));
 
  // quick return if possible
  if (alpha == ZERO && beta == ONE)
    return;
 
  // And if alpha equals zero...
  if (alpha == ZERO) {
    if (beta == ZERO) {
      for (IndexType i = 0; i < m; i++) {
        for (IndexType j = 0; j < k; j++) {
          C(i, j) = ZERO;
        }
      }
    } else {
      for (IndexType i = 0; i < m; i++) {
        for (IndexType j = 0; j < k; j++) {
          C(i, j) = beta * C(i, j);
        }
      }
    }
  }
 
  // Start the operations
  if (transB[0] == 'n' || transB[0] == 'N') {
    if (transA[0] == 'n' || transA[0] == 'N') {
      // Form C = alpha*A*B + beta*C
      for (IndexType j = 0; j < n; j++) {
        if (beta == ZERO) {
          for (IndexType i = 0; i < m; i++) {
            C(i, j) = ZERO;
          }
        } else if (beta != ONE) {
          for (IndexType i = 0; i < m; i++) {
            C(i, j) = beta * C(i, j);
          }
        }
        for (IndexType l = 0; l < k; l++) {
          Scalar temp = alpha * B(l, j);
          for (IndexType i = 0; i < m; i++) {
            C(i, j) = C(i, j) + temp * A(i, l);
          }
        }
      }
    } else {
      // Form C = alpha*A**T*B + beta*C
      for (IndexType j = 0; j < n; j++) {
        for (IndexType i = 0; i < m; i++) {
          Scalar temp = ZERO;
          for (IndexType l = 0; l < k; l++) {
            temp = temp + A(l, i) * B(l, j);
          }
          if (beta == ZERO) {
            C(i, j) = alpha * temp;
          } else {
            C(i, j) = alpha * temp + beta * C(i, j);
          }
        }
      }
    }
  } else {
    if (transA[0] == 'n' || transA[0] == 'N') {
      // Form C = alpha*A*B**T + beta*C
      for (IndexType j = 0; j < n; j++) {
        if (beta == ZERO) {
          for (IndexType i = 0; i < m; i++) {
            C(i, j) = ZERO;
          }
        } else if (beta != ONE) {
          for (IndexType i = 0; i < m; i++) {
            C(i, j) = beta * C(i, j);
          }
        }
        for (IndexType l = 0; l < k; l++) {
          Scalar temp = alpha * B(j, l);
          for (IndexType i = 0; i < m; i++) {
            C(i, j) = C(i, j) + temp * A(i, l);
          }
        }
      }
    } else {
      // Form C = alpha*A**T*B**T + beta*C
      for (IndexType j = 0; j < n; j++) {
        for (IndexType i = 0; i < m; i++) {
          Scalar temp = ZERO;
          for (IndexType l = 0; l < k; l++) {
            temp = temp + A(l, i) * B(j, l);
          }
          if (beta == ZERO) {
            C(i, j) = alpha * temp;
          } else {
            C(i, j) = alpha * temp + beta * C(i, j);
          }
        }
      }
    }
  }
}

template <class LittleBlockType,
          class LittleVectorType>
KOKKOS_INLINE_FUNCTION void
GETF2(const LittleBlockType& A, const LittleVectorType& ipiv, int& info) {
  // The type of an entry of ipiv is the index type.
  typedef typename std::decay<decltype(ipiv(0))>::type IndexType;
  static_assert(std::is_integral<IndexType>::value,
                "GETF2: The type of each entry of ipiv must be an integer type.");
  typedef typename std::remove_reference<decltype(A(0, 0))>::type Scalar;
  static_assert(!std::is_const<Scalar>::value,
                "GETF2: A must not be a const View (or LittleBlock).");
  static_assert(!std::is_const<std::remove_reference<decltype(ipiv(0))>>::value,
                "GETF2: ipiv must not be a const View (or LittleBlock).");
  static_assert(LittleBlockType::rank == 2, "GETF2: A must have rank 2 (be a matrix).");
  typedef KokkosKernels::ArithTraits<Scalar> STS;
  const Scalar ZERO = STS::zero();
 
  const IndexType numRows = static_cast<IndexType>(A.extent(0));
  const IndexType numCols = static_cast<IndexType>(A.extent(1));
  const IndexType pivDim  = static_cast<IndexType>(ipiv.extent(0));
 
  // std::min is not a CUDA device function
  const IndexType minPivDim = (numRows < numCols) ? numRows : numCols;
  if (pivDim < minPivDim) {
    info = -2;
    return;
  }
 
  // Initialize info
  info = 0;
 
  for (IndexType j = 0; j < pivDim; j++) {
    // Find pivot and test for singularity
    IndexType jp = j;
    for (IndexType i = j + 1; i < numRows; i++) {
      if (STS::abs(A(i, j)) > STS::abs(A(jp, j))) {
        jp = i;
      }
    }
    ipiv(j) = jp + 1;
 
    if (A(jp, j) != ZERO) {
      // Apply the interchange to columns 1:N
      if (jp != j) {
        for (IndexType i = 0; i < numCols; i++) {
          Scalar temp = A(jp, i);
          A(jp, i)    = A(j, i);
          A(j, i)     = temp;
        }
      }
 
      // Compute elements J+1:M of J-th column
      for (IndexType i = j + 1; i < numRows; i++) {
        A(i, j) = A(i, j) / A(j, j);
      }
    } else if (info == 0) {
      info = j;
    }
 
    // Update trailing submatrix
    for (IndexType r = j + 1; r < numRows; r++) {
      for (IndexType c = j + 1; c < numCols; c++) {
        A(r, c) = A(r, c) - A(r, j) * A(j, c);
      }
    }
  }
}
 
namespace Impl {

template <class LittleBlockType,
          class LittleIntVectorType,
          class LittleScalarVectorType,
          const int rank = LittleScalarVectorType::rank>
struct GETRS {
  static KOKKOS_INLINE_FUNCTION void
  run(const char mode[],
      const LittleBlockType& A,
      const LittleIntVectorType& ipiv,
      const LittleScalarVectorType& B,
      int& info);
};

template <class LittleBlockType,
          class LittleIntVectorType,
          class LittleScalarVectorType>
struct GETRS<LittleBlockType, LittleIntVectorType, LittleScalarVectorType, 1> {
  static KOKKOS_INLINE_FUNCTION void
  run(const char mode[],
      const LittleBlockType& A,
      const LittleIntVectorType& ipiv,
      const LittleScalarVectorType& B,
      int& info) {
    // The type of an entry of ipiv is the index type.
    typedef typename std::decay<decltype(ipiv(0))>::type IndexType;
    // IndexType must be signed, because this code does a countdown loop
    // to zero.  Unsigned integers are always >= 0, even on underflow.
    static_assert(std::is_integral<IndexType>::value &&
                      std::is_signed<IndexType>::value,
                  "GETRS: The type of each entry of ipiv must be a signed integer.");
    typedef typename std::decay<decltype(A(0, 0))>::type Scalar;
    static_assert(!std::is_const<std::remove_reference<decltype(B(0))>>::value,
                  "GETRS: B must not be a const View (or LittleBlock).");
    static_assert(LittleBlockType::rank == 2, "GETRS: A must have rank 2 (be a matrix).");
    static_assert(LittleIntVectorType::rank == 1, "GETRS: ipiv must have rank 1.");
    static_assert(LittleScalarVectorType::rank == 1, "GETRS: For this specialization, B must have rank 1.");
 
    typedef KokkosKernels::ArithTraits<Scalar> STS;
    const Scalar ZERO       = STS::zero();
    const IndexType numRows = static_cast<IndexType>(A.extent(0));
    const IndexType numCols = static_cast<IndexType>(A.extent(1));
    const IndexType pivDim  = static_cast<IndexType>(ipiv.extent(0));
 
    info = 0;
 
    // Ensure that the matrix is square
    if (numRows != numCols) {
      info = -2;
      return;
    }
 
    // Ensure that the pivot array is sufficiently large
    if (pivDim < numRows) {
      info = -3;
      return;
    }
 
    // No transpose case
    if (mode[0] == 'n' || mode[0] == 'N') {
      // Apply row interchanges to the RHS
      for (IndexType i = 0; i < numRows; i++) {
        if (ipiv(i) != i + 1) {
          Scalar temp    = B(i);
          B(i)           = B(ipiv(i) - 1);
          B(ipiv(i) - 1) = temp;
        }
      }
 
      // Solve Lx=b, overwriting b with x
      for (IndexType r = 1; r < numRows; r++) {
        for (IndexType c = 0; c < r; c++) {
          B(r) = B(r) - A(r, c) * B(c);
        }
      }
 
      // Solve Ux=b, overwriting b with x
      for (IndexType r = numRows - 1; r >= 0; r--) {
        // Check whether U is singular
        if (A(r, r) == ZERO) {
          info = r + 1;
          return;
        }
 
        for (IndexType c = r + 1; c < numCols; c++) {
          B(r) = B(r) - A(r, c) * B(c);
        }
        B(r) = B(r) / A(r, r);
      }
    }
    // Transpose case
    else if (mode[0] == 't' || mode[0] == 'T') {
      info = -1;  // NOT YET IMPLEMENTED
      return;
    }
    // Conjugate transpose case
    else if (mode[0] == 'c' || mode[0] == 'C') {
      info = -1;  // NOT YET IMPLEMENTED
      return;
    } else {  // invalid mode
      info = -1;
      return;
    }
  }
};

template <class LittleBlockType,
          class LittleIntVectorType,
          class LittleScalarVectorType>
struct GETRS<LittleBlockType, LittleIntVectorType, LittleScalarVectorType, 2> {
  static KOKKOS_INLINE_FUNCTION void
  run(const char mode[],
      const LittleBlockType& A,
      const LittleIntVectorType& ipiv,
      const LittleScalarVectorType& B,
      int& info) {
    // The type of an entry of ipiv is the index type.
    typedef typename std::decay<decltype(ipiv(0))>::type IndexType;
    static_assert(std::is_integral<IndexType>::value,
                  "GETRS: The type of each entry of ipiv must be an integer type.");
    static_assert(!std::is_const<std::remove_reference<decltype(B(0))>>::value,
                  "GETRS: B must not be a const View (or LittleBlock).");
    static_assert(LittleBlockType::rank == 2, "GETRS: A must have rank 2 (be a matrix).");
    static_assert(LittleIntVectorType::rank == 1, "GETRS: ipiv must have rank 1.");
    static_assert(LittleScalarVectorType::rank == 2, "GETRS: For this specialization, B must have rank 2.");
 
    // The current implementation iterates over one right-hand side at
    // a time.  It might be faster to do this differently, but this
    // should work for now.
    const IndexType numRhs = B.extent(1);
    info                   = 0;
 
    for (IndexType rhs = 0; rhs < numRhs; ++rhs) {
      auto B_cur = Kokkos::subview(B, Kokkos::ALL(), rhs);
      GETRS<LittleBlockType, LittleIntVectorType, decltype(B_cur), 1>::run(mode, A, ipiv, B_cur, info);
      if (info != 0) {
        return;
      }
    }
  }
};
 
}  // namespace Impl

template <class LittleBlockType,
          class LittleIntVectorType,
          class LittleScalarVectorType>
KOKKOS_INLINE_FUNCTION void
GETRS(const char mode[],
      const LittleBlockType& A,
      const LittleIntVectorType& ipiv,
      const LittleScalarVectorType& B,
      int& info) {
  Impl::GETRS<LittleBlockType, LittleIntVectorType, LittleScalarVectorType,
              LittleScalarVectorType::rank>::run(mode, A, ipiv, B, info);
}

template <class LittleBlockType,
          class LittleIntVectorType,
          class LittleScalarVectorType>
KOKKOS_INLINE_FUNCTION void
GETRI(const LittleBlockType& A,
      const LittleIntVectorType& ipiv,
      const LittleScalarVectorType& work,
      int& info) {
  // The type of an entry of ipiv is the index type.
  typedef typename std::decay<decltype(ipiv(0))>::type IndexType;
  // IndexType must be signed, because this code does a countdown loop
  // to zero.  Unsigned integers are always >= 0, even on underflow.
  static_assert(std::is_integral<IndexType>::value &&
                    std::is_signed<IndexType>::value,
                "GETRI: The type of each entry of ipiv must be a signed integer.");
  typedef typename std::remove_reference<decltype(A(0, 0))>::type Scalar;
  static_assert(!std::is_const<std::remove_reference<decltype(A(0, 0))>>::value,
                "GETRI: A must not be a const View (or LittleBlock).");
  static_assert(!std::is_const<std::remove_reference<decltype(work(0))>>::value,
                "GETRI: work must not be a const View (or LittleBlock).");
  static_assert(LittleBlockType::rank == 2, "GETRI: A must have rank 2 (be a matrix).");
  typedef KokkosKernels::ArithTraits<Scalar> STS;
  const Scalar ZERO = STS::zero();
  const Scalar ONE  = STS::one();
 
  const IndexType numRows = static_cast<IndexType>(A.extent(0));
  const IndexType numCols = static_cast<IndexType>(A.extent(1));
  const IndexType pivDim  = static_cast<IndexType>(ipiv.extent(0));
  const IndexType workDim = static_cast<IndexType>(work.extent(0));
 
  info = 0;
 
  // Ensure that the matrix is square
  if (numRows != numCols) {
    info = -1;
    return;
  }
 
  // Ensure that the pivot array is sufficiently large
  if (pivDim < numRows) {
    info = -2;
    return;
  }
 
  // Ensure that the work array is sufficiently large
  if (workDim < numRows) {
    info = -3;
    return;
  }
 
  // Form Uinv in place
  for (IndexType j = 0; j < numRows; j++) {
    if (A(j, j) == ZERO) {
      info = j + 1;
      return;
    }
 
    A(j, j) = ONE / A(j, j);
 
    // Compute elements 1:j-1 of j-th column
    for (IndexType r = 0; r < j; r++) {
      A(r, j) = A(r, r) * A(r, j);
      for (IndexType c = r + 1; c < j; c++) {
        A(r, j) = A(r, j) + A(r, c) * A(c, j);
      }
    }
    for (IndexType r = 0; r < j; r++) {
      A(r, j) = -A(j, j) * A(r, j);
    }
  }
 
  // Compute Ainv by solving A\L = Uinv
  for (IndexType j = numCols - 2; j >= 0; j--) {
    // Copy lower triangular data to work array and replace with 0
    for (IndexType r = j + 1; r < numRows; r++) {
      work(r) = A(r, j);
      A(r, j) = 0;
    }
 
    for (IndexType r = 0; r < numRows; r++) {
      for (IndexType i = j + 1; i < numRows; i++) {
        A(r, j) = A(r, j) - work(i) * A(r, i);
      }
    }
  }
 
  // Apply column interchanges
  for (IndexType j = numRows - 1; j >= 0; j--) {
    IndexType jp = ipiv(j) - 1;
    if (j != jp) {
      for (IndexType r = 0; r < numRows; r++) {
        Scalar temp = A(r, j);
        A(r, j)     = A(r, jp);
        A(r, jp)    = temp;
      }
    }
  }
}
 
// mfh 08 Nov 2015: I haven't tested this overload yet.  It also needs
// an implementation for trans != 'N' (the transpose and conjugate
// transpose cases).
#if 0
template<class LittleBlockType,
         class LittleVectorTypeX,
         class LittleVectorTypeY,
         class CoefficientType,
         class IndexType = int>
KOKKOS_INLINE_FUNCTION void
GEMV (const char trans,
      const CoefficientType& alpha,
      const LittleBlockType& A,
      const LittleVectorTypeX& x,
      const CoefficientType& beta,
      const LittleVectorTypeY& y)
{
  // y(0) returns a reference to the 0-th entry of y.  Remove that
  // reference to get the type of each entry of y.  It's OK if y has
  // zero entries -- this doesn't actually do y(i), it just returns
  // the type of that expression.
  typedef typename std::remove_reference<decltype (y(0)) >::type y_value_type;
  const IndexType numRows = static_cast<IndexType> (A.extent (0));
  const IndexType numCols = static_cast<IndexType> (A.extent (1));
 
  if (beta == 0.0) {
    if (alpha == 0.0) {
      for (IndexType i = 0; i < numRows; ++i) {
        y(i) = 0.0;
      }
    }
    else {
      for (IndexType i = 0; i < numRows; ++i) {
        y_value_type y_i = 0.0;
        for (IndexType j = 0; j < numCols; ++j) {
          y_i += A(i,j) * x(j);
        }
        y(i) = y_i;
      }
    }
  }
  else { // beta != 0
    if (alpha == 0.0) {
      if (beta == 0.0) {
        for (IndexType i = 0; i < numRows; ++i) {
          y(i) = 0.0;
        }
      }
      else {
        for (IndexType i = 0; i < numRows; ++i) {
          y(i) *= beta;
        }
      }
    }
    else {
      for (IndexType i = 0; i < numRows; ++i) {
        y_value_type y_i = beta * y(i);
        for (IndexType j = 0; j < numCols; ++j) {
          y_i += alpha * A(i,j) * x(j);
        }
        y(i) = y_i;
      }
    }
  }
}
 
#endif  // 0
 
}  // namespace Tpetra
 
#endif  // TPETRA_BLOCKVIEW_HPP