v2 backend architecture: 2-level IR with StableHLO cut point

[shinaoka]

Summary

Brainstorming results for the v2 backend architecture. This resolves the core design questions from #20 (semiring backend contracts, compile-cache, StdTensorOp consistency) by establishing a clean 2-level IR architecture.

Key decisions

1. StableHLO-compatible IR as the single cut point

All computation (standard and custom algebra) is lowered to a StableHLO-compatible IR (Rust structs, in-process). This is the sole interface boundary:

• XLA backend: receives the StableHLO IR directly. XLA handles all optimization internally.
• faer / custom backends: tenferro's optimizing compiler lowers StableHLO IR → low-level IR → generic execution engine.

Surface API → Fragment → CompiledProgram
    ↓
StableHLO IR (Rust struct, complete, standard-compliant)  ← single cut point
    ├─ XLA backend: pass directly to XLA
    └─ faer / custom backends:
         Optimizing compiler → Low-level IR → Generic engine → Backend trait

2. Optimizing compiler (algebra-agnostic)

Transforms StableHLO IR to low-level IR. Key passes (modeled after XLA):

Pass	Role
TransposeFolding	Absorb Transpose into DotGeneral's dimension_numbers
DotDecomposer	Canonicalize DotGeneral → Transpose + Reshape + BatchedGemm
Contiguous materialization	Insert physical copy where needed

These passes are algebra-agnostic — they work identically for standard and custom algebras.

3. Low-level IR: contiguous, column-major

All buffers in the low-level IR are contiguous column-major. No stride tricks at this level.

Key instructions:

• BatchedGemm(batch, m, n, k) — contiguous [batch, M, K] × [batch, K, N]
• ReduceSum(axes) — contiguous input
• Permute(perm) / Copy — physical memory reordering
• Reshape — noop if element count unchanged, else copy

4. Generic execution engine

Interprets low-level IR step by step. Generic over backend trait. No kernel fusion — simple sequential dispatch.

fn execute_instruction<Alg, B: SemiringCore<Alg>>(inst: &LowLevelInst, ...) {
    match inst {
        BatchedGemm { .. } => {
            // Try optional fast path first
            if let Some(fp) = fast_path {
                if fp.contract(&desc, a, b, c) { return; }
            }
            // Fallback to required trait
            backend.batched_gemm(plan, a, b, c);
        }
        ReduceSum { axes } => backend.reduce_sum(axes, input, output),
        Permute { perm } => { /* physical memory copy */ }
        Reshape { .. } => { /* noop or copy */ }
    }
}

5. Custom algebra backend: minimal contract

Required	Optional (fast path)
batched_gemm()	contract() (direct contraction, skips Transpose+Reshape+BatchedGemm)
reduce_sum()	elementwise_mul() (faster than degenerate K=1 BatchedGemm)

This matches tenferro v1's TensorSemiringCore + TensorSemiringFastPath design. v1 implementation should be reused as much as possible.

Why only these two:

• BatchedGemm handles all contraction (⊕ and ⊗ inside)
• ReduceSum handles all dimension reduction (⊕ inside)
• Add (standalone): not needed for einsum pipeline; AD cotangent accumulation doesn't apply to custom algebras
• Mul (standalone): expressible as degenerate BatchedGemm (K=1); optional fast path for efficiency
• DotGeneral: decomposed by the compiler into Transpose + Reshape + BatchedGemm
• Transpose, Reshape, BroadcastInDim: handled by the common infrastructure (algebra-independent)

6. Low-level IR and trait methods need NOT be 1:1

The engine can pattern-match sequences in the low-level IR and dispatch to higher-level backend methods:

• IR has Transpose + Reshape + BatchedGemm (3 instructions)
• If SemiringFastPath::contract() is available: dispatch all 3 as a single call
• Otherwise: execute each instruction sequentially

7. Tensor allows arbitrary strides; contiguous at IR boundary

• tenferro::Tensor supports arbitrary strides (zero-copy permute, slice, view)
• At the StableHLO IR entry point, non-contiguous inputs are materialized to contiguous
• All computation below the IR boundary operates on contiguous column-major buffers
• eval() output is always contiguous

8. Column-major throughout

tenferro standardizes on column-major (Fortran) layout at all levels:

• Tensor storage: column-major
• Low-level IR: column-major contiguous
• StableHLO IR: layout annotation specifies column-major
• XLA backend: XLA receives column-major layout specification and handles internal layout assignment (may use row-major internally for GPU, returns column-major output)
• BLAS/LAPACK/faer: column-major native — no conversion needed
• tropical-gemm: column-major native

Relationship to existing code

v1 (origin/main)	v2
einsum engine's lazy permutation	Optimizing compiler (TransposeFolding)
prepare_one_operand + fusability check	SemiringFastPath::contract()
SemiringCoreDescriptor	Low-level IR instruction set
TensorSemiringCore	SemiringCore trait (required)
TensorSemiringFastPath	SemiringFastPath trait (optional)
EinsumBackend = Core + FastPath	Generic execution engine

