使用模板技巧加快矩阵运算(宝宝辅食)
问题引入
模板是C++的一大精髓,本节内容以Eigen中的一个最简单的技巧来展示模板能为我们做的一些有趣的事。
本节内容不涉及:向量化指令
要看懂本节内容你需要
- C++基础
- 对模板有一定理解(非类型模板参数,模板偏特化等)
并且本节只考虑矩阵/向量加法, 矩阵的乘法显然可以如法炮制,不过会更复杂一些。
我们先给出一个矩阵(或向量)类及其加法的平凡实现,这是任何人都绝对会的,为了减少模板参数带来的混乱,我就不将Scalar写入模板参数了,你可以在类外部使用using Scalar = int;
或任何你喜欢的基本类型
template<long long M, long long N>
class Matrix {
public:
static constexpr long long Rows = M;
static constexpr long long Cols = N;
static constexpr long long size = Rows * Cols;
Scalar* data;
Matrix() {
data = new Scalar[Rows * Cols];
}
Matrix(const Matrix& other) {
data = new Scalar[Rows * Cols];
memcpy(data, other.data, Rows * Cols * sizeof(Scalar));
}
Matrix(Matrix&& other) noexcept {
data = other.data;
other.data = nullptr;
}
~Matrix() {
if (data != nullptr) {
delete[] data;
}
}
Scalar& coeff(long long i) {
return data[i];
}
Scalar coeff(long long i) const {
return data[i];
}
Scalar& operator()(long long i, long long j) {
return data[i * Cols + j];
}
Scalar operator()(long long i, long long j) const {
return data[i * Cols + j];
}
Matrix& operator=(const Matrix& other) {
if ((void*)&other == (void*)this) {
return *this;
}
for (long long i = 0; i < size; ++i) {
data[i] = other.data[i];
}
return *this;
}
Matrix operator+(const Matrix& other) {
Matrix res;
for (long long i = 0; i < size; ++i) {
res.data[i] = data[i] + other.data[i];
}
return res;
}
};
让我们从一个表达式开始:
m
3
=
m
1
+
m
2
+
m
3
w
h
e
r
e
:
m
n
∈
R
N
∗
N
m_3 = m_1 + m_2 + m_3 \\ where: m_n \in R^{N*N}
m3=m1+m2+m3where:mn∈RN∗N
在使用上述的实现计算这个表达式的过程可以表示为:
t
e
m
p
12
=
m
1
+
m
2
;
t
e
m
p
23
=
t
e
m
p
12
+
m
3
m
3
=
t
e
m
p
23
temp_{12} = m_1 + m_2; \\ temp_{23} = temp_{12} + m_3 \\ m_3 = temp_{23}
temp12=m1+m2;temp23=temp12+m3m3=temp23
问题显而易见,我们创建了从 t e m p 12 , t e m p 23 temp_{12}, temp_{23} temp12,temp23这2个临时变量来存储中间结果,并且我们对临时变量要进行额外的访问内存,而理论上来说,我们只需要遍历每个矩阵一次,总共 4 N N 4 NN 4NN次的访存就可以来完成这件事,可实际上一共执行了 8 N N 8NN 8NN次
所以我们理想中的情况其实是:
m3 = m1 + m2 + m3;
转化成下列代码:
for(int i = 0; i < size; ++ i) {
m3[i] = m1[i] + m2[i] + m3[i];
}
对每个位置进行4次访存重复 s i z e = N 2 size = N^2 size=N2次
如果这个表达式如此简短和简单, 那我们当然可以手动把上面展开成下面的表达式结果, 可是如果表达式再复杂一点呢?
m7 = m1 + m2 + m3 + m4 + m5 + ... + mn;
甚至加上乘法运算(不过不在本节讨论范围内)
m7 = m1 * m2 + m3 * m4 + m5 + m6;
解决方法
如果你知道计算图这个概念, 那理解下面的内容应该很简单, 计算图无处不在(并行算法, 深度学习等等), 在深度学习中的自动微分框架通常是基于计算图实现的。
我们先自顶向下分析一下这个表达式m3 = m1 + m2 + m3;
- 站在等号左侧的视角上,我们把
m1 + m2 + m3
看作一个表达式整体:m3 = Expression0;
- 对于等号右侧, 由于表达式一般情况下是从左往右解析的,先把
m1 + m2
看作整体:Expression1 + m3
;
我们自底向上再来看, 先看表达式Expression1
, 这个表达式代表两个矩阵/向量的加和。
表达式Expression0
是表达式Expression1
和m3的加和, 而总体的表达式是Expression0
对m3的赋值。
如果我们能够设计一些接口,这些接口能够把具体参与运算的表达式抽象出来,使用同样的形式表达各自的行为,而前端无需过分考虑具体的组成这个表达式的各个项的类型以及他们是如何组合的,不需要像前面那样手动地展开式子,那么我们就可以使用简单的表达式就可以隐式地对表达式进行线性展开,而不是创建中间变量。
我们需要有一个Evaluate<>
模板类,这个类的模板参数代表了这个表达式中的运算性质是怎么样的。 这个类统一提供一个接口coeff(long long i)
用来访问这个表达式所返回的结果在位置i
上的参数是什么。
template<typename ExpressT> class Evaluator;
我们需要抽象出3个操作:
-
将单个矩阵/向量抽象为一种表达式
-
将矩阵/向量/其他表达式之间的加法运算抽象为一种表达式
-
右侧表达式向左侧表达式赋值的操作
先介绍后两个操作的实现方式:
我们先需要实现一个控制加法行为的Sum
的类。
template<typename Express1, typename Express2>
class Sum {
public:
Express1 const& m1;
Express2 const& m2;
Sum(Express1 const& e1, Express2 const& e2) : m1(e1), m2(e2) {}
};
这个类实际上什么也不做,仅仅保存了加法两端的参数的常引用,但这个类型本身携带了当前处理的表达式信息, 我们可以通过模板偏特化的方式控制当前的生成的Evaluator类会选择我们为Evaluatir的特别实现, 如果你不是特明白,请看类模板的偏特化内容。
template<typename type_of_A, typename type_of_B>
class Evaluator <Sum<type_of_A, type_of_B>> {
public:
type_of_A const& evalA;
type_of_B const& evalB;
Evaluator(Sum<type_of_A, type_of_B>const & s) : evalA(s.m1), evalB(s.m2) {};
inline Scalar coeff(long long i) const { return evalA.coeff(i) + evalB.coeff(i); }
};
这个类的实现也很简单, 他保存表达式两侧的引用, 调用coeff
函数时返回表达式两侧参数(即evalA
,evalB
)在该位置的加和。
然后我们实现Assign类来控制赋值操作
template<typename Dest, typename Source>
class Assign {
public:
Dest& dst;
Source const& src;
Assign(Dest& d, Source const& s) : dst(d), src(s) {}
// Add a run method to perform the assignment
void run() {
for (long long i = 0; i < Dest::size; ++i) {
dst.coeff(i) = src.coeff(i);
}
}
};
也不需要多解释了, 保存目的操作数的引用和源操作数的常引用, 在run函数中逐位复制。
那么对于第一个要实现的操作:
template<typename Express>
class Evaluator {
public:
Express& express;
Evaluator(Express& e) : express(e) {}
inline Scalar& coeff(long long i) { return express.coeff(i); }
inline Scalar coeff(long long i) const { return express.coeff(i); }
};
很简单,我们保存任意表达式内容本身的引用(如果Express = Sum<. , .>
,编译器不会选择这个类实现),注意,Evaluator本身可能是常量或者常引用或者对临时对象的引用(Assign
类中就保存了一个表达式的常引用),也有可能是对非常量左值的引用(Assign
类中我们需要通过coeff
函数修改目的操作数的值)所以有必要为coeff
函数写这样的实现。
然后我们在类外定义operator+
:
template<typename Left, typename Right>
Evaluator<Sum<Left, Right>> operator+(const Left& left, const Right& right) {
return Evaluator<Sum<Left, Right>>(Sum<Left, Right>(left, right));
}
接受两端操作数的常引用,包装成一个Sum表达式。
然后就非常简单了,新的Matrix类的其他成员函数没有什么特别的变化。唯一要改的是operator=
template<typename Other>
Matrix& operator=(const Other& other) {
if ((void*)&other == (void*)this) {
return *this;
}
Assign<Matrix, Other> assign(*this, other);
assign.run();
return *this;
}
先判断自赋值, 构造Assign对象包装operator+
两端的操作数然后运行,返回自身的引用。
于是,表达式m3 = m1 + m2 + m3
被展开成了:
for(i = 0; i < size; ++ i) {
m3[i] = “Evaluator(m1+m2+m3)”.coeff(i);
// or Evaluate(Sum(Evaluate(Sum(m1, m2)), m3)).coeff(i);
}
Evaluator(m1 + m2 + m3)会生成这个类:
class Evaluator<Sum<Sum<Matrix,Matrix>, Matrix>>{
Evaluator<Sum<Matrix,Matrix>> evalA(“m1+m2”);
Evaluator<Matrix> evalB(“m3”);
Scalar coeff(i) {
return evalA.coeff(i) + evalB.coeff(i);
}
};
其中evalA
:
class Evaluator< Sum<Matrix,Matrix> > {
Evaluator<Matrix> evalA(“m1”);
Evaluator<Matrix> evalB(“m2”);
Scalar coeff(i) {
return evalA.coeff(i) + evalB.coeff(i);
}
};
总的来说,我们只进行了一次循环, 没有中间变量, 并且减少了一半的内存访问次数。实际上,这一技术被称为延迟计算。
下面是包含速度测试的完整代码
完整代码:
#include <iostream>
#include <chrono>
#include <random>
#include <algorithm>
#include <Eigen/core>
using Scalar = int;
constexpr int N = 4096 * 2;
using namespace std;
namespace naive_impl {
template<long long M, long long N>
class Matrix {
public:
static constexpr long long Rows = M;
static constexpr long long Cols = N;
static constexpr long long size = Rows * Cols;
Scalar* data;
Matrix() {
data = new Scalar[Rows * Cols];
}
Matrix(const Matrix& other) {
data = new Scalar[Rows * Cols];
memcpy(data, other.data, Rows * Cols * sizeof(Scalar));
}
Matrix(Matrix&& other) noexcept {
data = other.data;
other.data = nullptr;
}
~Matrix() {
if (data != nullptr) {
delete[] data;
}
}
Scalar& coeff(long long i) {
return data[i];
}
Scalar coeff(long long i) const {
return data[i];
}
Scalar& operator()(long long i, long long j) {
return data[i * Cols + j];
}
Scalar operator()(long long i, long long j) const {
return data[i * Cols + j];
}
Matrix& operator=(const Matrix& other) {
if ((void*)&other == (void*)this) {
return *this;
}
for (long long i = 0; i < size; ++i) {
data[i] = other.data[i];
}
return *this;
}
Matrix& operator=(Matrix&& other) noexcept {
if ((void*)&other == (void*)this) {
return *this;
}
if (data != nullptr) {
delete[] data;
}
data = other.data;
other.data = nullptr;
return *this;
}
Matrix operator+(const Matrix& other) {
Matrix res;
for (long long i = 0; i < size; ++i) {
res.data[i] = data[i] + other.data[i];
}
return res;
}
};
Matrix<N, N> getOne() {
Matrix<N, N> m;
for (int i = 0; i < Matrix<N, N>::size; i++) {
m.coeff(i) = rand();
}
return m;
}
}
namespace new_impl {
template<typename Dest, typename Source>
class Assign {
public:
Dest& dst;
Source const& src;
Assign(Dest& d, Source const& s) : dst(d), src(s) {}
// Add a run method to perform the assignment
void run() {
for (long long i = 0; i < Dest::size; ++i) {
dst.coeff(i) = src.coeff(i);
}
}
};
template<typename Express1, typename Express2>
class Sum {
public:
Express1 const& m1;
Express2 const& m2;
Sum(Express1 const& e1, Express2 const& e2) : m1(e1), m2(e2) {}
};
template<typename ExpressT> class Evaluator;
template<typename type_of_A, typename type_of_B>
class Evaluator <Sum<type_of_A, type_of_B>> {
public:
type_of_A const& evalA;
type_of_B const& evalB;
Evaluator(Sum<type_of_A, type_of_B>const & s) : evalA(s.m1), evalB(s.m2) {};
inline Scalar coeff(long long i) const { return evalA.coeff(i) + evalB.coeff(i); }
};
template<typename Express>
class Evaluator {
public:
Express& express;
Evaluator(Express& e) : express(e) {}
inline Scalar& coeff(long long i) { return express.coeff(i); }
inline Scalar coeff(long long i) const { return express.coeff(i); }
};
template<long long M, long long N>
class Matrix {
public:
static constexpr long long Rows = M;
static constexpr long long Cols = N;
static constexpr long long size = Rows * Cols;
Scalar* data;
public:
Matrix() {
data = new Scalar[Rows * Cols];
}
Matrix(const Matrix& other) {
data = new Scalar[Rows * Cols];
std::copy(other.data, other.data + Rows * Cols, data);
}
Matrix(Matrix&& other) noexcept {
data = other.data;
other.data = nullptr;
}
~Matrix() {
if (data != nullptr) {
delete[] data;
}
}
inline Scalar& coeff(long long i) {
return data[i];
}
inline Scalar coeff(long long i) const {
return data[i];
}
Scalar& operator()(long long i, long long j) {
return data[i * Cols + j];
}
Scalar operator()(long long i, long long j) const {
return data[i * Cols + j];
}
template<typename Other>
Matrix& operator=(const Other& other) {
if ((void*)&other == (void*)this) {
return *this;
}
Assign<Matrix, Other> assign(*this, other);
assign.run();
return *this;
}
Matrix& operator=(Matrix&& other) noexcept {
if ((void*)&other == (void*)this) {
return *this;
}
if (data != nullptr) {
delete[] data;
}
data = other.data;
other.data = nullptr;
return *this;
}
};
template<typename Left, typename Right>
Evaluator<Sum<Left, Right>> operator+(const Left& left, const Right& right) {
return Evaluator<Sum<Left, Right>>(Sum<Left, Right>(left, right));
}
Matrix<N, N> getOne() {
Matrix<N, N> m;
for (int i = 0; i < Matrix<N, N>::size; i++) {
m.coeff(i) = rand();
}
return m;
}
template<typename MatrixT>
void showOne(MatrixT& mat) {
for (int i = 0; i < MatrixT::size; ++i) {
cout << mat.coeff(i) << " ";
}
}
}
class Timer {
public:
std::chrono::time_point<std::chrono::high_resolution_clock> m_begin;
Timer() : m_begin(std::chrono::high_resolution_clock::now()) {}
void reset() {
m_begin = std::chrono::high_resolution_clock::now();
}
template<typename Duration = std::chrono::milliseconds>
long long elapsed() const {
return std::chrono::duration_cast<Duration>(std::chrono::high_resolution_clock::now() - m_begin).count();
}
};
int main() {
Timer timer;
{
naive_impl::Matrix<N, N> m1 = naive_impl::getOne();
naive_impl::Matrix<N, N> m2 = naive_impl::getOne();
naive_impl::Matrix<N, N> m3 = naive_impl::getOne();
timer.reset();
m3 = m1 + m2 + m3;
long long time = timer.elapsed();
cout << "Naive: " << time << "ms" << endl;
}
{
new_impl::Matrix<N, N> m1 = new_impl::getOne();
new_impl::Matrix<N, N> m2 = new_impl::getOne();
new_impl::Matrix<N, N> m3 = new_impl::getOne();
timer.reset();
m3 = m1 + m2 + m3;
long long time = timer.elapsed();
cout << "New: " << time << "ms" << endl;
new_impl::Matrix<N, N> ms[3] = { m1, m2, m3 };
timer.reset();
for (int i = 0; i < naive_impl::Matrix<N, N>::size; ++i) {
for (int j = 0; j < 3; ++j) {
ms[2].coeff(i) += ms[j].coeff(i);
}
}
time = timer.elapsed();
cout << "Ideal: " << time << "ms" << endl;
}
return 0;
}
在Ideal的实现中,为了防止编译器作弊开向量化指令特意不写成
for(int i = 0; i < size; ++ i) {
ms[2].coeff(i) = ms[0].coeff(i) + ms[1].coeff(i) + ms[2].coeff(i);
}
最后贴结果:
编译器: msvc -release版本
win10 Ryzen5800 16g
Naive implement takes: 146ms
New implement takes: 62ms
Ideal implement takes: 89ms
可以看到我们在同等规模下的矩阵加法运算相比于最简单的实现取得了相当不错的优化结果,这一点在更大规模的矩阵运算和更长的表达式下优势更明显。并且如果你观察运行时的内存占用可以发现后者在运算中几乎不产生额外空间开销。
这个例子并不难,也不支持更多的其他矩阵的复杂运算,但足以展示模板的强大,这使我们可以通过简洁抽象的表达式达到非常理想的结果。