Skip to content

Custom cuda and metal kernels for system validation#274

Open
Luthaf wants to merge 4 commits into
metatensor:metatomic-corefrom
Luthaf:kernels
Open

Custom cuda and metal kernels for system validation#274
Luthaf wants to merge 4 commits into
metatensor:metatomic-corefrom
Luthaf:kernels

Conversation

@Luthaf

@Luthaf Luthaf commented Jun 30, 2026

Copy link
Copy Markdown
Member

Add custom cuda and metal kernels for system validation. The intention is to add more kernels for unit conversion & any other work that needs to be done inside mta_execute_model.

No tests other than making sure the code compile for now, they will be added once we have Python bindings.

Contributor (creator of pull-request) checklist

  • Tests updated (for new features and bugfixes)?
  • Documentation updated (for new features)?
  • Issue referenced (for PRs that solve an issue)?

Reviewer checklist

  • CHANGELOG updated with public API or any other important changes?

@Luthaf Luthaf requested a review from RMeli June 30, 2026 09:19
@Luthaf Luthaf force-pushed the kernels branch 3 times, most recently from c2584f3 to 607c96a Compare July 8, 2026 13:10
Comment thread metatomic-core/src/kernels/cpu.rs
Comment thread metatomic-core/cmake/metatomic-config.in.cmake
Comment thread metatomic-core/Cargo.toml
Comment thread metatomic-core/src/kernels/cpu.rs Outdated
Comment thread metatomic-core/src/kernels/cuda.rs Outdated
Comment thread metatomic-core/src/kernels/cuda_kernels.cu Outdated
Comment thread metatomic-core/src/kernels/mod.rs
Comment thread metatomic-core/src/kernels/mod.rs
@Luthaf Luthaf force-pushed the kernels branch 3 times, most recently from 7d61507 to d25d254 Compare July 10, 2026 13:23

@RMeli RMeli left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A few final comments, but probably missed some other things without tests... We can revisit once the Python bindings are available.

})
.map_err(|e| Error::Internal(format!("kernel launch (f32): {e}")))?;
}
} else {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I missed this on the first pass, worth checking cell.dtype().bits == 64 explicitly?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

Comment thread metatomic-core/src/kernels/metal.rs Outdated
Comment on lines +68 to +69
// Safety: the cache lives for the lifetime of the process (Lazy static),
// so the reference is valid for 'static once initialized and never removed.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not very familiar with Rust, but from the documentation looks like HashMap can re-allocate. Is this taken into account?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yep, we don't actually need a 'static value here, I simplified it!

Comment thread metatomic-core/src/kernels/metal.rs Outdated

let values_buf = entry.device.new_buffer_with_data(values_ptr, values_bytes, opts);
let ref_buf = entry.device.new_buffer_with_data(ref_ptr, ref_bytes, opts);
let result_buf = entry.device.new_buffer(std::mem::size_of::<i32>() as u64, opts);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the CUDA version, the result buffer is initialized to 0. Here it doesn't seem to be initialized, is it?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

great catch, thanks!

Luthaf added 3 commits July 14, 2026 11:38
This way we don't have to keep uploading small values to the
device over and over again
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants