Skip to content

witness: add optional device backing metadata to RowMajorMatrix #33

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

Open
hero78119 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

witness: add optional device backing metadata to RowMajorMatrix #33

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e4ae45f
witness: add optional device backing metadata to RowMajorMatrix
hero78119 Apr 13, 2026
9e0b96e
Merge branch 'main' into feat/witness_gpu_buffer
hero78119 Apr 14, 2026
0d96edb
Merge branch 'main' into feat/witness_gpu_buffer
hero78119 Apr 14, 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
47 changes: 47 additions & 0 deletions crates/witness/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use rayon::{
prelude::ParallelSliceMut,
};
use std::{
any::Any,
ops::{Deref, DerefMut, Index},
slice::{Chunks, ChunksMut},
sync::Arc,
Expand Down Expand Up @@ -41,6 +42,19 @@ pub struct RowMajorMatrix<T: Sized + Sync + Clone + Send + Copy> {
log2_num_rotation: usize,
is_padded: bool,
padding_strategy: InstancePaddingStrategy,
device_backing: Option<DeviceMatrixBacking>,
}
Comment on lines 43 to +46
Copy link

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

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

RowMajorMatrix derives Clone, and the new device_backing: Option<DeviceMatrixBacking> will be cloned as well (cloning the inner Arc). Since inner is also cloned, the cloned matrix will share the same device backing metadata as the original even though the host-side values are a deep copy, which is very likely to make the backing inconsistent/wrong. Consider implementing a manual Clone for RowMajorMatrix that resets device_backing to None (or otherwise ensures the backing corresponds to the cloned contents).

Copilot uses AI. Check for mistakes.

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeviceMatrixLayout {
RowMajor,
ColMajor,
}

#[derive(Clone)]
struct DeviceMatrixBacking {
storage: Arc<dyn Any + Send + Sync>,
layout: DeviceMatrixLayout,
}

impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMatrix<T> {
Expand All @@ -56,6 +70,7 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
is_padded: true,
log2_num_rotation: 0,
padding_strategy: InstancePaddingStrategy::Default,
device_backing: None,
}
}
pub fn empty() -> Self {
Expand All @@ -65,6 +80,7 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
log2_num_rotation: 0,
is_padded: true,
padding_strategy: InstancePaddingStrategy::Default,
device_backing: None,
}
}

Expand Down Expand Up @@ -130,6 +146,7 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
log2_num_rotation,
is_padded: matches!(padding_strategy, InstancePaddingStrategy::Default),
padding_strategy,
device_backing: None,
}
}

Expand All @@ -148,6 +165,7 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
log2_num_rotation: 0,
is_padded: matches!(padding_strategy, InstancePaddingStrategy::Default),
padding_strategy,
device_backing: None,
}
}

Expand All @@ -166,6 +184,35 @@ impl<T: Sized + Sync + Clone + Send + Copy + Default + FieldAlgebra> RowMajorMat
next_pow2_instance_padding(self.num_instances()) - self.num_instances()
}

pub fn set_device_backing<D: Any + Send + Sync + 'static>(
&mut self,
storage: D,
layout: DeviceMatrixLayout,
) {
self.device_backing = Some(DeviceMatrixBacking {
storage: Arc::new(storage),
layout,
});
}
Comment on lines +187 to +196
Copy link

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

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

device_backing can become silently stale: RowMajorMatrix exposes multiple mutable access paths (e.g., DerefMut and the various *_mut iterators/padding methods) but setting a device backing doesn’t establish/maintain any invariant that it matches the current host-side inner.values. If callers set a device backing and then mutate the matrix, downstream consumers may read an out-of-sync device buffer. Consider defining the invariant explicitly and either (a) automatically invalidating device_backing on any host mutation (clear on deref_mut / mut iterators / padding / pad_to_height), or (b) restricting mutable access when a backing is present and forcing callers to clear/update it first.

Copilot uses AI. Check for mistakes.

pub fn clear_device_backing(&mut self) {
self.device_backing = None;
}

pub fn has_device_backing(&self) -> bool {
self.device_backing.is_some()
}
Comment on lines +198 to +204
Copy link

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

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

There’s no invalidation mechanism tying device_backing to the current host-side contents/dimensions: the matrix can be mutated via existing APIs (e.g., DerefMut, iter_mut, pad_to_height, padding_by_strategy, etc.) without clearing or updating device_backing. That can leave stale device metadata attached and risks later code using the wrong device buffer/layout. Consider clearing device_backing in mutating methods (or tracking a version/dirty flag and requiring callers to refresh the backing when dirty).

Copilot uses AI. Check for mistakes.

pub fn device_backing_layout(&self) -> Option<DeviceMatrixLayout> {
self.device_backing.as_ref().map(|backing| backing.layout)
}

pub fn device_backing_ref<D: Any + 'static>(&self) -> Option<&D> {
self.device_backing
.as_ref()
.and_then(|backing| backing.storage.downcast_ref::<D>())
}
Comment on lines +210 to +214
Copy link

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

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

device_backing_ref accepts D: Any + 'static, but set_device_backing only permits storing D: Any + Send + Sync + 'static. This mismatch makes it easy to ask for a non-Send + Sync type that can never be present, and is inconsistent with the setter’s contract. Consider mirroring the bounds (D: Any + Send + Sync + 'static) or documenting the intended usage to avoid confusing API calls that always return None.

Copilot uses AI. Check for mistakes.

// return raw num_instances without rotation
pub fn num_instances(&self) -> usize {
self.num_rows
Expand Down
Loading