Auto schedular performance on AMDGPU: the first attempt

merrymercy · December 6, 2020, 11:10pm

Thanks for the detailed comparison!

Yeah. AutoTVM and Ansor use different GFLOPS estimation. Autotvm uses this function https://github.com/apache/tvm/blob/3ec739c650b2b49c15528f301242e98f00f0493e/python/tvm/autotvm/task/task.py#L483-L484, while Ansor uses this visitor https://github.com/apache/tvm/blob/3ec739c650b2b49c15528f301242e98f00f0493e/src/auto_scheduler/compute_dag.cc#L547-L548. Ansor’s version is more conservative and more accurate.

Besides the final performance, you should also enjoy a faster search speed with Ansor (49879s vs 71122s)

The performance of the last dense layer is very strange. Could you print the task.compute_dag for this task? Do you have any thoughts on why it is slow?

merrymercy · December 6, 2020, 11:52pm

I checked the manual schedule for ROCM. The schedule uses rfactor for cross-thread reduction. Initially, I also implemented this for CUDA. But then I found this PR: warp level reduction for cuda. So I removed rfactor and generated schedules following this PR, which is much simpler.

The simpler version does not use rfactor and set_store_predicate. It achieves similar or better performance on CUDA. I don’t know how it works for ROCm, although I do find a related PR for ROCm ROCm warp reduction.

If you want to fix this, you can play with the schedule generated by ansor (help) to find the reason. If we really need rfactor, we can implement a new sketch rule for rfactor on GPU, and add it back to the search space similar to this one.

masahi · December 7, 2020, 12:30am

Interesting, the rocm dense schedule you mentioned is probably an old one copy-pasted from CUDA, so I don’t think it is supposed to be particularly fast.

Below is the compute DAG, and the final schedule found by Ansor, showing both TIR and equivalent python scheduling commands:

========== Task 1  (workload key: ["eca51cb8a8335304c6e670bdb115a9b7"]) ==========                                                                                           
placeholder = PLACEHOLDER [1, 2048]                                                                                                                                          
placeholder = PLACEHOLDER [1000, 2048]                                                                                                                                       
T_dense(i, j) += (placeholder[i, k]*placeholder[j, k])                                                                                                                       
placeholder = PLACEHOLDER [1000]                                                                                                                                             
T_add(ax0, ax1) = (T_dense[ax0, ax1] + placeholder[ax1])   

