use alloc::sync::Arc;
use core::sync::atomic::{AtomicBool, Ordering};
use crate::{
Result,
logging::LogLevel,
operator::{BoxedKernel, ComputeContext, Kernel, KernelContext, Operator, OperatorDomain, OperatorInput, OperatorOutput},
session::Session,
value::{Tensor, TensorElementType}
};
struct CustomOpOne;
impl Operator for CustomOpOne {
type Kernel<'attr> = BoxedKernel<'attr>;
const INPLACES: &[(u32, u32)] = &[(0, 0), (1, 0)];
fn name(&self) -> &str {
"CustomOpOne"
}
fn inputs(&self) -> impl IntoIterator<Item = OperatorInput> {
[OperatorInput::required(TensorElementType::Float32), OperatorInput::required(TensorElementType::Float32)]
}
fn outputs(&self) -> impl IntoIterator<Item = OperatorOutput> {
[OperatorOutput::required(TensorElementType::Float32)]
}
fn create_kernel<'attr>(&self, _: &KernelContext<'attr>) -> Result<Self::Kernel<'attr>> {
Ok(Box::new(|ctx: &ComputeContext| {
let logger = ctx.logger()?;
crate::log!(logger, Warning @ "$sentinel1$");
let x = ctx.input(0)?.ok_or_else(|| crate::Error::new("missing input"))?;
let y = ctx.input(1)?.ok_or_else(|| crate::Error::new("missing input"))?;
let (x_shape, x) = x.try_extract_tensor::<f32>()?;
let (y_shape, y) = y.try_extract_tensor::<f32>()?;
let mut z = ctx.output(0, x_shape.to_vec())?.ok_or_else(|| crate::Error::new("missing input"))?;
let (_, z_ref) = z.try_extract_tensor_mut::<f32>()?;
for i in 0..y_shape.iter().copied().reduce(|acc, e| acc * e).unwrap_or(0) as usize {
if i % 2 == 0 {
z_ref[i] = x[i];
} else {
z_ref[i] = y[i];
}
}
Ok(())
}))
}
}
struct CustomOpTwo;
impl Operator for CustomOpTwo {
type Kernel<'attr> = BoxedKernel<'attr>;
const INPLACES: &[(u32, u32)] = &[(0, 0)];
fn name(&self) -> &str {
"CustomOpTwo"
}
fn inputs(&self) -> impl IntoIterator<Item = OperatorInput> {
[OperatorInput::required(TensorElementType::Float32)]
}
fn outputs(&self) -> impl IntoIterator<Item = OperatorOutput> {
[OperatorOutput::required(TensorElementType::Int32)]
}
fn create_kernel<'attr>(&self, _: &KernelContext<'attr>) -> crate::Result<Self::Kernel<'attr>> {
Ok(Box::new(|ctx: &ComputeContext| {
let logger = ctx.logger()?;
crate::log!(logger, Verbose @ "$sentinel2$");
let x = ctx.input(0)?.ok_or_else(|| crate::Error::new("missing input"))?;
let (x_shape, x) = x.try_extract_tensor::<f32>()?;
let mut z = ctx.output(0, x_shape.to_vec())?.ok_or_else(|| crate::Error::new("missing input"))?;
let (_, z_ref) = z.try_extract_tensor_mut::<i32>()?;
for i in 0..x_shape.iter().copied().reduce(|acc, e| acc * e).unwrap_or(0) as usize {
z_ref[i] = (x[i] * i as f32) as i32;
}
Ok(())
}))
}
}
#[test]
fn test_custom_ops() -> crate::Result<()> {
let logged_values = Arc::new((AtomicBool::new(false), AtomicBool::new(false)));
let model = std::fs::read("tests/data/custom_op_test.onnx").expect("");
let mut session = Session::builder()?
.with_operators(OperatorDomain::new("test.customop")?.add(CustomOpOne)?.add(CustomOpTwo)?)?
.with_logger(Arc::new({
let logged_values = logged_values.clone();
move |_level: LogLevel, _category: &str, _id: &str, _code_location: &str, message: &str| match message {
"$sentinel1$" => logged_values.0.store(true, Ordering::Release),
"$sentinel2$" => logged_values.1.store(true, Ordering::Release),
_ => {}
}
}))?
.commit_from_memory(&model)?;
let allocator = session.allocator();
let mut value1 = Tensor::<f32>::new(allocator, [3_usize, 5])?;
{
let (_, data) = value1.extract_tensor_mut();
for datum in data {
*datum = 0.;
}
}
let mut value2 = Tensor::<f32>::new(allocator, [3_usize, 5])?;
{
let (_, data) = value2.extract_tensor_mut();
for datum in data {
*datum = 1.;
}
}
let values = session.run(crate::inputs![&value1, &value2])?;
assert_eq!(values[0].try_extract_tensor::<i32>()?.1, [0, 1, 0, 3, 0, 5, 0, 7, 0, 9, 0, 11, 0, 13, 0]);
assert!(logged_values.0.load(Ordering::Acquire));
assert!(logged_values.1.load(Ordering::Acquire));
Ok(())
}
struct AttrTesterIntFloat;
impl Operator for AttrTesterIntFloat {
type Kernel<'attr> = BoxedKernel<'attr>;
const INPLACES: &[(u32, u32)] = &[(0, 0)];
fn name(&self) -> &str {
"AttrTesterIntFloat"
}
fn inputs(&self) -> impl IntoIterator<Item = OperatorInput> {
[OperatorInput::required(TensorElementType::Float32)]
}
fn outputs(&self) -> impl IntoIterator<Item = OperatorOutput> {
[OperatorOutput::required(TensorElementType::Float32)]
}
fn infer_shape(&self, ctx: &mut super::ShapeInferenceContext) -> crate::Result<()> {
assert!(matches!(ctx.attr("a_int"), Ok(1_i64)));
assert!(matches!(ctx.attr("a_float"), Ok(2.0_f32)));
assert!(matches!(ctx.attr::<Vec<i64>>("ints").as_deref(), Ok(&[3, 4, 5])));
assert!(matches!(ctx.attr::<Vec<f32>>("floats").as_deref(), Ok(&[6., 7., 8.])));
ctx.set_output(0, &ctx.inputs()[0])?;
Ok(())
}
fn create_kernel<'attr>(&self, _: &KernelContext<'attr>) -> crate::Result<Self::Kernel<'attr>> {
Ok(Box::new(|ctx: &ComputeContext| {
let x = ctx.input(0)?.ok_or_else(|| crate::Error::new("missing input"))?;
let (x_shape, x) = x.try_extract_tensor::<f32>()?;
let mut z = ctx.output(0, x_shape.to_vec())?.ok_or_else(|| crate::Error::new("missing input"))?;
let (_, z_ref) = z.try_extract_tensor_mut::<f32>()?;
for i in 0..x_shape.iter().copied().reduce(|acc, e| acc * e).unwrap_or(0) as usize {
z_ref[i] = x[i] * 2.;
}
Ok(())
}))
}
}
struct AttrTesterString;
impl Operator for AttrTesterString {
type Kernel<'attr> = BoxedKernel<'attr>;
const INPLACES: &[(u32, u32)] = &[(0, 0)];
fn name(&self) -> &str {
"AttrTesterString"
}
fn inputs(&self) -> impl IntoIterator<Item = OperatorInput> {
[OperatorInput::required(TensorElementType::Float32)]
}
fn outputs(&self) -> impl IntoIterator<Item = OperatorOutput> {
[OperatorOutput::required(TensorElementType::Float32)]
}
fn infer_shape(&self, ctx: &mut super::ShapeInferenceContext) -> crate::Result<()> {
assert!(matches!(ctx.attr::<String>("a_string").as_deref(), Ok("iamastring")));
ctx.set_output(0, &ctx.inputs()[0])?;
Ok(())
}
fn create_kernel<'attr>(&self, ctx: &KernelContext<'attr>) -> crate::Result<Self::Kernel<'attr>> {
#[cfg(feature = "api-25")]
assert_eq!(ctx.get::<Vec<String>>("strings"), Some(vec![String::from("more"), String::from("strings")]));
Ok(Box::new(|ctx: &ComputeContext| {
let x = ctx.input(0)?.ok_or_else(|| crate::Error::new("missing input"))?;
let (x_shape, x) = x.try_extract_tensor::<f32>()?;
let mut z = ctx.output(0, x_shape.to_vec())?.ok_or_else(|| crate::Error::new("missing input"))?;
let (_, z_ref) = z.try_extract_tensor_mut::<f32>()?;
for i in 0..x_shape.iter().copied().reduce(|acc, e| acc * e).unwrap_or(0) as usize {
z_ref[i] = x[i] * 3.;
}
Ok(())
}))
}
}
#[test]
fn test_op_attrs() -> crate::Result<()> {
let mut session = Session::builder()?
.with_operators(OperatorDomain::new("test.customop")?.add(AttrTesterIntFloat)?.add(AttrTesterString)?)?
.commit_from_file("tests/data/attr_tester.onnx")?;
let value1 = Tensor::from_array(([5], vec![1.0_f32, 2.0, 3.0, 4.0, 5.0]))?;
let values = session.run(crate::inputs!["input_0" => &value1])?;
assert_eq!(values[0].try_extract_tensor::<f32>()?.1, [6.0, 12.0, 18.0, 24.0, 30.0]);
Ok(())
}
struct CopyTensorArrayAllVariadic;
impl Operator for CopyTensorArrayAllVariadic {
type Kernel<'attr> = CopyVariadic;
fn name(&self) -> &str {
"CopyTensorArrayAllVariadic"
}
fn inputs(&self) -> impl IntoIterator<Item = OperatorInput> {
[OperatorInput::variadic(1).homogenous(TensorElementType::Float32)]
}
fn outputs(&self) -> impl IntoIterator<Item = OperatorOutput> {
[OperatorOutput::variadic(1).homogenous(TensorElementType::Float32)]
}
fn infer_shape(&self, ctx: &mut super::ShapeInferenceContext) -> crate::Result<()> {
for (i, input) in ctx.inputs().into_iter().enumerate() {
ctx.set_output(i, &input)?;
}
Ok(())
}
fn create_kernel<'attr>(&self, _: &KernelContext<'attr>) -> crate::Result<Self::Kernel<'attr>> {
Ok(CopyVariadic { start: 0 })
}
}
struct CopyVariadic {
start: usize
}
impl Kernel for CopyVariadic {
fn compute(&mut self, ctx: &ComputeContext) -> crate::Result<()> {
for i in self.start..ctx.num_inputs()? {
let input = ctx.input(i)?.ok_or_else(|| crate::Error::new("missing input"))?;
let mut output = ctx
.output(i, input.shape().clone())?
.ok_or_else(|| crate::Error::new(format!("failed to allocate output {i}")))?;
output.try_extract_tensor_mut::<f32>()?.1.copy_from_slice(input.try_extract_tensor()?.1);
}
Ok(())
}
}
struct CopyTensorArrayCombined;
impl Operator for CopyTensorArrayCombined {
type Kernel<'attr> = BoxedKernel<'attr>;
fn name(&self) -> &str {
"CopyTensorArrayCombined"
}
fn inputs(&self) -> impl IntoIterator<Item = OperatorInput> {
vec![OperatorInput::optional(TensorElementType::Float32), OperatorInput::variadic(0).homogenous(TensorElementType::Float32)]
}
fn outputs(&self) -> impl IntoIterator<Item = OperatorOutput> {
vec![OperatorOutput::optional(TensorElementType::Float32), OperatorOutput::variadic(0).homogenous(TensorElementType::Float32)]
}
fn infer_shape(&self, ctx: &mut super::ShapeInferenceContext) -> crate::Result<()> {
for (i, input) in ctx.inputs().into_iter().enumerate() {
ctx.set_output(i, &input)?;
}
Ok(())
}
fn create_kernel<'attr>(&self, _info: &KernelContext<'attr>) -> crate::Result<Self::Kernel<'attr>> {
Ok(Box::new(move |ctx: &ComputeContext| {
if let Ok(Some(input)) = ctx.input(0) {
let mut output = ctx
.output(0, input.shape().clone())?
.ok_or_else(|| crate::Error::new("failed to allocate output 0"))?;
output.try_extract_tensor_mut::<f32>()?.1.copy_from_slice(input.try_extract_tensor()?.1);
}
CopyVariadic { start: 1 }.compute(ctx)
}))
}
}
#[test]
fn test_variadic_io() -> crate::Result<()> {
let ops = Arc::new(
OperatorDomain::new("test.customop")?
.add(CopyTensorArrayAllVariadic)?
.add(CopyTensorArrayCombined)?
);
let mut session = Session::builder()?
.with_operators(Arc::clone(&ops))?
.commit_from_file("tests/data/copy_2_inputs_2_outputs.onnx")?;
let input0 = Tensor::from_array(([15], vec![1.1_f32, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0, 11.1, 12.2, 13.3, 14.4, 15.5]))?;
let input1 = Tensor::from_array(([15], vec![15.5_f32, 14.4, 13.3, 12.2, 11.1, 10.0, 9.9, 8.8, 7.7, 6.6, 5.5, 4.4, 3.3, 2.2, 1.1]))?;
let values = session.run(crate::inputs!["input_0" => &input0, "input_1" => &input1])?;
assert_eq!(values[0].try_extract_tensor::<f32>()?.1, input0.extract_tensor().1);
assert_eq!(values[1].try_extract_tensor::<f32>()?.1, input1.extract_tensor().1);
let mut session = Session::builder()?
.with_operators(Arc::clone(&ops))?
.commit_from_file("tests/data/copy_3_inputs_3_outputs.onnx")?;
let input2 = Tensor::from_array(([15], vec![6.6_f32, 7.7, 8.8, 9.9, 10.0, 1.1, 2.2, 3.3, 4.4, 5.5, 11.1, 12.2, 13.3, 14.4, 15.5]))?;
let values = session.run(crate::inputs![
"input_0" => &input0,
"input_1" => &input1,
"input_2" => &input2
])?;
assert_eq!(values[0].try_extract_tensor::<f32>()?.1, input0.extract_tensor().1);
assert_eq!(values[1].try_extract_tensor::<f32>()?.1, input1.extract_tensor().1);
assert_eq!(values[2].try_extract_tensor::<f32>()?.1, input2.extract_tensor().1);
Ok(())
}