初探TVM--使用tensor engine在ADM GPU上编译生成优化算子

在AMD GPU上使用TE生成优化算子

这里我们会一起尝试在AMD GPU上面使用tvm生成优化算子，本节以MI100系列为例，架构gfx908.

AMD GPU环境配置

我使用的是AMD Radeon Instinct显卡，MI100，架构代号gfx908。
在rocm上面cmake要用：

cmake -DCMAKE_CXX_FLAGS="-Wno-deprecated-declarations" .. 

我是用的别人的docker来跑的：

docker pull mevermeulen/rocm-tvm:4.5

这个docker装了rocm 4.5，并且把llvm这些都预装好了，需要重新在新的机器上面编译，在config.cmake中打开下面几个：

mkdir build
sed -e 's/USE_ROCM OFF/USE_ROCM ON/g' -e 's?USE_LLVM OFF?USE_LLVM /opt/rocm/llvm/bin/llvm-config?g' -e 's/USE_MIOPEN OFF/USE_MIOPEN ON/g' -e 's/USE_ROCBLAS OFF/USE_ROCBLAS ON/g' ../cmake/config.cmake > config.cmake
cd build
cmake -DCMAKE_CXX_FLAGS="-Wno-deprecated-declarations" .. 

同时，由于好像llvm的某个函数的问题，还要在cc文件里面改改东西：

sed -e 's/getBaseName/getName/g' src/target/llvm/llvm_module.cc>src/target/llvm/llvm_module.cc.log
mv src/target/llvm/llvm_module.cc.log src/target/llvm/llvm_module.cc

sed -e 's/getBaseName/getName/g' src/target/llvm/codegen_llvm.cc>src/target/llvm/codegen_llvm.cc.log
mv src/target/llvm/codegen_llvm.cc.log src/target/llvm/codegen_llvm.cc

生成llvm端的code

这里的llvm不是指cpu，而是rocm使用了clang作为gpu的编译器，clang被包在了llvm里面。

然后使用上一章节的code，改动一行：

tgt_gpu = tvm.target.Target(target="rocm", host="llvm")

也可以直接用一个脚本跑：rocm_tvm_vector_add。哎，不过好像rocm的生成出来的module如果用get_source()方法的话，打印出来的是llvm的一堆IR，基本没法看懂。。不知道能不能干脆输出hip或者gcn的asm。
code我也直接贴在这里：

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import tvm
import tvm.testing
from tvm import te
import numpy as np

tgt_gpu = tvm.target.Target(target="rocm", host="llvm")

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
print(type(C))

s = te.create_schedule(C.op)

bx, tx = s[C].split(C.op.axis[0], factor=64)

s[C].bind(bx, te.thread_axis("blockIdx.x"))
s[C].bind(tx, te.thread_axis("threadIdx.x"))

fadd = tvm.build(s, [A, B, C], target=tgt_gpu, name="myadd")

dev = tvm.device(tgt_gpu.kind.name, 0)

n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)
fadd(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

if (
    tgt_gpu.kind.name == "cuda"
    or tgt_gpu.kind.name == "rocm"
    or tgt_gpu.kind.name.startswith("opencl")
):
    dev_module = fadd.imported_modules[0]
    print("-----GPU code-----")
    print(dev_module.get_source())
else:
    print(fadd.get_source())