perf: restore the nominal isinstance fast path for concrete nplike and backend classes

[ikrommyd]

PublicSingleton and ArrayLike inherit from Protocol, so everything descending from them (backends, nplikes, typetracer arrays) gets typing._ProtocolMeta.__instancecheck__. That one is written in python and costs about 4x a plain class check, even though for a non-protocol subclass it only does the ordinary nominal test. None of these classes is ever checked structurally, so they get a metaclass that puts type.__instancecheck__ back.

One thing to note: they can no longer be given virtual subclasses through abc's register(). Nothing in awkward, uproot, dask-awkward, coffea or vector does that.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.88%. Comparing base (60390ac) to head (1cac8a2).

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/_backends/backend.py	92.10% <100.00%> (ø)
src/awkward/_nplikes/array_like.py	100.00% <100.00%> (ø)
src/awkward/_nplikes/array_module.py	93.46% <100.00%> (ø)
src/awkward/_nplikes/typetracer.py	77.18% <100.00%> (ø)
src/awkward/_typing.py	88.88% <100.00%> (+2.22%)	⬆️

[ikrommyd]

🤖 AI text below 🤖

CPU timings, main vs this branch, median of 7 repeats per case.

case	base (us)	branch (us)	change
isinstance(nplike, ArrayModuleNumpyLike)	0.13	0.03	-75.4%
isinstance(backend, Backend)	0.12	0.03	-73.6%
isinstance(nplike, TypeTracer) [False]	0.12	0.04	-69.4%
isinstance(ndarray, TypeTracerArray) [False]	0.12	0.04	-68.4%
small_flat + small_flat	65.03	62.40	-4.0%

Median across all cases: -11.0% (17 cases total, showing the four best and the worst).

benchmark script

# Benchmark for PR #4339: nominal isinstance fast path for nplike/backend classes.
# Runs unchanged on both base and head.
from __future__ import annotations

import statistics
import timeit

import numpy as np

import awkward as ak
from awkward._backends.backend import Backend
from awkward._backends.numpy import NumpyBackend
from awkward._nplikes.array_module import ArrayModuleNumpyLike
from awkward._nplikes.numpy import Numpy
from awkward._nplikes.typetracer import TypeTracer, TypeTracerArray


def report(label, stmt, glb, *, number, repeats=7):
    ts = timeit.repeat(stmt, number=number, repeat=repeats, globals=glb)
    us = sorted(t / number * 1e6 for t in ts)
    print(f"{label:<45} {statistics.median(us):10.3f} us  min {us[0]:10.3f}  sd {statistics.stdev(us):8.3f}  (n={number}x{repeats})")


print(f"awkward {ak.__version__} from {ak.__file__}")

small = ak.Array([[1, 2, 3], [], [4, 5]])
small_flat = ak.Array([1.0, 2.0, 3.0, 4.0])
tracer = ak.Array(small.layout.to_typetracer(forget_length=True))

nplike = Numpy.instance()
backend = NumpyBackend.instance()
tt_data = tracer.layout.content.data
ndarray = np.arange(4, dtype=np.int64)

g = globals()

print("\n-- microbenchmarks: the isinstance checks themselves --")
report("isinstance(nplike, ArrayModuleNumpyLike)", "isinstance(nplike, ArrayModuleNumpyLike)", g, number=200_000)
report("isinstance(backend, Backend)", "isinstance(backend, Backend)", g, number=200_000)
report("isinstance(tt_data, TypeTracerArray)", "isinstance(tt_data, TypeTracerArray)", g, number=200_000)
report("isinstance(ndarray, TypeTracerArray) [False]", "isinstance(ndarray, TypeTracerArray)", g, number=200_000)
report("isinstance(nplike, TypeTracer) [False]", "isinstance(nplike, TypeTracer)", g, number=200_000)
report("baseline: isinstance(ndarray, np.ndarray)", "isinstance(ndarray, np.ndarray)", g, number=200_000)

print("\n-- high-level ops on small arrays (dispatch-dominated) --")
report("small[1]", "small[1]", g, number=20_000)
report("small[:, 0:1]", "small[:, 0:1]", g, number=5_000)
report("small_flat[2]", "small_flat[2]", g, number=20_000)
report("ak.num(small)", "ak.num(small)", g, number=5_000)
report("small * 2", "small * 2", g, number=5_000)
report("small_flat + small_flat", "small_flat + small_flat", g, number=5_000)
report("ak.sum(small, axis=-1)", "ak.sum(small, axis=-1)", g, number=2_000)
report("ak.Array([[1, 2, 3], [], [4, 5]])", "ak.Array([[1, 2, 3], [], [4, 5]])", g, number=2_000)
report("small.to_list()", "small.to_list()", g, number=5_000)

print("\n-- typetracer backend (all-python, no kernels) --")
report("ak.num(tracer)", "ak.num(tracer)", g, number=2_000)
report("tracer * 2", "tracer * 2", g, number=2_000)

[github-actions]

The documentation preview is ready to be viewed at https://awkward-array.org/doc/pr/4339/

[ikrommyd]

@henryiii do you know of a better way to do this? Those isinstance checks are a killer and something definitely needs to be done.