nu_plugin_vulkan_matrix 0.1.1

GPU matrix operations for Nushell using Vulkan
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

use super::*;
use std::path::PathBuf;

fn path(engine: &EngineInterface, call: &EvaluatedCall) -> Result<PathBuf, LabeledError> {
    let path = PathBuf::from(call.req::<String>(0)?);
    Ok(if path.is_absolute() {
        path
    } else {
        PathBuf::from(engine.get_current_dir()?).join(path)
    })
}

pub(super) struct FromNpy;
impl SimplePluginCommand for FromNpy {
    type Plugin = MatrixPlugin;
    fn name(&self) -> &str {
        "vulkan matrix from-npy"
    }
    fn description(&self) -> &str {
        "Import a C-order f32/f64 NumPy file without a userspace payload copy."
    }
    fn signature(&self) -> Signature {
        Signature::build(SimplePluginCommand::name(self))
            .input_output_type(Type::Nothing, Type::Custom("vulkan matrix".into()))
            .required("path", SyntaxShape::Filepath, "NumPy .npy input file.")
            .named("device", SyntaxShape::String, "Explicit GPU selector (dN or iN).", None)
            .switch("shared", "Keep a read-only file mapping as the GPU tensor; requires an aligned header. Keep the file unchanged while in use.", None)
    }
    fn run(
        &self,
        plugin: &MatrixPlugin,
        engine: &EngineInterface,
        call: &EvaluatedCall,
        _input: &Value,
    ) -> Result<Value, LabeledError> {
        let device: String = call
            .get_flag("device")?
            .ok_or_else(|| error("--device is required; select dN or iN", call.head))?;
        let gpu = plugin.device(&device, call.head)?;
        let tensor = crate::file_io::load(&gpu, &path(engine, call)?, call.has_flag("shared")?)
            .map_err(|e| error(e.to_string(), call.head))?;
        let dtype = if tensor.dtype == <f64 as cubecl::CubeElement>::cube_type() {
            "f64"
        } else {
            "f32"
        };
        plugin.register(engine, gpu, tensor, &device, dtype, call.head)
    }
}

pub(super) struct ToNpy;
impl SimplePluginCommand for ToNpy {
    type Plugin = MatrixPlugin;
    fn name(&self) -> &str {
        "vulkan matrix to-npy"
    }
    fn description(&self) -> &str {
        "Write GPU matrix data directly into NumPy file pages and sync the file."
    }
    fn signature(&self) -> Signature {
        Signature::build(SimplePluginCommand::name(self))
            .input_output_type(Type::Custom("vulkan matrix".into()), Type::Nothing)
            .required("path", SyntaxShape::Filepath, "NumPy .npy output file.")
            .switch("force", "Atomically replace an existing output file.", None)
    }
    fn run(
        &self,
        plugin: &MatrixPlugin,
        engine: &EngineInterface,
        call: &EvaluatedCall,
        input: &Value,
    ) -> Result<Value, LabeledError> {
        let resident = plugin
            .resources
            .lock()
            .map_err(|e| error(e.to_string(), call.head))?
            .get(input, call.head)?;
        crate::file_io::save(
            &resident.gpu,
            &resident.tensor,
            &path(engine, call)?,
            call.has_flag("force")?,
        )
        .map_err(|e| error(e.to_string(), call.head))?;
        Ok(Value::nothing(call.head))
    }
}