File size: 3,183 Bytes
a8baeed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | use super::{Add, GpuContext, GpuError, Negate, Result, UnaryOp};
use ocl::Buffer;
/// Expression tree of GPU operations, e.g. `Negate(Add(x, y))`.
/// Every node executes as an OpenCL kernel; nothing is computed on the host.
#[derive(Clone, Debug)]
pub enum Expr {
Input(usize),
Negate(Box<Expr>),
Add(Box<Expr>, Box<Expr>),
}
enum Val<'a> {
Borrowed(&'a Buffer<u8>),
Owned(Buffer<u8>),
}
impl Val<'_> {
fn buf(&self) -> &Buffer<u8> {
match self {
Val::Borrowed(b) => b,
Val::Owned(b) => b,
}
}
}
impl Expr {
pub fn input(i: usize) -> Expr {
Expr::Input(i)
}
pub fn negate(e: Expr) -> Expr {
Expr::Negate(Box::new(e))
}
// Constructor named after the node it builds; `Expr` is not an arithmetic type.
#[allow(clippy::should_implement_trait)]
pub fn add(a: Expr, b: Expr) -> Expr {
Expr::Add(Box::new(a), Box::new(b))
}
/// Evaluate on the device over the first `count` elements of `inputs`.
/// Inputs are never mutated; the result is a fresh buffer of `count` elements.
pub fn eval(
&self,
ctx: &GpuContext,
inputs: &[&Buffer<u8>],
count: usize,
) -> Result<Buffer<u8>> {
if count == 0 {
return Err(GpuError::InvalidCount(0));
}
match self.eval_val(ctx, inputs, count)? {
Val::Owned(b) => Ok(b),
Val::Borrowed(src) => {
let out = ctx.alloc(count)?;
Identity.launch(ctx, src, &out, count)?;
Ok(out)
}
}
}
fn eval_val<'a>(
&self,
ctx: &GpuContext,
inputs: &[&'a Buffer<u8>],
count: usize,
) -> Result<Val<'a>> {
match self {
Expr::Input(i) => {
inputs
.get(*i)
.map(|b| Val::Borrowed(b))
.ok_or(GpuError::MissingInput {
index: *i,
supplied: inputs.len(),
})
}
Expr::Negate(e) => {
let v = e.eval_val(ctx, inputs, count)?;
let out = ctx.alloc(count)?;
Negate.launch(ctx, v.buf(), &out, count)?;
Ok(Val::Owned(out))
}
Expr::Add(a, b) => {
let va = a.eval_val(ctx, inputs, count)?;
let vb = b.eval_val(ctx, inputs, count)?;
let out = ctx.alloc(count)?;
Add.launch(ctx, va.buf(), vb.buf(), &out, count)?;
Ok(Val::Owned(out))
}
}
}
}
/// Device-side copy, used only when an expression is a bare `Input`.
struct Identity;
impl UnaryOp for Identity {
fn launch(
&self,
ctx: &GpuContext,
input: &Buffer<u8>,
output: &Buffer<u8>,
count: usize,
) -> Result<()> {
ctx.work_sizes(count, &[("input", input.len()), ("output", output.len())])?;
input
.cmd()
.queue(ctx.proque().queue())
.copy(output, Some(0), Some(count))
.enq()?;
Ok(())
}
}
|