Make various updates and fixes: (#164)

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

csrc/jit/compiler.hpp

@@ -35,10 +35,10 @@ public:
     }
 
     static void prepare_init(const std::string& library_root_path,
-                             const std::string& cuda_home_path_by_torch) {
+                             const std::string& cuda_home_path_by_python) {
         Compiler::library_root_path = library_root_path;
         Compiler::library_include_path = Compiler::library_root_path / "include";
-        Compiler::cuda_home = cuda_home_path_by_torch;
+        Compiler::cuda_home = cuda_home_path_by_python;
         Compiler::library_version = get_library_version();
     }
 
@@ -64,6 +64,8 @@ public:
                             get_env<int>("DG_JIT_CPP_STANDARD", 20));
         if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0))
             flags += " --ptxas-options=--verbose";
+        if (get_env("DG_JIT_WITH_LINEINFO", 0))
+            flags += " -Xcompiler -rdynamic -lineinfo";
     }
 
     virtual ~Compiler() = default;

csrc/jit/handle.hpp

@@ -8,7 +8,7 @@
 
 namespace deep_gemm {
 
-#if CUDART_VERSION >= 12080 or defined(DG_JIT_USE_DRIVER_API)
+#if CUDART_VERSION >= 12080 and not defined(DG_JIT_USE_DRIVER_API)
 
 // Use CUDA runtime API
 using LibraryHandle = cudaLibrary_t;

csrc/jit/kernel_runtime.hpp

@@ -64,8 +64,8 @@ public:
         kernel = load_kernel(cubin_path, symbol_names[0], &library);
     }
 
-    static void prepare_init(const std::string& cuda_home_path_by_torch) {
-        cuda_home = cuda_home_path_by_torch;
+    static void prepare_init(const std::string& cuda_home_path_by_python) {
+        cuda_home = cuda_home_path_by_python;
     }
 
     static bool check_validity(const std::filesystem::path& dir_path) {