想请教一个使用稀疏数据结构加速 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 (反而更快一些)
最近自己也在学习matrix-free FEM相关知识,想请问一下taichi目前有没有matrix-free的相关示例代码?
比如渊鸣大大在SIGGRAPH Asia 2018上的论文“Narrow-Band Topology Optimization on a Sparsely Populated Grid” 有考虑推出一个taichi实现吗?
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))