fandhe_ai/inference/
batch.rs

1//! DataLoader 反復推論（`Sequential::predict_batches`）のフェーズ計測
2//! （イシュー #2192・#2582・親 #2131・#2581）。
3//!
4//! `compat::sequential::Sequential::predict_batches`（`compat/sequential.rs`）
5//! が本モジュールの型・thread-local アキュムレータを使って、
6//! `Sequential::predict` をバッチ単位で反復する際の内訳（データロード・
7//! tape 構築・forward・デバイス転送の各フェーズの累計時間・呼び出し回数）を
8//! 記録する。
9//!
10//! 本ファイルは `pub(crate)` モジュールで、公開項目（`PhaseMetrics`・
11//! `PhaseStat`・`InferencePhase`・`PredictBatchInput`・`get_phase_metrics`・
12//! `reset_phase_metrics`）は `crate::inference`（`pub mod`）が `pub use` で
13//! フラットに公開する（`docs/facade-predict-batches-phase-metrics-decision.md`
14//! §8.4 (e)）。`PhaseRecorder`・`NoopPhaseRecorder`・`TimingPhaseRecorder`・
15//! `merge_inference_phase_stats` は crate 内専用に留める。
16
17/// 推論の各フェーズの識別子（`PhaseMetrics::phase` の引数）。
18///
19/// `#[non_exhaustive]`（variant の追加は非破壊）。`Sequential::predict_batches`
20/// の内部実装が計測点ごとに指定する。
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22#[non_exhaustive]
23pub enum InferencePhase {
24    /// `DataLoader::iter`／`Batches::next` の呼び出し（バッチの切り出し）。
25    DataLoad,
26    /// tape 経路（`predict_via_tape`）の `crate::tape()`。tape 不要経路
27    /// では計上されない（`calls == 0`）。
28    TapeBuild,
29    /// 層の forward 実行（tape 不要経路は `forward_host` 連鎖、tape 経路
30    /// は `tape.var(..)`・`forward`・`to_tensor()`）。
31    Forward,
32    /// デバイス常駐経路のホスト⇔デバイス転送用に予約された識別子。CPU 固定
33    /// の `predict_batches` 経路では転送が発生しないため**常に `calls == 0`**
34    /// （`docs/facade-predict-batches-phase-metrics-decision.md` §2.1・§8.4 (c)）。
35    DeviceTransfer,
36}
37
38/// 推論フェーズの計測点を差し込むためのトレイト（`predict_recorded`・
39/// `predict_batches` の内部実装が経由する。crate 内専用）。単相化される
40/// ため [`NoopPhaseRecorder`] 使用時（既定の `Sequential::predict`）は
41/// オーバーヘッドを追加しない。
42pub(crate) trait PhaseRecorder {
43    /// `phase` の経過時間を計測しつつ `f` を実行し、その戻り値を返す。
44    fn record<T>(&mut self, phase: InferencePhase, f: impl FnOnce() -> T) -> T;
45}
46
47/// 既定の `Sequential::predict`（単体呼び出し）が使う no-op recorder。
48/// hot path への計測オーバーヘッドを避けるため、単体 `predict` は常時
49/// 計測しない設計（`docs/facade-predict-batches-phase-metrics-decision.md`
50/// §6 スコープ外）。
51pub(crate) struct NoopPhaseRecorder;
52
53impl PhaseRecorder for NoopPhaseRecorder {
54    #[inline]
55    fn record<T>(&mut self, _phase: InferencePhase, f: impl FnOnce() -> T) -> T {
56        f()
57    }
58}
59
60// 公開 5 名（`crate::inference` が `InferencePhase` と合わせて `pub use` する）。
61// `pub(crate) use` は `pub use` で再公開できない（E0364）ため、公開分と crate
62// 内専用分を分けて再エクスポートする。`use` は `mod recorded` 宣言より前に置く
63// （内部の `#[cfg(test)] mod tests` より後ろに項目を置くと clippy
64// `items_after_test_module` に抵触するため。順序は到達可能性に影響しない）。
65pub use recorded::{
66    PhaseMetrics, PhaseStat, PredictBatchInput, get_phase_metrics, reset_phase_metrics,
67};
68pub(crate) use recorded::{TimingPhaseRecorder, merge_inference_phase_stats};
69
70/// 計測値の型・thread-local アキュムレータ・入力トレイトの実体。
71mod recorded {
72    use std::cell::Cell;
73    use std::time::{Duration, Instant};
74
75    use super::{InferencePhase, PhaseRecorder};
76    use crate::Tensor;
77
78    /// 1 つの推論フェーズの累計時間・呼び出し回数。フィールドは非公開で、
79    /// accessor（[`Self::total_micros`]・[`Self::calls`]・[`Self::total`]）
80    /// で読む。加算・減算は飽和演算で panic しない（`Duration::MAX`／
81    /// `u64::MAX` 付近でも安全）。
82    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
83    pub struct PhaseStat {
84        total: Duration,
85        calls: u64,
86    }
87
88    impl PhaseStat {
89        /// 1 回分の経過時間を累積する（飽和加算。`Duration::MAX` 付近
90        /// でも panic しない）。
91        fn add(&mut self, elapsed: Duration) {
92            self.total = self.total.saturating_add(elapsed);
93            self.calls = self.calls.saturating_add(1);
94        }
95
96        /// 累計時間をマイクロ秒で返す（`u128`。`Duration::as_micros` と
97        /// 同じ精度で、呼び出し側が浮動小数へ変換しやすい形）。
98        pub fn total_micros(&self) -> u128 {
99            self.total.as_micros()
100        }
101
102        /// 累計呼び出し回数。
103        pub fn calls(&self) -> u64 {
104            self.calls
105        }
106
107        /// 累計時間を `Duration` で返す（`total_micros` と同じ値の
108        /// ナノ秒精度版）。
109        pub fn total(&self) -> Duration {
110            self.total
111        }
112
113        /// 飽和減算（`before` より小さくならない。before/after 差分用）。
114        fn saturating_sub(&self, before: &Self) -> Self {
115            Self {
116                total: self.total.saturating_sub(before.total),
117                calls: self.calls.saturating_sub(before.calls),
118            }
119        }
120
121        /// 2 つの `PhaseStat` を飽和加算で合算する
122        /// （`PhaseMetrics::merge` が各フィールドに適用する）。
123        fn merged(&self, other: &Self) -> Self {
124            Self {
125                total: self.total.saturating_add(other.total),
126                calls: self.calls.saturating_add(other.calls),
127            }
128        }
129    }
130
131    /// `Sequential::predict_batches` が計測した、フェーズ別の集計
132    /// （呼び出し 1 回分、または thread-local への蓄積後）。
133    ///
134    /// フィールドは非公開で accessor（[`Self::phase`]・[`Self::total`]・
135    /// [`Self::batches`]・[`Self::samples`]・[`Self::since`]）で読む。
136    /// `#[non_exhaustive]`（フィールド追加は非破壊）。`batches`／`samples`
137    /// は反復したバッチ数・総サンプル数（入力バッチの先頭軸から数える。
138    /// rank 0 バッチは 1 件として飽和加算）。
139    ///
140    /// # Examples
141    ///
142    /// ラベル付きデータセット（`(features, labels)`）でもラベルは無視され、特徴量だけが
143    /// 推論される。呼び出し前後のスナップショット差分は [`Self::since`] で取る
144    /// （時間値は環境依存のため、ここでは回数のみ確認する）。
145    ///
146    /// ```
147    /// use fandhe_ai::Tensor;
148    /// use fandhe_ai::compat::Sequential;
149    /// use fandhe_ai::data::{DataLoader, DataLoaderConfig, TensorDataset};
150    /// use fandhe_ai::inference::{InferencePhase, PhaseMetrics, get_phase_metrics};
151    ///
152    /// let model = Sequential::new().add_linear(2, 1, 7).unwrap();
153    /// let x = Tensor::<f32>::new(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]).unwrap();
154    /// let y = Tensor::<i32>::new(vec![0, 1], &[2]).unwrap();
155    /// let loader = DataLoader::new(
156    ///     (TensorDataset::new(x).unwrap(), TensorDataset::new(y).unwrap()),
157    ///     DataLoaderConfig::new(1),
158    /// )
159    /// .unwrap();
160    ///
161    /// let before: PhaseMetrics = get_phase_metrics();
162    /// let outputs = model.predict_batches(&loader).unwrap();
163    /// let delta = get_phase_metrics().since(&before);
164    ///
165    /// assert_eq!(outputs.len(), 2);
166    /// assert_eq!(delta.batches(), 2);
167    /// assert_eq!(delta.samples(), 2);
168    /// assert!(delta.phase(InferencePhase::DataLoad).calls() >= 2);
169    /// let _micros: u128 = delta.phase(InferencePhase::Forward).total_micros();
170    /// assert_eq!(delta.phase(InferencePhase::DeviceTransfer).calls(), 0);
171    /// ```
172    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
173    #[non_exhaustive]
174    pub struct PhaseMetrics {
175        data_load: PhaseStat,
176        tape_build: PhaseStat,
177        forward: PhaseStat,
178        device_transfer: PhaseStat,
179        /// `predict_batches` 呼び出し全体（成功・失敗を問わない）の
180        /// 経過時間。個々のフェーズの合計とは独立に計測する（フェーズ間
181        /// の計測されない隙間時間も含むため）。
182        total: PhaseStat,
183        batches: u64,
184        samples: u64,
185    }
186
187    impl PhaseMetrics {
188        fn phase_mut(&mut self, phase: InferencePhase) -> &mut PhaseStat {
189            match phase {
190                InferencePhase::DataLoad => &mut self.data_load,
191                InferencePhase::TapeBuild => &mut self.tape_build,
192                InferencePhase::Forward => &mut self.forward,
193                InferencePhase::DeviceTransfer => &mut self.device_transfer,
194            }
195        }
196
197        /// 指定フェーズの集計値を返す。
198        pub fn phase(&self, phase: InferencePhase) -> PhaseStat {
199            match phase {
200                InferencePhase::DataLoad => self.data_load,
201                InferencePhase::TapeBuild => self.tape_build,
202                InferencePhase::Forward => self.forward,
203                InferencePhase::DeviceTransfer => self.device_transfer,
204            }
205        }
206
207        /// `predict_batches` 呼び出し全体の集計値（`calls` は呼び出し回数）。
208        pub fn total(&self) -> PhaseStat {
209            self.total
210        }
211
212        /// 反復したバッチ数の累計。
213        pub fn batches(&self) -> u64 {
214            self.batches
215        }
216
217        /// 反復した総サンプル数の累計（飽和加算）。
218        pub fn samples(&self) -> u64 {
219            self.samples
220        }
221
222        /// `other` を自身へ飽和加算で合算する（thread-local への蓄積用。
223        /// `predict_batches` は途中で失敗しても、それまでに計測した分を
224        /// 失わないようこのメソッドで thread-local へ merge してから返す）。
225        fn merge(&mut self, other: &Self) {
226            self.data_load = self.data_load.merged(&other.data_load);
227            self.tape_build = self.tape_build.merged(&other.tape_build);
228            self.forward = self.forward.merged(&other.forward);
229            self.device_transfer = self.device_transfer.merged(&other.device_transfer);
230            self.total = self.total.merged(&other.total);
231            self.batches = self.batches.saturating_add(other.batches);
232            self.samples = self.samples.saturating_add(other.samples);
233        }
234
235        /// `self`（after）と `before` の飽和差分を返す（before/after
236        /// スナップショット比較。`before` が大きい項目は 0 になる）。
237        pub fn since(&self, before: &Self) -> Self {
238            Self {
239                data_load: self.data_load.saturating_sub(&before.data_load),
240                tape_build: self.tape_build.saturating_sub(&before.tape_build),
241                forward: self.forward.saturating_sub(&before.forward),
242                device_transfer: self.device_transfer.saturating_sub(&before.device_transfer),
243                total: self.total.saturating_sub(&before.total),
244                batches: self.batches.saturating_sub(before.batches),
245                samples: self.samples.saturating_sub(before.samples),
246            }
247        }
248
249        fn record_batch(&mut self, samples: u64) {
250            self.batches = self.batches.saturating_add(1);
251            self.samples = self.samples.saturating_add(samples);
252        }
253    }
254
255    /// `predict_batches` が使う実計測 recorder（crate 内専用）。内部に累積した
256    /// [`PhaseMetrics`] を [`Self::into_metrics`] で取り出す。
257    #[derive(Debug, Clone, Copy, Default)]
258    pub(crate) struct TimingPhaseRecorder {
259        metrics: PhaseMetrics,
260    }
261
262    impl TimingPhaseRecorder {
263        pub(crate) fn into_metrics(self) -> PhaseMetrics {
264            self.metrics
265        }
266
267        /// バッチ 1 件分の反復を記録する（`DataLoad`／`Forward` 等の
268        /// `record` とは独立に、バッチ数・サンプル数を積む）。
269        pub(crate) fn record_batch(&mut self, samples: u64) {
270            self.metrics.record_batch(samples);
271        }
272
273        /// `total`（呼び出し全体）の経過時間を記録する。
274        pub(crate) fn record_total(&mut self, elapsed: Duration) {
275            self.metrics.total.add(elapsed);
276        }
277    }
278
279    impl PhaseRecorder for TimingPhaseRecorder {
280        fn record<T>(&mut self, phase: InferencePhase, f: impl FnOnce() -> T) -> T {
281            let start = Instant::now();
282            let out = f();
283            self.metrics.phase_mut(phase).add(start.elapsed());
284            out
285        }
286    }
287
288    thread_local! {
289        /// 呼び出しスレッドごとの累計 phase 計測値。プロセス全体ではなく
290        /// スレッド単位にする理由: ホットパスにロックが不要でテスト間の
291        /// 干渉を避け、「このスレッドが `predict_batches` で計測した
292        /// 累計」という単純な意味論にするため（`docs/facade-predict-
293        /// batches-phase-metrics-decision.md` §2.3・§8.4 (f)）。
294        static INFERENCE_PHASE_METRICS: Cell<PhaseMetrics> =
295            Cell::new(PhaseMetrics::default());
296    }
297
298    /// 呼び出しスレッドが `Sequential::predict_batches` で計測した累計
299    /// phase 計測値のスナップショットを返す（スレッド単位。他スレッドの
300    /// 計測は含まない）。呼び出し前後の差分は [`PhaseMetrics::since`] で
301    /// 取る。
302    pub fn get_phase_metrics() -> PhaseMetrics {
303        INFERENCE_PHASE_METRICS.with(|cell| cell.get())
304    }
305
306    /// 呼び出しスレッドの累計 phase 計測値をゼロへ戻す（他スレッドの
307    /// 累計には影響しない）。
308    pub fn reset_phase_metrics() {
309        INFERENCE_PHASE_METRICS.with(|cell| cell.set(PhaseMetrics::default()));
310    }
311
312    /// `predict_batches` の呼び出し 1 回分の計測値を、呼び出し
313    /// スレッドの累計へ飽和加算で反映する（途中失敗した試行の計測値も
314    /// 失わない。crate 内専用）。
315    pub(crate) fn merge_inference_phase_stats(delta: &PhaseMetrics) {
316        INFERENCE_PHASE_METRICS.with(|cell| {
317            let mut current = cell.get();
318            current.merge(delta);
319            cell.set(current);
320        });
321    }
322
323    /// 封印用の private supertrait（`PredictBatchInput` の外部実装を
324    /// 禁止する。後から unseal するのは非破壊だが逆は破壊的なため。
325    /// 決定記録 §8.4 (a)）。
326    mod sealed {
327        pub trait Sealed {}
328    }
329
330    /// `DataLoader<D>` が生成するバッチ型から、推論対象の `Tensor<f32>` を
331    /// 取り出すための封印済みトレイト。`Sequential::predict_batches` の
332    /// `D::Batch` 境界に使う。ラベル付きデータセット
333    /// （`(Tensor<f32>, B)`／`(Tensor<f32>, B, C)`）ではラベルを無視する
334    /// （Keras `model.predict(dataset)` が特徴量のみを使うのと同じ扱い。
335    /// 決定記録 §3）。
336    ///
337    /// 実装は `Tensor<f32>`・`(Tensor<f32>, B)`・`(Tensor<f32>, B, C)` の
338    /// 3 形に固定され、利用者が実装することも、メソッドを呼ぶ必要もない。
339    pub trait PredictBatchInput: sealed::Sealed {
340        /// 推論入力の `Tensor<f32>` への参照（ラベルは含まない）。
341        fn inference_input(&self) -> &Tensor<f32>;
342    }
343
344    impl sealed::Sealed for Tensor<f32> {}
345    impl<B> sealed::Sealed for (Tensor<f32>, B) {}
346    impl<B, C> sealed::Sealed for (Tensor<f32>, B, C) {}
347
348    impl PredictBatchInput for Tensor<f32> {
349        fn inference_input(&self) -> &Tensor<f32> {
350            self
351        }
352    }
353
354    impl<B> PredictBatchInput for (Tensor<f32>, B) {
355        fn inference_input(&self) -> &Tensor<f32> {
356            &self.0
357        }
358    }
359
360    impl<B, C> PredictBatchInput for (Tensor<f32>, B, C) {
361        fn inference_input(&self) -> &Tensor<f32> {
362            &self.0
363        }
364    }
365
366    #[cfg(test)]
367    mod tests {
368        use super::*;
369
370        #[test]
371        fn phase_stat_add_accumulates_total_and_calls() {
372            let mut stat = PhaseStat::default();
373            stat.add(Duration::from_micros(10));
374            stat.add(Duration::from_micros(5));
375            assert_eq!(stat.total_micros(), 15);
376            assert_eq!(stat.calls(), 2);
377        }
378
379        #[test]
380        fn phase_stat_total_matches_total_micros() {
381            let mut stat = PhaseStat::default();
382            stat.add(Duration::from_micros(1500));
383            assert_eq!(stat.total(), Duration::from_micros(1500));
384            assert_eq!(stat.total().as_micros(), stat.total_micros());
385        }
386
387        #[test]
388        fn phase_stat_add_saturates_instead_of_panicking() {
389            let mut stat = PhaseStat {
390                total: Duration::MAX,
391                calls: u64::MAX,
392            };
393            stat.add(Duration::from_secs(1));
394            assert_eq!(stat.total, Duration::MAX);
395            assert_eq!(stat.calls, u64::MAX);
396        }
397
398        #[test]
399        fn phase_stat_saturating_sub_never_underflows() {
400            let before = PhaseStat {
401                total: Duration::from_micros(100),
402                calls: 10,
403            };
404            let after = PhaseStat {
405                total: Duration::from_micros(50),
406                calls: 3,
407            };
408            let delta = after.saturating_sub(&before);
409            assert_eq!(delta.total, Duration::ZERO);
410            assert_eq!(delta.calls, 0);
411        }
412
413        #[test]
414        fn phase_metrics_since_returns_only_the_later_delta() {
415            let mut a = PhaseMetrics::default();
416            a.phase_mut(InferencePhase::Forward)
417                .add(Duration::from_micros(10));
418            a.record_batch(4);
419
420            let mut b = a;
421            b.phase_mut(InferencePhase::Forward)
422                .add(Duration::from_micros(7));
423            b.record_batch(4);
424
425            let delta = b.since(&a);
426            assert_eq!(delta.phase(InferencePhase::Forward).calls(), 1);
427            assert_eq!(delta.phase(InferencePhase::Forward).total_micros(), 7);
428            assert_eq!(delta.batches(), 1);
429            assert_eq!(delta.samples(), 4);
430        }
431
432        #[test]
433        fn merge_accumulates_across_multiple_deltas() {
434            let mut acc = PhaseMetrics::default();
435            let mut delta1 = PhaseMetrics::default();
436            delta1
437                .phase_mut(InferencePhase::DataLoad)
438                .add(Duration::from_micros(3));
439            delta1.record_batch(2);
440            let mut delta2 = PhaseMetrics::default();
441            delta2
442                .phase_mut(InferencePhase::DataLoad)
443                .add(Duration::from_micros(4));
444            delta2.record_batch(3);
445
446            acc.merge(&delta1);
447            acc.merge(&delta2);
448
449            assert_eq!(acc.phase(InferencePhase::DataLoad).calls(), 2);
450            assert_eq!(acc.phase(InferencePhase::DataLoad).total_micros(), 7);
451            assert_eq!(acc.batches(), 2);
452            assert_eq!(acc.samples(), 5);
453        }
454
455        #[test]
456        fn thread_local_snapshot_is_isolated_per_thread() {
457            reset_phase_metrics();
458            let mut delta = PhaseMetrics::default();
459            delta
460                .phase_mut(InferencePhase::Forward)
461                .add(Duration::from_micros(1));
462            delta.record_batch(1);
463            merge_inference_phase_stats(&delta);
464
465            let main_snapshot = get_phase_metrics();
466            assert_eq!(main_snapshot.batches(), 1);
467
468            let other_thread_snapshot = std::thread::spawn(get_phase_metrics).join().unwrap();
469            // 別スレッドの thread-local は常にゼロから始まる
470            // （このスレッドの蓄積が漏れ伝わらない）。
471            assert_eq!(other_thread_snapshot.batches(), 0);
472
473            reset_phase_metrics();
474            assert_eq!(get_phase_metrics().batches(), 0);
475        }
476
477        #[test]
478        fn predict_batch_input_ignores_labels() {
479            let x = Tensor::<f32>::new(vec![1.0, 2.0], &[2]).unwrap();
480            let y = 42i32;
481            let pair = (x.clone(), y);
482            assert_eq!(pair.inference_input().shape(), x.shape());
483            let triple = (x.clone(), y, "ignored");
484            assert_eq!(triple.inference_input().shape(), x.shape());
485            assert_eq!(x.inference_input().shape(), x.shape());
486        }
487    }
488}
fandhe_ai/inference/batch.rs

fandhe_ai/inference/
batch.rs
