Also check cuda runtime version when using the ASYNC allocator (#4850)

Signed-off-by: Rong Ou <rong.ou@gmail.com>
NVIDIA · Feb 24, 2022 · d0e14f3 · d0e14f3
1 parent d0c1ccb
commit d0e14f3
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala
@@ -1492,8 +1492,9 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
 
   lazy val rmmPool: String = {
     val pool = get(RMM_POOL)
-    if ("ASYNC".equalsIgnoreCase(pool) && Cuda.getDriverVersion < 11020) {
-      logWarning("CUDA driver does not support the ASYNC allocator, falling back to ARENA")
+    if ("ASYNC".equalsIgnoreCase(pool) &&
+        (Cuda.getRuntimeVersion < 11020 || Cuda.getDriverVersion < 11020)) {
+      logWarning("CUDA runtime/driver does not support the ASYNC allocator, falling back to ARENA")
       "ARENA"
     } else {
       pool