[问题] 如何向 taichi 传递 uint32 数据

import taichi as ti

ti.init(arch=ti.gpu)


@ti.kernel
def func(x: ti.template()):
    print('x:', x)
    y = ti.cast(x, ti.uint32)
    print('y:', y)


# func(ti.Vector([0xffffffff], ti.uint32))
"""
taichi.lang.exception.TaichiTypeError: 
    print('x:', x)
    ^^^^^^^^^^^^^^
Constant 4294967295 has exceeded the range of i32: [-2147483648, 2147483647]
"""

func(ti.Vector([-1], ti.uint32))
"""
x: [-1]
y: [4294967295]
"""
  • 我用 0xffffffff 构造 uint32 Vector 时, 在 python scope 没有报错, 但是运行时在 taichi scope 里给我报错超出 i32 范围
  • 如果我用 -1 构造 uint32 Vector, 在 python scope 也没有报错, 运行时在 taichi scope 可以强制转换为 uint32

请问如何向taichi 传递 uint32 数据?

Hi @geekberry , 关于Taichi在各个后端类型的支持,可以看文档:here

你的例子我稍微修改了一下:

import taichi as ti

ti.init(arch=ti.gpu)


@ti.kernel
def func(x: ti.template()):
    print('x:', x)
    y = ti.cast(x, ti.u32)
    print('y:', y)

func()

func(ti.Vector([0xffffffff], ti.u32))
还是报错啊

出现错误的原因是,你打印的地址而不是里面的数字。

import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

# numpy
y = np.array([-1], np.uint32)
print(f"y: {y}")

# taichi
x = ti.field(ti.u32,shape=1)
@ti.kernel
def func(x: ti.template()):
    x[0] = -1
    y = ti.cast(x[0], ti.u32)
    print(y)

func(x)

问题在于我想传的是 uint32 啊 :joy:, 而 0xffffffff 明显符合 u32 的区间, 但是报错:

Constant 4294967295 has exceeded the range of i32: [-2147483648, 2147483647]
是不是 ti.template() 在展开时忽视了 Vector 类型, 默认使用了 i32 ?

@ti.kernel
def func(x: ti.template()):
    print('x:', x[0])
    y = ti.cast(x, ti.u32)
    print('y:', y[0])


x = ti.Vector([0xffffffff], ti.u32)
func(x)

我想问题不是出在u32上, 而是出在打印地址上.

@ti.kernel
def func(x: ti.template()):
    # print('x:', x)
    """
    x 不能直接打印, 否则报错:
    Constant 4294967295 has exceeded the range of i32: [-2147483648, 2147483647]
    """

    print(f'x.n:{x.n},x.m:{x.m}')  # x.n:1,x.m:1
    print(f'x[0]:{x[0]}')  # x[0]:4294967295


vector = ti.Vector([0xffffffff], ti.u32)
func(vector)  # 直接传 vector

为什么直接打印x报错是4294967295 has exceeded the range呢?

因为Taichi里面使用int32来表示地址的。

还是有问题, 比如

np_array = np.array([1, 0xffffffff, 0x80000000], np.uint32)
print(np_array)  # [1, 4294967295, 2147483648]

@ti.kernel
def func():
    print(np_array[0])  # 1
    
    # print(np_array[1])
    """
    Constant 4294967295 has exceeded the range of i32: [-2147483648, 2147483647]
    """
    
    print(np_array[2])
    """
    Constant 2147483648 has exceeded the range of i32: [-2147483648, 2147483647]
    """

func()

numpy 里把类型设置为 np.uint32, 但是在 taichi scope 引用时还是默认识别为 i32

在使用Taichi语言的时候记得和Python语言区分开来,在Taichi scope是不知道numpy lib的。所以可以通过参数传递进Taichi scope,可以尝试下面的方式运行。

import taichi as ti
import numpy as np
ti.init(arch=ti.gpu)


np_array = np.array([1, 0xffffffff, 0x80000000], np.uint32)
print(np_array)  # [1, 4294967295, 2147483648]

