utils.py

import torch
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.nn as nn
from AM.AM_layer import AMConv2d
import errorgrad
from tqdm import tqdm




def get_all_layers(module):
    layers = []
    for child in module.children():
        # If the child is a container like Sequential or another Module, recurse into it
        if isinstance(child, nn.Module):
            sub_layer=get_all_layers(child)
            if sub_layer:
                layers.extend(sub_layer)
            else:
                if isinstance(child, AMConv2d):
        # Append the child itself (this will be the leaf layers)
                    layers.append(child)
    return layers


def get_output_and_gradient(model,inputs,targets):
        layer_inputs = {}
        layer_outputs = {}

        # Function to save outputs
        def hook_fn(module, input, output):
            layer_inputs[module] = input[0]
            layer_outputs[module] = output
            
        # Create a dictionary to store gradients
        layer_gradients = {}

        # Function to save gradients
        def save_gradients(module, grad_input, grad_output):
            layer_gradients[module] = grad_output[0]

        all_layer=get_all_layers(model)
        # Register hooks for each layer
        for layer in all_layer:
            layer.register_forward_hook(hook_fn)
            layer.register_backward_hook(save_gradients)
        # Forward pass
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        inputs, targets = inputs.to(device), targets.to(device)
        
        output = model(inputs)
        
        criterion = nn.CrossEntropyLoss()
        loss = criterion(output, targets)
        loss.backward()
        return layer_inputs,layer_outputs,layer_gradients
    


def fake_quantize_channelwise(x,bits):
    qmin=0
    qmax=2**bits-1
    xmax=torch.amax(x,dim=(2,3),keepdim=True)
    xmin=torch.amin(x,dim=(2,3),keepdim=True)
    scale=((xmax-xmin)/qmax)
    zero_point = torch.round(torch.nan_to_num(- xmin / scale, nan=0.0))
    # we assume weight quantization is always signed
    x_int = torch.round(torch.nan_to_num(x / scale, nan=0.0))
    x_quant = torch.clamp(x_int + zero_point, qmin,qmax)
    x_float_q = (x_quant - zero_point) * scale

    return x_float_q

def FakeQuantize(x,bits):
        qmin=0
        qmax=2**bits-1
        xmax=x.max()
        xmin=x.min()
        scale= (xmax - xmin) / qmax
        zero_point = torch.round(torch.nan_to_num(- xmin / scale, nan=0.0))
        # we assume weight quantization is always signed
            
        x_int = torch.round(torch.nan_to_num(x / scale, nan=0.0))
        x_quant = torch.clamp(x_int + zero_point, qmin,qmax)
        x_float_q = (x_quant - zero_point) * scale
        return x_float_q

def sensitivity(input,weight,stride,grad,padding,input_bits,weight_bits):
    input=FakeQuantize(input,input_bits)
    weight=fake_quantize_channelwise(weight,weight_bits)
    input_min=input.min()
    input_max=input.max()
    input_scale=(input_max-input_min)/(2**input_bits-1)
    input_offset=input_min

    weight_max=torch.amax(weight,dim=(2,3))
    weight_min=torch.amin(weight,dim=(2,3))
    weight_scale=(weight_max-weight_min)/(2**weight_bits-1)
    weight_offset=weight_min
    _reversed_padding_repeated_twice = padding*2
    input_dim,weight_dim=2**input_bits,2**weight_bits
    
    return errorgrad.LUTgrad(F.pad(input,_reversed_padding_repeated_twice), weight, torch.tensor(stride).float(),
                                          grad,input_dim,weight_dim,input_scale.float(),input_offset.float(),weight_scale.float(),weight_offset.float())