Evidence from XLA / JAX

• JAX minimizes Transpose generation by trying both operand orders and using DotGeneral(dimension_numbers) to encode contraction axes without physical transpose
• XLA has a multi-pass pipeline: DotDimensionSorter → DotDecomposer → TransposeFolding (×3) → DotMerger → AlgebraicSimplifier to eliminate unnecessary transposes
• XLA's DotDecomposer canonicalizes arbitrary DotGeneral to [batch, M, K] × [batch, K, N] form — exactly BatchedGemm
• StableHLO's dot_general does NOT include standalone ReduceSum — non-contracting dimensions always appear in the output. ij,jk->i requires dot_general(ij,jk->ik) + reduce(k).

Evidence from existing libraries

• tropical-gemm: implements only BatchedGemm (pure GEMM library). No reduce, no elementwise ops. This covers the heaviest kernel.
• omeinsum-rs: N-ary einsum engine. Its contract_binary is reshape→GEMM→reshape. Reduction of non-contracting dims is handled by the einsum engine decomposing into GEMM + unary reduce steps.

Resolves

Resolves the design questions in #20:

1. Custom semiring minimum contract: batched_gemm + reduce_sum. Structural ops (Transpose, Reshape, BroadcastInDim) are common infrastructure.
2. Compile-cache identity: tenferro Engine owns a normalized compile-cache layer above computegraph (normalizes away InputKey/DiffPassId). computegraph's GlobalValKey-based cache stays unchanged.
3. StdTensorOp consistency: StableHLO IR is the single source of truth. StdTensorOp lowers to it via lower_to_stablehlo(). Enum definition and lowering tables must be unified against primitive-catalog.md.
4. Descriptor + plan/execute as v2 pattern: Yes, via SemiringCore + SemiringFastPath traits (same pattern as v1).

[shinaoka]

Addendum: detailed Op type design (from continued brainstorming)

Layer 1 → Layer 2 → Layer 3 → Layer 4 mapping

Layer 1: Graph Op types (Fragment construction)
  StdTensorOp       — flat, mirrors StableHLO 1:1, AD-capable
  SemiringOp<T>     — wraps SemiringOpKind, no AD
  SemiringOps trait  — shared interface for einsum builder

Layer 2: StableHLO IR (the single cut point)
  StableHloProgram / StableHloOp
  Both StdTensorOp and SemiringOp<T> lower here

Layer 3: Low-level IR (contiguous, column-major)
  LowLevelProgram / LowLevelOp
  Output of optimizing compiler

Layer 4: Backend trait
  SemiringCore (required) + SemiringFastPath (optional)

StdTensorOp (flat, no SemiringOpKind nesting)

enum StdTensorOp {
    // Semiring-compatible (implements SemiringOps trait)
    Add, Mul, Neg, Conj,
    DotGeneral(DotGeneralConfig),
    Transpose { perm: Vec<usize> },
    Reshape { shape: Vec<usize> },
    BroadcastInDim { shape: Vec<usize>, dims: Vec<usize> },
    ReduceSum { axes: Vec<usize> },

    // Standard arithmetic only
    Div, Abs, Sign, Maximum, Minimum,
    Compare(CompareDir), Select, Clamp,

    // Analytic
    Exp, Log, Sin, Cos, Tanh, Sqrt, Rsqrt, Pow, Expm1, Log1p,

    // Indexing & structure
    Gather(GatherConfig), Scatter(ScatterConfig),
    Slice(SliceConfig), DynamicSlice, Pad(PadConfig),
    Concatenate { axis: usize }, Reverse { axes: Vec<usize> },

    // Additional reductions
    ReduceProd { axes: Vec<usize> },
    ReduceMax { axes: Vec<usize> },
    ReduceMin { axes: Vec<usize> },

    // Linalg
    Cholesky, Svd, Qr, Eigh, Solve,
    CustomCall { target: String, n_inputs: usize, n_outputs: usize },
}

SemiringOpKind (custom algebra, inside SemiringOp only)

enum SemiringOpKind {
    Add, Mul,
    DotGeneral(DotGeneralConfig),
    ReduceSum { axes: Vec<usize> },
    Transpose { perm: Vec<usize> },
    Reshape { shape: Vec<usize> },
    BroadcastInDim { shape: Vec<usize>, dims: Vec<usize> },
}

Transpose/Reshape/BroadcastInDim remain in SemiringOpKind because the einsum builder needs to emit them as graph nodes. However, custom algebra authors do NOT implement these — see trait design below.

SemiringOps trait (generic einsum builder interface)

trait SemiringOps: GraphOp {
    fn add() -> Self;
    fn mul() -> Self;
    fn dot_general(config: DotGeneralConfig) -> Self;
    fn reduce_sum(axes: Vec<usize>) -> Self;
    fn transpose(perm: Vec<usize>) -> Self;
    fn reshape(shape: Vec<usize>) -> Self;
    fn broadcast_in_dim(shape: Vec<usize>, dims: Vec<usize>) -> Self;
}

