rlx_ir/ops/
conv2d.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 convolution builders (`conv2d`, `conv_transpose2d`).
6
7use crate::{Graph, NodeId, Op};
8
9impl Graph {
10    /// 2D convolution on NCHW tensors (`Op::Conv`). Weight `[C_out, C_in/g, kH, kW]`.
11    pub fn conv2d(
12        &mut self,
13        input: NodeId,
14        weight: NodeId,
15        kernel_size: [usize; 2],
16        stride: [usize; 2],
17        padding: [usize; 2],
18        dilation: [usize; 2],
19        groups: usize,
20    ) -> NodeId {
21        let in_s = self.node(input).shape.clone();
22        let w_s = self.node(weight).shape.clone();
23        let out = crate::shape::conv2d_output_shape(
24            &in_s,
25            &w_s,
26            kernel_size,
27            stride,
28            padding,
29            dilation,
30            groups,
31        )
32        .expect("conv2d shape inference");
33        self.push(
34            Op::Conv {
35                kernel_size: kernel_size.to_vec(),
36                stride: stride.to_vec(),
37                padding: padding.to_vec(),
38                dilation: dilation.to_vec(),
39                groups,
40            },
41            vec![input, weight],
42            out,
43            None,
44        )
45    }
46
47    /// NCHW im2col (`Op::Im2Col`). Output `[M, C·kH·kW]`.
48    pub fn im2col(
49        &mut self,
50        input: NodeId,
51        kernel_size: [usize; 2],
52        stride: [usize; 2],
53        padding: [usize; 2],
54        dilation: [usize; 2],
55    ) -> NodeId {
56        let in_s = self.node(input).shape.clone();
57        let out = crate::shape::im2col_output_shape(&in_s, kernel_size, stride, padding, dilation)
58            .expect("im2col shape inference");
59        self.push(
60            Op::Im2Col {
61                kernel_size: kernel_size.to_vec(),
62                stride: stride.to_vec(),
63                padding: padding.to_vec(),
64                dilation: dilation.to_vec(),
65            },
66            vec![input],
67            out,
68            None,
69        )
70    }
71
72    /// 2D transposed convolution on NCHW. Weight `[C_in, C_out/g, kH, kW]`.
73    pub fn conv_transpose2d(
74        &mut self,
75        input: NodeId,
76        weight: NodeId,
77        kernel_size: [usize; 2],
78        stride: [usize; 2],
79        padding: [usize; 2],
80        dilation: [usize; 2],
81        output_padding: [usize; 2],
82        groups: usize,
83    ) -> NodeId {
84        let in_s = self.node(input).shape.clone();
85        let w_s = self.node(weight).shape.clone();
86        let out = crate::shape::conv_transpose2d_output_shape(
87            &in_s,
88            &w_s,
89            kernel_size,
90            stride,
91            padding,
92            dilation,
93            output_padding,
94            groups,
95        )
96        .expect("conv_transpose2d shape inference");
97        self.push(
98            Op::ConvTranspose2d {
99                kernel_size: kernel_size.to_vec(),
100                stride: stride.to_vec(),
101                padding: padding.to_vec(),
102                dilation: dilation.to_vec(),
103                output_padding: output_padding.to_vec(),
104                groups,
105            },
106            vec![input, weight],
107            out,
108            None,
109        )
110    }
111}
rlx_ir/ops/conv2d.rs

rlx_ir/ops/
conv2d.rs
