code: def benchmark_matmul_tiling(
      dims: MatmulSize,
      tiling: MatmulTiling,
      s: torch.cuda.Stream,
      shared_stream: torch.cuda.Stream,
      a: torch.Tensor | triton.TensorWrapper,
      b: torch.Tensor | triton.TensorWrapper,
      c: torch.Tensor,
      scratchpad: torch.Tensor,  # Largest size: c * SPLIT_K
      repetitions_ms: int,
      debug=False,
  ) -> typing.Optional[MatmulTiming]:
      """Benchmarks a single matmul tiling."""
      grid = lambda META: (  # pylint: disable=g-long-lambda
          triton.cdiv(dims.M, tiling.BLOCK_M)
          * triton.cdiv(dims.N, tiling.BLOCK_N),
          tiling.SPLIT_K,
          1,  # batch
      )
      data_a = getattr(a, "base", a)
      data_b = getattr(b, "base", b)
  
      is_rm = lambda x: x == MatrixLayout.ROW_MAJOR
      m, n, k = int(dims.M), int(dims.N), int(dims.K)
  
      def run_matmul():
          used_output = c if tiling.SPLIT_K == 1 else scratchpad
          _matmul_kernel[grid](
              a,
              b,
              used_output,
              m=m,
              n=n,
              k=k,
              # This won't quite compute with the correct values, since we're not
              # transposing the data beforehand for column-major layout, but we're
              # never actually checking the correctness of the result here, so it's
              # fine for now.
              stride_am=k if is_rm(tiling.lhs_layout) else 1,
              stride_ak=1 if is_rm(tiling.lhs_layout) else m,
              stride_bk=n if is_rm(tiling.rhs_layout) else 1,
              stride_bn=1 if is_rm(tiling.rhs_layout) else k,
              stride_cm=n if is_rm(tiling.result_layout) else 1,
              stride_cn=1 if is_rm(tiling.result_layout) else m,
              block_m=int(tiling.BLOCK_M),
              block_n=int(tiling.BLOCK_N),
              block_k=int(tiling.BLOCK_K),
              group_m=8,
              split_k=tiling.SPLIT_K,
              num_warps=tiling.num_warps,
              num_stages=tiling.num_stages,
              force_num_warps=tiling.num_warps,
              force_num_stages=tiling.num_stages,
              acc_ty=tl.float32,
          )
          if tiling.SPLIT_K != 1:
              # Run reduction kernel.
              _reduce_kernel[(triton.cdiv(dims.M * dims.N, 1024),)](
                  scratchpad,
                  c,
                  row_size=int(dims.M),
                  col_size=tiling.SPLIT_K,
                  num_stages=1,
                  num_warps=1024 // 32,
                  row_block_size=1024,
              )
  
      for dim in ["M", "N", "K"]:
          next_pow2 = lambda v: 2 ** int(math.ceil(math.log2(v)))
          dim_size: int = getattr(dims, dim)
          if dim == "K":
              dim_size = math.ceil(dim_size / tiling.SPLIT_K)
          tile_size = getattr(tiling, f"BLOCK_{dim}")
          if next_pow2(dim_size) < tile_size:
              if debug:
                  LOG.error(
                      "Tile %s larger than the dimension %s (%s)",
                      tile_size,
                      dim,
                      dim_size,
                  )
              return None
  
      if tiling.BLOCK_M * tiling.BLOCK_N > 131072:
          if debug:
              LOG.error("Overly large tile")
          return None
  
      # TODO(cheshire): Compilation time is huge for such tiles.
      if tiling.BLOCK_M > 512 or tiling.BLOCK_N > 512:
          if debug:
              LOG.error("Overly large tile")
          return None
  
      max_shared_memory = (
          triton.runtime.driver.utils.get_device_properties(
              torch.cuda.current_device()
          )["max_shared_mem"]
      )
  
      required_shared_memory = (
          (
              tiling.BLOCK_M * data_a.element_size()
              + tiling.BLOCK_N * data_b.element_size()
          )
          * tiling.BLOCK_K
          * tiling.num_stages
      )
      if required_shared_memory > max_shared_memory:
          if debug:
              LOG.error(
                  "Skipping %s due to exceeding shmem bound", tiling
              )
          return None
      with torch.cuda.stream(s):
          try:
              run_matmul()  # Warmup on our own stream.
          except Exception as exc:
              LOG.error(
                  "%s for %s generated %s",
                  tiling,
                  dims,
                  exc,
                  exc_info=True,
              )
              raise
  
      # Use shared stream to take actual measurements.
      with torch.cuda.stream(shared_stream):
          try:
              percentiles = triton.testing.do_bench(
                  run_matmul,
                  warmup=0,
                  rep=repetitions_ms,
                  quantiles=(0.001, 0.1, 0.5, 0.9),
              )
              min_ms = percentiles[0]
          except Exception as exc:
              LOG.error(
                  "%s for %s generated %s",
                  tiling,
                  dims,
                  exc,
                  exc_info=True,
              )
              raise
          return MatmulTiming(dims, tiling, min_ms)
  
imports: import enum
  import itertools
  import logging
  import math
  import typing
  import torch
  import tqdm
  import triton
  import triton.language
label_desc: Use image node_id 9 for calling a global function with example usage: benchmark_matmul_tiling(dims, tiling, s, shared_stream, a, b, c, scratchpad, repetitions_ms, debug) and returns: None, None, None, MatmulTiming, None