update algo_reduce schedule (#296)

microsoft · Sep 1, 2021 · 15da6bb · 15da6bb
1 parent 9b7e2bb
commit 15da6bb
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/backends/c-rocm/schedule/standard/algo_reduce.py b/backends/c-rocm/schedule/standard/algo_reduce.py
@@ -11,8 +11,6 @@ def schedule_branch(attrs, output, prefix):
     sizes = cfg.define_split(f"{prefix}R{i}", attrs.get_extent(ax), num_outputs=2)
     if rax == i:
       r_range = max(2, sizes[1])
-      if not attrs.backend.startswith('c-cuda'):
-        r_range = r_range if r_range != 32 else 16
       ko, ki = s[output].split(ax, factor=r_range)
       BF = s.rfactor(output, ki)
 

diff --git a/engine/device-stub/tvm_v0.7.patch b/engine/device-stub/tvm_v0.7.patch
@@ -213,7 +213,7 @@ index 9cd29357f..e65c61149 100644
                     TVMArrayGet(DataType::UInt(8), handle, builtin::kArrTypeBits) ==
                         IntImm(DataType::UInt(8), dtype.bits()) &&
 diff --git a/src/tir/transforms/lower_thread_allreduce.cc b/src/tir/transforms/lower_thread_allreduce.cc
-index f6cb09672..fae978d03 100644
+index f6cb09672..9f39e841c 100644
 --- a/src/tir/transforms/lower_thread_allreduce.cc
 +++ b/src/tir/transforms/lower_thread_allreduce.cc
 @@ -34,13 +34,17 @@
@@ -235,6 +235,15 @@ index f6cb09672..fae978d03 100644
 
    Stmt VisitStmt_(const AttrStmtNode* op) final {
      if (op->attr_key == attr::thread_extent) {
+@@ -486,7 +490,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
+   // Also, the warp/wavefront size differs (64 on rocm, 32 on cuda).
+   bool is_warp_reduction(const std::vector<DataType>& types) const {
+     // Only cuda target supports warp reductions.
+-    if ((target_->kind->name != "cuda") && (target_->kind->name != "rocm")) return false;
++    if (strncmp("c-cuda", getenv("BACKEND"), 6) != 0) return false;
+
+     // rocm only supports 32 bit operands for shuffling at the moment
+     if ((target_->kind->name == "rocm") &&
 diff --git a/src/tir/transforms/split_host_device.cc b/src/tir/transforms/split_host_device.cc
 index 921c7ad79..d96f871de 100644
 --- a/src/tir/transforms/split_host_device.cc

diff --git a/engine/install_antares_host.sh b/engine/install_antares_host.sh
@@ -3,7 +3,7 @@
 cd $(dirname $0)/..
 ANTARES_ROOT=$(pwd)
 
-VERSION_TAG=v0.2dev10
+VERSION_TAG=v0.2dev11
 
 REQUIRED_CMDS="git python3 g++ make"