rlx_ir/ops/
pool2d.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//! NCHW pooling builder (`Op::Pool`).
6
7use crate::op::ReduceOp;
8use crate::{Graph, NodeId, Op};
9
10impl Graph {
11    /// 2-D pooling over NCHW (`Op::Pool`). Channels pass through; only the two
12    /// spatial axes shrink.
13    ///
14    /// `kind` selects the reduction — [`ReduceOp::Max`] and [`ReduceOp::Mean`]
15    /// are the ones backends lower.
16    pub fn pool2d(
17        &mut self,
18        input: NodeId,
19        kind: ReduceOp,
20        kernel_size: [usize; 2],
21        stride: [usize; 2],
22        padding: [usize; 2],
23    ) -> NodeId {
24        let in_s = self.node(input).shape.clone();
25        let out = crate::shape::pool2d_output_shape(&in_s, kernel_size, stride, padding)
26            .expect("pool2d shape inference");
27        self.push(
28            Op::Pool {
29                kind,
30                kernel_size: kernel_size.to_vec(),
31                stride: stride.to_vec(),
32                padding: padding.to_vec(),
33            },
34            vec![input],
35            out,
36            None,
37        )
38    }
39
40    /// Max pooling over NCHW — [`Graph::pool2d`] with [`ReduceOp::Max`].
41    pub fn max_pool2d(
42        &mut self,
43        input: NodeId,
44        kernel_size: [usize; 2],
45        stride: [usize; 2],
46        padding: [usize; 2],
47    ) -> NodeId {
48        self.pool2d(input, ReduceOp::Max, kernel_size, stride, padding)
49    }
50
51    /// Average pooling over NCHW — [`Graph::pool2d`] with [`ReduceOp::Mean`].
52    pub fn avg_pool2d(
53        &mut self,
54        input: NodeId,
55        kernel_size: [usize; 2],
56        stride: [usize; 2],
57        padding: [usize; 2],
58    ) -> NodeId {
59        self.pool2d(input, ReduceOp::Mean, kernel_size, stride, padding)
60    }
61}
rlx_ir/ops/pool2d.rs

rlx_ir/ops/
pool2d.rs
