Some thoughts on optimizing your taichi kernels (Lessons from ray tracing in one weekend)

[bsavery]

Hi there. First of all congrats on the 1.0 release! I was inspired to take another look at my Ray Tracing in one Weekend repo and had some thoughts on optimizing kernel code. In fact it would be interesting to post a longer post on doing RTOW in taichi if I had the time, but wanted to at least write down some thoughts and get feedback.

(Of note I mainly test on a MacBook Pro, using metal - more on that later. I have not tested on CPU much and CUDA less so, but more so on Vulkan)

Object Oriented code

This is more of style thing. If you like for your code to be object oriented using the @ti.data_oriented is an insignificant overhead. Of course you'll want to store some data in your objects, so let's take the case of an array of Spheres:

class Spheres:
    def __init__(self, n):
        self.center = [...]
        self.radius = [...]  # n spheres

Representing this optimally in Taichi - Array of structs or struct of arrays

Now there are few ways we could move the above data structure to taichi. Basically it comes down to which you prefer. I found little difference in performance. One is to use two arrays:

@ti.data_oriented
class Spheres:
    def __init__(self, n):
        self.center = ti.Vector.field(3, ti.f32)
        self.radius =  ti.field(ti.f32)
        ti.root.dense(ti.i, n).place(self.center, self.radius)

This last part is important. If you are going to use the center and radius of the sphere together (probably a good bet in this case) this will put the center[I] and radius[I] next to each other.

The other alternative I've found is using structs:

sphere_type ti.types.struct(center=ti.Vector(3, ti.f32), radius=ti.f32)

spheres = sphere_type.field(shape(n,))

Of course if you want to associate functions with the struct above you can wrap it in a class:

@ti.data_oriented
class Spheres:
    def __init__(self, n):
        self.data = sphere_type.field(shape(n,))

    @ti.func
    def do_something_with_sphere(self, i):
        ....

Optimizing memory usage of arrays of structs

Both of the above ways, Structs of arrays and an Array of Structs are roughly equivalent performance! Why is this? Well if you look closely at the array of structs spheres = sphere_type.field(shape(n,)) it appears that this is internally just a structure of arrays! (Devs is this true?). For example I can do bulk operations on one array in the struct.
spheres.center.from_numpy(...)

Caveat on performance

There is a large BUT here though. I have found that dereferencing a struct object after getting it is slower in multiple parts of the code. Take for example.

spheres = sphere_type.field(shape(n,))

@ti.kernel
def double_and_move_spheres():
    for i in spheres:
        sphere = spheres[i]
        sphere.center += Vector(1, 0, 0) 
        sphere.radius *= 2

Catch what is happening here? The above code is clean but doesn't behave nicely from a memory standpoint. Multiple threads will be addressing different points in the spheres array at the same time. This is a trivial example but in complex code it matters more. A better solution is something like:

spheres = sphere_type.field(shape(n,))

@ti.kernel
def double_and_move_spheres():
    for i in spheres:
        sphere = spheres[I]
        center, radius = sphere.center, sphere.radius # really only dereference the struct once
        center += Vector(1, 0, 0) 
        radius *= 2
        sphere.center, sphere.radius = center, radius # again only dereference once

Again, developers please correct me if I'm wrong but this seems to be what's happening of why this would affect performance.

Wrapping up the object oriented design

This is more a matter of style but here's how I like to keep my object oriented data structures somewhat pythonic and performing well. The key thing to note here is that you can override python "magic methods" with ti.funcs!

sphere_type ti.types.struct(center=ti.Vector(3, ti.f32), radius=ti.f32)

@ti.data_oriented
class Spheres:
    def __init__(self, n):
        self.data = sphere_type.field(shape(n,))

    @ti.func
    def __getitem__(self, n):
        sphere = self.data[n]
        return sphere.center, sphere.radius

    @ti.func
    def __setitem__(self, n, val):
        center, radius = val
        self.data[n] = sphere_type(center=center, radius=radius)

@ti.kernel
def double_and_move_spheres():
    for i in spheres.data:
        center, radius = spheres[i]
        center += Vector(1, 0, 0) 
        radius *= 2
        spheres[i] = (center, radius)

A couple random notes on performance

1. At least on metal / vulkan, I have found u32 types to be slightly slower (!) than i32. I'm not sure why, my suspicion is it has something to do with some casting taichi is doing internally.
2. Some point around version 0.8XX on metal performance suffered dramatically. However, I can get even better performance than before by using Vulkan / MoltenVK on macOS.

The latter two points are probably worth a developer commenting on! Let me know if you have other questions.

[k-ye]

Wow, thank you for this very detailed report!

Well if you look closely at the array of structs spheres = sphere_type.field(shape(n,)) it appears that this is internally just a structure of arrays! (Devs is this true?).

Not really. While using from_numpy makes centers appear to be a dedicated array, this is actually implemented as a Taichi kernel that knows how to read each cell of the centeres field. I'm not sure about the scale of your testing data, but in general, the difference between AOS and SOA can be more obvious for larger scale tests (cc @turbo0628 @qiao-bo).

Caveat on performance

