ariadnetor_tensor/
reorder.rs

1//! Memory-order conversion for `DenseTensorData<T>`.
2//!
3//! Provides [`reorder_data`] for explicit conversion and
4//! [`normalize_to_data`] for lazy `Cow`-borrowed normalization at the
5//! entry of operations that require a specific input layout.
6
7use crate::DenseTensorData;
8use ariadnetor_core::backend::MemoryOrder;
9use std::borrow::Cow;
10
11/// Reorder a `DenseTensorData<T>` to the requested memory layout.
12///
13/// If `tensor.order() == to`, returns a clone (zero-copy via Arc).
14/// Otherwise produces a new `DenseTensorData` whose layout `order()`
15/// matches `to`. Primary callers are the linalg kernels at their
16/// `&DenseTensorData<T>` entry points.
17pub fn reorder_data<T: Clone>(tensor: &DenseTensorData<T>, to: MemoryOrder) -> DenseTensorData<T> {
18    let from = tensor.order();
19    let shape = tensor.shape();
20    let rank = shape.len();
21    let total = tensor.len();
22    if from == to {
23        return tensor.clone();
24    }
25    if total == 0 {
26        return DenseTensorData::from_raw_parts(Vec::new(), shape.to_vec(), to);
27    }
28    let raw = tensor.storage().data();
29    let mut new_data = Vec::with_capacity(total);
30    let mut coords = vec![0usize; rank];
31
32    let axis_order: Vec<usize> = match to {
33        MemoryOrder::RowMajor => (0..rank).collect(),
34        MemoryOrder::ColumnMajor => (0..rank).rev().collect(),
35    };
36
37    for _ in 0..total {
38        let src_idx = flat_index(&coords, shape, from);
39        new_data.push(raw[src_idx].clone());
40        for &d in axis_order.iter().rev() {
41            coords[d] += 1;
42            if coords[d] < shape[d] {
43                break;
44            }
45            coords[d] = 0;
46        }
47    }
48    DenseTensorData::from_raw_parts(new_data, shape.to_vec(), to)
49}
50
51/// Normalize a `DenseTensorData<T>`'s memory order to `target`, with
52/// `Cow::Borrowed` when the tensor is already in the target order.
53///
54/// Use at the entry of any operation that requires a specific input
55/// layout (typically backend kernels expecting
56/// `backend.preferred_order()`). The returned `Cow` is `Borrowed`
57/// when no conversion is needed and `Owned` when a reorder was
58/// performed.
59pub fn normalize_to_data<T: Clone>(
60    tensor: &DenseTensorData<T>,
61    target: MemoryOrder,
62) -> Cow<'_, DenseTensorData<T>> {
63    if tensor.order() == target {
64        Cow::Borrowed(tensor)
65    } else {
66        Cow::Owned(reorder_data(tensor, target))
67    }
68}
69
70/// Compute flat index for given coordinates in the specified memory order.
71pub fn flat_index(coords: &[usize], shape: &[usize], order: MemoryOrder) -> usize {
72    let mut idx = 0;
73    let mut stride = 1;
74    match order {
75        MemoryOrder::RowMajor => {
76            for i in (0..shape.len()).rev() {
77                idx += coords[i] * stride;
78                stride *= shape[i];
79            }
80        }
81        MemoryOrder::ColumnMajor => {
82            for i in 0..shape.len() {
83                idx += coords[i] * stride;
84                stride *= shape[i];
85            }
86        }
87    }
88    idx
89}
ariadnetor_tensor/reorder.rs

ariadnetor_tensor/
reorder.rs
