!===============================================================================
!
! cg_loops_simd.F90
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2012-2016 Hinnerk Stueben
!
! This file is part of BQCD -- Berlin Quantum ChromoDynamics program
!
! BQCD is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! BQCD is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with BQCD.  If not, see <http://www.gnu.org/licenses/>.
!
!-------------------------------------------------------------------------------
# define LANGUAGE_FORTRAN
# define SIMD
# include "defs.h"
# include "cg_defs.h"
# include "simd.h"

!-------------------------------------------------------------------------------
# undef  TRACE_NAME
# define TRACE_NAME T(cg_loop1)
!-------------------------------------------------------------------------------
subroutine T(cg_loop1)(rtrold, r, p, b, volh_tot, varsBLOCK)

  use typedef_simd
  implicit none
  integer, intent(in)        :: volh_tot, varsBLOCK
  spincol_FIELD, intent(in)  :: b
  spincol_FIELD, intent(out) :: r, p
  DOUBLE, intent(out)        :: rtrold
  DOUBLE, dimension(SIMD_WIDTH) :: dvec1, dvec2
  DOUBLE                        :: dtmp1, dtmp2
  integer                    :: i

  TRACE_ENTRY
  TIMING_START(timing_bin_cg_loop1)

  rtrold = ZERO

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, r, p, b) &
  !$omp private(dvec1, dvec2, dtmp1, dtmp2) &
  !$omp reduction(+: rtrold)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     Real(Spincol(r, s1, s2, c, i)) = Real(Spincol(b, s1, s2, c, i)) - Real(Spincol(r, s1, s2, c, i))
     Imag(Spincol(r, s1, s2, c, i)) = Imag(Spincol(b, s1, s2, c, i)) - Imag(Spincol(r, s1, s2, c, i))
     Real(Spincol(p, s1, s2, c, i)) = Real(Spincol(r, s1, s2, c, i))
     Imag(Spincol(p, s1, s2, c, i)) = Imag(Spincol(r, s1, s2, c, i))
     dvec1 = Real(Spincol(r, s1, s2, c, i))
     dvec2 = Imag(Spincol(r, s1, s2, c, i))
     dtmp1 = dot_product(dvec1, dvec1)
     dtmp2 = dot_product(dvec2, dvec2)
     rtrold = rtrold + dtmp1
     rtrold = rtrold + dtmp2
     ENDDO
     ENDDO
     ENDDO
  enddo

  TIMING_STOP(timing_bin_cg_loop1)
  TRACE_RETURN
end

!-------------------------------------------------------------------------------
# undef  TRACE_NAME
# define TRACE_NAME T(cg_loop2)
!-------------------------------------------------------------------------------
subroutine T(cg_loop2)(paap, p, aap, volh_tot, varsBLOCK)

  use typedef_simd
  implicit none
  integer, intent(in)        :: volh_tot, varsBLOCK
  spincol_FIELD, intent(in)  :: p, aap
  DOUBLE, intent(out)        :: paap
  DOUBLE, dimension(SIMD_WIDTH) :: dvec1, dvec2
  DOUBLE                        :: dtmp
  integer                    :: i

  TRACE_ENTRY
  TIMING_START(timing_bin_cg_loop2)

  paap = ZERO

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, p, aap) &
  !$omp private(dvec1, dvec2, dtmp) &
  !$omp reduction(+: paap)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     dvec1 = Real(Spincol(p, s1, s2, c, i))
     dvec2 = Real(Spincol(aap, s1, s2, c, i))
     dtmp = dot_product(dvec1, dvec2)
     paap = paap + dtmp
     dvec1 = Imag(Spincol(p, s1, s2, c, i))
     dvec2 = Imag(Spincol(aap, s1, s2, c, i))
     dtmp = dot_product(dvec1, dvec2)
     paap = paap + dtmp
     ENDDO
     ENDDO
     ENDDO
  enddo

  TIMING_STOP(timing_bin_cg_loop2)
  TRACE_RETURN
