use crate::{Dims, Error, Output};
#[derive(Debug, Fail, Clone, Eq, PartialEq)]
pub enum WorkError {
#[fail(display = "Work does not allow a zero value for any of its specified dimenions")]
DimLengthCannotBeZero,
#[fail(display = "Work dimensions must be 1, 2, or 3.")]
InvalidDimsCount,
#[fail(display = "Work size dimensions cannot have any zero values")]
InvalidWorkSize,
}
const INVALID_WORK_SIZE: Error = Error::WorkError(WorkError::InvalidWorkSize);
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct NonZeroVolume3DArray([usize; 3]);
impl NonZeroVolume3DArray {
fn as_ptr(&self) -> *const usize {
self.0.as_ptr()
}
}
impl NonZeroVolume3DArray {
fn from_dims(d: &Dims) -> Output<NonZeroVolume3DArray> {
let work_size: [usize; 3] = match d {
Dims::One(x) => [*x, 1, 1],
Dims::Two(x, y) => [*x, *y, 1],
Dims::Three(x, y, z) => [*x, *y, *z],
};
match work_size {
[0, _, _] | [_, 0, _] | [_, _, 0] => Err(INVALID_WORK_SIZE),
w => Ok(NonZeroVolume3DArray(w)),
}
}
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct GlobalWorkSize(NonZeroVolume3DArray);
impl GlobalWorkSize {
pub fn from_dims(dims: &Dims) -> Output<GlobalWorkSize> {
let arr = NonZeroVolume3DArray::from_dims(dims)?;
Ok(GlobalWorkSize(arr))
}
pub fn as_ptr(&self) -> *const usize {
self.0.as_ptr()
}
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum LocalWorkSize {
WorkSize(NonZeroVolume3DArray),
Null,
}
impl LocalWorkSize {
pub fn from_dims(dims: &Dims) -> Output<LocalWorkSize> {
let arr = NonZeroVolume3DArray::from_dims(dims)?;
Ok(LocalWorkSize::WorkSize(arr))
}
pub fn as_ptr(&self) -> *const usize {
match self {
LocalWorkSize::WorkSize(arr) => arr.as_ptr(),
LocalWorkSize::Null => std::ptr::null::<usize>(),
}
}
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum GlobalWorkOffset {
Offset([usize; 3]),
Null,
}
impl GlobalWorkOffset {
pub fn from_dims(d: &Dims) -> GlobalWorkOffset {
let work_offset: [usize; 3] = match d {
Dims::One(x) => [*x, 0, 0],
Dims::Two(x, y) => [*x, *y, 0],
Dims::Three(x, y, z) => [*x, *y, *z],
};
GlobalWorkOffset::Offset(work_offset)
}
pub fn as_ptr(&self) -> *const usize {
match self {
GlobalWorkOffset::Offset(arr) => arr.as_ptr(),
GlobalWorkOffset::Null => std::ptr::null::<usize>(),
}
}
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Work {
pub global_size: Dims,
pub global_offset: Option<Dims>,
pub local_size: Option<Dims>,
}
impl Work {
pub fn new<D: Into<Dims>>(global_size: D) -> Work {
Work {
global_size: global_size.into(),
global_offset: None,
local_size: None,
}
}
pub fn with_global_offset<D: Into<Dims>>(mut self, offset: D) -> Work {
self.global_offset = Some(offset.into());
self
}
pub fn with_local_size<D: Into<Dims>>(mut self, local_size: D) -> Work {
self.local_size = Some(local_size.into());
self
}
pub fn work_dims(&self) -> u32 {
self.global_size.n_dimensions() as u32
}
pub fn global_work_size(&self) -> Output<GlobalWorkSize> {
GlobalWorkSize::from_dims(&self.global_size)
}
pub fn global_work_offset(&self) -> GlobalWorkOffset {
match &self.global_offset {
Some(dims) => GlobalWorkOffset::from_dims(dims),
None => GlobalWorkOffset::Null,
}
}
pub fn local_work_size(&self) -> Output<LocalWorkSize> {
match &self.local_size {
Some(dims) => LocalWorkSize::from_dims(dims),
None => Ok(LocalWorkSize::Null),
}
}
pub fn n_items(&self) -> usize {
self.global_size.n_items()
}
}