code: def benchmark(size, provider: str, mode: str):
      print(
          f"{provider} benchmarking {mode} pass on ({size}, {size})..."
      )
      torch_nn = TorchNN(size)
      linear_weights = torch_nn.linear.state_dict()[
          "weight"
      ].T.contiguous()
      x = torch.randn((size, size), device="cuda", dtype=torch.float16)
      if provider == "pytorch":
          forward_f = lambda: torch_nn.forward(x)
      elif provider == "triton":
          forward_f = lambda: partial_gelu(x, linear_weights)
      else:
          assert False, "Got invalid provider {}".format(provider)
      if mode == "forward":
          ms, min_ms, max_ms = triton.testing.do_bench(forward_f)
      elif mode == "backward":
          x.requires_grad_(True)
          forward = forward_f()
          dA = torch.randn(
              size, size * 8 // 2, device=x.device, dtype=x.dtype
          )
          ms, min_ms, max_ms = triton.testing.do_bench(
              lambda: forward.backward(dA, retain_graph=True)
          )
      else:
          assert False, "Got invalid mode {}".format(mode)
  
      # TODO: This perf function should be better estimated for the forward and backward passes
      def perf(ms):
          return 2 * size * size * size * 1e-12 / (ms * 1e-3)
  
      return perf(ms), perf(max_ms), perf(min_ms)
  
imports: from typing import Tuple
  import torch
  import triton
  import triton.language
  import numpy
label_desc: Use image node_id 8 for calling a global function with example usage: benchmark(size, provider, mode) and returns: perf,  perf,  perf, unknown