unified test

juda · May 17, 2022 · df7708f · df7708f
1 parent 754c83e
commit df7708f
Show file tree

Hide file tree

Showing 3 changed files with 403 additions and 2 deletions.
diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc
@@ -28,7 +28,6 @@
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/index_map.h>
 
-#include <algorithm>
 #include <cmath>
 #include <string>
 #include <utility>

diff --git a/src/tir/transforms/lower_warp_memory.cc b/src/tir/transforms/lower_warp_memory.cc
@@ -472,7 +472,7 @@ Pass LowerWarpMemory() {
   auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
     auto* n = f.CopyOnWrite();
     auto target = f->GetAttr<Target>(tvm::attr::kTarget);
-    int warp_size = 32;
+    int warp_size = target.value()->GetAttr<Integer>("thread_warp_size", 1).value();
     WarpMemoryRewriter warp_memory_rewriter(warp_size);
     auto stmt = warp_memory_rewriter.Rewrite(std::move(n->body));
     n->body = UpdatePointerStorageScope(warp_memory_rewriter.new_storage_scopes_)(stmt);