llvm::SmallDenseMap<unsigned, std::vector<Conv2DWithBlockSize>>
      argnum_and_convolutions;
  // A map maps block argument id to the number of users.
  llvm::SmallDenseMap<unsigned, int> argnum_num_users;

  // Find out the qualified convolutions and its block argument ids.
  auto conv2d_result = device_func.walk([&](TF::Conv2DOp conv2d) {
    std::optional<BlockArgumentInfo> arg_num_and_num_users =