gqlite 1.5.0 → 1.7.0

data/ext/graphcore/src/value/tensor.rs

@@ -0,0 +1,739 @@
+//! Dense tensor values and (de)serialization
+
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use std::convert::TryFrom;
+
+/// Element data type of a tensor.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[repr(u8)]
+pub enum DType
+{
+  /// 32-bit IEEE float
+  F32 = 0,
+  /// 64-bit IEEE float
+  F64 = 1,
+  /// 8-bit signed integer
+  I8 = 2,
+  /// 8-bit unsigned integer
+  U8 = 3,
+  /// 16-bit signed integer
+  I16 = 4,
+  // 5 and 6 reserved for F16/BF16
+}
+
+impl DType
+{
+  /// Number of bytes per element.
+  pub fn byte_size(self) -> usize
+  {
+    match self
+    {
+      DType::F32 => 4,
+      DType::F64 => 8,
+      DType::I8 | DType::U8 => 1,
+      DType::I16 => 2,
+    }
+  }
+
+  /// Returns `true` for floating-point dtypes (F32, F64).
+  pub fn is_float(self) -> bool
+  {
+    matches!(self, DType::F32 | DType::F64)
+  }
+
+  /// Returns the canonical ASCII name used in JSON serialization.
+  pub fn as_str(self) -> &'static str
+  {
+    match self
+    {
+      DType::F32 => "F32",
+      DType::F64 => "F64",
+      DType::I8 => "I8",
+      DType::U8 => "U8",
+      DType::I16 => "I16",
+    }
+  }
+
+  /// Parses the canonical ASCII name produced by [`DType::as_str`]. Returns `None` for unknown names.
+  pub fn from_name(s: &str) -> Option<DType>
+  {
+    match s
+    {
+      "F32" => Some(DType::F32),
+      "F64" => Some(DType::F64),
+      "I8" => Some(DType::I8),
+      "U8" => Some(DType::U8),
+      "I16" => Some(DType::I16),
+      _ => None,
+    }
+  }
+}
+
+use crate::Error;
+
+/// Validate shape, dims, and data length. This is the single place the checks run.
+pub fn validate_shape(dtype: DType, shape: &[usize], data_len: usize) -> Result<(), Error>
+{
+  if shape.len() > (u8::MAX as usize)
+  {
+    return Err(Error::TensorRankTooLarge);
+  }
+  for &dim in shape.iter()
+  {
+    if dim > (u32::MAX as usize)
+    {
+      return Err(Error::TensorDimensionTooLarge);
+    }
+  }
+  let mut numel: usize = 1;
+  for &d in shape.iter()
+  {
+    numel = numel
+      .checked_mul(d)
+      .ok_or(Error::TensorShapeNumelOverflow)?;
+  }
+  let bytes = numel
+    .checked_mul(dtype.byte_size())
+    .ok_or(Error::TensorShapeNumelOverflow)?;
+  if data_len != bytes
+  {
+    return Err(Error::TensorShapeDataLengthMismatch);
+  }
+  Ok(())
+}
+
+/// Dense tensor container
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Tensor
+{
+  dtype: DType,
+  shape: Vec<usize>,
+  data: Vec<u8>,
+}
+
+impl Tensor
+{
+  // Private constructor — caller must ensure the invariant (data length matches shape & dtype)
+  fn from_raw(dtype: DType, shape: Vec<usize>, data: Vec<u8>) -> Self
+  {
+    debug_assert!(validate_shape(dtype, &shape, data.len()).is_ok());
+    Self { dtype, shape, data }
+  }
+
+  /// Construct from an f32 slice as a 1-D tensor
+  pub fn from_f32_slice(v: &[f32]) -> Self
+  {
+    let bytes = bytemuck::cast_slice(v).to_vec();
+    Self::from_raw(DType::F32, vec![v.len()], bytes)
+  }
+
+  /// Construct from an f64 slice as a 1-D tensor
+  pub fn from_f64_slice(v: &[f64]) -> Self
+  {
+    let bytes = bytemuck::cast_slice(v).to_vec();
+    Self::from_raw(DType::F64, vec![v.len()], bytes)
+  }
+
+  /// Construct from an i8 slice as a 1-D tensor
+  pub fn from_i8_slice(v: &[i8]) -> Self
+  {
+    let bytes = bytemuck::cast_slice(v).to_vec();
+    Self::from_raw(DType::I8, vec![v.len()], bytes)
+  }
+
+  /// Construct from an u8 slice as a 1-D tensor
+  pub fn from_u8_slice(v: &[u8]) -> Self
+  {
+    let bytes = bytemuck::cast_slice(v).to_vec();
+    Self::from_raw(DType::U8, vec![v.len()], bytes)
+  }
+
+  /// Construct from an i16 slice as a 1-D tensor
+  pub fn from_i16_slice(v: &[i16]) -> Self
+  {
+    let bytes = bytemuck::cast_slice(v).to_vec();
+    Self::from_raw(DType::I16, vec![v.len()], bytes)
+  }
+
+  /// Construct an f32 tensor with an explicit shape. Shape vs values length is validated.
+  pub fn from_f32_shaped(shape: Vec<usize>, values: &[f32]) -> Result<Self, crate::Error>
+  {
+    let data_len = values
+      .len()
+      .checked_mul(4)
+      .ok_or(crate::Error::TensorShapeNumelOverflow)?;
+    match validate_shape(DType::F32, &shape, data_len)
+    {
+      Ok(()) => Ok(Self::from_raw(
+        DType::F32,
+        shape,
+        bytemuck::cast_slice(values).to_vec(),
+      )),
+      Err(crate::Error::TensorShapeDataLengthMismatch) => Err(crate::Error::TensorShapeMismatch),
+      Err(e) => Err(e),
+    }
+  }
+
+  /// Element data type.
+  pub fn dtype(&self) -> DType
+  {
+    self.dtype
+  }
+
+  /// Shape of the tensor (number of elements per dimension).
+  pub fn shape(&self) -> &[usize]
+  {
+    &self.shape
+  }
+
+  /// Number of dimensions.
+  pub fn rank(&self) -> usize
+  {
+    self.shape.len()
+  }
+
+  /// Number of elements (product of shape dimensions).
+  pub fn numel(&self) -> usize
+  {
+    self.shape.iter().product()
+  }
+
+  /// Number of elements (compat with slice/Vec API).
+  pub fn len(&self) -> usize
+  {
+    self.numel()
+  }
+
+  /// Returns `true` if the tensor has no elements.
+  pub fn is_empty(&self) -> bool
+  {
+    self.numel() == 0
+  }
+
+  /// Raw byte slice over the tensor data (little-endian).
+  pub fn as_bytes(&self) -> &[u8]
+  {
+    &self.data
+  }
+
+  /// Always returns `true`; tensors are always stored contiguously.
+  pub fn is_contiguous(&self) -> bool
+  {
+    true
+  }
+
+  /// Zero-copy F32 slice. Returns `None` on big-endian platforms or non-F32 dtypes.
+  #[cfg(target_endian = "little")]
+  pub fn as_f32_slice(&self) -> Option<&[f32]>
+  {
+    if self.dtype != DType::F32
+    {
+      return None;
+    }
+    bytemuck::try_cast_slice::<u8, f32>(&self.data).ok()
+  }
+
+  /// Zero-copy F32 slice. Returns `None` on big-endian platforms or non-F32 dtypes.
+  #[cfg(not(target_endian = "little"))]
+  pub fn as_f32_slice(&self) -> Option<&[f32]>
+  {
+    None
+  }
+
+  /// Zero-copy F64 slice. Returns `None` on big-endian platforms or non-F64 dtypes.
+  #[cfg(target_endian = "little")]
+  pub fn as_f64_slice(&self) -> Option<&[f64]>
+  {
+    if self.dtype != DType::F64
+    {
+      return None;
+    }
+    bytemuck::try_cast_slice::<u8, f64>(&self.data).ok()
+  }
+
+  /// Zero-copy F64 slice. Returns `None` on big-endian platforms or non-F64 dtypes.
+  #[cfg(not(target_endian = "little"))]
+  pub fn as_f64_slice(&self) -> Option<&[f64]>
+  {
+    None
+  }
+
+  /// Produce a Vec<f32> view of the tensor, converting other dtypes as necessary.
+  pub fn to_f32_vec(&self) -> Vec<f32>
+  {
+    match self.dtype
+    {
+      DType::F32 =>
+      {
+        #[cfg(target_endian = "little")]
+        {
+          if let Ok(s) = bytemuck::try_cast_slice::<u8, f32>(&self.data)
+          {
+            return s.to_vec();
+          }
+        }
+        // fallback: parse little-endian bytes
+        self
+          .data
+          .as_chunks::<4>()
+          .0
+          .iter()
+          .map(|c| {
+            let mut arr = [0u8; 4];
+            arr.copy_from_slice(c);
+            f32::from_le_bytes(arr)
+          })
+          .collect()
+      }
+      DType::F64 => self
+        .data
+        .as_chunks::<8>()
+        .0
+        .iter()
+        .map(|c| {
+          let mut arr = [0u8; 8];
+          arr.copy_from_slice(c);
+          f64::from_le_bytes(arr) as f32
+        })
+        .collect(),
+      DType::I8 => self.data.iter().map(|b| (*b as i8) as f32).collect(),
+      DType::U8 => self.data.iter().map(|b| *b as f32).collect(),
+      DType::I16 => self
+        .data
+        .as_chunks::<2>()
+        .0
+        .iter()
+        .map(|c| {
+          let mut arr = [0u8; 2];
+          arr.copy_from_slice(c);
+          i16::from_le_bytes(arr) as f32
+        })
+        .collect(),
+    }
+  }
+
+  /// Provide an iterator over f32 elements (only supported for F32 tensors on LE targets)
+  pub fn iter(&self) -> std::slice::Iter<'_, f32>
+  {
+    #[cfg(target_endian = "little")]
+    {
+      if let Some(s) = self.as_f32_slice()
+      {
+        return s.iter();
+      }
+      panic!(
+        "Tensor::iter only supported for F32 tensors with proper alignment on little-endian targets"
+      );
+    }
+    #[cfg(not(target_endian = "little"))]
+    {
+      panic!("Tensor::iter not supported on non-little-endian targets");
+    }
+  }
+
+  /// Binary encoding:
+  /// [version:1][flags:1][dtype:1][rank:1][dims: rank * u32][data...]
+  pub fn to_bytes(&self) -> Result<Vec<u8>, crate::Error>
+  {
+    validate_shape(self.dtype, &self.shape, self.data.len())?;
+    let rank = self.shape.len();
+    let mut out = Vec::with_capacity(4 + rank * 4 + self.data.len());
+    out.push(1u8); // version
+    out.push(0u8); // flags
+    out.push(self.dtype as u8);
+    out.push(u8::try_from(rank).map_err(|_| crate::Error::TensorRankTooLarge)?);
+    for &d in self.shape.iter()
+    {
+      let d32 = u32::try_from(d).map_err(|_| crate::Error::TensorDimensionTooLarge)?;
+      out.extend_from_slice(&d32.to_le_bytes());
+    }
+    out.extend_from_slice(&self.data);
+    Ok(out)
+  }
+
+  /// Deserialise a tensor from the binary wire format produced by [`Tensor::to_bytes`].
+  pub fn from_bytes(bytes: &[u8]) -> Result<Self, crate::Error>
+  {
+    if bytes.len() < 4
+    {
+      return Err(crate::Error::TensorTooShort);
+    }
+    let version = bytes[0];
+    if version != 1
+    {
+      return Err(crate::Error::TensorUnknownVersion(version));
+    }
+    let flags = bytes[1];
+    if flags != 0
+    {
+      return Err(crate::Error::TensorUnknownFlags(flags));
+    }
+    let dtype = match bytes[2]
+    {
+      0 => DType::F32,
+      1 => DType::F64,
+      2 => DType::I8,
+      3 => DType::U8,
+      4 => DType::I16,
+      v => return Err(crate::Error::TensorUnknownDType(v)),
+    };
+    let rank = bytes[3] as usize;
+    let header_len = 4usize;
+    let shape_bytes_len = rank
+      .checked_mul(4)
+      .ok_or(crate::Error::TensorNumelOverflow)?;
+    if bytes.len() < header_len + shape_bytes_len
+    {
+      return Err(crate::Error::TensorTooShort);
+    }
+    let mut shape = Vec::with_capacity(rank);
+    let mut pos = header_len;
+    for _ in 0..rank
+    {
+      let chunk = &bytes[pos..pos + 4];
+      let d = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]) as usize;
+      shape.push(d);
+      pos += 4;
+    }
+    let data_len = bytes.len() - pos;
+    match validate_shape(dtype, &shape, data_len)
+    {
+      Ok(()) => Ok(Self::from_raw(dtype, shape, bytes[pos..].to_vec())),
+      Err(crate::Error::TensorShapeDataLengthMismatch) =>
+      {
+        Err(crate::Error::TensorDataLengthMismatch)
+      }
+      Err(_) => Err(crate::Error::TensorNumelOverflow),
+    }
+  }
+}
+
+impl Serialize for Tensor
+{
+  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+  where
+    S: Serializer,
+  {
+    if serializer.is_human_readable()
+    {
+      use serde::ser::SerializeMap;
+      let mut map = serializer.serialize_map(Some(4))?;
+      map.serialize_entry("__gql_type", "tensor")?;
+      map.serialize_entry("dtype", self.dtype.as_str())?;
+      map.serialize_entry("shape", &self.shape)?;
+
+      match self.dtype
+      {
+        DType::F32 =>
+        {
+          let data: Vec<f64> = self
+            .data
+            .as_chunks::<4>()
+            .0
+            .iter()
+            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]) as f64)
+            .collect();
+          map.serialize_entry("data", &data)?;
+        }
+        DType::F64 =>
+        {
+          let data: Vec<f64> = self
+            .data
+            .as_chunks::<8>()
+            .0
+            .iter()
+            .map(|b| f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]))
+            .collect();
+          map.serialize_entry("data", &data)?;
+        }
+        DType::I8 =>
+        {
+          let data: Vec<i64> = self.data.iter().map(|&b| b as i8 as i64).collect();
+          map.serialize_entry("data", &data)?;
+        }
+        DType::U8 =>
+        {
+          let data: Vec<i64> = self.data.iter().map(|&b| b as i64).collect();
+          map.serialize_entry("data", &data)?;
+        }
+        DType::I16 =>
+        {
+          let data: Vec<i64> = self
+            .data
+            .as_chunks::<2>()
+            .0
+            .iter()
+            .map(|b| i16::from_le_bytes([b[0], b[1]]) as i64)
+            .collect();
+          map.serialize_entry("data", &data)?;
+        }
+      }
+
+      map.end()
+    }
+    else
+    {
+      let bs = self.to_bytes().map_err(serde::ser::Error::custom)?;
+      serializer.serialize_bytes(&bs)
+    }
+  }
+}
+
+struct TensorJsonVisitor;
+
+impl<'de> serde::de::Visitor<'de> for TensorJsonVisitor
+{
+  type Value = Tensor;
+
+  fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result
+  {
+    write!(f, "a tensor JSON object with \"__gql_type\": \"tensor\"")
+  }
+
+  fn visit_map<A: serde::de::MapAccess<'de>>(self, mut map: A) -> Result<Tensor, A::Error>
+  {
+    let mut gql_type: Option<String> = None;
+    let mut dtype_str: Option<String> = None;
+    let mut shape: Option<Vec<usize>> = None;
+    let mut data: Option<Vec<f64>> = None;
+
+    while let Some(key) = map.next_key::<String>()?
+    {
+      match key.as_str()
+      {
+        "__gql_type" => gql_type = Some(map.next_value()?),
+        "dtype" => dtype_str = Some(map.next_value()?),
+        "shape" => shape = Some(map.next_value()?),
+        "data" => data = Some(map.next_value()?),
+        _ =>
+        {
+          map.next_value::<serde::de::IgnoredAny>()?;
+        }
+      }
+    }
+
+    match gql_type.as_deref()
+    {
+      Some("tensor") =>
+      {}
+      Some(t) =>
+      {
+        return Err(serde::de::Error::custom(format!(
+          "__gql_type is {}, expected \"tensor\"",
+          t
+        )));
+      }
+      None => return Err(serde::de::Error::missing_field("__gql_type")),
+    }
+
+    let dtype_str = dtype_str.ok_or_else(|| serde::de::Error::missing_field("dtype"))?;
+    let dtype = DType::from_name(&dtype_str)
+      .ok_or_else(|| serde::de::Error::custom(format!("unknown dtype {}", dtype_str)))?;
+    let shape = shape.ok_or_else(|| serde::de::Error::missing_field("shape"))?;
+    let data = data.ok_or_else(|| serde::de::Error::missing_field("data"))?;
+
+    let numel: usize = shape.iter().product();
+    if data.len() != numel
+    {
+      return Err(serde::de::Error::custom(format!(
+        "data length {} does not match shape product {}",
+        data.len(),
+        numel
+      )));
+    }
+
+    let mut bytes: Vec<u8> = Vec::with_capacity(numel * dtype.byte_size());
+    match dtype
+    {
+      DType::F32 =>
+      {
+        for &v in &data
+        {
+          bytes.extend_from_slice(&(v as f32).to_le_bytes());
+        }
+      }
+      DType::F64 =>
+      {
+        for &v in &data
+        {
+          bytes.extend_from_slice(&v.to_le_bytes());
+        }
+      }
+      DType::I8 =>
+      {
+        for &v in &data
+        {
+          bytes.push((v as i8) as u8);
+        }
+      }
+      DType::U8 =>
+      {
+        for &v in &data
+        {
+          bytes.push(v as u8);
+        }
+      }
+      DType::I16 =>
+      {
+        for &v in &data
+        {
+          bytes.extend_from_slice(&(v as i16).to_le_bytes());
+        }
+      }
+    }
+
+    validate_shape(dtype, &shape, bytes.len()).map_err(serde::de::Error::custom)?;
+    Ok(Tensor::from_raw(dtype, shape, bytes))
+  }
+}
+
+impl<'de> Deserialize<'de> for Tensor
+{
+  fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+  where
+    D: Deserializer<'de>,
+  {
+    if deserializer.is_human_readable()
+    {
+      deserializer.deserialize_map(TensorJsonVisitor)
+    }
+    else
+    {
+      struct BytesVisitor;
+      impl<'de> serde::de::Visitor<'de> for BytesVisitor
+      {
+        type Value = Tensor;
+        fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result
+        {
+          write!(f, "a byte string")
+        }
+        fn visit_bytes<E: serde::de::Error>(self, v: &[u8]) -> Result<Tensor, E>
+        {
+          Tensor::from_bytes(v).map_err(E::custom)
+        }
+        fn visit_byte_buf<E: serde::de::Error>(self, v: Vec<u8>) -> Result<Tensor, E>
+        {
+          Tensor::from_bytes(&v).map_err(E::custom)
+        }
+      }
+      deserializer.deserialize_bytes(BytesVisitor)
+    }
+  }
+}
+
+impl IntoIterator for Tensor
+{
+  type Item = f32;
+  type IntoIter = std::vec::IntoIter<f32>;
+  fn into_iter(self) -> Self::IntoIter
+  {
+    self.to_f32_vec().into_iter()
+  }
+}
+
+impl std::ops::Deref for Tensor
+{
+  type Target = [f32];
+  fn deref(&self) -> &Self::Target
+  {
+    #[cfg(target_endian = "little")]
+    {
+      if self.dtype == DType::F32
+      {
+        bytemuck::try_cast_slice::<u8, f32>(&self.data)
+          .expect("Tensor data is not a properly aligned f32 slice")
+      }
+      else
+      {
+        panic!("Tensor deref to [f32] only supported for DType::F32");
+      }
+    }
+    #[cfg(not(target_endian = "little"))]
+    {
+      panic!("Tensor deref to [f32] not supported on non-little-endian targets");
+    }
+  }
+}
+
+impl<'a> IntoIterator for &'a Tensor
+{
+  type Item = &'a f32;
+  type IntoIter = std::slice::Iter<'a, f32>;
+  fn into_iter(self) -> Self::IntoIter
+  {
+    self.iter()
+  }
+}
+
+impl FromIterator<f32> for Tensor
+{
+  fn from_iter<I: IntoIterator<Item = f32>>(iter: I) -> Self
+  {
+    let v: Vec<f32> = iter.into_iter().collect();
+    Tensor::from_f32_slice(&v)
+  }
+}
+
+impl From<Vec<f32>> for Tensor
+{
+  fn from(v: Vec<f32>) -> Self
+  {
+    Tensor::from_f32_slice(&v)
+  }
+}
+
+#[cfg(test)]
+mod tests
+{
+  use crate::Value;
+
+  #[test]
+  fn test_tensor_from_vec_roundtrip()
+  {
+    let v: Vec<f32> = vec![1.0f32, 2.5f32, -3.25f32];
+    let val: Value = v.clone().into();
+    let out: Vec<f32> = Vec::try_from(val).expect("Should convert back to Vec<f32>");
+    assert_eq!(out, v);
+  }
+
+  #[test]
+  fn test_optimize_array_to_tensor()
+  {
+    let arr = Value::Array(vec![
+      Value::Float(1.0),
+      Value::Integer(2),
+      Value::Float(3.5),
+    ]);
+    let opt = arr.optimize();
+    let out: Vec<f32> =
+      Vec::try_from(opt).expect("optimize should produce a Tensor convertible to Vec<f32>");
+    assert_eq!(out, vec![1.0f32, 2.0f32, 3.5f32]);
+  }
+
+  #[test]
+  fn test_tensor_add_tensor_and_scalar()
+  {
+    let a: Value = vec![1.0f32, 2.0f32, 3.0f32].into();
+    let b: Value = vec![4.0f32, 5.0f32, 6.0f32].into();
+    let sum = (a + b).expect("tensor + tensor should succeed");
+    let out: Vec<f32> = Vec::try_from(sum).expect("result should be a tensor");
+    assert_eq!(out, vec![5.0f32, 7.0f32, 9.0f32]);
+
+    // tensor + float scalar
+    let a: Value = vec![1.0f32, 2.0f32].into();
+    let scalar = Value::Float(3.0);
+    let sum = (a + scalar).expect("tensor + float should succeed");
+    let out: Vec<f32> = Vec::try_from(sum).expect("result should be a tensor");
+    assert_eq!(out, vec![4.0f32, 5.0f32]);
+  }
+
+  #[test]
+  fn test_tensor_pow_elementwise()
+  {
+    let a: Value = vec![2.0f32, 3.0f32].into();
+    let b: Value = vec![3.0f32, 2.0f32].into();
+    let res = a.pow(b).expect("tensor pow tensor should succeed");
+    let out: Vec<f32> = Vec::try_from(res).expect("result should be a tensor");
+    assert_eq!(out, vec![8.0f32, 9.0f32]);
+  }
+}