Skip to content

Commit 9f456ae

Browse files
committed
More simplifications to the understanding jit decorator optim tip
1 parent ca601fb commit 9f456ae

1 file changed

Lines changed: 8 additions & 12 deletions

File tree

doc/guides/optimization_tips.md

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -130,27 +130,26 @@ The benchmark times four variants — the Python loop above, a vectorized NumPy
130130

131131
![Mandelbrot escape times across jit modes](optim_tips/tip_15a_jit_control_flow.png)
132132

133-
The default `@blosc2.jit` is **~155x faster than the Python loop** and **~2.3x faster than the best plain-NumPy version** (with `jit_backend="cc"`, ~317x and ~4.7x), and it is simply the loop you would have written anyway. The NumPy version, by contrast, costs ~15 lines of `alive`/`escaped` mask bookkeeping plus an overflow trap for the pixels that already escaped.
133+
The default `@blosc2.jit` is **~160x faster than the Python loop** and **~2.3x faster than the best plain-NumPy version** (with `jit_backend="cc"`, ~320x and ~4.7x).
134134

135135
Operands may equally be on-disk {class}`~blosc2.NDArray` objects: the same kernel over `blosc2.asarray()` views of the two grids runs at the same speed and returns a plain NumPy array.
136136

137137
#### Gotcha: silent fallback to tracing
138138

139-
The DSL grammar is narrower than Python, and a body that misses it falls back to tracing *silently* — you only find out when the call fails. Two rules bite in this example:
140-
141-
- **Simple assignments only.** The tuple assignment `zr, zi = ...` of the Python reference is not valid DSL; that is why the jit version uses a `zr2` temporary.
142-
- **No docstring in the kernel body.**
139+
The DSL grammar is narrower than Python, and a body that misses it falls back to tracing *silently* — you only find out when the call fails. For example, the tuple assignment `zr, zi = ...` of the Python reference is not valid DSL program; that is why the jit version uses a `zr2` temporary.
143140

144141
The [DSL syntax reference](../reference/dsl_syntax.md) has the full grammar. To turn that silent fallback into a hard error, decorate with `@blosc2.jit(strict=True)`: it forces the DSL route and raises `DSLSyntaxError` at *decoration* time if the function cannot be compiled.
145142

146143
The opposite knob, `strict=False`, forces tracing even when there is control flow. On this function it fails loudly — tracing evaluates `if zr * zr + zi * zi > 4.0` on an array, which raises `ValueError`. The dangerous case is a branch on a plain Python value instead (say `if max_iter > 100`): tracing records only the path that one call happened to take, and that path is then reused for every element, quietly. So use `strict=False` only when the branches depend on Python values, never on the arrays.
147144

148145
### Element-wise functions still trace by default
149146

150-
If compiling the whole function is this good, why doesn't `jit` do it for everything? Because without control flow, tracing wins: the traced expression is evaluated as one vectorized miniexpr over whole chunks, while a compiled kernel has to loop element by element. Here is a heavy elementwise mix of transcendental functions:
147+
If compiling the whole function is this good, why doesn't `jit` do it for everything? Because without control flow, tracing usually wins: the traced expression is evaluated as one vectorized miniexpr over whole chunks, while a compiled kernel has to loop element by element.
148+
149+
Here is a heavy elementwise mix of transcendental functions:
151150

152151
```python
153-
@blosc2.jit
152+
@blosc2.jit # using jit(strict=True) forces the DSL route
154153
def heavy(x):
155154
return (
156155
np.sin(x)
@@ -159,9 +158,6 @@ def heavy(x):
159158
+ np.sqrt(np.abs(x))
160159
+ np.log1p(np.abs(x))
161160
)
162-
163-
164-
# The same body decorated with @blosc2.jit(strict=True) forces the DSL route.
165161
```
166162

167163
Measured over 8M float32 values:
@@ -172,7 +168,7 @@ The first bar is the same expression in plain NumPy, as a scale anchor: all thre
172168

173169
Among the three, the default `@blosc2.jit` (which traces) is the faster route here: forcing the DSL route with `strict=True` is **~1.37x slower** with the bundled tcc, and gets close to tracing when compiled with `jit_backend="cc"` (still ~1.07x slower; the plot shows the exact times). Note that `jit_backend="cc"` alone does *not* switch an elementwise function to the compiled route; it keeps tracing, at the same speed.
174170

175-
So, use `strict=True` when you want the compiled-kernel guarantee, even if sometimes it may cost you speed.
171+
So, use `strict=True` when you want the compiled-kernel guarantee; but be aware that usually it may cost you speed.
176172

177173
### Pros and cons of forcing the system compiler
178174

@@ -185,7 +181,7 @@ The price is the one-time compile. The benchmark measures it as the first call m
185181
| mandelbrot | 3.2 ms | 233 ms | 1.3 ms |
186182
| elementwise | 4.0 ms | 253 ms | 2.9 ms |
187183

188-
[tcc](https://bellard.org/tcc/) compiles in memory, so every process pays those few milliseconds again. `cc` writes a shared object into `$TMPDIR/miniexpr-jit`, keyed by a fingerprint of the kernel, its dtypes and the toolchain: only the *first* process on a machine pays the compiler, and later ones just load the cached artifact — cheaper even than tcc's in-memory compile. And that cold compile is repaid by the faster steady state after ~54 calls of the mandelbrot kernel (~42 of the elementwise one).
184+
[tcc](https://bellard.org/tcc/) compiles in memory, so every process pays those few milliseconds again. `cc` writes a shared object into `$TMPDIR/miniexpr-jit`, keyed by a fingerprint of the kernel, its dtypes and the toolchain: only the *first* process on a machine pays the compiler.
189185

190186
So `cc` pays off for kernels you call repeatedly in the same run (or across runs), or when run time is much larger than compile time. It also requires a C compiler and a writable cache directory.
191187

0 commit comments

Comments
 (0)