code: def test_tma_load_store(
      M, N, K, NUM_CTAS, NUM_WARPS, TRANS_A, TRANS_B, OUTPUT_F16
  ):
      if TRANS_A:
          a = torch.randn((K, M), device="cuda", dtype=torch.float16).T
      else:
          a = torch.randn((M, K), device="cuda", dtype=torch.float16)
      if TRANS_B:
          b = torch.randn((N, K), device="cuda", dtype=torch.float16).T
      else:
          b = torch.randn((K, N), device="cuda", dtype=torch.float16)
  
      c = torch.empty((M, N), device=a.device, dtype=torch.float32)
      if OUTPUT_F16:
          c = torch.empty((M, N), device=a.device, dtype=torch.float16)
  
      matmul_tma_load_store[(1, 1)](
          a_ptr=a,
          b_ptr=b,
          c_ptr=c,  #
          M=M,
          N=N,
          K=K,  #
          stride_am=a.stride(0),
          stride_ak=a.stride(1),  #
          stride_bk=b.stride(0),
          stride_bn=b.stride(1),  #
          stride_cm=c.stride(0),
          stride_cn=c.stride(1),  #
          BLOCK_M=M,
          BLOCK_N=N,
          BLOCK_K=K,  #
          num_warps=NUM_WARPS,
          num_ctas=NUM_CTAS,  #
          OUTPUT_F16=OUTPUT_F16,
      )
      golden = torch.matmul(a, b)
      torch.set_printoptions(profile="full")
      assert_close(c, golden, rtol=1e-2, atol=1e-3, check_dtype=False)
  
imports: import pytest
  import torch
  from torch.testing import assert_close
  import triton
  import triton.language
label_desc: Use image node_id 2 for calling a global function with example usage: test_tma_load_store(M, N, K, NUM_CTAS, NUM_WARPS, TRANS_A, TRANS_B, OUTPUT_F16) without return types