From 2a875a9ec0bc729605186abd01f4df2d5303fd1e Mon Sep 17 00:00:00 2001
From: Jan Ferdinand Sauer <ferdinand@neptune.cash>
Date: Tue, 15 Oct 2024 13:20:49 +0200
Subject: [PATCH] perf(RAM): Clear unusable caches

It is possible that the quotient domain codewords are cached for the
main table, but not for the auxiliary table. To compute the quotients,
either both or neither are needed.[^1] For peak memory consumption, it
is beneficial to clear any unused cache.

Throwing away any cache here incurs a performance penalty later, when
revealing the rows indicated by FRI. This is deemed an acceptable
tradeoff when peak memory usage is a larger issue than compute time.

[^1]: Code using exactly one cache could exist, but oh! the engineering.
---
 triton-vm/src/stark.rs              | 15 +++++++++++++++
 triton-vm/src/table/master_table.rs | 19 +++++++++++--------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/triton-vm/src/stark.rs b/triton-vm/src/stark.rs
index dde1bcb2..a54ef6b0 100644
--- a/triton-vm/src/stark.rs
+++ b/triton-vm/src/stark.rs
@@ -160,6 +160,21 @@ impl Stark {
         profiler!(stop "Fiat-Shamir");
         profiler!(stop "aux tables");
 
+        // It is possible that the quotient domain codewords are cached for the main
+        // table, but not for the auxiliary table. To compute the quotients, either both
+        // or neither are needed.[^1] For peak memory consumption, it is beneficial to
+        // clear any unused cache.
+        //
+        // Throwing away any cache here incurs a performance penalty later, when
+        // revealing the rows indicated by FRI. This is deemed an acceptable tradeoff
+        // when peak memory usage is a larger issue than compute time.
+        //
+        // [^1]: Code using exactly one cache _could_ exist, but oh! the engineering.
+        if master_main_table.cache_is_empty() || master_aux_table.cache_is_empty() {
+            master_main_table.clear_cache();
+            master_aux_table.clear_cache();
+        }
+
         let (fri_domain_quotient_segment_codewords, quotient_segment_polynomials) =
             Self::compute_quotient_segments(
                 &master_main_table,
diff --git a/triton-vm/src/table/master_table.rs b/triton-vm/src/table/master_table.rs
index 9e2ff475..f17d2f60 100644
--- a/triton-vm/src/table/master_table.rs
+++ b/triton-vm/src/table/master_table.rs
@@ -294,8 +294,13 @@ where
         low_degree_extended_columns: Array2<Self::Field>,
     );
 
-    /// Return the cached low-degree-extended table, if any.
-    fn low_degree_extended_table(&self) -> Option<ArrayView2<Self::Field>>;
+    #[doc(hidden)]
+    fn cache_is_empty(&self) -> bool {
+        self.fri_domain_table().is_none()
+    }
+
+    #[doc(hidden)]
+    fn clear_cache(&mut self);
 
     /// Return the FRI domain view of the cached low-degree-extended table, if any.
     ///
@@ -711,9 +716,8 @@ impl MasterTable for MasterMainTable {
         self.low_degree_extended_table = Some(low_degree_extended_columns);
     }
 
-    fn low_degree_extended_table(&self) -> Option<ArrayView2<BFieldElement>> {
-        let low_degree_extended_table = self.low_degree_extended_table.as_ref()?;
-        Some(low_degree_extended_table.view())
+    fn clear_cache(&mut self) {
+        drop(self.low_degree_extended_table.take());
     }
 
     fn fri_domain_table(&self) -> Option<ArrayView2<BFieldElement>> {
@@ -784,9 +788,8 @@ impl MasterTable for MasterAuxTable {
         self.low_degree_extended_table = Some(low_degree_extended_columns);
     }
 
-    fn low_degree_extended_table(&self) -> Option<ArrayView2<XFieldElement>> {
-        let low_degree_extended_table = self.low_degree_extended_table.as_ref()?;
-        Some(low_degree_extended_table.view())
+    fn clear_cache(&mut self) {
+        drop(self.low_degree_extended_table.take());
     }
 
     fn fri_domain_table(&self) -> Option<ArrayView2<XFieldElement>> {