Taichi Python Runtime Overhead Optimization (Ndarray)

[qiao-bo]

Currently, Taichi experiences an noticeable amount of python overhead when executing light kernels using Field or Ndarray. Using cProfile and/or pyinstrument, we can identify a few hot spots that incur this python runtime overhead. Here is an example using pyinstrument:

import taichi as ti
from pyinstrument import Profiler
from pyinstrument.renderers import ConsoleRenderer

ti.init(arch=ti.cuda)

N = 4096
a = ti.ndarray(ti.i32, N)

@ti.kernel
def fill_ones(a: ti.any_arr()):
    for i in a:
        a[i] = 1

def repeat_execution(n):
    for _ in range(n):
        fill_ones(a)

# initial run
repeat_execution(100)

# start profiling
profiler = Profiler()
profiler.start()
repeat_execution(10000)
session = profiler.stop()
profile_renderer = ConsoleRenderer(unicode=True, color=True, show_all=True)
print(profile_renderer.render(session))

On my PC with rtx 3080+i9 11900K, the output is

[Taichi] version 0.8.9, llvm 10.0.0, commit 345f0f04, linux, python 3.9.7
[Taichi] Starting on arch=cuda
[I 12/23/21 14:57:27.289 296821] [llvm_program.cpp:LlvmProgramImpl@84] CUDA max blocks per SM = 16

  _     ._   __/__   _ _  _  _ _/_   Recorded: 14:57:27  Samples:  205
 /_//_/// /_\ / //_// / //_'/ //     Duration: 0.205     CPU time: 0.206
/   _/                      v4.1.1

Program: test.py

0.205 <module>  test1.py:1
└─ 0.205 repeat_execution  test1.py:16
   └─ 0.203 wrapped  taichi/lang/kernel_impl.py:712
      ├─ 0.198 __call__  taichi/lang/kernel_impl.py:634
      │  ├─ 0.134 func__  taichi/lang/kernel_impl.py:485
      │  │  ├─ 0.123 [self]  
      │  │  ├─ 0.005 current_cfg  taichi/lang/impl.py:906
      │  │  │  └─ 0.005 PyCapsule.current_compile_config  <built-in>:0
      │  │  └─ 0.003 match_ext_arr  taichi/lang/kernel_impl.py:619
      │  ├─ 0.056 ensure_compiled  taichi/lang/kernel_impl.py:626
      │  │  ├─ 0.043 lookup  taichi/lang/kernel_impl.py:311
      │  │  │  ├─ 0.034 extract  taichi/lang/kernel_impl.py:305
      │  │  │  │  ├─ 0.024 extract_arg  taichi/lang/kernel_impl.py:257
      │  │  │  │  │  ├─ 0.013 [self]  
      │  │  │  │  │  ├─ 0.005 dtype  taichi/lang/_ndarray.py:53
      │  │  │  │  │  └─ 0.004 shape  taichi/lang/_ndarray.py:251
      │  │  │  │  └─ 0.009 [self]  
      │  │  │  ├─ 0.006 [self]  
      │  │  │  └─ 0.003 __hash__  enum.py:774
      │  │  └─ 0.011 materialize  taichi/lang/kernel_impl.py:430
      │  │     └─ 0.010 materialize  taichi/lang/impl.py:355
      │  │        ├─ 0.007 materialize_root_fb  taichi/lang/impl.py:303
      │  │        │  ├─ 0.004 __getattr__  taichi/lang/impl.py:544
      │  │        │  └─ 0.003 [self]  
      │  │        └─ 0.003 [self]  
      │  └─ 0.008 [self]  
      └─ 0.005 [self]  

Potential hot spots: func__ which executes the kernel (in the [self] time), or ensure_compiled. @ailzhang has commited a fix to remove some property used in Taichi (PR 3910). In addition, we can also check the cost of is_instance(). And also the performance of Matrix/Vector Ndarray, which may expose further bottlenecks.

comments/suggestions are welcome!

[ailzhang]

cc: @turbo0628 ;)

[turbo0628]

Tested on my laptop with LLVM backend. My evaluations confirm that hot spot resides in the lookup method invoked by ensure_compiled. The func__ overhead looks fine to me. I'm thinking about a faster look up method, or if we can merge / skip this look up for small kernel invokes. Any ideas?

Appendix - my profiling log:

[Taichi] version 0.8.9, llvm 10.0.0, commit 61120146, osx, python 3.8.12
[Taichi] Starting on arch=x64

  _     ._   __/__   _ _  _  _ _/_   Recorded: 14:01:34  Samples:  1582
 /_//_/// /_\ / //_// / //_'/ //     Duration: 1.620     CPU time: 5.498
/   _/                      v4.1.1

Program: runtime_overhead.py

