Tensor 与 Kind

#[derive(new, Clone, Debug)]
pub struct Tensor<const D: usize, K = Float>
where
    K: Basic,
{
    pub(crate) primitive: BridgeTensor,
    _kind: PhantomData<K>,
}

println!("Tensor<2>: size {} align {}", size_of::<Tensor<2>>(), align_of::<Tensor<2>>());
println!("Tensor<4, Int>: size {} align {}", size_of::<Tensor<4, Int>>(), align_of::<Tensor<4, Int>>());
println!("Device: size {} align {}", size_of::<Device>(), align_of::<Device>());

== sizes
Tensor<2>: size 384 align 64
Tensor<4, Int>: size 384 align 64
Device: size 128 align 64

/// A type-level representation of the kind of a tensor.
/// Metadata access is lazy.
///
/// # Notes
/// This trait is intentionally sealed to keep the set of tensor kinds closed.
///
/// Although exposed publicly, tensor kinds are not meant to be extensible:
/// the backend dispatch system, `DType`, and all tensor ops assume a fixed,
/// closed set of tensor kinds (e.g. Float, Int, Bool), each mapping directly to a
/// corresponding backend implementation.
pub trait TensorKind: sealed::Sealed + Clone + Send + Sync + core::fmt::Debug {
    /// The tensor kind identifier.
    const KIND: Kind;

    /// The name of the tensor kind.
    fn name() -> &'static str {
        Self::KIND.as_str()
    }
}

/// The base trait for any tensor kind.
pub trait Basic: crate::ops::BasicOps {}
impl<K: crate::ops::BasicOps> Basic for K {}

/// Kinds that support numeric operations.
pub trait Numeric: Basic + crate::ops::Numeric {}
impl<K: Basic + crate::ops::Numeric> Numeric for K {}

/// Kinds that support ordered operations.
pub trait Ordered: Numeric + crate::ops::Ordered {}
impl<K: Numeric + crate::ops::Ordered> Ordered for K {}

/// Kinds that support float math operations.
pub trait FloatMath: Numeric + crate::ops::FloatMathOps {}
impl<K: Numeric + crate::ops::FloatMathOps> FloatMath for K {}
// ...

use burn::tensor::{Bool, Device, Tensor};

fn main() {
  let device = Device::flex();
  let a = Tensor::<2, Bool>::empty([2, 2], &device);
  let _ = a.clone().matmul(a);
  let _s: Tensor<2, String>;
}

error[E0599]: the method `matmul` exists for struct `burn::Tensor<2, burn::tensor::Bool>`, but its trait bounds were not satisfied
  --> src/bin/kinderr.rs:6:21
   |
 6 |   let _ = a.clone().matmul(a);
   |                     ^^^^^^ method cannot be called on `burn::Tensor<2, burn::tensor::Bool>` due to unsatisfied trait bounds
   |
  ::: ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/bridge/kind.rs:16:1
   |
16 | pub struct Bool;
   | --------------- doesn't satisfy `_: Numeric` or `burn::tensor::Bool: burn::tensor::kind::Numeric`
   |
   = note: the following trait bounds were not satisfied:
           `burn::tensor::Bool: burn_tensor::bridge::ops::numeric::Numeric`
           which is required by `burn::tensor::Bool: burn::tensor::kind::Numeric`

error[E0277]: the trait bound `std::string::String: burn_tensor::bridge::ops::base::BasicOps` is not satisfied
  --> src/bin/kinderr.rs:7:11
   |
 7 |   let _s: Tensor<2, String>;
   |           ^^^^^^^^^^^^^^^^^ the trait `burn_tensor::bridge::ops::base::BasicOps` is not implemented for `std::string::String`
...
   = note: required for `std::string::String` to implement `burn::tensor::kind::Basic`
note: required by a bound in `burn::Tensor`
...

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(try_from = "TensorDataDe")]
pub struct TensorData {
    /// The values of the tensor (as bytes).
    pub(in crate::data) bytes: Bytes,

    /// The shape of the tensor.
    #[serde(with = "shape_inner")]
    pub(in crate::data) shape: Shape,

    /// The data type of the tensor.
    pub(in crate::data) dtype: DType,
}

