code: def test_gemm_fusion():
      M, N, K = 4096, 4096, 64
      BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 64
      A = torch.empty(
          (M, K), dtype=torch.float16, device="cuda"
      ).normal_(mean=0.1, std=0.2)
      B = torch.empty(
          (N, K), dtype=torch.float16, device="cuda"
      ).normal_(mean=0.1, std=0.2)
      C = torch.empty(
          (N, K), dtype=torch.float16, device="cuda"
      ).normal_(mean=0.1, std=0.2)
      E = torch.empty((M, K), dtype=torch.float16, device="cuda")
      ref_out = torch.matmul(torch.matmul(A, B.T), C)
      num_warps = 4
      grid = (triton.cdiv(M, BLOCK_M), 1)
      gemm_fusion_kernel[grid](
          A,
          B,
          C,
          E,
          M,
          N,
          K,  #
          A.stride(0),
          A.stride(1),  #
          B.stride(0),
          B.stride(1),  #
          C.stride(0),
          C.stride(1),  #
          E.stride(0),
          E.stride(1),  #
          BLOCK_M,
          BLOCK_N,
          BLOCK_K,  #
          num_warps=num_warps,
      )
  
      torch.testing.assert_close(ref_out, E, atol=1e-2, rtol=0)
  
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 2 for calling a global function with example usage: test_gemm_fusion() without return types