AVX2 整数混洗与字节以外的类型？

Question

用于洗牌的 api 仅支持

byte

和

sbyte

 //
 // Summary:
 //     __m256i _mm256_shuffle_epi8 (__m256i a, __m256i b)
 //
 //     VPSHUFB ymm, ymm, ymm/m256
 //
 // Parameters:
 //   value:
 //
 //   mask:
 public static Vector256<sbyte> Shuffle(Vector256<sbyte> value, Vector256<sbyte> mask);
 //
 // Summary:
 //     __m256i _mm256_shuffle_epi8 (__m256i a, __m256i b)
 //
 //     VPSHUFB ymm, ymm, ymm/m256
 //
 // Parameters:
 //   value:
 //
 //   mask:
 public static Vector256<byte> Shuffle(Vector256<byte> value, Vector256<byte> mask);

您会如何进行其他类型的洗牌？例如，假设我有一个

Vector256<short>

，并且想要使用类似

[0, 1, 7, 7, 3, 3, 2, 0]

之类的掩码进行随机播放？

我是否必须在字节级别进行操作？即，将上述掩码转换为其等效字节？

Answer 1

我是否必须在字节级别进行操作？即，将上述掩码转换为其等效字节？

对于

(u)short

的向量，通常是的（但它更复杂），除非你可以使用AVX512（对于

VPERMW

）或者索引成对排列，以便你可以将其作为

(u)int

的向量进行洗牌.

对于

(u)int

的向量，有

PermuteVar8x32

，反正一般来说更方便。

顺便说一句，

Vector256.Shuffle

确实有一个重载来洗牌短裤向量，但在我的测试中至少它只是调用一些后备方法，所以你可能不想依赖它。

一般来说，使用 AVX2 对 Shorts 向量进行洗牌比仅仅将其洗牌为字节向量要困难一些 - 对字节向量进行洗牌通常比调用

Avx2.Shuffle

更复杂，这确实是这里的问题。

Avx2.Shuffle

是解决方案的一部分，但

VPSHUFB

不会在 256 位向量的两个 128 位半部分之间移动字节。根据索引的外观，有多种解决方案，但一般来说，主要依靠混洗字节，并分别处理两个 128 位部分之间的移动。

例如，您可以创建一个具有两个下半部分数据副本的 256 位向量，另一个具有两个上半部分数据副本的 256 位向量，对每个向量进行混洗，然后根据无论您想要来自下部还是上部的字节。一般来说，您可以用它进行任何 32 字节的洗牌，并且可以在其上构建单词洗牌。

Answer 2

Avx2指令集还支持32位索引，并且提供了vpermd（_mm256_permutevar8x32_epi32，

Avx2.PermuteVar8x32

）指令。

Avx512系列指令集支持16~64位索引，并提供vpermw (_mm256_permutexvar_epi16,

Avx512BW.VL.PermuteVar16x16

)、vpermq (_mm256_permutexvar_epi64,

Avx512F.VL.PermuteVar4x64

)指令。

对于不支持Avx512指令集的情况，需要对索引进行转换。随后可以使用vpshufb来实现16~64的shuffle。源码如下

private static readonly Vector256<ushort> Shuffle_UInt16_Multiplier = Vector256.Create((ushort)0x202);
private static readonly Vector256<byte> Shuffle_UInt16_ByteOffset = Vector256.Create<byte>(0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1);

public static Vector256<ushort> Shuffle(Vector256<ushort> vector, Vector256<ushort> indices) {
    Vector256<ushort> mask, raw, rt;
    mask = Avx2.CompareEqual(Avx2.ShiftRightLogical(indices, 4), Vector256<ushort>.Zero); // Unsigned compare: (i < 16)
    raw = YShuffleKernel(vector, indices);
    rt = Avx2.And(raw, mask);
    return rt;
}

public static Vector256<ushort> YShuffleKernel(Vector256<ushort> vector, Vector256<ushort> indices) {
    Vector256<byte> indices2 = Avx2.Add(Multiply(indices, Shuffle_UInt16_Multiplier).AsByte(), Shuffle_UInt16_ByteOffset);
    return YShuffleKernel(vector.AsByte(), indices2).AsUInt16();
}

private static readonly Vector256<byte> Shuffle_Byte_LaneAdd_K0 = Vector256.Create(0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0);
private static readonly Vector256<byte> Shuffle_Byte_LaneAdd_K1 = Vector256.Create(0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70);

// Cross lane
public static Vector256<byte> YShuffleKernel(Vector256<byte> vector, Vector256<byte> indices) {
    // Format: Code; //Latency, Throughput(references IceLake)
    Vector256<byte> vector1 = Avx2.Permute4x64(vector.AsInt64(), (byte)0x4E).AsByte(); // 3,1. _MM_SHUFFLE(1, 0, 3, 2) = (1 << 6) | (0 << 4) | (3 << 2) | 2 = 0x4E = 78
    Vector256<byte> indices0 = Avx2.Add(indices, Shuffle_Byte_LaneAdd_K0); // 1,0.33
    Vector256<byte> indices1 = Avx2.Add(indices, Shuffle_Byte_LaneAdd_K1); // 1,0.33
    Vector256<byte> v0 = Avx2.Shuffle(vector, indices0); // 1,0.5
    Vector256<byte> v1 = Avx2.Shuffle(vector1, indices1); // 1,0.5
    Vector256<byte> rt = Avx2.Or(v0, v1); // 1,0.33
    return rt; //total latency: 8, total throughput CPI: 3
}

注意：

Avx2.Shuffle

在每个 128 位通道中进行随机播放。但 YShuffleKernel 可以跨车道并打乱整个向量。

为了方便使用，我开发了VectorTraits库，它集成了上述算法。其 Shuffle 方法支持 8-64 位整数索引，并在这些架构上具有硬件加速功能。

• X86：使用

  _mm256_shuffle_epi8

  和其他说明。
• 手臂：使用

  vqvtbl1q_u8

  说明。
• Wasm：使用

  i8x16.swizzle

  说明。

NuGet：https://www.nuget.org/packages/VectorTraits （披露：我是回购协议的所有者）