code: def gemm_test(
      hidden_dim,
      inner_dim,
      vocab_size,
      min_bsz,
      max_bsz,
      dir_name="/tmp/igemm_configs",
  ):
      check_args(hidden_dim, inner_dim, vocab_size, min_bsz, max_bsz)
      sm = cuda_module.get_sm_version()
      if sm < 75:
          raise RuntimeError(
              "int8 gemm is only supported on GPUs with SM >= 75."
          )
  
      if min_bsz > BORDER:
          min_bsz = (min_bsz // STRIDE) * STRIDE
          print("Adjust the min_bsz to {}.".format(min_bsz))
      if max_bsz > BORDER:
          max_bsz = ((max_bsz + STRIDE - 1) // STRIDE) * STRIDE
          print("Adjust the max_bsz to {}.".format(max_bsz))
  
      # All (m, n, k) which may be searched.
      mnk_set = set()
      for bsz in range(min_bsz, max_bsz + 1):
          m = (
              bsz
              if bsz < BORDER
              else ((bsz + STRIDE - 1) // STRIDE) * STRIDE
          )
          if hidden_dim is not None and inner_dim is not None:
              nk = base_nk(hidden_dim, inner_dim)
              for n, k in nk:
                  mnk_set.add((m, n, k))
          if hidden_dim is not None and vocab_size is not None:
              mnk_set.add((m, vocab_size, hidden_dim))
          elif hidden_dim is not None:
              pass
  
      # Existing (m, n, k).
      mkdir(dir_name)
      gpu_name = cuda_module.get_gpu_name()
      output_cfg_file = "{}/igemm_{}.cfg".format(dir_name, gpu_name)
      exist_mnk_set = set()
      output_cfg_str = []
      if os.path.exists(output_cfg_file):
          with open(output_cfg_file, "r") as fin:
              for line in fin:
                  m, n, k = [int(x) for x in line.split()[:3]]
                  data_order = line.split()[3]
                  exist_mnk_set.add((m, n, k))
                  output_cfg_str.append(
                      (m, n, k, data_order, line.rstrip())
                  )
  
      # (m, n, k) to be searched.
      mnk_set -= exist_mnk_set
      if len(mnk_set) <= 0:
          print("No gemm shapes need to be searched.")
          return
      search(mnk_set, output_cfg_file, output_cfg_str)
  
file_path: lightseq/lightseq/training/ops/pytorch/gemm_test.py
imports: import os
  import functools
  from dataclasses import dataclass
  from copy import deepcopy
  from tqdm import tqdm
  from lightseq.training.ops.pytorch.builder import KernelBuilder
label_desc: Use image node_id 8 for calling a global function with example usage: gemm_test(hidden_dim, inner_dim, vocab_size, min_bsz, max_bsz, dir_name) without return types