Skip to content

Remove mallocs and pass single buffer in run_op_integration #542

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
felixhekhorn merged 9 commits into from
Jun 26, 2026
Merged

Remove mallocs and pass single buffer in run_op_integration #542

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
da770fc
Removed mallocs in ekore
AkshatRai07 Jun 24, 2026
3255762
Single buffer for integration and removed python malloc loops
AkshatRai07 Jun 24, 2026
5083adc
Changed cache size const name and added related tests
AkshatRai07 Jun 25, 2026
ec36567
Add MAX_ORDER_QED/QCD and rm magic numbers
AkshatRai07 Jun 25, 2026
d48e819
Changed dimensions in docstring and res array
AkshatRai07 Jun 25, 2026
dc43f3a
Add return dimensions to docs in ekore
AkshatRai07 Jun 26, 2026
2f8656e
Rename `n_max`
felixhekhorn Jun 26, 2026
02d23f1
Removed all branches from LHA Bench
AkshatRai07 Jun 26, 2026
9bfbbc9
Update performance.md with #542
AkshatRai07 Jun 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 75 additions & 61 deletions crates/eko/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,71 +7,69 @@ use std::ffi::c_void;
pub mod bib;
pub mod mellin;

/// Wrapper to pass arguments back to Python.
struct RawCmplx {
re: Vec<f64>,
im: Vec<f64>,
}

