# Copyright (C) 2001-2016 Quantum ESPRESSO Foundation
#####
#
# SYNOPSIS
#
# AX_CHECK_CUDA
#
# DESCRIPTION
#
# Figures out if CUDA Driver API/nvcc is available, i.e. existence of:
#   nvcc
#   cuda.h
#   libcuda.a
#
# If something isn't found, fails straight away.
#
# The following variables are substituted in the makefile:
# NVCC        : the nvcc compiler command.
# NVCCFLAGS   : nvcc specific flags
# CUDA_CFLAGS : CUDA includes
# CUDA_LDLIBS : CUDA libraries
#
# Defines HAVE_CUDA in config.h
#
# LICENCE
# Public domain
#
#####

AC_DEFUN([X_AC_QE_GPU], [

# Variables
NVCC=no
CUDA_CFLAGS=
CUDA_LDLIBS=
ROCM_CFLAGS=
ROCM_LDLIBS=
cuda_fflags=
cuda_intlibs=
cuda_extlibs=
cuda_int_line=
cuda_ext_line=
rocm_ext_line=

AC_ARG_WITH([cuda-cc],
   [AS_HELP_STRING([--with-cuda-cc=VAL],[GPU architecture (Kepler: 35, Pascal: 60, Volta: 70) @<:@default=35@:>@])],
   [],
   [with_cuda_cc=35])
   
AC_ARG_WITH([cuda-runtime],
   [AS_HELP_STRING([--with-cuda-runtime=VAL],[CUDA runtime (Pascal: 8+, Volta: 9+) @<:@default=none, checks if the NVHPC_CUDA_HOME variable is set.@:>@])],
   [],
   [with_cuda_runtime=none])

# 
AC_ARG_WITH([cuda-int-libs],
   [AS_HELP_STRING([--with-cuda-int-libs=VAL],[CUDA internal libraries () @<:@default=cuda,cufft,cublas,cusolver,cudart@:>@])],
   [],[with_cuda_int_libs=cufft,cublas,cusolver])

AC_ARG_ENABLE([cuda-env-check],
   [AS_HELP_STRING([--enable-cuda-env-check=yes],[The configure script will check CUDA installation and report problems @<:@default=yes@:>@])],
   [],
   [enable_cuda_env_check=yes])
# 
AC_ARG_ENABLE([cublas],
   [AS_HELP_STRING([--enable-cublas=yes],[Activates cuBLAS support])], [], [])
AC_ARG_ENABLE([rocblas],
   [AS_HELP_STRING([--enable-rocblas=yes],[Activates rocBLAS support])], [], [])
AC_ARG_ENABLE([mkl-gpu],
   [AS_HELP_STRING([--enable-mkl-gpu=yes],[Activates MKL-GPU support])], [], [])


#
# LinAlg defaults
#
if test "x$enable_cublas" = "xyes" ; then 
   if test "x$enable_rocblas" = "x" ; then enable_rocblas=no ; fi
   if test "x$enable_mkl_gpu" = "x" ; then enable_mkl_gpu=no ; fi
fi
if test "x$enable_rocblas" = "xyes" ; then 
   if test "x$enable_cublas"  = "x" ; then enable_cublas=no  ; fi
   if test "x$enable_mkl_gpu" = "x" ; then enable_mkl_gpu=no ; fi
fi
if test "x$enable_mkl_gpu" = "xyes" ; then 
   if test "x$enable_cublas"  = "x" ; then enable_cublas=no  ; fi
   if test "x$enable_rocblas" = "x" ; then enable_rocblas=no ; fi
fi

if test "x$enable_cuda_fortran" = "xyes" || test "x$enable_openacc" = "xyes" ||  test "x$enable_openmp5" = "xyes" ; then
   # -----------------------------------------
   # Setup CUDA paths
   # -----------------------------------------
   if test "x$LIBCUDA_LIBS" != "x"; then
     #
     # Mapping from yambo to QE syntax
     #
     CUDA_CFLAGS="$LIBCUDA_INCS"
     CUDA_LDLIBS="$LIBCUDA_LIBS"
   else
      AC_CHECK_FILE(/usr/local/cuda/,[CUDAPATH="/usr/local/cuda"],[])
      AC_CHECK_FILE(/usr/local/cuda/include,[CUDA_CFLAGS+=" -I/usr/local/cuda/include"],[CUDA_CFLAGS=""])
      AC_CHECK_FILE(/usr/local/cuda/lib64,[CUDA_LDLIBS+=" -L/usr/local/cuda/lib64"],[])
      case "${FCVERSION}" in
    	  *GNU* | *gnu* )
	          CUDA_LDLIBS+=" -lcuda -lcudart -lcublas -lcufft"
		  ;;
          *nvfortran* | *pgfortran* | *pgf* )
	          CUDA_LDLIBS+=" "
		  ;;
          *ifx* )
	          CUDA_LDLIBS+=" "
		  ;;
          *Cray* )
	          CUDA_LDLIBS+=" "
		  ;;
    	  * )
	          CUDA_LDLIBS+=" "
      esac
   fi
   #
   cuda_extlibs="$CUDA_LDLIBS"