1.620 <module>  runtime_overhead.py:1
└─ 1.620 repeat_execution  runtime_overhead.py:15
   └─ 1.604 wrapped  taichi/lang/kernel_impl.py:726
      └─ 1.594 __call__  taichi/lang/kernel_impl.py:651
         ├─ 1.226 func__  taichi/lang/kernel_impl.py:476
         │  ├─ 1.144 [self]  
         │  ├─ 0.038 current_cfg  taichi/lang/impl.py:879
         │  │  ├─ 0.020 [self]  
         │  │  └─ 0.018 PyCapsule.current_compile_config  <built-in>:0
         │  └─ 0.018 match_ext_arr  taichi/lang/kernel_impl.py:636
         ├─ 0.316 ensure_compiled  taichi/lang/kernel_impl.py:643
         │  ├─ 0.173 lookup  taichi/lang/kernel_impl.py:307
         │  │  ├─ 0.101 extract  taichi/lang/kernel_impl.py:301
         │  │  │  ├─ 0.057 extract_arg  taichi/lang/kernel_impl.py:253
         │  │  │  │  └─ 0.040 [self]  
         │  │  │  └─ 0.043 [self]  
         │  │  └─ 0.060 [self]  
         │  ├─ 0.119 materialize  taichi/lang/kernel_impl.py:423
         │  │  └─ 0.103 materialize  taichi/lang/impl.py:346
         │  │     ├─ 0.056 materialize_root_fb  taichi/lang/impl.py:290
         │  │     │  ├─ 0.032 [self]  
         │  │     │  └─ 0.024 __getattr__  taichi/lang/impl.py:523
         │  │     │     └─ 0.018 [self]  
         │  │     └─ 0.036 [self]  
         │  └─ 0.024 [self]  
         └─ 0.052 [self]  

[qiao-bo]

i don't think we can skip the look up, but it seems like the materialize only needs the key to search cached kernels, not the arg_features. something to optimize maybe?
There are a bunch of isinstance in def extract_arg(), they account for the 0.040s i think.
Have you tried GPUs? I wonder if func_ overhead will increase ;)

[ailzhang]

@turbo0628 For func__ it's actually pretty suspicious that we spend 1.144 in itself (it's mostly python checks and I'd expect them to run much much faster. :P

[qiao-bo]

@turbo0628 For func__ it's actually pretty suspicious that we spend 1.144 in itself (it's mostly python checks and I'd expect them to run much much faster. :P

that 1.114 includes the kernel time i think. The call -> func -> [self] in pyintrument output includes the Taichi kernel execution ;/

[ailzhang]

For isinstance I wonder if there's a way to refactor extract_arg so that we don't call isinstance any more.

[turbo0628]

@alizhang I get your point: the function body has a bunch of python checks, and it will always get invoked in func__. However, the [self] time (the 1.144 one) doesn't include the checks, it is the kernel time as @qiao-bo said. The python check overhead hides in the differences of func__ and [self].

[ailzhang]

@turbo0628 yup sorry I was wrong - 1.144 is indeed kernel time as you said.
It's a bit tricky here since we test on GPU backend where the kernel time is much smaller (so that python overhead is more visible).

[qiao-bo]

@alizhang I get your point: the function body has a bunch of python checks, and it will always get invoked in func__. However, the [self] time (the 1.144 one) doesn't include the checks, it is the kernel time as @qiao-bo said. The python check overhead hides in the differences of func__ and [self].

It would be interesting to know the difference though. is there a way to separate these two parts? either we let the profiler somehow recognizes t_kernel(launch_ctx) call, or manually do a timing breakdown? WDYT?

[turbo0628]

I think we have well picked the low-hanging fruits. Further optimization requires major adjustment to the ndarray launch methods and data structures, so let's put a brief summary here and think what can be done in the future.

Cause of the problem:

Every Taichi Ndarray sets an object in the args buffer, which are processed at kernel launch. For kernels that runs for hundreds or thousands milliseconds, this is fine because the access overhead is insignificant. However, if kernels are lightweight or even empty, the launch overhead becomes significant and unbearable. After investigating into the problem, we have found that the major overhead incurs in the Python scope, with two parts: launch context building and args buffer lookup.

What have been done

[PR #3947] has added two other C++ functions in order to collect the Ndarray related frequent tiny Python loops and pybind calls into single invoke to build the launch context. This PR almost cut the context building overhead in half.
[PR #3984] conservatively refactored the if clauses so as to skip some expensive checks and makes the code more readable. This PR brings another perf bump of ~10% with 6 Ndarrays.

What else can we do

After the above 2 PRs, the major overhead now moves to the lookup func. However, this function cannot easily get refactored as involved data structures are written in python. Therefore, if we want to optimize this function, we have to first refactor the Annotation python class into the taichi core part. This Annotation is indeed a utility class for frontend, we actually don't want this part of logics invade into taichi core. The perf gain is also minor.

Another possible opt is that we can try to skip the raw args buffer traversal, at least not for every kernel launch. Apparently, we can cache the keys somewhere in the code (maybe along with the cached kernels).

Lesson learnt

1. Pybind doesn't come for free. There's noticeable overhead especially when you are using some python-related funcs in the C++ code.
2. Python loops and if clauses are expensive. Carefully inspect all of them when you care about python overhead.
3. The line_profiler cannot deliver a precise overhead report. However, it is still useful as the relevant overhead numbers can reveal the hotspots. The profiled data are sometimes against our instinct, but the profiler never lies!

The optimization work is now pended. We might look back at this issue when we care about the python overhead again. Hope this brief summary helps :)