############################ begin_copyright_notice ############################
### Copyright (C) 2018-2025 Intel Corporation
### 
### SPDX-License-Identifier: MIT
############################ end_copyright_notice ##############################

#!/usr/bin/python

import sys, re, tempfile, os, os.path, shutil, time, struct
import array

from array import *
from argparse import ArgumentParser

def Usage(progName):
    return '\nUsage: %s --input_dir c:\profile_dir\Session_Final\PS_asm5782fb10f48402ef_simd8\draw__1359' %(progName)

def ParseArgs():
    progName = sys.argv[0]

    if (len(sys.argv) != 2):
        Error(Usage(progName))

    profilerDir = sys.argv[1]

    return (profilerDir)


def OperandWidthString(elementSize):
    return {
         1  : "  B",
         2  : "  W",
         4  : " DW",
         8  : " QW",
         16 : " OW",
         32 : " HW",
         64 : "DHW",
    }.get(elementSize, "UKN")

class SEND:
    offset            = 0
    isWrite           = 0
    isBlock2D         = 0
    isScatter         = 0
    isBTS             = 0
    isSLM             = 0
    isScratch         = 0
    isAtomic          = 0
    isFence           = 0
    addressWidth      = 0
    simdWidth         = 0
    bti               = 0
    addrPayloadLength = 0
    dataPort          = 0
    isEOT             = 0
    isMedia           = 0
    elementSize       = 0
    numElements       = 0
    execSize          = 0
    channelOffset     = 0

    def read_from_file(self, fIn):
        self.offset            = int(struct.unpack('I', fIn.read(4))[0])
        self.isWrite           = int(struct.unpack('I', fIn.read(4))[0])
        self.isBlock2D         = int(struct.unpack('I', fIn.read(4))[0])
        self.isScatter         = int(struct.unpack('I', fIn.read(4))[0])
        self.isBTS             = int(struct.unpack('I', fIn.read(4))[0])
        self.isSLM             = int(struct.unpack('I', fIn.read(4))[0])
        self.isScratch         = int(struct.unpack('I', fIn.read(4))[0])
        self.isAtomic          = int(struct.unpack('I', fIn.read(4))[0])
        self.isFence           = int(struct.unpack('I', fIn.read(4))[0])
        self.addressWidth      = int(struct.unpack('I', fIn.read(4))[0])
        self.simdWidth         = int(struct.unpack('I', fIn.read(4))[0])
        self.bti               = int(struct.unpack('I', fIn.read(4))[0])
        self.addrPayloadLength = int(struct.unpack('I', fIn.read(4))[0])
        self.dataPort          = int(struct.unpack('I', fIn.read(4))[0])
        self.isEOT             = int(struct.unpack('I', fIn.read(4))[0])
        self.isMedia           = int(struct.unpack('I', fIn.read(4))[0])
        self.elementSize       = int(struct.unpack('I', fIn.read(4))[0])
        self.numElements       = int(struct.unpack('I', fIn.read(4))[0])
        self.execSize          = int(struct.unpack('I', fIn.read(4))[0])
        self.channelOffset     = int(struct.unpack('I', fIn.read(4))[0])

    def get_offset(self):
        return self.offset

    def is_write(self):
        return self.isWrite
    
    def is_block_2d(self):
        return self.isBlock2D

    def is_scatter(self):
        return self.isScatter

    def is_bts(self):
        return self.isBTS

    def is_slm(self):
        return self.isSLM

    def is_scratch_block(self):
        return self.isScratch

    def is_atomic(self):
        return self.isAtomic

    def is_fence(self):
        return self.isFence

    def addr_width(self):
        return self.addressWidth

    def simd_width(self):
        return self.simdWidth

    def binding_table_index(self):
        return self.bti

    def payload_len(self):
        return self.addrPayloadLength

    def data_port(self):
        return {
             0 : " DP0",
             1 : " DP1",
             2 : " UGM",
             3 : "UGML",
             4 : " TGM",
             5 : " SLM",
        }.get(self.dataPort, "OTHR")

    def is_eot(self):
        return self.isEOT

    def element_size(self):
        return self.elementSize

    def num_of_elements(self):
        return self.numElements

    def access_size(self):
        return (self.element_size() * self.num_of_elements())

    def is_media(self):
        return self.isMedia

    def channel_offset(self):
        return self.channelOffset

    def exec_size(self):
        return self.execSize

    def parse(self):        
        isWrite                = self.is_write()
        is_block_2d            = self.is_block_2d()
        isScatter              = self.is_scatter()
        isBTS                  = self.is_bts()
        isSLM                  = self.is_slm()
        isScratch              = self.is_scratch_block()
        isAtomic               = self.is_atomic()
        isFence                = self.is_fence()
        addressWidth           = self.addr_width()
        simdWidth              = self.simd_width()
        bti                    = self.binding_table_index()
        addrPayloadLength      = self.payload_len()
        isEOT                  = self.is_eot()
        elementSize            = self.element_size()
        numElements            = self.num_of_elements()
        isMedia                = self.is_media()

        return [isWrite, is_block_2d, isScatter, isBTS, isSLM, isScratch, isAtomic, isFence, addressWidth, simdWidth,bti, elementSize, numElements, addrPayloadLength, isEOT, isMedia]

    def show(self):

        [isWrite, isBlock2D, isScatter,isBTS,isSLM,isScratch,isAtomic,isFence,addressWidth,simdWidth,bti,elementSize,numElements,addrPayloadLength, isEOT, isMedia] = self.parse()

        WR = "W" if isWrite else "R"
        Scatter = "Scatter" if isScatter else "       "
        AddresingModel = "BTS" if isBTS else "A{}".format(addressWidth)
        AddrWidth = "{}".format(addressWidth)
        SIMD = "SIMD{:<2d}".format(simdWidth)
        BTI  = "BTI = {:02x}".format(bti)
        OperWidth = OperandWidthString(elementSize)
        AccessSize = "{}".format(elementSize*numElements)
        Special = "Fence" if isFence else "Media Block" if isMedia else ("SLM" if isSLM else ("Scratch Block" if isScratch else ""))
        Atomic = "A" if isAtomic  else " "
        Fence = "Fence" if isFence else " " 
        msg = WR + "  " + Scatter + "  " + AddresingModel + "   "+AddrWidth+"-bit access  " + SIMD + "  " + BTI + "  Operand = " + OperWidth
        if not isBlock2D:
            msg += "  Access size = " + AccessSize + "  " + Atomic + "  " + Special
        
        return msg