end

!-------------------------------------------------------------------------------
# undef  TRACE_NAME
# define TRACE_NAME T(cg_loop3)
!-------------------------------------------------------------------------------
subroutine T(cg_loop3)(rtr, x, r, p, aap, ak, volh_tot, varsBLOCK)

  use typedef_simd
  implicit none
  integer, intent(in)          :: volh_tot, varsBLOCK
  REAL, intent(in)             :: ak
  spincol_FIELD, intent(in)    :: p, aap
  spincol_FIELD, intent(inout) :: x, r
  DOUBLE, intent(out)          :: rtr
  DOUBLE, dimension(SIMD_WIDTH) :: dvec1, dvec2
  DOUBLE                        :: dtmp1, dtmp2
  integer                      :: i

  TRACE_ENTRY
  TIMING_START(timing_bin_cg_loop3)

  rtr = ZERO

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, x, r, p, aap, ak) &
  !$omp private(dvec1, dvec2, dtmp1, dtmp2) &
  !$omp reduction(+: rtr)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     Real(Spincol(x, s1, s2, c, i)) = Real(Spincol(x, s1, s2, c, i)) + ak * Real(Spincol(p, s1, s2, c, i))
     Imag(Spincol(x, s1, s2, c, i)) = Imag(Spincol(x, s1, s2, c, i)) + ak * Imag(Spincol(p, s1, s2, c, i))
     Real(Spincol(r, s1, s2, c, i)) = Real(Spincol(r, s1, s2, c, i)) - ak * Real(Spincol(aap, s1, s2, c, i))
     Imag(Spincol(r, s1, s2, c, i)) = Imag(Spincol(r, s1, s2, c, i)) - ak * Imag(Spincol(aap, s1, s2, c, i))
     dvec1 = Real(Spincol(r, s1, s2, c, i))
     dvec2 = Imag(Spincol(r, s1, s2, c, i))
     dtmp1 = dot_product(dvec1, dvec1)
     dtmp2 = dot_product(dvec2, dvec2)
     rtr = rtr + dtmp1
     rtr = rtr + dtmp2
     ENDDO
     ENDDO
     ENDDO
  enddo

  TIMING_STOP(timing_bin_cg_loop3)
  TRACE_RETURN
end

!-------------------------------------------------------------------------------
# undef  TRACE_NAME
# define TRACE_NAME T(cg_axpy)
!-------------------------------------------------------------------------------
subroutine T(cg_axpy)(inout, in, a, volh_tot, varsBLOCK)  ! inout = inout + a * in

  use typedef_simd
  implicit none
  integer, intent(in)          :: volh_tot, varsBLOCK
  REAL, intent(in)             :: a
  spincol_FIELD, intent(in)    :: in
  spincol_FIELD, intent(inout) :: inout
  integer                      :: i

  TRACE_ENTRY
  TIMING_START(timing_bin_cg_axpy)

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, inout, in, a)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     Real(Spincol(inout, s1, s2, c, i)) = Real(Spincol(inout, s1, s2, c, i)) + a * Real(Spincol(in, s1, s2, c, i))
     Imag(Spincol(inout, s1, s2, c, i)) = Imag(Spincol(inout, s1, s2, c, i)) + a * Imag(Spincol(in, s1, s2, c, i))
     ENDDO
     ENDDO
     ENDDO
  enddo

  TIMING_STOP(timing_bin_cg_axpy)
  TRACE_RETURN
end

