使用稀疏数据结构加速B^T@A@B

综合讨论区
#1

想请教一个使用稀疏数据结构加速 B.transpose() @ A @ B的问题，这里A、B都是稀疏矩阵，但是因为维度比较大，都是用field去存储的

为了加速矩阵计算我尝试使用bitmasked，一个可以复现的最小demo如下：

import taichi as ti
ti.init(arch=ti.cpu, kernel_profiler=True)

n = 500
m = 200
A = ti.field(ti.f32)
B = ti.field(ti.f32)
A_bm = ti.root.bitmasked(ti.ij, n) # A: (n x n)
A_bm.place(A)
B_bm = ti.root.bitmasked(ti.ij, (n, m)) # B: (n x m)
B_bm.place(B)
result = ti.field(ti.f32, shape=(m,m))

for i in range(n):
    A[i, i] = 1. 

for i in range(m):
    B[i, i] = 1. 

@ti.kernel
def calculate():
    for i, j in ti.ndrange(m, m): # 计算 B^T @ A @ B, 结果为result
        result[i, j] = 0.
        for t in range(n):
            s = 0.
            for k in range(n):
                if ti.is_active(A_bm, [k,j]) and ti.is_active(B_bm, [t, k]):
                    s += A[t, k] * B[k, j]
                    s = s * B[t, i]
            result[i, j] += s

calculate()
ti.print_kernel_profile_info()

但是结果并没有加速：
使用ti.is_active： 8.501 s
不使用：6.807 s （反而更快一些）

请问正确使用稀疏数据结构的方式是怎么样的？谢谢。

1 Like
#2

Hi,

想了解下这里A和B的SNode是啥？

我的理解是这里如果A和B 是sparse Snode的话，可以直接ti.is_active(A, [i, j, ...])来判断index这里是不是sparse, 这样不需要访问额外的A_bm, 节省memory bandwidth。（还是你这里的self.A就是直接运算的矩阵，如果这样的话当我没说 :sweat_smile:）

不谈稀疏矩阵乘法的话，Taichi自己免费提供的优化方向应该有

  1. block_shared_memory(Cuda上的shared memory)，可以看看这里
  2. 调整一下ti.block_dim(Cuda上的block number), 也可能有提升
  3. 调整一下A, B 的SNode结构，如何合理的访问内存也很关键
    1. 如果A, B不大的话，甚至可以直接粗暴的用ti.dense，有时候反而比sparse来的快
  4. 用async_mode

这里推荐看一下Taichi的sparse documentation.

这里建议如果不麻烦的话可以把这部分拆出来做成一个能复制粘贴一键run的脚本，这样方便感兴趣的同学可以本地调试，参与的同学会多一些😂

1 Like
#3

可以期待一下 @YuPeng 的 Sparse matrix系统，预计10月中旬上线。

（另外补充一下，稀疏线性代数和matrix-free的技术计算是两个逻辑，二者的稀疏性和计算模式都有很大不同。）

1 Like
#4

非常感谢你的帮助！
我已经修改了之前的问题，贴上了一个可复现的最简demo。这里的A和B都是直接运算的矩阵。

#5

感谢渊鸣大大的回复！期待Sparse matrix系统！ :star_struck:

最近自己也在学习matrix-free FEM相关知识，想请问一下taichi目前有没有matrix-free的相关示例代码？
比如渊鸣大大在SIGGRAPH Asia 2018上的论文“Narrow-Band Topology Optimization on a Sparsely Populated Grid” 有考虑推出一个taichi实现吗？ :stuck_out_tongue_closed_eyes:

1 Like
#6

你可以自己实现一个哈哈，其实那篇文章里面的Linear FEM就是个巨大的stencil，有3x3x3x3x2个系数。实现一个Matrix-free CG不难的，preconditioner有点挑战。

1 Like
#7

好的！我参考C++的开源版本试一试 :wink:

#8

由于实际中B是很稀疏的，因此把之前的矩阵乘法换了个写法，变成了以B主导循环的矩阵乘法，
速度快了一倍多，
但是原子加又带来了并行时数据正确性的问题。

新的矩阵乘法的思想如下demo：

import numpy as np
x = 10
y = 5

K = np.random.rand(10,10)
I = np.random.rand(10, 5)

A0 = I.transpose() @ K @ I

A1 = np.zeros((y, y))
for i in range(x):
    for j in range(y): # 前两个循环在使用taichi的时候换成struct for
        if I[i,j] != 0.:
            for t in range(x):
                for m in range(y):
                    if I[t, m] != 0.:
                        A1[j, m] += I[i,j] * K[i,t] * I[t, m]

print(np.linalg.norm(A1 - A0))