fault_simulation.py

import torch
import torch.nn as nn
import torch.sparse as sparse

class FISparseConv(nn.Conv2d):
    def __init__(self, FI_location, *args, **kwargs):
        super(FISparseConv, self).__init__(*args, **kwargs)
        self.FI_location = FI_location
        self.injection_activate = 0

    def forward(self, 
                input_act: torch.tensor) -> torch.tensor:
        Unfold = nn.Unfold(kernel_size=self.kernel_size, dilation=self.dilation, padding=self.padding, stride=self.stride)
        unfolded_input = input_act.to(torch.float16).detach()
        unfolded_input = Unfold(unfolded_input)
        unfolded_input_size = unfolded_input.size()

        flattened_weight = self.weight.half().detach()
        flattened_weight = flattened_weight.view(self.weight.size(0), -1).t()
        fw_complete = self._complete(flattened_weight.T).contiguous()
        fw_sparse = sparse.to_sparse_semi_structured(fw_complete)
        ui_trans = self._complete3D(unfolded_input)
        
        fw_sparse_size = fw_sparse.size()
        flattened_weight_size = flattened_weight.size()

        del unfolded_input
        del flattened_weight 
        del fw_complete

        #inject a single random fault into weight or metadata
        if self.injection_activate == 1:
            if self.FI_location == "metadata":
                self._metadata_FI(fw_sparse.indices()[:])

            elif self.FI_location == "weight":
                self._weight_FI(fw_sparse.values()[:])
        
        out_unf = torch.zeros((unfolded_input_size[0], fw_sparse_size[0], ui_trans.size(2)), dtype=torch.float16, device='cuda')
        for ind in range(unfolded_input_size[0]):
            out_unf[ind] = torch.mm(fw_sparse, ui_trans[ind]).detach()

        if self.bias is not None:
            out_unf = out_unf + self.bias.view(1, -1, 1)

        out_unf_sliced = out_unf[:, :flattened_weight_size[1], :unfolded_input_size[2]].detach()

        del fw_sparse
        del ui_trans
        del out_unf

        output_size = torch.sqrt(torch.tensor([out_unf_sliced.size(2)])).to(torch.int)
        Fold = nn.Fold(output_size=(output_size, output_size), kernel_size=(1, 1))
        output_act = Fold(out_unf_sliced)
        
        del out_unf_sliced
        del Fold
        del Unfold
        torch.cuda.empty_cache()

        return output_act.to(torch.float).detach()
    

    def _complete(self, tensor, MS=64, NS=64, KS=64):
        shape = tensor.size()
        if (shape[0] % MS) > 0:
            new_shape_a = MS * (shape[0] // MS) + MS
        else:
            new_shape_a = shape[0]
        if (shape[1] % NS) > 0:
            new_shape_b = NS * (shape[1] // NS) + NS
        else:
            new_shape_b = shape[1]
        new_shape = (new_shape_a, new_shape_b)
        new_tensor = torch.zeros(new_shape, dtype=torch.float16, device='cuda') #.half().cuda()
        new_tensor[: shape[0], : shape[1]] = tensor     #.half()
        
        del tensor
        del shape
        del new_shape
        torch.cuda.empty_cache()
        # gc.collect()

        return new_tensor.detach()

    def _complete3D(self, tensor, MS=64, NS=64, KS=64):
        shape = tensor.size()
        if (shape[1] % MS) > 0:
            new_shape_a = MS * (shape[1] // MS) + MS
        else:
            new_shape_a = shape[1]
        if (shape[2] % NS) > 0:
            new_shape_b = NS * (shape[2] // NS) + NS
        else:
            new_shape_b = shape[2]
        new_shape = (shape[0], new_shape_a, new_shape_b)
        new_tensor = torch.zeros(new_shape,  dtype=torch.float16, device='cuda')  #.half().cuda() 
        new_tensor[:, : shape[1], : shape[2]] = tensor      #.half()
        
        del tensor
        del shape
        del new_shape
        torch.cuda.empty_cache()
        # gc.collect()

        return new_tensor.detach()

    # Applies a random bitflip into a 2D metadata array with torch.int16 dtype
    def _metadata_FI(self, indices):
        (dim0, dim1) = indices.size()
        rand_dim1 = torch.randint(low=0, high=dim0, size=(1,))
        rand_dim2 = torch.randint(low=0, high=dim1, size=(1,))
        rand_bit  = torch.randint(low=0, high=16, size=(1,))

        mask = torch.tensor(1 << rand_bit, dtype=torch.int16).cuda()
        indices[rand_dim1, rand_dim2] = mask ^ indices[rand_dim1, rand_dim2]

        del rand_bit
        del rand_dim1
        del rand_dim2
        del mask
        torch.cuda.empty_cache()
        

    # Applies a random bitflip into a 2D weight array with torch.float16 dtype
    def _weight_FI(self, weights):
        (dim0, dim1) = weights.size()
        rand_dim1 = torch.randint(low=0, high=dim0, size=(1,))
        rand_dim2 = torch.randint(low=0, high=dim1, size=(1,))
        rand_bit  = torch.randint(low=0, high=16, size=(1,))

        mask = torch.tensor(1 << rand_bit, dtype=torch.int16).cuda()
        weights[rand_dim1, rand_dim2] = (mask ^ weights[rand_dim1, rand_dim2].view(torch.int16)).view(torch.float16)
        
        del rand_bit
        del rand_dim1
        del rand_dim2
        del mask
        torch.cuda.empty_cache()

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, injection_activate={self.injection_activate})"


def FI_spconv_replacement(model: nn.Module,
                    FI_location: str='metadata',
                    ) -> nn.Module:
    # assert(target_layer_name is not None and target_layer_name != "")
    
    for name1, layer in model.named_children():
        if list(layer.children()) == []:
            if isinstance(layer, nn.Conv2d):    # and name1 == target_layer_name:
                sparse_layer = FISparseConv(
                            FI_location=FI_location,
                            in_channels=layer.in_channels,
                            out_channels=layer.out_channels,
                            kernel_size=layer.kernel_size,
                            stride=layer.stride,
                            padding=layer.padding,
                            dilation=layer.dilation,
                            groups=layer.groups,
                            bias=(layer.bias is not None),
                            padding_mode=layer.padding_mode,
                            )
                sparse_layer.weight.data = layer.weight.data    #.clone()
                if layer.bias is not None:
                    sparse_layer.bias.data = layer.bias.data    #.clone()
                # Replace the module in the model
                setattr(model, name1, sparse_layer)
        else:
            FI_spconv_replacement(layer, FI_location)
    
    # return model




class FIConv(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        super(FIConv, self).__init__(*args, **kwargs)
        self.injection_activate = 0

    def forward(self, 
                input_act: torch.tensor) -> torch.tensor:
        #inject a single random fault into weight 
        if self.injection_activate == 1:
            self.weight = self._weight_FI(self.weight)
        
        output_act = super(FIConv, self).forward(input_act)
        torch.cuda.empty_cache()

        return output_act.to(torch.float).detach()
    
    # Applies a random bitflip into a 2D weight array with torch.float16 dtype
    def _weight_FI(self, weights):
        flattened_weights = weights.flatten().detach().to(torch.half)
        dim0 = flattened_weights.size(0)
        rand_dim0 = torch.randint(low=0, high=dim0, size=(1,))
        rand_bit  = torch.randint(low=0, high=16, size=(1,))

        #print(flattened_weights.size(), dim0, rand_dim0.size(), rand_dim0)

        mask = torch.tensor(1 << rand_bit, dtype=torch.int16).cuda()
        flattened_weights[rand_dim0] = (mask ^ flattened_weights[rand_dim0].view(torch.int16)).view(torch.float16)
        weights = nn.Parameter(flattened_weights.view(weights.size()).to(torch.float32))
        
        del rand_bit
        del rand_dim0
        del mask
        del flattened_weights
        torch.cuda.empty_cache()

        return weights

    def __repr__(self):
        base_repr = super().__repr__()
        return f"{base_repr[:-1]}, injection_activate={self.injection_activate})"





def FI_conv_replacement(model: nn.Module,
                    FI_location: str='metadata',
                    ) -> nn.Module:    
    for name1, layer in model.named_children():
        if list(layer.children()) == []:
            if isinstance(layer, nn.Conv2d):
                sparse_layer = FIConv(
                            in_channels=layer.in_channels,
                            out_channels=layer.out_channels,
                            kernel_size=layer.kernel_size,
                            stride=layer.stride,
                            padding=layer.padding,
                            dilation=layer.dilation,
                            groups=layer.groups,
                            bias=(layer.bias is not None),
                            padding_mode=layer.padding_mode,
                            )
                sparse_layer.weight.data = layer.weight.data
                if layer.bias is not None:
                    sparse_layer.bias.data = layer.bias.data
                # Replace the module in the model
                setattr(model, name1, sparse_layer)
        else:
            FI_conv_replacement(layer, FI_location)


def activate_FI_layer(model, layer_name):
    for name1, layer in model.named_children():
        if list(layer.children()) == []:
            if isinstance(layer, FISparseConv) or isinstance(layer, FIConv):
                if layer_name == name1:
                    layer.injection_activate = 1
                else:
                    layer.injection_activate = 0
        else:
            activate_FI_layer(layer, layer_name)


def get_layers_names(model: nn.Module, name_list: list=[]) -> list:
    for name1, layer in model.named_children():
        if list(layer.children()) == []:
            if isinstance(layer, nn.Conv2d):
                name_list.append(name1)
        else:
            get_layers_names(layer, name_list)
    
    return name_list