Struct Unary

pub struct Unary(pub ComputePipeline, pub UnaryOp);

Shader implementing various unary operations selected with UnaryOp.

Tuple Fields

0: ComputePipeline

1: UnaryOp

Implementations

impl Unary

pub const SRC: &'static str = "// This was mostly ported from ggml-cuda/unary.cu\n// (MIT license).\n\n#import wgblas::shape as Shape;\n\n@group(0) @binding(0)\nvar<uniform> shape_dst: Shape::Shape;\n@group(0) @binding(1)\nvar<uniform> shape_src: Shape::Shape;\n@group(0) @binding(2)\nvar<storage, read_write> dst: array<f32>;\n@group(0) @binding(3)\nvar<storage, read_write> src: array<f32>;\n@group(0) @binding(4)\nvar<uniform> args: vec4<f32>;\n\nconst GELU_COEF_A: f32 = 0.044715f;\nconst SQRT_2_OVER_PI: f32 = 0.79788456080286535587989211986876f;\nconst GELU_QUICK_COEF: f32 = -1.702f;\n\nfn abs_f32(x: f32) -> f32 {\n return abs(x);\n}\n\nfn sgn_f32(x: f32) -> f32 {\n return sign(x);\n}\n\nfn neg_f32(x: f32) -> f32 {\n return -x;\n}\n\nfn step_f32(x: f32) -> f32 {\n if x > 0.0 {\n return 1.0;\n } else {\n return 0.0;\n }\n}\n\nfn elu_f32(x: f32) -> f32 {\n if x > 0.0 {\n return x;\n } else {\n return exp(x) - 1.0;\n }\n}\n\nfn gelu_f32(x: f32) -> f32 {\n return 0.5f * x * (1.0f + tanh(SQRT_2_OVER_PI * x * (1.0f + GELU_COEF_A * x * x)));\n}\n\nfn gelu_quick_f32(x: f32) -> f32 {\n return x * (1.0f / (1.0f + exp(GELU_QUICK_COEF * x)));\n}\n\nfn silu_f32(x: f32) -> f32 {\n return x / (1.0 + exp(-x));\n}\n\nfn tanh_f32(x: f32) -> f32 {\n return tanh(x);\n}\n\nfn relu_f32(x: f32) -> f32 {\n return max(x, 0.0);\n}\n\nfn sigmoid_f32(x: f32) -> f32 {\n return 1.0 / (1.0 + exp(-x));\n}\n\nfn hard_sigmoid_f32(x: f32) -> f32 {\n return min(1.0, max(0.0, (x + 3.0) / 6.0));\n}\n\nfn hard_swish_f32(x: f32) -> f32 {\n return x * min(1.0, max(0.0, (x + 3.0) / 6.0));\n}\n\nfn sqr_f32(x: f32) -> f32 {\n return x * x;\n}\n\nfn sqrt_f32(x: f32) -> f32 {\n return sqrt(x);\n}\n\nfn log_f32(x: f32) -> f32 {\n return log(x);\n}\n\nfn placeholder(x: f32) -> f32 {\n return x;\n}\n\n// Unary ops with extra arguments passed through the args uniform.\nfn leaky_relu_f32(x: f32) -> f32 {\n return max(x, 0.0) + min(x, 0.0) * args.x;\n}\n\nfn clamp_f32(x: f32) -> f32 {\n return min(max(args.x, x), args.y);\n}\n\nfn scale_f32(x: f32) -> f32 {\n return x * args.x;\n}\n\nfn add_scalar_f32(x: f32) -> f32 {\n return x + args.x;\n}\n\n@compute @workgroup_size(64, 1, 1)\nfn main(@builtin(global_invocation_id) invocation_id: vec3<u32>) {\n if (invocation_id.x < shape_src.nrows) {\n let isrc = Shape::iv(shape_src, invocation_id.x);\n let idst = Shape::iv(shape_dst, invocation_id.x);\n dst[idst] = placeholder(src[isrc]);\n }\n}\n"

pub const FILE_PATH: &'static str = "wgml/src/unary.wgsl"

pub fn new(device: &Device, op: UnaryOp) -> Result<Self, ComposerError>

pub fn queue<'a, 'b, T: Pod + Scalar>( &'a self, queue: &mut KernelInvocationQueue<'a>, dest: impl Into<GpuVectorView<'b, T>>, src: impl Into<GpuVectorView<'b, T>>, args: Option<&'b GpuScalar<Vector4<T>>>, )

pub fn run_cpu<S: StorageMut<f32, Dyn>>( &self, vals: &mut Vector<f32, Dyn, S>, args: Vector4<f32>, )