fi

if test "x$enable_cuda_fortran" = "xyes"
then
   # -----------------------------------------
   # Check compiler is PGI
   # -----------------------------------------
   AC_LANG_PUSH([Fortran])
   AC_FC_SRCEXT([f90])
   AX_CHECK_COMPILE_FLAG([-cuda], [have_cudafor=yes], [have_cudafor=no], [], [MODULE test; use cudafor; END MODULE])
   AC_LANG_POP([Fortran])
   if test "x$have_cudafor" != "xyes"
   then
      AC_MSG_ERROR([You do not have the cudafor module. Are you using a PGI compiler?])
   fi

   # -----------------------------------------
   # Checking for nvcc
   # -----------------------------------------
   AC_PATH_PROG([NVCC],[nvcc],[no],[$PATH:$CUDAPATH/bin])
   if test "x$NVCC" = "xno" && test "x$enable_cuda_env_check" = "xyes"
   then
      AC_MSG_ERROR([Cannot find nvcc compiler. To enable CUDA, please add path to
                    nvcc in the PATH environment variable and/or specify the path
                    where CUDA is installed using: --with-cuda=PATH])
   fi


   # -----------------------------------------
   # Setup nvcc flags
   # -----------------------------------------
   AC_ARG_VAR(NVCCFLAGS,[Additional nvcc flags (example: NVCCFLAGS="-arch=compute_30 -code=sm_30")])
   if test x$DEBUG = xtrue
   then
      NVCCFLAGS+=" -g -arch=compute_$with_cuda_cc -code=sm_$with_cuda_cc"
   else
      NVCCFLAGS+=" -O3 -arch=compute_$with_cuda_cc -code=sm_$with_cuda_cc"
   fi
   

   # -----------------------------------------
   # Check if nvcc works
   # -----------------------------------------
   ac_compile_nvcc=no
   AC_MSG_CHECKING([whether nvcc works])
   cat>conftest.cu<<EOF
   __global__ static void test_cuda() {
      const int tid = threadIdx.x;
      const int bid = blockIdx.x;
      __syncthreads();
   }
EOF

   if test "x$NVCC" != "xno"
   then
      if $NVCC -c $NVCCFLAGS conftest.cu &> /dev/null
      then
         ac_compile_nvcc=yes
      fi
   fi
   rm -f conftest.cu conftest.o
   AC_MSG_RESULT([$ac_compile_nvcc])

   if test "x$ac_compile_nvcc" = "xno"
   then
      AC_MSG_WARN([CUDA compiler has problems.])
   fi


   # -----------------------------------------
   # Check for headers and libraries
   # -----------------------------------------
   ax_save_CXXFLAGS="${CXXFLAGS}"
   ax_save_LIBS="${LIBS}"
   

   FCFLAGS="$CUDA_CFLAGS $CXXFLAGS"
   LIBS="$CUDA_LDLIBS $LIBS"

   # And the header and the lib
   if test "x$enable_cuda_env_check" = "xyes" && test -z "$NVHPC"
   then
      AC_CHECK_LIB([cuda], [cuInit], [], AC_MSG_FAILURE([Couldn't find libcuda]))
      AC_CHECK_LIB([cudart], [cudaMalloc], [], AC_MSG_FAILURE([Couldn't find libcudart]))
      AC_CHECK_LIB([cublas], [cublasInit], [], AC_MSG_FAILURE([Couldn't find libcublas]))
      AC_CHECK_LIB([cufft], [cufftPlanMany], [], AC_MSG_FAILURE([Couldn't find libcufft]))
   fi

   
   # Returning to the original flags
   CXXFLAGS=${ax_save_CXXFLAGS}
   LIBS=${ax_save_LIBS}
fi

if test "x$enable_cuda_fortran" = "xyes" || test "x$enable_openacc" = "xyes" ||  test "x$enable_openmp5" = "xyes" ; then

   AC_DEFINE(HAVE_CUDA,1,[Define if we have CUDA])
   gpu_arch="$with_cuda_cc"
   gpu_runtime="$with_cuda_runtime"

   cuda_cflags=$CUDA_CFLAGS
   cuda_ext_line="CUDA_EXTLIBS=$CUDA_LDLIBS CUDA_EXTINCS=$CUDA_CFLAGS"

fi

if test "x$enable_cuda_fortran" = "xyes" ; then
   try_dflags="$try_dflags -D__DXL_CUDAF"
   if test "x$enable_cublas" = "x" ; then enable_cublas=yes ; fi
   if test "x$with_cuda_runtime" = "xnone" ; then
     cuda_fflags="-cuda -gpu=cc$with_cuda_cc"
   else
     cuda_fflags="-cuda -gpu=cc$with_cuda_cc,cuda$with_cuda_runtime"
   fi
   cuda_intlibs=""
   if test "x$with_cuda_int_libs" != "x"; then
     cuda_intlibs="-cudalib=$with_cuda_int_libs"
     cuda_int_line="CUDA_INTLIBS=$cuda_intlibs"
   fi
   if test "x$with_cuda_runtime" = "xnone" ; then
     ldflags="$ldflags -cuda -gpu=cc$with_cuda_cc $cuda_intlibs"
   else
     ldflags="$ldflags -cuda -gpu=cc$with_cuda_cc,cuda$with_cuda_runtime $cuda_intlibs"
   fi
fi

if test "x$enable_openacc" = "xyes"      ; then
   try_dflags="$try_dflags -D__DXL_OPENACC"
   if test "x$enable_cublas" = "x" ; then enable_cublas=yes ; fi
   case "${FCVERSION}" in
    *nvfortran* | *pgfortran* | *pgf* )
      if test "x$with_cuda_runtime" = "xnone" ; then
        cuda_fflags="-acc=gpu,multicore -acclibs -gpu=cc${with_cuda_cc} -cudalib=${with_cuda_int_libs}"
      else
        cuda_fflags="-acc=gpu,multicore -acclibs -gpu=cc${with_cuda_cc},cuda${with_cuda_runtime} -cudalib=${with_cuda_int_libs}"
      fi
      ;;
    *GNU* | *gnu*)
      cuda_fflags="-fopenacc"
   esac
   ldflags="$ldflags $cuda_fflags"
fi

if test "x$enable_openmp5" = "xyes"      ; then
   try_dflags="$try_dflags -D__DXL_OPENMP_GPU"
   #if test "x$enable_mkl_gpu" = "x" ; then enable_mkl_gpu=yes ; fi
   #
   # cuda_fflags="-mp gpu -ta=tesla:cc${with_cuda_cc},cuda${with_cuda_runtime} "
   #
   case "${FCVERSION}" in
    *ifx* )
      cuda_fflags="-fiopenmp -fopenmp-targets=spir64"
      cuda_ldflags="$cuda_fflags -fsycl -lsycl -lOpenCL -liomp5 -lpthread"
      ;;
    *"Cray Fortran"* )
      cuda_fflags="-fopenmp"
      cuda_ldflags="$cuda_fflags "
      ;;
    *nvfortran* | *pgfortran* | *pgf* )
      if test "x$with_cuda_runtime" = "xnone" ; then
        cuda_fflags="-mp=gpu -ta=tesla:cc${with_cuda_cc} -cudalib=$with_cuda_int_libs"
      else
        cuda_fflags="-mp=gpu -ta=tesla:cc${with_cuda_cc},cuda${with_cuda_runtime} -cudalib=$with_cuda_int_libs"
      fi
   esac
   ldflags="$ldflags $cuda_ldflags"
   #
   if test "x$enable_mkl_gpu" = "xyes" ; then 
     ldflags="$ldflags -lmkl_sycl"
   fi
   if test "x$enable_rocblas" = "xyes" ; then 
     #ldflags="$ldflags"
     rocm_ext_line="ROCM_EXTLIBS=$LIBROCM_LIBS ROCM_EXTINCS=$LIBROCM_INCS"
     rocm_extlibs="$LIBROCM_LIBS"
   fi
   # TODO
   echo "todo";
fi

if test "x$enable_openmp" = "xyes"      ; then
   if test "$f90" = "ifx" ; then
     ldflags="$ldflags -liomp5"
   fi
   ldflags="$ldflags -lpthread"
fi

#
# LinAlg libs
#
if test "x$enable_cublas" = "xyes" ; then 
   try_dflags="$try_dflags -D__DXL_CUBLAS"
fi
if test "x$enable_rocblas" = "xyes" ; then 
   try_dflags="$try_dflags -D__DXL_ROCBLAS"
fi
if test "x$enable_mkl_gpu" = "xyes" ; then 
   try_dflags="$try_dflags -D__DXL_MKL_GPU"
fi

# Announcing the new variables
# For C (maybe needed in the future)
AC_SUBST([NVCC])
AC_SUBST([NVCCFLAGS])
AC_SUBST([CUDA_CFLAGS])
AC_SUBST([CUDA_LDLIBS])
AC_SUBST([ROCM_CFLAGS])
AC_SUBST([ROCM_LDLIBS])
# And for Fortran
AC_SUBST(gpu_arch)
AC_SUBST(gpu_runtime)
AC_SUBST(cuda_fflags)
AC_SUBST(cuda_cflags)
AC_SUBST(cuda_intlibs)
AC_SUBST(cuda_extlibs)
AC_SUBST(cuda_int_line)
AC_SUBST(cuda_ext_line)
AC_SUBST(rocm_ext_line)
AC_SUBST(rocm_extlibs)
])
