spirv_backend: use RuntimeArray to remove all this UnsizedArray stuff

This commit is contained in:
2022-07-26 15:58:23 -07:00
parent 09f7c8acb9
commit d93d14d260
7 changed files with 149 additions and 323 deletions

View File

@@ -144,11 +144,11 @@ fn flat_idx(dim: Vec3u, idx: Vec3u) -> usize {
}
fn sample_pos<R: Copy + Default>(arr: &[Vec3<R>], dim: Vec3u, idx: Vec3u) -> VolumeSamplePos<R> {
VolumeSamplePos::from_indexable(DimIndexer::new(arr, dim), dim, idx)
VolumeSamplePos::from_indexable(&DimIndexer::new(arr, dim), dim, idx)
}
fn sample_neg<R: Copy + Default>(arr: &[Vec3<R>], dim: Vec3u, idx: Vec3u) -> VolumeSampleNeg<R> {
VolumeSampleNeg::from_indexable(DimIndexer::new(arr, dim), idx)
VolumeSampleNeg::from_indexable(&DimIndexer::new(arr, dim), idx)
}
#[cfg(test)]

View File

@@ -1,5 +1,5 @@
use futures::FutureExt as _;
use log::{info, warn};
use log::info;
use std::borrow::Cow;
use std::marker::PhantomData;
use std::num::NonZeroU64;
@@ -40,14 +40,6 @@ impl<R, M: 'static> WgpuBackend<R, M> {
let max_array_size = volume * max_elem_size as u64;
let max_buf_size = max_array_size + 0x1000; // allow some overhead
if max_array_size >= spirv_backend::support::MAX_UNSIZED_ARRAY as u64 {
warn!(
"array ({:?} elements) is too large to be used by SPIRV (limit: {})!",
volume,
spirv_backend::support::MAX_UNSIZED_ARRAY,
);
}
let entry_names = entry_points::<M>().unwrap_or(("invalid_mat", "invalid_mat"));
let (device, queue) = futures::executor::block_on(open_device(max_buf_size));
let shader_binary = get_shader();

View File

