如何實(shí)現(xiàn)比PyTorch快6倍的Permute/Transpose算子？

x?=?flow.randn(2,?3)
y?=?x.permute(1,?0)
y.shape?
(3,?2)

for?row?in?x.shape[0]:?
??for?col?in?x.shape[1]:?
????y[row][col]?=?x[col][row]

template<size_t?num_dims,?size_t?movement_size,?typename?IndexType>
__global__?void?PermuteKernel(PermuteKernelParams?params)?{
??using?T?=?typename?std::aligned_storage::type;
??const?T*?src?=?reinterpret_cast<const?T*>(params.src);
??T*?dst?=?reinterpret_cast(params.dst);
??IndexType?src_index[num_dims];
??IndexType?dst_index[num_dims];
??CUDA_1D_KERNEL_LOOP_T(IndexType,?i,?params.count)?{
????params.dst_index_helper.OffsetToNdIndex(i,?dst_index);
#pragma?unroll
????for?(size_t?dim?=?0;?dim???????src_index[params.permutation[dim]]?=?dst_index[dim];
????}
????IndexType?src_offset?=?params.src_index_helper.NdIndexToOffset(src_index);
????dst[i]?=?src[src_offset];
??}
}

#?0,?1,?2,?3)?->?(2,?3,?0,?1)
x?=?flow.randn(3,?4,?5,?6)
y?=?x.permute(2,?3,?0,?1)
y.shape?
(5,?6,?3,?4)

#?(0,?1,?2,?3)?->?((2,?3),?(0,?1))
x?=?x.reshape(x.shape[0]*x.shape[1],?x.shape[2]*x.shape[3])
y?=?x.permute(1,?0)
y?=?y.reshape(x.shape[2],?x.shape[3],?x.shape[0],?x.shape[1])

(0,?1)?->?(1,?0)
(0,?1,?2)?->?(0,?2,?1)

