code: def profile_bandwidth(path):
      s, h = 512, 512
      path_dir = os.path.dirname(path)
      os.makedirs(path_dir, exist_ok=True)
  
      links = [
          ("cpu", "gpu"),
          ("gpu", "cpu"),
          ("gpu", "gpu"),
          ("cpu", "cpu"),
          ("cpu", "disk"),
          ("disk", "cpu"),
      ]
  
      for dst, src in links:
          for b in [1, 128, 512]:
              if dst == "cpu":
                  dst_tensor = torch.ones(
                      (b, s, h), dtype=torch.int8, pin_memory=True
                  )
              elif dst == "gpu":
                  dst_tensor = torch.ones(
                      (b, s, h), dtype=torch.int8, device="cuda:0"
                  )
              elif dst == "disk":
                  np.lib.format.open_memmap(
                      path, mode="w+", shape=((b, s, h)), dtype=np.int8
                  )
                  dst_tensor = path
  
              if src == "cpu":
                  src_tensor = torch.ones(
                      (b, s, h), dtype=torch.int8, pin_memory=True
                  )
              elif src == "gpu":
                  src_tensor = torch.ones(
                      (b, s, h), dtype=torch.int8, device="cuda:0"
                  )
              elif src == "disk":
                  np.lib.format.open_memmap(
                      path, mode="w+", shape=((b, s, h)), dtype=np.int8
                  )
                  src_tensor = path
  
              dst_indices = (slice(0, b), slice(0, s), slice(0, h))
              src_indices = (slice(0, b), slice(0, s), slice(0, h))
  
              def func():
                  if isinstance(src_tensor, str):
                      src_tensor_ = torch.from_numpy(
                          np.lib.format.open_memmap(src_tensor)
                      )
                  else:
                      src_tensor_ = src_tensor
                  if isinstance(dst_tensor, str):
                      dst_tensor_ = torch.from_numpy(
                          np.lib.format.open_memmap(dst_tensor)
                      )
                  else:
                      dst_tensor_ = dst_tensor
                  dst_tensor_[dst_indices].copy_(
                      src_tensor_[src_indices]
                  )
  
              size = np.prod(
                  [
                      (x.stop - x.start) / (x.step or 1)
                      for x in dst_indices
                  ]
              )
              cost = np.mean(benchmark_func(func, number=5, repeat=3))
              bandwidth = size / cost / GB
  
              print(
                  f"size: {size / MB:6.2f} MB, {src}-to-{dst} bandwidth: {bandwidth:.3f} GB/s"
              )
          print()
  
file_path: H2O/h2o_flexgen/flexgen/profile_bandwidth.py
imports: import argparse
  import numpy
  import os
  import time
  import torch
  from flexgen.utils import GB, MB, KB
label_desc: Use image node_id 2 for calling a global function with example usage: profile_bandwidth(path) without return types