code: def bench_layer_norm(
      M, N, dtype, provider, mode="backward", eps=1e-5, device="cuda"
  ):
      # create data
      x_shape = (M, N)
      w_shape = (x_shape[-1],)
      weight = torch.rand(
          w_shape, dtype=dtype, device="cuda", requires_grad=True
      )
      bias = torch.rand(
          w_shape, dtype=dtype, device="cuda", requires_grad=True
      )
      x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
      dy = 0.1 * torch.randn_like(x)
      x.requires_grad_(True)
      quantiles = [0.5, 0.2, 0.8]
      # utility functions
      if provider == "triton":
  
          def y_fwd():
              return layer_norm(
                  x, w_shape, weight, bias, eps
              )  # noqa: F811, E704
  
      if provider == "torch":
  
          def y_fwd():
              return torch.nn.functional.layer_norm(
                  x, w_shape, weight, bias, eps
              )  # noqa: F811, E704
  
      if provider == "apex":
          apex_layer_norm = (
              apex.normalization.FusedLayerNorm(w_shape)
              .to(x.device)
              .to(x.dtype)
          )
  
          def y_fwd():
              return apex_layer_norm(x)  # noqa: F811, E704
  
      # forward pass
      if mode == "forward":
          gbps = lambda ms: 2 * x.numel() * x.element_size() / ms * 1e-6
          ms, min_ms, max_ms = triton.testing.do_bench(
              y_fwd, quantiles=quantiles, rep=500
          )
      # backward pass
      if mode == "backward":
  
          def gbps(ms):
              return (
                  3 * x.numel() * x.element_size() / ms * 1e-6
              )  # noqa: F811, E704
  
          y = y_fwd()
          ms, min_ms, max_ms = triton.testing.do_bench(
              lambda: y.backward(dy, retain_graph=True),
              quantiles=quantiles,
              grad_to_none=[x],
              rep=500,
          )
      return gbps(ms), gbps(max_ms), gbps(min_ms)
  
file_path: triton/python/tutorials/05-layer-norm.py
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 5 for calling a global function with example usage: bench_layer_norm(M, N, dtype, provider, mode, eps, device) and returns: gbps,  gbps,  gbps, layer_norm, torch, apex_layer_norm, unknown