code: def main() -> None:
      dims = MatmulSize(
          M=_M.value,
          N=_N.value,
          K=_K.value,
          quantized_lhs=_QUANTIZED_LHS.value,
          quantized_rhs=_QUANTIZED_RHS.value,
      )
      s = torch.cuda.Stream()
      tilings = generate_tiling_configs(
          parse_int_list(_TILINGS_M.value),
          parse_int_list(_TILINGS_N.value),
          parse_int_list(_TILINGS_K.value),
          parse_int_list(_SPLIT_KS.value),
          parse_layout_list(_LHS_LAYOUTS.value),
          parse_layout_list(_RHS_LAYOUTS.value),
          parse_layout_list(_RESULT_LAYOUTS.value),
          parse_int_list(_NUM_STAGES.value),
          parse_int_list(_NUM_WARPS.value),
      )
      pbar = tqdm.tqdm(total=len(tilings), ncols=0)
      timings = sorted(
          benchmark_matmul(
              dims,
              pbar,
              s,
              tilings,
              repetitions_ms=300,
              debug=_DEBUG.value,
          ),
          key=lambda t: t.min_time_ms,
      )
      fastest: MatmulTiming = timings[0]
      print(f"Fastest configuration: {fastest}")
  
      features_list = [
          "BLOCK_M",
          "BLOCK_N",
          "BLOCK_K",
          "SPLIT_K",
          "lhs_layout",
          "rhs_layout",
          "result_layout",
          "num_stages",
          "num_warps",
      ]
      features = frozenset(features_list)
      for f in features:
          other_features = features - {f}
  
          def other_features_equal_to_best(t):
              return all(
                  getattr(fastest.tiling, of) == getattr(t.tiling, of)
                  for of in other_features  # pylint: disable=cell-var-from-loop
              )
  
          # Keep everyting but the currently evaluated feature fixed to the best
          # value.
          others_fixed = [
              t for t in timings if other_features_equal_to_best(t)
          ]
  
          # TODO(cheshire): Visualize.
          print(
              f"Varying feature {f}:",
              ", ".join(
                  f"{t.min_time_ms:0.4f} @ {f}={getattr(t.tiling, f)}"
                  for t in others_fixed
              ),
          )
  
      print_roofline_performance(dims, fastest.min_time_ms)
      cublas_time = benchmark_cublas(dims)
      print(
          f"Reference cuBLAS time (bf16xbf16->bf16): {cublas_time:0.4f}ms"
      )
  
      if _APPEND_TO_CSV.value:
          fields = (
              ["M", "N", "K", "quantized_lhs", "quantized_rhs"]
              + features_list
              + ["min_time_ms", "cublas_time_ms"]
          )
          with open(_APPEND_TO_CSV.value, "a") as f:
              writer = csv.DictWriter(f, fieldnames=fields)
              if f.tell() == 0:
                  writer.writeheader()
              writer.writerow(
                  dict(
                      fastest.dims._asdict(),
                      **fastest.tiling._asdict(),
                      min_time_ms=fastest.min_time_ms,
                      cublas_time_ms=cublas_time,
                  )
              )
  
file_path: tensorflow-upstream/third_party/xla/xla/experiments/triton_autotuning/tune_single_matmul.py
imports: import csv
  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 generate_tiling_configs
  from matmul_lib import MatmulSize
  from matmul_lib import MatmulTiming
  from matmul_lib import parse_int_list
  from matmul_lib import parse_layout_list
  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() and returns: all