text: #!/tools/lm-venv/py3.6-tf-1.3.0-svail/bin/python
  
  import click
  import math
  import os
  import sys
  import config
  import shutil
  import itertools
  import numpy as np
  
  from parallelism import Parallelism
  from topology import Topology
  from simulate import Graph
  import util
  from hw_component import Core, MemoryHierarchy, Network
  from model import Model
  
  algByte=False #algorithmic ops false
  proj=False #consider projection layer, turn off for end-2-end validation, as baeline model does not have projection layer
  validating_v100=True
  
  class TimeCalculation:
      def __init__(self, exp_config):
          #Model Parameters
          self.model              = Model(exp_config)
          self.B                  = self.model.batch_size
          self.V                  = self.model.vocab_size
          self.L                  = self.model.num_layers
          self.D                  = self.model.hidden_dim
          self.projection         = self.model.projection
          self.S                  = self.model.seq_len
          self.G                  = self.model.num_gates
          self.NL                 = self.model.num_non_linear
          self.A                  = self.model.num_add
          self.P                  = self.model.num_pointwise
         
          #Software Parameters
          self.O                  = exp_config.sw_config.kernel_launch_overhead
          self.precision          = exp_config.sw_config.precision
          self.attached           = True
          
  
          #Hardware Parameters
          self.core               = Core(exp_config)
          self.th                 = self.core.getThroughput()
          self.FMA_width          = self.core.FMA_width
          self.dataflow          = self.core.dataflow
  
          self.memoryHierarchy     = MemoryHierarchy(exp_config)
          self.num_levels          = self.memoryHierarchy.num_levels
          self.memLayer            = self.memoryHierarchy.memLayer
          self.tileSpace           = self.generateTileSpace()
        
          #TODO: move this to config file
          self.H2Dbw               = 12.4*1024*1024*1024 
  
          #System Parameters
          self.num_wafer          = exp_config.system_config.num_wafers
          self.num_workers        = exp_config.system_config.num_workers
  
          self.network            = Network(exp_config)
  
          intra_throughput, inter_throughput = self.network.calcThroughput()
          intra_latency, inter_latency       = self.network.calcLatency()
                  
          inter_derate            = exp_config.system_config.inter_derate
          intra_derate            = exp_config.system_config.intra_derate
          par2cross               = exp_config.system_config.par2cross
          
          derated_inter_throughput = -1
          derated_intra_throughput = -1
          
          #cross-wafercommunications will pass through intra links too
          if self.num_wafer > 1 and self.num_workers > 1:
            if intra_derate != 0:
              derated_inter_throughput = min(intra_throughput/intra_derate, 
                                             inter_throughput/inter_derate)
            else:
              derated_inter_throughput = inter_throughput/inter_derate
          else:
              derated_inter_throughput = 0
  
          if self.num_workers > 1 and intra_derate != 0:
                derated_intra_throughput = intra_throughput/intra_derate
          else:
              derated_intra_throughput = 0
  
          self.IBK1, self.LLK1    = ((derated_inter_throughput, inter_latency) if par2cross["kp1"] else 
                                     (derated_intra_throughput, intra_latency)) 
          self.IBK2, self.LLK2    = ((derated_inter_throughput, inter_latency) if par2cross["kp2"] else 
                                     (derated_intra_throughput, intra_latency)) 
          self.IBD, self.LLD      = ((derated_inter_throughput, inter_latency) if par2cross["dp"] else 
                                     (derated_intra_throughput, intra_latency)) 
          self.IBL, self.LLL      = ((derated_inter_throughput, inter_latency) if par2cross["lp"] else 
                                     (derated_intra_throughput, intra_latency)) 
          
  
          #Scheduling Parameters
          par                     = Parallelism(exp_config)
          par.findParallelStrategy()
          self.autoPar            = par.autoPar
          self.lp                 = par.lp
          self.kp_hidden_dim1     = par.kp_hidden_dim1
          self.kp_softmax_dim1    = par.kp_softmax_dim1
          self.kp_embedding_dim1  = par.kp_embedding_dim1
          self.kp_projection_dim1 = par.kp_projection_dim1
          self.kp_hidden_dim2     = par.kp_hidden_dim2
          self.kp_softmax_dim2    = par.kp_softmax_dim2
          self.kp_embedding_dim2  = par.kp_embedding_dim2
          self.kp_projection_dim2 = par.kp_projection_dim2
          self.dp                 = par.dp
          self.kp_hidden_type     = par.kp_hidden_type #1: CR, 2: RC
          self.kp_softmax_type    = par.kp_softmax_type #1: CR, 2: RC
          self.kp_embedding_type  = par.kp_embedding_type #1: CR, 2: RC
          self.kp_projection_type = par.kp_projection_type #1: CR, 2: RC
  
  
  
          #Define miniBatch size
          self.miniB              = math.ceil(self.B / self.dp)
  
          #Statistics Param
          self.tot_flop           = 0
          self.tot_mem            = 0
          self.tot_time           = 0
          self.debug              = False
          self.validating_GEMM    = False
      
      def updateParams(self, debug, m, n, k, t, kp1, kp2, dp, lp, gemm,
                        batch_size, hidden_dim, seq_len, vocab_size, num_layer):
  
          self.B = batch_size
          self.D = hidden_dim
          self.S = seq_len
          self.V = vocab_size
          self.L = num_layer
          
          #Define miniBatch size
          self.dp = dp if dp != None else self.dp
          self.miniB              = math.ceil(self.B / self.dp)
          
          self.debug = debug
          self.validating_GEMM = gemm
          self.lp = lp if lp != None else self.lp
          self.kp_hidden_dim1 = kp1 if kp1 != None else self.kp_hidden_dim1
          self.kp_hidden_dim1 = kp1 if kp1 != None else self.kp_hidden_dim1
          self.kp_hidden_dim2 = kp2 if kp2 != None else self.kp_hidden_dim2
          self.kp_hidden_type = (2 if t == 'RC' else (1 if t == 'CR' else self.kp_hidden_type))
          
          #TODO: decide if we want kp1, kp2 to control other layers besides hidden layer
          self.kp_softmax_dim1 = kp1 if kp1 != None else self.kp_softmax_dim1
          self.kp_softmax_dim2 = kp2 if kp2 != None else self.kp_softmax_dim2
          self.kp_softmax_type = (2 if t == 'RC' else (1 if t == 'CR' else self.kp_softmax_type))
      
          self.kp_embedding_dim1 = kp1 if kp1 != None else self.kp_embedding_dim1
          self.kp_embedding_dim2 = kp2 if kp2 != None else self.kp_embedding_dim2
          self.kp_embedding_type = (2 if t == 'RC' else (1 if t == 'CR' else self.kp_embedding_type))
          
          self.kp_projection_dim1 = kp1 if kp1 != None else self.kp_projection_dim1
          self.kp_projection_dim2 = kp2 if kp2 != None else self.kp_projection_dim2
          self.kp_projection_type = (2 if t == 'RC' else (1 if t == 'CR' else self.kp_projection_type))
          
          #TODO: need to change all equations to be a function of m,n and k
          #self.D              = n//4
          
          print("kp1: {}".format(self.kp_hidden_dim1))
          print("kp2: {}".format(self.kp_hidden_dim2))
          #TODO: It is a hacky way of capturing assymetry across links within V100
          #move this to network topology and distinguish between inter and intra network
          if validating_v100:
            self.IBK1 = util.scale_down(self.IBK1, self.kp_hidden_dim1, "kp1")
            self.IBK2 = util.scale_down(self.IBK2, self.kp_hidden_dim2, "kp2")
            self.IBD  = util.scale_down(self.IBD, self.dp, "dp")
            self.IBL  = util.scale_down(self.IBL, self.lp, "lp")
               
      #Number of parameters
      def tot_param(self):
          embedding = self.V * self.D
          hidden = (2 * self.D + 1) * (self.G * self.D) * self.L
          projection = self.D * self.projection
          softmax = ((self.projection if proj else self.D) + 1) * self.V
  
          tot_param = embedding + hidden + projection + softmax
          return tot_param
  
      def printSysConfig(self, exp_config, output_file):
  
        kiloByte = 1024
        megaByte = kiloByte * 1024
        gigaByte = megaByte * 1024
        teraByte = gigaByte * 1024
  
        with open(output_file, "w") as f:
            f.write("==========================\n")
            f.write("Hardware Configuration\n")
            f.write("==========================\n")
  
            f.write("Throughput: {:.5f} Tflops\n".format(self.core.operating_throughput/1e12))
            for i in range(self.num_levels-1, -1, -1):
                mem_bw    = self.memLayer[i].dynamic_throughput
                mem_size  = self.memLayer[i].size
  
                if mem_bw < 1e3 * gigaByte:
                    f.write("L{:} Bandwidth: {:.1f} GB/s\n".format(i, mem_bw/(gigaByte)))
                else:
                    f.write("L{:} Bandwidth: {:.1f} TB/s\n".format(i, mem_bw/(teraByte)))
  
                if mem_size < 1e3 * megaByte:
                    f.write("L{:} Size: {:.10f} MB\n".format(i, mem_size/(megaByte)))
                elif mem_size < 1e3 * gigaByte:
                    f.write("L{:} Size: {:.1f} GB\n".format(i, mem_size/(gigaByte)))
                else:
                    f.write("L{:} Size: {:.1f} TB\n".format(i, mem_size/(teraByte)))
            
            f.write("Intra-node Bandwidth: {:.1f} GB/s\n".format(self.network.intra_network.throughput/(gigaByte)))
            f.write("Inter-node Bandwidth: {:.1f} GB/s\n".format(self.network.inter_network.throughput/(gigaByte)))
            
            M = self.memLayer[self.num_levels - 1].size
            tot_mem, embedding_mem, hidden_mem, softmax_mem, projection_mem, wt_mem, act_mem, point_mem = util.getTotMemReq(exp_config, 
                                                                                                                            batch_size = self.B,
                                                                                                                            hidden_dim = self.D,
                                                                                                                            vocab_size = self.V,
                                                                                                                            seq_len = self.S,
                                                                                                                            num_layer = self.L,
                                                                                                                            dp = self.dp,
                                                                                                                            lp = self.lp,
                                                                                                                            kp1 = self.kp_hidden_dim1,
                                                                                                                            kp2 = self.kp_hidden_dim2,
                                                                                                                            kp_type = self.kp_hidden_type)
            f.write("\n\n===========================================\n")
            f.write("Memory Requirement Breakdown per Data Shard\n")
            f.write("===========================================\n")
            f.write("Total Memory: {:.1f} GB\n"
                    "Embedding Memory: {:.1f} GB\n"
                    "Hidden Memory: {:.1f} GB\n"
                    "Softmax Memory: {:.1f} GB\n"
                    "Projection Memory: {:.1f} GB\n"
                    .format(tot_mem/gigaByte, 
                            embedding_mem/gigaByte, 
                            hidden_mem/gigaByte, 
                            softmax_mem/gigaByte, 
                            projection_mem/gigaByte))
            
            f.write("\nTotal Memory: {:.1f} GB\n"
                    "Weight Memory: {:.1f} GB\n"
                    "Activation Memory: {:.1f} GB\n"
                    "Pointwise Memory: {:.1f} GB\n"
                    .format(tot_mem/gigaByte, 
                            wt_mem/gigaByte, 
                            act_mem/gigaByte, 
                            point_mem/gigaByte))
  
  
            f.write("\nMemory Overflow Rate (Total Memory Required per Data Shard / Memory capacity per node): {:.1f}\n".format(float("inf") if M==0 else tot_mem/M))
  
            
            tot_mem, embedding_mem, hidden_mem, softmax_mem, projection_mem, wt_mem, act_mem, point_mem = util.getMemUsagePerCore(exp_config,
                                                                                                                                  batch_size = self.B,
                                                                                                                                  hidden_dim = self.D,
                                                                                                                                  vocab_size = self.V,
                                                                                                                                  seq_len = self.S,
                                                                                                                                  num_layer = self.L,
                                                                                                                                  dp = self.dp,
                                                                                                                                  lp = self.lp,
                                                                                                                                  kp1 = self.kp_hidden_dim1,
                                                                                                                                  kp2 = self.kp_hidden_dim2,
                                                                                                                                  kp_type = self.kp_hidden_type)
            f.write("\n\n===========================================================\n")
            f.write("Memory Requirement Breakdown per Data Shard Per Model Shard\n")
            f.write("===========================================================\n")
            f.write("Total Memory: {:.1f} GB\n"
                    "Embedding Memory: {:.1f} GB\n"
                    "Hidden Memory: {:.1f} GB\n"
                    "Softmax Memory: {:.1f} GB\n"
                    "Projection Memory: {:.1f} GB"
                       .format(tot_mem/gigaByte, 
                               embedding_mem/gigaByte, 
                               hidden_mem/gigaByte, 
                               softmax_mem/gigaByte, 
                               projection_mem/gigaByte))
  
            f.write("\nMemory Overflow Rate (Total Memory Required per Data Shard Per Model Shard/ Memory capacity per node): {:.5f}\n"
                  .format(float("inf") if M == 0 else tot_mem/M))
  
            f.write("\nTotal Memory: {:.1f} GB\n"
                    "Weight Memory: {:.1f} GB\n"
                    "Activation Memory: {:.1f} GB\n"
                    "Pointwise Memory: {:.1f} GB\n"
                    .format(tot_mem/gigaByte, 
                            wt_mem/gigaByte, 
                            act_mem/gigaByte, 
                            point_mem/gigaByte))
            
            f.write("\n\n====================\n")
            f.write("Parallelism Strategy\n")
            f.write("====================\n")
            f.write("dp: {}, lp: {}, kp_hidden_dim1: {}, kp_hidden_dim2: {}," 
                    "kp_softmax_dim1: {}, kp_softmax_dim2: {}, kp_embedding1: {}, kp_embedding2: {}," 
                    "kp_projection_dim1: {}, kp_proejction_dim2: {}\n"
                    .format(self.dp, self.lp, self.kp_hidden_dim1, self.kp_hidden_dim2, 
                     self.kp_softmax_dim1, self.kp_softmax_dim2, self.kp_embedding_dim1, self.kp_embedding_dim2,
                     self.kp_projection_dim1, self.kp_projection_dim2))   
  
  
            f.write("\n\n==============================================================================\n")
            f.write("Hardware Component Stats\n")
            f.write("==============================================================================\n")
            self.core.printStats(f)
            for i in range(0, self.num_levels):
                self.memLayer[i].printStats(f)
  
            self.network.printStats(f) 
  
      def roofline(self, flop, mem_access_, name=''):
  
          #print("Roofline: entered {}".format(name))
          mem_access = []
          if isinstance(mem_access_, int):
              mem_access.append(mem_access_)
          elif isinstance(mem_access_, float):
              mem_access.append(int(mem_access_))
          elif isinstance(mem_access_, list):
              mem_access = mem_access_
          else:
              print(mem_access_)
              print("mem_access_ should be inetger or list, wrong input", flush=True)
              sys.exit(0)
               
  
            
          num_level        = len(mem_access)
          time             = [0] * num_level
          comp_int         = [0] * num_level
          inflection_point = [0] * num_level
  
          try:
              assert(mem_access[num_level - 1] > 0) , "last_level_mem = 0"
          except Exception as e:
              print("{}: Number of accesses to the last level of memory hierarchy cannot be zero:\n {}".format(name, e), flush=True)
              sys.exit(0)
  
          for i in range(0, num_level):
              time[i]      = 0
              mem_bw       = self.memLayer[i].getThroughput()
              mem_latency  = self.memLayer[i].getLatency()
              num_mem      = mem_access[i] 
              inflection_point[i]   = float("inf") if mem_bw == 0 else self.th / mem_bw
              comp_int[i]  = 0 if num_mem == 0 else flop / num_mem
       
              if comp_int[i] < inflection_point[i]: #mem-bound
                  time[i] = (float("inf") if (mem_bw == 0 or num_mem == 0) else (num_mem / mem_bw)) + mem_latency
              else: #compute-bound
                  time[i] = float("inf") if (self.th == 0) else (flop / self.th)
          
         
          max_time = max(time)
          
          
          if self.debug:
              print('{}: {}'.format(name, max_time))
              print('GEMM flops: {:,}'.format(flop))
              for i in range(0, num_level):
                  print("L{}".format(i))
                  print("inflection_point: {:.2f}".format(inflection_point[i]))
                  print("comp_int: {:.2f}".format(comp_int[i]))
                  print("time: {}".format(time[i]))
                  print()
          
          #print("Roofline: exited {}".format(name))
          return max_time
  
      
  
      #Convert GEMM into sqaure tiles
     # def getGEMMTime(self, A_, B_, C_, name):
     #     
     #     #A = util.power2RoundUp(A_)
     #     #B = util.power2RoundUp(B_)
     #     #C = util.power2RoundUp(C_)
     #     A = A_
     #     B = B_
     #     C = C_
  
     #     #return False, self.GEMM_wrapper(A, B, C, name)
     #     dim        = min(min(A, B), C)
     #     Af         = math.ceil(A / dim)
     #     Bf         = math.ceil(B / dim)
     #     Cf         = math.ceil(C / dim)
  
     #     time       = (Af * Bf * Cf) * self.GEMM_Strassen(dim, name) + (Af * Cf * (Bf-1)) * self.getAddTime(dim, dim, name)
     #     return False, time
  
     # def GEMM_Strassen(self, dim, name):
     #     if dim <= 512:
     #         time = self.GEMM_wrapper(dim, dim, dim, name)
     #         return time
     #     else:
     #         time = 7 * self.GEMM_Strassen(dim // 2, name) #+ 18 * self.getAddTime(dim // 2, dim // 2, name)
     #         return time
     #  
     # def getAddTime(self, A, B, name):
     #     ADD_flop  = A * B
     #     ADD_gmem  = 3 * A * B * self.precision
     #     ADD_time  = self.roofline(ADD_flop, ADD_gmem, name='FMA addition') + self.O
     #     return ADD_time
  
      def getGEMMTime(self, dim1, dim2, dim3, name):
         tile2time = {}
         orderSpace = self.generateOrder(dim1, dim2, dim3, name)
         for order_dims in orderSpace:
            if self.debug:
              print("===============================================================")
              print("order: {}".format(order_dims))
              print("===============================================================")
            for tile_dims in self.tileSpace:
              if self.debug:
                print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
                print("tile: {}".format(tile_dims))
                print("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
              GEMM_flop, mem_access = self.GEMM(order_dims, tile_dims, name)
              GEMM_time = self.roofline(GEMM_flop,mem_access, name) + self.O
              tile2time[(order_dims, tile_dims)] = GEMM_time
  
         
         best_tile = min(tile2time, key=tile2time.get)
         best_time = tile2time[best_tile]
  
         if self.debug:
            print("{}: Best Time: {:,}, Best Order: {}, Best Tile: {}\n".format(name, best_time, best_tile[0], best_tile[1]))
  
         return best_time, best_tile[0], best_tile[1]
      
      def generateOrder(self, dim1, dim2, dim3, name):
  
          if self.dataflow =="best": # best stationary
             if dim1 >= max(dim2, dim3):
                self.dataflow = "wst"
             elif dim2 >= max(dim1, dim3):
                self.dataflow = "ost"
             elif dim3 >= max(dim1, dim2):
                self.dataflow = "ast"
  
          order=[]
          if self.dataflow == "wst": #weight stationary
              order.append((dim2, dim3, dim1))
              if dim2 != dim3:
                  order.append((dim3, dim2, dim1))
          elif self.dataflow == "ast": #activation stationary
              order.append((dim1, dim2, dim3))
              if dim2 != dim1:
                  order.append((dim2, dim1, dim3))
          elif self.dataflow == "ost": #output stationary
              order.append((dim1, dim3, dim2))
              if dim1 != dim3:
                  order.append((dim3, dim1, dim2))
          elif self.dataflow == "none": # not stationary
              if dim1 != dim2 and dim2 != dim3 and dim1 != dim3:
                  order=list(itertools.permutations([dim1, dim2, dim3]))
              elif dim1 == dim2 and dim2 != dim3:
                  order = [(dim1, dim2, dim3), (dim1, dim3, dim2), (dim3, dim1, dim2)]
              elif dim1 == dim3 and dim2 != dim1:
                  order = [(dim1, dim2, dim3), (dim1, dim3, dim2), (dim2, dim1, dim3)]
              elif dim2 == dim3 and dim1 != dim2:
                  order = [(dim1, dim2, dim3), (dim2, dim1, dim3), (dim2, dim3, dim1)]
  
          return order
  
      def generateTileSpace(self):
          tile_space = []
          tiles = [None] * self.num_levels 
          
          for level in range(0, self.num_levels-1):
              memory = self.memLayer[level]
              #tiles[level] = self.getTileDims(memory)
              tiles[level] = memory.getTileDims()
          
          if self.num_levels == 1:
              tile_space = []
          elif self.num_levels == 2:
              tile_space = tiles[0]
          elif self.num_levels == 3:
              tile_space = [(x,y) for x in tiles[0] for y in tiles[1]]
          elif self.num_levels == 4:
              tile_space = [(x,y,z) for x in tiles[0] for y in tiles[1] for z in tiles[2]]
          else: 
            raise NotImplementedError()
  
          return tile_space
  
      def getTileSize(self, lid):
          memory = self.memLayer[lid]
          memory.calcTileDim()
          tile_dim  = memory.getTileDim()
          return tile_dim, tile_dim, tile_dim
  
      #Count the number of accesses from level-1 to level
      # input matrix A(dim1, dim2) and B(dim2, dim3)
      # output matrix C(dim1, dim3)
      def getNumAccesses(self, level, dim1, dim2, dim3, tile_dim, num_repeat, name):
          #tile1,tile2,tile3 = self.getTileSize(level-1)
          tile1, tile2, tile3 = tile_dim
  
          orig_size = tile1*tile2 + tile1*tile3 + tile2*tile3
          short_tile_cond = [0,0,0]
  
          if tile1 > dim1:
              tile1 = dim1
              short_tile_cond[0] = 1  
          if tile2 > dim2:
              tile2 = dim2
              short_tile_cond[1] = 1
          if tile3 > dim3:
              tile3 = dim3
              short_tile_cond[2] = 1
  
          if short_tile_cond[2] == 0 and (short_tile_cond[0] | short_tile_cond[1]) == 1:
              if level <= 1:
                tile3 = math.floor((orig_size - tile1 * tile2) / (tile1 + tile2))
              else:
              #store bypasses cache, directly goes to memory 
                tile3 = math.floor((orig_size - tile1 * tile2) / tile2)
              if tile3 > dim3:
                tile3 = dim3
              #Uncomment if tile3 needs to be pow of 2
              #tile3 = int(math.pow(2, math.floor(math.log2(tile3))))
          elif short_tile_cond[0] == 0 and (short_tile_cond[1] | short_tile_cond[2]) == 1:
              if level <= 1:
                tile1 = math.floor((orig_size - tile3 * tile2) / (tile3 + tile2))
              else:
              #store bypasses cache, directly goes to memory 
                tile1 = math.floor((orig_size - tile3 * tile2) / tile2)
              if tile1 > dim1:
                tile1 = dim1
          elif short_tile_cond[1] == 0 and (short_tile_cond[0] & short_tile_cond[2]) == 1:
              if level <= 1:
                tile2 = math.floor((orig_size - tile3 * tile1) / (tile3 + tile1))
              else:
                tile2 = math.floor((orig_size) / (tile1 + tile3))
              if tile2 > dim2:
                tile2 = dim2
  
          reload_A = 1
          reload_B = 1
          reload_C = 1
  
          if tile1 > 0 and tile2 > 0 and tile3 > 0:
             reload_A = math.ceil(dim3 / tile3)
             reload_B = math.ceil(dim1 / tile1)
             #do not access the slow memory on every write,acculmuate in fast memory 
             reload_C = (1 if level > 1 else math.ceil(dim2 / tile2))
             
           
          num_mem = num_repeat * (dim1 * dim2 * reload_A + dim2 * dim3 * reload_B + dim1 * dim3 * reload_C) * self.precision
          if self.debug:
             print(name)
             print("Matrix dimension at Level {}: {:,} x {:,} x {:,}".format(level, dim1, dim2, dim3))
             print("Tile dimension at Level {}: {:,} x {:,} x {:,}".format(level-1, tile1, tile2, tile3))
             print("reload_A: {}, reload_B: {}, reload_C: {}".format(reload_A, reload_B, reload_C))
             print("num_repeat: {}".format(num_repeat))
             print("Bytes Accessed: {:,}".format(num_mem))
             print("")
  
          return num_mem, tile1, tile2, tile3
          
  
      #This is the main function that captures the memory hierarchy impact
      #on the number of accesses to global memory considering not everything fits in 
      #L2 cache and also captures the effect of shared memory
      def GEMM(self, order_dims, tile_dims, name):
          dim1_ = order_dims[0]
          dim2_ = order_dims[1]
          dim3_ = order_dims[2]
          #dim1 = util.power2RoundUp(dim1_)
          #dim2 = util.power2RoundUp(dim2_)
          #dim3 = util.power2RoundUp(dim3_)
          dim1 = dim1_
          dim2 = dim2_
          dim3 = dim3_
  
          GEMM_flop = dim1 * dim3 * (dim2 + dim2 - 1)
          #dim2 multiply
          #dim2-1 add
  
          #X1 = self.L2_tile_dim
          #X2 = self.shared_mem_tile_dim
          #X3 = self.reg_tile_dim
  
          num_accesses = [0] * self.num_levels
  
          if (algByte):
              num_accesses[self.num_levels - 1] = (dim1 * dim2 + dim2 * dim3 + dim1 * dim3) * self.precision
          else:
              num_repeat = 1
              for level in range(self.num_levels - 1, 0, -1):
                  num_accesses[level], tile1, tile2, tile3  = self.getNumAccesses(level, dim1, dim2, dim3, tile_dims[level-1], num_repeat, name)
                  try:
                    num_repeat           *= math.ceil(dim1/tile1) * math.ceil(dim2/tile2) * math.ceil(dim3/tile3)
                  except:
                    num_repeat           *= 1
  
                  dim1                   = tile1 if tile1 != 0 else dim1
                  dim2                   = tile2 if tile2 != 0 else dim2
                  dim3                   = tile3 if tile3 != 0 else dim3
  
  
              #Number of accesses to level0 (for every 2N^3 computation, 3N^2 memory accesses happen, where N is the width of the systolic engine)
              reuse = 1
              dim1 = dim1_
              dim2 = dim2_
              dim3 = dim3_
              
              if self.dataflow == "none":
                  reuse = 1
              elif self.dataflow == "best":
                  reuse = max(math.ceil(dim1/self.FMA_width), math.ceil(dim3/self.FMA_width), math.ceil(dim2/self.FMA_width))
              elif self.dataflow == "wst": #wt stationary
                  reuse = math.ceil(dim1/self.FMA_width)
              elif self.dataflow == "ast": #act statinary
                  reuse = math.ceil(dim3/self.FMA_width)
              elif self.dataflow == "ost": #output stationary
                  reuse = math.ceil(dim2/self.FMA_width)
              else:
                  raise NotImplementedError()
                 
              #TODO: make sure to model underutilized systolic array
              #TODO: support FMA_width_x and FMA_width_y
              num_accesses[0]    = GEMM_flop * ((2 * reuse + 1) / (2 * reuse)) * 1/self.FMA_width * self.precision
              #num_accesses[0]    = GEMM_flop * ((2 * reuse + self.FMA_width) / (2 * reuse)) * 1/self.FMA_width * self.precision
               
              #TODO: do we still need these in new hierarchical version?
              #  if X3 == 0:
              #    GEMM_smem  = GEMM_rmem
              #    GEMM_rmem  = 0
              #  if X2 == 0:
              #    GEMM_l2mem = GEMM_smem
              #    GEMM_smem = 0
              #  if X1 == 0:
              #    GEMM_gmem  = GEMM_l2mem
              #    GEMM_l2mem = 0
  
              #  try:
              #    GEMM_l2mem = GEMM_smem
              #    GEMM_smem = 0
              #  if X1 == 0:
              #    GEMM_gmem  = GEMM_l2mem
              #    GEMM_l2mem = 0
  
          return GEMM_flop, num_accesses
  
  
      #Column-Row MM
      def getCf_kp1(self):
          #Multiply
          assert(self.kp_hidden_type == 1)
          assert(self.kp_hidden_dim1 > 1)
          assert(self.kp_hidden_dim1 % 4 == 0 or self.kp_hidden_dim1 == 2) #4 bc it is LSTM cell
          assert((2 * self.D) % self.kp_hidden_dim1 == 0)
          GEMM_time, reduction_time = self.getDistGEMM_f_kp1(self.miniB, 2 * self.D, self.G * self.D, self.kp_hidden_dim1, "Cf_kp1")
          
         
          #Pointwise ops: all the linear/non-linear ops after MM
          point_flop = self.miniB * (self.G * self.D / self.kp_hidden_dim1) * 5
          #4 refers to the number of pointwise ops (mul + add +tanh + mul + tanh) on 
          #the critical path 
          point_mem  = (self.precision * self.miniB * (self.G * self.D / self.kp_hidden_dim1) *
                       (3 * 3 + 2 * 2 ))
          # 3(3 memory access per operation with two input and one output)
          # 3(mul +  add + mul) on critical path
          # 2(2 memory access per operation with one input and one output)
          # 1(tanh) on critical path
          
          data_size = 4 * self.miniB * (self.G * self.D / self.kp_hidden_dim1) * self.precision
          # 4 refers to the number of pointwise ops (mul + add + mul + tanh) on the
          # critical path whose inputs are located across different GPUs
          #NOTE:Assuming all communications can happpen in parallel
          mem_transfer = self.roofline(0,  2 * data_size, name="Cf_kp1: memory accesses before and after data transfer over network")
          #2:  one read from the source and one write to the destination memory
          data_transfer =  data_size / self.IBK1 
          point_comm = mem_transfer + data_transfer
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise_cf_kp1') + 5 * self.O + point_comm
  
  
          return GEMM_time + reduction_time + point_time
   
      def getCb_kp1(self):
          #TODO:Add local accumulation of weights at every time step
          #Pointwise
          point_flop = ((self.miniB) * (self.G * self.D / self.kp_hidden_dim1) * 5
                       + (2 * self.D * self.G * self.D / self.kp_hidden_dim1)) # local accumulation of wts
          #4 refers to the number of pointwise ops (mul + add +tanh + mul) on 
          #the critical path 
          point_mem  = (self.precision * self.miniB * 
                        (self.G * self.D / self.kp_hidden_dim1) * (3 * 3 + 2 * 2)
                       + (2 * self.precision * self.D * self.G * self.D / self.kp_hidden_dim1) * 3) # local accumulation of wts
          # 3(3 memory access per operation with two input and one output)
          # 3(mul +  add + mul) on critical path
          return GEMM_time + reduction_time + point_time
   
      def getCb_kp1(self):
          #TODO:Add local accumulation of weights at every time step
          #Pointwise
          point_flop = ((self.miniB) * (self.G * self.D / self.kp_hidden_dim1) * 5
                       + (2 * self.D * self.G * self.D / self.kp_hidden_dim1)) # local accumulation of wts
          #4 refers to the number of pointwise ops (mul + add +tanh + mul) on 
          #the critical path 
          point_mem  = (self.precision * self.miniB * 
                        (self.G * self.D / self.kp_hidden_dim1) * (3 * 3 + 2 * 2)
                       + (2 * self.precision * self.D * self.G * self.D / self.kp_hidden_dim1) * 3) # local accumulation of wts
          # 3(3 memory access per operation with two input and one output)
          # 3(mul +  add + mul) on critical path
          # 2(2 memory access per operation with one input and one output)
          # 1(tanh) on critical path
    
          data_size =  4 * self.miniB * (self.G * self.D / self.kp_hidden_dim1) * self.precision
          mem_transfer = self.roofline(0,  2 * data_size, name='Cb_kp1: memory accesses before and after data transfer over network') 
          data_transfer =  data_size / self.IBK1
          point_comm = mem_transfer + data_transfer
          #3 refers to the number of pointwise ops (mul + tanh + mul) on
          # critical path whose inputs are located across different GPUs
          #NOTE:Assuming all communications can happpen in parallel
  
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise_Cb_kp1') + 5 * self.O + point_comm
  
          #GEMM_wrt_act and wt is calculated under getDistGEMM_b_kp1
          GEMM_time, reduction_time = self.getDistGEMM_b_kp1(self.miniB, 2 * self.D, self.G * self.D, self.kp_hidden_dim1, "Cb_kp1")
          
          if self.debug:
              print("(gr) Hidden point_flop: {:,}, point_mem: {:,}\n".format(int(point_flop/1e9), int(point_mem/1e9)))
  
          return GEMM_time + reduction_time + point_time
  
      #Row-Column MM
      def getCf_kp2(self):
          #Multiply
          assert(self.kp_hidden_type == 2)
          assert(self.kp_hidden_dim1 > 1 or self.kp_hidden_dim2 > 1)
          assert(self.kp_hidden_dim2 % self.G == 0 or self.kp_hidden_dim2 == 2 or self.kp_hidden_dim2 == 1)
          assert(self.miniB % self.kp_hidden_dim1 == 0)
          assert(self.G * self.D % self.kp_hidden_dim2 == 0)
  
          GEMM_time, reduction_time = self.getDistGEMM_f_kp2(self.miniB, 2 * self.D, self.G * self.D, self.kp_hidden_dim1,self.kp_hidden_dim2, "Cf_kp2")
          
          #Pointwise ops
          point_flop = (self.miniB/self.kp_hidden_dim1) * (self.G * self.D / self.kp_hidden_dim2) * 5
          #4 refers to the number of pointwise ops (mul + add +tanh + mul) on 
          #the critical path 
          point_mem  = int(self.precision * (self.miniB / self.kp_hidden_dim1) * 
                          (self.G * self.D / self.kp_hidden_dim2) *
                          (3 * 3 + 2 * 2 ))
          # 3(3 memory access per operation with two input and one output)
          # 3(mul +  add + mul) on critical path
          # 2(2 memory access per operation with one input and one output)
          # 1(tanh) on critical path
          data_size =  ((self.miniB / self.kp_hidden_dim1) * 
                         (self.G * self.D / self.kp_hidden_dim2) * 4 * self.precision)
          #4 refers to the number of pointwise ops (mul + add + tanh + mul) whose inputs
          #across different GPU
          
          point_comm = 0
          if (self.kp_softmax_dim2 > 1):
            mem_transfer = self.roofline(0,  2 * data_size, name='Cf_kp2: memory accesses before and after data transfer over network') 
            data_transfer =  data_size / self.IBK2
            point_comm = mem_transfer + data_transfer
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise_Cf_kp2') + 5 * self.O + point_comm
  
          
          return GEMM_time + reduction_time + point_time
      
      def getCb_kp2(self):
        
          #Pointwise ops
          point_flop = ((self.miniB / self.kp_hidden_dim1) * (self.G * self.D / self.kp_hidden_dim2) * 5
                       + (2 * self.D * self.G * self.D / self.kp_hidden_dim2)) # local accumulation of wts
          #4 refers to the number of pointwise ops (mul + add +tanh + mul) on 
          #the critical path 
          # kp_hidden_dim2 is for the reduction sum operation after doing outer product
          # for (B,4D)x(4D,2D).This is outerproduct due to the data distribution.
          point_mem  = int((self.precision * (self.miniB / self.kp_hidden_dim1) * 
                                         (self.G * self.D / self.kp_hidden_dim2) *
                                         (3 * 3 + 2 * 2))
                       + (2 * self.precision * self.D * self.G * self.D / self.kp_hidden_dim2) * 3) # local accumulation of wts
          # 3(3 memory access per operation with two input and one output)
          # 3(mul +  add + mul) on critical path
          # 2(2 memory access per operation with one input and one output)
          # 1(tanh) on critical path
     
          data_size =  int(self.miniB * (self.G * self.D / self.kp_hidden_dim2) * 4 * self.precision)
          #3 refers to the number of pointwise ops (mul + add +tanh + mul) on 
          #3 refers to the number of hops to gather i,f, o and c in each GPU
          #in order to perform (B,4D)x(4D,2D)
  
          point_comm = 0
          if (self.kp_softmax_dim2 > 1):
            mem_transfer = self.roofline(0,  2 * data_size, name='Cb_kp2:memory accesses before and after data transfer over network') 
            data_transfer =  data_size / self.IBK2
            point_comm = mem_transfer + data_transfer
          
          point_time = self.roofline(point_flop, point_mem, name='pointwise_Cb_kp2') + 5 * self.O + point_comm
  
  
  
          GEMM_time, reduction_time = self.getDistGEMM_b_kp2(self.miniB, 2 * self.D, self.G * self.D, self.kp_hidden_dim1,self.kp_hidden_dim2, "Cb_kp2")
        
          if self.debug:
              print("(gr) Hidden point_flop: {:,}, point_mem: {:,}\n".format(int(point_flop/1e9), int(point_mem/1e9)))
  
          return GEMM_time + reduction_time + point_time
  
  
  
      def getCf(self, m, n, k):
          #Add Biad adds
          """Get LSTM Cell Time on Forward Path"""
          GEMM_time = self.getGEMMTime(m, k, n, "Cf")
  
          
          point_flop = m * n * 5
          #1: add bias
          #5: add nonlinearities, there is one more than the number of gates (self.G)
          #1: pointwise muliply and add
          point_mem  = (self.precision * m * n *
                       (3 * 3 + 2 * 2 ))
          #3: 3 memory accesses for operands with two inputs and one output
          #2: 1 for bias add + 1 for pointwise mul
          #2: 2 memory accesses for operands with one input and one output
          #1: 5/4 non-linearities per gate
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise_Cf') + 5 * self.O
  
    
          if self.debug:
            gigaByte = 1024 * 1024 * 1024
            print("Hidden point_flop: {:,}, point_mem: {:,}\n".format(int(point_flop/1e9), int(point_mem/gigaByte)))
            print("Hidden point_time: {:,}\n".format(point_time))
  
          if self.validating_GEMM:
            return GEMM_time
          else:
            return GEMM_time[0] + point_time
  
  
     
      def getCb(self):
          """Get LSTM Cell Time on Backward Path"""
          grad_act_time,_,_ = self.getGEMMTime(self. miniB, self.G * self.D, 2 * self.D, "Cb_act") 
          grad_wt_time,_,_   = self.getGEMMTime(2 * self.D, self.miniB, self.G * self.D, "Cb_wt")
  
          
          GEMM_time = grad_act_time + grad_wt_time
  
          point_flop = ((self.miniB * self.D * 5) + 
                       (2 * self.D * self.G * self.D)) # local accumulation of wts
          point_mem  = ((self.precision * self.miniB * self.D * (3 * 3 + 2 * 2)) + 
                       (2 * self.precision * self.D * self.G * self.D) * 3) #local accumulation of wts
          point_time = self.roofline(point_flop, point_mem, name='pointwise_Cb') + 5 * self.O
     
          if self.debug:
              print("(gr) Hidden/ point_flop: {:,}, point_mem: {:,} ".format(int(point_flop/1e9), int(point_mem/1e9)))
              print("Hidden point_time: {:,}\n".format(point_time))
  
          return GEMM_time + point_time
  
  
      #Reduction and all-gather time estimation
      def getR(self, Dim0 = None, Dim1 = None, p = None, ib = None, ll = None, partial = None,
              allReduce = None, name = None):
          """Get partail or full reduction or allGather latency"""
          """Partial reduction means each gpu is only collecting a shard of 
          reduced data"""
          """allReduce= False measures allGather latency otherwise allReduce"""
          """Partial: True, All-reduce:True, half All-reduce"""
          """Partial: True, All-reduce:False, All-gather"""
          """Partial: False, All-reduce:True, All-reduce"""
          """Partial: False, All-reduce:False, All-gather"""
          if (Dim0 == None):
              #for data parallel reduction, 
              Dim0 = (2 * self.D // self.kp_hidden_dim) if (self.kp_hidden_type == 1) else (2 * self.D // self.kp_hidden_dim2 if (self.kp_hidden_type == 2) else (2 * self.D))
          if (Dim1 == None):
              Dim1 = self.G * self.D
          if (p == None):
              p = self.dp
          if (ib == None): 
              ib = self.IBD
          if (ll == None): 
              ll = self.LLD
          if (partial == None): 
              partial = False
          if (allReduce == None):
              allReduce = True
          if (p == 1):
              return 0
          #If small data transfers, just broadcast
          #NOTE: Keep threshold zero to avoid if loop
          threshold = 0
          data_tranfer = 0
          data_prep = 0
  
          #FIXME: Here I assumed point-2-point links exist across all nodes
          #Implement brodcast timing under ring topology
          if (self.precision * Dim0 * Dim1 < threshold):
              factor = (1/p if partial else 1)
              data_transfer = (((self.precision * Dim0 * Dim1)/ib + ll) * factor if p > 1 else 0)
              data_prep_comp = Dim0 * Dim1 * (p-1) * factor
              data_prep_mem = int((3 * self.precision * Dim0 * Dim1) * (p - 1) * factor)
              data_prep = self.roofline(data_prep_comp, data_prep_mem, name='R-prepTime')
          else:
              #Assuming point-2-point link between consecutive data partitions
              #In other words, the network topology assumed is Ring, 
              #therefore all (dp-1) transfers can happen in parallel,
              #To assume different toplogy data_transfer formulation should change
              #e.g. assuming bus, data_transfer formulation would change as follows:
              #data_transfer  = ((self.precision * self.D * self.D) * (self.dp /self.dp)) * 
              #                 (self.G * 2) * (2 * (self.dp - 1))) / self.IBD
              factor = (1 if partial or not allReduce else 2)
              mem_access  = self.roofline(0, int(2 * self.precision * Dim0 * Dim1 / p), name='Reduction: memory accesses before and after data transfer over network')
              data_transfer  = float("inf") if (ib == 0) else ((((self.precision * Dim0 * Dim1) / p) / ib) + mem_access + ll) * factor * (p - 1)
              #dt = ((self.precision * Dim0 * Dim1) / p) * factor * (p - 1)
              
              #First round accumlates the updates as going around the ring
              data_prep_comp = (Dim0 * Dim1) / p
              data_prep_mem  = int(3 * self.precision * Dim0 * Dim1 / p)
              data_prep = ((self.roofline(data_prep_comp, data_prep_mem, name='R-prepTime') + self.O) * (p - 1))
              
              #all-gather-concat
  
              data_concat_mem = 3 * Dim0 * Dim1 * self.precision
              concat_time = ((self.roofline(0, data_concat_mem, name='all-gather-concat') + self.O))
  
              
              #print("R1: {}, factor: {}\n".format(dt,factor))
          if self.debug:
              print("Bandwidth: {:,} GB/s".format(ib/(1024*1024*1024)))
              print("data_transfer_time: {:,}, data_prep_time: {:,}, concat_time: {:,}".format(data_transfer, (data_prep if allReduce else 0), (concat_time if not allReduce else 0)))
              print("(data_prep) allReduce_flop: {:,}, allReduce_mem: {:,}".format(int(data_prep_comp), int(data_prep_mem)))
              print("(data_transfer) {:,}".format(int(self.precision * Dim0 * Dim1 / (p))))
  
          return data_transfer + (data_prep if allReduce else 0) + (concat_time if not allReduce else 0)
      
      def gradClipping(self, Dim0 = None, Dim1 = None, name = None):
          if (Dim0 == None):
              Dim0 = 2 * self.D
          if (Dim1 == None):
              Dim1 = self.G * self.D
          if (name == None):
              name = "Hidden"
           #t_list[i] * clip_norm / max(global_norm, clip_norm)
           #where:
           #global_norm = sqrt(sum([l2norm(t)**2 for t in t_list]))
            
          norm_comp = Dim0 * Dim1 * 2
          #1: power 2
          #1: summ
          norm_mem = (Dim0 * Dim1 * 1) * self.precision
          #1: one read per element and power it by 2 in local registers  anfd 
          #summing to local acc
  
        
          clip_comp = Dim0 * Dim1 * 2
          #1: pointwise mul
          #1: pointwise div
  
          clip_mem = (Dim0 * Dim1 * 2) * self.precision
          #1: one read for pointwise mul
          #1: one write for pointwise div
  
          gradclip_mem = norm_mem + clip_mem
          gradclip_comp = norm_comp + clip_comp
  
          gradclip_time = self.roofline(gradclip_comp, gradclip_mem, name='pointwise-grad-clipping')
  
          if self.debug:
              print("({}) gradclip_flop: {:,}, gradclip_mem: {:,}".format(name, gradclip_comp, gradclip_mem))
              print("({}) gradclip_time: {:,}\n".format(name, gradclip_time))
  
          return gradclip_time
  
  
      def applyGrad(self, Dim0 = None, Dim1 = None, name = None):
          if (Dim0 == None):
              Dim0 = 2 * self.D
          if (Dim1 == None):
              Dim1 = self.G * self.D
          if (name == None):
              name = "Hidden"
          
          applyGrad_comp = Dim0 * Dim1 * 3
          #3: one pointwise division by  scalar after reducing all the gradients, 
          #   one final addition of gradients to the weights
          #   one multiply by learning rate
          applyGrad_mem = ((1 * Dim0 * Dim1 * self.precision) +
                           (2 * Dim0 * Dim1 * self.precision) +
                           (1 * Dim0 * Dim1 * self.precision))
          #1: read for pointiwse div
          #2: 1 reads and one write for pointwise add
          #1: one write for multiplication by lr
          applyGrad_time = self.roofline(applyGrad_comp, applyGrad_mem, name='pointwise-applyGrad')
         
          clip_time = self.gradClipping(Dim0, Dim1, name)
          
          grad_time = applyGrad_time + clip_time
  
          if self.debug:
              print("({}) applyGrad_flop: {:,}, applyGrad_mem: {:,}".format(name, applyGrad_comp, applyGrad_mem))
              print("({}) applyGrad_time: {:,}\n".format(name, applyGrad_time))
         
          return grad_time
  
      def getDistGEMM_f_kp1(self, m, k, n, dim1, name):
          GEMM_time = self.getGEMMTime(m , k//dim1, n , name)
          
          #Sum-Reduce within each row for use in the next time step
          reduction_time = self.getR(Dim0 = m, 
                                     Dim1 = n,
                                     p = dim1,
                                     ib = self.IBK1,
                                     ll = self.LLK1,
                                     partial = True,
                                     allReduce = True,
                                     name = name)
          if self.validating_GEMM:
            print("GEMM_time: {}, Reduction_time:{}".format(GEMM_time[0], reduction_time))
            return GEMM_time[0] + reduction_time, GEMM_time[1], GEMM_time[2]
          else:
            return GEMM_time[0], reduction_time
  
  
      def getDistGEMM_b_kp1(self, m, k, n, dim1, name):
          #calculate grad wrt. act (A'. W^T)
          #gather whole(A') before MM
          #A' is distibuted as columns across different nodes
          reduction_time = self.getR(Dim0 = m, 
                                     Dim1 = n,
                                     p = dim1,
                                     ib = self.IBK1,
                                     ll = self.LLK1,
                                     partial = False,
                                     allReduce = False,
                                     name = name)
  
          #Multiply full grad_activation with shards of weights
          grad_wt_time,_,_  = self.getGEMMTime(k, (m // dim1), n, name + "wt")
          #Multiply full grad-activation with shards of activations 
          grad_act_time,_,_ = self.getGEMMTime(m, (n // dim1), k, name + "act")
  
  
          GEMM_time = grad_wt_time + grad_act_time
  
          return GEMM_time, reduction_time
  
      def getDistGEMM_f_kp2(self, m, k, n, dim1, dim2, name):
          GEMM_time = self.getGEMMTime(m // dim1, k, n // dim2, name)
          reduction_time = self.getR(Dim0 = m // dim1, 
                                     Dim1 = n,
                                     p = dim2,
                                     ib = self.IBK2,
                                     ll = self.LLK2,
                                     partial = False,
                                     allReduce = False,
                                     name = name)
          if self.validating_GEMM:
            print("GEMM_time: {}, Reduction_time:{}".format(GEMM_time[0], reduction_time))
            return GEMM_time[0] + reduction_time, GEMM_time[1], GEMM_time[2]
          else:
            return GEMM_time[0], reduction_time
  
      def getDistGEMM_b_kp2(self, m, k, n, dim1, dim2, name):
          ######################################################################################
          #calculate grad wrt. weights (A^T. grad(A')) 
          #gather row(A^T)
          reduction_time_wt1 = self.getR(Dim0 = k, 
                                         Dim1 = m,
                                         p = dim1,
                                         ib = self.IBK1,
                                         ll = self.LLK1,
                                         partial = False,
                                         allReduce = False,
                                         name = name)/2
          #To calculate grad wrt weights (A^T, grad(A')), 
          #gather column grad(A')
          reduction_time_wt2 = self.getR(Dim0 = m, 
                                         Dim1 = n / dim2,
                                         p = dim1,
                                         ib = self.IBK1,
                                         ll = self.LLK1,
                                         partial = False,
                                         allReduce = False,
                                         name = name)
  
          ########################################################################################
          #calculate grad wrt. act (grad(A'). w^T)
          #gather row grad(A')
          reduction_time_act1 = self.getR(Dim0 = m / dim1, 
                                          Dim1 = n,
                                          p = dim2,
                                          ib = self.IBK2,
                                          ll = self.LLK2,
                                          partial = False,
                                          allReduce = False,
                                          name = name)
          #calculate grad wrt. act (grad(A'). w^T)
          #gather col(w^T)
          reduction_time_act2 = self.getR(Dim0 = k, 
                                         Dim1 = n,
                                         p = dim2,
                                         ib = self.IBK2,
                                         ll = self.LLK2,
                                         partial = False,
                                         allReduce = False,
                                         name = name)/2
          
          reduction_time = reduction_time_wt1 + reduction_time_wt2 + reduction_time_act1 +reduction_time_act2
  
          #Multiply full grad_activation with shards of weights
          grad_wt_time,_,_  = self.getGEMMTime(k / dim1, m, n / dim2, name + "wt")
          #Multiply full grad-activation with shards of activations 
          grad_act_time,_,_ = self.getGEMMTime(m / dim1, n, k / dim2, name + "act")
  
          GEMM_time = grad_wt_time + grad_act_time
  
          return GEMM_time, reduction_time
  
      def getDataParallelReduction(self, k, n, dim1, dim2, name):
          #k = 2 * self.D
          #n = 4 * self.D
          #dim1 = self.kp_hidden_dim1
          #dim2 = self.kp_hidden_dim2
  
          reduction_time_wt_kp = 0
          reduction_time_wt_dp = 0
          apply_grad_time = 0
  
          if self.kp_hidden_type == 1: #CR
              reduction_time_wt_kp = 0
              reduction_time_wt_dp = self.getR(Dim0 = k/dim1, 
                                               Dim1 = n,
                                               p = self.dp,
                                               ib = self.IBD,
                                               ll = self.LLD,
                                               partial = False,
                                               allReduce = True,
                                               name = name)
              apply_grad_time = self.applyGrad(Dim0 = k/dim1, Dim1 = n, name = name)
          
          elif self.kp_hidden_type == 2: #RC
              reduction_time_wt_dp = self.getR(Dim0 = k/dim1, 
                                               Dim1 = n/dim2,
                                               p = self.dp,
                                               ib = self.IBD,
                                               ll = self.LLD,
                                               partial = False,
                                               allReduce = True,
                                               name = name)
  
              #gather col(w)
              reduction_time_wt_kp = self.getR(Dim0 = k, 
                                               Dim1 = n/dim2,
                                               p = dim1,
                                               ib = self.IBK1,
                                               ll = self.LLK1,
                                               partial = False,
                                               allReduce = False,
                                               name = name)
              apply_grad_time = self.applyGrad(Dim0 = k, Dim1 = n/dim2, name = name)
          else:
              reduction_time_wt_kp = 0
              reduction_time_wt_dp = self.getR(Dim0 = k, 
                                               Dim1 = n,
                                               p = self.dp,
                                               ib = self.IBD,
                                               ll = self.LLD,
                                               partial = False,
                                               allReduce = True,
                                               name = name)
              apply_grad_time = self.applyGrad(Dim0 = k, Dim1 = n, name = name)
  
          reduction_time =  reduction_time_wt_kp + reduction_time_wt_dp + apply_grad_time 
          return reduction_time
      
      
      def getProjection_f(self):
          GEMM_time,_,_ = self.getGEMMTime(self.miniB, self.D, self.projection, "projection")
          return GEMM_time
  
      def getProjection_b(self):
          grad_wt_time,_,_  = self.getGEMMTime(self.projection, self.miniB, self.D, "projection_b_wt")
          grad_act_time,_,_ = self.getGEMMTime(self.miniB, self.projection, self.D, "projection_b_act")
  
          GEMM_time = grad_wt_time + grad_act_time
          return GEMM_time
  
      def getProjection_f_kp1(self):
          assert(self.kp_projection_type == 1)
          assert(self.kp_projection_dim1 > 1)
          assert(self.D % self.kp_projection_dim1 == 0)
  
          GEMM_time, reduction_time = self.getDistGEMM_f_kp1(self.miniB, self.D, self.projection, self.kp_projection_dim1, "projection_f")
          return GEMM_time + reduction_time
  
      def getProjection_b_kp1(self):
          assert(self.kp_projection_type == 1)
          assert(self.kp_projection_dim1 > 1)
          assert(self.D % self.kp_projection_dim1 == 0)
          
          GEMM_time, reduction_time = self.getDistGEMM_b_kp1(self.miniB, self.D, self.projection, self.kp_projection_dim1, "projection_b")
          return GEMM_time + reduction_time
  
      def getProjection_f_kp2(self):
          assert(self.kp_projection_type == 2)
          assert(self.kp_projection_dim1 > 1 or self.kp_projection_dim2 > 1)
          assert((self.miniB) % self.kp_projection_dim1 == 0)
          assert(self.projection % self.kp_projection_dim2 == 0)
  
          GEMM_time, reduction_time = self.getDistGEMM_f_kp2(self.miniB, self.D, self.projection, self.kp_projection_dim1, self.kp_projection_dim2, "projection_f")
          return GEMM_time + reduction_time
  
      def getProjection_b_kp2(self):
          assert(self.kp_projection_type == 2)
          assert(self.kp_projection_dim1 > 1 or self.kp_projection_dim2 > 1)
          assert((self.miniB) % self.kp_projection_dim1 == 0)
          assert(self.projection % self.kp_projection_dim2 == 0)
          
          GEMM_time, reduction_time = self.getDistGEMM_f_kp2(self.miniB, self.D, self.projection, self.kp_projection_dim1, self.kp_projection_dim2, "projection_b")
          return GEMM_time + reduction_time
  
      def getSoftmax_f(self):
          GEMM_time,_,_ = self.getGEMMTime(self.miniB, (self.projection if proj else self.D), self.V, "softmax_f")
  
          #Final matrix after GEMM has (B, V) dimensionality
          #We get exponential on each of the elements in a row
          #and then normalize them across the row
          #Therefore for each row we do V sigmoids and V-1 additions and V divisions
          #For each row, we read V sigmoids and write one
          #Up to here is 3 operations
          point_flop = (self.miniB * (3 * self.V - 1))
          
          point_mem  = self.precision * self.miniB * (7 * self.V)
          #2: one read and one write for sigmoid
          #1: one read for reduction
          #1: one write for extension
          #2: for pointwise division
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise-softmax-f') + 4 * self.O 
     
          if self.debug:
              print("Softmax point_flop: {:,}, point_mem: {:,}".format(int(point_flop/1e9), int(point_mem/1e9)))
              print("point_time: {:,}\n".format( point_time))
  
          return GEMM_time + point_time
  
      #FIXME: where is the reduction time?
      def getSoftmax_b(self):
          grad_wt_time,_,_  = self.getGEMMTime((self.projection if proj else self.D), self.miniB, self.V, "softmax_b_wt")
          grad_act_time,_,_ = self.getGEMMTime(self.miniB, self.V, (self.projection if proj else self.D), "softmax_b_act")
  
          GEMM_time = grad_wt_time + grad_act_time
          point_flop = self.miniB * self.V * 5
          #1: one for one of the divisions, grad(A) (y=A/B)
          #2: one for division and multiplication, grad(B)
          #1: one for addition, copies turn into add
          #1: one for sigmoid
  
          point_mem  = self.precision * self.miniB * self.V * 11
          #3: grad(A) in pointwise division
          #3: grad(B) in pointwise division
          #3: addition in copy backprop
          #2: sigmoid
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise-softmax-b') + 4 * self.O 
      
          if self.debug:
              print("(gr) Softmax point_flop: {:,}, point_mem: {:,}".format(int(point_flop/1e9), int(point_mem/1e9)))
              print("(gr) Softmax point_time: {:,}\n".format(point_time))
  
  
          return GEMM_time + point_time
  
      #Column-Row MM
      def getSoftmax_f_kp1(self):
          #Multiply
          assert(self.kp_softmax_type == 1)
          assert(self.kp_softmax_dim1 > 1)
          assert((self.projection if proj else self.D) % self.kp_softmax_dim1 == 0)
  
          GEMM_time, reduction_time = self.getDistGEMM_f_kp1(self.miniB, self.projection if proj else self.D, self.V, self.kp_softmax_dim1, "softmax_f_kp1")
  
          #Final matrix after GEMM has (B, V) dimensionality
          #We get exponential on each of the elements in a row
          #and then reduce all elements in the row to one
          #Therefore for each row we do V sigmoids and V-1 additions and V divisions
          #For each row, we read V sigmoids and write one
          #After GEMM reduction, each matrix has the full (B,V) 
          #but each needs to only operate on 1/dim1 rows to get the reduction
          point_flop = (self.miniB / self.kp_softmax_dim1) * self.V * 3
          point_mem  = self.precision * (self.miniB / self.kp_softmax_dim1) * self.V * 7
          #2: sigmoid
          #1: one read for reduction, the accumulate is a register
          #1: one for write/extend the reduction result into all cells 
          #3: division needs one for read and one for write. 
  
          point_comm = self.getR(Dim0 = self.miniB, 
                                 Dim1 = 1,
                                 p = self.kp_softmax_dim1,
                                 ib = self.IBK1,
                                 ll = self.LLK1,
                                 partial = False,
                                 allReduce = False,
                                 name="getSoftmax_f_kp1")
          #communicating partail sum per row from one GPU to all others to perform sum reduce
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise-softmax-f-kp1') + self.O + point_comm
     
          if self.debug:
              print("Softmax point_flop: {:,}, point_mem: {:,}".format(int(point_flop/1e9), int(point_mem/1e9)))
              print("Softmax GEMM_time: {:,}, point_time: {:,}\n".format(GEMM_time, point_time))
  
          return GEMM_time + reduction_time + point_time
  
      def getSoftmax_b_kp1(self):
  
          point_flop = (self.miniB) * (self.V / self.kp_softmax_dim1) * 5
          #1: one for one of the divisions, grad(A) (y=A/B)
          #2: one for division and multiplication, grad(B)
          #1: one for addition, copies turn into add
          #1: one for sigmoid
  
          point_mem  = self.precision * (self.miniB) * ((11 * self.V )/ self.kp_softmax_dim1)
          #3: grad(A) in pointwise division
          #3: grad(B) in pointwise division
          #3: addition in copy backprop
          #2: sigmoid
  
          point_comm = 0
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise-softmax-b-kp1') + self.O + point_comm
  
          GEMM_time, reduction_time = self.getDistGEMM_b_kp1(self.miniB, self.projection if proj else self.D, self.V, self.kp_softmax_dim1, "softmax_b_kp1")
  
          if self.debug:
              print("(gr) Softmax point_flop: {:,}, point_mem: {:,}\n".format(int(point_flop/1e9), int(point_mem/1e9)))
          
          return reduction_time + GEMM_time + point_time
  
      #Row-Column MM
      def getSoftmax_f_kp2(self):
          #Multiply
          assert(self.kp_softmax_type == 2)
          assert(self.kp_softmax_dim1 > 1 or self.kp_softmax_dim2 > 1)
          assert((self.miniB) % self.kp_softmax_dim1 == 0)
          assert((self.projection if proj else self.D) % self.kp_softmax_dim2 == 0)
  
          GEMM_time, reduction_time = self.getDistGEMM_f_kp2(self.miniB, self.projection if proj else self.D, self.V, self.kp_softmax_dim1, self.kp_softmax_dim2, "softmax_f_kp2")
          
  
          #Final matrix after GEMM has (B X S, V) dimensionality
          #We get exponential on each of the elements in a row
          #and then reduce all elements in the row to one
          #Therefore for each row we do V sigmoids and V-1 additions and V divisions
          #For each row, we read V sigmoids and write one
          
          point_flop = (self.miniB / self.kp_softmax_dim1) * (self.V / self.kp_softmax_dim2) * 3
          point_mem  = self.precision * (self.miniB / self.kp_softmax_dim1) * (self.V / self.kp_softmax_dim2) * 7
          #2: sigmoid
          #1: one read for reduction, the accumulate is a register
          #1: one for write/broadcast the reduction result into all cells 
          #3: division needs one for read and one for write. 
  
          data_size = self.precision * (self.miniB / self.kp_softmax_dim1) * (self.kp_softmax_dim2)
         
          point_comm = 0
          if (self.kp_softmax_dim2 > 1):
            mem_transfer = self.roofline(0,  2 * data_size, name='memory accesses before and after data transfer over network') 
            data_transfer =  data_size / self.IBK2
            point_comm = mem_transfer + data_transfer
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise-Softmax_f_kp2') + self.O + point_comm
    
  
          if self.debug:
              print("Softmax point_flop: {:,}, point_mem: {:,}".format(int(point_flop/1e9), int(point_mem/1e9)))
              print("Softmax GEMM_time: {:,}, point_time: {:,}\n".format(GEMM_time, point_time))
  
          return GEMM_time + point_time + reduction_time
  
      def getSoftmax_b_kp2(self):
  
          point_flop = (self.miniB / self.kp_softmax_dim1) * (self.V / self.kp_softmax_dim2) * 5
          #1: one for one of the divisions, grad(A) (y=A/B)
          #2: one for division and multiplication, grad(B)
          #1: one for addition, copies turn into add
          #1: one for sigmoid
  
          point_mem  = self.precision * (self.miniB / self.kp_softmax_dim1) * ((11 * self.V )/ self.kp_softmax_dim2)
          #3: grad(A) in pointwise division
          #3: grad(B) in pointwise division
          #3: addition in copy backprop
          #2: sigmoid
  
          point_comm = 0
  
          point_time = self.roofline(point_flop, point_mem, name='pointwise-Softmax_b_kp2') + self.O + point_comm
     
  
          GEMM_time, reduction_time = self.getDistGEMM_b_kp2(self.miniB, self.projection if proj else self.D, self.V, self.kp_softmax_dim1, self.kp_softmax_dim2, "softmax_b_kp2")
      
          if self.debug:
              print("(gr) Softmax point_flop: {:,}, point_mem: {:,}\n".format(int(point_flop/1e9), int(point_mem/1e9)))
          
          return reduction_time + GEMM_time + point_time
  
  
      def getEmbedding_f(self):
          embedding_mem = 2 * (self.miniB * self.D * self.precision)
          #embedding_time = (embedding_mem)/ (self.mem_bw) + self.mem_latency + self.O
          embedding_time = self.roofline(0, embedding_mem, name='embedding_f') + self.O
          embedding_transfer_time = 2 * self.miniB * self.D * self.precision / self.H2Dbw
          if self.debug:
              print("Embedding_mem: {:,}".format(int(embedding_mem/1e9)))
          return embedding_time + embedding_transfer_time
  
  
      def getEmbedding_b(self):
          #p2p_data_transfer = (self.precision * self.miniB * self.D)
          #data_transfer_time  = 0 if (self.dp == 1) else (float("inf") if (self.IBD == 0) else (((p2p_data_transfer) / self.IBD + self.LLD) * 2 * (self.dp -1 ))) 
          
          embedding_mem = 2 * self.miniB * self.D * self.precision
          #embedding_mem_time = (embedding_mem / self.mem_bw) + self.mem_latency + self.O
          embedding_mem_time = self.roofline(0, embedding_mem, name='embedding_b') + self.O
  
          if self.debug:
              print("(gr) Embedding_mem: {:,}".format(int(embedding_mem/1e9)))
          #return data_transfer_time + embedding_mem_time
          return embedding_mem_time
  
      def getEmbedding_f_kp1(self):
          #Each GPU has only a portion of the activations since each GPU had only a row of the weights
          reduction_time_act= self.getR(Dim0 = self.miniB, 
                                        Dim1 = self.D,
                                        p = self.kp_embedding_dim1,
                                        ib = self.IBK1,
                                        ll = self.LLK1,
                                        partial = False,
                                        allReduce = False,
                                        name="getEmbedding_f_kp1")
          embedding_mem = 2 * (self.miniB * self.D * self.precision)
          #embedding_time = (embedding_mem)/ (self.mem_bw) + self.mem_latency + self.O
          embedding_time = self.roofline(0, embedding_mem, name='embedding_f') + self.O
          if self.debug:
              print("Embedding_mem: {:,}".format(int(embedding_mem/1e9)))
          return embedding_time + reduction_time_act
  
  
      def getEmbedding_b_kp1(self):
          #Activations from previous row arrive in column fasion, they need to be gathered
          #before applying them to the local portion of the embeddings
          reduction_time_act= self.getR(Dim0 = self.miniB, 
                                        Dim1 = self.D,
                                        p = self.kp_embedding_dim1,
                                        ib = self.IBK1,
                                        ll = self.LLK1,
                                        partial = False,
                                        allReduce = False,
                                        name="getEmbedding_b_kp1")
          #Each GPU would read through the entire actication and write as many at most as many of B rows
          embedding_mem = 2 * self.miniB * self.D * self.precision
          embedding_mem_time = self.roofline(0, embedding_mem, name='embedding_b') + self.O
  
          if self.debug:
              print("(gr) Embedding_mem: {:,}".format(int(embedding_mem/1e9)))
          return embedding_mem_time + reduction_time_act
  
      def getEmbedding_f_kp2(self):
          embedding_mem = 2 * ((self.miniB / self.kp_embedding_dim1) * (self.D / self.kp_embedding_dim2) * self.precision)
          embedding_time = self.roofline(0, embedding_mem, name='embedding_f') + self.O
          if self.debug:
              print("Embedding_mem: {:,}".format(int(embedding_mem/1e9)))
          return embedding_time
  
  
      def getEmbedding_b_kp2(self):
          #Every GPU will update a little tile of the embedding
          #need to be gathered after the update across the rows of each column
          reduction_time_act= self.getR(Dim0 = self.miniB, 
                                        Dim1 = self.D / self.kp_embedding_dim2,
                                        p = self.kp_embedding_dim1,
                                        ib = self.IBK1,
                                        ll = self.LLK1,
                                        partial = False,
                                        allReduce = False,
                                        name="getEmbedding_b_kp2")
          
          embedding_mem = 2 * (self.miniB / self.kp_embedding_dim1) * (self.D / self.kp_embedding_dim2) * self.precision
          embedding_mem_time = self.roofline(0, embedding_mem, name='embedding_b') + self.O
  
          if self.debug:
              print("(gr) Embedding_mem: {:,}".format(int(embedding_mem/1e9)))
          return embedding_mem_time + reduction_time_act
  
      def getInterLayerCommLatency(self, dim1, dim2):
          w = 0
          if self.lp > 1:
            w_size         = self.precision * dim1 * dim2
            transfer_time  = w_size / self.IBL + self.LLL
            mem_time       = self.roofline(0, 2 * w_size, name='inter_layer')
            #2: read from memory of previous layer and write to the memory of the next layer
            w              = mem_time + transfer_time
          return w
  
  
      def dprint(self, string):
          if self.debug:
              print(string)
      
      def readjust_type(self):
        if self.kp_hidden_dim1 == 1 and self.kp_hidden_dim2 == 1:
          self.kp_hidden_type = -1
        
        if self.kp_softmax_dim1 == 1 and self.kp_softmax_dim2 == 1:
          self.kp_softmax_type = -1
  
        if self.kp_embedding_dim1 == 1 and self.kp_embedding_dim2 == 1:
          self.kp_embedding_type = -1
  
      def calcTime(self):
          B = self.miniB
          D = self.D
          G = self.G
          L = self.L
          S = self.S
          V = self.V
          lp = self.lp
          dp = self.dp
  
          self.readjust_type()
  
          if self.kp_hidden_type == -1:
              Cf = self.getCf(m = B, k = 2*D,  n = G*D)
              Cb = self.getCb()
              Tf = self.getInterLayerCommLatency(B, D) 
          elif self.kp_hidden_type == 1: #CR
              Cf = self.getCf_kp1()
              Cb = self.getCb_kp1()
              Tf = self.getInterLayerCommLatency(B, D / self.kp_hidden_dim1) 
          elif self.kp_hidden_type == 2: #RC
              Cf = self.getCf_kp2()
              Cb = self.getCb_kp2()
              Tf = self.getInterLayerCommLatency(B / self.kp_hidden_dim1, D / self.kp_hidden_dim2) 
          else:
              print("Incorrect distributed GEMM type, 1: Column-Row, 2: Row-Column")
              sys.exit()
          
          if self.lp == 1:
              Tf = 0
              
          Tb = Tf
         
          if self.kp_softmax_type == -1:
              Sf = self.getSoftmax_f()
              Sb = self.getSoftmax_b()
          elif self.kp_softmax_type == 1: #CR
              Sf = self.getSoftmax_f_kp1()
              Sb = self.getSoftmax_b_kp1()
          elif self.kp_softmax_type== 2: #RC
              Sf = self.getSoftmax_f_kp2()
              Sb = self.getSoftmax_b_kp2()
          else:
              print("Incorrect distributed GEMM type, 1: Column-Row, 2: Row-Column")
              sys.exit()
          
          if self.kp_embedding_type == -1:
              Ef = self.getEmbedding_f()
              Eb = self.getEmbedding_b()
          elif self.kp_embedding_type == 1: #CR
              Ef = self.getEmbedding_f_kp1()
              Eb = self.getEmbedding_b_kp1()
          elif self.kp_embedding_type== 2: #RC
              Ef = self.getEmbedding_f_kp2()
              Eb = self.getEmbedding_b_kp2()
          else:
              print("Incorrect distributed GEMM type, 1: Column-Row, 2: Row-Column")
              sys.exit()
         
  
          Rc = self.getDataParallelReduction(k = 2 * D, n = G * D, dim1 = self.kp_hidden_dim1, dim2 = self.kp_hidden_dim2, name = "Hidden Reduction")
          Rs = self.getDataParallelReduction(k = D, n = V, dim1 = self.kp_softmax_dim1, dim2 = self.kp_softmax_dim2, name = "Softmax Reduction")
          Re = self.getDataParallelReduction(k = V, n = D, dim1 = self.kp_embedding_dim1, dim2 = self.kp_embedding_dim2, name = "Embedding Reduction")
  
                          
          if self.debug:
              print("dp: {}, lp: {}, kp_hidden_dim1: {}, kp_hidden_dim2: {}, kp_softmax_dim1: {}, kp_softmax_dim2: {}, kp_embedding_dim1: {}, kp_embedding_dim2: {},  kp_hidden_type: {}, kp_softmax_type: {}, kp_embedding_type: {}\n".
                      format(dp, 
                            lp, 
                            self.kp_hidden_dim1, 
                            self.kp_hidden_dim2, 
                            self.kp_softmax_dim1, 
                            self.kp_softmax_dim2, 
                            self.kp_embedding_dim1, 
                            self.kp_embedding_dim2, 
                            self.kp_hidden_type, 
                            self.kp_softmax_type, 
                            self.kp_embedding_type))
  
            
              print("Cf: {} Cb: {} softmax_f: {} softmax_b: {} embedding_f: {} embedding_b: {} " 
                     "Rs: {} Rc: {} Re: {}\n".format(Cf, 
                                                     Cb,
                                                     Sf, 
                                                     Sb, 
                                                     Ef,
                                                     Eb,
                                                     Rs,
                                                     Rc,
                                                     Re))
  
  
          g = Graph(num_seq = S, num_layer = L, lp = lp, Ef = Ef, Cf = Cf, Sf = Sf, Tf = Tf, Eb = Eb, Cb = Cb, Sb = Sb, Tb = Tb, Re = Re, Rc = Rc, Rs = Rs)
              
          fw_roots = g.construct_fwd_graph()
          bw_roots = g.construct_bwd_graph()
  
          time_fw = g.simulate(fw_roots[0], 0)
          time_bw = g.simulate(bw_roots[g.num_seq - 1], g.lp - 1)
         
          self.tot_time  = time_fw + time_bw
          tot_param = self.tot_param()
          
          return self.tot_time, tot_param
  
      def getTime(self):
          return self.tot_time
  def callPerf(exp_config, exp_dir, debug):
      exp_path = os.path.expandvars(os.path.expanduser(exp_config))
      exp_config = config.parse_config(exp_path)
  
      #try:
      #    #print("Removing directory:" + exp_dir)
      #    shutil.rmtree(exp_dir)
      #except:
      #    pass
      #os.makedirs(exp_dir)
  
      TC = TimeCalculation(exp_config)
      TC.debug = debug
      tot_time, tot_param = TC.calcTime()
  
      output_file = exp_dir + "/summary.txt"
      
      TC.printSysConfig(exp_config, output_file)
  
      with open(output_file, "a+") as f:
          f.write("Time: {0:.8f}\n".format(tot_time))
          f.write("Params (Billion): {0:.8f}\n".format(tot_param/1e9))
  
  @click.command("standalone")
  @click.option("--args_input", help="Shall it read the args from the input command (True) or from exp_config (False)", default=False, type=bool, required=False)
  @click.option("--exp_config", help="Path to experiment config", required=True)
  @click.option("--exp_dir", help="Checkpoint/log directory", required=True)
  @click.option("--debug", help="debug", default=False, type=bool)
  @click.option("--m", help="input dimension", default=32768, type=int, required=False) #only use for GEMM validation. This allows arbitrary choice of dimension. For LSTM, dimensions are fixed at m=mini_batch, k=2*D and n=4*D.
  @click.option("--n", help="output dimension", default=32768, type=int, required=False) #only use for GEMM validation
  @click.option("--k", help="input dimension", default=32768, type=int, required=False) #only use for GEMM validation
  @click.option("--t", help="parallelism strategy (RC or CR)", default='None', type=str, required=False) #only use for GEMM validation
  @click.option("--kp1", help="RC:parallelism along input dimension, CR: parallelism along inner dimension", default=None, type=int, required=False) #only use for GEMM validation
  @click.option("--kp2", help="RC:parallelism along output dimension", default=None, type=int, required=False) #only use for GEMM validation
  @click.option("--gemm", help="report ONLY GEMM time", default=False, type=bool, required=False) #only use for GEMM validation
  @click.option("--batch_size", help="Total Batch Size", default=2048, type=int, required=False)
  @click.option("--hidden_dim", help="Hidden Dimension per LSTM layer", default=19968, type=int, required=False)
  @click.option("--seq_len", help="Number of times to unroll LSTM", default=20, type=int, required=False)
  @click.option("--vocab_size", help="Vocabulary Size", default=800000, type=int, required=False)
  @click.option("--num_layer", help="number of lstm layers", default=2, type=int, required=False)
  @click.option("--dp", help="data parallelism", default=None, type=int, required=False) #only use for GEMM validation
  @click.option("--lp", help="layer parallelism", default=None, type=int, required=False) #only use for GEMM validation
  
  def main(exp_config, exp_dir, debug, m, n, k, t, kp1, kp2, gemm, batch_size, hidden_dim, seq_len, vocab_size, num_layer, dp, lp, args_input=False):
      exp_path = os.path.expandvars(os.path.expanduser(exp_config))
      exp_config = config.parse_config(exp_path)
      output_file = exp_dir + "/summary.txt" ##Output dir should be created manually
  
  
      TC = TimeCalculation(exp_config)
      if args_input:
          TC.updateParams(debug, m, n, k, t, kp1, kp2, dp, lp, gemm, 
                      batch_size, hidden_dim, seq_len, vocab_size, num_layer)
  
      #Report GEMM time on fw path
      if TC.validating_GEMM:
          
          if kp1 == 1 and kp2 ==1: #no parallelism
            gemm_time = TC.getCf(m, k, n)
          elif t == 'CR':
            gemm_time = TC.getDistGEMM_f_kp1(m, k, n, kp1, "Cf_CR")
          elif t == 'RC':
            gemm_time = TC.getDistGEMM_f_kp2(m, k, n, kp1, kp2, "Cf_RC")
          else:
            print("Incorrect parallelism type, CR: Column-Row, RC: Row-Column")
            sys.exit()
            
          with open(output_file, "w") as f:
            f.write("Best Order: {}\n".format(gemm_time[1]))
            f.write("Best Tile: {}\n".format(gemm_time[2]))
            f.write("Time: {}\n".format(gemm_time[0]))
          return
  
      tot_time, tot_param = TC.calcTime()
      TC.printSysConfig(exp_config, output_file)
      
      with open(output_file, "a+") as f:
          f.write("\n\n==============================================\n")
          f.write("Performance Results\n")
          f.write("==============================================\n")
          f.write("Time: {0:.8f}\n".format(tot_time))
          f.write("Params (Billion): {0:.8f}\n".format(tot_param/1e9))
  
     
  if __name__ == "__main__":
      main()
  