[TOPI] Rewrite GPU argwhere using exclusive scan (apache#7314)

* use ex scan to write argwhere * add doc
neo-ai · Mar 2, 2021 · 160036f · 160036f
1 parent f5c9b84
commit 160036f
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 428 deletions.
diff --git a/python/tvm/contrib/nvcc.py b/python/tvm/contrib/nvcc.py
@@ -186,7 +186,7 @@ def find_libdevice_path(arch):
     selected_ver = 0
     selected_path = None
     cuda_ver = get_cuda_version(cuda_path)
-    if cuda_ver in (9.0, 9.1, 10.0, 10.1, 10.2, 11.0, 11.1):
+    if cuda_ver in (9.0, 9.1, 10.0, 10.1, 10.2, 11.0, 11.1, 11.2):
         path = os.path.join(lib_path, "libdevice.10.bc")
     else:
         for fn in os.listdir(lib_path):