填充张量的高效逻辑

Question

我正在尝试填充某种形状的张量，使得张量使用的总内存始终是 512 的倍数，例如 SI32 类型的张量形状 16x1x1x4（乘以 4 以获得总大小）

The total elements are 16x4x1x1 = 64
Total Memory required 64x**4** = 256 (Not multiple of 512)
Padded shape would be 32x1x1x4 = 512

下面的逻辑适用于基本形状，但会破坏形状，例如16x51x1x4 SI32或随机的东西80x240x1x1 U8 填充逻辑如下所示

from functools import reduce

DATA_TYPE_MULTIPLYER = 2 # This would change at runtime with different type e.g. 8 with U8 16 with F16 32 with SI32

ALIGNMENT = 512 #Always Constant
CHAR_BIT = 8    # Always Const for given fixed Arch

def approachOne(tensor):
    totalElements = reduce((lambda x, y: x * y), tensor)
    totalMemory = totalElements * DATA_TYPE_MULTIPLYER
    
    divisor = tensor[1] * tensor[2] * tensor[3]
    tempDimToPad = totalElements/divisor
    orgDimToPad = totalElements/divisor
    while (True):
        if ((tempDimToPad * divisor * DATA_TYPE_MULTIPLYER) % ALIGNMENT == 0):
            return int(tempDimToPad - orgDimToPad)
        tempDimToPad = tempDimToPad + 1;
    
def getPadding(tensor):
    totalElements = reduce((lambda x, y: x * y), tensor)
    totalMemory = totalElements * DATA_TYPE_MULTIPLYER
    newSize = totalMemory + (ALIGNMENT - (totalMemory % ALIGNMENT))
    newTotalElements = (newSize * CHAR_BIT) / (CHAR_BIT * DATA_TYPE_MULTIPLYER)
    
    # Any DIM can be padded, using first for now
    paddingValue = tensor[0] 
    padding =  int(((newTotalElements * paddingValue) / totalElements) - paddingValue)
    return padding
    
tensor = [11, 7, 3, 5]
print(getPadding(tensor))
print(approachOne(tensor))

tensorflow包可能会在这里有所帮助，但我最初是用 C++ 编码的，所以只需在 python 中发布一个最小的工作示例任何帮助都是值得赞赏的，谢谢

方法 1： 暴力方法是在任何选定的维度上继续递增 1，并检查总内存是否是 512 的倍数。暴力方法可以工作，但不会给出最小填充并使张量膨胀

更新条件 最初的方法是填充第一个暗淡的区域。由于总是填充第一个维度我不是最好的解决方案，只需摆脱这个约束

Answer 1

如果您希望总内存是 的倍数，512那么张量中的元素数量必须是 的倍数512 // DATA_TYPE_MULTIPLIER，例如128在您的情况下。无论该数字是什么，它都将具有以下形式的质因数分解2**n。张量中元素的数量由下式给出，s[0]*s[1]*...*s[d-1]其中s是包含张量形状的序列，并且d是整数，即维数。该乘积还具有一些素因数分解，并且当且仅当它包含这些素因数时，s[0]*s[1]*...*s[d-1]它才是倍数。2**n即，任务是填充各个维度，s[i]使得乘积的质因数分解结果s[0]*s[1]*...*s[d-1]包含2**n。

如果目标是达到填充张量的最小可能大小，则可以简单地迭代给定目标元素数量的所有倍数，以找到第一个可以通过填充（增加）张量的各个维度来满足的元素⁽¹⁾ . 只要维度包含至少一个未包含在目标倍数大小中的素因子，就必须增加该维度。当所有维度都增加到它们的主要因子包含在目标倍数尺寸中后，我们可以检查候选形状的最终尺寸：如果它与目标倍数尺寸匹配，我们就完成了；否则，我们就完成了。如果它的素因子是目标多个素因子的严格子集，我们可以将缺失的素因子添加到任何维度（例如第一个）；否则，我们可以使用多余的质因数来存储未来（更大）乘数的候选形状。第一个这样的未来乘数然后标记所有可能乘数的迭代的上限，即算法将终止。但是，如果候选形状（调整所有尺寸后）对于目标倍数尺寸有过多的质因数，并且错过了一些其他质因数，则唯一的方法是迭代所有可能的填充形状，其大小受目标限制多个尺寸。

