code: def main():
      s = torch.cuda.Stream()
      pbar = tqdm.tqdm(ncols=0)
      dims = MatmulSize(
          M=_M.value,
          N=_N.value,
          K=_K.value,
          quantized_lhs=_QUANTIZED_LHS.value,
          quantized_rhs=_QUANTIZED_RHS.value,
      )
      timing = benchmark_matmul(
          dims=dims,
          pbar=pbar,
          shared_stream=s,
          tilings=[
              MatmulTiling(
                  _BLOCK_M.value,
                  _BLOCK_N.value,
                  _BLOCK_K.value,
                  _SPLIT_K.value,
                  _LHS_LAYOUT.value,
                  _RHS_LAYOUT.value,
                  _RESULT_LAYOUT.value,
                  _NUM_STAGES.value,
                  _NUM_WARPS.value,
              )
          ],
          repetitions_ms=300,
          debug=_DEBUG.value,
      )
      if len(timing) != 1:
          print("Failed to find working configuration")
          sys.exit(1)
      t = timing[0]
      print(f"Timing: {t}")
      print_roofline_performance(dims, t.min_time_ms)
      cublas_time = benchmark_cublas(dims)
      print(
          f"Reference cuBLAS time (bf16xbf16->bf16): {cublas_time:0.4f}ms"
      )
  
file_path: tensorflow/third_party/xla/xla/experiments/triton_autotuning/run_single_matmul.py
imports: import sys
  from absl import app
  from absl import flags
  from matmul_lib import benchmark_cublas
  from matmul_lib import benchmark_matmul
  from matmul_lib import MatmulSize
  from matmul_lib import MatmulTiling
  from matmul_lib import MatrixLayout
  from matmul_lib import print_roofline_performance
  from matmul_lib import QuantizedInputType
  import torch
  import tqdm
label_desc: Use image node_id 1 for calling a global function with example usage: main() without return types