Skip to main content

rlx_ir/
capability.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! One table describing what every [`OpKind`] *is*.
6//!
7//! Passes constantly need to ask categorical questions about an op — is it
8//! elementwise, is it a fusion boundary, can it be folded on the host, does it
9//! carry a nested body. Historically each caller answered with its own
10//! `matches!` list: [`Op::is_elementwise`] and friends here in `rlx-ir`,
11//! `is_pure` in `rlx-compile`'s constant folder, `is_leaf` in `param_hoist`
12//! *and* again in `rlx-distributed`'s partitioner, `FUSED_KINDS` in the
13//! backend rewriter. Every one of those lists is open — a `matches!` with no
14//! wildcard arm silently answers "no" for a variant nobody remembered to add.
15//!
16//! That is the same failure mode the opcode tables in [`crate::opcodes`] were
17//! written to design out, and the fix is the same: classify every variant
18//! **exactly once**, in an exhaustive `match`. Adding an [`OpKind`] without
19//! classifying it is a compile error, not a silently wrong answer.
20//!
21//! ```
22//! use rlx_ir::{Op, OpKind, capability::OpCaps};
23//!
24//! assert!(OpKind::MatMul.caps().contains(OpCaps::FUSION_BOUNDARY));
25//! assert!(Op::MatMul.is_blas());
26//! assert!(!OpKind::Reshape.caps().contains(OpCaps::ELEMENTWISE));
27//! ```
28//!
29//! # Adding a capability
30//!
31//! Add the flag, then extend the arms below. Because the match is exhaustive
32//! and grouped by capability *set*, a new flag usually means moving a handful
33//! of kinds between existing arms rather than touching all 184 of them.
34
35use crate::op::{Op, OpKind};
36
37/// Number of [`OpKind`] variants.
38pub const N_KINDS: usize = 186;
39
40/// Static properties of an op, independent of its operands or attributes.
41///
42/// A bit set rather than an enum: most questions are independent
43/// (a `FusedMatMulBiasAct` is simultaneously BLAS, a fusion boundary, and a
44/// fused op with a decomposition).
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
46pub struct OpCaps(u32);
47
48impl OpCaps {
49    /// No static capability. The honest answer for most ops: they are neither
50    /// elementwise nor BLAS nor anything else the optimizer special-cases.
51    pub const NONE: Self = Self(0);
52
53    /// A graph root with no tensor inputs (`Input` / `Param` / `Constant`).
54    pub const LEAF: Self = Self(1 << 0);
55
56    /// Same shape in, same shape out, computed independently per element —
57    /// the prime fusion candidate.
58    pub const ELEMENTWISE: Self = Self(1 << 1);
59
60    /// Evaluable on the host from constant operands, so the constant folder
61    /// may replace it outright.
62    ///
63    /// Narrower than "deterministic": it also requires that the folder has an
64    /// implementation. An op can be perfectly deterministic and still not
65    /// carry this flag.
66    pub const CONST_FOLDABLE: Self = Self(1 << 2);
67
68    /// Compute-intensive and dispatched to a tuned library / kernel
69    /// (matmul, conv, solves).
70    pub const BLAS: Self = Self(1 << 3);
71
72    /// Elementwise fusion must not span across this op.
73    ///
74    /// Implied by [`BLAS`](Self::BLAS) — every BLAS op carries both flags, so
75    /// this can be tested directly without also testing `BLAS`.
76    pub const FUSION_BOUNDARY: Self = Self(1 << 4);
77
78    /// Collapses one or more axes; drives loop iteration in a fused kernel.
79    pub const REDUCTION: Self = Self(1 << 5);
80
81    /// May appear in an [`Op::TransformRegion`] chain.
82    pub const TRANSFORM: Self = Self(1 << 6);
83
84    /// A composite op with a decomposition into primitives — the backend
85    /// rewriter can unfuse it when a target lacks the native kernel.
86    pub const FUSED: Self = Self(1 << 7);
87
88    /// Carries one or more nested [`Graph`](crate::Graph) bodies. Consumers
89    /// that walk the IR recursively must descend via [`Op::subgraphs`].
90    pub const NESTED_BODY: Self = Self(1 << 8);
91
92    /// Output is not a function of the inputs alone (RNG / sampling). Never
93    /// const-foldable, never CSE-able, and a hard root for hoisting.
94    pub const NONDETERMINISTIC: Self = Self(1 << 9);
95
96    /// Union of two sets (`const`-friendly, so the table can be built in a
97    /// `const fn`).
98    pub const fn union(self, other: Self) -> Self {
99        Self(self.0 | other.0)
100    }
101
102    /// Does this set contain every flag in `other`?
103    pub const fn contains(self, other: Self) -> bool {
104        self.0 & other.0 == other.0
105    }
106
107    /// Does this set share any flag with `other`?
108    pub const fn intersects(self, other: Self) -> bool {
109        self.0 & other.0 != 0
110    }
111
112    /// True when no capability is set.
113    pub const fn is_empty(self) -> bool {
114        self.0 == 0
115    }
116}
117
118impl std::ops::BitOr for OpCaps {
119    type Output = Self;
120    fn bitor(self, rhs: Self) -> Self {
121        self.union(rhs)
122    }
123}
124
125impl OpKind {
126    /// Every [`OpKind`], in declaration order.
127    ///
128    /// Exists so tests can fan out exhaustively over the whole op set — the
129    /// same role `BinaryOp::ALL` and friends play in [`crate::opcodes`].
130    pub const ALL: [OpKind; N_KINDS] = [
131        OpKind::Input,
132        OpKind::Param,
133        OpKind::Constant,
134        OpKind::Activation,
135        OpKind::Cast,
136        OpKind::StopGradient,
137        OpKind::Quantize,
138        OpKind::Dequantize,
139        OpKind::FakeQuantize,
140        OpKind::FakeQuantizeLSQ,
141        OpKind::FakeQuantizeLSQBackwardX,
142        OpKind::FakeQuantizeLSQBackwardScale,
143        OpKind::Binary,
144        OpKind::Compare,
145        OpKind::Where,
146        OpKind::Fma,
147        OpKind::ElementwiseRegion,
148        OpKind::TransformRegion,
149        OpKind::BatchElementwiseRegion,
150        OpKind::MatMul,
151        OpKind::DotGeneral,
152        OpKind::DenseSolve,
153        OpKind::BatchedDenseSolve,
154        OpKind::Cholesky,
155        OpKind::TriangularSolve,
156        OpKind::Det,
157        OpKind::LogDet,
158        OpKind::Sort,
159        OpKind::ArgSort,
160        OpKind::Svd,
161        OpKind::Qr,
162        OpKind::LayerNorm,
163        OpKind::LayerNorm2d,
164        OpKind::GroupNorm,
165        OpKind::BatchNormInference,
166        OpKind::RmsNorm,
167        OpKind::ResizeNearest2x,
168        OpKind::Interpolate3d,
169        OpKind::Attention,
170        OpKind::Rope,
171        OpKind::AxialRope2d,
172        OpKind::Reshape,
173        OpKind::Transpose,
174        OpKind::Narrow,
175        OpKind::Concat,
176        OpKind::KvAppend,
177        OpKind::Expand,
178        OpKind::Gather,
179        OpKind::Reverse,
180        OpKind::Pad,
181        OpKind::Slice,
182        OpKind::Roll,
183        OpKind::Clamp,
184        OpKind::Tile,
185        OpKind::Trilu,
186        OpKind::Reduce,
187        OpKind::Histogram,
188        OpKind::Softmax,
189        OpKind::Cumsum,
190        OpKind::CumProd,
191        OpKind::CumMax,
192        OpKind::ArgMax,
193        OpKind::ArgMin,
194        OpKind::TopK,
195        OpKind::Sample,
196        OpKind::RngNormal,
197        OpKind::RngUniform,
198        OpKind::Conv,
199        OpKind::Im2Col,
200        OpKind::ConvTranspose2d,
201        OpKind::Conv3d,
202        OpKind::ConvTranspose3d,
203        OpKind::Pool,
204        OpKind::ReluBackward,
205        OpKind::ActivationBackward,
206        OpKind::FakeQuantizeBackward,
207        OpKind::ComplexNormSq,
208        OpKind::ComplexNormSqBackward,
209        OpKind::Conjugate,
210        OpKind::MaxPool2dBackward,
211        OpKind::Conv2dBackwardInput,
212        OpKind::Conv2dBackwardWeight,
213        OpKind::MaxPool3dBackward,
214        OpKind::Conv3dBackwardInput,
215        OpKind::Conv3dBackwardWeight,
216        OpKind::SoftmaxCrossEntropy,
217        OpKind::SoftmaxCrossEntropyWithLogits,
218        OpKind::SoftmaxCrossEntropyBackward,
219        OpKind::AttentionBackward,
220        OpKind::AttentionBackwardAll,
221        OpKind::LayerNormBackwardInput,
222        OpKind::LayerNormBackwardGamma,
223        OpKind::RmsNormBackwardInput,
224        OpKind::RmsNormBackwardGamma,
225        OpKind::RmsNormBackwardBeta,
226        OpKind::RopeBackward,
227        OpKind::GroupNormBackwardInput,
228        OpKind::GroupNormBackwardGamma,
229        OpKind::GroupNormBackwardBeta,
230        OpKind::BatchNormInferenceBackwardInput,
231        OpKind::BatchNormInferenceBackwardGamma,
232        OpKind::BatchNormInferenceBackwardBeta,
233        OpKind::CumsumBackward,
234        OpKind::GatherBackward,
235        OpKind::GroupedMatMul,
236        OpKind::DequantGroupedMatMul,
237        OpKind::DequantGroupedMatMulMlx,
238        OpKind::DequantMoEWeights,
239        OpKind::ScaledGroupedMatMul,
240        OpKind::ScatterAdd,
241        OpKind::ScatterNd,
242        OpKind::ScatterElements,
243        OpKind::GatherNd,
244        OpKind::GatherElements,
245        OpKind::LoraMatMul,
246        OpKind::PartitionedConv,
247        OpKind::DequantMatMul,
248        OpKind::SynthMatMul,
249        OpKind::SynthMatMulBackward,
250        OpKind::SynthReconstruct,
251        OpKind::SplineActivation,
252        OpKind::SplineActivationBackwardX,
253        OpKind::SplineActivationBackwardCoeff,
254        OpKind::QMatMul,
255        OpKind::QConv2d,
256        OpKind::ScaledMatMul,
257        OpKind::ScaledQuantize,
258        OpKind::ScaledQuantScale,
259        OpKind::ScaledDequantize,
260        OpKind::SelectiveScan,
261        OpKind::GatedDeltaNet,
262        OpKind::Lstm,
263        OpKind::Gru,
264        OpKind::Rnn,
265        OpKind::Mamba2,
266        OpKind::FusedSwiGLU,
267        OpKind::FusedMatMulBiasAct,
268        OpKind::FusedMatMulResidual,
269        OpKind::FusedConvBiasAct,
270        OpKind::FusedResidualLN,
271        OpKind::FusedResidualRmsNorm,
272        OpKind::FusedAttentionBlock,
273        OpKind::FusedTransformerLayer,
274        OpKind::If,
275        OpKind::While,
276        OpKind::Scan,
277        OpKind::ScanBackward,
278        OpKind::ScanBackwardXs,
279        OpKind::GaussianSplatRender,
280        OpKind::GaussianSplatRenderBackward,
281        OpKind::GaussianSplatPrepare,
282        OpKind::GaussianSplatRasterize,
283        OpKind::Custom,
284        OpKind::CustomFn,
285        OpKind::Fft,
286        OpKind::FftQ,
287        OpKind::FftButterflyStage,
288        OpKind::LogMel,
289        OpKind::LogMelBackward,
290        OpKind::WelchPeaks,
291        OpKind::BiMap,
292        OpKind::ReEig,
293        OpKind::LogEig,
294        OpKind::SpdBatchNorm,
295        OpKind::SpdKarcherMean,
296        OpKind::ReEigBackward,
297        OpKind::LogEigBackward,
298        OpKind::SpdBatchNormBackwardX,
299        OpKind::SpdBatchNormBackwardG,
300        OpKind::SpdKarcherMeanWeighted,
301        OpKind::SpdLogMap,
302        OpKind::SpdExpMap,
303        OpKind::SpdParallelTransport,
304        OpKind::SpdMatrixFnBatch,
305        OpKind::SpdLogMapBackward,
306        OpKind::SpdExpMapBackward,
307        OpKind::SpdParallelTransportBackward,
308        OpKind::SpdMatrixFnBatchBackward,
309        OpKind::Eigh,
310        OpKind::EighBackward,
311        OpKind::EighBatch,
312        OpKind::EighBatchBackward,
313        OpKind::AdaLayerNorm,
314        OpKind::AdaLayerNormBackward,
315        OpKind::GatedResidual,
316        OpKind::GatedResidualBackward,
317    ];
318
319    /// Static capabilities of this op kind.
320    ///
321    /// Exhaustive by construction: a new [`OpKind`] variant will not compile
322    /// until it is classified here. If it fits no category, add it to the
323    /// final `OpCaps::NONE` arm — that is a deliberate statement that the
324    /// optimizer treats it as an opaque node, not an oversight.
325    pub const fn caps(self) -> OpCaps {
326        use OpKind::*;
327        match self {
328            Input | Param | Constant => OpCaps::LEAF,
329            Activation | Cast | Binary | Compare => {
330                OpCaps::ELEMENTWISE.union(OpCaps::CONST_FOLDABLE)
331            }
332            StopGradient | Where | Fma | ElementwiseRegion | BatchElementwiseRegion => {
333                OpCaps::ELEMENTWISE
334            }
335            MatMul
336            | DotGeneral
337            | DenseSolve
338            | BatchedDenseSolve
339            | Conv
340            | Im2Col
341            | ConvTranspose2d
342            | Conv3d
343            | ConvTranspose3d
344            | GroupedMatMul
345            | DequantGroupedMatMul
346            | DequantGroupedMatMulMlx
347            | DequantMoEWeights
348            | ScaledGroupedMatMul
349            | DequantMatMul
350            | SynthMatMul
351            | QMatMul
352            | QConv2d
353            | ScaledMatMul
354            | FusedMatMulResidual => OpCaps::BLAS.union(OpCaps::FUSION_BOUNDARY),
355            ResizeNearest2x => OpCaps::TRANSFORM,
356            Reshape | Expand => OpCaps::CONST_FOLDABLE,
357            Reduce | Softmax | TopK => OpCaps::REDUCTION,
358            Sample | RngNormal | RngUniform => OpCaps::NONDETERMINISTIC,
359            LoraMatMul | FusedMatMulBiasAct | FusedConvBiasAct => {
360                OpCaps::BLAS.union(OpCaps::FUSION_BOUNDARY.union(OpCaps::FUSED))
361            }
362            PartitionedConv
363            | SelectiveScan
364            | GatedDeltaNet
365            | GatedDeltaNetBackward
366            | Lstm
367            | Gru
368            | Rnn
369            | Mamba2
370            | FusedSwiGLU
371            | FusedResidualLN
372            | FusedResidualRmsNorm
373            | FusedAttentionBlock
374            | FusedTransformerLayer
375            | AdaLayerNorm
376            | GatedResidual => OpCaps::FUSED,
377            If | While | Scan | ScanBackward | ScanBackwardXs | CustomFn => OpCaps::NESTED_BODY,
378            GaussianSplatRender
379            | GaussianSplatRenderBackward
380            | GaussianSplatPrepare
381            | GaussianSplatRasterize => OpCaps::FUSION_BOUNDARY,
382            Quantize
383            | Dequantize
384            | FakeQuantize
385            | FakeQuantizeLSQ
386            | FakeQuantizeLSQBackwardX
387            | FakeQuantizeLSQBackwardScale
388            | TransformRegion
389            | Cholesky
390            | TriangularSolve
391            | Det
392            | LogDet
393            | Sort
394            | ArgSort
395            | Svd
396            | Qr
397            | LayerNorm
398            | LayerNorm2d
399            | GroupNorm
400            | BatchNormInference
401            | RmsNorm
402            | Interpolate3d
403            | Attention
404            | Rope
405            | AxialRope2d
406            | Transpose
407            | Narrow
408            | Concat
409            | KvAppend
410            | Gather
411            | Reverse
412            | Pad
413            | Slice
414            | Roll
415            | Clamp
416            | Tile
417            | Trilu
418            | Histogram
419            | Cumsum
420            | CumProd
421            | CumMax
422            | ArgMax
423            | ArgMin
424            | Pool
425            | ReluBackward
426            | ActivationBackward
427            | FakeQuantizeBackward
428            | ComplexNormSq
429            | ComplexNormSqBackward
430            | Conjugate
431            | MaxPool2dBackward
432            | Conv2dBackwardInput
433            | Conv2dBackwardWeight
434            | MaxPool3dBackward
435            | Conv3dBackwardInput
436            | Conv3dBackwardWeight
437            | SoftmaxCrossEntropy
438            | SoftmaxCrossEntropyWithLogits
439            | SoftmaxCrossEntropyBackward
440            | AttentionBackward
441            | AttentionBackwardAll
442            | LayerNormBackwardInput
443            | LayerNormBackwardGamma
444            | RmsNormBackwardInput
445            | RmsNormBackwardGamma
446            | RmsNormBackwardBeta
447            | RopeBackward
448            | GroupNormBackwardInput
449            | GroupNormBackwardGamma
450            | GroupNormBackwardBeta
451            | BatchNormInferenceBackwardInput
452            | BatchNormInferenceBackwardGamma
453            | BatchNormInferenceBackwardBeta
454            | CumsumBackward
455            | GatherBackward
456            | ScatterAdd
457            | ScatterNd
458            | ScatterElements
459            | GatherNd
460            | GatherElements
461            | SynthMatMulBackward
462            | SynthReconstruct
463            | SplineActivation
464            | SplineActivationBackwardX
465            | SplineActivationBackwardCoeff
466            | ScaledQuantize
467            | ScaledQuantScale
468            | ScaledDequantize
469            | Custom
470            | Fft
471            | FftQ
472            | FftButterflyStage
473            | LogMel
474            | LogMelBackward
475            | WelchPeaks
476            | BiMap
477            | ReEig
478            | LogEig
479            | SpdBatchNorm
480            | SpdKarcherMean
481            | ReEigBackward
482            | LogEigBackward
483            | SpdBatchNormBackwardX
484            | SpdBatchNormBackwardG
485            | SpdKarcherMeanWeighted
486            | SpdLogMap
487            | SpdExpMap
488            | SpdParallelTransport
489            | SpdMatrixFnBatch
490            | SpdLogMapBackward
491            | SpdExpMapBackward
492            | SpdParallelTransportBackward
493            | SpdMatrixFnBatchBackward
494            | Eigh
495            | EighBackward
496            | EighBatch
497            | EighBatchBackward
498            | AdaLayerNormBackward
499            | GatedResidualBackward => OpCaps::NONE,
500        }
501    }
502
503    /// Convenience: does this kind have all of `caps`?
504    pub const fn has(self, caps: OpCaps) -> bool {
505        self.caps().contains(caps)
506    }
507}
508
509// ── Predicates ──────────────────────────────────────────────────
510//
511// These read off the table. They are kept as methods on `Op` because that is
512// where callers already reach for them; the classification itself lives in
513// one place.
514
515impl Op {
516    /// True if this op is element-wise (same shape in, same shape out).
517    /// Element-wise ops are prime fusion candidates.
518    pub fn is_elementwise(&self) -> bool {
519        self.kind().has(OpCaps::ELEMENTWISE)
520    }
521
522    /// True if this op may appear in a [`Op::TransformRegion`] chain.
523    pub fn is_transform_eligible(&self) -> bool {
524        self.kind().has(OpCaps::TRANSFORM)
525    }
526
527    /// True if this op is a BLAS/compute-intensive op that forms a fusion boundary.
528    pub fn is_blas(&self) -> bool {
529        self.kind().has(OpCaps::BLAS)
530    }
531
532    /// True if element-wise fusion must not span across this op.
533    pub fn is_fusion_boundary(&self) -> bool {
534        self.kind().has(OpCaps::FUSION_BOUNDARY)
535    }
536
537    /// True if this op is a reduction (drives loop iteration in fused kernels).
538    pub fn is_reduction(&self) -> bool {
539        self.kind().has(OpCaps::REDUCTION)
540    }
541
542    /// True if this op is a graph root with no tensor inputs.
543    pub fn is_leaf(&self) -> bool {
544        self.kind().has(OpCaps::LEAF)
545    }
546
547    /// True if the constant folder can evaluate this op on the host.
548    pub fn is_const_foldable(&self) -> bool {
549        self.kind().has(OpCaps::CONST_FOLDABLE)
550    }
551
552    /// True if this is a composite op with a decomposition into primitives.
553    pub fn is_fused(&self) -> bool {
554        self.kind().has(OpCaps::FUSED)
555    }
556
557    /// True if this op carries nested [`Graph`](crate::Graph) bodies.
558    /// Equivalent to `!self.subgraphs().is_empty()`, without the walk.
559    pub fn has_nested_body(&self) -> bool {
560        self.kind().has(OpCaps::NESTED_BODY)
561    }
562
563    /// True if the output is not a function of the inputs alone (RNG /
564    /// sampling), so it must never be folded, CSE'd or hoisted.
565    pub fn is_nondeterministic(&self) -> bool {
566        self.kind().has(OpCaps::NONDETERMINISTIC)
567    }
568}
569
570/// Every kind carrying [`OpCaps::FUSED`] — the set the backend rewriter tries
571/// to unfuse when a target lacks the native kernel.
572pub fn fused_kinds() -> Vec<OpKind> {
573    OpKind::ALL
574        .into_iter()
575        .filter(|k| k.has(OpCaps::FUSED))
576        .collect()
577}
578
579#[cfg(test)]
580mod tests {
581    use super::*;
582
583    #[test]
584    fn all_is_complete_and_duplicate_free() {
585        let mut seen = std::collections::HashSet::new();
586        for kind in OpKind::ALL {
587            assert!(seen.insert(kind), "{kind:?} appears twice in OpKind::ALL");
588        }
589        assert_eq!(seen.len(), N_KINDS);
590    }
591
592    #[test]
593    fn blas_implies_fusion_boundary() {
594        // The predicate tests `FUSION_BOUNDARY` directly, so the table must
595        // set both flags rather than relying on caller-side implication.
596        for kind in OpKind::ALL {
597            if kind.has(OpCaps::BLAS) {
598                assert!(
599                    kind.has(OpCaps::FUSION_BOUNDARY),
600                    "{kind:?} is BLAS but not a fusion boundary"
601                );
602            }
603        }
604    }
605
606    #[test]
607    fn leaves_are_not_elementwise() {
608        for kind in OpKind::ALL {
609            if kind.has(OpCaps::LEAF) {
610                assert!(!kind.has(OpCaps::ELEMENTWISE), "{kind:?}");
611                assert!(!kind.has(OpCaps::FUSION_BOUNDARY), "{kind:?}");
612            }
613        }
614    }
615
616    #[test]
617    fn nondeterministic_ops_are_never_const_foldable() {
618        // Folding an RNG op would bake one draw into the graph.
619        for kind in OpKind::ALL {
620            if kind.has(OpCaps::NONDETERMINISTIC) {
621                assert!(!kind.has(OpCaps::CONST_FOLDABLE), "{kind:?}");
622            }
623        }
624    }
625
626    #[test]
627    fn nested_body_flag_agrees_with_the_subgraph_accessor() {
628        // The flag is the cheap answer; `Op::subgraphs` is the real one. They
629        // must not drift — `has_nested_body` is used to decide whether to
630        // recurse at all.
631        let g = || Box::new(crate::Graph::new("b"));
632        let cases: Vec<Op> = vec![
633            Op::MatMul,
634            Op::StopGradient,
635            Op::If {
636                then_branch: g(),
637                else_branch: g(),
638            },
639            Op::While {
640                cond: g(),
641                body: g(),
642                max_iterations: None,
643            },
644            Op::Scan {
645                body: g(),
646                length: 1,
647                save_trajectory: false,
648                num_bcast: 0,
649                num_xs: 0,
650                num_checkpoints: 0,
651            },
652            Op::CustomFn {
653                fwd_body: g(),
654                vjp_body: None,
655                jvp_body: None,
656                num_inputs: 1,
657            },
658        ];
659        for op in cases {
660            assert_eq!(
661                op.has_nested_body(),
662                !op.subgraphs().is_empty(),
663                "{:?} disagrees with its subgraph accessor",
664                op.kind()
665            );
666        }
667    }
668
669    #[test]
670    fn spot_check_known_classifications() {
671        assert!(OpKind::Input.has(OpCaps::LEAF));
672        assert!(OpKind::Activation.has(OpCaps::ELEMENTWISE.union(OpCaps::CONST_FOLDABLE)));
673        assert!(OpKind::MatMul.has(OpCaps::BLAS));
674        assert!(OpKind::Reduce.has(OpCaps::REDUCTION));
675        assert!(OpKind::Scan.has(OpCaps::NESTED_BODY));
676        assert!(OpKind::RngNormal.has(OpCaps::NONDETERMINISTIC));
677        assert!(OpKind::FusedSwiGLU.has(OpCaps::FUSED));
678        // Reshape is host-foldable but not elementwise — a shape-only op.
679        assert!(OpKind::Reshape.has(OpCaps::CONST_FOLDABLE));
680        assert!(!OpKind::Reshape.has(OpCaps::ELEMENTWISE));
681        // Transpose is opaque to the optimizer's categories.
682        assert!(OpKind::Transpose.caps().is_empty());
683    }
684
685    #[test]
686    fn fused_kinds_matches_the_flag() {
687        let listed = fused_kinds();
688        assert!(!listed.is_empty());
689        for kind in &listed {
690            assert!(kind.has(OpCaps::FUSED));
691        }
692        assert_eq!(
693            listed.len(),
694            OpKind::ALL.iter().filter(|k| k.has(OpCaps::FUSED)).count()
695        );
696    }
697}