class_docstr: Manage Embedding Weights on CPU and CUDA memory uses a software cache.
  CPU maintains the entire original weight.
  CUDA maintains a fraction of the weights used in the upcoming computation. The row number in CUDA is controlled by `cuda_row_num`.
  During training, GPU needs to transmit embedding rows between CPU and GPU.
  Args:
      weight (torch.Tensor): the weight of the Embedding layer.
      cuda_row_num (int, optional): the number of rows cached in CUDA memory. Defaults to 0.
      buffer_size (int, optional): the number of rows in a data transmitter buffer. Defaults to 50_000.
      pin_weight (bool, optional): use pin memory to store the cpu weight. If set `True`, the cpu memory usage will increase largely. Defaults to False.
      evict_strategy (EvictionStrategy, optional): the eviction strategy. There are two options.
      `EvictionStrategy.LFU`: use the least frequently used cache.
      `EvictionStrategy.DATASET`: use the stats collected from the target dataset. It usually leads to less cpu-gpu communication volume.
      Defaults to EvictionStrategy.DATASET.
class_docstr_tok: ["Manage","Embedding","Weights","on","CPU","and","CUDA","memory","uses","a","software","cache",".","CPU","maintains","the","entire","original","weight",".","CUDA","maintains","a","fraction","of","the","weights","used","in","the","upcoming","computation",".","The","row","number","in","CUDA","is","controlled","by","`","cuda_row_num","`",".","During","training",",","GPU","needs","to","transmit","embedding","rows","between","CPU","and","GPU",".","Args",":","weight","(","torch.Tensor",")",":","the","weight","of","the","Embedding","layer",".","cuda_row_num","(","int",",","optional",")",":","the","number","of","rows","cached","in","CUDA","memory",".","Defaults","to","0.","buffer_size","(","int",",","optional",")",":","the","number","of","rows","in","a","data","transmitter","buffer",".","Defaults","to","50_000",".","pin_weight","(","bool",",","optional",")",":","use","pin","memory","to","store","the","cpu","weight",".","If","set","`","True","`",",","the","cpu","memory","usage","will","increase","largely",".","Defaults","to","False",".","evict_strategy","(","EvictionStrategy",",","optional",")",":","the","eviction","strategy",".","There","are","two","options",".","`","EvictionStrategy.LFU","`",":","use","the","least","frequently","used","cache",".","`","EvictionStrategy.DATASET","`",":","use","the","stats","collected","from","the","target","dataset",".","It","usually","leads","to","less","cpu-gpu","communication","volume",".","Defaults","to","EvictionStrategy.DATASET","."]
docstr: prepare rows in cpu_row_idxs on CUDA memory
  Args:
      cpu_row_idxs (torch.Tensor): the rows to be placed on CUDA
