JNI Bindings to fetch CUDA compute capability versions.

Signed-off-by: MithunR <[email protected]>
rapidsai · Apr 1, 2022 · dfd9975 · dfd9975
1 parent a02b7c2
commit dfd9975
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 2 deletions.
diff --git a/java/src/main/java/ai/rapids/cudf/Cuda.java b/java/src/main/java/ai/rapids/cudf/Cuda.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -386,6 +386,42 @@ static void asyncMemcpy(long dst, long src, long count, CudaMemcpyKind kind) {
    */
   static native int getNativeComputeMode() throws CudaException;
 
+  /**
+   * Gets the major CUDA compute capability of the current device.
+   * 
+   * For reference: https://developer.nvidia.com/cuda-gpus
+   * Hardware Generation	Compute Capability
+   *     Ampere	                8.x
+   *     Turing	                7.5
+   *     Volta	                7.0, 7.2
+   *     Pascal	                6.x
+   *     Maxwell                5.x
+   *     Kepler	                3.x
+   *     Fermi	                2.x
+   * 
+   * @return The Major compute capability version number of the current CUDA device
+   * @throws CudaException on any error
+   */
+  public static native int getCurrentComputeCapabilityMajor() throws CudaException;  
+
+  /**
+   * Gets the minor CUDA compute capability of the current device.
+   * 
+   * For reference: https://developer.nvidia.com/cuda-gpus
+   * Hardware Generation	Compute Capability
+   *     Ampere	                8.x
+   *     Turing	                7.5
+   *     Volta	                7.0, 7.2
+   *     Pascal	                6.x
+   *     Maxwell                5.x
+   *     Kepler	                3.x
+   *     Fermi	                2.x
+   * 
+   * @return The Minor compute capability version number of the current CUDA device
+   * @throws CudaException on any error
+   */
+  public static native int getCurrentComputeCapabilityMinor() throws CudaException;
+
   /**
    * Calls cudaFree(0). This can be used to initialize the GPU after a setDevice()
    * @throws CudaException on any error

diff --git a/java/src/main/native/src/CudaJni.cpp b/java/src/main/native/src/CudaJni.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -195,6 +195,32 @@ JNIEXPORT jint JNICALL Java_ai_rapids_cudf_Cuda_getNativeComputeMode(JNIEnv *env
   CATCH_STD(env, -2);
 }
 
+namespace {
+
+jint cuda_device_get_attribute_impl(JNIEnv *env, cudaDeviceAttr const &attribute_label) {
+  try {
+    cudf::jni::auto_set_device(env);
+    int device;
+    JNI_CUDA_TRY(env, -2, ::cudaGetDevice(&device));
+    int attribute_value;
+    JNI_CUDA_TRY(env, -2, ::cudaDeviceGetAttribute(&attribute_value, attribute_label, device));
+    return attribute_value;
+  }
+  CATCH_STD(env, -2);
+}
+
+} // namespace
+
+JNIEXPORT jint JNICALL Java_ai_rapids_cudf_Cuda_getCurrentComputeCapabilityMajor(JNIEnv *env,
+                                                                                 jclass) {
+  return cuda_device_get_attribute_impl(env, ::cudaDevAttrComputeCapabilityMajor);
+}
+
+JNIEXPORT jint JNICALL Java_ai_rapids_cudf_Cuda_getCurrentComputeCapabilityMinor(JNIEnv *env,
+                                                                                 jclass) {
+  return cuda_device_get_attribute_impl(env, ::cudaDevAttrComputeCapabilityMinor);
+}
+
 JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cuda_freeZero(JNIEnv *env, jclass) {
   try {
     cudf::jni::auto_set_device(env);