np.addのトレース、すべてを下まで

Veit's Blog np.addのトレース、すべてを下まで 2026-09-02 この記事のメモは、半年間ドラフトフォルダにありました。過去1年間、NumPy自体に少し手を加えましたが、目立ったものではありませんが、ソースコードを調べるのに十分でした。それがこの記事を書くきっかけになったのですが、他の義務が私のNumPyへの貢献を凌駕してしまい、静かに腐り始めていました。NumPyのリリースが1つか2つ後に出て、ようやく再び手に取り、自分の足跡をたどり直し、ここにたどり着きました。 前提はこうです。np.add(a, b) は、世界で最も実行される数値Pythonの行の1つである可能性があり、ほとんどの人は「Cで配列を素早く加算する」という同等のメンタルモデルを持っています。そのモデルは正しいですが、Pythonの呼び出しと加算を行うループの間には多くの仕組みがあり、私はそれを分解するのが楽しい機械だと考えています。そこで今日は、2つのfloat64配列を持つ単一の呼び出し、np.add(a, b) を、PythonのエントリポイントからSIMDカーネルまで、実際のNumPyソースを読みながらトレースします。うまくいけば、多くのことを学べるでしょう！ 以下のすべては、私がこれを書いている時点での現在のリリースであるNumPy 2.5.2に固定されており、すべてのリンクはそのタグを指しています。内部はバージョン間で移動します1。したがって、自宅で一緒に探求している場合は、対応するタグを確認してください。Cを読むことに少なくともある程度慣れていることを前提としますが、NumPyの内部知識は必要ありません。それが私たちの目的です。 マップ 飛び込む前に、宝の地図を用意しました。これで、常にどこにいるかを知ることができます。 np.add(a, b) (Python) │ ▼ ufunc_generic_fastcall (C: 引数を解析) │ ▼ __array_ufunc__ オーバーライドチェック (他のライブラリに分岐する可能性あり) │ ▼ プロモーション & ディスパッチ (float64ループを見つけ、キャッシュする) │ ▼ トrivialループまたはNpyIter (イテレーション戦略) │ ▼ DOUBLE_add (実際の内部ループ、SIMD) これらはすべて以下のセクションです。一番上から始めましょう。 np.addはオブジェクトです 最初に知っておくべきことは、np.addは通常のPython関数ではないということです。それはnumpy.ufuncのインスタンスであり、Cで定義された型です2。 >>> type(np.add) <class 'numpy.ufunc'> >>> np.add.nin, np.add.nout (2, 1) >>> len(np.add.types) 22 >>> np.add.types[11:14] ['ee->e', 'ff->f', 'dd->d'] ufuncは、その核となる部分では、内部ループのバンドルです。サポートされている型シグネチャごとに1つの小さなC関数と、入力と出力の数に関するメタデータが含まれています。np.addは22個を出荷しますが、typesはクラシックなものだけをリストしていると言うべきです。最新の方法で登録されたループ（後でその区別について説明します）は、Pythonが決して見ることのないufuncの内部マッピングにあります。今日私たちが追跡しているのはdd->d、つまりdouble, doubleからdoubleへの変換です。 この投稿の残りの部分はすべて、NumPyがあなたの呼び出しからそのエントリに到達するまで、そしてそれが一度見つかったら何が起こるかについてです。他のパスは異なる場合があります。それは大きなキットです！洞窟に飛び込みましょう。 C言語へ Pythonがnp.add(a, b) を認識すると、ufuncオブジェクトを呼び出します。ufunc型はvectorcallプロトコルを実装しているため、呼び出しはufunc_generic_vectorcallに着地し、すぐに実際の作業馬であるufunc_generic_fastcallに転送されます。その関数は長いですが、チェックリストのように読め、それは操作全体の骨格です。大幅に省略されています。 static PyObject * ufunc_generic_fastcall(PyUFuncObject *ufunc, PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames, npy_bool outer) { /* ... 入力、出力、およびキーワード引数を抽出 ... */ /* ここでオーバーライドをチェックするために必要なすべての情報が得られました */ PyObject *override = NULL; errval = PyUFunc_CheckOverride(ufunc, method, full_args.in, full_args.out, where_obj, args, len_args, kwnames, &override); /* ... オーバーライドが見つかった場合は、その結果を返します ... */ /* ... 引数を配列に変換し、それらのDTypeを抽出します ... */ PyArrayMethodObject *ufuncimpl = promote_and_get_ufuncimpl(ufunc, operands, signature, operand_DTypes, ...); /* 正しいディスクリプタを操作のために見つけます */ if (resolve_descriptors(nop, ufunc, ufuncimpl, ...) < 0) { goto fail; } /* * 最終準備を行い、内部ループを呼び出します。 */ errval = PyUFunc_GenericFunctionInternal(ufunc, ufuncimpl, operation_descrs, operands, casting, order, wheremask); /* ... 出力をラップして返します ... */ } 解析、オーバーライドのチェック、ループの選択、実行、結果のラップ。計画ができたようです！では、興味深い部分に移りましょう。洞窟の光が弱まってきました。 エスケープハッチ NumPyが作業を行うことをコミットする前に、引数に自分でやりたいかどうかを尋ねます（常に良い操作方法です）。PyUFunc_CheckOverrideはすべての入力と出力をウォークし、デフォルト以外の__array_ufunc__メソッド（NEP 133で定義されているプロトコル）を探します。引数のいずれかにそれがある場合、NumPyはそれを呼び出し、それが生成したものを返します。そして、以下で説明する仕組みは決して実行されません。これは、np.add(dask_array, cupy_array) がサードパーティと連携して動作するフックです！DaskやCuPyのようなライブラリは__array_ufunc__を実装し、引き継ぎます。私たちは4行でそれを自分で演奏できます。 class Diverted: def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): return f"intercepted {ufunc.__name__}.{method}" >>> np.add(np.arange(3), Diverted()) 'intercepted add.__call__' 私たちのオブジェクトは、誰が計算するかについての議論に勝つこと以外は何もしません。私たちのトレースでは、両方の引数がプレーンなndarrayであると仮定するため、何もオーバーライドされず、ラペリングを続けます。 ループの選択 次に、NumPyは「2つのfloat64配列」から「そのdd->dエントリ」に到達する必要があります。これはプロモーションとディスパッチであり、dispatching.cppにあります。そのヘッダーコメントは、私がこれまでに見つけたプロセスに関する最高のドキュメントなので、そのまま引用します。タイポも含めて。 ディスパッチとプロモーションのプロセスは、次のステップに要約できます。 1. `signature`から`operand_DTypes`をオーバーライドします。 2. 新しい`operand_Dtypes`がキャッシュされているか確認します（キャッシュされている場合は、ステップ4に進みます）。 3. 最も一致する「ループ」を見つけます。これは、すべての`operand_DTypes`とループの`dtypes`に対する複数のディスパッチを使用して行われます。一致するループは、その`DTypes`が（定義されている）`operand_DTypes`のスーパーセットである必要があります。最も一致するループは、他のどのマッチングループよりも優れている必要があります。この結果はキャッシュされます。 4. 見つかったループがプロモーターの場合: プロモーターを呼び出します。それは現在の`operand_DTypes`を変更できます。その後、ステップ2に戻ります。 5. 最終的な`ArrayMethod`が見つかり、その登録された`dtypes`が`signature`にコピーされ、ufuncループで利用可能になります。 いくつかの翻訳が必要です。「シグネチャ」とは、np.add(a, b, dtype=...) を呼び出すときに明示的に固定するものです。私たちの呼び出しでは空です。「プロモーター」とは、ループが直接一致しないケースを処理するために、要求された型を書き直し、ディスパッチを再度実行させる登録済みのヘルパーです。 紛らわしいことに、日常的な混合ケース、np.add(int32_array, float64_array) はそれを使用さえしません。ディスパッチが空の場合、ufuncの古い型解決メカニズム（この場合はPyUFunc_AdditionTypeResolver）にフォールバックして共通の型を選択し、それらでディスパッチを再開してdd->dに着地します。そしてステップ2のキャッシュは非常に重要です！完全な解決は、特定の型の組み合わせでufuncを最初に呼び出したときにのみ発生します。通常のキャッシュ可能なケース（私たちのケース）では、同じ型での後続のすべての呼び出しは、DTypeクラスに対する単一のハッシュルックアップです。それは純粋な仕組みですが、便利です。 今持っているものを理解するために、考古学と言葉遊びに従事する必要があります。promote_and_get_ufuncimplが返すのはPyArrayMethodObject、つまり「具体的なDTypeに対する操作の具体的な実装」の最新（NEP 43以降）の表現です。しかし、float64加算の場合、ArrayMethodははるかに古いものの薄いラッパーです。実際のループが必要になると、get_wrapped_legacy_ufunc_loopはPyUFunc_DefaultLegacyInnerLoopSelectorを呼び出します。これは、最初の試行で書くであろうことと全く同じことをします！それはufuncの型テーブルをエントリごとにウォークし、dd->dを見つけるまで続行し、ufunc->functions[i]、つまりプレーンなC関数ポインタを返します。それを現在の（mo