arrayfire · 9prady9 · Jul 16, 2019 · Jul 15, 2019 · Jul 15, 2019
diff --git a/include/af/cuda.h b/include/af/cuda.h
@@ -10,13 +10,33 @@
 #pragma once
 #include <af/defines.h>
 #include <af/exception.h>
+
+/// This file contain functions that apply only to the CUDA backend. It will
+/// include cuda headers when it is built with NVCC. Otherwise the you can
+/// define the AF_DEFINE_CUDA_TYPES before including this file and it will
+/// define the cuda types used in this header.
+
+#ifdef __NVCC__
 #include <cuda.h>
 #include <cuda_runtime.h>
+#include <cublas_v2.h>
+#else
+#ifdef AF_DEFINE_CUDA_TYPES
+typedef struct CUstream_st *cudaStream_t;
+
+/*Enum for default math mode/tensor operation*/
+typedef enum {
+  CUBLAS_DEFAULT_MATH = 0,
+  CUBLAS_TENSOR_OP_MATH = 1
+} cublasMath_t;
+#endif
+#endif
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+
 #if AF_API_VERSION >= 31
 /**
    Get the stream for the CUDA device with \p id in ArrayFire context
@@ -55,6 +75,20 @@ AFAPI af_err afcu_get_native_id(int* nativeid, int id);
 AFAPI af_err afcu_set_native_id(int nativeid);
 #endif
 
+#if AF_API_VERSION >= 37
+/**
+    Sets the cuBLAS math mode for the internal handle
+
+    See the cuBLAS documentation for additional details
+
+    \param[in] mode The cublasMath_t type to set
+    \returns \ref af_err error code
+
+    \ingroup cuda_mat
+*/
+AFAPI af_err afcu_cublasSetMathMode(cublasMath_t mode);
+#endif
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/src/api/unified/CMakeLists.txt b/src/api/unified/CMakeLists.txt
@@ -9,6 +9,7 @@ target_sources(af
     ${CMAKE_CURRENT_SOURCE_DIR}/arith.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/blas.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/cuda.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/data.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/error.cpp

diff --git a/src/api/unified/cuda.cpp b/src/api/unified/cuda.cpp
@@ -0,0 +1,52 @@
+/*******************************************************
+ * Copyright (c) 2019, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include "symbol_manager.hpp"
+#include <af/backend.h>
+
+#define AF_DEFINE_CUDA_TYPES
+#include <af/cuda.h>
+
+af_err afcu_get_stream(cudaStream_t* stream, int id) {
+    af_backend backend;
+    af_get_active_backend(&backend);
+    if(backend == AF_BACKEND_CUDA) {
+        return CALL(stream, id);
+    }
+    return AF_ERR_NOT_SUPPORTED;
+}
+
+af_err afcu_get_native_id(int* nativeid, int id) {
+    af_backend backend;
+    af_get_active_backend(&backend);
+    if(backend == AF_BACKEND_CUDA) {
+        return CALL(nativeid, id);
+    }
+    return AF_ERR_NOT_SUPPORTED;
+}
+
+
+af_err afcu_set_native_id(int nativeid) {
+    af_backend backend;
+    af_get_active_backend(&backend);
+    if(backend == AF_BACKEND_CUDA) {
+        return CALL(nativeid);
+    }
+    return AF_ERR_NOT_SUPPORTED;
+}
+
+
+af_err afcu_cublasSetMathMode(cublasMath_t mode) {
+    af_backend backend;
+    af_get_active_backend(&backend);
+    if(backend == AF_BACKEND_CUDA) {
+        return CALL(mode);
+    }
+    return AF_ERR_NOT_SUPPORTED;
+}
diff --git a/src/backend/cuda/device_manager.cpp b/src/backend/cuda/device_manager.cpp
@@ -23,6 +23,7 @@
 #include <platform.hpp>
 #include <spdlog/spdlog.h>
 #include <version.hpp>
+#include <cublas_v2.h> // needed for af/cuda.h
 #include <af/cuda.h>
 #include <af/version.h>
 // cuda_gl_interop.h does not include OpenGL headers for ARM

diff --git a/src/backend/cuda/platform.cpp b/src/backend/cuda/platform.cpp
@@ -435,3 +435,11 @@ af_err afcu_set_native_id(int nativeid) {
     CATCHALL;
     return AF_SUCCESS;
 }
+
+af_err afcu_cublasSetMathMode(cublasMath_t mode) {
+    try {
+        CUBLAS_CHECK(cublasSetMathMode(cuda::blasHandle(), mode));
+    }
+    CATCHALL;
+    return AF_SUCCESS;
+}