linalg.cuh

Go to the documentation of this file.

#pragma once
#include <color_spinor.h> // vector container

namespace quda {

  namespace linalg {

    template <template<typename,int> class Mat, typename T, int N, bool fast=true>
    class Cholesky {

      Mat<T,N> L_;

    public:
      __device__ __host__ inline Cholesky(const Mat<T,N> &A) {
  const Mat<T,N> &L = L_;

#pragma unroll
  for (int i=0; i<N; i++) {
#pragma unroll
    for (int j=0; j<N; j++) if (j<i+1) {
      complex<T> s = 0;
#pragma unroll
      for (int k=0; k<N; k++) {
        if (k==0) {
    s.x  = L(i,k).real()*L(j,k).real();
    s.x += L(i,k).imag()*L(j,k).imag();
    s.y  = L(i,k).imag()*L(j,k).real();
    s.y -= L(i,k).real()*L(j,k).imag();
        } else if (k<j) {
    s.x += L(i,k).real()*L(j,k).real();
    s.x += L(i,k).imag()*L(j,k).imag();
    s.y += L(i,k).imag()*L(j,k).real();
    s.y -= L(i,k).real()*L(j,k).imag();
        }
      }
      if (!fast) { // traditional Cholesky with sqrt and division
        L_(i,j) = (i == j) ? sqrt((A(i,i)-s).real()) : (A(i,j) - s) / L(j,j).real();
      } else { // optimized - since fwd/back subsitition only need inverse diagonal elements, avoid division and use rsqrt
        L_(i,j) = (i == j) ? rsqrt((A(i,i)-s).real()) : (A(i,j)-s) * L(j,j).real();
      }
    }
  }
      }

      __device__ __host__ inline const T D(int i) const {
  const auto &L = L_;
  if (!fast) return L(i,i).real();
  else return static_cast<T>(1.0) / L(i,i).real();
      }

      template <class Vector>
      __device__ __host__ inline Vector forward(const Vector &b) {
  const Mat<T,N> &L = L_;
  Vector x;
#pragma unroll
  for (int i=0; i<N; i++) {
    x(i) = b(i);
#pragma unroll
    for (int j=0; j<N; j++) if (j<i) {
      x(i).x -= L(i,j).real()*x(j).real();
      x(i).x += L(i,j).imag()*x(j).imag();
      x(i).y -= L(i,j).real()*x(j).imag();
      x(i).y -= L(i,j).imag()*x(j).real();
    }
    if (!fast) x(i) /= L(i,i).real(); // traditional
    else x(i) *= L(i,i).real();       // optimized
  }
  return x;
      }

      template <class Vector>
      __device__ __host__ inline Vector backward(const Vector &b) {
  const Mat<T,N> &L = L_;
  Vector x;
#pragma unroll
  for (int i=N-1; i>=0; i--) {
    x(i) = b(i);
#pragma unroll
    for (int j=0; j<N; j++) if (j>=i+1) {
      x(i).x -= L(i,j).real()*x(j).real();
      x(i).x += L(i,j).imag()*x(j).imag();
      x(i).y -= L(i,j).real()*x(j).imag();
      x(i).y -= L(i,j).imag()*x(j).real();
    }
    if (!fast) x(i) /=L(i,i).real(); // traditional
    else x(i) *= L(i,i).real();      // optimized
  }
  return x;
      }

      __device__ __host__ inline Mat<T,N> invert() {
  const Mat<T,N> &L = L_;
  Mat<T,N> Ainv;
  ColorSpinor<T,1,N> v;

#pragma unroll
  for (int k=0;k<N;k++) {

    // forward substitute
    if (!fast) v(k) = complex<T>(static_cast<T>(1.0)/L(k,k).real());
    else v(k) = L(k,k).real();

#pragma unroll
    for (int i=0; i<N; i++) if (i>k) {
      v(i) = complex<T>(0.0);
#pragma unroll
      for (int j=0; j<N; j++) if (j>=k && j<i) {
        v(i).x -= L(i,j).real() * v(j).real();
        v(i).x += L(i,j).imag() * v(j).imag();
        v(i).y -= L(i,j).real() * v(j).imag();
        v(i).y -= L(i,j).imag() * v(j).real();
      }
      if (!fast) v(i) *= static_cast<T>(1.0) / L(i,i);
      else v(i) *= L(i,i);
    }

    // backward substitute
    if (!fast) v(N-1) *= static_cast<T>(1.0) / L(N-1,N-1);
    else v(N-1) *= L(N-1,N-1);

#pragma unroll
    for (int i=N-2; i>=0; i--) if (i>=k) {
#pragma unroll
      for (int j=0; j<N; j++) if (j>i) {
        v(i).x -= L(i,j).real() * v(j).real();
        v(i).x += L(i,j).imag() * v(j).imag();
        v(i).y -= L(i,j).real() * v(j).imag();
        v(i).y -= L(i,j).imag() * v(j).real();
      }
      if (!fast) v(i) *= static_cast<T>(1.0) / L(i,i);
      else v(i) *= L(i,i);
    }

    // Overwrite column k
    Ainv(k,k) = v(k);

#pragma unroll
    for(int i=0;i<N;i++) if (i>k) Ainv(i,k) = v(i);
  }

  return Ainv;
      }

    };

  } // namespace linalg

} // namespace quda