@@ -1,159 +1,101 @@
use spirv_std::RuntimeArray;
use coremem_types::step::{StepEContext, StepHContext, VolumeSampleNeg, VolumeSamplePos};
use crate::support::{Array3, Array3Mut, UnsizedArray};
use coremem_types::mat::Material;
use coremem_types::real::Real;
use coremem_types::step::SimMeta;
use coremem_types::vec::{Vec3, Vec3u};
use crate::support::{DimensionedSlice, SizedArray};
pub(crate) fn step_h<R: Real, M: Material<R>>(
id: Vec3u,
idx: Vec3u,
meta: &SimMeta<R>,
stimulus_h: &UnsizedArray<Vec3<R>>,
material: &UnsizedArray<M>,
e: &UnsizedArray<Vec3<R>>,
h: &mut UnsizedArray<Vec3<R>>,
m: &mut UnsizedArray<Vec3<R>>,
stimulus_h: &RuntimeArray<Vec3<R>>,
material: &RuntimeArray<M>,
e: &RuntimeArray<Vec3<R>>,
h: &mut RuntimeArray<Vec3<R>>,
m: &mut RuntimeArray<Vec3<R>>,
) {
if id.x() < meta.dim.x() && id.y() < meta.dim.y() && id.z() < meta.dim.z() {
let sim_state = SerializedStepH::new(meta, stimulus_h, material, e, h, m);
let update_state = sim_state.index(id);
let (new_h, new_m) = update_state.step_h();
sim_state.write_output(id, new_h, new_m);
}
}
if idx.x() < meta.dim.x() && idx.y() < meta.dim.y() && idx.z() < meta.dim.z() {
let dim = meta.dim;
let len = dim.product_sum_usize();
pub(crate) fn step_e<R: Real, M: Material<R>>(
id: Vec3u,
meta: &SimMeta<R>,
stimulus_e: &UnsizedArray<Vec3<R>>,
material: &UnsizedArray<M>,
e: &mut UnsizedArray<Vec3<R>>,
h: &UnsizedArray<Vec3<R>>,
) {
if id.x() < meta.dim.x() && id.y() < meta.dim.y() && id.z() < meta.dim.z() {
let sim_state = SerializedStepE::new(meta, stimulus_e, material, e, h);
let stim_h_array = unsafe { SizedArray::new(stimulus_h, len) };
let stim_h_matrix = DimensionedSlice::new(dim, &stim_h_array);
let update_state = sim_state.index(id);
let new_e = update_state.step_e();
sim_state.write_output(id, new_e);
}
}
let mat_array = unsafe { SizedArray::new(material, len) };
let mat_matrix = DimensionedSlice::new(dim, &mat_array);
let e_array = unsafe { SizedArray::new(e, len) };
let e_matrix = DimensionedSlice::new(dim, &e_array);
/// Whatever data we received from the host in their call to step_h
struct SerializedStepH<'a, R, M> {
meta: &'a SimMeta<R>,
stimulus_h: &'a UnsizedArray<Vec3<R>>,
material: &'a UnsizedArray<M>,
e: &'a UnsizedArray<Vec3<R>>,
h: &'a mut UnsizedArray<Vec3<R>>,
m: &'a mut UnsizedArray<Vec3<R>>,
}
let mut h_array = unsafe { SizedArray::new(h, len) };
let mut h_matrix = DimensionedSlice::new(dim, &mut h_array);
impl<'a, R, M> SerializedStepH<'a, R, M> {
fn new(
meta: &'a SimMeta<R>,
stimulus_h: &'a UnsizedArray<Vec3<R>>,
material: &'a UnsizedArray<M>,
e: &'a UnsizedArray<Vec3<R>>,
h: &'a mut UnsizedArray<Vec3<R>>,
m: &'a mut UnsizedArray<Vec3<R>>,
) -> Self {
Self {
meta, stimulus_h, material, e, h, m
}
}
}
let mut m_array = unsafe { SizedArray::new(m, len) };
let mut m_matrix = DimensionedSlice::new(dim, &mut m_array);
impl<'a, R: Real, M> SerializedStepH<'a, R, M> {
/// returns a context which the user can call `step_h` on
fn index(&self, idx: Vec3u) -> StepHContext<'a, R, M> {
let dim = self.meta.dim;
let stim_h_matrix = Array3::new(self.stimulus_h, dim);
let mat_matrix = Array3::new(self.material, dim);
let e = Array3::new(self.e, dim);
let h = Array3::new(self.h, dim);
let m = Array3::new(self.m, dim);
let stim_h = stim_h_matrix[idx];
let mat = &mat_matrix[idx];
let in_e = VolumeSamplePos::from_indexable(&e_matrix, dim, idx);
let in_h = h_matrix[idx];
let in_m = m_matrix[idx];
let in_e = VolumeSamplePos::from_indexable(e, dim, idx);
let in_h = h.get(idx).unwrap();
let in_m = m.get(idx).unwrap();
let mat = mat_matrix.into_ref(idx);
StepHContext {
inv_feature_size: self.meta.inv_feature_size,
time_step: self.meta.time_step,
stim_h: stim_h_matrix.get(idx).unwrap(),
let update_state = StepHContext {
inv_feature_size: meta.inv_feature_size,
time_step: meta.time_step,
stim_h,
mat,
in_e,
in_h,
in_m,
};
let (new_h, new_m) = update_state.step_h();
h_matrix[idx] = new_h;
m_matrix[idx] = new_m;
}
}
fn write_output(self, idx: Vec3u, h: Vec3<R>, m: Vec3<R>) {
let dim = self.meta.dim;
let arr_h = Array3Mut::new(self.h, dim);
let arr_m = Array3Mut::new(self.m, dim);
let mut out_h = arr_h.into_mut_handle(idx);
let mut out_m = arr_m.into_mut_handle(idx);
out_h.write(h);
out_m.write(m);
}
}
pub(crate) fn step_e<R: Real, M: Material<R>>(
idx: Vec3u,
meta: &SimMeta<R>,
stimulus_e: &RuntimeArray<Vec3<R>>,
material: &RuntimeArray<M>,
e: &mut RuntimeArray<Vec3<R>>,
h: &RuntimeArray<Vec3<R>>,
) {
if idx.x() < meta.dim.x() && idx.y() < meta.dim.y() && idx.z() < meta.dim.z() {
let dim = meta.dim;
let len = dim.product_sum_usize();
/// Whatever data we received from the host in their call to step_e
struct SerializedStepE<'a, R, M> {
meta: &'a SimMeta<R>,
stimulus_e: &'a UnsizedArray<Vec3<R>>,
material: &'a UnsizedArray<M>,
e: &'a mut UnsizedArray<Vec3<R>>,
h: &'a UnsizedArray<Vec3<R>>,
}
let stim_e_array = unsafe { SizedArray::new(stimulus_e, len) };
let stim_e_matrix = DimensionedSlice::new(dim, &stim_e_array);
impl<'a, R, M> SerializedStepE<'a, R, M> {
pub fn new(
meta: &'a SimMeta<R>,
stimulus_e: &'a UnsizedArray<Vec3<R>>,
material: &'a UnsizedArray<M>,
e: &'a mut UnsizedArray<Vec3<R>>,
h: &'a UnsizedArray<Vec3<R>>,
) -> Self {
Self {
meta, stimulus_e, material, e, h
}
}
}
let mat_array = unsafe { SizedArray::new(material, len) };
let mat_matrix = DimensionedSlice::new(dim, &mat_array);
impl<'a, R: Real, M> SerializedStepE<'a, R, M> {
fn index(&self, idx: Vec3u) -> StepEContext<'a, R, M> {
let dim = self.meta.dim;
let stim_e_matrix = Array3::new(self.stimulus_e, dim);
let mat_matrix = Array3::new(self.material, dim);
let e = Array3::new(self.e, dim);
let h = Array3::new(self.h, dim);
let mut e_array = unsafe { SizedArray::new(e, len) };
let mut e_matrix = DimensionedSlice::new(dim, &mut e_array);
let in_e = e.get(idx).unwrap();
let in_h = VolumeSampleNeg::from_indexable(h, idx);
let h_array = unsafe { SizedArray::new(h, len) };
let h_matrix = DimensionedSlice::new(dim, &h_array);
let mat = mat_matrix.into_ref(idx);
let stim_e = stim_e_matrix[idx];
let mat = &mat_matrix[idx];
let in_e = e_matrix[idx];
let in_h = VolumeSampleNeg::from_indexable(&h_matrix, idx);
StepEContext {
inv_feature_size: self.meta.inv_feature_size,
time_step: self.meta.time_step,
stim_e: stim_e_matrix.get(idx).unwrap(),
let update_state = StepEContext {
inv_feature_size: meta.inv_feature_size,
time_step: meta.time_step,
stim_e,
mat,
in_h,
in_e,
}
}
fn write_output(self, idx: Vec3u, e: Vec3<R>) {
let dim = self.meta.dim;
let arr_e = Array3Mut::new(self.e, dim);
let mut out_e = arr_e.into_mut_handle(idx);
out_e.write(e);
};
let new_e = update_state.step_e();
e_matrix[idx] = new_e;
}
}

