Crate kn_runtime

source ·
Expand description

A wrapper crate around kn-graph and kn-cuda-eval that allows selecting whether to use a CPU or GPU at runtime through the Device type.

By default this crate only includes the Device::Cpu device, and only depends on kn-graph. To enable the Device::Cuda device, enable the cuda cargo feature. This adds a dependency on kn-cuda-eval and the cuda libraries.

This crate is part of the Kyanite project, see its readme for more information. See system-requirements for how to set up the cuda libraries.

§Quick demo

// load and optimize a graph
let graph = load_graph_from_onnx_path("test.onnx", false)?;
let graph = optimize_graph(&graph, Default::default());

// select a device, at runtime
let device_str = "cpu";
let device = match device_str {
    "auto" => Device::best(),
    "cpu" => Device::Cpu,
    "cuda" => Device::Cuda(CudaDevice::all().next().unwrap()),
    _ => panic!("unknown device"),
};

// prepare the graph for execution
let batch_size = 8;
let mut  prepared = device.prepare(graph, batch_size);

// evaluate the graph with some inputs, get the outputs back
let inputs = [DTensor::F32(Tensor::zeros(vec![batch_size, 16]))];
let outputs: Vec<DTensor> = prepared.eval(&inputs);

Structs§

Enums§

  • Device
    A device that can be used to evaluate a graph.
  • PreparedGraph
    A graph that has been prepared for evaluation on a device.

Functions§