@ti.kernel
def func(np: ti.types.ndarray()):
    print("print in taichi scope")
    for i in range(3):
        print(np[i])

func(np_array)

我想访问的是全局变量, 不是参数!
print(np_array[0]) 来看, 是可以在 taichi scope 里访问 numpy 数组的, 但类型转换有问题!

import taichi as ti
import numpy as np
ti.init(arch=ti.gpu)


np_array = np.array([1, 0xffffffff, 0x80000000], np.uint32)
print(np_array)  # [1, 4294967295, 2147483648]

@ti.kernel
def func(np: ti.types.ndarray()):
    print("print in taichi scope")
    for i in range(3):
        print(np[i])
    """
    1
    4294967295
    2147483648
    """

    # for i in range(3):
    #     print(np_array[i])
    """
    IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices
    """

    print(np_array[0])

    # print(np_array[1])
    """
    Constant 4294967295 has exceeded the range of i32: [-2147483648, 2147483647]
    """

    # print(np_array[2])
    """
    Constant 2147483648 has exceeded the range of i32: [-2147483648, 2147483647]
    """

func(np_array)

@geekberry I think there’s a gap in understanding what’s handled in python and what’s handled in taichi.
Let’s first make it clear that if you want to interact with numpy arrrays or torch tensors in taichi scope, you’ll have to pass them in explicitly as function arguments as @YuPeng did. And if you turn on ti.init(..., print_ir=True) you’ll see taichi correctly recognize it as uint32 and print is done without any problem.

kernel {
  $0 = offloaded
  body {
    print "print in taichi scope\n"
  }
  $2 = offloaded range_for(0, 3) grid_dim=2624 block_dim=128
  body {
    <i32> $3 = loop $2 index 0
    <*u32> $4 = arg[0]
    <*u32> $5 = external_ptr <$4>, [$3] element_dim=0 layout=AOS
    <u32> $6 = global load $5
    print $6, "\n"
  }
}

The you may wonder why your script above doesn’t work. Let’s comment out the line and see the IR for print(np_array[0]).

import taichi as ti
import numpy as np

ti.init(ti.cuda, print_ir=True)

np_array = np.array([1, 0xffffffff, 0x80000000], np.uint32)
print(np_array)  # [1, 4294967295, 2147483648]

@ti.kernel
def func():
    print(np_array[0])  # 1

func()

The Taichi IR is

kernel {
  $0 = offloaded
  body {
    <i32> $1 = const [1]
    print $1, "\n"
  }
}

See the i32 and const[1] there? It means np_array[0] is actually evaluated by Python (instead of Taichi) when you compile the Taichi kernel. So what Taichi saw is actually just a print with python number 1. Similarly for 0xffffffff, Taichi sees it as a literal number but tries to assign it to a i32 to print, thus the error thrown.
So to properly do what you asked for above you’ll need to give them an explicit type:

import taichi as ti
import numpy as np

ti.init(ti.cuda, print_ir=True)

np_array = np.array([1, 0xffffffff, 0x80000000], np.uint32)
print(np_array)  # [1, 4294967295, 2147483648]

@ti.kernel
def func():
    print(ti.u32(np_array[0]))
    print(ti.u32(np_array[1]))
    print(ti.u32(np_array[2]))

func()

Hopefully this can answer your question above. Thanks!

  1. 问题不只是出现在函数调用的地方, 赋值也会先转换为 i32, 如下:
import taichi as ti

ti.init(arch=ti.gpu, print_ir=True)

@ti.kernel
def kernel():
    a = ti.Vector([0x7fff_ffff], ti.u32)  # work

    # b = ti.Vector([0x8000_0000], ti.u32)
    """
    Constant 2147483648 has exceeded the range of i32: [-2147483648, 2147483647]
    """
    
    b = ti.Vector([ti.u32(0x8000_0000)], ti.u32) # work

kernel()

