code: def test_gemm(
      BLOCK_M,
      BLOCK_N,
      BLOCK_K,
      NUM_WARPS,
      NUM_CTAS,
      M,
      N,
      K,
      TRANS_A,
      TRANS_B,
      TRANS_OUTPUT,
      epilogue,
      out_dtype,
      USE_TMA_STORE,
      NUM_STAGES,
      ENABLE_WS,
  ):
      if "-".join(
          map(
              str,
              [
                  BLOCK_M,
                  BLOCK_N,
                  BLOCK_K,
                  NUM_WARPS,
                  NUM_CTAS,
                  M,
                  N,
                  K,
                  TRANS_A,
                  TRANS_B,
              ],
          )
      ) in [
          "16-32-64-4-4-512-256-64-True-False",
          "16-32-64-4-4-512-256-64-True-True",
          "16-32-64-4-4-512-256-64-False-False",
          "16-32-64-4-4-512-256-64-False-True",
      ]:
          pytest.skip("shapePerCTA[1] < 16 not supported")
  
      if "-".join(
          map(
              str,
              [
                  BLOCK_M,
                  BLOCK_N,
                  BLOCK_K,
                  NUM_WARPS,
                  NUM_CTAS,
                  M,
                  N,
                  K,
                  TRANS_B,
              ],
          )
      ) in [
          "16-32-64-4-1-256-256-256-False",
          "16-32-64-4-2-256-256-256-False",
          "16-32-64-4-2-256-256-256-True",
          "16-32-64-8-2-256-256-256-False",
          "16-32-64-8-2-256-256-256-True",
      ]:
          pytest.skip(
              "Known legacy issue, ldmatrix can only support x4"
          )
      enable_tma = os.environ.get("ENABLE_TMA", "not found").lower()
      if NUM_CTAS > 1 and enable_tma in ["on", "true", "1"]:
          pytest.skip(
              "multi-CTA with TMA not supported in MaterializeLoadStore"
          )
  
      M = BLOCK_M if M is None else M
      N = BLOCK_N if N is None else N
      K = BLOCK_K if K is None else K
  
      if TRANS_A:
          a = torch.randn((K, M), device="cuda", dtype=torch.float16).T
          a_order = [0, 1]
      else:
          a = torch.randn((M, K), device="cuda", dtype=torch.float16)
          a_order = [1, 0]
  
      if TRANS_B:
          b = torch.randn((N, K), device="cuda", dtype=torch.float16).T
          b_order = [0, 1]
      else:
          b = torch.randn((K, N), device="cuda", dtype=torch.float16)
          b_order = [1, 0]
  
      if out_dtype == "float16" and epilogue != "softmax":
          # TODO: for out_dtype == 'float16' and epilogue == 'softmax', it will
          # fail with the following error: 'llvm.fmul' op requires the same type
          # for all operands and results
          out_dtype = tl.float16
          torch_out_dtype = torch.float16
      else:
          out_dtype = tl.float32
          torch_out_dtype = torch.float32
  
      # avoid out of memory
      if epilogue in ["add-matrix", "add-rows", "add-cols"]:
          if TRANS_OUTPUT:
              bias = torch.randn(
                  (N, M), device="cuda", dtype=torch_out_dtype
              ).T
          else:
              bias = torch.randn(
                  (M, N), device="cuda", dtype=torch_out_dtype
              )
      else:
          bias = torch.randn(
              (1, 1), device="cuda", dtype=torch_out_dtype
          )
  
      # for chain-dot only
      w = torch.randn((N, N), device="cuda", dtype=torch.float16).T
      w_order = [0, 1]
  
      if TRANS_OUTPUT:
          z = torch.full(
              (N, M), 1.0, device="cuda", dtype=torch_out_dtype
          ).T
          z_order = [0, 1]
      else:
          z = torch.full(
              (M, N), 1.0, device="cuda", dtype=torch_out_dtype
          )
          z_order = [1, 0]
  
      # torch result
      a_f32 = a.to(torch.float32)
      b_f32 = b.to(torch.float32)
      dot = torch.matmul(a_f32, b_f32)
  
      def process_epilogue(d, bias, w, epilogue):
          if epilogue == "add-matrix":
              ref = d + bias
          elif epilogue == "add-rows":
              ref = d + bias[:, 0][:, None]
          elif epilogue == "add-cols":
              ref = d + bias[0, :][None, :]
          elif epilogue == "softmax":
              num = torch.exp(
                  d - torch.max(d, dim=-1, keepdims=True)[0]
              )
              denom = torch.sum(num, dim=-1, keepdims=True)
              ref = num / denom
              # ref = torch.softmax(d, 1)
          elif epilogue == "chain-dot":
              ref = torch.matmul(d, w.to(torch.float32))
          else:
              ref = d
          return ref
  
      golden = process_epilogue(dot, bias, w, epilogue)
  
      def grid(META):
          return (
              triton.cdiv(M, META["BLOCK_M"])
              * triton.cdiv(N, META["BLOCK_N"]),
          )
  
      pgm = matmul_kernel[grid](
          a_ptr=a,
          b_ptr=b,
          w_ptr=w,
          bias_ptr=bias,
          z_ptr=z,  #
          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_wm=w.stride(0),
          stride_wn=w.stride(1),  #
          stride_zm=z.stride(0),
          stride_zn=z.stride(1),  #
          BLOCK_M=BLOCK_M,
          BLOCK_N=BLOCK_N,
          BLOCK_K=BLOCK_K,
          GROUP_SIZE_M=8,  #
          out_dtype=out_dtype,  #
          USE_TMA_STORE=USE_TMA_STORE,  #
          ADD_MATRIX=epilogue == "add-matrix",  #
          ADD_ROWS=epilogue == "add-rows",  #
          ADD_COLS=epilogue == "add-cols",  #
          DO_SOFTMAX=epilogue == "softmax",  #
          CHAIN_DOT=epilogue == "chain-dot",  #
          A_ORDER_0=a_order[0],
          A_ORDER_1=a_order[1],  #
          B_ORDER_0=b_order[0],
          B_ORDER_1=b_order[1],  #
          W_ORDER_0=w_order[0],
          W_ORDER_1=w_order[1],  #
          Z_ORDER_0=z_order[0],
          Z_ORDER_1=z_order[1],  #
          num_warps=NUM_WARPS,
          num_ctas=NUM_CTAS,
          num_stages=NUM_STAGES,  #
          enable_warp_specialization=ENABLE_WS,
      )
  
      torch.set_printoptions(profile="full")
      golden = torch.nn.functional.normalize(golden)
      z = torch.nn.functional.normalize(z)
      assert_close(z, golden, rtol=1e-2, atol=1e-3, check_dtype=False)
  
      disable_mmav3 = os.environ.get(
          "DISABLE_MMA_V3", "not found"
      ).lower()
      if (
          disable_mmav3 not in ["on", "true", "1"]
          and BLOCK_M >= 64
          and NUM_CTAS == 1
          and BLOCK_N <= 256
      ):
          ptx = pgm.asm["ptx"]
          assert re.search(
              r"wgmma.mma_async.sync.aligned.m\d+n{}k16(?:.row.col)?.f32.f16.f16".format(
                  BLOCK_N
              ),
              ptx,
          )
  
file_path: triton/python/test/unit/hopper/test_gemm.py
imports: import itertools
  import os
  import re
  import pytest
  import torch
  from torch.testing import assert_close
  import triton
  import triton.language
label_desc: Use image node_id 4 for calling a global function with example usage: test_gemm(BLOCK_M, BLOCK_N, BLOCK_K, NUM_WARPS, NUM_CTAS, M, N, K, TRANS_A, TRANS_B, TRANS_OUTPUT, epilogue, out_dtype, USE_TMA_STORE, NUM_STAGES, ENABLE_WS) and returns: ref, unknown