具有位旋转技巧的快速整数矩阵乘法

Question

我在问是否有可能使用bitwise operations 改进显着整数矩阵乘法。矩阵很小，元素是小的非负整数（小意味着最多 20 个）。

为了让我们保持专注，让我们非常具体，假设我有两个 3x3 矩阵，整数条目 0

以下简单的 C++ 实现执行一百万次执行大约 1 秒，用 linux time 测量。

#include <random>

int main() {
//Random number generator
std::random_device rd;
std::mt19937 eng(rd());
std::uniform_int_distribution<> distr(0, 15);

int A[3][3];
int B[3][3];
int C[3][3];
for (int trials = 0; trials <= 1000000; trials++) {
    //Set up A[] and B[]
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 3; ++j) {
            A[i][j] = distr(eng);
            B[i][j] = distr(eng);
            C[i][j] = 0;
        }
    }
    //Compute C[]=A[]*B[]
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 3; ++j) {
            for (int k = 0; k < 3; ++k) {
                C[i][j] = C[i][j] + A[i][k] * B[k][j];
            }
        }
    }
}
return 0;
}

注意事项：

1. 矩阵不一定是稀疏的。
2. Strassen-like cmets 在这里没有帮助。
3. 我们尽量不要使用间接观察，在这个特定问题中矩阵A[]和B[]可以被编码为单个 64 位整数。想想稍微大一点的矩阵会发生什么。
4. 计算是单线程的。

   ---

相关：Binary matrix multiplication bit twiddling hack 和 What is the optimal algorithm for the game 2048?

Answer 1

您链接的问题是关于每个元素都是一个位的矩阵。对于一位值 a 和 b，a * b 完全等同于 a & b。

对于添加 2 位元素，使用 XOR（无进位加法）从头开始添加可能是合理的（并且比解包更快），然后使用 AND、移位和屏蔽跨元素边界的进位生成进位.

当添加进位产生另一个进位时，需要检测第 3 位。与使用 SIMD 相比，我认为即使模拟 3 位加法器或乘法器也不会是一种胜利。如果没有 SIMD（即在带有 uint64_t 的纯 C 中），它可能是有意义的。对于加法，您可以尝试使用普通加法，然后尝试撤消元素边界之间的进位，而不是通过 XOR/AND/shift 操作自己构建加法器。

---

打包与解包到字节的存储格式

如果您有很多这样的小矩阵，将它们以压缩形式（例如打包的 4 位元素）存储在内存中有助于缓存占用空间/内存带宽。 4bit 元素很容易解包，将每个元素放在向量的单独字节元素中。

否则，将它们存储为每个字节一个矩阵元素。从那里，如果需要，您可以轻松地将它们解压缩为每个元素 16 位或 32 位，具体取决于目标 SIMD 指令集提供的元素大小。您可以将一些矩阵以解包格式保存在局部变量中，以便在乘法中重复使用，但将它们打包回每个元素 4 位以存储在数组中。

---

编译器在 x86 的标量 C 代码中使用 uint8_t 来解决这个问题。请参阅@Richard 的答案中的 cmets：gcc 和 clang 都喜欢将 mul r8 用于 uint8_t，这迫使它们将数据移动到 eax（单操作数乘法的隐式输入/输出），而不是 @987654321 @。

uint8_t 版本的运行速度实际上比uint16_t 版本慢，尽管它的缓存占用空间只有一半。

---

您可能会从某种 SIMD 中获得最佳结果。

英特尔 SSSE3 有一个vector byte multiply, but only with adding of adjacent elements。使用它需要将您的矩阵解压缩成一个在行之间有一些零的向量或其他东西，因此您不会从一行中获取数据与另一行中的数据混合。幸运的是，pshufb 可以将元素归零并复制它们。

SSE2 PMADDWD 更有用，如果您将每个矩阵元素解压缩到单独的 16 位向量元素中。因此，给定一个向量中的一行，以及另一个向量中的转置列，pmaddwd (_mm_madd_epi16) 与为您提供C[i][j] 所需的点积结果相差一个水平add。

您可以将多个pmaddwd 结果打包到一个向量中，这样您就可以一次性存储C[i][0..2]，而不是单独进行这些添加。