b = ti.Vector([0x8000_0000], ti.u32) 而言, 无论哪一门语言的用户都无法理解, 为什么会涉及到 i32 的范围的错误.


  1. 有些情况直接引用全局常量, 而非参数传递是更合理的写法.
import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

COLOR_MASK = np.array([0xff000000, 0x00ff0000, 0x0000ff00, 0x000000ff], np.uint32)

@ti.kernel
def kernel(color: ti.u32):
    red = color & COLOR_MASK[0]
    """
    Constant 4278190080 has exceeded the range of i32: [-2147483648, 2147483647]
    """

kernel(0xa53f7600)

既然支持外部变量使用, 用户在定义时也注意符合 uint32 的返回, 用户将不会意识到会有问题出现(直到外部变量中包含大于 i32 范围数字时出现…)


  1. 如果要将全局变量(如 COLOR_MASK) 作为参数传递, 不但可能让参数表变得非常长, 还会代理变量名污染的问题
import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

COLOR_MASK = np.array([0xff000000, 0x00ff0000, 0x0000ff00, 0x000000ff], np.uint32)

@ti.kernel
def kernel(color: ti.u32, COLOR_MASK: ti.types.ndarray()):
    red = color & COLOR_MASK[0]  # work

kernel(0xa53f7600, COLOR_MASK)

  1. 在 ti.kernel 内部定义也不是一个好办法, 首先 python scope 无法使用, ti.func使用需传参, 最后还是要手动处理类型转换
import taichi as ti

ti.init(arch=ti.gpu)

@ti.kernel
def kernel(color: ti.u32):
    COLOR_MASK = ti.Vector([ti.u32(0xff000000), 0x00ff0000, 0x0000ff00, 0x000000ff], ti.u32)

    red = color & COLOR_MASK[0]  # work

kernel(0xf53f7600)

  1. 对于上述代码, 官方是否有推荐的实践方式?
    直接引用外部的 ti.Vector 同样会在运行时报错
    直接引用外部的 ti.field 可以实现目的, 但是 field 却只能一一赋值, 不封装的话非常丑
import taichi as ti

ti.init(arch=ti.gpu)

# COLOR_MASK = ti.Vector([0xff000000, 0x00ff0000, 0x0000ff00, 0x000000ff], ti.u32)
"""
On line 14 of file "D:\PYTHON_PATH\test_taichi\@dev\learn.py", in kernel:
    red = color & COLOR_MASK[0]
          ^^^^^^^^^^^^^^^^^^^^^
Constant 4278190080 has exceeded the range of i32: [-2147483648, 2147483647]
"""

COLOR_MASK = ti.field(ti.u32, 4)
COLOR_MASK[0] = 0xff000000
COLOR_MASK[1] = 0x00ff0000
COLOR_MASK[2] = 0x0000ff00
COLOR_MASK[3] = 0x000000ff

@ti.kernel
def kernel(color: ti.u32):
    red = color & COLOR_MASK[0]  # only work for ti.field
    green = color & COLOR_MASK[1]  # work for ti.field and ti.Vector

kernel(0xf53f7600)

谢谢解答

Hey @geekberry!

  • About passing argument v.s. global, yea I totally agree that not everything should be passed as args, but you should pass them as args if you want to write into numpy arrays or torch tensors from a taichi kernel. But I agree that in your use case above simply reading from it doesn’t need to be passed in. (in fact it’s just a python value)
  • The error message about out of range for i32 isn’t helpful at all. (i32 comes from default int dtype set in Taichi fyi). We’ll take a look to see if we can fix it or improve the error message. In the meanwhile I’d suggest ti.init(..., default_ip=ti.i64) to temporarily work around it.(I know it’s not perfect) Is it okay for you?
  • 我最后使用了 ti.field(ti.u32, 4) 方案
  • ti.init(..., default_ip=ti.i64) 可能可以解决本项目问题, 但潜在可能遇到 u64 超 i64 的问题

@geekberry Yup u64 issue does exist (in fact I think we currently just clamp to i64 range which is brutal). Will try to fix that as well. Thanks a lot for report!