1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only
//! Loop merging optimization.
//!
//! This module provides loop merging optimizations for kernels,
//! which merge nested loops into single loops when possible.
//!
//! Loop merging can improve performance by:
//!
//! - Reducing loop overhead
//! - Enabling better instruction scheduling
//! - Improving vectorization opportunities
use std::collections::BTreeMap;
use super::autotune::Optimization;
use crate::{
dtype::Constant,
kernel::{BOp, Kernel, Op, OpId},
};
impl Kernel {
/// Get last op in the given loop scope
pub(crate) fn get_last_dim_op(&self, loop_id: OpId) -> OpId {
match self.ops[loop_id].op {
Op::Index { .. } => return self.tail,
Op::Loop { .. } => {}
_ => unreachable!(),
}
let mut loop_depth = 0;
let mut op_id = loop_id;
while !op_id.is_null() {
match self.ops[op_id].op {
Op::Loop { .. } => {
loop_depth += 1;
}
Op::EndLoop => {
loop_depth -= 1;
if loop_depth == 0 {
return op_id;
}
}
_ => {}
}
op_id = self.next_op(op_id);
}
op_id
}
/// Merge nested Op::Loops into a single loop.
///
/// Takes a chain of nested loops (outermost first) and merges them into one
/// loop whose length is the product of all lengths. After merging, each
/// original loop is replaced with arithmetic that decomposes the merged loop
/// variable back into the original loop variables (via `/` and `%` chain,
/// like `merge_indices`). The existing address computation continues to
/// work correctly because it still references the same OpIds.
pub(crate) fn merge_nested_loops(&mut self, loop_ids: &[OpId]) {
if loop_ids.len() < 2 {
return;
}
let mut total_len: u64 = 1;
for &id in loop_ids {
if let Op::Loop { len: len_id } = self.ops[id].op {
total_len *= self.loop_len_dim(len_id);
}
}
// Replace original loops with merged loop, removing inner EndLoops
let anchor = loop_ids[0];
let merged_len = self.insert_const_idx_before(anchor, total_len);
let mut x = self.insert_before(anchor, Op::Loop { len: merged_len });
// Single pass: remove inner EndLoops (keep only the last one)
let mut op_id = self.next_op(anchor);
let mut depth: u32 = 1;
while !op_id.is_null() {
let next = self.next_op(op_id);
match self.ops[op_id].op {
Op::Loop { .. } => depth += 1,
Op::EndLoop => {
depth -= 1;
if depth > 0 {
self.remove_op(op_id);
} else {
break;
}
}
_ => {}
}
op_id = next;
}
// Decompose the merged loop variable back into original loop variables.
// Process innermost to outermost (reverse order of loop_ids).
// Insert all new ops before the anchor so all definitions
// precede all uses (avoids backward-reference verification errors).
for i in (0..loop_ids.len()).rev() {
let Op::Loop { len: len_id } = self.ops[loop_ids[i]].op else {
unreachable!()
};
let len = self.loop_len_dim(len_id);
let y = self.insert_before(anchor, Op::Const(Constant::idx(len)));
self.ops[loop_ids[i]].op = Op::Binary { x, y, bop: BOp::Mod };
x = self.insert_before(anchor, Op::Binary { x, y, bop: BOp::Div });
}
}
/// Merges two or more indices together
pub(crate) fn merge_indices(&mut self, loops: &[OpId]) {
let mut acc = 1;
let mut axes = BTreeMap::default();
let mut first_id = None;
let mut op_id = self.head;
while axes.len() != loops.len() {
if loops.contains(&op_id) {
// TODO check all scopes are the same
let Op::Index { len, axis, .. } = self.ops[op_id].op else {
unreachable!()
};
acc *= len;
axes.insert(axis, (op_id, len));
if first_id.is_none() {
first_id = Some(op_id);
}
}
op_id = self.next_op(op_id);
}
let Op::Index { axis, scope, .. } = self.ops[first_id.unwrap()].op else {
unreachable!()
};
let mut x = self.insert_before(first_id.unwrap(), Op::Index { len: acc, axis, scope });
for (.., (loop_id, len)) in axes {
let y = self.insert_before(loop_id, Op::Const(Constant::idx(len as u64)));
self.ops[loop_id].op = Op::Binary { x, y, bop: BOp::Mod };
x = self.insert_after(loop_id, Op::Binary { x, y, bop: BOp::Div });
}
self.verify();
}
/// Returns the Optimization for merging nested loops and the number of nested loop groups.
/// Each group is a chain of nested loops that can be merged into one loop.
pub(crate) fn opt_merge_nested_loops(&self) -> (Optimization, usize) {
let groups = self.find_nested_loop_groups();
let n = groups.len();
(Optimization::MergeNestedLoops { groups }, n)
}
/// Find all groups of nested loops in the kernel.
/// Each group is a chain of consecutive nested loops (outermost first).
fn find_nested_loop_groups(&self) -> Vec<Vec<OpId>> {
let mut groups: Vec<Vec<OpId>> = Vec::new();
let mut current_group: Vec<OpId> = Vec::new();
let mut depth: u32 = 0;
let mut in_group = false;
let mut op_id = self.head;
while !op_id.is_null() {
match self.ops[op_id].op {
Op::Loop { .. } => {
if depth == 0 {
// Start a new group
if in_group {
groups.push(std::mem::take(&mut current_group));
}
in_group = true;
}
current_group.push(op_id);
depth += 1;
}
Op::EndLoop => {
depth -= 1;
if depth == 0 {
// End of this group
if !current_group.is_empty() {
groups.push(std::mem::take(&mut current_group));
}
in_group = false;
}
}
_ => {}
}
op_id = self.next_op(op_id);
}
// Flush any remaining group
if !current_group.is_empty() {
groups.push(current_group);
}
groups.retain(|g| g.len() >= 2);
groups
}
}