Answer 2

如果您对大量矩阵执行此计算，您可能会发现减少数据大小可以显着提高性能：

#include <cstdint>
#include <cstdlib>

using T = std::uint_fast8_t;

void mpy(T A[3][3], T B[3][3], T C[3][3])
{
for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 3; ++j) {
            for (int k = 0; k < 3; ++k) {
                C[i][j] = C[i][j] + A[i][k] * B[k][j];
            }
        }
    }
}

奔腾可以在一条指令中移动和符号扩展一个 8 位值。这意味着每个缓存行可以获得 4 倍的矩阵。

更新：好奇心被激起，我写了一个测试：

#include <random>
#include <utility>
#include <algorithm>
#include <chrono>
#include <iostream>
#include <typeinfo>

template<class T>
struct matrix
{
    static constexpr std::size_t rows = 3;
    static constexpr std::size_t cols = 3;
    static constexpr std::size_t size() { return rows * cols; }

    template<class Engine, class U>
    matrix(Engine& engine, std::uniform_int_distribution<U>& dist)
    : matrix(std::make_index_sequence<size()>(), engine, dist)
    {}

    template<class U>
    matrix(std::initializer_list<U> li)
    : matrix(std::make_index_sequence<size()>(), li)
    {

    }

    matrix()
    : _data { 0 }
    {}

    const T* operator[](std::size_t i) const {
        return std::addressof(_data[i * cols]);
    }

    T* operator[](std::size_t i) {
        return std::addressof(_data[i * cols]);
    }

private:

    template<std::size_t...Is, class U, class Engine>
    matrix(std::index_sequence<Is...>, Engine& eng, std::uniform_int_distribution<U>& dist)
    : _data { (void(Is), dist(eng))... }
    {}

    template<std::size_t...Is, class U>
    matrix(std::index_sequence<Is...>, std::initializer_list<U> li)
    : _data { ((Is < li.size()) ? *(li.begin() + Is) : 0)... }
    {}


    T _data[rows * cols];
};

template<class T>
matrix<T> operator*(const matrix<T>& A, const matrix<T>& B)
{
    matrix<T> C;
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 3; ++j) {
            for (int k = 0; k < 3; ++k) {
                C[i][j] = C[i][j] + A[i][k] * B[k][j];
            }
        }
    }
    return C;
}

static constexpr std::size_t test_size = 1000000;
template<class T, class Engine>
void fill(std::vector<matrix<T>>& v, Engine& eng, std::uniform_int_distribution<T>& dist)
{
    v.clear();
    v.reserve(test_size);
    generate_n(std::back_inserter(v), test_size,
               [&] { return matrix<T>(eng, dist); });
}

template<class T>
void test(std::random_device& rd)
{
    std::mt19937 eng(rd());
    std::uniform_int_distribution<T> distr(0, 15);

    std::vector<matrix<T>> As, Bs, Cs;
    fill(As, eng, distr);
    fill(Bs, eng, distr);
    fill(Cs, eng, distr);

    auto start = std::chrono::high_resolution_clock::now();
    auto ia = As.cbegin();
    auto ib = Bs.cbegin();
    for (auto&m : Cs)
    {
        m = *ia++ * *ib++;
    }
    auto stop = std::chrono::high_resolution_clock::now();

    auto diff = stop - start;
    auto millis = std::chrono::duration_cast<std::chrono::microseconds>(diff).count();
    std::cout << "for type " << typeid(T).name() << " time is " << millis << "us" << std::endl;

}

int main() {
    //Random number generator
    std::random_device rd;
    test<std::uint64_t>(rd);
    test<std::uint32_t>(rd);
    test<std::uint16_t>(rd);
    test<std::uint8_t>(rd);
}

示例输出（最近的 macbook pro，64 位，使用 -O3 编译）

for type y time is 32787us
for type j time is 15323us
for type t time is 14347us
for type h time is 31550us

总结：

在这个平台上，int32 和 int16 被证明是一样快的。 int64 和 int8 同样慢（8 位结果让我吃惊）。

结论：

与以往一样，向编译器表达意图并让优化器完成它的工作。如果程序在生产环境中运行速度过慢，请进行测量并优化最严重的问题。