/// Map tensor with shape (o,d,d) to c-ordered list.
/// Map tensor with shape (o,d,d) into pre-allocated output buffers.
///
/// This is needed for the QCD singlet.
fn unravel<const DIM: usize>(res: Vec<[[Complex<f64>; DIM]; DIM]>, order_qcd: usize) -> RawCmplx {
let mut target = RawCmplx {
re: Vec::<f64>::new(),
im: Vec::<f64>::new(),
};
fn unravel<const DIM: usize>(
res: &[[[Complex<f64>; DIM]; DIM]],
Comment thread
AkshatRai07 marked this conversation as resolved.
Outdated
order_qcd: usize,
out_re: &mut [f64],
out_im: &mut [f64],
) {
let mut i = 0;
for obj in res.iter().take(order_qcd) {
for col in obj.iter().take(DIM) {
for el in col.iter().take(DIM) {
target.re.push(el.re);
target.im.push(el.im);
for col in obj.iter() {
for el in col.iter() {
out_re[i] = el.re;
out_im[i] = el.im;
i += 1;
}
}
}
target
}

/// Map tensor with shape (o,o',d,d) to c-ordered list.
/// Map tensor with shape (o,o',d,d) into pre-allocated output buffers.
///
/// This is needed for the QED singlet and valence.
fn unravel_qed<const DIM: usize>(
res: Vec<Vec<[[Complex<f64>; DIM]; DIM]>>,
res: &[[[[Complex<f64>; DIM]; DIM]; 3]],
order_qcd: usize,
order_qed: usize,
) -> RawCmplx {
let mut target = RawCmplx {
re: Vec::<f64>::new(),
im: Vec::<f64>::new(),
};
out_re: &mut [f64],
out_im: &mut [f64],
) {
let mut i = 0;
for obj_ in res.iter().take(order_qcd + 1) {
for obj in obj_.iter().take(order_qed + 1) {
for col in obj.iter().take(DIM) {
for el in col.iter().take(DIM) {
target.re.push(el.re);
target.im.push(el.im);
for col in obj.iter() {
for el in col.iter() {
out_re[i] = el.re;
out_im[i] = el.im;
i += 1;
}
}
}
}
target
}

/// Map tensor with shape (o,o',d) to c-ordered list.
/// Map tensor with shape (o,o',d) into pre-allocated output buffers.
///
/// This is needed for the QED non-singlet.
fn unravel_qed_ns(res: Vec<Vec<Complex<f64>>>, order_qcd: usize, order_qed: usize) -> RawCmplx {
let mut target = RawCmplx {
re: Vec::<f64>::new(),
im: Vec::<f64>::new(),
};
fn unravel_qed_ns(
res: &[[Complex<f64>; 3]],
order_qcd: usize,
order_qed: usize,
out_re: &mut [f64],
out_im: &mut [f64],
) {
let mut i = 0;
for col in res.iter().take(order_qcd + 1) {
for el in col.iter().take(order_qed + 1) {
target.re.push(el.re);
target.im.push(el.im);
out_re[i] = el.re;
out_im[i] = el.im;
i += 1;
}
}
target
}

/// Integration kernel inside quad.
Expand All @@ -93,42 +91,46 @@ pub unsafe extern "C" fn rust_quad_ker(u: f64, rargs: *mut c_void) -> f64 {
let path = mellin::TalbotPath::new(u, args.logx, is_singlet);
let jac = path.jac() * path.prefactor();
let mut c = Cache::new(path.n());
let mut raw = RawCmplx {
re: Vec::<f64>::new(),
im: Vec::<f64>::new(),
};
let n3lo_ad_variation = std::slice::from_raw_parts(args.n3lo_ad_variation, 7)
let n3lo_ad_variation: [u8; 7] = std::slice::from_raw_parts(args.n3lo_ad_variation, 7)
.try_into()
.unwrap();

let n_max = (args.order_qcd + 1) * (args.order_qed + 1) * 16;
let out_re = std::slice::from_raw_parts_mut(args.re_gamma, n_max);
let out_im = std::slice::from_raw_parts_mut(args.im_gamma, n_max);
Comment thread
felixhekhorn marked this conversation as resolved.
Outdated

if args.is_ome {
if is_singlet {
raw = unravel(
ekore::operator_matrix_elements::unpolarized::spacelike::A_singlet(
unravel(
&ekore::operator_matrix_elements::unpolarized::spacelike::A_singlet(
args.order_qcd,
&mut c,
args.nf,
args.L,
),
args.order_qcd,
out_re,
out_im,
);
} else {
raw = unravel(
ekore::operator_matrix_elements::unpolarized::spacelike::A_non_singlet(
unravel(
&ekore::operator_matrix_elements::unpolarized::spacelike::A_non_singlet(
args.order_qcd,
&mut c,
args.nf,
args.L,
),
args.order_qcd,
out_re,
out_im,
);
}
} else if is_singlet {
if args.order_qed > 0 {
let gamma_singlet_qed =
ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qed;
raw = unravel_qed(
gamma_singlet_qed(
unravel_qed(
&gamma_singlet_qed(
args.order_qcd,
args.order_qed,
&mut c,
Expand All @@ -137,28 +139,32 @@ pub unsafe extern "C" fn rust_quad_ker(u: f64, rargs: *mut c_void) -> f64 {
),
args.order_qcd,
args.order_qed,
out_re,
out_im,
);
} else {
let gamma_singlet_qcd = match args.is_polarized {
true => ekore::anomalous_dimensions::polarized::spacelike::gamma_singlet_qcd,
false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qcd,
};
raw = unravel(
gamma_singlet_qcd(
unravel(
&gamma_singlet_qcd(
args.order_qcd,
&mut c,
args.nf,
n3lo_ad_variation[0..4].try_into().unwrap(),
),
args.order_qcd,
out_re,
out_im,
);
}
} else if args.order_qed > 0 {
if is_qed_valence {
let gamma_valence_qed =
ekore::anomalous_dimensions::unpolarized::spacelike::gamma_valence_qed;
raw = unravel_qed(
gamma_valence_qed(
unravel_qed(
&gamma_valence_qed(
args.order_qcd,
args.order_qed,
&mut c,
Expand All @@ -167,11 +173,13 @@ pub unsafe extern "C" fn rust_quad_ker(u: f64, rargs: *mut c_void) -> f64 {
),
args.order_qcd,
args.order_qed,
out_re,
out_im,
);
} else {
let gamma_ns_qed = ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qed;
raw = unravel_qed_ns(
gamma_ns_qed(
unravel_qed_ns(
&gamma_ns_qed(
args.order_qcd,
args.order_qed,
args.mode0,
Expand All @@ -181,10 +189,11 @@ pub unsafe extern "C" fn rust_quad_ker(u: f64, rargs: *mut c_void) -> f64 {
),
args.order_qcd,
args.order_qed,
out_re,
out_im,
);
}
} else {
// we can not do 1D
let gamma_ns_qcd = match args.is_polarized {
true => ekore::anomalous_dimensions::polarized::spacelike::gamma_ns_qcd,
false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qcd,
Expand All @@ -196,16 +205,16 @@ pub unsafe extern "C" fn rust_quad_ker(u: f64, rargs: *mut c_void) -> f64 {
args.nf,
n3lo_ad_variation[4..7].try_into().unwrap(),
);
for el in res.iter().take(args.order_qcd) {
raw.re.push(el.re);
raw.im.push(el.im);
for (i, el) in res.iter().take(args.order_qcd).enumerate() {
out_re[i] = el.re;
out_im[i] = el.im;
}
}

// pass on
(args.py)(
raw.re.as_ptr(),
raw.im.as_ptr(),
args.re_gamma as *const f64,
args.im_gamma as *const f64,
c.n().re,
c.n().im,
jac.re,
Expand Down Expand Up @@ -317,6 +326,9 @@ pub struct QuadArgs {
pub alphaem_running: bool,
// additional param required for N3LO
pub n3lo_ad_variation: *const u8,
// pre-allocated gamma output buffers (owned by Python, written by Rust)
pub re_gamma: *mut f64,
pub im_gamma: *mut f64,
}

/// Empty placeholder function for python callback.
Expand Down Expand Up @@ -403,5 +415,7 @@ pub unsafe extern "C" fn empty_args() -> QuadArgs {
a_half_y: 0,
alphaem_running: false,
n3lo_ad_variation: [0, 0, 0, 0, 0, 0, 0].as_ptr(),
re_gamma: std::ptr::null_mut(),
im_gamma: std::ptr::null_mut(),
}
}
14 changes: 4 additions & 10 deletions crates/ekore/src/anomalous_dimensions/polarized/spacelike.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ pub fn gamma_ns_qcd(
c: &mut Cache,
nf: u8,
_n3lo_variation: [u8; 3],
) -> Vec<Complex<f64>> {
let mut gamma_ns = vec![Complex::<f64>::zero(); order_qcd];
) -> [Complex<f64>; 4] {
let mut gamma_ns = [Complex::<f64>::zero(); 4];
gamma_ns[0] = as1::gamma_ns(c, nf);
// NLO and beyond
if order_qcd >= 2 {
Expand Down Expand Up @@ -47,14 +47,8 @@ pub fn gamma_singlet_qcd(
c: &mut Cache,
nf: u8,
_n3lo_variation: [u8; 4],
) -> Vec<[[Complex<f64>; 2]; 2]> {
let mut gamma_S = vec![
[
[Complex::<f64>::zero(), Complex::<f64>::zero()],
[Complex::<f64>::zero(), Complex::<f64>::zero()]
];
order_qcd
];
) -> [[[Complex<f64>; 2]; 2]; 4] {
let mut gamma_S = [[[Complex::<f64>::zero(); 2]; 2]; 4];
gamma_S[0] = as1::gamma_singlet(c, nf);
// NLO and beyond
if order_qcd >= 2 {
Expand Down
37 changes: 10 additions & 27 deletions crates/ekore/src/anomalous_dimensions/unpolarized/spacelike.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ pub fn gamma_ns_qcd(
c: &mut Cache,
nf: u8,
n3lo_variation: [u8; 3],
) -> Vec<Complex<f64>> {
let mut gamma_ns = vec![Complex::<f64>::zero(); order_qcd];
) -> [Complex<f64>; 4] {
let mut gamma_ns = [Complex::<f64>::zero(); 4];
gamma_ns[0] = as1::gamma_ns(c, nf);
// NLO and beyond
if order_qcd >= 2 {
Expand Down Expand Up @@ -71,14 +71,8 @@ pub fn gamma_singlet_qcd(
c: &mut Cache,
nf: u8,
n3lo_variation: [u8; 4],
) -> Vec<[[Complex<f64>; 2]; 2]> {
let mut gamma_S = vec![
[
[Complex::<f64>::zero(), Complex::<f64>::zero()],
[Complex::<f64>::zero(), Complex::<f64>::zero()]
];
order_qcd
];
) -> [[[Complex<f64>; 2]; 2]; 4] {
let mut gamma_S = [[[Complex::<f64>::zero(); 2]; 2]; 4];
gamma_S[0] = as1::gamma_singlet(c, nf);
// NLO and beyond
if order_qcd >= 2 {
Expand Down Expand Up @@ -107,9 +101,8 @@ pub fn gamma_ns_qed(
c: &mut Cache,
nf: u8,
n3lo_variation: [u8; 3],
) -> Vec<Vec<Complex<f64>>> {
let col = vec![Complex::<f64>::zero(); order_qed + 1];
let mut gamma_ns = vec![col; order_qcd + 1];
) -> [[Complex<f64>; 3]; 5] {
let mut gamma_ns = [[Complex::<f64>::zero(); 3]; 5];

// QCD corrections
let qcd_mode = match mode {
Expand Down Expand Up @@ -158,17 +151,8 @@ pub fn gamma_singlet_qed(
c: &mut Cache,
nf: u8,
n3lo_variation: [u8; 7],
) -> Vec<Vec<[[Complex<f64>; 4]; 4]>> {
let col = vec![
[[
Complex::<f64>::zero(),
Complex::<f64>::zero(),
Complex::<f64>::zero(),
Complex::<f64>::zero()
]; 4];
order_qed + 1
];
let mut gamma_s = vec![col; order_qcd + 1];
) -> [[[[Complex<f64>; 4]; 4]; 3]; 5] {
let mut gamma_s = [[[[Complex::<f64>::zero(); 4]; 4]; 3]; 5];

// QCD corrections
let gamma_qcd_s = gamma_singlet_qcd(order_qcd, c, nf, n3lo_variation[0..4].try_into().unwrap());
Expand Down Expand Up @@ -229,9 +213,8 @@ pub fn gamma_valence_qed(
c: &mut Cache,
nf: u8,
n3lo_variation: [u8; 3],
) -> Vec<Vec<[[Complex<f64>; 2]; 2]>> {
let col = vec![[[Complex::<f64>::zero(), Complex::<f64>::zero(),]; 2]; order_qed + 1];
let mut gamma_v = vec![col; order_qcd + 1];
) -> [[[[Complex<f64>; 2]; 2]; 3]; 5] {
let mut gamma_v = [[[[Complex::<f64>::zero(); 2]; 2]; 3]; 5];

// QCD corrections
let gamma_qcd_nsv = gamma_ns_qcd(order_qcd, PID_NSV, c, nf, n3lo_variation);
Expand Down
Loading
Loading