ort 2.0.0-rc.15

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,
	test_util::test_env,
	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(test_env())?
		.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(test_env())?
		.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 env = test_env().clone();

	let mut session = Session::builder(&env)?
		.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(&env)?
		.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(())
}

/// Stands in for `CustomOpTwo`, recording the shape of the preallocated output it sees on each run.
#[cfg(feature = "api-30")]
struct PreallocatedOutputOp(Arc<crate::util::Mutex<alloc::vec::Vec<Option<alloc::vec::Vec<i64>>>>>);

#[cfg(feature = "api-30")]
impl Operator for PreallocatedOutputOp {
	type Kernel<'attr> = BoxedKernel<'attr>;

	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>> {
		let seen = Arc::clone(&self.0);
		Ok(Box::new(move |ctx: &ComputeContext| {
			let x = ctx.input(0)?.ok_or_else(|| crate::Error::new("missing input"))?;
			let (x_shape, _) = x.try_extract_tensor::<f32>()?;

			let preallocated = ctx.preallocated_output(0)?;
			seen.lock().push(preallocated.as_ref().map(|v| v.shape().to_vec()));

			let mut z = match preallocated {
				Some(z) => {
					assert_eq!(z.shape(), x_shape);
					z
				}
				None => ctx.output(0, x_shape.to_vec())?.ok_or_else(|| crate::Error::new("missing output"))?
			};
			let (_, z_ref) = z.try_extract_tensor_mut::<i32>()?;
			for (i, z) in z_ref.iter_mut().enumerate() {
				*z = i as i32 * 10;
			}
			Ok(())
		}))
	}
}

#[test]
#[cfg(feature = "api-30")]
fn test_preallocated_output() -> crate::Result<()> {
	use crate::{
		memory::Allocator,
		session::{OutputSelector, RunOptions}
	};

	let seen = Arc::new(crate::util::Mutex::new(alloc::vec::Vec::new()));
	let mut session = Session::builder(test_env())?
		.with_operators(
			OperatorDomain::new("test.customop")?
				.add(CustomOpOne)?
				.add(PreallocatedOutputOp(Arc::clone(&seen)))?
		)?
		.commit_from_file("tests/data/custom_op_test.onnx")?;

	let x = Tensor::from_array(([3_usize, 5], alloc::vec![0.0_f32; 15]))?;
	let y = Tensor::from_array(([3_usize, 5], alloc::vec![1.0_f32; 15]))?;
	let expected: alloc::vec::Vec<i32> = (0..15).map(|i| i * 10).collect();

	// without a preallocated output, the kernel has to allocate one
	let outputs = session.run(crate::inputs![&x, &y])?;
	assert_eq!(outputs[0].try_extract_tensor::<i32>()?.1, expected);
	drop(outputs);

	// with a preallocated output, the kernel writes straight into it
	let output_name = session.outputs()[0].name().to_string();
	let preallocated = Tensor::<i32>::new(&Allocator::default(), [3_usize, 5])?;
	let preallocated_ptr = preallocated.data_ptr();
	let options = RunOptions::new()?.with_outputs(OutputSelector::default().preallocate(output_name, preallocated));
	let outputs = session.run_with_options(crate::inputs![&x, &y], &options)?;
	assert_eq!(outputs[0].try_extract_tensor::<i32>()?.1, expected);
	assert_eq!(outputs[0].view().downcast::<crate::value::DynTensorValueType>()?.data_ptr(), preallocated_ptr);

	assert_eq!(*seen.lock(), [None, Some(alloc::vec![3, 5])]);
	Ok(())
}