Rust-related/luminal
1
0
Fork 1
mirror of https://github.com/jafioti/luminal.git synced 2026-06-01 21:49:47 +09:00
Code Issues Packages Projects Releases Wiki Activity

loop unrolling

Browse Source
This commit is contained in:
Joe Fioti
2025-09-17 07:37:37 -07:00
parent 8a6afcc28c
commit 355378ddda
2 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/luminal_2/src/code.lisp
View File

@@ -399,7 +399,6 @@
; Swapping
; Fission
; TensorCore
(ruleset tc)

9
crates/luminal_2/src/codegen.rs
View File

@@ -469,6 +469,9 @@ fn make_kernel(
.next()
.unwrap();
// Use a single loop with correct striding from the input
if loads.to_usize().map(|i| i <= 16).unwrap_or_default() {
kernel_lines.push(format!("{spacing}#pragma unroll")); // Loop unroll short loops
}
kernel_lines.push(format!(
"{spacing}for (int load = 0; load < {}; ++load) {{",
loads.to_kernel()
@@ -546,6 +549,9 @@ fn make_kernel(
// for
*prev_max_var += 1;
let loop_var = var_to_char(*prev_max_var);
if range.to_usize().map(|i| i <= 16).unwrap_or_default() {
kernel_lines.push(format!("{spacing}#pragma unroll")); // Loop unroll short loops
}
kernel_lines.push(format!("{spacing}for (int loop_{loop_var} = 0; loop_{loop_var} < {}; ++loop_{loop_var}) {{", range.to_kernel()));
};
let loop_var = var_to_char(*prev_max_var);
@@ -737,6 +743,9 @@ fn make_kernel(
);
current_elem_size *= range;
}
if size.to_usize().map(|i| i <= 16).unwrap_or_default() {
kernel_lines.push(format!("{spacing}#pragma unroll")); // Loop unroll short loops
}
kernel_lines.push(format!(
"{spacing}for (int save = 0; save < {}; ++save) {{",
size.to_kernel()