[翻译]张量表达式入门 — tvm 0.7.dev1文档

张量表达式入门

作者：Tianqi Chen这是TVM中Tensor表达式语言的入门教程。TVM使用特定领域的张量表达式来进行有效的内核构建。在本教程中，我们将演示使用张量表达式语言的基本工作流程。fromfutureimportabsolute_import,print_function
importtvmfromtvmimportteimportnumpyasnp
# Global declarations of environment.
tgt_host=”llvm”# Change it to respective GPU if gpu is enabled Ex: cuda, opencl, rocmtgt=”cuda”## 向量添加示例 在本教程中，我们将使用向量加法示例演示工作流程。## 描述计算 第一步，我们需要描述我们的计算。TVM采用张量语义，每个中间结果表示为多维数组。用户需要描述生成张量的计算规则。我们首先定义一个符号变量n来表示形状。然后我们定义给定形状（n，）的两个占位符张量A和B然后，我们使用compute运算来描述结果张量C。该compute函数采用张量的形状，以及描述张量每个位置的计算规则的lambda函数。在此阶段没有任何计算发生，因为我们只是在声明应该如何进行计算。n=te.var(“n”)A=te.placeholder((n,),name=’A’)B=te.placeholder((n,),name=’B’)C=te.compute(A.shape,lambdai:A[i]+B[i],name=”C”)print(type(C))输出：## 安排计算 尽管以上几行描述了计算规则，但由于可以以数据并行方式计算C轴，因此我们可以用多种方式计算C。TVM要求用户提供对计算的描述，称为schedule。schedule是一组计算转换，可转换程序中的计算循环。构造完schedule后，默认情况下，schedule将按行优先顺序以串行方式计算C。for(inti=0;i<n;++i){C[i]=A[i]+B[i];}s=te.create_schedule(C.op)我们使用split结构拆分了C的第一个轴，这会将原始迭代轴拆分为两个迭代的乘积。这等效于以下代码。for(intbx=0;bx<ceil(n/64);++bx){for(inttx=0;tx<64;++tx){inti=bx*64+tx;if(i<n){C[i]=A[i]+B[i];}}}bx,tx=s[C].split(C.op.axis[0],factor=64)最后，我们将迭代轴bx和tx绑定到GPU计算网格中的线程。这些是GPU特定的结构，可让我们生成在GPU上运行的代码。iftgt==”cuda”ortgt==”rocm”ortgt.startswith(‘opencl’):s[C].bind(bx,te.thread_axis(“blockIdx.x”))s[C].bind(tx,te.thread_axis(“threadIdx.x”))## 汇编 完成时间表的指定后，我们可以将其编译为TVM函数。默认情况下，TVM会编译为可从python端直接调用的类型擦除函数。在下面的代码行中，我们使用tvm.build创建一个函数。构建函数获取时间表，函数所需的签名（包括输入和输出）以及我们要编译为的目标语言。编译fadd的结果是GPU设备功能（如果涉及GPU）以及调用该GPU功能的主机包装。fadd是生成的主机包装器函数，它在内部包含对生成的设备函数的引用。fadd=tvm.build(s,[A,B,C],tgt,target_host=tgt_host,name=”myadd”)## 运行功能 编译的TVM函数公开了可以从任何语言调用的简洁C API。我们在python中提供了最小的数组API，以帮助快速测试和原型设计。阵列API基于DLPack标准。

  • 我们首先创建一个GPU上下文。
  • 然后tvm.nd.array将数据复制到GPU。
  • fadd运行实际的计算。
  • asnumpy（）将GPU阵列复制回CPU，我们可以使用它来验证正确性 ctx=tvm.context(tgt,0)
    n=1024a=tvm.nd.array(np.random.uniform(size=n).astype(A.dtype),ctx)b=tvm.nd.array(np.random.uniform(size=n).astype(B.dtype),ctx)c=tvm.nd.array(np.zeros(n,dtype=C.dtype),ctx)fadd(a,b,c)tvm.testing.assertallclose(c.asnumpy(),a.asnumpy()+b.asnumpy())## 检查生成的代码 您可以在TVM中检查生成的代码。tvm.build的结果是一个TVM模块。fadd是包含主机包装程序的主机模块，它还包含用于CUDA（GPU）功能的设备模块。以下代码获取设备模块并打印内容代码。iftgt==”cuda”ortgt==”rocm”ortgt.startswith(‘opencl’):devmodule=fadd.importedmodules[0]print(“——-GPU code——-“)print(devmodule.getsource())else:print(fadd.getsource())输出：——-GPU code——-extern “C” global void myadd_kernel0(void* __restrict C, void* __restrict A, void* __restrict B, int n, int stride, int stride1, int stride2) {if (((int)blockIdx.x) < (n >> 6)) {((float)C)[((((((int)blockIdx.x) 64) + ((int)threadIdx.x)) stride2))] = (((float)A)[((((((int)blockIdx.x) 64) + ((int)threadIdx.x)) stride))] + ((float)B)[((((((int)blockIdx.x) 64) + ((int)threadIdx.x)) stride1))]);} else {if (((((int)blockIdx.x) 64) + ((int)threadIdx.x)) < n) {((float)C)[((((((int)blockIdx.x) 64) + ((int)threadIdx.x)) stride2))] = (((float)A)[((((((int)blockIdx.x) 64) + ((int)threadIdx.x)) stride))] + ((float)B)[((((((int)blockIdx.x) 64) + ((int)threadIdx.x)) * stride1))]);}}}注意代码专业化您可能已经注意到，A，B和C的声明都采用相同的shape参数n。TVM将利用此优势，仅将单个shape参数传递给内核，如您在印刷设备代码中所发现的那样。这是专业化的一种形式。在主机端，TVM将自动生成检查代码，以检查参数中的约束。因此，如果将具有不同形状的数组传递给fadd，则会引发错误。我们可以做更多的专业化。例如，我们可以在计算声明中编写而不是。生成的函数将仅采用长度为1024的向量。n=tvm.runtime.convert(1024)n=te.var(“n”)## 保存编译的模块 除了运行时编译外，我们还可以将编译后的模块保存到文件中，并在以后加载它们。这称为提前编译。以下代码首先执行以下步骤：
  • 它将已编译的主机模块保存到目标文件中。
  • 然后将设备模块保存到ptx文件中。
  • cc.create_shared调用编译器（gcc）创建共享库 fromtvm.contribimportccfromtvm.contribimportutil
    temp=util.tempdir()fadd.save(temp.relpath(“myadd.o”))iftgt==”cuda”:fadd.imported_modules[0].save(temp.relpath(“myadd.ptx”))iftgt==”rocm”:fadd.imported_modules[0].save(temp.relpath(“myadd.hsaco”))iftgt.startswith(‘opencl’):fadd.imported_modules[0].save(temp.relpath(“myadd.cl”))cc.create_shared(temp.relpath(“myadd.so”),[temp.relpath(“myadd.o”)])print(temp.listdir())输出：[‘myadd.tvm_meta.json’, ‘myadd.so’, ‘myadd.ptx’, ‘myadd.o’]注意模块存储格式CPU（主机）模块直接保存为共享库（.so）。设备代码可以有多种自定义格式。在我们的示例中，设备代码以及元数据json文件存储在ptx中。可以通过导入分别加载和链接它们。## 加载编译模块 我们可以从文件系统中加载已编译的模块并运行代码。以下代码分别加载主机和设备模块，并将它们重新链接在一起。我们可以验证新加载的函数是否起作用。fadd1=tvm.runtime.load_module(temp.relpath(“myadd.so”))iftgt==”cuda”:fadd1_dev=tvm.runtime.load_module(temp.relpath(“myadd.ptx”))fadd1.import_module(fadd1_dev)
    iftgt==”rocm”:fadd1_dev=tvm.runtime.load_module(temp.relpath(“myadd.hsaco”))fadd1.import_module(fadd1_dev)
    iftgt.startswith(‘opencl’):fadd1_dev=tvm.runtime.load_module(temp.relpath(“myadd.cl”))fadd1.import_module(fadd1_dev)
    fadd1(a,b,c)tvm.testing.assert_allclose(c.asnumpy(),a.asnumpy()+b.asnumpy())## 将所有内容打包到一个库中 在上面的示例中，我们分别存储了设备代码和主机代码。TVM还支持将所有内容导出为一个共享库。在后台，我们将设备模块打包到二进制Blob中，并将它们与主机代码链接在一起。目前，我们支持打包Metal，OpenCL和CUDA模块。fadd.export_library(temp.relpath(“myadd_pack.so”))fadd2=tvm.runtime.load_module(temp.relpath(“myadd_pack.so”))fadd2(a,b,c)tvm.testing.assert_allclose(c.asnumpy(),a.asnumpy()+b.asnumpy())注意运行时API和线程安全TVM的已编译模块不依赖于TVM编译器。相反，它们仅取决于最小的运行时库。TVM运行时库包装了设备驱动程序，并在编译后的函数中提供了线程安全和与设备无关的调用。这意味着您可以从任何线程，任何GPU上调用已编译的TVM函数。## 生成OpenCL代码 TVM将代码生成功能提供给多个后端，我们还可以生成在CPU后端运行的OpenCL代码或LLVM代码。以下代码块生成OpenCL代码，在OpenCL设备上创建数组，并验证代码的正确性。iftgt.startswith(‘opencl’):fadd_cl=tvm.build(s,[A,B,C],tgt,name=”myadd”)print(“———opencl code———“)print(fadd_cl.imported_modules[0].get_source())ctx=tvm.cl(0)n=1024a=tvm.nd.array(np.random.uniform(size=n).astype(A.dtype),ctx)b=tvm.nd.array(np.random.uniform(size=n).astype(B.dtype),ctx)c=tvm.nd.array(np.zeros(n,dtype=C.dtype),ctx)fadd_cl(a,b,c)tvm.testing.assert_allclose(c.asnumpy(),a.asnumpy()+b.asnumpy())## 摘要 本教程使用矢量添加示例逐步介绍了TVM工作流程。常规工作流程是
  • 通过一系列操作来描述您的计算。
  • 描述我们如何计算使用时间表原语。
  • 编译为我们想要的目标函数。
  • （可选）保存要在以后加载的功能。 非常欢迎您结帐其他示例和教程，以了解有关TVM中支持的操作，调度原语和其他功能的更多信息。DownloadPythonsourcecode:tensor_expr_get_started.pyDownloadJupyternotebook:tensor_expr_get_started.ipynb