template<size_t?num_dims,?size_t?movement_size,?size_t?tile_size,?typename?IndexType>
__global__?void?BatchTransposeKernel(const?void*?src_ptr,?void*?dst_ptr,?IndexType?H,?IndexType?W,
?????????????????????????????????????IndexType?num_tile_rows,?IndexType?num_tile_cols,
?????????????????????????????????????int32_t?block_nums)?{
??using?T?=?typename?std::aligned_storage::type;
??__shared__?T?tile[tile_size][tile_size?+?1];??//?To?avoid?bank?conflict.

??const?T*?src?=?reinterpret_cast<const?T*>(src_ptr);
??T*?dst?=?reinterpret_cast(dst_ptr);

??IndexType?batch_num_tile?=?num_tile_rows?*?num_tile_cols;
??for?(int?i?=?blockIdx.x,?step?=?gridDim.x;?i?????const?IndexType?batch_index?=?i?/?batch_num_tile;??//?the?index?of?batch.
????const?IndexType?flatten_index?=
????????i?-?batch_index?*?batch_num_tile;??
????const?IndexType?row_index?=?flatten_index?/?num_tile_cols;??//?the?row?index?of?tile?in?a?batch.
????const?IndexType?col_index?=
????????flatten_index
????????-?row_index
??????????????*?num_tile_cols;??//?the?col?index?of?tile?in?a?batch.
????const?IndexType?offset?=?batch_index?*?H?*?W;
????IndexType?x?=?col_index?*?tile_size?+?threadIdx.x;
????IndexType?y?=?row_index?*?tile_size?+?threadIdx.y;
????if?(x???????IndexType?y_range?=
??????????((tile_size?-?threadIdx.y)?#pragma?unroll
??????for?(int?i?=?0;?i?????????tile[threadIdx.y?+?i][threadIdx.x]?=?src[offset?+?(y?+?i)?*?W?+?x];
??????}
????}
????__syncthreads();
????x?=?row_index?*?tile_size?+?threadIdx.x;
????y?=?col_index?*?tile_size?+?threadIdx.y;
????if?(x???????IndexType?x_range?=
??????????((tile_size?-?threadIdx.y)?#pragma?unroll
??????//?`i?
??????for?(int?i?=?0;?i?????????dst[offset?+?(y?+?i)?*?H?+?x]?=?tile[threadIdx.x][threadIdx.y?+?i];
??????}
????}
????__syncthreads();
??}
}

IndexType?y_range?=?((tile_size?-?threadIdx.y)?#pragma?unroll
for?(int?i?=?0;?i???...
}

template<size_t?num_dims,?size_t?tile_size,?typename?IndexType>
__global__?void?BatchTransposeMovement2Kernel(const?void*?src_ptr,?void*?dst_ptr,?IndexType?rows,
??????????????????????????????????????????????IndexType?cols,?IndexType?num_tile_rows,
??????????????????????????????????????????????IndexType?num_tile_cols,?int32_t?block_nums)?{
??static_assert(tile_size?%?2?==?0);
??using?T_MOV2?=?typename?std::aligned_storage<2,?2>::type;
??using?T_MOV4?=?typename?std::aligned_storage<4,?4>::type;

??const?T_MOV4*?src?=?reinterpret_cast<const?T_MOV4*>(src_ptr);
??T_MOV4*?dst?=?reinterpret_cast(dst_ptr);

??//?Use?union?structure?to?process?Load?and?Store.
??__shared__?union?{
????T_MOV2?tile_m2[tile_size][tile_size?+?2];??????//?half?[64][66]
????T_MOV4?tile_m4[tile_size][tile_size?/?2?+?1];??//?half2?[64][33]
??}?tile_mem;

??IndexType?batch_num_tile?=?num_tile_rows?*?num_tile_cols;
??for?(int?i?=?blockIdx.x,?step?=?gridDim.x;?i?????const?IndexType?batch_index?=?i?/?batch_num_tile;??//?the?index?of?batch.
????const?IndexType?flatten_index?=
????????i?-?batch_index?*?batch_num_tile;??//?the?flatten?index?of?tile?in?a?batch.

????const?IndexType?row_index?=?flatten_index?/?num_tile_cols;??//?the?row?index?of?tile?in?a?batch.
????const?IndexType?col_index?=
????????flatten_index
????????-?row_index
??????????????*?num_tile_cols;??//?equal?to?k?%?num_tile_cols.?the?col?index?of?tile?in?a?batch.
????const?IndexType?offset?=?batch_index?*?rows?*?cols;
????IndexType?x?=
????????col_index?*?tile_size?+?threadIdx.x?*?2;??//?cause?each?thread?process?a?half2?element,?we?need?to?multiply?2?for?threadIdx.x.
????IndexType?y?=?row_index?*?tile_size?+?threadIdx.y;
????if?(x???????//?each?thread?process?4?elements.
??????IndexType?y_range?=
??????????((tile_size?-?threadIdx.y)?#pragma?unroll
??????//?`i?
??????for?(int?i?=?0;?i?????????//?each?thread?load?a?half2.
????????tile_mem.tile_m4[threadIdx.y?+?i][threadIdx.x]?=?src[(offset?+?(y?+?i)?*?cols?+?x)?/?2];
??????}
????}
????__syncthreads();
????x?=?row_index?*?tile_size?+?threadIdx.x?*?2;??//?cause?each?thread?process?a?half2?element,?we?need?to?multiply?2?for?threadIdx.x.
????y?=?col_index?*?tile_size?+?threadIdx.y;
????if?(x???????IndexType?x_range?=
??????????((tile_size?-?threadIdx.y)?#pragma?unroll
??????//?`i?
??????for?(int?i?=?0;?i?????????/*
????????When?write?back?as?column,?it?cannot?be?stored?as?half2?directly.
????????So?we?split?as?2?half?elements,?and?write?back?separately.
????????*/
????????union?{
??????????T_MOV4?m4;
??????????T_MOV2?m2[2];
????????}?tmp_storage;
????????tmp_storage.m2[0]?=?tile_mem.tile_m2[threadIdx.x?*?2][threadIdx.y?+?i];
????????tmp_storage.m2[1]?=?tile_mem.tile_m2[threadIdx.x?*?2?+?1][threadIdx.y?+?i];
????????dst[(offset?+?(y?+?i)?*?rows?+?x)?/?2]?=?tmp_storage.m4;
??????}
????}
????__syncthreads();
??}
}