寻求张量内存对齐至512倍数的高效填充逻辑

需要实现一种高效逻辑，对任意形状的张量进行填充，使其占用的总内存始终为512的倍数。例如：

形状为16x1x1x4的SI32类型张量（需乘以4计算总内存）：总元素数为16×4×1×1=64，总内存为64×4=256（非512的倍数），填充后形状应为32x1x1x4，总内存为512。

现有两种填充逻辑在部分张量形状（如16x51x1x4 SI32、80x240x1x1 U8）下失效，代码示例如下：

from functools import reduce

DATA_TYPE_MULTIPLYER = 2 # 会随数据类型动态变化，如U8为8、F16为16、SI32为32

ALIGNMENT = 512 # 固定常量
CHAR_BIT = 8    # 给定架构下固定常量

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)
    
    # 可对任意维度填充，当前暂用第一维度
    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示例。

现有方法存在的问题：

• ApproachOne为暴力法，通过逐次递增目标维度并检查内存是否对齐，虽可行但无法得到最小填充量，会过度膨胀张量；
• 最初仅针对第一维度填充，现需取消该限制，支持对任意维度进行填充，寻求更优的高效填充逻辑。

内容的提问来源于stack exchange，提问作者CMouse