É difícil escrever um kernel que seja ótimo em todos os cenários de transposição ao mesmo tempo. Vários cenários comuns de transposição podem ser resumidos e otimizados de maneira direcionada. Aqui estão apenas alguns casos de transposição de transformação do eixo, sem considerar o tamanho da forma. Portanto, nesses cenários, a otimização direcionada do tamanho da forma do eixo transposto também deve ser considerada.
Cena 1: lote 2D, perm: 021
Para tensores 2D, 3D ou de dimensão superior, troque as duas dimensões mais internas. Tudo isso pode ser unificado em lote 2D. A transposição de uma matriz 2D é equivalente a lote=1, e um tensor maior que três dimensões pode combinar todas as dimensões que não sejam as duas dimensões mais internas em uma dimensão.
Um método de transposição de uma grande matriz 021 em uma pequena matriz:
Por exemplo, a entrada é uma grande matriz de [4096, 4096] para transposição, mas cada thread lê a posição adjacente da mesma linha de entrada, mas escreve a posição adjacente da mesma coluna ao escrever e, para escrever, o mesmo spans de coluna Se o comprimento da linha for muito longo, não é amigável para o hit do cache de write-back.Você pode considerar a redução do comprimento da linha do span para otimizar o desempenho.
Considere dividir uma grande transposição em várias menores:
Por exemplo, transponha MN para NM, considerando que tanto M quanto N podem ser divididos em dimensões menores: M é dividido em M1M0 e N é dividido em N1N0
Então o problema tem MN transposto para NM torna-se M1M0N1N0 transposto para N1N0M1M0
Assumindo que apenas dois eixos podem ser trocados por vez, isso pode ser alcançado através das três etapas a seguir:
M1M0N1N0->N1M1M0N0->N1M1N0M0->N1N0M1M0
A primeira etapa e a terceira parte transpõem cada uma um tensor (como a cena de transposição 0213 abaixo), que geralmente pode ser alcançada com muita eficiência.
Na segunda etapa, a dimensão da transposição de elemento único é reduzida de MN para M0N0, o que é mais propício à utilização do cache. Por exemplo, [4096,4096] pode ser dividido em [64,64,64,64] para transposição.
A exibição visual desse método é a seguinte: é equivalente a dividir uma matriz grande em matrizes pequenas, transpor cada matriz pequena independentemente e, em seguida, transpor a matriz pequena como um todo.
Observe que este esquema mostra M1M0N1N0->M1N1M0N0->M1N1N0M0->N1M1N0M0, que está a um passo do objetivo final.
Cena 2: 0213
Sua característica é que as duas dimensões adjacentes internas interagem e não contêm uma ou mais dimensões internas. Como acima, as dimensões adjacentes que não são trocadas podem ser combinadas como um todo, e a dimensão mais externa pode ser complementada por 1 se for insuficiente.
Nesta cena perm=[2, 0, 3, 1, 4], parece que vários eixos são transpostos ao mesmo tempo, mas devido à particularidade do elemento de forma 1, ele pode ser espremido, para que possa ser convertido para [784, 3, 4, 12] para [3, 4, 784, 12] transpor, você pode usar o método 0213 para resolver.
Exclua o algoritmo cuja forma de transposição é 1
perm = [2, 0, 3, 1, 4]
in_shape = [1, 784, 1, 4, 12]
rm_axes = []
for idx, elem in enumerate(in_shape):
if elem == 1:
rm_axes.append(idx)
print("rm_axes:", rm_axes)
def remove_axis(in_shape, perm, rm_axis):
del in_shape[rm_axis]
perm_rm_idx = -1
for idx, elem in enumerate(perm):
if elem == rm_axis:
perm_rm_idx = idx
if elem > rm_axis:
perm[idx] = perm[idx]-1
del perm[perm_rm_idx]
for rm_axis in reversed(rm_axes):
remove_axis(in_shape, perm, rm_axis)
print("perm:", perm)
print("in_shape:", in_shape)Cenário 3: Troque dois eixos adjacentes, mas a forma correspondente a um dos eixos é 1
Esta cena não requer transposição, apenas remodelação é necessária.
Depois de usar o algoritmo acima para excluir a forma de transposição como 1, o perm desta transposição se tornará [0,1,2,3,...] É muito simples julgar que esta transposição não precisa executar nenhuma operação, e excluí-lo diretamente É isso.
Cenário 4: Troque vários eixos, mas algumas permissões são adjacentes
Aqui perm=[1, 2, 0], parece que 3 eixos são trocados, mas na verdade os dois 1x64 são trocados juntos e podem ser mesclados em uma dimensão. Esse problema se torna o cenário 1 acima. Portanto, uma solução poderia ser mesclar eixos adjacentes que são transformados juntos, simplificando assim o problema.
Cena 5: Outro
Claro, ainda há um pequeno número de cenários que não podem ser resolvidos usando o método acima, por exemplo, o caso em que a primeira dimensão da forma de entrada não é 1.