写一个迷你 einsum

[dependencies]

# 只有对照测试 tests/tenferro_compare.rs 用 tenferro 和 omeco；库本身没有依赖
[dev-dependencies]
tenferro-einsum = "=0.8.0"
tenferro-cpu = "=0.8.0"
tenferro-tensor = "=0.8.0"
omeco = "=0.2.6"

/// 解析后的下标：每个输入一串标签，加输出标签。
#[derive(Debug, Clone, PartialEq)]
pub struct Equation {
  pub inputs: Vec<Vec<char>>,
  pub output: Vec<char>,
}

/// 解析 `"ij,jk,kl->il"`。必须有且只有一个 `->`；标签只能是 ASCII 字母；
/// 不支持省略号，也不支持同一个输入里重复的标签（对角线）。
pub fn parse(notation: &str) -> Result<Equation> {
  let Some((lhs, rhs)) = notation.split_once("->") else {
    return err(format!("missing '->' in {notation:?}"));
  };
  if rhs.contains("->") {
    return err(format!("more than one '->' in {notation:?}"));
  }
  let inputs = lhs.split(',').map(labels).collect::<Result<Vec<_>>>()?;
  let output = labels(rhs)?;
  // ... 同一输入里重复的标签、输出里重复的标签、输出里有而输入里没有的标签，都报错
  Ok(Equation { inputs, output })
}

  /// 每个输入在收缩之前先求和掉的标签：只在这一个输入里出现、输出里也没有。
  /// 返回每个输入剩下的标签。
  pub(crate) fn reduced_inputs(&self) -> Vec<Vec<char>> {
    let count = |label: &char| self.inputs.iter().filter(|op| op.contains(label)).count();
    self
      .inputs
      .iter()
      .map(|op| op.iter().copied().filter(|l| self.output.contains(l) || count(l) > 1).collect())
      .collect()
  }

  fn sum_except(self, keep: &[char]) -> Labeled {
    let (kept, dropped): (Vec<char>, Vec<char>) = self.labels.iter().partition(|l| keep.contains(l));
    if dropped.is_empty() {
      return self;
    }
    let rows = self.product(&kept);
    let shape = kept.iter().map(|&l| self.size(l)).collect();
    let moved = self.arrange(&[&kept[..], &dropped].concat());
    let mut data = vec![0.0; rows];
    if rows > 0 {
      for column in moved.data.chunks(rows) {
        data.iter_mut().zip(column).for_each(|(acc, x)| *acc += x);
      }
    }
    Labeled { labels: kept, tensor: Tensor { shape, data } }
  }

/// 两个操作数一起出现的标签分成四类，决定置换之后的布局：
/// 左边 `[left, contracted, batch]`，右边 `[contracted, right, batch]`。
pub(crate) struct Split {
  pub left: Vec<char>,
  pub contracted: Vec<char>,
  pub right: Vec<char>,
  pub batch: Vec<char>,
}

impl Split {
  pub(crate) fn new(a: &[char], b: &[char], keep: &[char]) -> Split {
    let (shared, left): (Vec<char>, Vec<char>) = a.iter().partition(|l| b.contains(l));
    let (batch, contracted) = shared.into_iter().partition(|l| keep.contains(l));
    let right = b.iter().copied().filter(|l| !a.contains(l)).collect();
    Split { left, contracted, right, batch }
  }

  pub(crate) fn result_labels(&self) -> Vec<char> {
    [&self.left[..], &self.right, &self.batch].concat()
  }
}

fn contract_pair(a: &Labeled, b: &Labeled, labels: &[char], split: &Split) -> Labeled {
  let m = a.product(&split.left);
  let k = a.product(&split.contracted);
  let n = b.product(&split.right);
  let batch = a.product(&split.batch);
  let lhs = a.arrange(&[&split.left[..], &split.contracted, &split.batch].concat());
  let rhs = b.arrange(&[&split.contracted[..], &split.right, &split.batch].concat());
  let mut out = vec![0.0; m * n * batch];
  if m * n * k > 0 {
    let blocks = out.chunks_mut(m * n).zip(lhs.data.chunks(m * k)).zip(rhs.data.chunks(k * n));
    for ((c, a), b) in blocks {
      matmul(a, b, c, m, k);
    }
  }
  // ... 由 labels 得到结果形状
}