def get_each_layer_sensitivity_matrix(model,images,targets,mul):
    all_layers=get_all_layers(model)
    layer_inputs,layer_outputs,layer_gradients=get_output_and_gradient(model,images,targets)
    sensitivity_matrix_list=[]
    power_sensitivity_list=[]
    batch_size=len(images)
    for idx,layer in enumerate(all_layers):
        outputs=layer_outputs[layer]
        inputs=layer_inputs[layer]
        gradients=layer_gradients[layer]
        padding=layer.padding
        weights=layer.weight
        
        OC,IC,KW,KH=weights.size()
        B,W,H,OC=outputs.size()
        power_sensitivity_list.append(OC*IC*KW*KH*W*H)
        
        stride=torch.tensor(layer.stride)
        input_bits,weight_bits=int(mul[idx][3]),int(mul[idx][4])
        layer_sensitivity_matrix=sensitivity(inputs,weights,stride,gradients,padding,input_bits,weight_bits).cpu()/batch_size
        sensitivity_matrix_list.append(layer_sensitivity_matrix)
    return sensitivity_matrix_list,power_sensitivity_list



def hessian_dzde(model,images,targets,mul):
    grad=[]
    hessian=[]
    dzde=[]
    inputs_list=torch.split(images,1,dim=0)
    targets_list=torch.split(targets,1,dim=0)
    for sample_input,sample_target in tqdm(zip(inputs_list,targets_list)):
        
        sample_layer_inputs={}
        sample_layer_outputs={}
        sample_dzde={}
        def hook_fn(module, input, output):
            sample_layer_outputs[module]=output
            sample_layer_inputs[module]=input[0]
        
        for name, layer in model.named_modules():
            if 'conv' in name:
                layer.register_forward_hook(hook_fn)
        
        all_layer=get_all_layers(model)
        # Register hooks for each layer
        for layer in all_layer:
            layer.register_forward_hook(hook_fn)


        sample_output = model(sample_input)
        sample_output.retain_grad()
        criterion = nn.CrossEntropyLoss()
        loss = criterion(sample_output, sample_target)
        loss.backward(retain_graph=True)
        
        for layer_idx,layer in enumerate(all_layer):   
            sample_layer_output=sample_layer_outputs[layer]
            sample_layer_input=sample_layer_inputs[layer]
            lut=mul[layer_idx]
            input_bits,weight_bits=int(lut[3]),int(lut[4])
            weights=layer.weight
            stride=torch.tensor(layer.stride)
            sample_layer_dzde=torch.zeros(sample_output.size(1),2**(input_bits+weight_bits))
            padding=layer.padding
            for idx,sample_output_s in enumerate(sample_output[0]):
                current_dzdy=torch.autograd.grad(sample_output_s,sample_layer_output,retain_graph=True)[0]
                sample_layer_dzde[idx]=sensitivity(sample_layer_input,weights,stride,current_dzdy,padding,input_bits,weight_bits).flatten()
            sample_dzde[layer]=sample_layer_dzde

        
        dzde.append(sample_dzde)
        dldz=torch.autograd.grad(loss,sample_output,retain_graph=True,create_graph=True)[0].squeeze()
        grad.append(dldz)
        sample_hessian=torch.zeros(len(dldz),len(dldz))
        for idx,dldz_s in enumerate(dldz):
            hessian_row=torch.autograd.grad(dldz_s,sample_output,retain_graph=True)[0]
            sample_hessian[idx]=hessian_row.squeeze()
        hessian.append(sample_hessian)
    return hessian, dzde




