asyms
commented
2 weeks ago Proposed change (will add this later):
def limit_unrolling_to_mem_capacity(self, mapping: SpatialMapping) -> SpatialMapping:
"""! Scale the given unroll factors such that they do not exceed the capacity of the memory structure"""
def limit_loop_unrolling(
spatial_mapping: SpatialMapping, dims_to_limit: set[LayerDim], max_unrolling: float
) -> SpatialMapping:
def adjust_unrolling_factors(factors: list[UnrollFactor], max_unrolling: float) -> list[UnrollFactor]:
product = math.prod(factors)
while product > max_unrolling:
max_factor = max(factors)
max_index = factors.index(max_factor)
factors[max_index] -= 1
product = math.prod(factors)
return factors
# Extract the unrolling factors for the limited dimensions
unrolling_factors: dict[tuple[OADimension, LayerDim], UnrollFactor] = {}
for dim in dims_to_limit:
for oa_dim in spatial_mapping:
loop_unrollings = spatial_mapping[oa_dim]
if dim in loop_unrollings:
unrolling_factors[(oa_dim, dim)] = loop_unrollings[dim]
# Adjust the unrolling factors
adjusted_factors = adjust_unrolling_factors(list(unrolling_factors.values()), max_unrolling)
adjusted_unrolling_factors = {key: adjusted_factors[i] for i, key in enumerate(unrolling_factors)}
# Generate the limited SpatialMapping
limited_spatial_mapping = SpatialMapping({})
for oa_dim in spatial_mapping:
loop_unrollings = spatial_mapping[oa_dim]
mapping_single_oa_dim_dict = {}
for loop_dim, factor in loop_unrollings.items():
mapping_single_oa_dim_dict[loop_dim] = adjusted_unrolling_factors.get((oa_dim, loop_dim), factor)
limited_mapping_single_oa_dim = MappingSingleOADim(mapping_single_oa_dim_dict)
limited_spatial_mapping[oa_dim] = limited_mapping_single_oa_dim
return limited_spatial_mapping
There is a bug in the SpatialMappingGeneration stage when limiting a UserSpatialMapping to limited memory capacity. Currently, a dictionary called 'factor_map' is used to map the old unrolling factors to the new (limited) factors: https://github.com/KULeuven-MICAS/zigzag/blob/920ee6957b030461494ca5ed475205721a5b6ad8/zigzag/stages/mapping/spatial_mapping_generation.py#L248
If e.g. two old unrollings are of the same size, this dictionary will only contain one entry, incorrectly using the second new unrolling for both old ones. More detailed: if unrolling_factors = [4, 4] and adjusted_factors = [2, 3], this dict will be: {4: 3}, incorrectly limiting the 4 factor to 3 for both dimensions instead of 2 and 3.