class BBL:
  id = 0
  sends = None


def find_bbl_from_id(BBLs, id):

    for i in list(range(len(BBLs))):
        if id == BBLs[i].id:
            return i

    return -1

def ProcessProfileDir(profileDir, verbose, single_thread):

    read_size = 0

    infilePath = os.path.join(profileDir, 'memorytrace_compressed.bin')

    if not os.path.isfile(infilePath):
        Error('File "%s" does not exist'%infilePath)

    try:
        fIn = open(infilePath, 'rb')
    except (IOError, e):
        Error(e)
    else:
        bbl_num = int(struct.unpack('I', fIn.read(4))[0])
        read_size += 4
        
        #print("bbl_num = {}".format(bbl_num))
                
        bbls = list(range(bbl_num)) #array('I', (0 for i in range(0, bbl_num-1)))
        #bbls.range(bbl_num) # = [0] * bbl_num
        
        BBLs = [BBL() for i in range(bbl_num)]

        for i in bbls:
            BBLs[i].id = int(struct.unpack('I', fIn.read(4))[0])
            read_size += 4
            num_of_sends = int(struct.unpack('I', fIn.read(4))[0])
            read_size += 4

            #print("BBL[{}] ID = {} num of sends = {}".format(i,BBLs[i].id,num_of_sends))

            BBLs[i].sends = [SEND() for j in range(num_of_sends)]

            for j in range(num_of_sends):
                BBLs[i].sends[j].read_from_file(fIn)
                read_size += 4 * 19

                #print("bbl[{:2d}]: offset = 0x{:08x} {}".format(j ,BBLs[i].sends[j].offset, BBLs[i].sends[j].show()))

        #print("read_size = 0x{:08x}".format(read_size))
        register_size  =  int(struct.unpack('I', fIn.read(4))[0])
        #print("register_size = {}".format(register_size))
        read_size += 4

        timestamp_included  =  int(struct.unpack('I', fIn.read(4))[0])
        #print("timestamp_included = {}".format(timestamp_included))
        read_size += 4
        timestamp = 0

        num_of_tiles = int(struct.unpack('I', fIn.read(4))[0])
        #print("num_of_tiles = {}".format(num_of_tiles))
        read_size += 4

        while num_of_tiles:

            TileID = int(struct.unpack('I', fIn.read(4))[0])
            read_size += 4

            num_of_threads = int(struct.unpack('I', fIn.read(4))[0])
            read_size += 4

            #print("read_size = 0x{:08x}".format(read_size))
            #print("TileID = {} num_of_threads = {}".format(TileID, num_of_threads))

            while num_of_threads:
                SliceID = int(struct.unpack('I', fIn.read(4))[0])
                DualSubSliceID = int(struct.unpack('I', fIn.read(4))[0])
                SubSliceID = int(struct.unpack('I', fIn.read(4))[0])
                EUID = int(struct.unpack('I', fIn.read(4))[0])
                ID = int(struct.unpack('I', fIn.read(4))[0])
                num_of_records = int(struct.unpack('I', fIn.read(4))[0])
                read_size += 4*6

                #print("num_of_records = {}".format(num_of_records))
                #print("tile_{}_s_{}_dss_{}_ss_{}_eu_{}_tid_{}".format(TileID, SliceID, DualSubSliceID, SubSliceID, EUID, ID))

                if DualSubSliceID != 0xFFFFFFFF:
                    outFile = os.path.join(profileDir, 'memtrace___tile_{}_s_{}_dss_{}_ss_{}_eu_{}_tid_{}.out'.format(TileID, SliceID, DualSubSliceID, SubSliceID, EUID, ID))
                else:
                    outFile = os.path.join(profileDir, 'memtrace___tile_{}_s_{}_ss_{}_eu_{}_tid_{}.out'.format(TileID, SliceID, SubSliceID, EUID, ID))

                fOut = open(outFile,"w")

                fOut.write(" BBL-ID   DP  SEND-OFFSET    RW      ADDRESS      SIZE\n"
                           "======================================================\n")

                while num_of_records:

                    bbl_id = int(struct.unpack('I', fIn.read(4))[0])
                    read_size += 4
                    #print("bbl_id = {}".format(bbl_id))
                    exec_mask = int(struct.unpack('I', fIn.read(4))[0])
                    read_size += 4
                    #print("exec_mask = {:08x}".format(exec_mask))

                    if timestamp_included != 0:
                        timestamp = int(struct.unpack('<Q', fIn.read(8))[0])
                        read_size += 8
                        #print("timestamp = {:#016x}".format(timestamp))

                    bbl = find_bbl_from_id(BBLs, bbl_id)
                    #print("bbl = {}".format(bbl))

                    num_of_sends = len(BBLs[bbl].sends)
                    #print("num_of_sends = {}".format(num_of_sends))

                    for s in range(num_of_sends):
                        addr_width             = BBLs[bbl].sends[s].addr_width()
                        is_scatter             = BBLs[bbl].sends[s].is_scatter()
                        access_size            = BBLs[bbl].sends[s].access_size()
                        payload_len            = BBLs[bbl].sends[s].payload_len()
                        is_write               = BBLs[bbl].sends[s].is_write()
                        is_block_2d            = BBLs[bbl].sends[s].is_block_2d()
                        simd_width             = BBLs[bbl].sends[s].simd_width()
                        data_port              = BBLs[bbl].sends[s].data_port()
                        is_eot                 = BBLs[bbl].sends[s].is_eot()
                        element_size           = BBLs[bbl].sends[s].element_size()
                        channel_offset         = BBLs[bbl].sends[s].channel_offset()
                        is_media               = BBLs[bbl].sends[s].is_media()
                        exec_size              = BBLs[bbl].sends[s].exec_size()

                        if timestamp_included != 0 and s != 0:
                            _ = fIn.read(4*5)
                            timestamp = int(struct.unpack('<Q', fIn.read(8))[0])
                            #print("timestamp = {:#016x}".format(timestamp))
                            _ = fIn.read(register_size - 4*7)

                        #print("exec_mask = 0x{:08x} channel_offset = 0x{:08x}".format(exec_mask, channel_offset))
                        #print("payload_len = {} read_size = 0x{:08x}".format(payload_len, read_size))

                        str = ""

                        if payload_len > 0:

                            if timestamp_included != 0:
                                fOut.write("timestamp = {:#016x}\n".format(timestamp));

                            if is_media:
                                offsetX = int(struct.unpack('I', fIn.read(4))[0])
                                offsetY = int(struct.unpack('I', fIn.read(4))[0])
                                control = int(struct.unpack('I', fIn.read(4))[0])
                                _ = fIn.read(4*5) 
                                read_size += 4*8
                                width = (control & 0x3F) + 1
                                heigth = ((control>>16) & 0x3F) + 1
                            
                                str = "  {:4d}   {}    0x{:08x}     {}      0x{:08x},0x{:08x}  {:>9}  // {}".format(bbl_id,data_port,BBLs[bbl].sends[s].offset,"W" if is_write else "R",offsetX,offsetY,'({}X{})'.format(width,heigth),BBLs[bbl].sends[s].show() if verbose else "")
                                fOut.write("{}\n".format(str))

                            elif is_scatter:
                                num_of_bytes_in_address = int(addr_width / 8)
                                num_of_addresses_in_payload = int(payload_len * register_size / num_of_bytes_in_address)
                                if addr_width == 32:
                                    for a in range(num_of_addresses_in_payload):

                                        addr = int(struct.unpack('I', fIn.read(4))[0])
                                        read_size += 4

                                        active_channel = ((exec_mask & (1 << (channel_offset + a))) != 0)
                                    
                                        #print("32bit: a = {:2d}  active_channel = {}".format(a, active_channel))

                                        if a < exec_size and active_channel:
                                            #print("addr = {:08x}".format(addr))
                                            str = "  {:4d}   {}    0x{:08x}     {}      0x{:08x}            {:3d}  // {}".format(bbl_id,data_port,BBLs[bbl].sends[s].offset,"W" if is_write else "R",addr,access_size,BBLs[bbl].sends[s].show() if verbose else "")
                                            fOut.write("{}\n".format(str));
                                else:
                                    for a in range(num_of_addresses_in_payload):

                                        addr = int(struct.unpack('Q', fIn.read(8))[0])
                                        read_size += 8

                                        active_channel = ((exec_mask & (1 << (channel_offset + a))) != 0)

                                        #print("64bit: a = {:2d}  active_channel = {}".format(a, active_channel))

                                        if a < exec_size and active_channel:
                                            #print("addr = {:016x}".format(addr))
                                            str = "  {:4d}   {}    0x{:08x}     {}      0x{:016x}    {:3d}  // {}".format(bbl_id,data_port,BBLs[bbl].sends[s].offset,"W" if is_write else "R",addr,access_size,BBLs[bbl].sends[s].show() if verbose else "")
                                            fOut.write("{}\n".format(str));
                            else:
                                assert (payload_len == 1)
                                num_of_bytes_in_address = int(addr_width / 8)
                                num_of_addresses_in_payload = int(payload_len * register_size / num_of_bytes_in_address)
                                if addr_width == 32:
                                    for a in range(num_of_addresses_in_payload):
                                        addr = int(struct.unpack('I', fIn.read(4))[0])
                                        read_size += 4
                                        if a == 2:
                                           str = "  {:4d}   {}    0x{:08x}     {}      0x{:08x}            {:3d}  // {}".format(bbl_id,data_port,BBLs[bbl].sends[s].offset,"W" if is_write else "R",addr * element_size,access_size,BBLs[bbl].sends[s].show() if verbose else "")
                                           fOut.write("{}\n".format(str));
                                else:
                                    if is_block_2d:
                                        surface_base_address = int(struct.unpack('Q', fIn.read(8))[0])
                                        surface_width        = 1 + int(struct.unpack('I', fIn.read(4))[0])
                                        surface_height       = 1 + int(struct.unpack('I', fIn.read(4))[0])
                                        surface_pitch        = 1 + int(struct.unpack('I', fIn.read(4))[0])
                                        block_start_x        = access_size * int(struct.unpack('I', fIn.read(4))[0])
                                        block_start_y        = int(struct.unpack('I', fIn.read(4))[0])
                                        block_width          = 1 + int(struct.unpack('B', fIn.read(1))[0])
                                        block_height         = 1 + int(struct.unpack('B', fIn.read(1))[0])
                                        array_length         = 1 + int(struct.unpack('B', fIn.read(1))[0])
                                        block2d_read_size    = 8 + 4*5 + 3
                                        
                                        if(register_size > block2d_read_size):
                                            _ = fIn.read(register_size - block2d_read_size)
                                            
                                        line_access_size = array_length * block_width * access_size
                                            
                                        for y in range(block_start_y, block_start_y + block_height):
                                            addr = surface_base_address + y * surface_pitch + block_start_x * access_size
                                            str = "  {:4d}   {}    0x{:08x}     {}      0x{:016x}    {:3d}  // ".format(bbl_id,data_port,BBLs[bbl].sends[s].offset,"W" if is_write else "R",addr,line_access_size)
                                            if verbose:
                                                str += "{} Access size = {} Data size = {} X = {} Y = {} Surface width = {} Surface pitch = {}".format(BBLs[bbl].sends[s].show(),line_access_size,access_size,block_start_x,y,surface_width,surface_pitch)
                                            
                                            fOut.write("{}\n".format(str))
                                    else:
                                        for a in range(num_of_addresses_in_payload):
                                            addr = int(struct.unpack('Q', fIn.read(8))[0])
                                            read_size += 8
                                            if a == 0:
                                                str = "  {:4d}   {}    0x{:08x}     {}      0x{:016x}    {:3d}  // {}".format(bbl_id,data_port,BBLs[bbl].sends[s].offset,"W" if is_write else "R",addr * element_size,access_size,BBLs[bbl].sends[s].show() if verbose else "")
                                                fOut.write("{}\n".format(str))

                                #str = "Not yet supported unscattered #print"
                                #fOut.write("{}\n".format(str));

                        if is_eot:
                           fOut.write("\n----EOT----\n\n");

                    num_of_records = num_of_records - 1

                fOut.close()
            
                if single_thread:
                    break

                num_of_threads = num_of_threads - 1

            num_of_tiles = num_of_tiles - 1

        fIn.close()

def Main():
    parser = ArgumentParser()
    parser.add_argument("--input_dir",  dest="input_dir",default='',
                        help="write report to FILE", metavar="FILE")
    parser.add_argument("--extract-one-file",  dest="extract_one_file",default=False,
                        action='store_true',
                        help="forcing extraction single file from trace")
    parser.add_argument("-v",
                        action='store_true', default=False,
                        help="Enable verbose")

    args = parser.parse_args()

    if (not os.path.exists(args.input_dir)):
        #print("the directory {} doesn't exist".format(args.input_dir))
        Error(Usage(sys.argv[0]))

    ProcessProfileDir(args.input_dir, args.v, args.extract_one_file)

Main()
