Eerie 👻

Eerie 是 IREE 库的 Rust 绑定。它的目标是提供一个安全且健壮的 API，接近 IREE C API。

顺便说一下，这个包是实验性的，远未完成。它可能包含一些不稳定的代码，API 在未来可能会有破坏性的变化。如果您遇到任何问题，请随时提交问题。

我能用这个做什么？

Rust 有几个 ML 框架，如 Candle 和 Burn，但它们都不支持 JIT/AOT 模型编译。通过使用 IREE 编译器/运行时，可以构建一个完整的 ML 库，可以在运行时生成编译后的模型。

此外，IREE 的运行时可以在裸机环境中小到约 30KB，因此将来可以用于在嵌入式 Rust 中部署 ML 模型。

支持的操作系统

• MacOS
• Linux
• Windows
• 裸机 (thumb, rv32)

示例

将 MLIR 代码编译到 IREE VM 图形缓冲区

// simple_mul.mlir

module @arithmetic {
  func.func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
    %0 = arith.mulf %arg0, %arg1 : tensor<4xf32>
    return %0 : tensor<4xf32>
  }
}

fn output_vmfb() -> Vec<u8> {
    use eerie::compiler::*;
    use std::path::Path;
    let compiler = Compiler::new().unwrap();
    let mut session = compiler.create_session();
    session
        .set_flags(vec![
            "--iree-hal-target-backends=llvm-cpu".to_string(),
        ])
        .unwrap();
    let source = Source::from_file(&session, Path::new("simple_mul.mlir")).unwrap();
    let mut invocation = session.create_invocation();
    let mut output = MemBufferOutput::new(&compiler).unwrap();
    invocation
        .parse_source(source)
        .unwrap()
        .set_verify_ir(true)
        .set_compile_to_phase("end")
        .unwrap()
        .pipeline(Pipeline::Std)
        .unwrap()
        .output_vm_byte_code(&mut output)
        .unwrap();
    Vec::from(output.map_memory().unwrap())
}

在 IREE 运行时环境中运行张量操作

fn run_vmfb(vmfb: &[u8]) -> Vec<f32> {
    use eerie::runtime::*;
    use eerie::runtime::vm::{List, ToRef};
    let instance = api::Instance::new(
        &api::InstanceOptions::new(&mut hal::DriverRegistry::new())
            .use_all_available_drivers(),
    )
    .unwrap();
    let device = instance
        .try_create_default_device("local-task")
        .expect("Failed to create device");
    let session = api::Session::create_with_device(
        &instance,
        &api::SessionOptions::default(),
        &device,
    )
    .unwrap();
    unsafe { session.append_module_from_memory(vmfb) }.unwrap();
    let function = session.lookup_function("arithmetic.simple_mul").unwrap();
    let input_list = vm::DynamicList::<vm::Ref<hal::BufferView<f32>>>::new(
        2, &instance,
        )
        .unwrap();
    let input_buffer = hal::BufferView::<f32>::new(
        &session,
        &[4],
        hal::EncodingType::DenseRowMajor,
        &[1.0, 2.0, 3.0, 4.0]
    ).unwrap();
    let input_buffer_ref = input_buffer.to_ref(&instance).unwrap();
    input_list.push_ref(&input_buffer_ref).unwrap();
    let output_list = vm::DynamicList::<vm::Ref<hal::BufferView<f32>>>::new(
        1, &instance,
        )
        .unwrap();
    function.invoke(&input_list, &output_list).unwrap();
    let output_buffer_ref = output_list.get_ref(0).unwrap();
    let output_buffer: hal::BufferView<f32> = output_buffer_ref.to_buffer_view(&session);
    let output_mapping = hal::BufferMapping::new(output_buffer).unwrap();
    let out = output_mapping.data().to_vec();
    out
}

更多示例 这里

安装

该包分为两个部分：编译器和运行时。您可以通过切换 "编译器" 和 "运行时" 功能标志来启用每个功能。

运行时

Eerie 在编译时从源代码构建 IREE 运行时。需要 CMake 和 Clang。

MacOS

安装 XCode 和 MacOS SDK。

无 std

运行时库可以不使用默认的 std 功能进行编译。这需要一个 C/C++ 嵌入式工具链 (arm-none-eabi-gcc/riscv64-unknown-elf-gcc)，并在 sysroot 中预编译的 Newlib 二进制文件。

编译器

用户必须源预编译的共享库。（这是必要的，因为构建编译器需要大约 20 分钟）可以从 iree-compiler 的 Python 包安装中源共享库。

pip3 install iree-compiler

为了导出已安装库的位置，运行此脚本

python -c "import iree.compiler as _; print(f'{_.__path__[0]}/_mlir_libs/')"

然后，根据需要设置 rpath 和环境变量。这可以通过向您的项目目录添加以下 .cargo/config.toml 来完成。

MacOS

[build]
rustflags = ["-C", "link-arg=-Wl,-rpath,/path/to/library/"]
rustdocflags = ["-C", "link-arg=-Wl,-rpath,/path/to/library"]
[env]
LIB_IREE_COMPILER = "/path/to/library"

Linux

[build]
rustflags = ["-C", "link-arg=-Wl,-rpath=/path/to/library/"]
rustdocflags = ["-C", "link-arg=-Wl,-rpath=/path/to/library"]
[env]
LIB_IREE_COMPILER = "/path/to/library"

参考文献

• 也可以查看 SamKG/iree-rs
• Rustic MLIR 绑定 raviqqe/melior

许可证

Apache 2.0