Merge pull request #157 from mit-han-lab/dev

Fix missing third_party & merge main branch of dev repo

setup.py

@@ -129,7 +129,7 @@ def cond(s) -> list:
     for target in sm_targets:
         NVCC_FLAGS += ["-gencode", f"arch=compute_{target},code=sm_{target}"]
 
-    NVCC_MSVC_FLAGS = ["-Xcompiler", "/Zc:__cplusplus", "-Xcompiler", "/FS"]
+    NVCC_MSVC_FLAGS = ["-Xcompiler", "/Zc:__cplusplus", "-Xcompiler", "/FS", "-Xcompiler", "/bigobj"]
 
     nunchaku_extension = CUDAExtension(
         name="nunchaku._C",

src/Module.h

@@ -215,6 +215,11 @@ struct LayerOffloadHelper {
         if (offload) {
             streamCompute = std::make_unique<CUDAStreamWrapper>();
             streamLoad = std::make_unique<CUDAStreamWrapper>();
+
+            needWorkaround = checkWorkaround();
+            if (needWorkaround) {
+                spdlog::debug("Offloading helper: use WDDM workaround");
+            }
         }
     }
 
@@ -240,6 +245,7 @@ struct LayerOffloadHelper {
                 funcCompute(layer);
                 nextComputeDone = std::make_unique<CUDAEventWrapper>();
                 checkCUDA(cudaEventRecord(nextComputeDone->event, getCurrentCUDAStream()));
+                workaroundFlush();
             }
 
             {
@@ -253,10 +259,13 @@ struct LayerOffloadHelper {
                 }
                 nextLoadDone = std::make_unique<CUDAEventWrapper>();
                 checkCUDA(cudaEventRecord(nextLoadDone->event, getCurrentCUDAStream()));
+                workaroundFlush();
             }
 
             eventComputeDone = std::move(nextComputeDone);
             eventLoadDone = std::move(nextLoadDone);
+
+            workaroundSynchronize();
         }
     }
 
@@ -266,4 +275,35 @@ struct LayerOffloadHelper {
         }
         checkCUDA(cudaStreamWaitEvent(getCurrentCUDAStream(), event->event));
     }
+
+    // WDDM prevents multiple streams run concurrently
+    // use flush and synchronize to work around
+    bool needWorkaround;
+    static bool checkWorkaround() {
+        if (char *env = getenv("NUNCHAKU_OFFLOAD_WDDM_WORKAROUND")) {
+            if (std::string(env) == "1") {
+                return true;
+            } else if (std::string(env) == "0") {
+                return false;
+            }
+        }
+
+    #ifdef _WIN32
+        return true;
+    #else
+        return false;
+    #endif
+    }
+    void workaroundFlush() {
+        if (!needWorkaround) {
+            return;
+        }
+        cudaStreamQuery(getCurrentCUDAStream());
+    }
+    void workaroundSynchronize() {
+        if (!needWorkaround) {
+            return;
+        }
+        checkCUDA(cudaEventSynchronize(eventComputeDone->event));
+    }
 };