def hessian_quick(model,images,targets,mul):
    

    from tqdm import tqdm
    approx_dzde=[]
    approx_top_eigen=[]
    inputs_list=torch.split(images,1,dim=0)
    targets_list=torch.split(targets,1,dim=0)
    for sample_input,sample_target in tqdm(zip(inputs_list,targets_list)):
        
        sample_layer_inputs={}
        sample_layer_outputs={}
        sample_dzde={}
        def hook_fn(module, input, output):
            sample_layer_outputs[module]=output
            sample_layer_inputs[module]=input[0]
        
        for name, layer in model.named_modules():
            if 'conv' in name:
                layer.register_forward_hook(hook_fn)
        
        all_layer=get_all_layers(model)
        # Register hooks for each layer
        for layer in all_layer:
            layer.register_forward_hook(hook_fn)


        sample_output = model(sample_input)
        sample_output.retain_grad()
        criterion = nn.CrossEntropyLoss()
        loss = criterion(sample_output, sample_target)
        loss.backward(retain_graph=True)


        dldz=torch.autograd.grad(loss,sample_output,retain_graph=True,create_graph=True)[0].squeeze()
        b = torch.rand_like(dldz)
        num_simulations=100
        for _ in range(num_simulations):
            # Multiply the matrix with the vector
            b = torch.autograd.grad(dldz,sample_output,grad_outputs=b,create_graph=True,retain_graph=True)[0].squeeze()
            
            # Normalize the vector
            b = b / torch.linalg.norm(b)
        
        # Return the dominant eigenvalue and eigenvector
        eigenvalue = torch.dot(b.T, torch.autograd.grad(dldz,sample_output,grad_outputs=b,create_graph=True,retain_graph=True)[0].squeeze())
        approx_top_eigen.append(eigenvalue)
        v_max=b
        
        for layer_idx,layer in enumerate(all_layer):   
            sample_layer_output=sample_layer_outputs[layer]
            sample_layer_input=sample_layer_inputs[layer]
            lut=mul[layer_idx]
            input_bits,weight_bits=int(lut[3]),int(lut[4])
            weights=layer.weight
            stride=torch.tensor(layer.stride)
            padding=layer.padding
            dzdy=torch.autograd.grad(sample_output[0],sample_layer_output,grad_outputs=v_max.T,retain_graph=True)[0]
            sample_layer_dzde=sensitivity(sample_layer_input,weights,stride,dzdy,padding,input_bits,weight_bits).flatten()
            sample_dzde[layer]=sample_layer_dzde
        approx_dzde.append(sample_dzde)
    return approx_dzde,approx_top_eigen




import struct

def get_bits(file_path):
    input_bits=int(file_path[3])
    weight_bits=int(file_path[4])
    return (input_bits,weight_bits)

def generate_mul_matrix(mul_path):
    data = []
    with open(mul_path, 'rb') as file:
        while True:
            # Read 2 bytes at a time since each result is a uint16_t
            chunk = file.read(2)
            if not chunk:
                break
            # Unpack the 2 bytes into a uint16_t
            result = struct.unpack('H', chunk)[0]
            data.append(result)
    return torch.tensor(data)
    

def generate_bit_error_matrix_list(bit_mul_path):
    input_bits,weight_bits=get_bits(bit_mul_path[0])
    result=torch.zeros([len(bit_mul_path),2**input_bits,2**weight_bits])
    for i in range(len(bit_mul_path)):
        mul_path='lut/'+bit_mul_path[i]+'.bin'
        mul_matrix=generate_mul_matrix(mul_path)
        mul_matrix=mul_matrix.view([2**input_bits,2**weight_bits])
        mul_error_matrix=torch.tensor([[mul_matrix[a][w]-a*w for w in range(mul_matrix.size(1)) ] for a in range(mul_matrix.size(0))])
        result[i]=mul_error_matrix
    return result

def approx_info(file_path='approxmul.xlsx'):
    import pandas as pd

    data = pd.read_excel(file_path)

    approx_mul=data['name'].to_list()
    # Set the 'Name' column as the index to easily create dictionaries
    data.set_index('name', inplace=True)

    # Create dictionaries for power and delay
    power_dict = data['power'].to_dict()
    delay_dict = data['delay'].to_dict()

    # Optional: You can also create a dictionary for area if needed
    area_dict = data['area'].to_dict()
    bit_approx_list_dict={}
    for mul in approx_mul:
        if mul[0:5] not in bit_approx_list_dict.keys():
            bit_approx_list_dict[mul[0:5]]=[mul]
        else:
            bit_approx_list_dict[mul[0:5]].append(mul)
            
    

    bit_error_matrix_list_dict={}
    for key in bit_approx_list_dict.keys():
        bit_error_matrix_list_dict[key]=generate_bit_error_matrix_list(bit_approx_list_dict[key])

    return bit_approx_list_dict,bit_error_matrix_list_dict,power_dict,delay_dict




