code: def _matmul_kernel(
      lhs,
      rhs,
      out,
      m: tl.constexpr,
      n: tl.constexpr,
      k: tl.constexpr,
      stride_am: tl.constexpr,
      stride_ak: tl.constexpr,
      stride_bk: tl.constexpr,
      stride_bn: tl.constexpr,
      stride_cm: tl.constexpr,
      stride_cn: tl.constexpr,
      block_m: tl.constexpr,
      block_n: tl.constexpr,
      block_k: tl.constexpr,
      group_m: tl.constexpr,
      split_k: tl.constexpr,
      acc_ty: tl.constexpr,
      # Workaround for a bug in Triton cache:
      # force recompilation on different num_warps/num_stages.
      force_num_warps: tl.constexpr,  # pylint: disable=unused-argument
      force_num_stages: tl.constexpr,  # pylint: disable=unused-argument
  ):
      """Computes a block-level matmul."""
      even_k = k % (block_k * split_k) == 0
      pid0 = tl.program_id(0)
      pid1 = tl.program_id(1)
      pid2 = tl.program_id(2)
      grid_m = (m + block_m - 1) // block_m
      grid_n = (n + block_n - 1) // block_n
      # re-order program ID for better L2 performance
      width = group_m * grid_n
      group_id = pid0 // width
      group_size = min(grid_m - group_id * group_m, group_m)
      pid_m = group_id * group_m + pid0 % group_size
      pid_n = (pid0 % width) // group_size
      rm = pid_m * block_m + tl.arange(0, block_m)
      rn = pid_n * block_n + tl.arange(0, block_n)
      ram = tl.max_contiguous(tl.multiple_of(rm % m, block_m), block_m)
      rbn = tl.max_contiguous(tl.multiple_of(rn % n, block_n), block_n)
      rk = pid1 * block_k + tl.arange(0, block_k)
      lhs += (
          ram[:, None] * stride_am
          + rk[None, :] * stride_ak
          + pid2 * m * k
      )
      rhs += rk[:, None] * stride_bk + rbn[None, :] * stride_bn
      acc = tl.zeros((block_m, block_n), dtype=acc_ty)
      # for ki in range(0, k, block_k * split_k):  # pytype: disable=wrong-arg-types
      for ki in range(
          k, 0, -block_k * split_k
      ):  # pytype: disable=wrong-arg-types
          if even_k:
              a = tl.load(_fix_type_for_load(lhs))
              b = tl.load(_fix_type_for_load(rhs))
          else:
              a = tl.load(
                  _fix_type_for_load(lhs),
                  mask=rk[None, :] < ki,
                  other=0,
              )
              b = tl.load(
                  _fix_type_for_load(rhs),
                  mask=rk[:, None] < ki,
                  other=0,
              )
          casted_a = a.to(lhs.dtype.element_ty, bitcast=True).to(
              out.dtype.element_ty
          )
          casted_b = b.to(out.dtype.element_ty)
          acc += tl.dot(casted_a, casted_b, allow_tf32=True)
          lhs += block_k * split_k * stride_ak
          rhs += block_k * split_k * stride_bk
      acc = acc.to(out.dtype.element_ty)
      # rematerialize rm and rn to save registers
      rm = pid_m * block_m + tl.arange(0, block_m)
      rn = pid_n * block_n + tl.arange(0, block_n)
      out += (
          rm[:, None] * stride_cm
          + rn[None, :] * stride_cn
          + pid2 * m * n
      )
      out += m * n * pid1
      mask = (rm < m)[:, None] & (rn < n)[None, :]
      tl.store(out, acc, mask=mask)
  
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 5 for calling a global function with example usage: _matmul_kernel(lhs, rhs, out, m, n, k, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, block_m, block_n, block_k, group_m, split_k, acc_ty, force_num_warps, force_num_stages) without return types