let f = Tensor::<2>::zeros([2, 3], &device);       // device = Device::flex()
let i = Tensor::<2, Int>::zeros([2, 3], &device);
let b = Tensor::<2, Bool>::empty([2, 3], &device);
// ... 同样三个张量建在 Device::metal(DeviceKind::DefaultDevice) 上

== kind / dtype / shape
float dtype F32 shape Shape { dims: [2, 3] }
int   dtype I32
bool  dtype Bool(Native)
metal float F32 int I32 bool Bool(U8)

let a = Tensor::<2>::ones([2, 2], &device);
let b = a.matmul(a.clone());

error[E0382]: borrow of moved value: `a`
 --> src/bin/moveerr.rs:6:20
  |
5 |   let a = Tensor::<2>::ones([2, 2], &device);
  |       - move occurs because `a` has type `burn::Tensor<2>`, which does not implement the `Copy` trait
6 |   let b = a.matmul(a.clone());
  |           -        ^ value borrowed here after move
  |           |
  |           value moved here
  |
help: consider cloning the value if the performance cost is acceptable
  |
6 |   let b = a.clone().matmul(a.clone());
  |            ++++++++
...

    // In-place SIMD fast path: lhs unique contiguous at offset 0, rhs
    // contiguous (no broadcast).
    if let Some(simd_op) = simd_hint
        && lhs.is_unique()
        && let (Some((0, l_end)), Some((r_start, r_end))) = (
            lhs.layout().contiguous_offsets(),
            rhs.layout().contiguous_offsets(),
        )
    {
        let r_slice: &[f32] = &rhs.storage()[r_start..r_end];
        let lhs_storage: &mut [f32] = lhs.storage_mut();
        let l_slice = &mut lhs_storage[..l_end];

        match simd_op {
            BinaryOp::Add => simd::add_inplace_f32(l_slice, r_slice),
            // ...
        }
        return lhs;
    }

for (name, device) in [("flex", Device::flex()), ("metal", metal.clone())] {
  let a = Tensor::<2>::ones([4, 4], &device);
  let before = a.can_mut();
  let a2 = a.clone();
  let shared = a.can_mut();
  drop(a2);
  let after = a.can_mut();
  println!("{name}: unique {before}, after clone {shared}, after drop of clone {after}");
}

let n = 4096;
let x = Tensor::<2>::random([n, n], Distribution::Default, &device); // Flex
let y = x.clone(); // 计时、计分配字节

let rhs = Tensor::<2>::ones([n, n], &device);
for round in 0..3 {
  let lhs = x.clone() + 0.0; // 新建一个唯一持有的缓冲
  let out = lhs + rhs.clone(); // 计时、计分配字节
  drop(out);
  let lhs = x.clone() + 0.0;
  let keep = lhs.clone(); // 让 lhs 不再唯一
  let out = lhs + rhs.clone(); // 计时、计分配字节
  drop((out, keep));
}

== clone shares the buffer (can_mut)
flex: unique true, after clone false, after drop of clone true
metal: unique true, after clone false, after drop of clone true
== clone vs copy cost on flex, 4096x4096 f32 (64 MiB)
clone: 584ns, allocated 16 bytes
== by-value add on flex: unique lhs vs shared lhs
round 0 unique lhs: 2.864292ms, allocated 0.091 MiB
round 0 shared lhs: 6.228667ms, allocated 64.000 MiB
round 1 unique lhs: 1.250541ms, allocated 0.000 MiB
round 1 shared lhs: 2.676833ms, allocated 64.000 MiB
round 2 unique lhs: 1.253792ms, allocated 0.000 MiB
round 2 shared lhs: 2.333834ms, allocated 64.000 MiB

round 0 unique lhs: 1.444416ms, allocated 0.091 MiB
round 0 shared lhs: 1.35475ms, allocated 0.000 MiB
round 1 unique lhs: 1.212167ms, allocated 0.000 MiB
round 1 shared lhs: 1.27275ms, allocated 0.000 MiB
round 2 unique lhs: 1.326834ms, allocated 0.000 MiB
round 2 shared lhs: 1.203291ms, allocated 0.000 MiB