code: def test_user_defined_persistent_non_warp_specialized_gemm(
      M,
      N,
      K,
      BLOCK_M,
      BLOCK_N,
      BLOCK_K,
      NUM_WARPS,
      NUM_CTAS,
      TRANS_A,
      TRANS_B,
      USE_TMA,
  ):
      if TRANS_A:
          a = (
              0.1
              * torch.randn(
                  (K, M), device="cuda", dtype=torch.float16
              ).T
          )
      else:
          a = 0.1 * torch.randn(
              (M, K), device="cuda", dtype=torch.float16
          )
  
      if TRANS_B:
          b = (
              0.1
              * torch.randn(
                  (N, K), device="cuda", dtype=torch.float16
              ).T
          )
      else:
          b = 0.1 * torch.randn(
              (K, N), device="cuda", dtype=torch.float16
          )
      c = torch.empty((M, N), device=a.device, dtype=torch.float32)
  
      NUM_SMS = torch.cuda.get_device_properties(
          "cuda"
      ).multi_processor_count
      grid = lambda META: (
          min(
              META["NUM_SMS"],
              triton.cdiv(M, META["BLOCK_M"])
              * triton.cdiv(N, META["BLOCK_N"]),
          ),
      )
  
      if USE_TMA:
          static_persistent_tma_matmul_kernel[grid](
              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=BLOCK_M,
              BLOCK_N=BLOCK_N,
              BLOCK_K=BLOCK_K,
              NUM_SMS=NUM_SMS,
              num_warps=NUM_WARPS,
              num_ctas=NUM_CTAS,
          )
      else:
          static_persistent_matmul_kernel[grid](
              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=BLOCK_M,
              BLOCK_N=BLOCK_N,
              BLOCK_K=BLOCK_K,
              NUM_SMS=NUM_SMS,
              num_warps=NUM_WARPS,
              num_ctas=NUM_CTAS,
          )
  
      th_c = torch.matmul(a, b)
      torch.testing.assert_close(
          th_c, c, atol=1e-2, rtol=0, check_dtype=False
      )
  
imports: import itertools
  import pytest
  import torch
  from torch.testing import assert_close
  import triton
  import triton.language
  from triton.runtime import driver
label_desc: Use image node_id 4 for calling a global function with example usage: test_user_defined_persistent_non_warp_specialized_gemm(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_WARPS, NUM_CTAS, TRANS_A, TRANS_B, USE_TMA) without return types