/// 列主序 `c[m x n] += a[m x k] * b[k x n]`。最内层沿 `a` 和 `c` 的一列走，都是连续内存。
fn matmul(a: &[f64], b: &[f64], c: &mut [f64], m: usize, k: usize) {
  for (c_col, b_col) in c.chunks_mut(m).zip(b.chunks(k)) {
    for (a_col, &b_pj) in a.chunks(m).zip(b_col) {
      c_col.iter_mut().zip(a_col).for_each(|(c, &a)| *c += a * b_pj);
    }
  }
}

    /// Target subscript order for A: [lo..., sum..., batch...].
    pub(crate) target_a: Vec<u32>,
    /// Target subscript order for B: [sum..., ro..., batch...].
    pub(crate) target_b: Vec<u32>,
    /// Shape of GEMM output: [m, n, batch...].
    pub(crate) c_gemm_shape: Vec<usize>,

  /// 新的第 `a` 轴是原来的第 `perm[a]` 轴。按输出顺序走，
  /// 用一个里程表式的多重下标在输入里跳。
  fn permute(&self, perm: &[usize]) -> Tensor {
    let strides = col_major_strides(&self.shape);
    let shape: Vec<usize> = perm.iter().map(|&p| self.shape[p]).collect();
    let steps: Vec<usize> = perm.iter().map(|&p| strides[p]).collect();
    let len = self.data.len();
    let mut data = Vec::with_capacity(len);
    let mut index = vec![0; shape.len()];
    let mut offset = 0;
    for _ in 0..len {
      data.push(self.data[offset]);
      for axis in 0..shape.len() {
        index[axis] += 1;
        offset += steps[axis];
        if index[axis] < shape[axis] {
          break;
        }
        offset -= steps[axis] * shape[axis];
        index[axis] = 0;
      }
    }
    Tensor { shape, data }
  }

/// 贪心：每一步在所有还活着的操作数对里挑乘加次数最少的一对，相同时取先遇到的。
pub fn greedy_plan(eq: &Equation, sizes: &HashMap<char, usize>) -> Result<Plan> {
  let mut builder = Builder::new(eq);
  while builder.live.len() > 1 {
    let mut best: Option<(u128, usize, usize)> = None;
    for (i, &a) in builder.live.iter().enumerate() {
      for &b in &builder.live[i + 1..] {
        let cost = step_cost(&builder.ops[a], &builder.ops[b], sizes)?;
        if best.is_none_or(|(c, _, _)| cost < c) {
          best = Some((cost, a, b));
        }
      }
    }
    let Some((_, a, b)) = best else { break };
    builder.push(a, b, sizes)?;
  }
  Ok(builder.finish())
}

