add tensorcore support

1 files changed, 14 insertions, 4 deletions

diff --git a/src/amun/gpu/mblas/handles.cu b/src/amun/gpu/mblas/handles.cu
index 8f17e232..13c99f8c 100644
--- a/src/amun/gpu/mblas/handles.cu
+++ b/src/amun/gpu/mblas/handles.cu
@@ -23,14 +23,24 @@ CublasHandler::CublasHandler()
   cublasStatus_t stat;
   stat = cublasCreate(&handle_);
   if (stat != CUBLAS_STATUS_SUCCESS) {
-  printf ("cublasCreate initialization failed\n");
-  abort();
+    printf ("cublasCreate initialization failed\n");
+    abort();
   }
 
+#if CUDA_VERSION >= 9000
+  ///*
+    stat = cublasSetMathMode(handle_, CUBLAS_TENSOR_OP_MATH);
+    if (stat != CUBLAS_STATUS_SUCCESS) {
+      printf ("cublasSetMathMode failed\n");
+      abort();
+    }
+  //*/
+#endif
+		  
   stat = cublasSetStream(handle_, CudaStreamHandler::GetStream());
   if (stat != CUBLAS_STATUS_SUCCESS) {
-  printf ("cublasSetStream initialization failed\n");
-  abort();
+    printf ("cublasSetStream initialization failed\n");
+    abort();
   }
 }