下面是一个示例实现：

from collections import Counter
import itertools as it
import math
from typing import Iterator, Sequence


def pad(shape: Sequence[int], target: int) -> tuple[int,...]:
    """Pad the given `shape` such that the total number of elements
       is a multiple of the given `target`.
    """
    size = math.prod(shape)
    if size % target == 0:
        return tuple(shape)

    target_prime_factors = get_prime_factors(target)

    solutions: dict[int, tuple[int,...]] = {}  # maps `target` multipliers to corresponding padded shapes

    for multiplier in it.count(math.ceil(size / target)):

        if multiplier in solutions:
            return solutions[multiplier]

        prime_factors = [*get_prime_factors(multiplier), *target_prime_factors]
        
        def good(x):
            return all(f in prime_factors for f in get_prime_factors(x))

        candidate = list(shape)
        for i, x in enumerate(candidate):
            while not good(x):
                x += 1
            candidate[i] = x

        if math.prod(candidate) == multiplier*target:
            return tuple(candidate)

        candidate_prime_factor_counts = Counter(f for x in candidate for f in get_prime_factors(x))
        target_prime_factor_counts = Counter(prime_factors)

        missing = target_prime_factor_counts - candidate_prime_factor_counts
        excess = candidate_prime_factor_counts - target_prime_factor_counts

        if not excess:
            return (
                candidate[0] * math.prod(k**v for k, v in missing.items()),
                *candidate[1:],
            )
        elif not missing:
            solutions[multiplier * math.prod(k**v for k, v in excess.items())] = tuple(candidate)
        else:
            for padded_shape in generate_all_padded_shapes(shape, bound=multiplier*target):
                padded_size = math.prod(padded_shape)
                if padded_size == multiplier*target:
                    return padded_shape
                elif padded_size % target == 0:
                    solutions[padded_size // target] = padded_shape


def generate_all_padded_shapes(shape: Sequence[int], *, bound: int) -> Iterator[tuple[int,...]]:
    head, *tail = shape
    if bound % head == 0:
        max_value = bound // math.prod(tail)
    else:
        max_value = math.floor(bound / math.prod(tail))
    for x in range(head, max_value+1):
        if tail:
            yield from ((x, *other) for other in generate_all_padded_shapes(tail, bound=math.floor(bound/x)))
        else:
            yield (x,)


def get_prime_factors(n: int) -> list[int]:
    """From: https://stackoverflow.com/a/16996439/3767239
       Replace with your favorite prime factorization method.
    """
    primfac = []
    d = 2
    while d*d <= n:
        while (n % d) == 0:
            primfac.append(d)  # supposing you want multiple factors repeated
            n //= d
        d += 1
    if n > 1:
       primfac.append(n)
    return primfac

这里有一些例子：

pad((16, 1, 1), 128) = (128, 1, 1)
pad((16, 51, 1, 4), 128) = (16, 52, 1, 4)
pad((80, 240, 1, 1), 128) = (80, 240, 1, 1)
pad((3, 5, 7, 11), 128) = (3, 5, 8, 16)
pad((3, 3, 3, 1), 128) = (8, 4, 4, 1)
pad((7, 7, 7, 7), 128) = (7, 8, 8, 8)
pad((9, 9, 9, 9), 128) = (10, 10, 10, 16)

_{脚注： (1) 事实上，我们需要在整数域上求多项式的根。但是，我不知道有任何算法可以解决这个问题。(s[0]+x[0])*(s[1]+x[1])*...*(s[d-1]+x[d-1]) - multiple*targetx[i] >= 0}