【萌新求教】如何在taichi scope下优雅地交换两个field

综合讨论区
#1

如题，在python scope下，交换两个field易于实现：

a = ti.field(ti.int32, shape=(3))
b = ti.field(ti.int32, shape=(3))
a.fill(1)
b.fill(2)
print(a) //"[1,1,1]"
print(b) //"[2,2,2]"
a,b = b,a
print(a) //"[2,2,2]"
print(b) //"[1,1,1]"

但是在taichi scope下，同样的操作将导致错误：

Traceback (most recent call last):
  File "z:\test.py", line 18, in <module>
    test()
  File "C:\Python310\lib\site-packages\taichi\lang\kernel_impl.py", line 944, in wrapped
    raise type(e)('\n' + str(e)) from None
taichi.lang.exception.TaichiTypeError:
File "z:\test.py", line 16, in test:
    a,b = b,a
    ^^^^^^^^^
Invalid constant scalar data type: <class 'taichi.lang.field.ScalarField'>
#2

最后我用一个trick绕过了这一问题：将a、b合并为field(2,3)，使用a[i1,..]和a[i2,..]代替a[..]、b[..]，交换时只需i1,i2=i2,i1。

但这将导致field的fill()方法不再可用——当我想独立地将a、b中的某一个全部赋值为0时，我只能手动用循环完成。

所以，有任何更加优雅的方法吗？

#3

你可以尝试用struct for来实现field间数据拷贝，比如：

import taichi as ti

ti.init(ti.cpu)

a = ti.field(ti.int32, shape=(3))
b = ti.field(ti.int32, shape=(3))
a.fill(1)
b.fill(2)

@ti.kernel
def test():
    for i in b:
        print(a[i])
        a[i] = b[i]
        print(a[i])

test()

直接支持field assignment比如a = b会有一些歧义：如果a在之前定义过则表示要做数值拷贝，而如果a在之前没定义过则表示定义并初始化a，这样语义的一致性比较差。

我个人感觉也许可以给ti.field做一个函数库，支持一些比如Field::copy_from()一类的方法?

#4

类似这样，把field copy包到一个ti.func里（你可能需要加一些检查）

import taichi as ti

ti.init(ti.cpu)

a = ti.field(ti.int32, shape=(3))
b = ti.field(ti.int32, shape=(3))
a.fill(1)
b.fill(2)

@ti.func
def field_copy(src: ti.template(), dst: ti.template()):
    for I in ti.grouped(src):
        dst[I] = src[I]

@ti.kernel
def test():
    # copy b to a
    field_copy(b, a)
    print(a[0])

test()
#5

可以用alias：

@ti.kernel
def swap():
    print(a[0]) #1
    print(b[0]) #2
    a,b = ti.static(b,a)
    print(a[0]) #2
    print(b[0]) #1

但这只是看起来像，你并不能再用一次a,b = ti.static(b,a)把它们换回来。真要换还是在python scope吧。