Note that by default, these augassignments on global fields are atomic (see https://docs.taichi-lang.org/lang/articles/operator#supported-atomic-operations). So by rewriting them into local variables, you can avoid this problem.

Taichi does try to demote unnecessary atomics during JIT if it can prove that it is safe to do so. This might need to take a look.

The key thing to note here is that you can override python "magic methods" with ti.funcs!

That is an interesting idea! cc @lin-hitonami who is working on the real function support.

Some point around version 0.8XX on metal performance suffered dramatically.

Ouch. Do you have a concrete repro :-( We will look into this.

[bsavery]

Well if you look closely at the array of structs spheres = sphere_type.field(shape(n,)) it appears that this is internally just a structure of arrays! (Devs is this true?).

Not really. While using from_numpy makes centers appear to be a dedicated array, this is actually implemented as a Taichi kernel that knows how to read each cell of the centeres field. I'm not sure about the scale of your testing data, but in general, the difference between AOS and SOA can be more obvious for larger scale tests (cc @turbo0628 @qiao-bo).

What I mean is that an struct field of spheres seems looking at it to be pretty equivalent (in performance and under the hood) to:

center = ti.Vector.field(3, ti.f32)
radius =  ti.field(ti.f32)
ti.root.dense(ti.i, n).place(self.center, self.radius)

# these two seem to perform about the same and I suspect they're similar in layout on memory?
sphere_type ti.types.struct(center=ti.Vector(3, ti.f32), radius=ti.f32)
spheres = sphere_type.field(shape(n,))

Caveat on performance

Note that by default, these augassignments on global fields are atomic (see https://docs.taichi-lang.org/lang/articles/operator#supported-atomic-operations). So by rewriting them into local variables, you can avoid this problem.

Taichi does try to demote unnecessary atomics during JIT if it can prove that it is safe to do so. This might need to take a look.

The key thing to note here is that you can override python "magic methods" with ti.funcs!

So I want to be clear about what you mean here, and I think my testing said about the same thing. This is a simple example but looking at what happens a lot in the raytracer. Do you mean:

spheres = sphere_type.field(shape(n,))

@ti.kernel
def do_something():
    for i in spheres:
        spheres[I].center += Vector(1, 0, 0)   # forces atomic.  slower?
        
        center, radius = spheres[I].center, spheres[I].radius
        sphere.center
        center += Vector(1, 0, 0) 
        spheres[I].center = center   # copy locally, not bad performance

        # what about this case when it's just reads of the structure?  This doesn't seem to make a difference to copy locally
        # in this we're calculating 'c' of b^2 - 4 a c to intersect ray with sphere
        oc = ray_origin - spheres[I].center
        c = (oc.norm_sqr - spheres.radius**2)

        # vs 
        center, radius = spheres[I].center, spheres[I].radius
        oc = ray_origin - center
        c = (oc.norm_sqr - radius**2)

That is an interesting idea! cc @lin-hitonami who is working on the real function support.

Some point around version 0.8XX on metal performance suffered dramatically.

Ouch. Do you have a concrete repro :-( We will look into this.

It's not a simple example I have but just my tests in https://github.com/bsavery/ray-tracing-one-weekend-taichi

[k-ye]

Some point around version 0.8XX on metal performance suffered dramatically.

Suspicious PR: #3090

• After testing https://github.com/taichi-dev/taichi/releases/tag/v0.8.2 (which included [Metal] Support multiple SNode roots in codegen/runtime #3090), there wasn't any perf issue.
• v0.9.0 was slow, but v0.8.8 was OK. Starting from v0.8.9, it started throwing compile-time errors (around for i in ti.static(range(3)))

[bsavery]

BTW if you want a succinct issue (perhaps I should open this if you think it's valid). THIS IS VERY SLOW But passing the deconstructed struct is much faster.

I'm not sure if this is expected or not @k-ye

spheres = sphere_type.field(shape(n,))

@ti.kernel
def do_something():
    for i in spheres:
        sphere = spheres[I]
        intersect_sphere(sphere, ray)


@ti.func
def intersect_sphere(sphere, ray):
    oc = ray.origin - ray.center
    a = ray.direction()
    half_b = oc.dot(ray.direction)
    c = (oc.norm_sqr() - sphere.radius**2)
    discriminant = (half_b**2) - a * c

    hit = discriminant >= 0.0
    root = -1.0
    if hit:
        sqrtd = discriminant**0.5
        root = (-half_b - sqrtd) / a

    return hit, root

[victoriacity]

My speculation is that not deconstructing the results in many global loads from the memory and thus slowing the performance. I tested with my own minimal example but it seemed that whether passing the whole struct element or not resulted in the same IR.

Could you provide a small fully-working example which shows the performance difference? Printing their IR (ti.init(print_ir=True)) will give insights into the underlying cause.

[bsavery]

Sure. Here's a good example. I'll make an issue too.

import taichi as ti
import time

ti.init(arch=ti.gpu)

n_spheres = 100000

# just a ray pointing in Z direction
ray_org = ti.Vector([0.0, 0.0, 0.0])
ray_dir = ti.Vector([0.0, 0.0, 1.0])

# random set of spheres
sphere_type = ti.types.struct(center=ti.types.vector(3, ti.f32), radius=ti.f32)
spheres = sphere_type.field(shape=(n_spheres,))
@ti.kernel
def fill_spheres():
    # generate random spheres in 0-100 xyz with radius 0-10
    for i in spheres:
        spheres[i].radius = ti.random() * 10.0
        spheres[i].center = ti.Vector([ti.random(), ti.random(), ti.random()]) * 100.0
fill_spheres()


@ti.kernel
def pass_reference() -> ti.i32:
    num_hit = 0
    for i in spheres:
        sphere = spheres[i]
        if intersect_sphere(sphere, ray_org, ray_dir):
            num_hit += 1
    return num_hit


@ti.kernel
def pass_decomposed() -> ti.i32:
    num_hit = 0
    for i in spheres:
        sphere = spheres[i]
        center, radius = sphere.center, sphere.radius
        if intersect_center_radius(center, radius, ray_org, ray_dir):
            num_hit += 1
    return num_hit


@ti.func
def intersect_sphere(sphere, ray_origin, ray_direction):
    # return if the ray hits the sphere
    oc = ray_origin - sphere.center
    a = ray_direction.norm_sqr()
    half_b = oc.dot(ray_direction)
    c = (oc.norm_sqr() - sphere.radius**2)
    discriminant = (half_b**2) - a * c

    return discriminant >= 0.0

@ti.func
def intersect_center_radius(center, radius, ray_origin, ray_direction):
    # return if the ray hits the decomposed sphere
    oc = ray_origin - center
    a = ray_direction.norm_sqr()
    half_b = oc.dot(ray_direction)
    c = (oc.norm_sqr() - radius**2)
    discriminant = (half_b**2) - a * c

    return discriminant >= 0.0


t = time.time()
pass_reference()
print('Passing reference', time.time() - t)

t = time.time()
pass_decomposed()
print('Passing decomposed', time.time() - t)

Also note metal vs vulkan effects quite a bit:
Metal
Passing reference 0.23202180862426758
Passing decomposed 0.09195494651794434

Vulkan
Passing reference 0.07849979400634766
Passing decomposed 0.05303597450256348

CPU
Passing reference 0.08902120590209961
Passing decomposed 0.06803393363952637

[victoriacity]

I checked the generated IR of both kernels and they indeed produce the exact same IR. Besides, the kernel profiler results did not give significant difference in performance between them either. Now my speculation is that the performance difference mainly results from Python overhead.

[bsavery]

Also regarding performance. I see significantly better performance in macOS with https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/rendering/cornell_box.py in vulkan vs metal.

[bsavery]

So back to one of my original questions. Is a field of a compound struct type stored internally as a structure of arrays?

I.e

sphere_type ti.types.struct(center=ti.Vector(3, ti.f32), radius=ti.f32)
spheres = sphere_type.field(shape(n,))

vs

center = ti.Vector.field(3, ti.f32)
radius =  ti.field(ti.f32)
ti.root.dense(ti.i, n).place(center, radius)

[victoriacity]

A field of compound struct is stored in the same way as a vector or matrix field. By default, they are stored as an array of structures. In fact, the backends do not distinguish a struct field from a collection of fields of each struct member, and the member access logic is completely done in the Python frontend.

taichi/python/taichi/lang/struct.py

Lines 298 to 339 in 8a3a2d7

	for key, dtype in members.items():
	field_name = name + '.' + key
	if isinstance(dtype, CompoundType):
	field_dict[key] = dtype.field(shape=None,
	name=field_name,
	offset=offset,
	needs_grad=needs_grad)
	else:
	field_dict[key] = impl.field(dtype,
	shape=None,
	name=field_name,
	offset=offset,
	needs_grad=needs_grad)
	if shape is not None:
	if isinstance(shape, numbers.Number):
	shape = (shape, )
	if isinstance(offset, numbers.Number):
	offset = (offset, )
	if offset is not None and len(shape) != len(offset):
	raise TaichiSyntaxError(
	f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
	)
	dim = len(shape)
	if layout == Layout.SOA:
	for e in field_dict.values():
	impl.root.dense(impl.index_nd(dim),
	shape).place(e, offset=offset)
	if needs_grad:
	for e in field_dict.values():
	impl.root.dense(impl.index_nd(dim),
	shape).place(e.grad, offset=offset)
	else:
	impl.root.dense(impl.index_nd(dim),
	shape).place(*tuple(field_dict.values()),
	offset=offset)
	if needs_grad:
	grads = tuple(e.grad for e in field_dict.values())
	impl.root.dense(impl.index_nd(dim),
	shape).place(*grads, offset=offset)
	return StructField(field_dict, name=name)

That is, these two snippets generate the exact same data structure:

sphere_type = ti.types.struct(center=ti.Vector(3, ti.f32), radius=ti.f32)
spheres = sphere_type.field(shape(n,))

center = ti.Vector.field(3, ti.f32)
radius =  ti.field(ti.f32)
ti.root.dense(ti.i, n).place(center, radius)

[bsavery]

Oh @victoriacity I didn't see that optional parameter! That would be good to have in the docs!