Better TVM schedule for ROCm

adityaatluri · May 25, 2018, 8:07pm

Hi,
This gist has 2 IR for 4k x 4k gemm kernel generated for ROCm (Vega10 GPU)

asm_gemm_4k.ll, this kernel perform at peak, generated form HIP+inline asm
tvm_gemm_4k.ll, this kernel does not perform well and here are the issues I think are causing it

Issues:

Non vectorized loads. Compiler can find some vectorized loads but not all the time. TVM code gen for ROCm should express float4s rather than floats. (Using floats work good for CUDA as nvcc can find places to vectorize) In order to get peak throughput, explicit declaration as float4 can help
Irregular load patterns from lds. LDS loads are not properly coalesced in IR, using structs or multiples of float4s will get rid of bank conflicts in LDS.

With these fixed, we can judge how good is llvm in translating ir to asm.

gist.github.com

https://gist.github.com/adityaatluri/48002899ce9c7fc5fff76251a9cfcc60

asm_gemm_4k.ll


; Function Attrs: alwaysinline
define weak void @_Z5SGEMMPDv4_fS0_S0_PiS1_(<4 x float>* %A, <4 x float>* %B, <4 x float>* %C, i32* %getGlobalAId, i32* %getGlobalCId) #1 {
  %call = tail call i32 @hc_get_workitem_id(i32 0) #7
  %call1 = tail call i32 @hc_get_workitem_id(i32 1) #7
  %call2 = tail call i32 @hc_get_group_id(i32 0) #7
  %call3 = tail call i32 @hc_get_group_id(i32 1) #7
  %rem = srem i32 %call1, 2
  %mul = shl nsw i32 %rem, 4
  %add = add i32 %mul, %call

This file has been truncated. show original

tvm_gemm_4k.ll

source_filename = "default_function__kernel0"
target datalayout = "e-p:32:32-p1:64:64-p2:64:64-p3:32:32-p4:64:64-p5:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64"
target triple = "amdgcn-amd-amdhsa-hcc"

@.shared = private unnamed_addr addrspace(3) global [1024 x float], align 16
@.shared.1 = private unnamed_addr addrspace(3) global [1024 x float], align 16

; Function Attrs: nounwind
define dllexport amdgpu_kernel void @default_function__kernel0(float addrspace(1)* noalias nocapture readonly, float addrspace(1)* noalias nocapture readonly, float addrspace(1)* noalias) local_unnamed_addr #0 {
entry:

This file has been truncated. show original

FrozenGene · May 26, 2018, 7:30am

We also find this schedule is not good at Intel GPU OpenCL. I will keep an eye of this progress.