View File

@@ -8,14 +8,12 @@
extern crate spirv_std;
pub use spirv_std::glam;
pub use spirv_std::{glam, RuntimeArray};
#[cfg(not(target_arch = "spirv"))]
use spirv_std::macros::spirv;
mod adapt;
pub mod support;
pub use support::UnsizedArray;
mod support;
use coremem_types::compound::Optional;
use coremem_types::mat::{Ferroxcube3R1MH, FullyGenericMaterial, IsoConductorOr};
@@ -56,12 +54,12 @@ macro_rules! steps {
#[spirv(global_invocation_id)] id: glam::UVec3,
#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] meta: &SimMeta<$flt>,
// XXX: delete this input?
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] _unused_stimulus_e: &UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 2)] stimulus_h: &UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &UnsizedArray<$mat>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &mut UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 6)] m: &mut UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] _unused_stimulus_e: &RuntimeArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 2)] stimulus_h: &RuntimeArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &RuntimeArray<$mat>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &RuntimeArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &mut RuntimeArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 6)] m: &mut RuntimeArray<Vec3<$flt>>,
) {
adapt::step_h(glam_vec_to_internal(id), meta, stimulus_h, material, e, h, m)
}
@@ -70,14 +68,14 @@ macro_rules! steps {
pub fn $step_e(
#[spirv(global_invocation_id)] id: glam::UVec3,
#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] meta: &SimMeta<$flt>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] stimulus_e: &UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] stimulus_e: &RuntimeArray<Vec3<$flt>>,
// XXX: delete this input?
#[spirv(storage_buffer, descriptor_set = 0, binding = 2)] _unused_stimulus_h: &UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &UnsizedArray<$mat>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &mut UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 2)] _unused_stimulus_h: &RuntimeArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &RuntimeArray<$mat>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &mut RuntimeArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &RuntimeArray<Vec3<$flt>>,
// XXX: can/should this m input be deleted?
#[spirv(storage_buffer, descriptor_set = 0, binding = 6)] _unused_m: &UnsizedArray<Vec3<$flt>>,
#[spirv(storage_buffer, descriptor_set = 0, binding = 6)] _unused_m: &RuntimeArray<Vec3<$flt>>,
) {
adapt::step_e(glam_vec_to_internal(id), meta, stimulus_e, material, e, h)
}

View File

@@ -1,200 +1,90 @@
use core::ops::Index;
use core::ops::{Index, IndexMut};
use spirv_std::RuntimeArray;
use coremem_types::compound::Optional;
use coremem_types::vec::Vec3u;
/// This struct allows doing things like *(ptr + offset) = value.
/// Such code wouldn't ordinarily compile with the spirv target because of
/// unsupported pointer math. RuntimeArray exists to overcome this, however
/// it emits invalid code for non-primitive types. Hence, this hack.
// XXX: maximum bytes an array may occupy is 0x7fff_ffff
// We don't know the element size, so assume it to be <= 44 bytes
// pub const MAX_UNSIZED_ARRAY: usize = 0x2e7_ffff;
pub const MAX_UNSIZED_ARRAY: usize = 0x1ff_ffff;
pub struct UnsizedArray<T>(pub [T; MAX_UNSIZED_ARRAY]);
pub struct SizedArray<T> {
items: T,
len: usize,
}
impl<T: Copy> UnsizedArray<T> {
/// user is responsible for bounds-checking the index before calling this.
/// unsafe because this can trigger OOB reads if used incorrectly.
pub unsafe fn index(&self, index: usize) -> T {
// *self.0.index(index)
self.0[index]
// *self.0.get_unchecked(index)
// &self.0
//asm! {
// "%result = OpAccessChain _ {arr} {index}",
// "OpReturnValue %result",
// "%unused = OpLabel",
// arr = in(reg) self,
// index = in(reg) index,
//}
impl<T> SizedArray<T> {
pub unsafe fn new(items: T, len: usize) -> Self {
Self { items, len }
}
}
impl<T> UnsizedArray<T> {
/// user is responsible for bounds-checking the index before calling this.
/// unsafe because this can trigger OOB writes if used incorrectly.
pub unsafe fn write(&mut self, index: usize, value: T) {
self.0[index] = value;
}
}
impl<'a, T> Index<usize> for SizedArray<&'a RuntimeArray<T>> {
type Output=T;
impl<T> UnsizedArray<T> {
/// user is responsible for bounds-checking the index before calling this.
/// unsafe because this can trigger OOB access if used incorrectly.
pub unsafe fn index_ref(&self, index: usize) -> &T {
&self.0[index]
}
pub unsafe fn get_handle<'a>(&'a self, index: usize) -> ArrayHandle<'a, T> {
ArrayHandle {
base: self,
index,
}
}
pub unsafe fn get_handle_mut<'a>(&'a mut self, index: usize) -> ArrayHandleMut<'a, T> {
ArrayHandleMut {
base: self,
index,
}
}
}
pub struct ArrayHandle<'a, T> {
base: &'a UnsizedArray<T>,
index: usize,
}
impl<'a, T: Copy> ArrayHandle<'a, T> {
pub fn get(&self) -> T {
fn index(&self, idx: usize) -> &Self::Output {
assert!(idx < self.len);
unsafe {
self.base.index(self.index)
self.items.index(idx)
}
}
pub fn get_into(&self, out: &mut T) {
core::clone::Clone::clone_from(out, self.get_ref());
}
}
impl<'a, T> ArrayHandle<'a, T> {
pub fn get_ref(&self) -> &'a T {
impl<'a, T> Index<usize> for SizedArray<&'a mut RuntimeArray<T>> {
type Output=T;
fn index(&self, idx: usize) -> &Self::Output {
assert!(idx < self.len);
unsafe {
self.base.index_ref(self.index)
self.items.index(idx)
}
}
}
pub struct ArrayHandleMut<'a, T> {
base: &'a mut UnsizedArray<T>,
index: usize,
}
impl<'a, T: Copy> ArrayHandleMut<'a, T> {
pub fn write(&mut self, value: T) {
impl<'a, T> IndexMut<usize> for SizedArray<&'a mut RuntimeArray<T>> {
fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
assert!(idx < self.len);
unsafe {
self.base.write(self.index, value);
}
}
pub fn get(&self) -> T {
unsafe {
self.base.index(self.index)
self.items.index_mut(idx)
}
}
}
/// 3d dynamically-sized array backed by a borrowed buffer.
#[derive(Clone, Copy)]
pub struct Array3<'a, T> {
data: &'a UnsizedArray<T>,
pub struct DimensionedSlice<T> {
dim: Vec3u,
items: T,
}
impl<T> DimensionedSlice<T> {
pub fn new(dim: Vec3u, items: T) -> Self {
Self { dim, items }
}
}
impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> for DimensionedSlice<&'a T> {
type Output=T::Output;
fn index(&self, idx: Vec3u) -> &Self::Output {
let idx = index(idx, self.dim);
&self.items[idx]
}
}
impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> for DimensionedSlice<&'a mut T> {
type Output=T::Output;
fn index(&self, idx: Vec3u) -> &Self::Output {
let idx = index(idx, self.dim);
&self.items[idx]
}
}
impl<'a, T: IndexMut<usize> + ?Sized> IndexMut<Vec3u> for DimensionedSlice<&'a mut T> {
fn index_mut(&mut self, idx: Vec3u) -> &mut Self::Output {
let idx = index(idx, self.dim);
&mut self.items[idx]
}
}
fn index(loc: Vec3u, dim: Vec3u) -> usize {
((loc.z()*dim.y() + loc.y())*dim.x() + loc.x()) as usize
}
fn checked_index(idx: Vec3u, dim: Vec3u) -> Optional<usize> {
if idx.x() < dim.x() && idx.y() < dim.y() && idx.z() < dim.z() {
let flat_idx = index(idx, dim);
Optional::some(flat_idx)
} else {
Optional::none()
}
}
impl<'a, T> Array3<'a, T> {
pub fn new(data: &'a UnsizedArray<T>, dim: Vec3u) -> Self {
Self {
data,
dim,
}
}
pub fn into_ref(self, idx: Vec3u) -> &'a T {
let idx = checked_index(idx, self.dim).unwrap();
unsafe {
self.data.index_ref(idx)
}
}
}
impl<'a, T: Copy + Default> Array3<'a, T> {
pub fn get(&self, idx: Vec3u) -> Optional<T> {
let idx = checked_index(idx, self.dim);
if idx.is_some() {
Optional::some(unsafe {
self.data.index(idx.unwrap())
})
} else {
Optional::none()
}
}
}
impl<'a, T> Index<Vec3u> for Array3<'a, T> {
type Output = T;
fn index(&self, idx: Vec3u) -> &Self::Output {
let idx = checked_index(idx, self.dim).unwrap();
unsafe {
self.data.index_ref(idx)
}
}
}
/// 3d dynamically-sized array backed by a mutably borrowed buffer.
pub struct Array3Mut<'a, T> {
data: &'a mut UnsizedArray<T>,
dim: Vec3u,
}
impl<'a, T> Array3Mut<'a, T> {
pub fn new(data: &'a mut UnsizedArray<T>, dim: Vec3u) -> Self {
Self {
data,
dim,
}
}
pub fn index(&self, idx: Vec3u) -> Optional<usize> {
if idx.x() < self.dim.x() && idx.y() < self.dim.y() && idx.z() < self.dim.z() {
let flat_idx = index(idx, self.dim);
Optional::some(flat_idx)
} else {
Optional::none()
}
}
}
impl<'a, T: Copy> Array3Mut<'a, T> {
pub fn into_mut_handle(self, idx: Vec3u) -> ArrayHandleMut<'a, T> {
let idx = self.index(idx).unwrap();
unsafe {
self.data.get_handle_mut(idx)
}
}
}
#[cfg(test)]
mod test {

View File

@@ -45,7 +45,7 @@ pub struct VolumeSampleNeg<R> {
}
impl<R: Copy + Default> VolumeSampleNeg<R> {
pub fn from_indexable<I: Index<Vec3u, Output=Vec3<R>>>(i: I, idx: Vec3u) -> Self {
pub fn from_indexable<I: Index<Vec3u, Output=Vec3<R>>>(i: &I, idx: Vec3u) -> Self {
VolumeSampleNeg {
mid: i[idx],
xm1: prev_x(idx).map(|idx| i[idx]),
@@ -132,7 +132,7 @@ pub struct VolumeSamplePos<R> {
}
impl<R: Copy + Default> VolumeSamplePos<R> {
pub fn from_indexable<I: Index<Vec3u, Output=Vec3<R>>>(i: I, dim: Vec3u, idx: Vec3u) -> Self {
pub fn from_indexable<I: Index<Vec3u, Output=Vec3<R>>>(i: &I, dim: Vec3u, idx: Vec3u) -> Self {
VolumeSamplePos {
mid: i[idx],
xp1: next_x(dim, idx).map(|idx| i[idx]),

View File

@@ -51,6 +51,10 @@ impl Vec3u {
pub fn product_sum(&self) -> u32 {
self.x * self.y * self.z
}
pub fn product_sum_usize(&self) -> usize {
self.x as usize * self.y as usize * self.z as usize
}
}
impl From<(u32, u32, u32)> for Vec3u {