!===============================================================================
!
! cg_loops.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
# include "defs.h"
# include "cg_defs.h"

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

  !// first loop in cg:
  !//
  !// rtrold = ZERO
  !// do i = 1, size_sc_field
  !//    r(i) = b(i) - r(i)
  !//    p(i) = r(i)
  !//    dtmp = r(i)
  !//    rtrold = rtrold + dtmp**2
  !// enddo

  implicit none
  integer, intent(in)          :: volh_tot, varsBLOCK
  spincol_FIELD, intent(in)    :: b
  spincol_FIELD, intent(out)   :: p
  spincol_FIELD, intent(inout) :: r
  DOUBLE, intent(out)          :: rtrold
  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(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
     Spincol(r, s1, s2, c, i) = Spincol(b, s1, s2, c, i) - Spincol(r, s1, s2, c, i)
     Spincol(p, s1, s2, c, i) = Spincol(r, s1, s2, c, i)
     dtmp1 = Re(Spincol(r, s1, s2, c, i)); dtmp1 = dtmp1**2; rtrold = rtrold + dtmp1
     dtmp2 = Im(Spincol(r, s1, s2, c, i)); dtmp2 = dtmp2**2; 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)

  !// second loop in cg:
  !//
  !// paap = ZERO
  !// do i = 1, size_sc_field
  !//    dtmp1 = p(i)
  !//    dtmp2 = aap(i)
  !//    paap = paap + dtmp1 * dtmp2
  !// enddo

  implicit none
  integer, intent(in)        :: volh_tot, varsBLOCK
  spincol_FIELD, intent(in)  :: p, aap
  DOUBLE, intent(out)        :: paap
  DOUBLE                     :: dtmp1, dtmp2
  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(dtmp1, dtmp2) &
  !$omp reduction(+: paap)
  do i = 1, varsBLOCK
#else
  do i = varsBLOCK
#endif   
     DO c = 1, 3
     DO s2 = 1, 2
     DO s1 = 1, 2
     dtmp1 = Re(Spincol(p, s1, s2, c, i))
     dtmp2 = Im(Spincol(p, s1, s2, c, i))

     dtmp1 = dtmp1 * Re(Spincol(aap, s1, s2, c, i))
     dtmp2 = dtmp2 * Im(Spincol(aap, s1, s2, c, i))

     paap = paap + dtmp1
     paap = paap + dtmp2

     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)

  !// third loop in cg:
  !//
  !// rtr = ZERO
  !// do i = 1, size_sc_field
  !//    x(i) = x(i) + ak * p(i)
  !//    r(i) = r(i) - ak * aap(i)
  !//    dtmp = r(i) 
  !//    rtr = rtr + dtmp**2
  !// enddo 

  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                       :: 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(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
     Spincol(x, s1, s2, c, i) = Spincol(x, s1, s2, c, i) + ak * Spincol(p, s1, s2, c, i)
     Spincol(r, s1, s2, c, i) = Spincol(r, s1, s2, c, i) - ak * Spincol(aap, s1, s2, c, i)
     dtmp1 = Re(Spincol(r, s1, s2, c, i))
     dtmp2 = Im(Spincol(r, s1, s2, c, i))
     dtmp1 = dtmp1**2
     dtmp2 = dtmp2**2
     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

  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
     Spincol(inout, s1, s2, c, i) = Spincol(inout, s1, s2, c, i) + a * 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

  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
     Spincol(inout, s1, s2, c, i) = b * Spincol(inout, s1, s2, c, i) + 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

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

  integer :: i
  DOUBLE  :: dsum, dtmp1, dtmp2

  TRACE_ENTRY

  dsum = ZERO

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, in) &
  !$omp private(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
     dtmp1 = Re(Spincol(in, s1, s2, c, i))
     dtmp2 = Im(Spincol(in, s1, s2, c, i))
     dtmp1 = dtmp1**2
     dtmp2 = dtmp2**2
     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)

  !// second loop in cg with complex(!) result:
  !//
  !// paap = ZERO
  !// do i = 1, size_sc_field
  !//    paap = paap + conjg(p(i)) * aap(i)
  !// enddo

  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                       :: dtmp1, dtmp2
  integer                    :: i

  TRACE_ENTRY

  dtmp1 = ZERO
  dtmp2 = ZERO

#ifndef BLOCK
  !$omp parallel do &
  !$omp default(none) &
  !$omp shared(volh, p, aap) &
  !$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
     dtmp1 = dtmp1 + Re(Spincol(p, s1, s2, c, i)) * Re(Spincol(aap, s1, s2, c, i))
     dtmp1 = dtmp1 + Im(Spincol(p, s1, s2, c, i)) * Im(Spincol(aap, s1, s2, c, i))

     dtmp2 = dtmp2 + Re(Spincol(p, s1, s2, c, i)) * Im(Spincol(aap, s1, s2, c, i))
     dtmp2 = dtmp2 - Im(Spincol(p, s1, s2, c, i)) * Re(Spincol(aap, s1, s2, c, i))
     ENDDO
     ENDDO
     ENDDO
  enddo

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

  TRACE_RETURN
end

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