ApplyPhaseC2C.hpp

Go to the documentation of this file.

//////////////////////////////////////////////////////////////////////////////////////
// This file is distributed under the University of Illinois/NCSA Open Source License.
// See LICENSE file in top directory for details.
//
// Copyright (c) 2023 QMCPACK developers.
//
// File developed by: Ye Luo, yeluo@anl.gov, Argonne National Laboratory
//
// File created by: Ye Luo, yeluo@anl.gov, Argonne National Laboratory
//////////////////////////////////////////////////////////////////////////////////////


#include "OMPTarget/OMPTargetMath.hpp"

namespace qmcplusplus
{
namespace C2C
{
template<typename ST, typename TT>
inline void assign_v(ST x,
                     ST y,
                     ST z,
                     TT* restrict results_scratch_ptr,
                     const ST* restrict offload_scratch_ptr,
                     const ST* restrict myKcart_ptr,
                     size_t myKcart_padded_size,
                     size_t first_spo,
                     int index)
{
  const ST* restrict kx = myKcart_ptr;
  const ST* restrict ky = myKcart_ptr + myKcart_padded_size;
  const ST* restrict kz = myKcart_ptr + myKcart_padded_size * 2;

  const ST* restrict val = offload_scratch_ptr;
  TT* restrict psi       = results_scratch_ptr;

  //phase
  ST s, c, p = -(x * kx[index] + y * ky[index] + z * kz[index]);
  omptarget::sincos(p, &s, &c);

  const ST val_r         = val[index * 2];
  const ST val_i         = val[index * 2 + 1];
  psi[first_spo + index] = TT(val_r * c - val_i * s, val_i * c + val_r * s);
}

/** assign_vgl
   */
template<typename ST, typename TT>
inline void assign_vgl(ST x,
                       ST y,
                       ST z,
                       TT* restrict results_scratch_ptr,
                       size_t orb_padded_size,
                       const ST* mKK_ptr,
                       const ST* restrict offload_scratch_ptr,
                       size_t spline_padded_size,
                       const ST G[9],
                       const ST* myKcart_ptr,
                       size_t myKcart_padded_size,
                       size_t first_spo,
                       int index)
{
  constexpr ST two(2);
  const ST &g00 = G[0], &g01 = G[1], &g02 = G[2], &g10 = G[3], &g11 = G[4], &g12 = G[5], &g20 = G[6], &g21 = G[7],
           &g22 = G[8];

  const ST* restrict k0 = myKcart_ptr;
  const ST* restrict k1 = myKcart_ptr + myKcart_padded_size;
  const ST* restrict k2 = myKcart_ptr + myKcart_padded_size * 2;

  const ST* restrict val   = offload_scratch_ptr + spline_padded_size * SoAFields3D::VAL;
  const ST* restrict g0    = offload_scratch_ptr + spline_padded_size * SoAFields3D::GRAD0;
  const ST* restrict g1    = offload_scratch_ptr + spline_padded_size * SoAFields3D::GRAD1;
  const ST* restrict g2    = offload_scratch_ptr + spline_padded_size * SoAFields3D::GRAD2;
  const ST* restrict lcart = offload_scratch_ptr + spline_padded_size * SoAFields3D::LAPL;

  const size_t jr = index << 1;
  const size_t ji = jr + 1;

  const ST kX    = k0[index];
  const ST kY    = k1[index];
  const ST kZ    = k2[index];
  const ST val_r = val[jr];
  const ST val_i = val[ji];

  //phase
  ST s, c, p = -(x * kX + y * kY + z * kZ);
  omptarget::sincos(p, &s, &c);

  //dot(PrimLattice.G,myG[j])
  const ST dX_r = g00 * g0[jr] + g01 * g1[jr] + g02 * g2[jr];
  const ST dY_r = g10 * g0[jr] + g11 * g1[jr] + g12 * g2[jr];
  const ST dZ_r = g20 * g0[jr] + g21 * g1[jr] + g22 * g2[jr];

  const ST dX_i = g00 * g0[ji] + g01 * g1[ji] + g02 * g2[ji];
  const ST dY_i = g10 * g0[ji] + g11 * g1[ji] + g12 * g2[ji];
  const ST dZ_i = g20 * g0[ji] + g21 * g1[ji] + g22 * g2[ji];

  // \f$\nabla \psi_r + {\bf k}\psi_i\f$
  const ST gX_r = dX_r + val_i * kX;
  const ST gY_r = dY_r + val_i * kY;
  const ST gZ_r = dZ_r + val_i * kZ;
  const ST gX_i = dX_i - val_r * kX;
  const ST gY_i = dY_i - val_r * kY;
  const ST gZ_i = dZ_i - val_r * kZ;

  const ST lap_r = lcart[jr] + mKK_ptr[index] * val_r + two * (kX * dX_i + kY * dY_i + kZ * dZ_i);
  const ST lap_i = lcart[ji] + mKK_ptr[index] * val_i - two * (kX * dX_r + kY * dY_r + kZ * dZ_r);

  TT* restrict psi    = results_scratch_ptr;
  TT* restrict dpsi_x = results_scratch_ptr + orb_padded_size;
  TT* restrict dpsi_y = results_scratch_ptr + orb_padded_size * 2;
  TT* restrict dpsi_z = results_scratch_ptr + orb_padded_size * 3;
  TT* restrict d2psi  = results_scratch_ptr + orb_padded_size * 4;

  const size_t psiIndex = first_spo + index;
  psi[psiIndex]         = TT(c * val_r - s * val_i, c * val_i + s * val_r);
  d2psi[psiIndex]       = TT(c * lap_r - s * lap_i, c * lap_i + s * lap_r);
  dpsi_x[psiIndex]      = TT(c * gX_r - s * gX_i, c * gX_i + s * gX_r);
  dpsi_y[psiIndex]      = TT(c * gY_r - s * gY_i, c * gY_i + s * gY_r);
  dpsi_z[psiIndex]      = TT(c * gZ_r - s * gZ_i, c * gZ_i + s * gZ_r);
}
} // namespace C2C
} // namespace qmcplusplus