Lowered TIR:                                                                                                                                                                 
primfn(placeholder_3: handle, placeholder_4: handle, placeholder_5: handle, T_add_1: handle) -> ()                                                                           
  attr = {"global_symbol": "main", "tir.noalias": True}                                                                                                                      
  buffers = {T_add: Buffer(T_add_2: Pointer(float32), float32, [1, 1000], []),                                                                                               
             placeholder_2: Buffer(placeholder_6: Pointer(float32), float32, [1000], []),                                                                                    
             placeholder_1: Buffer(placeholder_7: Pointer(float32), float32, [1000, 2048], []),                                                                              
             placeholder: Buffer(placeholder_8: Pointer(float32), float32, [1, 2048], [])}                                                                                   
  buffer_map = {placeholder_3: placeholder, placeholder_4: placeholder_1, placeholder_5: placeholder_2, T_add_1: T_add} {                                                    
  attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 1000;                                                                          
  attr [normal_reduce_temp0: handle] "storage_scope" = "local";
  allocate(normal_reduce_temp0, float32, [1]);
  attr [reduce_temp0: handle] "storage_scope" = "local";
  allocate(reduce_temp0, float32, [1]);
  attr [T_dense: Pointer(float32)] "storage_scope" = "shared";
  allocate(T_dense, float32, [1]) {
    attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 1 {
      normal_reduce_temp0[0] = 0f32
      for (k.outer: int32, 0, 2048) {
        normal_reduce_temp0[0] = ((float32*)normal_reduce_temp0[0] + ((float32*)placeholder_8[k.outer]*(float32*)placeholder_7[((blockIdx.x*2048) + k.outer)]))
      }
      attr [meta[tir.CommReducer][0]] "reduce_scope" = @tir.reinterpret(0u64, dtype=handle);
      @tir.tvm_thread_allreduce(1u32, (float32*)normal_reduce_temp0[0], True, reduce_temp0, 0, dtype=handle)
      T_dense[0] = (float32*)reduce_temp0[0]
    }
    attr [IterVar(threadIdx.x_1: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 1;
    T_add_2[blockIdx.x] = ((float32*)T_dense[0] + (float32*)placeholder_6[blockIdx.x])
  }
}

Equivalent python schedule:
T_dense_i, T_dense_j, T_dense_k = tuple(T_dense.op.axis) + tuple(T_dense.op.reduce_axis)
T_add_ax0, T_add_ax1 = tuple(T_add.op.axis) + tuple(T_add.op.reduce_axis)
T_add_ax1_o, T_add_ax1_i = s[T_add].split(T_add_ax1, factor=1)
s[T_add].bind(T_add_ax1_i, te.thread_axis("threadIdx.x"))
T_dense_k_o, T_dense_k_i = s[T_dense].split(T_dense_k, factor=1)
s[T_dense].bind(T_dense_k_i, te.thread_axis("threadIdx.x"))
s[T_dense].compute_at(s[T_add], T_add_ax1_o)
T_add_ax0_ax1_o_fused = s[T_add].fuse(T_add_ax0, T_add_ax1_o)
s[T_add].bind(T_add_ax0_ax1_o_fused, te.thread_axis("blockIdx.x"))
s[T_dense].pragma(T_dense_i, "auto_unroll_max_step", 1024)
s[T_dense].pragma(T_dense_i, "unroll_explicit", True)

I’ve also checked the generated asm, the warp shuffle operation (ds_bpermute_b32) for rocm introduced in https://github.com/apache/tvm/pull/5727 is used in softmax, but not in dense. Are you saying CUDA backend uses warp reduction for dense too, and Ansor is relying on that assumption?

cc @t-vi

merrymercy · December 7, 2020, 5:18am

The schedule does not make sense to me. It should not be the optimal one. I don’t know what happened to other candidates.

I made a unit test for this layer. It uses a high verbosity level so it will print all schedules and their GFLOPs. Could you run it and share the full output?

gist.github.com

https://gist.github.com/merrymercy/c867582a59d0994845ea22b7c35ae524

tune_dense.py

import numpy as np
import tvm
from tvm import te, auto_scheduler, topi
from tvm.topi.testing import conv2d_nchw_python

@auto_scheduler.register_workload
def dense_layer(in_dim, out_dim):
    data = te.placeholder((1, in_dim), name="data")
    weight = te.placeholder((out_dim, in_dim), name="weight")
    bias = te.placeholder((out_dim,), name="bias")

This file has been truncated. show original

masahi · December 7, 2020, 8:02am

ok the full log posted to https://gist.github.com/masahi/005cc165f1c3bd5179baee6bdb49e674

The final IR is a bit different from before but it is still slow:

Lowered TIR:
primfn(data_1: handle, weight_1: handle, bias_1: handle, T_add_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {T_add: Buffer(T_add_2: Pointer(float32), float32, [1, 1000], []),
             bias: Buffer(bias_2: Pointer(float32), float32, [1000], []),
             weight: Buffer(weight_2: Pointer(float32), float32, [1000, 2048], []),
             data: Buffer(data_2: Pointer(float32), float32, [1, 2048], [])}
  buffer_map = {data_1: data, weight_1: weight, bias_1: bias, T_add_1: T_add} {
  attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 250;
  attr [T_dense: Pointer(float32)] "storage_scope" = "local";
  allocate(T_dense, float32, [1]);
  attr [data.shared: Pointer(float32)] "storage_scope" = "shared";
  allocate(data.shared, float32, [8]);
  attr [weight.shared: Pointer(float32)] "storage_scope" = "shared";
  allocate(weight.shared, float32, [32]);
  attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 4 {
    T_dense[0] = 0f32
    for (k.outer.outer: int32, 0, 256) {
      attr [IterVar(threadIdx.x_1: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 4;
      data.shared[ramp((threadIdx.x_1*2), 1, 2)] = (float32x2*)data_2[ramp(((k.outer.outer*8) + (threadIdx.x_1*2)), 1, 2)]
      attr [IterVar(threadIdx.x_2: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 4;
      weight.shared[ramp((threadIdx.x_2*4), 1, 4)] = (float32x4*)weight_2[ramp(((((blockIdx.x*8192) + (floordiv(threadIdx.x_2, 2)*2048)) + (k.outer.outer*8)) + (floormod(threadIdx.x_2, 2)*4)), 1, 4)]
      attr [IterVar(threadIdx.x_2, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 4;
      weight.shared[ramp(((threadIdx.x_2*4) + 16), 1, 4)] = (float32x4*)weight_2[ramp((((((blockIdx.x*8192) + (floordiv(threadIdx.x_2, 2)*2048)) + (k.outer.outer*8)) + (floormod(threadIdx.x_2, 2)*4)) + 4096), 1, 4)]
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[0]*(float32*)weight.shared[(threadIdx.x*8)]))
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[1]*(float32*)weight.shared[((threadIdx.x*8) + 1)]))
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[2]*(float32*)weight.shared[((threadIdx.x*8) + 2)]))
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[3]*(float32*)weight.shared[((threadIdx.x*8) + 3)]))
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[4]*(float32*)weight.shared[((threadIdx.x*8) + 4)]))
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[5]*(float32*)weight.shared[((threadIdx.x*8) + 5)]))
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[6]*(float32*)weight.shared[((threadIdx.x*8) + 6)]))
      T_dense[0] = ((float32*)T_dense[0] + ((float32*)data.shared[7]*(float32*)weight.shared[((threadIdx.x*8) + 7)]))
    }
    T_add_2[((blockIdx.x*4) + threadIdx.x)] = ((float32*)T_dense[0] + (float32*)bias_2[((blockIdx.x*4) + threadIdx.x)])
  }
}


