AV1代码学习：av1_foreach_transformed_block_in_plane函数

在AV1中，进行预测变换都是基于Transform Block（变换块）进行的，变换块一共19种尺寸，并且其尺寸通常是小于或者等于编码块尺寸的，如下代码所示。

enum {
  TX_4X4,             // 4x4 transform
  TX_8X8,             // 8x8 transform
  TX_16X16,           // 16x16 transform
  TX_32X32,           // 32x32 transform
  TX_64X64,           // 64x64 transform
  TX_4X8,             // 4x8 transform
  TX_8X4,             // 8x4 transform
  TX_8X16,            // 8x16 transform
  TX_16X8,            // 16x8 transform
  TX_16X32,           // 16x32 transform
  TX_32X16,           // 32x16 transform
  TX_32X64,           // 32x64 transform
  TX_64X32,           // 64x32 transform
  TX_4X16,            // 4x16 transform
  TX_16X4,            // 16x4 transform
  TX_8X32,            // 8x32 transform
  TX_32X8,            // 32x8 transform
  TX_16X64,           // 16x64 transform
  TX_64X16,           // 64x16 transform
  TX_SIZES_ALL,       // Includes rectangular transforms
  TX_SIZES = TX_4X8,  // Does NOT include rectangular transforms
  TX_SIZES_LARGEST = TX_64X64,
  TX_INVALID = 255  // Invalid transform size
} UENUM1BYTE(TX_SIZE);

如果编码块小于或等于64x64，则变换块划分只能应用于亮度分量，对于色度块，变换块的大小与编码块的大小相同，不使用变换块划分。否则，如果编码块的宽度或高度大于64，则亮度和色度编码块都将隐式拆分为min（W，64）x min（H，64）和min（W，32）x min（H，32）的变换块。

av1_foreach_transformed_block_in_plane函数就是循环遍历编码块中全部的变换块，对每一个变换块进行预测、变换和量化等操作的。

首先，通过av1_get_tx_size函数获取变换块尺寸，并获得每个变换块的4x4单元数；这里，亮度的变换块尺寸是由上层函数决定传递下来的，此处可以直接通过mbmi获取到

static INLINE TX_SIZE av1_get_tx_size(int plane, const MACROBLOCKD *xd) {
  const MB_MODE_INFO *mbmi = xd->mi[0];
  if (xd->lossless[mbmi->segment_id]) return TX_4X4; 
  if (plane == 0) return mbmi->tx_size;//返回亮度的变换块尺寸
  const MACROBLOCKD_PLANE *pd = &xd->plane[plane];
  return av1_get_max_uv_txsize(mbmi->bsize, pd->subsampling_x,
                               pd->subsampling_y); //返回色度最大变换块尺寸
}

获取当前编码块尺寸，此时会判断当前块的右侧/下侧是否在边界外，如果是在边界外，则会减去边界外的区域；后面也不会访问完全位于边界外的子块。

static INLINE int max_block_wide(const MACROBLOCKD *xd, BLOCK_SIZE bsize,
                                 int plane) {
  assert(bsize < BLOCK_SIZES_ALL);
  int max_blocks_wide = block_size_wide[bsize];

  if (xd->mb_to_right_edge < 0) { //当前编码块右侧到帧边界的距离，为负表示超出边界
    const struct macroblockd_plane *const pd = &xd->plane[plane];
    max_blocks_wide += xd->mb_to_right_edge >> (3 + pd->subsampling_x);
    // 右移3原因是mb_to_right_edge是以1/8像素精度存储的
  }

  // Scale the width in the transform block unit.
  return max_blocks_wide >> MI_SIZE_LOG2;
}
static INLINE int max_block_high(const MACROBLOCKD *xd, BLOCK_SIZE bsize,
                                 int plane) {
  int max_blocks_high = block_size_high[bsize];

  if (xd->mb_to_bottom_edge < 0) {
    const struct macroblockd_plane *const pd = &xd->plane[plane];
    max_blocks_high += xd->mb_to_bottom_edge >> (3 + pd->subsampling_y);
  }

  // Scale the height in the transform block unit.
  return max_blocks_high >> MI_SIZE_LOG2;
}

如果块尺寸是大于64x64的，会进行隐式划分，对于亮度块会以min（W，64）x min（H，64）进行预测变换操作，对于色度块会以min（W，32）x min（H，32）进行预测变换量化操作。

最后，遍历编码块中的变换块，进行预测变换等操作。

代码及其注释如下：

void av1_foreach_transformed_block_in_plane(
    const MACROBLOCKD *const xd, BLOCK_SIZE plane_bsize, int plane,
    foreach_transformed_block_visitor visit, void *arg) {
  const struct macroblockd_plane *const pd = &xd->plane[plane];
  // block and transform sizes, in number of 4x4 blocks log 2 ("*_b")
  // 4x4=0, 8x8=2, 16x16=4, 32x32=6, 64x64=8
  // transform size varies per plane, look it up in a common way.
  const TX_SIZE tx_size = av1_get_tx_size(plane, xd); //变换块尺寸
  const uint8_t txw_unit = tx_size_wide_unit[tx_size]; // 变换块宽度的单元数tx_size_wide/4
  const uint8_t txh_unit = tx_size_high_unit[tx_size]; // 变换块高度的单元数tx_size_high/4
  const int step = txw_unit * txh_unit; // 按照变换块尺寸遍历BLOCK

  // If mb_to_right_edge is < 0 we are in a situation in which
  // the current block size extends into the UMV and we won't
  // visit the sub blocks that are wholly within the UMV.
  //如果mb_to_right_edge<0，则当前块大小扩展到UMV中，并且我们不会访问完全在UMV中的子块
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane); //宽度的最大单元数 width/4
  const int max_blocks_high = max_block_high(xd, plane_bsize, plane); // 高度的最大单元数 height/4
  // 这里，使用BLOCK_64X64，是为了由于进行变换量化时，其最大尺寸为64x64
  // 如果当前Block尺寸大于64时，在这里会进行隐式划分
  const BLOCK_SIZE max_unit_bsize =
      get_plane_block_size(BLOCK_64X64, pd->subsampling_x, pd->subsampling_y); // 64
  const int mu_blocks_wide =
      AOMMIN(mi_size_wide[max_unit_bsize], max_blocks_wide);
  const int mu_blocks_high =
      AOMMIN(mi_size_high[max_unit_bsize], max_blocks_high);

  // Keep track of the row and column of the blocks we use so that we know
  // if we are in the unrestricted motion border.
  // 跟踪我们使用的块的行和列，以便知道我们是否在不受限制的运动边界中。
  int i = 0;
  for (int r = 0; r < max_blocks_high; r += mu_blocks_high) {
    const int unit_height = AOMMIN(mu_blocks_high + r, max_blocks_high);
    // Skip visiting the sub blocks that are wholly within the UMV.
    // 跳过访问整个UMV中的子块。
    for (int c = 0; c < max_blocks_wide; c += mu_blocks_wide) {
      const int unit_width = AOMMIN(mu_blocks_wide + c, max_blocks_wide);
      // 前两个循环是相当于隐式划分
      for (int blk_row = r; blk_row < unit_height; blk_row += txh_unit) {
        for (int blk_col = c; blk_col < unit_width; blk_col += txw_unit) {
          // 变换全部的变换块
          // 以变换块尺寸为单元进行预测变换量化
          visit(plane, i, blk_row, blk_col, plane_bsize, tx_size, arg);
          i += step;
        }
      }
    }
  }
}