遇到一个关于kernel参数的奇怪问题

昨天根据example代码,写了一个用“半拉格朗日对流+maccormack”的程序,程序只是一个方块从左到右平移,还没涉及流体的运动。然后就遇到一个奇怪的问题,之前写的时候没遇到过。
写了两种maccormack的kernel程序,表达的意思是一样的,但是结果不一样。
第一种,kernel没有参数,直接当全局变量访问速度的field,代码如下图:


运行的结果,就是maccormack的效果很不明显(我没写截断,应该会出现artificial的现象)
结果如下:
image
第二种,把速度field当作参数送进kernel,代码如下:

运行结果,maccormack的效果就比较明显,有artificial的现象
结果如下:
image

两种写法的代码只有标红的地方不一样,其他地方一致,我以为两种写法应该一样的233333 :joy:. 我去找了文档,但是没找到相关线索,现在就没有头绪不知道问题出在哪,想请各位老师看看,非常感谢 :blush:

整个代码如下:

import taichi as ti


ti.init(arch=ti.gpu)

n_grid = 512
dt = 0.03
use_mc = True

shape = ti.field(ti.u8, shape=(n_grid, n_grid))
_dye_buffer = ti.field(float, shape=(n_grid, n_grid))
_new_dye_buffer = ti.field(float, shape=(n_grid, n_grid))
_new_dye_buffer_aux = ti.field(float, shape=(n_grid, n_grid))
_velocities = ti.Vector.field(2, float, shape=(n_grid, n_grid))
_new_velocities = ti.Vector.field(2, float, shape=(n_grid, n_grid))



        
@ti.kernel
def init_dens_vel():
    for i,j in ti.ndrange(n_grid, n_grid):
        _velocities[i,j] = ti.Vector([10.0 ,0.0])
        if i>0 and i<80 and j>200 and j<400:
             _dye_buffer[i,j] = 1.0


@ti.func
def lerp(vl, vr, frac):
    # frac: [0.0, 1.0]
    return vl + frac * (vr - vl)

@ti.func
def sample_min(qf, p):
    u, v = p
    s, t = u - 0.5, v - 0.5      
    # floor
    iu, iv = ti.floor(s), ti.floor(t)  
    a = sample(qf, iu, iv)        
    b = sample(qf, iu + 1, iv)
    c = sample(qf, iu, iv + 1)
    d = sample(qf, iu + 1, iv + 1)
    return min(a,  b, c, d)

@ti.func
def sample_max(qf, p):
    u, v = p
    s, t = u - 0.5, v - 0.5      
    # floor
    iu, iv = ti.floor(s), ti.floor(t)  
    a = sample(qf, iu, iv)         
    b = sample(qf, iu + 1, iv)
    c = sample(qf, iu, iv + 1)
    d = sample(qf, iu + 1, iv + 1)
    return max(a,  b, c, d)
    
@ti.func
def sample(qf, u, v):     
    I = ti.Vector([int(u), int(v)])   
    I = max(0, min(n_grid - 1, I))  
    return qf[I]                   

@ti.func
def bilerp(vf, p):   
    u, v = p
    s, t = u - 0.5, v - 0.5     
    # floor
    iu, iv = ti.floor(s), ti.floor(t) 
    # fract
    fu, fv = s - iu, t - iv
    a = sample(vf, iu, iv)         
    b = sample(vf, iu + 1, iv)
    c = sample(vf, iu, iv + 1)
    d = sample(vf, iu + 1, iv + 1)
    return lerp(lerp(a, b, fu), lerp(c, d, fu), fv)   


@ti.func
def backtrace(vf, p, dt):
    v1 = bilerp(vf, p)
    p1 = p - 0.5 * dt * v1
    v2 = bilerp(vf, p1)
    p2 = p - 0.75 * dt * v2
    v3 = bilerp(vf, p2)
    p -= dt * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3)
    return p


@ti.kernel
def advect(vf: ti.template(), qf: ti.template(), new_qf: ti.template()):
    for i, j in ti.ndrange(n_grid, n_grid):
        p = ti.Vector([i, j]) + 0.5  
        p = backtrace(vf, p, dt)
        new_qf[i, j] = bilerp(qf, p)

