Add cudaDeviceProp and compute capability numbers into cuda handle

Author: Ziminli

src/devices/cuda/cuda_handle.cc

@@ -23,7 +23,25 @@ infiniopStatus_t createCudaHandle(CudaHandle_t *handle_ptr, int device_id) {
     checkCudnnError(cudnnCreate(&cudnn_handle));
     cudnn_pool->push(std::move(cudnn_handle));
 
-    *handle_ptr = new CudaContext{DevNvGpu, device_id, std::move(pool), std::move(cudnn_pool)};
+    // set CUDA device property
+    cudaDeviceProp prop;
+    cudaGetDeviceProperties(&prop, device_id);
+
+    // set device compute capability numbers
+    int capability_major;
+    int capability_minor;
+    cudaDeviceGetAttribute(&capability_major, cudaDevAttrComputeCapabilityMajor, device_id);
+    cudaDeviceGetAttribute(&capability_minor, cudaDevAttrComputeCapabilityMinor, device_id);
+
+    *handle_ptr = new CudaContext{
+        DevNvGpu,
+        device_id,
+        std::move(pool),
+        std::move(cudnn_pool),
+        std::move(prop),
+        capability_major,
+        capability_minor,
+    };
 
     return STATUS_SUCCESS;
 }

src/devices/cuda/cuda_handle.h

@@ -15,6 +15,9 @@ struct CudaContext {
     int device_id;
     std::shared_ptr<Pool<cublasHandle_t>> cublas_handles_t;
     std::shared_ptr<Pool<cudnnHandle_t>> cudnn_handles_t;
+    cudaDeviceProp prop;
+    int compute_capability_major;
+    int compute_capability_minor;
 };
 typedef struct CudaContext *CudaHandle_t;
 
@@ -35,12 +38,13 @@ void use_cublas(std::shared_ptr<Pool<cublasHandle_t>> cublas_handles_t, int devi
 }
 
 template<typename T>
-cudnnStatus_t use_cudnn(std::shared_ptr<Pool<cudnnHandle_t>> cudnn_handles_t, int device_id, T const &f) {
+cudnnStatus_t use_cudnn(std::shared_ptr<Pool<cudnnHandle_t>> cudnn_handles_t, int device_id, cudaStream_t stream, T const &f) {
     auto handle = cudnn_handles_t->pop();
     if (!handle) {
         cudaSetDevice(device_id);
         cudnnCreate(&(*handle));
     }
+    cudnnSetStream(*handle, stream);
     cudnnStatus_t status = f(*handle);
     cudnn_handles_t->push(std::move(*handle));
     return status;

src/ops/expand/cuda/expand.cc

@@ -22,9 +22,6 @@ infiniopStatus_t cudaCreateExpandDescriptor(CudaHandle_t handle,
         x_strides[i] = (i < ndim - x->ndim || y->shape[i] != x->shape[i + x->ndim - ndim]) ? 0 : x->strides[i + x->ndim - ndim];
     }
 
-    cudaDeviceProp prop;
-    cudaGetDeviceProperties(&prop, handle->device_id);
-
     int64_t *x_strides_d, *y_strides_d;
     char *strides_and_shape_d;
     checkCudaErrorWithCode(cudaMalloc(&strides_and_shape_d, ndim * (2 * sizeof(int64_t) + sizeof(uint64_t))), STATUS_MEMORY_NOT_ALLOCATED);
@@ -38,7 +35,7 @@ infiniopStatus_t cudaCreateExpandDescriptor(CudaHandle_t handle,
         handle->device_id,
         ndim,
         y_data_size,
-        static_cast<uint64_t>(prop.maxGridSize[0]),
+        static_cast<uint64_t>(handle->prop.maxGridSize[0]),
         strides_and_shape_d,
     };