From b6ca4c8ff28096116de0542c58f2d7d86bcc8c2a Mon Sep 17 00:00:00 2001
From: Alejandro Gallo <aamsgallo@gmail.com>
Date: Wed, 27 Jul 2022 18:05:20 +0200
Subject: [PATCH] Update autoconf tests

---
 bench/Makefile.am      |  2 +-
 configure.ac           |  3 +-
 cuda.nix               |  2 +-
 etc/m4/atrip_cublas.m4 | 81 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 etc/m4/atrip_cublas.m4

diff --git a/bench/Makefile.am b/bench/Makefile.am
index d6e591e..39ce3c3 100644
--- a/bench/Makefile.am
+++ b/bench/Makefile.am
@@ -18,5 +18,5 @@ endif
 
 if WITH_CUDA
 test_main_CXXFLAGS = $(CUDA_CXXFLAGS)
-test_main_LDADD += $(CUDA_LDFLAGS) -lcublas -lcudadevrt -lthrust
+test_main_LDADD += $(CUDA_LDFLAGS)
 endif
diff --git a/configure.ac b/configure.ac
index cb25ce6..cc7293c 100644
--- a/configure.ac
+++ b/configure.ac
@@ -55,7 +55,7 @@ AC_ARG_ENABLE([cuda],
                               [WITH_CUDA=yes],
                               [WITH_CUDA=no])
 AC_ARG_VAR([NVCC], [Path to the nvidia cuda compiler.])
-AC_ARG_VAR([CUDA_LDFLAGS], [LDFLAGS to find libraries -lcuda and -lcudart.])
+AC_ARG_VAR([CUDA_LDFLAGS], [LDFLAGS to find libraries -lcuda, -lcudart, -lcublas.])
 AC_ARG_VAR([CUDA_CXXFLAGS], [CXXFLAGS to find the CUDA headers])
 
 
@@ -145,6 +145,7 @@ if test x${WITH_CUDA} = xyes; then
   AC_DEFINE([HAVE_CUDA],1,[Wether we are using CUDA])
   # TODO: make sure to find cuda and cudart
   # AC_SEARCH_LIBS([cudaMalloc], [cuda cudart], [FOUND_CUDA=yes])
+  ATRIP_CUBLAS
 fi
 
 dnl -----------------------------------------------------------------------
diff --git a/cuda.nix b/cuda.nix
index cf5e454..9daaa5c 100644
--- a/cuda.nix
+++ b/cuda.nix
@@ -13,7 +13,7 @@
 
     # in here is libcuda
     export CUDA_X11LIB="${pkgs.linuxPackages.nvidia_x11}/lib"
-    export CUDA_LDFLAGS="-L$CUDA_X11LIB -L$CUDA_ROOT_PATH/lib -L$CUDA_LIB_PATH/lib -lcuda -lcudart"
+    export CUDA_LDFLAGS="-L$CUDA_X11LIB -L$CUDA_ROOT_PATH/lib -L$CUDA_LIB_PATH/lib -lcuda -lcudart -lcublas"
   '';
 
 }
diff --git a/etc/m4/atrip_cublas.m4 b/etc/m4/atrip_cublas.m4
new file mode 100644
index 0000000..cfc182c
--- /dev/null
+++ b/etc/m4/atrip_cublas.m4
@@ -0,0 +1,81 @@
+
+
+dnl Create some tests for the compilers to work with CUDA
+
+
+
+
+
+
+
+AC_DEFUN([ATRIP_CUBLAS],
+[
+
+AC_LANG_PUSH([C++])dnl
+atrip_success=no
+ac_save_CXX="$CXX"
+ac_save_CXXFLAGS="$CXXFLAGS"
+ac_save_LDFLAGS="$LDFLAGS"
+
+
+CXXFLAGS="${CXXFLAGS} $CUDA_CXXFLAGS"
+LDFLAGS="${LDFLAGS} $CUDA_LDFLAGS"
+
+AC_MSG_CHECKING([that cublas works with $CXX])
+
+AC_COMPILE_IFELSE([AC_LANG_SOURCE([_ATRIP_CUBLAS_SOURCE])],
+	[
+	atrip_success=yes
+	AC_MSG_RESULT([yes])
+	],
+	[
+	atrip_success=no
+	AC_MSG_ERROR([Does not work!])
+	])
+
+CXX="$ac_save_CXX"
+CXXFLAGS="$ac_save_CXXFLAGS"
+LDFLAGS="$ac_save_LDFLAGS"
+
+AC_LANG_POP([C++])dnl
+])dnl DEFUN
+
+m4_define([_ATRIP_CUBLAS_SOURCE], [[
+#include <mpi.h>
+#include <mpi.h>
+#include <iostream>
+#include <string.h>
+#include <cublas_v2.h>
+#include <cuda.h>
+
+int main() {
+  MPI_Init(NULL, NULL);
+  int rank, np;
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm_size(MPI_COMM_WORLD, &np);
+
+  cublasStatus_t stat;
+  cublasHandle_t handle;
+  stat = cublasCreate(&handle);
+  cudaError_t error;
+
+  int64_t els = 100, oo, No, Nv;
+  double *PH_d, *HHP_d, *HHH_d, *F_d;
+  double one(1.0);
+
+  cuMemAlloc((CUdeviceptr*)&PH_d, els*sizeof(double));
+  cuMemAlloc((CUdeviceptr*)&HHP_d, els*sizeof(double));
+  cuMemAlloc((CUdeviceptr*)&HHH_d, els*sizeof(double));
+  cuMemAlloc((CUdeviceptr*)&F_d, els*sizeof(double));
+
+  stat = cublasDgemm(handle,
+	 CUBLAS_OP_N,
+	 CUBLAS_OP_N,
+	 oo, No, Nv,
+	 &one,
+	 HHP_d, oo, PH_d, Nv, &one, F_d, oo);
+  //cudaSetDevice(rank);
+
+  return 0;
+}
+]])