@ti.kernel
def maccormack(vf: ti.template()):
   
    for i, j in ti.ndrange(n_grid, n_grid):
        p = ti.Vector([i, j]) + 0.5   
        p = backtrace( vf, p, -dt )
        _new_dye_buffer_aux[i, j] = bilerp(_new_dye_buffer, p)

    for i, j in ti.ndrange(n_grid, n_grid):
        _new_dye_buffer[i, j] += 0.5 * (_dye_buffer[i,j]-_new_dye_buffer_aux[i,j])
    
       

def step():
    
    advect(_velocities, _velocities, _new_velocities)   
    advect(_velocities, _dye_buffer, _new_dye_buffer)   
    maccormack(_velocities)
    swap_v()
    swap_d()
    

def swap_v():
    global _velocities
    global _new_velocities 
    _velocities,_new_velocities = _new_velocities,_velocities

def swap_d():
    global _dye_buffer
    global _new_dye_buffer 
    _dye_buffer,_new_dye_buffer = _new_dye_buffer,_dye_buffer



init_dens_vel()
gui = ti.GUI('shape control', (n_grid, n_grid))

while True:
    step()
    
    gui.set_image(_dye_buffer)
    gui.show()

谢谢,@kongge。我运行了一下代码,发现确实有这样的问题存在。我们查找一下原因。

BTW,只要带上ti.template就会出现 artifacts。即使没有用到参数 vf

@ti.kernel
def maccormack(vf: ti.template()):
    for i, j in ti.ndrange(n_grid, n_grid):
        p = ti.Vector([i, j]) + 0.5   
        p = backtrace( _velocities, p, -dt )
        _new_dye_buffer_aux[i, j] = bilerp(_new_dye_buffer, p)
1 Like

问题出在使用python语法swap两个taichi field。如果你把代码改成这样:

def swap_v():
    global _velocities
    global _new_velocities 
    swap_buffers(_velocities, _new_velocities)
    # _velocities,_new_velocities = _new_velocities,_velocities

def swap_d():
    global _dye_buffer
    global _new_dye_buffer 
    swap_buffers(_dye_buffer, _new_dye_buffer)
    # _dye_buffer,_new_dye_buffer = _new_dye_buffer,_dye_buffer

@ti.kernel
def swap_buffers(f1:ti.template(), f2:ti.template()):
    for I in ti.grouped(f1):
        temp = f1[I]
        f1[I] = f2[I]
        f2[I] = temp

两种实现的表现就是完全一样的了。

这个问题藏得比较深,严肃解释起来也会比较晦涩:

taichi field和数组一样,field的名字记录的是指向内存块头部的指针。taichi会在编译某个kernel时(在JIT的情况下,就是你第一次调用这个kernel的时候),记录下这个kernel里面所有field的指针对应的内存空间是多少。当你重复调用同样的kernel时,taichi不会重复编译,所以即使你把field的指针头指向了别的内存空间,taichi也是不知道的。

我们再来看看python swap做了什么:python swap没有进行深拷贝,他只是会把两个field的指针头交换位置,所以在你的第一种实现方式下,maccormack函数只有在第一次执行的时候被编译了,在后续调用的时候并没有访问到swap之后的内存地址。

那么为什么加了ti.template()之后运行就正确了呢?因为加上了ti.template()的kernel会依据传入的field的地址实例化成不同的kernels,所以当你在加上了ti.template()之后,你的maccormack()函数会被实例化成两种不同的kernel,分别对应swap前和swap后的内存访问方式。在这种实现下,你的代码运行是正确的。这种实现方式节省了swap的时间(因为浅拷贝),但是增加了编译的时间(maccormac()函数会被实例化成许多分,取决于你调用函数时放入的参数field指向的不同内存空间的数目)。但是这样的实现可维护性非常低,所以个人并不推荐采用这种实现方式。

希望有解答到你的问题。

4 Likes

哦哦,是这样啊!我大概看懂了,这没讲解我肯定发现不了 :joy:
感谢老师能花时间做这么详细的解答,受益匪浅,非常非常感谢 :grinning: