Using Taichi from Julia

[lucifer1004]

Concisely describe the proposed feature Using Taichi from Julia will be a very nice-to-have feature.

Describe the solution you'd like (if any) Currently, there is the PythonCall.jl package, which solves most of the problems, however, a syntax check in Taichi becomes the main obstacle.

The following code cannot run at the moment. ti.init, ti.Vector.field and ti.GUI all work well. However, for ti.func and ti.kernel, there will be an error thrown ("Taichi functions do not support variable positional parameters (i.e., *args)"), which is related to https://github.com/taichi-dev/taichi/blob/13983c783129dcb39a714b82a217f69955c80722/python/taichi/lang/kernel_impl.py#L448. I think it is highly possible that if that check is removed or disabled temporarily, the following code can run fluently.

using CondaPkg
CondaPkg.add("python"; version="3.9.12")
CondaPkg.add_pip("taichi")
CondaPkg.resolve()

using PythonCall
ti = pyimport("taichi")
ti.init(arch=ti.gpu)
n = 640
pixels = ti.Vector.field(3, dtype=pytype(1.0), shape=(n * 2, n))
complex_sqr = ti.func(z -> ti.Vector([z[1]^2 - z[2]^2, z[1] * z[0] * 2]))
paint = ti.kernel(t ->
    for (i, j) in pixels
        c = ti.Vector([(1 + ti.sin(t)) * 0.285, (1 + ti.cos(t)) * 0.1])
        z = ti.Vector([i / n - 1, j / n - 0.5]) * 2
        rgb = ti.Vector([0, 1, 1])
        iterations = 0
        while z.norm() < 20 && iterations < 50
            z = complex_sqr(z) + c
            iterations += 1
        end
        pixels[i, j] = (1 - iterations * 0.02) * rgb
    end
)

gui = ti.GUI("Julia Set", res=(n * 2, n))
i = 0
flag = 0
while true
    global flag
    global i
    if flag == 0
        i -= 1
        if i * 0.02 <= 0.2
            flag = 1
        end
    else
        i += 1
        if i * 0.02 > (π * 1.2)
            flag = 0
        end
    end

    paint(i * 0.02)
    gui.set_image(pixels)
    gui.show()
end

[lucifer1004]

The reason this error is triggered is that PythonCall.jl uses CPython's PyObject_CallObject to call Python functions, which passes *args to the callable.

[yuanming-hu]

I feel like fixing "Taichi kernels do not support variable positional parameters" should not be too challenging. However, to use Taichi in Julia, Taichi needs to inspect the AST of kernels and functions written in Julia. Not sure if this can lead to further issues.

For people who wish to fix the variable positional parameter issue, please take a close look at https://github.com/taichi-dev/taichi/blob/13983c783129dcb39a714b82a217f69955c80722/python/taichi/lang/kernel_impl.py#L434.

[lucifer1004]

Removing the check is easy, but I do not know why there are those checks and whether there will be some downsides if they are removed.

[lucifer1004]

Yes, the AST part might be another obstacle. In PythonCall.jl, there is no AST conversion from Julia to Python, and Julia functions are pointers from Python's side.

[yuanming-hu]

Removing the check is easy, but I do not know why there are those checks and whether there will be some downsides if they are removed.

I don't think there are any drawbacks :-) Providing that flexibility has no harm.

[k-ye]

cc @lin-hitonami @strongoier in case there are some considerations behind this check...

[lucifer1004]

Maybe related:

https://discourse.julialang.org/t/idea-for-a-ast-frontend-for-parsing-between-julia-and-python/57963/3

[lin-hitonami]

IIRC, Taichi kernels do not yet support *args in the parameter list. It would be nice if you could add support for it. FYI, the AST transformer handles kernel arguments at https://github.com/taichi-dev/taichi/blob/08f6d60a9ba17591285e90b6d18a59afa4e5fc93/python/taichi/lang/ast/ast_transformer.py#L496

[lucifer1004]

IIRC, Taichi kernels do not yet support *args in the parameter list. It would be nice if you could add support for it. FYI, the AST transformer handles kernel arguments at

https://github.com/taichi-dev/taichi/blob/08f6d60a9ba17591285e90b6d18a59afa4e5fc93/python/taichi/lang/ast/ast_transformer.py#L496

Maybe I can look at this sometime. I just get a bit more familiar with Python AST while writing this toy project https://github.com/lucifer1004/Jl2Py.jl which converts Julia AST to Python AST.

[strongoier]

Another thing to note here is that Taichi requires type annotations of kernel parameters, so you may need to figure out a way to pass them in.

[lucifer1004]

Can this be done by passing a Python AST object to Taichi? I can convert the Julia kernel function to Python AST at the Julia side.

[strongoier]

Can this be done by passing a Python AST object to Taichi? I can convert the Julia kernel function to Python AST at the Julia side.

Yes. You can read the code path starting from https://github.com/taichi-dev/taichi/blob/08f6d60a9ba17591285e90b6d18a59afa4e5fc93/python/taichi/lang/kernel_impl.py#L928 and figure out the process before passing a Python AST. As a start, I think you can ignore code for advanced features like autodiff, ODOP, ..., and get a simple kernel working.

[lucifer1004]

The following code ran successfully @strongoier @lin-hitonami @k-ye @yuanming-hu

using CondaPkg
using PythonCall
CondaPkg.add_pip("taichi")
ti = pyimport("taichi")

