code: def benchmark_(S, B, hidden_size, itype, wtype, runs=100):
      epsilon = 1e-5
  
      x = torch.randn((S * B, hidden_size), dtype=itype, device=device)
      beta = torch.randn(hidden_size, dtype=wtype, device=device)
      gamma = torch.randn(hidden_size, dtype=wtype, device=device)
      dz = torch.randn(x.shape, dtype=wtype, device=device)
  
      stream = torch.cuda.Stream()
      with torch.cuda.stream(stream):
          timer = GPUTimer(stream)
  
          # warmup
          for r in range(runs):
              z, mu, rsigma = fln.ln_fwd(x, gamma, beta, epsilon)
  
          timer.start()
          for r in range(runs):
              z, mu, rsigma = fln.ln_fwd(x, gamma, beta, epsilon)
          timer.stop()
          timer.sync()
  
          total_bytes_fwd = sum(
              [
                  size_in_bytes(t)
                  for t in [x, z, gamma, beta, mu, rsigma]
              ]
          )
  
          ms_fwd = timer.millis() / runs
  
          print(
              "[FWD] Time: {:.4f}ms Throughput: {:.4f} GB/sec".format(
                  ms_fwd, total_bytes_fwd * 1e-6 / ms_fwd
              )
          )
  
          timer.start()
          for r in range(runs):
              dx, dgamma, dbeta, dbp, dgp = fln.ln_bwd(
                  dz, z, mu, rsigma, gamma, beta, True
              )
          timer.stop()
          timer.sync()
  
          total_bytes_bwd = sum(
              [
                  size_in_bytes(t)
                  for t in [
                      dz,
                      x,
                      mu,
                      rsigma,
                      gamma,
                      dx,
                      dgamma,
                      dbeta,
                      dbp,
                      dbp,
                      dgp,
                      dgp,
                  ]
              ]
          )
  
          ms_bwd = timer.millis() / runs
  
          print(
              "[BWD] Time: {:.4f}ms Throughput: {:.4f} GB/sec".format(
                  ms_bwd, total_bytes_bwd * 1e-6 / ms_bwd
              )
          )
  
file_path: apex/apex/contrib/test/layer_norm/test_fast_layer_norm.py
imports: import itertools
  import unittest
  import torch
label_desc: Use image node_id 4 for calling a global function with example usage: benchmark_(S, B, hidden_size, itype, wtype, runs) without return types