aot-executor-codegen

源码：tvm/src/relay/backend/aot_executor_codegen.cc

/*!
   * \brief ceil(size/word_size) to get number of words.
   * \param size The original size.
   * \param word_size The element size.
   */
  static size_t DivRoundUp(size_t size, size_t word_size) {
    return (size + word_size - 1) / word_size;
  }
  /*!
   * \brief Get the memory requirement.
   * \param prototype The prototype token.
   * \return The required memory size.
   *
   * TODO(mbs): Cf CalculateRelayExprSizeBytes in utils.cc, GetMemorySize is graph_plan_memory.cc
   */
  size_t GetMemorySizeBytes(const TensorType& ttype) {
    size_t size = 1;
    for (IndexExpr dim : ttype->shape) {
      const int64_t* pval = tir::as_const_int(dim);
      ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << ttype->shape;
      ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
      size *= static_cast<size_t>(pval[0]);
    }
    size *= DivRoundUp(ttype->dtype.bits() * ttype->dtype.lanes(), 8);
    return size;
  }

这段代码主要包含两个函数：DivRoundUp 和 GetMemorySizeBytes。

1. DivRoundUp 函数用于计算将一个数（size）除以另一个数（word_size）后向上取整的结果。这个函数接收两个参数，分别是原始大小（size）和元素大小（word_size）。它通过将 size 加上 word_size - 1，然后除以 word_size 来实现向上取整的效果。

2. GetMemorySizeBytes 函数用于计算给定张量类型（TensorType）所需的内存大小。它首先初始化一个名为 size 的变量为 1，然后遍历 ttype 的形状（shape），将每个维度的大小相乘。接着，它将 size 乘以 ttype 的数据类型（dtype）所占用的位数（bits）与通道数（lanes）之积，再除以 8（因为 1 字节 = 8 位），最后返回计算出的内存大小。在这个过程中，它还检查了维度是否为负数以及是否可以将形状转换为整数。