1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
//! `TensorGraph` — a directed weighted graph whose algorithms are matrix powers.
//!
//! The adjacency matrix is stored as a [`PaddedTileLattice`], and classic graph
//! questions are answered by raising it to a power over the right semiring
//! ([`crate::semiring`]):
//!
//! - **Reachability / transitive closure**: boolean matrix powers.
//! - **All-pairs shortest paths**: tropical (min-plus) matrix powers.
//! - **Walk counting**: ordinary matrix powers.
//!
//! Each "power" is computed by repeated squaring, so a graph of `n` nodes needs
//! only `⌈log₂ n⌉` dense matmuls — turning a traversal into a handful of GEMMs.
use crate::geometry::Geometry;
use crate::lattice::PaddedTileLattice;
use crate::semiring::{semiring_matmul, Boolean, Counting, Semiring, Tropical};
/// A directed, weighted graph backed by a dense adjacency lattice.
#[derive(Clone)]
pub struct TensorGraph {
n: usize,
geom: Geometry,
edges: Vec<(usize, usize, f32)>,
}
impl TensorGraph {
/// Create an `n`-node graph with no edges.
pub fn new(n: usize) -> Self {
TensorGraph {
n,
geom: Geometry::TPU_V,
edges: Vec::new(),
}
}
/// Number of nodes.
#[inline]
pub fn nodes(&self) -> usize {
self.n
}
/// Number of edges added.
#[inline]
pub fn edge_count(&self) -> usize {
self.edges.len()
}
/// Add a weighted directed edge `u → v`.
pub fn add_edge(&mut self, u: usize, v: usize, weight: f32) {
assert!(u < self.n && v < self.n, "endpoints must be in range");
self.edges.push((u, v, weight));
}
/// Add an unweighted directed edge `u → v` (weight `1`).
pub fn add_edge_unweighted(&mut self, u: usize, v: usize) {
self.add_edge(u, v, 1.0);
}
/// Build the boolean adjacency lattice (`1` where an edge exists), with the
/// diagonal set so the closure is reflexive.
fn reflexive_boolean(&self) -> PaddedTileLattice<f32> {
let mut dense = vec![0.0f32; self.n * self.n];
for i in 0..self.n {
dense[i * self.n + i] = 1.0;
}
for &(u, v, _) in &self.edges {
dense[u * self.n + v] = 1.0;
}
PaddedTileLattice::from_dense(self.n, self.n, &dense, self.geom).unwrap()
}
/// Build the tropical adjacency lattice: edge weights off-diagonal, `0` on the
/// diagonal, `+∞` where there is no edge.
fn tropical_adjacency(&self) -> PaddedTileLattice<f32> {
let mut dense = vec![f32::INFINITY; self.n * self.n];
for i in 0..self.n {
dense[i * self.n + i] = 0.0;
}
for &(u, v, w) in &self.edges {
let slot = &mut dense[u * self.n + v];
*slot = slot.min(w);
}
PaddedTileLattice::from_dense(self.n, self.n, &dense, self.geom).unwrap()
}
/// Build the `0/1` adjacency lattice with no self-loops, for walk counting.
fn plain_adjacency(&self) -> PaddedTileLattice<f32> {
let mut dense = vec![0.0f32; self.n * self.n];
for &(u, v, _) in &self.edges {
dense[u * self.n + v] = 1.0;
}
PaddedTileLattice::from_dense(self.n, self.n, &dense, self.geom).unwrap()
}
/// Number of squarings needed to reach any path length up to `n`.
fn squarings(&self) -> u32 {
// ⌈log₂(max(1, n))⌉
let mut steps = 0;
let mut reach = 1usize;
while reach < self.n.max(1) {
reach *= 2;
steps += 1;
}
steps
}
/// Raise `m` to a closure power over semiring `S` by repeated squaring.
fn close<S: Semiring>(&self, mut m: PaddedTileLattice<f32>) -> PaddedTileLattice<f32> {
for _ in 0..self.squarings() {
m = semiring_matmul::<S>(&m, &m).unwrap();
}
m
}
/// The reflexive-transitive closure: `reachable[u][v]` is `1` iff `v` is
/// reachable from `u`. Computed as boolean matrix powers.
pub fn reachability(&self) -> PaddedTileLattice<f32> {
self.close::<Boolean>(self.reflexive_boolean())
}
/// True if `v` is reachable from `u` (including `u == v`).
pub fn reachable(&self, u: usize, v: usize) -> bool {
*self.reachability().get(u, v).unwrap() != 0.0
}
/// All-pairs shortest path distances, computed as tropical matrix powers.
/// Unreachable pairs hold `+∞`.
pub fn shortest_paths(&self) -> PaddedTileLattice<f32> {
self.close::<Tropical>(self.tropical_adjacency())
}
/// The shortest-path distance from `u` to `v`, or `None` if unreachable.
pub fn distance(&self, u: usize, v: usize) -> Option<f32> {
let d = *self.shortest_paths().get(u, v).unwrap();
if d.is_finite() {
Some(d)
} else {
None
}
}
/// The number of distinct walks of exactly `k` edges from `u` to `v`, as the
/// `k`-th ordinary matrix power of the adjacency.
pub fn walk_counts(&self, k: usize) -> PaddedTileLattice<f32> {
let adj = self.plain_adjacency();
if k == 0 {
// The 0th power is the identity.
let mut dense = vec![0.0f32; self.n * self.n];
for i in 0..self.n {
dense[i * self.n + i] = 1.0;
}
return PaddedTileLattice::from_dense(self.n, self.n, &dense, self.geom).unwrap();
}
let mut acc = adj.clone();
for _ in 1..k {
acc = semiring_matmul::<Counting>(&acc, &adj).unwrap();
}
acc
}
}
impl core::fmt::Debug for TensorGraph {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(
f,
"TensorGraph {{ nodes: {}, edges: {} }}",
self.n,
self.edges.len()
)
}
}