code: def test_matmul(
      MODE,
      TRANS_A,
      TRANS_B,
      BLOCK,
      DTYPE,
      Z=3,
      H=2,
      M=512,
      N=384,
      K=256,
  ):
      seed = 0
      torch.manual_seed(seed)
      is_sdd = MODE == "sdd"
      is_dsd = MODE == "dsd"
      is_dds = MODE == "dds"
      do_sparsify = lambda x: sparsify_tensor(x, layout, BLOCK)
      do_mask = lambda x: mask_tensor(x, layout, BLOCK)
      # create inputs
      # create op
      a_shape = (Z, H, K, M) if TRANS_A else (Z, H, M, K)
      b_shape = (Z, H, N, K) if TRANS_B else (Z, H, K, N)
      c_shape = (Z, H, M, N)
      shape = {
          "sdd": (M, N),
          "dsd": (a_shape[2], a_shape[3]),
          "dds": (b_shape[2], b_shape[3]),
      }[MODE]
      layout = torch.randint(
          2, (H, shape[0] // BLOCK, shape[1] // BLOCK)
      )
      layout[1, 2, :] = 0
      layout[1, :, 1] = 0
      # create data
      a_ref, a_tri = make_pair(a_shape, alpha=0.1, dtype=DTYPE)
      b_ref, b_tri = make_pair(b_shape, alpha=0.1, dtype=DTYPE)
      dc_ref, dc_tri = make_pair(c_shape, dtype=DTYPE)
      # compute [torch]
      dc_ref = do_mask(dc_ref) if is_sdd else dc_ref
      a_ref = do_mask(a_ref) if is_dsd else a_ref
      b_ref = do_mask(b_ref) if is_dds else b_ref
      a_ref.retain_grad()
      b_ref.retain_grad()
      c_ref = torch.matmul(
          a_ref.transpose(2, 3) if TRANS_A else a_ref,
          b_ref.transpose(2, 3) if TRANS_B else b_ref,
      )
      c_ref.backward(dc_ref)
      c_ref = do_sparsify(c_ref) if is_sdd else c_ref
      da_ref = do_sparsify(a_ref.grad) if is_dsd else a_ref.grad
      db_ref = do_sparsify(b_ref.grad) if is_dds else b_ref.grad
      # triton result
      dc_tri = do_sparsify(dc_tri) if is_sdd else dc_tri
      a_tri = do_sparsify(a_tri) if is_dsd else a_tri
      b_tri = do_sparsify(b_tri) if is_dds else b_tri
      a_tri.retain_grad()
      b_tri.retain_grad()
      op = triton.ops.blocksparse.matmul(
          layout,
          BLOCK,
          MODE,
          trans_a=TRANS_A,
          trans_b=TRANS_B,
          device="cuda",
      )
      c_tri = op(a_tri, b_tri)
      c_tri.backward(dc_tri)
      da_tri = a_tri.grad
      db_tri = b_tri.grad
      # compare
      torch.testing.assert_close(c_ref, c_tri)
      torch.testing.assert_close(da_ref, da_tri)
      torch.testing.assert_close(db_ref, db_tri)
  
file_path: triton/python/test/unit/operators/test_blocksparse.py
imports: import pytest
  import torch
  import triton
  import triton.ops
label_desc: Use image node_id 4 for calling a global function with example usage: test_matmul(MODE, TRANS_A, TRANS_B, BLOCK, DTYPE, Z, H, M, N, K) without return types