sim-lib-compute-wgpu 0.1.0

Evidence-based wgpu tensor compute site discovery for SIM.
Documentation 
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

//! Portable kernel planning and host-equivalent execution.

use sim_kernel::{Cx, Symbol};
use sim_lib_numbers_tensor::{
    Tensor, TensorExecError, TensorRequest, add_op_symbol, bounded_element_count, cos_op_symbol,
    div_op_symbol, dot_op_symbol, exp_op_symbol, matmul_exec_op_symbol, max_op_symbol,
    min_op_symbol, mul_op_symbol, norm_op_symbol, sin_op_symbol, sqrt_op_symbol, sub_op_symbol,
    sum_op_symbol, transpose_exec_op_symbol,
};

use crate::{
    kernel_elementwise::{execute_binary, execute_unary},
    kernel_linalg::{execute_dot, execute_matmul, execute_transpose},
    kernel_reductions::execute_reduction,
    kernel_support::unsupported,
    pipeline::{WgpuKernelDType, WgpuKernelOp},
};

pub use crate::kernel_wgsl::{
    PORTABLE_ELEMENTWISE_WGSL, PORTABLE_LINALG_WGSL, PORTABLE_REDUCTION_WGSL,
};

/// Returns the portable kernel operation for a tensor request.
pub fn kernel_op(symbol: &Symbol) -> Option<WgpuKernelOp> {
    if *symbol == add_op_symbol() {
        Some(WgpuKernelOp::Add)
    } else if *symbol == sub_op_symbol() {
        Some(WgpuKernelOp::Sub)
    } else if *symbol == mul_op_symbol() {
        Some(WgpuKernelOp::Mul)
    } else if *symbol == div_op_symbol() {
        Some(WgpuKernelOp::Div)
    } else if *symbol == sqrt_op_symbol() {
        Some(WgpuKernelOp::Sqrt)
    } else if *symbol == exp_op_symbol() {
        Some(WgpuKernelOp::Exp)
    } else if *symbol == sin_op_symbol() {
        Some(WgpuKernelOp::Sin)
    } else if *symbol == cos_op_symbol() {
        Some(WgpuKernelOp::Cos)
    } else if *symbol == sum_op_symbol() {
        Some(WgpuKernelOp::Sum)
    } else if *symbol == min_op_symbol() {
        Some(WgpuKernelOp::Min)
    } else if *symbol == max_op_symbol() {
        Some(WgpuKernelOp::Max)
    } else if *symbol == norm_op_symbol() {
        Some(WgpuKernelOp::Norm)
    } else if *symbol == transpose_exec_op_symbol() {
        Some(WgpuKernelOp::Transpose)
    } else if *symbol == dot_op_symbol() {
        Some(WgpuKernelOp::Dot)
    } else if *symbol == matmul_exec_op_symbol() {
        Some(WgpuKernelOp::Matmul)
    } else {
        None
    }
}

/// Validates and evaluates one portable f32 kernel on host-equivalent values.
pub fn execute_portable_kernel(
    cx: &mut Cx,
    request: &TensorRequest,
    dtype: WgpuKernelDType,
) -> Result<Tensor, TensorExecError> {
    let op = kernel_op(&request.operation.symbol).ok_or_else(|| {
        unsupported(
            request.operation.symbol.clone(),
            "wgpu portable kernel does not support this operation",
        )
    })?;
    let shape = request.output.shape().to_vec();
    bounded_element_count(&shape).map_err(TensorExecError::from)?;
    let cells = if op.is_reduction() {
        execute_reduction(cx, op, request, dtype)?
    } else if op == WgpuKernelOp::Transpose {
        execute_transpose(request)?
    } else if op == WgpuKernelOp::Dot {
        execute_dot(cx, request, dtype)?
    } else if op == WgpuKernelOp::Matmul {
        execute_matmul(cx, request, dtype)?
    } else if op.is_binary() {
        execute_binary(cx, op, request, dtype, &shape)?
    } else {
        execute_unary(cx, op, request, dtype, &shape)?
    };
    sim_lib_numbers_tensor::build_tensor_value(
        cx,
        shape,
        Some(request.output.dtype().clone()),
        cells,
    )
    .map_err(TensorExecError::from)?
    .object()
    .downcast_ref::<Tensor>()
    .cloned()
    .ok_or_else(|| crate::kernel_support::invalid("wgpu kernel produced a non-tensor value"))
}