tiling_utils.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import os
from Schedulers import SchedulerDynamicBehaviourCNNs as scheduler
import SparseConv
import torch.sparse as sparse


class tiling_utils():
    def __init__(self, policy, gpu, MS, NS, KS):
        self.scheduler_info_dict = {}
        self.policy = policy
        self.MS = MS
        self.NS = NS
        self.KS = KS
        self.gpu = gpu

    def CNN_to_tiles(self, model: nn.Module, inputs: torch.tensor):
        for name, layer in model.named_modules():
            if isinstance(layer, nn.Conv2d):
                print(name)
                handle = layer.register_forward_hook(self.MM_tiling_hook(name))
                _ = model(inputs)
                handle.remove()

    def MM_tiling_hook(self, name):
        def hook(module, input, output):
            kernel_size = module.kernel_size
            stride = module.stride
            padding = module.padding
            dilation = module.dilation

            Unfold = nn.Unfold(kernel_size=kernel_size, dilation=dilation, padding=padding, stride=stride)
            unfolded_input = Unfold(input[0])
            flattened_weight = module.weight.view(module.weight.size(0), -1).t()

            for image_ind in range(unfolded_input.size(0)):
                # tiling(unfolded_input.transpose(1, 2)[image_ind], flattened_weight, 
                #        torch.zeros((unfolded_input.size(2), flattened_weight.size(1))), 
                #        16, 16, 16, file_dir=f"./tiles/{name}/{image_ind}")
                scheduler_info = scheduler.dynamic_scheduler(unfolded_input[image_ind], flattened_weight, \
                                                              self.policy, self.gpu, self.MS, self.NS, self.KS)
                if name not in self.scheduler_info_dict:
                    self.scheduler_info_dict[name] = {image_ind: scheduler_info}
                else:
                    self.scheduler_info_dict[name].update({image_ind: scheduler_info})

        return hook

    def conv_replacement(self, 
                         model: nn.Module,
                         name: str='') -> nn.Module:
        for name1, layer in model.named_children():
            if list(layer.children()) == []:
                if isinstance(layer, nn.Conv2d):
                    sparse_layer = SparseConv.SparseConv(
                                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()
                    name_ = name + name1
                    # Replace the module in the model
                    setattr(model, name1, sparse_layer)
            else:
                name += name1 + "."
                self.conv_replacement(layer, name)
        
        return model


def conv2d_im2col(self, layer: nn.Module,
            input_tensor: torch.tensor,
            ) -> torch.Tensor:    
        kernel_size = layer.kernel_size
        stride = layer.stride
        padding = layer.padding
        dilation = layer.dilation

        Unfold = nn.Unfold(kernel_size=kernel_size, dilation=dilation, padding=padding, stride=stride)
        unfolded_input = Unfold(input_tensor)

        flattened_weight = layer.weight.view(layer.weight.size(0), -1).t()
        out_unf = unfolded_input.transpose(1, 2).matmul(flattened_weight).transpose(1, 2)
        if layer.bias:
            out_unf += layer.bias.view(1, -1, 1)

        S = int(math.sqrt(out_unf.size(2)))
        Fold = nn.Fold(output_size=(S, S), kernel_size=(1, 1))
        folded_out = Fold(out_unf)

        return folded_out

def save_tensor(input_tensor: torch.tensor, file_name):
    tensor_file = open(file_name, "w")
    for i in range(input_tensor.size(0)):
        for j in range(input_tensor.size(1)):
            tensor_file.write(f"{input_tensor[i, j]}, ")
        tensor_file.write("\n")


#adopted and modified from https://github.com/TheColombianTeam/GPUScheduling
#source code: https://github.com/TheColombianTeam/GPUScheduling/blob/main/Kernel/basic_scheduler.py
def complete(matrix, MS=16, NS=16, KS=16):
    shape = matrix.size() #matrix.shape
    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_matrix = torch.zeros(new_shape).cuda() #np.zeros(new_shape)
    new_matrix[: shape[0], : shape[1]] = matrix
    return new_matrix

def complete3D(matrix, MS=16, NS=16, KS=16):
    shape = matrix.size() #matrix.shape
    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_matrix = torch.zeros(new_shape).cuda() #np.zeros(new_shape)
    new_matrix[:, : shape[1], : shape[2]] = matrix
    return new_matrix

def scheduler_sm(a, b, c, MS, NS, KS, file_dir="", tile_id=0):
    c_shape_original = c.size()
    a_shape = a.size()
    b_shape = b.size()
    c_shape = c.size()
    # mm_id = 0
    for row_c in range(c_shape[0] // KS):
        new_row_c_start = KS * row_c
        new_row_c_end = KS * row_c + KS
        for column_c in range(c_shape[1] // KS):
            new_column_c_start = KS * column_c
            new_column_c_end = KS * column_c + KS
            new_row_a_start = new_row_c_start
            new_row_a_end = new_row_c_end
            new_column_b_start = new_column_c_start
            new_column_b_end = new_column_c_end
            c_tensor = torch.zeros((KS, KS)).half().cuda()
            for column_a in range(a_shape[1] // KS):
                new_column_a_start = KS * column_a
                new_column_a_end = KS * column_a + KS
                new_row_b_start = new_column_a_start
                new_row_b_end = new_column_a_end
                a_tensor = a[
                    new_row_a_start:new_row_a_end, new_column_a_start:new_column_a_end
                ]
                b_tensor = b[
                    new_row_b_start:new_row_b_end, new_column_b_start:new_column_b_end
                ]
                c_tensor = c[
                    new_row_c_start:new_row_c_end, new_column_c_start:new_column_c_end
                ]

                ### To apply sparsity
                # b_tensor_sparse = b_tensor.half().cuda()
                # b_tensor_sparse = sparse.to_sparse_semi_structured(b_tensor_sparse.contiguous())
                # a_tensor = a_tensor.half().cuda()
                # c_tensor = torch.mm(a_tensor, b_tensor_sparse)
                
                ### only for dense operations
                
                c_tensor = torch.matmul(a_tensor, b_tensor) + c_tensor
                #print(a_tensor.size(), b_tensor.size(), c_tensor.size())

                c[
                    new_row_c_start:new_row_c_end, new_column_c_start:new_column_c_end
                ] = c_tensor
                
                # if file_dir:
                #     save_tensor(a_tensor, file_name=f"{file_dir}/MM_input_{tile_id}_{mm_id}.txt")
                #     save_tensor(b_tensor, file_name=f"{file_dir}/MM_weight_{tile_id}_{mm_id}.txt")
                
                # mm_id += 1
    #return a_tensor, b_tensor
    return c[: c_shape_original[0], : c_shape_original[1]]


def tiling(a, b, c, MS=16, NS=16, KS=16, file_dir=""):
    # if file_dir:
    #     if not os.path.exists(file_dir):
    #         os.makedirs(file_dir)

    c_shape_original = c.size()
    a = complete(a, MS, NS, KS).half().cuda()
    b = complete(b, MS, NS, KS).half().cuda()
    c = complete(c, MS, NS, KS).half().cuda()
    a_shape = a.size()
    b_shape = b.size()
    c_shape = c.size()
    tile_id = 0
    for row_c in range(c_shape[0] // MS):  
        new_row_c_start = MS * row_c  # ----- X
        new_row_c_end = new_row_c_start + MS
        for column_c in range(c_shape[1] // NS):
            new_column_c_start = NS * column_c  # ----- Y
            new_column_c_end = new_column_c_start + NS
            new_row_a_start = new_row_c_start
            new_row_a_end = new_row_c_end
            new_column_b_start = new_column_c_start
            new_column_b_end = new_column_c_end
            c_block = torch.zeros((MS, NS)).cuda()
            for column_a in range(a_shape[1] // KS):
                new_column_a_start = KS * column_a
                new_column_a_end = KS * column_a + KS
                new_row_b_start = new_column_a_start
                new_row_b_end = new_column_a_end
                a_block = a[
                    new_row_a_start:new_row_a_end, new_column_a_start:new_column_a_end
                ]
                b_block = b[
                    new_row_b_start:new_row_b_end, new_column_b_start:new_column_b_end
                ]
                c_block = c[
                    new_row_c_start:new_row_c_end, new_column_c_start:new_column_c_end
                ]
                c_block = scheduler_sm(a_block, b_block, c_block, MS, NS, KS, file_dir, tile_id)
                c[
                    new_row_c_start:new_row_c_end, new_column_c_start:new_column_c_end
                ] = c_block

                tile_id += 1
    return c[: c_shape_original[0], : c_shape_original[1]]