From 87d8f2d0edaa4ed431d9ca0091bc53742fdd8dcb Mon Sep 17 00:00:00 2001 From: lrubens Date: Mon, 28 Oct 2024 14:27:15 -0700 Subject: [PATCH] Fixed issue with data generator for block sparse --- sam/onyx/generate_matrices.py | 168 ++++++++++++++++++++++++---------- 1 file changed, 120 insertions(+), 48 deletions(-) diff --git a/sam/onyx/generate_matrices.py b/sam/onyx/generate_matrices.py index d2a80286..e7df7ca3 100644 --- a/sam/onyx/generate_matrices.py +++ b/sam/onyx/generate_matrices.py @@ -5,6 +5,9 @@ import random import scipy.sparse as ss import tempfile + +import torch + from sam.onyx.fiber_tree import * import argparse import math @@ -13,21 +16,23 @@ from sam.sim.test.test import * -class MatrixGenerator(): +class MatrixGenerator: def __init__(self, name='B', shape=None, mode_ordering=None, block_size=1, block_naive=True, sparsity=0.6, - format='CSF', dump_dir=None, tensor=None, value_cap=None) -> None: + format='CSF', dump_dir=None, tensor=None, value_cap=None, transpose=False) -> None: # assert dimension is not None # self.dimension = dimension self.shape = shape self.array = None + self.blocked_array = None self.sparsity = sparsity self.format = format self.name = name self.mode_ordering = mode_ordering self.block_size = block_size self.block_naive = block_naive + self.transpose = transpose if value_cap is None: self.value_cap = int(math.pow(2, 8)) - 1 else: @@ -67,23 +72,18 @@ def _create_matrix(self, value_cap=int(math.pow(2, 8)) - 1): ''' Routine to create the actual matrix from the dimension/shape ''' - if self.block_size > 1: - # self.shape = [*self.shape[:len(self.shape) - 2], self.shape[len(self.shape) - 2] * self.block_size, - # self.shape[len(self.shape) - 2] * self.block_size] - self.shape[len(self.shape) - 2] *= self.block_size - self.shape[len(self.shape) - 1] *= self.block_size - self.array = numpy.random.randint(low=1, high=value_cap, size=self.shape) + self.array = numpy.random.randint(low=0, high=value_cap, size=self.shape) + for idx, x in numpy.ndenumerate(self.array): + rand = random.random() if self.block_size > 1: if (idx[-2] % self.block_size == 0) and (idx[-1]) % self.block_size == 0: - if random.random() < self.sparsity: + if rand < self.sparsity: self.array[..., idx[-2]:idx[-2]+self.block_size, idx[-1]:idx[-1]+self.block_size] = 0 else: - if random.random() < self.sparsity: + if rand < self.sparsity: self.array[idx] = 0 - # print(self.array[...,:self.shape[-2]:self.block_size,:self.shape[-1]:self.block_size]) - def _create_fiber_tree(self): # self.fiber_tree = FiberTree(tensor=self.array if self.block_naive else self.array[..., # self.shape[-2]:self.block_size,:self.shape[-1]:self.block_size]) @@ -123,16 +123,17 @@ def dump_outputs(self, format=None, tpose=False, dump_shape=True, glb_override=F print(f"Using dump directory - {use_dir}") - all_zeros = not np.any(self.array) + # all_zeros = not np.any(self.array, out=self.array) + all_zeros = False debug_coord_arr = [] debug_seg_arr = [] debug_val_arr = [] # Transpose it first if necessary - if tpose is True: - self.array = numpy.transpose(self.array) - self.shape = self.array.shape - self.fiber_tree = FiberTree(tensor=self.array) + # if tpose is True: + # self.array = numpy.transpose(self.array) + # self.shape = self.array.shape + # self.fiber_tree = FiberTree(tensor=self.array) if format is not None: self.format = format @@ -171,18 +172,12 @@ def dump_outputs(self, format=None, tpose=False, dump_shape=True, glb_override=F # In CSF format, need to iteratively create seg/coord arrays tmp_lvl_list = [] small_tmp_lvl_list = [] - # if self.block_size == 1: tmp_lvl_list.append(self.fiber_tree.get_root()) - # else: - # tmp_lvl_list.append(self.tmp_fiber_tree.get_root()) small_tmp_lvl_list.append(self.tmp_fiber_tree.get_root()) - # print(small_tmp_lvl_list) - # print(tmp_lvl_list) seg_arr, coord_arr = None, None if self.block_size > 1: seg_arr, coord_arr = self._dump_csf(small_tmp_lvl_list) - # print(seg_arr, coord_arr) else: seg_arr, coord_arr = self._dump_csf(tmp_lvl_list) if glb_override: @@ -190,8 +185,6 @@ def dump_outputs(self, format=None, tpose=False, dump_shape=True, glb_override=F self.write_array(lines, name=f"tensor_{self.name}_mode_{self.mode_ordering[0]}", dump_dir=use_dir, hex=print_hex) else: - print(self.mode_ordering) - # print(seg_arr, coord_arr) self.write_array(seg_arr, name=f"tensor_{self.name}_mode_{self.mode_ordering[0]}_seg", dump_dir=use_dir, hex=print_hex) self.write_array(coord_arr, name=f"tensor_{self.name}_mode_{self.mode_ordering[0]}_crd", @@ -222,24 +215,27 @@ def dump_outputs(self, format=None, tpose=False, dump_shape=True, glb_override=F # self.write_array(tmp_lvl_list, name=f"tensor_{self.name}_mode_vals" dump_dir=use_dir) self.write_array(lines, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir, hex=print_hex) else: - reached_full_vals = False - if self.block_size > 1: - tmp_lst = tmp_lvl_list - # Retrieve values from full tensor - while not reached_full_vals: - next_tmp_lvl_list = [] - for fib in tmp_lst: - crd_payloads_tmp = fib.get_coord_payloads() - if type(crd_payloads_tmp[0][1]) is not FiberTreeFiber: - reached_full_vals = True - for crd, pld in crd_payloads_tmp: - next_tmp_lvl_list.append(pld) - else: - for crd, pld in crd_payloads_tmp: - next_tmp_lvl_list.append(pld) - tmp_lst = next_tmp_lvl_list - self.write_array(tmp_lst, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir, - hex=print_hex) + # reached_full_vals = False + # if self.block_size > 1: + # tmp_lst = tmp_lvl_list + # # Retrieve values from full tensor + # while not reached_full_vals: + # next_tmp_lvl_list = [] + # for fib in tmp_lst: + # crd_payloads_tmp = fib.get_coord_payloads() + # if type(crd_payloads_tmp[0][1]) is not FiberTreeFiber: + # reached_full_vals = True + # for crd, pld in crd_payloads_tmp: + # next_tmp_lvl_list.append(pld) + # else: + # for crd, pld in crd_payloads_tmp: + # next_tmp_lvl_list.append(pld) + # tmp_lst = next_tmp_lvl_list + if self.block_size == 1: + self.write_array(tmp_lst, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir, + hex=print_hex) + else: + self.write_blocked_array(self.array, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir, skip_zeros=True) else: seg_arr, coord_arr = self._dump_csf(tmp_lst) if glb_override: @@ -260,7 +256,10 @@ def dump_outputs(self, format=None, tpose=False, dump_shape=True, glb_override=F lines = [len(flat_array), *flat_array] self.write_array(lines, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir, hex=print_hex) else: - self.write_array(flat_array, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir, hex=print_hex) + if self.block_size == 1: + self.write_array(flat_array, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir, hex=print_hex) + else: + self.write_blocked_array(self.array, name=f"tensor_{self.name}_mode_vals", dump_dir=use_dir) elif self.format == "COO": crd_dict = dict() order = len(self.array.shape) @@ -293,8 +292,14 @@ def dump_outputs(self, format=None, tpose=False, dump_shape=True, glb_override=F hex=print_hex) if dump_shape: - final_shape = [x // self.block_size for x in self.array.shape] - self.write_array(final_shape, name=f"tensor_{self.name}_mode_shape", dump_dir=use_dir, hex=print_hex) + # final_shape = self.array.shape + final_shape = [x if i < len(self.array.shape) - 2 else x // self.block_size for i, x in enumerate(self.array.shape)] + # final_shape = final_shape[self.mode_ordering] + shape = [] + for elem in self.mode_ordering: + shape.append(final_shape[elem]) + + self.write_array(shape, name=f"tensor_{self.name}_mode_shape", dump_dir=use_dir, hex=print_hex) # Transpose it back if tpose is True: @@ -323,7 +328,7 @@ def _dump_csf(self, level_list): return seg_arr, coord_arr - def write_array(self, str_list, name, dump_dir=None, hex=False): + def write_array(self, str_list, name, dump_dir=None, hex=False, integer=False, num_repeats=1): """Write an array/list to a file Args: @@ -336,7 +341,35 @@ def write_array(self, str_list, name, dump_dir=None, hex=False): full_path = dump_dir + "/" + name with open(full_path, "w+") as wr_file: for item in str_list: - item_int = int(item) + if integer: + item_int = int(item) + else: + item_int = item + if hex: + wr_file.write(f"{item_int:04X}\n") + else: + for _ in range(num_repeats): + wr_file.write(f"{item_int}\n") + + def write_blocked_array(self, str_list, name, dump_dir=None, integer=False, hex=False, skip_zeros=False): + if dump_dir is None: + dump_dir = self.dump_dir + + tiles = tile_and_unroll_nd(str_list, [self.block_size, self.block_size]) + + if skip_zeros: + print("Skipping zeros") + tiles = [tile for tile in tiles if numpy.sum(tile) != 0] + + tiles = np.concatenate(tiles) + + full_path = dump_dir + "/" + name + with open(full_path, "w+") as wr_file: + for item in tiles: + if integer: + item_int = int(item) + else: + item_int = item if hex: wr_file.write(f"{item_int:04X}\n") else: @@ -361,6 +394,45 @@ def __setitem__(self, key, val): self.array[key] = val +def tile_and_unroll_nd(array, tile_size): + """ + Generalized function to tile and unroll a multidimensional array. + + Parameters: + - array: Input numpy array of any dimensionality. + - tile_size: A tuple that specifies the size of the tiles for the last dimensions. + + Returns: + - A list of flattened tiles. + """ + array_shape = array.shape + num_non_tiled_dims = len(array_shape) - len(tile_size) # Non-tiled dimensions + tiled_dims = array_shape[num_non_tiled_dims:] # Dimensions that will be tiled + + # Calculate the number of tiles along the tiling dimensions + tiled_sizes = [dim // tile for dim, tile in zip(tiled_dims, tile_size)] + + # Prepare a list to collect the tiles + tiles = [] + + # Iterate over the non-tiled dimensions + for index in np.ndindex(*array_shape[:num_non_tiled_dims]): + # Iterate over the tiled dimensions and extract each tile + for tile_index in np.ndindex(*tiled_sizes): + # Slice the array to extract the current tile + slices = tuple( + slice(idx * size, (idx + 1) * size) for idx, size in zip(tile_index, tile_size) + ) + full_slices = index + slices # Combine non-tiled and tiled indices + tile = array[full_slices] + if type(tile) == torch.Tensor: + tile = tile.numpy() + # Flatten the tile and append to the list + tiles.append(tile.flatten()) + + return tiles + + def get_runs(v1, v2): """Get the average run length/runs of each vector