docstr_tok: ["prepare","rows","in","cpu_row_idxs","on","CUDA","memory","Args",":","cpu_row_idxs","(","torch.Tensor",")",":","the","rows","to","be","placed","on","CUDA"]
code: def _prepare_rows_on_cuda(self, cpu_row_idxs: torch.Tensor) -> None:
      """prepare rows in cpu_row_idxs on CUDA memory
      Args:
          cpu_row_idxs (torch.Tensor): the rows to be placed on CUDA
      """
      evict_num = cpu_row_idxs.numel() - self.cuda_available_row_num
  
      cpu_row_idxs_copy = cpu_row_idxs.cpu()
  
      # move evict in rows to gpu
      if self._async_copy:
          if self.buffer_size == 0:
              evict_in_rows_gpu = (
                  self.weight.view(self.num_embeddings, -1)
                  .index_select(0, cpu_row_idxs_copy)
                  .pin_memory()
              )
              with torch.cuda.stream(self._memcpy_stream):
                  evict_in_rows_gpu = evict_in_rows_gpu.to(
                      torch.cuda.current_device(), non_blocking=True
                  )
          else:
              raise NotImplemented
  
      if evict_num > 0:
          with self.timer("2_identify_cuda_row_idxs") as timer:
              mask_cpu_row_idx = torch.isin(
                  self.cached_idx_map, self.evict_backlist
              )
              invalid_idxs = torch.nonzero(mask_cpu_row_idx).squeeze(1)
              if self._evict_strategy == EvictionStrategy.DATASET:
                  # mask method.
                  # set cached_idx_map[invalid_idxs] to -2.
                  # so those idxs will be sorted to end, therefore not being chosen as victim
                  backup_idxs = self.cached_idx_map[
                      mask_cpu_row_idx
                  ].clone()
                  self.cached_idx_map.index_fill_(0, invalid_idxs, -2)
  
                  with self.timer("2_1_find_evict_gpu_idxs") as timer:
                      evict_gpu_row_idxs = self._find_evict_gpu_idxs(
                          evict_num
                      )
  
                  # move evict out rows to cpu
                  if self._async_copy:
                      evict_out_rows_gpu = self.cuda_cached_weight.view(
                          self.cuda_row_num, -1
                      ).index_select(0, evict_gpu_row_idxs)
                      evict_out_rows_cpu = torch.empty_like(
                          evict_out_rows_gpu,
                          device="cpu",
                          pin_memory=True,
                      )
                      with torch.cuda.stream(None):
                          evict_out_rows_cpu.copy_(
                              evict_out_rows_gpu, non_blocking=True
                          )
                  self.cached_idx_map.index_copy_(
                      0, invalid_idxs, backup_idxs
                  )
  
              elif self._evict_strategy == EvictionStrategy.LFU:
                  with self.timer("2_1_backup_freqs") as timer:
                      backup_freqs = self.freq_cnter[
                          invalid_idxs
                      ].clone()
                      self.freq_cnter.index_fill_(
                          0, invalid_idxs, sys.maxsize
                      )
  
                  with self.timer("2_2_find_evict_gpu_idxs") as timer:
                      evict_gpu_row_idxs = self._find_evict_gpu_idxs(
                          evict_num
                      )
  
                  if self._async_copy:
                      evict_out_rows_gpu = self.cuda_cached_weight.view(
                          self.cuda_row_num, -1
                      ).index_select(0, evict_gpu_row_idxs)
                      evict_out_rows_cpu = torch.empty_like(
                          evict_out_rows_gpu,
                          device="cpu",
                          pin_memory=True,
                      )
                      with torch.cuda.stream(None):
                          evict_out_rows_cpu.copy_(
                              evict_out_rows_gpu, non_blocking=True
                          )
  
                  with self.timer("2_3_revert_freqs") as timer:
                      self.freq_cnter.index_copy_(
                          0, invalid_idxs, backup_freqs
                      )
  
              evict_info = self.cached_idx_map[evict_gpu_row_idxs]
  
          with self.timer("3_evict_out") as timer:
              if self.buffer_size > 0:
                  self.limit_buff_index_copyer.index_copy(
                      0,
                      src_index=evict_gpu_row_idxs,
                      tgt_index=evict_info.cpu(),
                      src=self.cuda_cached_weight.view(
                          self.cuda_row_num, -1
                      ),
                      tgt=self.weight.view(self.num_embeddings, -1),
                  )
              else:
                  # allocate tmp memory on CPU and copy rows on CUDA to CPU.
                  # TODO async gpu -> cpu
                  if self._async_copy:
                      _wait_for_data(evict_out_rows_cpu, None)
                  else:
                      with self.timer(
                          "3_1_evict_out_index_select"
                      ) as timer:
                          evict_out_rows_cpu = (
                              self.cuda_cached_weight.view(
                                  self.cuda_row_num, -1
                              ).index_select(0, evict_gpu_row_idxs)
                          )
                      with self.timer(
                          "3_2_evict_out_gpu_to_cpu_copy"
                      ) as timer:
                          evict_out_rows_cpu = evict_out_rows_cpu.cpu()
  
                  with self.timer("3_2_evict_out_cpu_copy") as timer:
                      self.weight.view(
                          self.num_embeddings, -1
                      ).index_copy_(
                          0, evict_info.cpu(), evict_out_rows_cpu
                      )
  
              self.cached_idx_map.index_fill_(0, evict_gpu_row_idxs, -1)
              self.inverted_cached_idx.index_fill_(0, evict_info, -1)
              # self.freq_cnter.index_fill(0, evict_gpu_row_idxs, sys.maxsize) # unnecessary
              self._cuda_available_row_num += evict_num
  
              weight_size = (
                  evict_gpu_row_idxs.numel() * self.embedding_dim
              )
              self._cuda_to_cpu_numel += weight_size
          # print(f"evict embedding weight: {weight_size*self.elem_size_in_byte/1e6:.2f} MB")
  
      # slots of cuda weight to evict in
      with self.timer("4_identify_cuda_slot") as timer:
          slots = torch.nonzero(self.cached_idx_map == -1).squeeze(1)[
              : cpu_row_idxs.numel()
          ]
  
      # TODO wait for optimize
      with self.timer("5_evict_in") as timer:
          # Here also allocate extra memory on CUDA. #cpu_row_idxs
          if self.buffer_size > 0:
              self.limit_buff_index_copyer.index_copy(
                  0,
                  src_index=cpu_row_idxs_copy,
                  tgt_index=slots,
                  src=self.weight.view(self.num_embeddings, -1),
                  tgt=self.cuda_cached_weight.view(
                      self.cuda_row_num, -1
                  ),
              )
          else:
              if self._async_copy:
                  _wait_for_data(evict_in_rows_gpu, self._memcpy_stream)
              else:
                  with self.timer("5_1_evict_in_index_select") as timer:
                      # narrow index select to a subset of self.weight
                      # tmp = torch.narrow(self.weight.view(self.num_embeddings, -1), 0, min(cpu_row_idxs).cpu(), max(cpu_row_idxs) - min(cpu_row_idxs) + 1)
                      # evict_in_rows_gpu = tmp.index_select(0, cpu_row_idxs_copy - min(cpu_row_idxs).cpu())
                      evict_in_rows_gpu = (
                          self.weight.view(self.num_embeddings, -1)
                          .index_select(0, cpu_row_idxs_copy)
                          .pin_memory()
                      )
  
                  with self.timer(
                      "5_2_evict_in_gpu_to_cpu_copy"
                  ) as timer:
                      evict_in_rows_gpu = evict_in_rows_gpu.cuda()
  
                  with self.timer("5_3_evict_in_index_copy") as timer:
                      self.cuda_cached_weight.view(
                          self.cuda_row_num, -1
                      ).index_copy_(0, slots, evict_in_rows_gpu)
  
      with self.timer("6_update_cache") as timer:
          self.cached_idx_map[slots] = cpu_row_idxs
          self.inverted_cached_idx.index_copy_(0, cpu_row_idxs, slots)
          if self._evict_strategy == EvictionStrategy.LFU:
              self.freq_cnter.index_fill_(0, slots, 0)
          self._cuda_available_row_num -= cpu_row_idxs.numel()
  
      weight_size = cpu_row_idxs.numel() * self.embedding_dim
      self._cpu_to_cuda_numel += weight_size
  
imports: import sys
  from contextlib import contextmanager
  from enum import Enum
  from typing import List, Optional
  import numpy
  import torch
  from contexttimer import Timer
  from torch.profiler import record_function
  from .copyer import LimitBuffIndexCopyer
label_desc: Use image node_id 14 for calling the CachedParamMgr obj's underlying member method code with example usage: obj._prepare_rows_on_cuda(cpu_row_idxs) without return types