def talyor_estimation(model,images,targets,mul_list,error_matrix_list_dict):
    sensitivity_matrix_list,power_sensitivity_list=get_each_layer_sensitivity_matrix(model,images,targets,mul_list)
    hessian,dzde=hessian_dzde(model,images,targets,mul_list)
    first_order_perturbation_list=[]
    second_order_perturbation_list = []
    dl2de2 = {}
    batch_size = len(dzde)

    # Batch processing using torch.einsum
    for idx,layer in enumerate(tqdm(get_all_layers(model))):
        lut_feature=mul_list[idx][0:5]
        #first order
        current_sensitivity_matrix=sensitivity_matrix_list[idx].cpu()
        error_list=error_matrix_list_dict[lut_feature]
        layer_perturbation=torch.zeros(len(error_list))
        for i in range(len(error_list)):
            layer_perturbation[i]=torch.sum(current_sensitivity_matrix*error_list[i])
        first_order_perturbation_list.append(layer_perturbation)
        
        
        # second order
        layer_dzde_stack = torch.stack([dzde[b][layer] for b in range(batch_size)])  # Shape: (batch_size, z, e)
        hessian_stack = torch.stack(hessian)  # Shape: (batch_size, m, m)

        dl2de2[layer] = torch.einsum('bmj,bmn,bnk->jk', layer_dzde_stack, hessian_stack, layer_dzde_stack) / batch_size
        error_matrix_stack = torch.stack([em.flatten() for em in error_matrix_list_dict[lut_feature]])  # Shape: (num_errors, num_elements)
        perturbations = torch.einsum('ni,ij,nj->n', error_matrix_stack,dl2de2[layer],error_matrix_stack)  # Shape: (num_errors,)
        second_order_perturbation_list.append(perturbations.tolist())
    
    perturbation_list=[]
    batch_size = len(dzde)
    for i in range(len(first_order_perturbation_list)):
        perturbation_list.append(first_order_perturbation_list[i]+0.5*torch.tensor(second_order_perturbation_list[i]))
        
    return perturbation_list,first_order_perturbation_list,second_order_perturbation_list,power_sensitivity_list





def quick_talyor_estimation(model,images,targets,mul_list,error_matrix_list_dict):
    sensitivity_matrix_list,power_sensitivity_list=get_each_layer_sensitivity_matrix(model,images,targets,mul_list)
    approx_dzde,approx_top_eigen=hessian_quick(model,images,targets,mul_list)
    first_order_perturbation_list=[]
    second_order_perturbation_list = []

    batch_size = len(images)

    # Batch processing using torch.einsum
    for idx,layer in enumerate(tqdm(get_all_layers(model))):
        lut_feature=mul_list[idx][0:5]
        #first order
        current_sensitivity_matrix=sensitivity_matrix_list[idx].cpu()
        error_list=error_matrix_list_dict[lut_feature]
        layer_perturbation=torch.zeros(len(error_list))
        for i in range(len(error_list)):
            layer_perturbation[i]=torch.sum(current_sensitivity_matrix*error_list[i])
        first_order_perturbation_list.append(layer_perturbation)
        
        # Perform the batched computation over `dl2de2` and `error_matrix_stack`
        perturbations=torch.zeros(len(error_list))
        error_matrix_stack = torch.stack([em.flatten() for em in error_list])
        for b in range(batch_size):
            perturbations += approx_top_eigen[b].cpu()*torch.einsum('ni,i,j,nj->n', error_matrix_stack,approx_dzde[b][layer].cpu(),approx_dzde[b][layer].cpu(),error_matrix_stack)  # Shape: (num_errors,)

        # Convert perturbations to list and add to results
        second_order_perturbation_list.append((perturbations/batch_size).tolist())
        
        perturbation_list=[]
        for i in range(len(first_order_perturbation_list)):
            perturbation_list.append(first_order_perturbation_list[i]+0.5*torch.tensor(second_order_perturbation_list[i]))
        
    return perturbation_list,first_order_perturbation_list,second_order_perturbation_list,power_sensitivity_list