Latency: 0.263 ms, GFLOPS: 15.57

Note that the same script is working fine on CUDA. On my NV GPU I get

Latency: 0.056 ms, GFLOPS: 73.04

So probably this is AMDGPU or LLVM specific issue? Or there could be another code path that takes only kDLGPU into account?

merrymercy · December 7, 2020, 8:26am

I checked the source code and did not find other “kDLGPU”-only path. I can only guess it is an AMDGPU or LLVM specific issue. But I would like to try two additional tests.

How is the performance of the default ROCm manual schedule?
How is the performance if we change the out_dim in the unit test from 1000 to 1024?

masahi · December 7, 2020, 11:32am

Using the default schedule, in_dim = 2048, out_dim = 1000: Latency: 0.081 ms, GFLOPS: 50.65. Seems a reasonable number.
out_dim = 1024, no improvement. Latency: 0.287 ms, GFLOPS: 14.64

merrymercy · December 7, 2020, 3:39pm

The default schedule seems to work well. Auto-scheduler can fallback to default TOPI if it does not find any tuned records. To use the default TOPI schedule, we can either

Before tuning, delete that task after task extraction by
```
del tasks[1]
del task_weights[1]
```
After tuning, delete the records for this task in the log file. Get the hash key of this task by
```
print(task[1].workload_key)
```
use sed to delete all lines containing this hash key
```
sed -i '/hash_string/d' log_file.json
```
For example, if the hash key is 4275f42916bcd87979a7b95925e01ab8, use
```
sed -i '/4275f42916bcd87979a7b95925e01ab8/d' log_file.json
```

By using the default TOPI schedule, you should be able to reduce about 0.2ms in the final network latency.

merrymercy · December 7, 2020, 3:46pm

I would like to try a final test. If it does not work, I will stop here. I don’t have time for debugging this ROCm-specific issue.

Test the performance of the unit test with out_dim = 1024 and this patch

gist.github.com

https://gist.github.com/merrymercy/6b494376c169210938ddc4e06940d8ee

only_use_cross_thread.patch

diff --git a/src/auto_scheduler/search_policy/sketch_policy_rules.cc b/src/auto_scheduler/search_policy/sketch_policy_rules.cc
index 814e72a94..636aea91b 100644
--- a/src/auto_scheduler/search_policy/sketch_policy_rules.cc
+++ b/src/auto_scheduler/search_policy/sketch_policy_rules.cc
@@ -352,7 +352,7 @@ SketchGenerationRule::ConditionKind RuleCrossThreadReduction::MeetCondition(
         return ConditionKind::kSkip;
       }
 
-      return cum_space_len < cum_reduce_len ? ConditionKind::kApply : ConditionKind::kSkip;
+      return cum_space_len < cum_reduce_len ? ConditionKind::kApplyAndSkipRest : ConditionKind::kSkip;

This file has been truncated. show original

masahi · December 7, 2020, 10:12pm

It didn’t help either. For some reason the tuning finished early after only 35 round.

Thanks for all the help. I’ll keep this issue in mind as I go through Ansor codebase.

I just want to let you know that, I found Ansor codebase very well organized: I haven’t even read the paper, but just after brief look I have a good idea of how pieces fit together overall. Kudos to the team @merrymercy @comaniac @jcf94

t-vi · December 9, 2020, 7:38pm

Hi,

so I implemented warp shuffles when the CUDA schedule for softmax started using it and ROCm inherited this change, but did not adapt dense to the new capabilities. Thank you for spotting this! Will you submit a PR?

Best regards

Thomas

masahi · December 9, 2020, 9:46pm

Nope it’s not obvious to me how to do it. Also whether or not cuda backend is using warp shuffles for dense is not clear to me.

Mei · January 5, 2021, 5:51pm

Has anyone tried removing the call to MIOpen and only use LLVM GPU codegen? That way you can use NHWC.

masahi · January 5, 2021, 8:18pm

I didn’t use MIOpen here, all results obtained by NHWC layout + LLVM AMDGPU backend.