let
    ti.init(; arch=ti.gpu)
    n = 640
    pixels = ti.Vector.field(3; dtype=pytype(1.0), shape=(n * 2, n))

    ti_str = """
@ti.kernel
def paint(t: float):
    for (i, j) in pixels:
        c = ti.Vector([(1 + ti.sin(t)) * 0.285, (1 + ti.cos(t)) * 0.1])
        z = ti.Vector([i / n - 1, j / n - 0.5]) * 2
        rgb = ti.Vector([0, 1, 1])
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = ti.Vector([z[0]**2 + z[1]**2, z[0] * z[1] * 2]) + c
            iterations += 1
        pixels[i, j] = (1 - iterations * 0.02) * rgb
    """
    namespace = pydict(["ti" => ti, "n" => n, "pixels" => pixels])
    write("__tmp.py", ti_str)
    code = pycompile(ti_str; filename="__tmp.py", mode="exec")
    pyexec(code, namespace)
    paint = namespace.get("paint")

    gui = ti.GUI("Julia Set"; res=(n * 2, n))
    i = 0
    flag = 0
    while pyconvert(Bool, gui.running)
        if flag == 0
            i -= 1
            if i * 0.02 <= 0.2
                flag = 1
            end
        else
            i += 1
            if i * 0.02 > (π * 1.2)
                flag = 0
            end
        end

        paint(i * 0.02)
        gui.set_image(pixels)
        gui.show()
    end
end

Screenshot:

And with my Jl2Py.jl package, the kernel function written in Python could be replaced by a Julia function transpiled to Python.

[lucifer1004]

I am also trying PythonCall.pyfunc, but this one requires removal of the positional args check as mentioned above.

Update

I think this is very promising. I have solved the signature problem using pyfunc, but the next problem is that Taichi also requires sourcelines, as in

https://github.com/taichi-dev/taichi/blob/a917fd068361580de370f5e14f6f05f9ba8deb1f/python/taichi/lang/kernel_impl.py#L110

In this version there is no explicit Python code

using CondaPkg
using PythonCall
CondaPkg.add_pip("taichi")
ti = pyimport("taichi")
inspect = pyimport("inspect")

let
    ti.init(; arch=ti.gpu)
    n = 640
    pixels = ti.Vector.field(3; dtype=pytype(1.0), shape=(n * 2, n))

    function paint(t)
        for (i, j) in pixels
            c = ti.Vector([(1 + ti.sin(t)) * 0.285, (1 + ti.cos(t)) * 0.1])
            z = ti.Vector([i / n - 1, j / n - 0.5]) * 2
            rgb = ti.Vector([0, 1, 1])
            iterations = 0
            while z.norm() < 20 && iterations < 50
                z = ti.Vector([z[0]^2 + z[1]^2, z[0] * z[1] * 2]) + c
                iterations += 1
                pixels[i, j] = (1 - iterations * 0.02) * rgb
            end
        end
    end

    param = inspect.Parameter("t"; annotation=pytype(1.0), kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
    sig = inspect.Signature([param])
    paint_py = pyfunc(paint; signature=sig)
    paint_ti = ti.kernel(paint_py)
    gui = ti.GUI("Julia Set"; res=(n * 2, n))
    i = 0
    flag = 0
    while pyconvert(Bool, gui.running)
        if flag == 0
            i -= 1
            if i * 0.02 <= 0.2
                flag = 1
            end
        else
            i += 1
            if i * 0.02 > (π * 1.2)
                flag = 0
            end
        end

        paint_ti(i * 0.02)
        gui.set_image(pixels)
        gui.show()
    end
end

[lucifer1004]

I found that in _get_tree_and_ctx(), the source code is just used for ast.parse(). Is there an equivalent function that directly takes the AST as input?

[lin-hitonami]

I found that in _get_tree_and_ctx(), the source code is just used for ast.parse(). Is there an equivalent function that directly takes the AST as input?

Currently there isn't. However, you can write one if you want. You can split the function into two functions: a get tree part and a get ctx part.

[lucifer1004]

Just made the following code work! Taichi.jl is an unpublished package written by me, using Jl2Py.jl under the hood.

using Taichi

let
    ti.init(; arch=ti.gpu)
    n = 640
    pixels = ti.Vector.field(3; dtype=pytype(1.0), shape=(n * 2, n))
    locals = map(x -> string(x.first) => x.second, collect(Base.@locals))

    paint = Taichi.@ti_kernel(function f(t::Float64)
                                  for (i, j) in pixels
                                      c = ti.Vector([(1 + ti.sin(t)) * 0.285, (1 + ti.cos(t)) * 0.1])
                                      z = ti.Vector([i / n - 1, j / n - 0.5]) * 2
                                      rgb = ti.Vector([0, 1, 1])
                                      iterations = 0
                                      while z.norm() < 20 && iterations < 50
                                          z = ti.Vector([z[0]^2 + z[1]^2, z[0] * z[1] * 2]) + c
                                          iterations += 1
                                          pixels[i, j] = (1 - iterations * 0.02) * rgb
                                      end
                                  end
                              end, locals)

    gui = ti.GUI("Julia Set"; res=(n * 2, n))
    i = 0
    flag = 0
    while pyconvert(Bool, gui.running)
        if flag == 0
            i -= 1
            if i * 0.02 <= 0.2
                flag = 1
            end
        else
            i += 1
            if i * 0.02 > (π * 1.2)
                flag = 0
            end
        end

        paint(i * 0.02)
        gui.set_image(pixels)
        gui.show()
    end
end

[yuanming-hu]

Amazing... Did you convert Julia AST to Python AST, or simply constructed the Taichi CHI IR using the Julia AST?

[lucifer1004]

Amazing... Did you convert Julia AST to Python AST, or simply constructed the Taichi CHI IR using the Julia AST?

I converted them to Python AST

[lucifer1004]

I will publish the code later, and maybe a blog.

[yuanming-hu]

I will publish the code later, and maybe a blog.

That's unbelievable progress, and I'm sure the community is eager to know more about how you made it!

[lucifer1004]

See https://github.com/lucifer1004/Taichi.jl