【问题】Sparse Matrix 的 '*' 操作在 ti.cuda 下会出现 unsupported operand type 吗

报错内容如下,在 ti.cpu 下就没有问题:

TypeError: unsupported operand type(s) for *: 'taichi._lib.core.taichi_python.CuSparseMatrix' and 'taichi._lib.core.taichi_python.CuSparseMatrix'

今天运行文档中的案例的时候出现了这个错误,我也不知道是哪里的问题,而且没有第二个 cuda 设备,我也不知道是不是我环境装得不对导致的。

运行环境:

  • Windows 11 22H2
  • RTX 3060 Laptop CUDA 12.0
  • Python 3.8.15
  • Taichi 1.4.0 commit fbe92fd8
  • llvm 15.0.1

运行的整个代码:

import taichi as ti

arch = ti.cuda
ti.init(arch=arch, debug=True)

n = 3
K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)

@ti.kernel
def fill(A: ti.types.sparse_matrix_builder()):
    for i in range(n):
        A[i, i] += 1

fill(K)

print(">>>> K.print_triplets()")
K.print_triplets()

A = K.build()
print(">>>> A = K.build()")
print(A)

print(">>>> Multiplication with a scalar on the right: E = A * 3.0")
E = A * 3.0
print(E)

print(">>>> Element-wise multiplication: H = E * A")
H = E * A
print(H)

print(f">>>> Element Access: A[0,0] = {A[0,0]}")

整个报错片段:

Traceback (most recent call last):
  File ".\test.py", line 28, in <module>
    H = E * A
  File "C:\Users\m\miniconda3\envs\taichi\lib\site-packages\taichi\linalg\sparse_matrix.py", line 92, in __mul__
    sm = self.matrix * other.matrix
TypeError: unsupported operand type(s) for *: 'taichi._lib.core.taichi_python.CuSparseMatrix' and 'taichi._lib.core.taichi_python.CuSparseMatrix'

欢迎来到 Taichi 论坛!
这里报错的原因是 CUDA 后端上 SparseMatrix 的 element-wise 乘法还没有实装,而 CPU 后端上是有 element-wise 乘法的实现的因此不会报错。关于这些运算符的实现可以参考这里的源代码:taichi/sparse_matrix.h at master · taichi-dev/taichi · GitHub

哦哦,感谢回复,原来还没有实装吗_(:з)∠)_

事实上好像由于 CuSparse 中并没有给出两个 SparseMatrix 的 element-wise 乘法实现,而只有一个 SparseMatrix 和一个 DenseMatrix 的实现,所以可能短时间内 Taichi 里也不会考虑去实现这个接口。你是有具体的场景的确需要使用到这个运算吗,还是说只是测试时刚好发现?如果有需求的话可以去 Taichi 的 Github 仓库提一个 feature request~