// StdTensorOp: maps to flat variants
impl SemiringOps for StdTensorOp {
    fn add() -> Self { StdTensorOp::Add }
    fn dot_general(c) -> Self { StdTensorOp::DotGeneral(c) }
    // ... direct mapping
}

// SemiringOp<T>: wraps in SemiringOpKind
impl<T: Operand> SemiringOps for SemiringOp<T> {
    fn add() -> Self { SemiringOp { kind: SemiringOpKind::Add, .. } }
    // ...
}

Operand and TensorData traits (custom algebra author's contract)

// Operand: pure algebra contract (custom author MUST implement)
trait Operand: Clone + Send + Sync + 'static {
    type Scalar: Scalar;
    fn zero(shape: &[usize]) -> Self;
    fn one(shape: &[usize]) -> Self;
    fn add(&self, other: &Self) -> Self;          // ⊕
    fn multiply(&self, other: &Self) -> Self;      // ⊗
    fn dot_general(&self, other: &Self, config: &DotGeneralConfig) -> Self;
    fn reduce_sum(&self, axes: &[usize]) -> Self;
}

// TensorData: buffer access (custom author implements, but trivial)
trait TensorData {
    type Scalar: Scalar;
    fn shape(&self) -> &[usize];
    fn strides(&self) -> &[isize];
    fn data(&self) -> &[Self::Scalar];
    fn from_data(data: Vec<Self::Scalar>, shape: Vec<usize>) -> Self;
}

// Structural ops: generic implementations, NOT on Operand trait
fn generic_transpose<T: TensorData>(t: &T, perm: &[usize]) -> T { /* ... */ }
fn generic_reshape<T: TensorData>(t: &T, shape: &[usize]) -> T { /* ... */ }
fn generic_broadcast<T: TensorData>(t: &T, shape: &[usize], dims: &[usize]) -> T { /* ... */ }

GraphOp::eval for SemiringOp uses Operand methods for algebraic ops and generic functions for structural ops:

impl<T: Operand + TensorData> GraphOp for SemiringOp<T> {
    fn eval(&self, ctx, inputs) -> Vec<T> {
        match &self.kind {
            SemiringOpKind::Add => vec![inputs[0].add(inputs[1])],
            SemiringOpKind::Mul => vec![inputs[0].multiply(inputs[1])],
            SemiringOpKind::DotGeneral(c) => vec![inputs[0].dot_general(inputs[1], c)],
            SemiringOpKind::ReduceSum { axes } => vec![inputs[0].reduce_sum(axes)],
            SemiringOpKind::Transpose { perm } => vec![generic_transpose(&inputs[0], perm)],
            SemiringOpKind::Reshape { shape } => vec![generic_reshape(&inputs[0], shape)],
            SemiringOpKind::BroadcastInDim { s, d } => vec![generic_broadcast(&inputs[0], s, d)],
        }
    }
}

Note: GraphOp::eval is the reference implementation for testing only. Production execution goes through StableHLO IR → optimizing compiler → low-level IR → backend trait.

StableHLO lowering (both Op types → shared IR)

// StdTensorOp → StableHloOp: mostly 1:1
fn lower_std(op: &StdTensorOp) -> Vec<StableHloInstruction> { /* ... */ }

// SemiringOp<T> → StableHloOp: semiring subset only
fn lower_semiring<T>(op: &SemiringOp<T>) -> Vec<StableHloInstruction> { /* ... */ }

Important: StableHloOp::Add in the IR means standard addition for StdTensorOp, but the custom algebra backend interprets it as ⊕. The IR structure is identical; only execution semantics differ.

Low-level IR (contiguous, column-major)

enum LowLevelOp {
    // Algebraic (dispatched to backend trait)
    BatchedGemm { batch_dims: Vec<usize>, m: usize, n: usize, k: usize },
    ReduceSum { axes: Vec<usize> },

    // Structural (common infrastructure, algebra-independent)
    Permute { perm: Vec<usize> },  // physical memory copy
    Reshape { shape: Vec<usize> }, // noop if element count unchanged
    Copy,                          // contiguify

    // Optional (emitted when fast path available)
    ElementwiseMul,
    ElementwiseAdd,
}

All buffers are contiguous column-major. No stride tricks at this level.

Additional decisions from the session

• Tensor allows arbitrary strides: zero-copy permute/slice/view at the user level. Contiguous materialization happens at the StableHLO IR entry point.
• Column-major (Fortran) layout throughout: BLAS/LAPACK/faer/tropical-gemm are all column-major native. XLA backend specifies column-major layout; XLA handles internal layout assignment.
• Compile cache normalization: tenferro Engine normalizes InputKey/DiffPassId to anonymous slot indices for cache lookup. computegraph-rs cache stays GlobalValKey-based (unchanged).
• Cholesky: direct StableHLO op. SVD/QR/Eigh/Solve: CustomCall.