!-------------------------------------------------------------------------------
# undef  TRACE_NAME
# define TRACE_NAME T(cg_xpby)
!-------------------------------------------------------------------------------
subroutine T(cg_xpby)(inout, in, b, volh_tot, varsBLOCK)  ! inout = b * inout + in

  use typedef_simd
  implicit none
  integer, intent(in)          :: volh_tot, varsBLOCK
  REAL, intent(in)             :: b
  spincol_FIELD, intent(in)    :: in
  spincol_FIELD, intent(inout) :: inout
  integer                      :: i

  TRACE_ENTRY
  TIMING_START(timing_bin_cg_xpby)

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, inout, in, b)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     Real(Spincol(inout, s1, s2, c, i)) = b * Real(Spincol(inout, s1, s2, c, i)) + Real(Spincol(in, s1, s2, c, i))
     Imag(Spincol(inout, s1, s2, c, i)) = b * Imag(Spincol(inout, s1, s2, c, i)) + Imag(Spincol(in, s1, s2, c, i))
     ENDDO
     ENDDO
     ENDDO
  enddo

  TIMING_STOP(timing_bin_cg_xpby)
  TRACE_RETURN
end

!-------------------------------------------------------------------------------
# undef  TRACE_NAME
# define TRACE_NAME T(cg_norm2)
!-------------------------------------------------------------------------------
subroutine T(cg_norm2)(norm2, in, volh_tot, varsBLOCK)  ! Sum_i abs(in_i)**2

  use typedef_simd
  implicit none
  integer, intent(in)          :: volh_tot, varsBLOCK
  DOUBLE, intent(out)          :: norm2
  spincol_FIELD, intent(in)    :: in
  integer                      :: i

  DOUBLE, dimension(SIMD_WIDTH) :: dvec1, dvec2
  DOUBLE                        :: dsum, dtmp1, dtmp2

  TRACE_ENTRY

  dsum = ZERO

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, in) &
  !$omp private(dvec1, dvec2, dtmp1, dtmp2) &
  !$omp reduction(+: dsum)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     dvec1 = Real(Spincol(in, s1, s2, c, i))
     dvec2 = Imag(Spincol(in, s1, s2, c, i))
     dtmp1 = dot_product(dvec1, dvec1)
     dtmp2 = dot_product(dvec2, dvec2)
     dsum = dsum + dtmp1
     dsum = dsum + dtmp2
     ENDDO
     ENDDO
     ENDDO
  enddo

  norm2 = dsum

  TRACE_RETURN
end

!-------------------------------------------------------------------------------
# undef  TRACE_NAME
# define TRACE_NAME T(cg_cdotc)
!-------------------------------------------------------------------------------
subroutine T(cg_cdotc)(paap, p, aap, volh_tot, varsBLOCK)

  use typedef_simd
  implicit none
  integer, parameter         :: re = 0
  integer, parameter         :: im = 1

  integer, intent(in)        :: volh_tot, varsBLOCK
  spincol_FIELD, intent(in)  :: p, aap
  REAL, intent(out)          :: paap(re:im)
  REAL, dimension(SIMD_WIDTH) :: dvec1, dvec2
  REAL                       :: dtmp1, dtmp2
  integer                    :: i

  TRACE_ENTRY

  dtmp1 = ZERO
  dtmp2 = ZERO

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, p, aap) &
  !$omp private(dvec1, dvec2) &
  !$omp reduction(+: dtmp1, dtmp2)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     dvec1 = Real(Spincol(p, s1, s2, c, i))
     dvec2 = Real(Spincol(aap, s1, s2, c, i))
     dtmp1 = dtmp1 + dot_product(dvec1, dvec2)

     dvec1 = Imag(Spincol(p, s1, s2, c, i))
     dvec2 = Imag(Spincol(aap, s1, s2, c, i))
     dtmp1 = dtmp1 + dot_product(dvec1, dvec2)

     dvec1 = Real(Spincol(p, s1, s2, c, i))
     dvec2 = Imag(Spincol(aap, s1, s2, c, i))
     dtmp2 = dtmp2 + dot_product(dvec1, dvec2)

     dvec1 = Imag(Spincol(p, s1, s2, c, i))
     dvec2 = Real(Spincol(aap, s1, s2, c, i))
     dtmp2 = dtmp2 - dot_product(dvec1, dvec2)
     ENDDO
     ENDDO
     ENDDO
  enddo

  paap(re) = dtmp1
  paap(im) = dtmp2

  TRACE_RETURN
end

!===============================================================================
