TVM 常用数据结构

TVM 常用数据结构#

DLDeviceType#

枚举类型 DLDeviceType，用于表示不同的设备类型。

该枚举类型包含以下成员：

  • kDLCPU：表示 CPU 设备。

  • kDLCUDA：表示 CUDA GPU 设备。

  • kDLCUDAHost：表示通过 cudaMallocHost 固定分配的 CUDA CPU 内存。

  • kDLOpenCL：表示 OpenCL 设备。

  • kDLVulkan：表示下一代图形的 Vulkan 缓冲区。

  • kDLMetal：表示 Apple GPU 的 Metal。

  • kDLVPI：表示 Verilog 模拟器缓冲区。

  • kDLROCM：表示 AMD GPU 的 ROCm GPU。

  • kDLExtDev：保留的扩展设备类型，用于快速测试扩展设备。语义可能因实现而异。

通过使用这个枚举类型，我们可以在程序中指定要使用的设备类型，以便根据不同的设备执行相应的操作或加载适当的库。

DLDevice#

结构体 DLDevice，用于表示一个设备，该设备可以处理张量和算子。

该结构体包含以下成员：

  • device_type：设备中使用的设备类型，使用前面定义的枚举类型 DLDeviceType 表示。

  • device_id：设备的索引或标识符，使用整数类型 int 表示。

通过使用这个结构体，我们可以创建一个表示特定设备的对象，并指定该设备的类型和索引。这样的对象可以用于在张量和运算符的处理中指定要使用的设备，以便执行相应的算子或加载适当的库。

DLDataTypeCode#

这段代码定义了一个枚举类型 DLDataTypeCode，用于表示不同的数据类型代码选项。

该枚举类型包含以下成员：

  • kDLInt：有符号整数类型代码。

  • kDLUInt：无符号整数类型代码。

  • kDLFloat：IEEE浮点数类型代码。

  • kDLOpaqueHandle：不透明句柄类型代码，保留给测试目的使用。框架需要就交换的句柄数据类型达成一致，才能定义交换。

  • kDLBfloat：bfloat16 类型代码。

  • kDLComplex：复数类型代码（C/C++/Python布局：每个复数都有一个紧凑的结构体）。

通过使用这个枚举类型，我们可以为不同的数据类型指定相应的类型代码。这样的代码可以用于在张量和运算符的处理中指定要使用的特定数据类型，以便执行相应的算子或加载适当的库。

DLDataType#

结构体 DLDataType 用于表示张量可以保存的数据类型：

  • code：基础类型的类型代码，使用 uint8_t 类型表示。为了最小化内存占用，我们使用 uint8_t 而不是 DLDataTypeCode，但值应该是 DLDataTypeCode 枚举值之一。

  • bits：位数，通常选择 8、16 和 32。

  • lanes：类型中的通道数，用于向量类型。

使用这个结构体，可以表示不同的数据类型，并指定它们在内存中占用的位数和通道数。例如，对于浮点数类型，type_code 为 2，bits 为 32，lanes 为 1；对于向量化 4 个浮点数的类型，type_code 为 2，bits 为 32，lanes 为 4；对于 8 位整数类型，type_code 为 0，bits 为 8，lanes 为 1；对于 std::complex<float>，type_code 为 5，bits 为 64，lanes 为 1。

DLTensor#

结构体 DLTensor，用于表示一个普通的 C 张量对象，该对象不管理内存。

该结构体包含以下成员：

  • data：指向分配的数据的不透明数据指针。在 CUDA 中，这将是设备指针或 OpenCL 中的 cl_mem 句柄。此指针始终按 256 字节对齐，就像在 CUDA 中一样。给定的 DLTensor，存储数据内容所需的内存大小计算如下：

    • GetDataSize 是一个静态内联函数，用于计算给定 DLTensor 所需的内存大小：

      static inline size_t GetDataSize(const DLTensor* t) {
      size_t size = 1;
      for (tvm_index_t i = 0; i < t->ndim; ++i) {
          size *= t->shape[i];
      }
      size *= (t->dtype.bits * t->dtype.lanes + 7) / 8;
      return size;
  }

  • device：张量的设备类型，使用前面定义的枚举类型 DLDevice 表示。

  • ndim：张量的维数。

  • dtype：指针的数据类型，使用前面定义的枚举类型 DLDataType 表示。

  • shape：张量的形状，使用指向整数的指针表示。

  • strides：张量的步幅（以元素数量而不是字节为单位），可以为 NULL，表示张量紧凑且采用行优先方式。

  • byte_offset：指向数据的起始指针的字节偏移量。

通过使用这个结构体，我们可以创建表示特定张量的对象，并指定该张量的数据、设备类型、维数、数据类型、形状和步幅等信息。这样的对象可以用于在张量处理中执行相应的算子或加载适当的库。

DLManagedTensor#

结构体 DLManagedTensor，用于表示托管张量对象。该结构体管理 DLTensor 的内存，并提供了在另一个框架中借用张量的便利性。它并不转移张量，当借用框架不再需要张量时，应调用删除器通知主机资源不再需要。

该结构体包含以下成员：

  • dl_tensor：正在被内存管理的 DLTensor。

  • manager_ctx：DLManagedTensor 使用的原始主机框架的上下文。如果 DLManagedTensor 在框架中使用，则可以为 NULL。

  • deleter：用于销毁 manager_ctx 的析构函数签名。应调用此函数以销毁持有 DLManagedTensor 的 manager_ctx。如果调用者无法提供合理的析构函数，则可以将其设置为 NULL。析构函数将删除参数 self。

通过使用这个结构体，我们可以创建一个托管张量对象，并指定该张量的数据、原始主机框架的上下文以及用于销毁 manager_ctx 的析构函数。这样的对象可以在框架之间方便地借用和管理张量，以便进行相应的操作或加载适当的库。

TVMValue#

联合类型（union type）TVMValue 允许在同一内存位置存储不同的数据类型，但一次只能存储其中一种类型的值：

  • v_int64：64 位整数（int64_t）类型。

  • v_float64：64 位浮点数（double）类型。

  • v_handle：指向句柄（handle）的指针类型，用于引用各种对象或资源。

  • v_str：指向字符串的指针类型，用于表示文本数据。

  • v_type：DLDataType 类型，表示数据类型。

  • v_device：DLDevice类型，表示设备信息。

通过使用这个联合类型，可以在 API 和函数调用中传递不同类型的值，而无需显式地进行类型转换。根据需要，可以将 TVMValue 实例中的特定成员赋值为相应的值，并使用该实例来传递数据。

TVMArgs#

结构体 TVMArgs，用于表示 TVM 函数的参数。

该结构体包含以下成员：

  • values：一个长度为 TVM_CRT_MAX_ARGS 的 TVMValue 数组，用于存储函数的参数值。

  • tcodes：一个长度为 TVM_CRT_MAX_ARGS 的整数数组，用于存储函数参数的数据类型码。数据类型应与 TVMPackedCFunc 中的 type_codes 相同。

  • values_count：一个无符号 32 位整数，表示 values 数组中的元素数量。

通过使用这个结构体，我们可以创建表示 TVM 函数参数的对象，并指定该函数的参数值、数据类型码和参数数量。这样的对象可以用于在 TVM 编译过程中传递参数给函数，以便进行相应的算子或加载适当的库。