Make various updates and fixes: (#164)

- Add BF16 support for SM90 and SM100
- Refactor Python APIs
- Other fixes and code refactoring

deep_gemm/__init__.py

@@ -1,6 +1,5 @@
 import os
-import torch
-import torch.utils.cpp_extension
+import subprocess
 
 # Set some default environment provided at setup
 try:
@@ -12,14 +11,8 @@ try:
 except ImportError:
     pass
 
-# Import functions from the CPP module
-import deep_gemm_cpp
-deep_gemm_cpp.init(
-    os.path.dirname(os.path.abspath(__file__)), # Library root directory path
-    torch.utils.cpp_extension.CUDA_HOME         # CUDA home
-)
-
 # Configs
+import deep_gemm_cpp
 from deep_gemm_cpp import (
     set_num_sms,
     get_num_sms,
@@ -34,13 +27,48 @@ from deep_gemm_cpp import (
     fp8_gemm_tn, fp8_gemm_tt,
     m_grouped_fp8_gemm_nt_contiguous,
     m_grouped_fp8_gemm_nn_contiguous,
-    fp8_m_grouped_gemm_nt_masked,
+    m_grouped_fp8_gemm_nt_masked,
     k_grouped_fp8_gemm_tn_contiguous,
+    # BF16 GEMMs
+    bf16_gemm_nt, bf16_gemm_nn,
+    bf16_gemm_tn, bf16_gemm_tt,
+    m_grouped_bf16_gemm_nt_contiguous,
+    m_grouped_bf16_gemm_nt_masked,
     # Layout kernels
     transform_sf_into_required_layout
 )
 
+# Some alias for legacy supports
+# TODO: remove these later
+fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked
+bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked
+
 # Some utils
 from . import testing
 from . import utils
 from .utils import *
+
+
+# Initialize CPP modules
+def _find_cuda_home() -> str:
+    # TODO: reuse PyTorch API later
+    # For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks
+    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+    if cuda_home is None:
+        # noinspection PyBroadException
+        try:
+            with open(os.devnull, 'w') as devnull:
+                nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n')
+                cuda_home = os.path.dirname(os.path.dirname(nvcc))
+        except Exception:
+            cuda_home = '/usr/local/cuda'
+            if not os.path.exists(cuda_home):
+                cuda_home = None
+    assert cuda_home is not None
+    return cuda_home
+
+
+deep_gemm_cpp.init(
+    os.path.dirname(os.path.abspath(__file__)), # Library root directory path
+    _find_cuda_home()                           # CUDA home
+)