Merge pull request #10575 from rapidsai/branch-22.04

[gpuCI] Forward-merge branch-22.04 to branch-22.06 [skip gpuci]
rapidsai · Apr 1, 2022 · f66c99a · f66c99a
2 parents ca952f8 + 2c81bed
commit f66c99a
Show file tree

Hide file tree

Showing 2 changed files with 66 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 getComputeCapabilityMajor() 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 getComputeCapabilityMinor() 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,34 @@ JNIEXPORT jint JNICALL Java_ai_rapids_cudf_Cuda_getNativeComputeMode(JNIEnv *env
   CATCH_STD(env, -2);
 }
 
+JNIEXPORT jint JNICALL Java_ai_rapids_cudf_Cuda_getComputeCapabilityMajor(JNIEnv *env, jclass) {
+  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, ::cudaDevAttrComputeCapabilityMajor, device));
+    return attribute_value;
+  }
+  CATCH_STD(env, -2);
+}
+
+JNIEXPORT jint JNICALL Java_ai_rapids_cudf_Cuda_getComputeCapabilityMinor(JNIEnv *env, jclass) {
+  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, ::cudaDevAttrComputeCapabilityMinor, device));
+    return attribute_value;
+  }
+  CATCH_STD(env, -2);
+}
+
 JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cuda_freeZero(JNIEnv *env, jclass) {
   try {
     cudf::jni::auto_set_device(env);