running 4 tests
test any_order_gives_same_result ... [(0, 1), (5, 2), (6, 3), (7, 4)] cost  189 diff 1.2e-15
[(3, 4), (2, 5), (1, 6), (0, 7)] cost  198 diff 1.5e-15
[(0, 4), (1, 2), (5, 6), (3, 7)] cost  258 diff 1.6e-15
ok
test errors ... ij,jk        Error("missing '->' in \"ij,jk\"")
ii->i        Error("label 'i' repeats in input 0 (diagonal is not supported)")
...ij->ij    Error("ellipsis '...' is not supported")
ij->ik       Error("output label 'k' appears in no input")
ij,jk->ik    Error("label 'j': size 2 vs 3 (input 1)")
ij,jk->ik    Error("2 inputs in equation, 1 tensors given")
ok
test matches_naive_reference ... ij,jk->ik              shape [3, 4]     plan [(0, 1)] max |mini - naive| = 0.0e0
ij,jk,kl->il           shape [3, 3]     plan [(0, 1), (2, 3)] max |mini - naive| = 1.1e-16
bij,bjk->bik           shape [3, 3, 4]  plan [(0, 1)] max |mini - naive| = 0.0e0
ij,ij->ij              shape [3, 2]     plan [(0, 1)] max |mini - naive| = 0.0e0
i,j->ij                shape [3, 2]     plan [(0, 1)] max |mini - naive| = 0.0e0
i,i->                  shape []         plan [(0, 1)] max |mini - naive| = 0.0e0
ij->ji                 shape [2, 3]     plan [] max |mini - naive| = 0.0e0
ijk->kj                shape [4, 2]     plan [] max |mini - naive| = 0.0e0
ij,jk->k               shape [4]        plan [(0, 1)] max |mini - naive| = 1.1e-16
ab,bc,cd,da->          shape []         plan [(0, 1), (2, 3), (4, 5)] max |mini - naive| = 8.9e-16
abc,cd,dbe->ae         shape [2, 2]     plan [(0, 1), (2, 3)] max |mini - naive| = 8.9e-16
ijk,jl,klm,mn,in->     shape []         plan [(3, 4), (1, 5), (0, 2), (6, 7)] max |mini - naive| = 1.0e-15
pqr,rs,st,tq,pu->us    shape [5, 4]     plan [(2, 3), (1, 5), (0, 6), (4, 7)] max |mini - naive| = 1.4e-16
ok
test zero_length_dimension ... [2,0] x [0,3] -> [2, 3] [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
ok

test result: ok. 4 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.02s

(3, 4) -> ['m', 'i'] cost 30
(1, 5) -> ['j', 'l', 'm', 'i'] cost 36
(0, 2) -> ['i', 'j', 'l', 'm'] cost 144
(6, 7) -> [] cost 36
total 246

fn tenferro_einsum(notation: &str, inputs: &[&Tensor]) -> Vec<f64> {
  let typed: Vec<TypedTensor<f64>> = inputs
    .iter()
    .map(|t| TypedTensor::<f64>::from_vec_col_major(t.shape().to_vec(), t.data().to_vec()).unwrap())
    .collect();
  let refs: Vec<&TypedTensor<f64>> = typed.iter().collect();
  let mut engine = CpuEngine::serial();
  let out = refs.as_slice().einsum(notation, &mut engine).unwrap();
  out.as_slice().unwrap().to_vec()
}

test values_match_tenferro ... ij,jk->ik              max |mini - tenferro| = 0.0e0
bij,bjk->bik           max |mini - tenferro| = 0.0e0
ij,jk->k               max |mini - tenferro| = 2.8e-17
ijk->kj                max |mini - tenferro| = 0.0e0
ab,bc,cd,da->          max |mini - tenferro| = 0.0e0
ijk,jl,klm,mn,in->     max |mini - tenferro| = 2.2e-16
pqr,rs,st,tq,pu->us    max |mini - tenferro| = 4.4e-16
ok

impl Default for ContractionOptimizerOptions {
    fn default() -> Self {
        Self {
            betas: Vec::new(),
            ntrials: 1,
            niters: 0,
            score: ScoreFunction::default(),
        }
    }
}

/// Each step contracts the connected pair (sharing at least one label) whose
/// result is smallest, the size counting every label still needed after the
/// pair: by another operand or by the output. Ties go to the smallest
/// `(left, right)` operand index pair, where an intermediate takes the next
/// index after all existing operands. When no remaining operands share a
/// label, the two smallest indices form an outer product.

running 2 tests
test plan_cost_vs_tenferro ... mps6: pa,aqb,br,pc,cqd,dr->
  mini greedy            137216  [(0, 3), (2, 5), (1, 6), (4, 7), (8, 9)]
  tenferro default         8712  [(0, 1), (2, 6), (3, 4), (5, 8), (7, 9)]
  tenferro TreeSA          8712  [(3, 4), (5, 6), (1, 2), (8, 0), (7, 9)]
  omeco exhaustive         8712  [(0, 1), (6, 2), (3, 4), (8, 5), (7, 9)]  (omeco flops 8712)
mps8: pa,aqb,brc,cs,pw,wqx,xry,ys->
  mini greedy            267264  [(0, 4), (3, 7), (1, 8), (2, 9), (5, 10), (6, 11), (12, 13)]
  tenferro default        17424  [(0, 1), (2, 3), (8, 9), (4, 5), (6, 7), (11, 12), (10, 13)]
  tenferro TreeSA         17424  [(5, 4), (7, 6), (8, 9), (3, 2), (0, 1), (11, 12), (10, 13)]
  omeco exhaustive        17424  [(0, 1), (2, 3), (8, 9), (4, 5), (6, 7), (11, 12), (10, 13)]  (omeco flops 17424)
grid6: ad,abe,bf,dg,egh,fh->
  mini greedy              9728  [(0, 3), (2, 5), (1, 6), (4, 7), (8, 9)]
  tenferro default         9280  [(0, 3), (2, 5), (1, 6), (4, 8), (7, 9)]
  tenferro TreeSA          9280  [(3, 0), (6, 4), (7, 1), (8, 5), (2, 9)]
  omeco exhaustive         9280  [(0, 3), (6, 4), (7, 1), (8, 5), (9, 2)]  (omeco flops 9280)
chain8: ab,bc,cd,de,ef,fg,gh,hi->ai
  mini greedy              2740  [(0, 1), (2, 8), (3, 9), (4, 10), (5, 11), (6, 7), (12, 13)]
  tenferro default         2820  [(0, 1), (6, 7), (2, 3), (8, 10), (4, 5), (9, 12), (11, 13)]
  tenferro TreeSA          2656  [(2, 3), (7, 6), (5, 9), (4, 10), (8, 11), (1, 0), (12, 13)]
  omeco exhaustive         2656  [(0, 1), (2, 3), (6, 7), (5, 10), (4, 11), (9, 12), (8, 13)]  (omeco flops 2656)
ok

== process 1
default plan, 100 calls: 1 distinct
  [(0, 1), (2, 6), (3, 4), (5, 8), (7, 9)] cost 8712 x100
annealed plan, 20 calls: 4 distinct
  [(0, 1), (2, 6), (4, 5), (8, 3), (7, 9)] cost 8712 x10
  [(2, 1), (0, 6), (4, 3), (8, 5), (7, 9)] cost 8712 x4
  [(3, 4), (5, 6), (1, 2), (8, 0), (7, 9)] cost 8712 x3
  [(5, 4), (3, 6), (1, 0), (8, 2), (7, 9)] cost 8712 x3
einsum(), 100 calls: 1 distinct
  -2.91600766218305019e0  bits c00753fbd33fb164  x100
prepared plan, 100 runs: 1 distinct results

ij,jk->ik
  mini, greedy plan         19.49 ms  cost    134217728  max diff 2.7e-15
  mini, tenferro's plan     19.51 ms  cost    134217728  max diff 2.7e-15
  tenferro, 8 threads        0.88 ms
  tenferro, serial           5.59 ms
pa,aqb,brc,cs,pw,wqx,xry,ys->
  mini, greedy plan          2.80 ms  cost     16859136  max diff 1.7e-12
  mini, tenferro's plan      0.30 ms  cost       266256  max diff 1.4e-13
  tenferro, 8 threads        0.10 ms
  tenferro, serial           0.09 ms

ii->i
  mini     Err(Error("label 'i' repeats in input 0 (diagonal is not supported)"))
  tenferro Ok([1.0, 4.0])
ii->
  mini     Err(Error("label 'i' repeats in input 0 (diagonal is not supported)"))
  tenferro Ok([5.0])
...ij->ji
  mini     Err(Error("ellipsis '...' is not supported"))
  tenferro Ok([1.0, 3.0, 2.0, 4.0])
ij,jk
  mini     Err(Error("missing '->' in \"ij,jk\""))
  tenferro Err(invalid einsum subscripts: einsum notation must contain exactly one '->', got: ij,jk)
(ij,jk),kl->il
  mini     Err(Error("invalid label '('"))
  tenferro Err(invalid einsum subscripts: EinsumNotation::parse does not accept parentheses; use NestedEinsum::parse for parenthesized contraction order)

        // 动脑筋：代价换成结果的元素个数
        let keep: Vec<char> = builder.output.iter().chain(builder.live.iter().filter(|&&x| x != a && x != b).flat_map(|&x| builder.ops[x].iter())).copied().collect();
        let cost = step_cost(&Split::new(&builder.ops[a], &builder.ops[b], &keep).result_labels(), &[], sizes)?;

running 1 test
mps6: pa,aqb,br,pc,cqd,dr->
  mini greedy              8712  [(0, 1), (2, 6), (3, 4), (5, 8), (7, 9)]
  tenferro default         8712  [(0, 1), (2, 6), (3, 4), (5, 8), (7, 9)]
  tenferro TreeSA          8712  [(2, 1), (0, 6), (4, 3), (8, 5), (7, 9)]
  omeco exhaustive         8712  [(0, 1), (6, 2), (3, 4), (8, 5), (7, 9)]  (omeco flops 8712)
mps8: pa,aqb,brc,cs,pw,wqx,xry,ys->
  mini greedy             17424  [(0, 1), (2, 3), (8, 9), (4, 5), (6, 7), (11, 12), (10, 13)]
  tenferro default        17424  [(0, 1), (2, 3), (8, 9), (4, 5), (6, 7), (11, 12), (10, 13)]
  tenferro TreeSA         17424  [(6, 7), (4, 5), (8, 9), (0, 1), (3, 2), (11, 12), (10, 13)]
  omeco exhaustive        17424  [(0, 1), (2, 3), (8, 9), (4, 5), (6, 7), (11, 12), (10, 13)]  (omeco flops 17424)
grid6: ad,abe,bf,dg,egh,fh->
  mini greedy              9280  [(0, 3), (2, 5), (1, 6), (4, 8), (7, 9)]
  tenferro default         9280  [(0, 3), (2, 5), (1, 6), (4, 8), (7, 9)]
  tenferro TreeSA          9280  [(5, 2), (6, 4), (7, 1), (8, 5), (3, 9)]
  omeco exhaustive         9280  [(0, 3), (6, 4), (7, 1), (8, 5), (9, 2)]  (omeco flops 9280)
chain8: ab,bc,cd,de,ef,fg,gh,hi->ai
  mini greedy              2820  [(0, 1), (6, 7), (2, 3), (8, 10), (4, 5), (9, 12), (11, 13)]
  tenferro default         2820  [(0, 1), (6, 7), (2, 3), (8, 10), (4, 5), (9, 12), (11, 13)]
  tenferro TreeSA          2656  [(2, 3), (7, 6), (5, 9), (4, 10), (8, 11), (1, 0), (12, 13)]
  omeco exhaustive         2656  [(0, 1), (2, 3), (6, 7), (5, 10), (4, 11), (9, 12), (8, 13)]  (omeco flops 2656)
test plan_cost_vs_tenferro ... ok