code: def bench(shape, layer_id, p, fusion_types=[""]):
      dtype = torch.float16
      M, K = shape[0]
      _, N = shape[1]
      torch.manual_seed(0)
      # allocate inputs
      a = torch.randn(shape[0], device="cuda", dtype=dtype)
      b = torch.randn(shape[1], device="cuda", dtype=dtype)
  
      def tflops(ms):
          return M * K * N / ms * 1e-9
  
      row = [layer_id]
      for fusion_type in fusion_types:
          if fusion_type == "":
              fn_mm = getattr(Func, "mm")
          else:
              fn_mm = getattr(Func, f"mm_{fusion_type}")
  
          if "add" in fusion_type:
              bias = torch.randn((M, N), dtype=dtype, device="cuda")
          else:
              bias = None
  
          args = (a, b, bias)
  
          def fn():
              return fn_mm(*args)
  
          torch._inductor.config.triton.mm = "aten"
          torch_mm_ms, _, _ = triton.testing.do_bench(fn)
          torch._inductor.config.triton.mm = "triton"
          # reset to force code gen new python code
          torch._dynamo.reset()
          torch._inductor.metrics.reset()
          triton_mm_ms, _, _ = triton.testing.do_bench(fn)
          assert (
              torch._inductor.metrics.generated_kernel_count == 1
          ), "codegen #kernel != 1"
          row.extend([tflops(torch_mm_ms), tflops(triton_mm_ms)])
  
      p.add_row(row)
  
file_path: pytorch/benchmarks/dynamo/microbenchmarks/bench_mm_fusion.py
imports: import torch
  import torch._dynamo
  import torch._inductor.config
  import triton
  from prettytable import PrettyTable
label_desc: Use image node_id 1 for calling a global function with example usage: bench(shape, layer_id, p, fusion_types) and returns: unknown, fn_mm