FlexChunk Demo: Complete interactive application for sparse matrix-vector multiplication benchmarking

app.py

@@ -83,7 +83,7 @@ def generate_sparse_matrix(size, density, challenging=False):
 
 # --- Benchmark Function (Placeholder) ---
 
-def run_benchmark(size, density, num_chunks, challenging, progress=gr.Progress()):
+def run_benchmark(size, density, num_chunks, challenging, flex_only=False, progress=gr.Progress()):
     # This function will contain the main logic from test_vs_scipy.py
     # Adapted for Gradio inputs and outputs
     progress(0, desc="Starting Benchmark...")
@@ -122,60 +122,94 @@ def run_benchmark(size, density, num_chunks, challenging, progress=gr.Progress()
     flex_compute_time = time.time() - flex_compute_start
     flex_total_time = load_time + flex_compute_time
     
-    # --- SciPy Run ---
-    progress(0.7, desc="Saving SciPy data...")
-    scipy_temp_dir = storage_dir / "scipy_temp"
-    scipy_temp_dir.mkdir(exist_ok=True)
-    matrix_file = scipy_temp_dir / "matrix.npz"
-    vector_file = scipy_temp_dir / "vector.npy"
-    
-    scipy_save_start = time.time()
-    sparse.save_npz(matrix_file, matrix)
-    np.save(vector_file, vector)
-    scipy_save_time = time.time() - scipy_save_start
-    
-    progress(0.8, desc="Loading SciPy data...")
-    scipy_load_start = time.time()
-    loaded_matrix = sparse.load_npz(matrix_file)
-    loaded_vector = np.load(vector_file)
-    scipy_load_time = time.time() - scipy_load_start
-    
-    progress(0.9, desc="Running SciPy SpMV...")
-    scipy_compute_start = time.time()
-    scipy_result = loaded_matrix @ loaded_vector
-    scipy_compute_time = time.time() - scipy_compute_start
-    scipy_total_time = scipy_load_time + scipy_compute_time
-    
-    # --- Comparison ---
-    progress(0.95, desc="Comparing results...")
-    diff = np.abs(scipy_result - flex_result)
-    max_diff = np.max(diff) if len(diff) > 0 else 0
-    mean_diff = np.mean(diff) if len(diff) > 0 else 0
-    is_close = np.allclose(scipy_result, flex_result, atol=1e-9) # Increased tolerance slightly
-    comparison_result = f"✅ Results Match! (Max Diff: {max_diff:.2e}, Mean Diff: {mean_diff:.2e})" if is_close else f"❌ Results Differ! (Max Diff: {max_diff:.2e}, Mean Diff: {mean_diff:.2e})"
-
+    # Estimate FlexChunk memory usage
+    max_chunk_size = max(chunk.data.nbytes + chunk.col_indices.nbytes + chunk.row_offsets.nbytes for chunk in chunks)
+    flex_operational_memory = max_chunk_size + vector.nbytes + (size * 8)  # Chunk + vector + result vector
+    flex_memory_mb = flex_operational_memory / (1024*1024)
+    
+    # --- SciPy Run (Optional) ---
+    if not flex_only:
+        progress(0.7, desc="Saving SciPy data...")
+        scipy_temp_dir = storage_dir / "scipy_temp"
+        scipy_temp_dir.mkdir(exist_ok=True)
+        matrix_file = scipy_temp_dir / "matrix.npz"
+        vector_file = scipy_temp_dir / "vector.npy"
+        
+        scipy_save_start = time.time()
+        sparse.save_npz(matrix_file, matrix)
+        np.save(vector_file, vector)
+        scipy_save_time = time.time() - scipy_save_start
+        
+        progress(0.8, desc="Loading SciPy data...")
+        scipy_load_start = time.time()
+        loaded_matrix = sparse.load_npz(matrix_file)
+        loaded_vector = np.load(vector_file)
+        scipy_load_time = time.time() - scipy_load_start
+        
+        progress(0.9, desc="Running SciPy SpMV...")
+        scipy_compute_start = time.time()
+        scipy_result = loaded_matrix @ loaded_vector
+        scipy_compute_time = time.time() - scipy_compute_start
+        scipy_total_time = scipy_load_time + scipy_compute_time
+        
+        # Estimate SciPy memory usage
+        scipy_memory = loaded_matrix.data.nbytes + loaded_matrix.indices.nbytes + loaded_matrix.indptr.nbytes + loaded_vector.nbytes
+        scipy_memory_mb = scipy_memory / (1024*1024)
+        
+        # --- Comparison ---
+        progress(0.95, desc="Comparing results...")
+        diff = np.abs(scipy_result - flex_result)
+        max_diff = np.max(diff) if len(diff) > 0 else 0
+        mean_diff = np.mean(diff) if len(diff) > 0 else 0
+        is_close = np.allclose(scipy_result, flex_result, atol=1e-9) # Increased tolerance slightly
+        comparison_result = f"✅ Results Match! (Max Diff: {max_diff:.2e}, Mean Diff: {mean_diff:.2e})" if is_close else f"❌ Results Differ! (Max Diff: {max_diff:.2e}, Mean Diff: {mean_diff:.2e})"
+    
     # --- Cleanup ---
     shutil.rmtree(storage_dir)
     
     progress(1.0, desc="Benchmark Complete")
     
     # --- Format Output ---
-    results_summary = f"""
+    if flex_only:
+        results_summary = f"""
+## Matrix Information
 {matrix_info}
 
-**FlexChunk Performance:**
-- Prepare Chunks Time: {prepare_time:.4f}s
-- Load Chunks Time:    {load_time:.4f}s
-- Compute Time:        {flex_compute_time:.4f}s
-- **Total (Load+Compute): {flex_total_time:.4f}s**
+## FlexChunk Performance
+| Stage | Time |
+|-------|------|
+| Prepare Chunks | {prepare_time:.4f}s |
+| Load Chunks | {load_time:.4f}s |
+| Compute | {flex_compute_time:.4f}s |
+| **Total (Load+Compute)** | **{flex_total_time:.4f}s** |
+
+## Memory Usage
+| Metric | Value |
+|--------|-------|
+| Peak RAM Usage | {flex_memory_mb:.2f} MB |
+| Chunks | {num_chunks} |
+"""
+    else:
+        results_summary = f"""
+## Matrix Information
+{matrix_info}
 
-**SciPy Performance (Out-of-Core Emulation):**
-- Save Data Time:   {scipy_save_time:.4f}s (For reference)
-- Load Data Time:   {scipy_load_time:.4f}s
-- Compute Time:     {scipy_compute_time:.4f}s
-- **Total (Load+Compute): {scipy_total_time:.4f}s**
+## Performance Comparison
 
-**Comparison:**
+| Stage | FlexChunk | SciPy (Out-of-Core) |
+|-------|-----------|---------------------|
+| Data Preparation | {prepare_time:.4f}s | {scipy_save_time:.4f}s |  
+| Load Time | {load_time:.4f}s | {scipy_load_time:.4f}s |
+| Compute Time | {flex_compute_time:.4f}s | {scipy_compute_time:.4f}s |
+| **Total (Load+Compute)** | **{flex_total_time:.4f}s** | **{scipy_total_time:.4f}s** |
+
+## Memory Usage
+| Metric | FlexChunk | SciPy |
+|--------|-----------|-------|
+| Peak RAM Usage | {flex_memory_mb:.2f} MB | {scipy_memory_mb:.2f} MB |
+| Memory Ratio | 1.0x | {scipy_memory_mb/flex_memory_mb:.2f}x |
+
+## Comparison
 {comparison_result}
 """
     
@@ -183,38 +217,146 @@ def run_benchmark(size, density, num_chunks, challenging, progress=gr.Progress()
 
 # --- Gradio Interface ---
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # FlexChunk: Out-of-Core Sparse Matrix-Vector Multiplication (SpMV) Demo
+    # FlexChunk: Out-of-Core Sparse Matrix-Vector Multiplication
     
-    This demo compares the performance of FlexChunk against standard SciPy for SpMV, 
-    simulating an out-of-core scenario where the matrix doesn't fit entirely in memory.
+    This interactive demo showcases **FlexChunk**, an algorithm for performing Sparse Matrix-Vector Multiplication (SpMV) on matrices that may be too large to fit entirely in memory.
     
-    FlexChunk splits the matrix into smaller chunks, processing them sequentially to reduce peak memory usage. 
-    SciPy performance includes the time to save and load the matrix from disk to mimic this out-of-core access.
+    **Key Benefits:**
+    * Process matrices up to 100M×100M using only ~1.7GB RAM
+    * Near-linear scaling in both time and memory usage
+    * Outperforms traditional approaches for large out-of-core matrices
     """)
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("**Benchmark Parameters**")
-            size_input = gr.Slider(label="Matrix Size (N x N)", minimum=100, maximum=50000, value=10000, step=100)
-            # Max density adjusted to prevent excessive nnz for large matrices in demo
-            density_input = gr.Slider(label="Matrix Density", minimum=0.00001, maximum=0.01, value=0.0001, step=0.00001, format="%.5f") 
-            chunks_input = gr.Slider(label="Number of Chunks", minimum=1, maximum=32, value=4, step=1)
-            challenging_input = gr.Checkbox(label="Use Challenging Matrix (Extreme Values)", value=False)
-            run_button = gr.Button("Run Benchmark", variant="primary")
+    with gr.Tabs() as tabs:
+        # Standard mode tab
+        with gr.TabItem("Standard Mode"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Matrix Parameters")
+                    standard_size = gr.Slider(
+                        label="Matrix Size (N×N)", 
+                        minimum=1000, 
+                        maximum=200000, 
+                        value=10000, 
+                        step=1000,
+                        info="Square matrix dimension (N×N)"
+                    )
+                    standard_density = gr.Slider(
+                        label="Matrix Density", 
+                        minimum=0.00001, 
+                        maximum=0.01, 
+                        value=0.0001, 
+                        step=0.00001,
+                        info="Fraction of non-zero elements (0.0001 = 0.01%)"
+                    )
+                    standard_chunks = gr.Slider(
+                        label="Number of Chunks", 
+                        minimum=1, 
+                        maximum=32, 
+                        value=4, 
+                        step=1,
+                        info="More chunks = less memory but more overhead"
+                    )
+                    standard_challenging = gr.Checkbox(
+                        label="Use Challenging Matrix",
+                        info="Includes extreme values and special patterns"
+                    )
+                    standard_flexonly = gr.Checkbox(
+                        label="FlexChunk Only",
+                        info="Skip SciPy comparison for better performance"
+                    )
+                    standard_button = gr.Button("Run Benchmark", variant="primary")
             
-        with gr.Column(scale=2):
-            gr.Markdown("**Results**")
-            output_textbox = gr.Markdown(label="Benchmark Summary")
-
-    run_button.click(
-        fn=run_benchmark, 
-        inputs=[size_input, density_input, chunks_input, challenging_input], 
-        outputs=[output_textbox]
+            standard_output = gr.Markdown()
+        
+        # Advanced mode tab
+        with gr.TabItem("Advanced Mode"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Large Matrix Parameters")
+                    gr.Markdown("""
+                    ⚠️ **Warning**: Processing time varies with matrix size:
+                    - 1M×1M matrices: ~1 second
+                    - 10M×10M matrices: ~10 seconds
+                    - 100M×100M matrices: ~1 minute 47 seconds
+                    
+                    For large matrices, FlexChunk-only mode is automatically enabled.
+                    """)
+                    
+                    advanced_size = gr.Slider(
+                        label="Matrix Size (N×N)", 
+                        minimum=50000, 
+                        maximum=300000000, 
+                        value=100000, 
+                        step=50000,
+                        info="Square matrix dimension - up to 300M×300M (extremely large values will take significant time)"
+                    )
+                    advanced_density = gr.Slider(
+                        label="Matrix Density", 
+                        minimum=0.0000001, 
+                        maximum=0.001, 
+                        value=0.000001, 
+                        step=0.0000001,
+                        info="Use lower density for very large matrices"
+                    )
+                    advanced_chunks = gr.Slider(
+                        label="Number of Chunks", 
+                        minimum=4, 
+                        maximum=100, 
+                        value=10, 
+                        step=1,
+                        info="More chunks recommended for larger matrices"
+                    )
+                    advanced_challenging = gr.Checkbox(
+                        label="Use Challenging Matrix",
+                        info="Includes extreme values and special patterns"
+                    )
+                    
+                    # Force FlexChunk only for advanced mode
+                    gr.Markdown("*SciPy comparison disabled for large matrices*")
+                    advanced_button = gr.Button("Run Advanced Benchmark", variant="primary")
+            
+            advanced_output = gr.Markdown()
+    
+    # Event handlers
+    standard_button.click(
+        fn=run_benchmark,
+        inputs=[standard_size, standard_density, standard_chunks, standard_challenging, standard_flexonly],
+        outputs=standard_output
     )
     
-    gr.Markdown("--- Developed based on the [FlexChunk concept](https://www.lesswrong.com/posts/zpRhsdDkWygTDScxb/flexchunk-enabling-100m-100m-out-of-core-spmv-1-8-min-1-7-gb).")
+    advanced_button.click(
+        fn=lambda size, density, chunks, challenging: run_benchmark(size, density, chunks, challenging, True),
+        inputs=[advanced_size, advanced_density, advanced_chunks, advanced_challenging],
+        outputs=advanced_output
+    )
+    
+    gr.Markdown("""
+    ---
+    ### About FlexChunk
+    
+    FlexChunk enables processing matrices that would normally exceed RAM capacity by dividing them into manageable chunks.
+    
+    **Links:**
+    - Read more in the [original article](https://www.lesswrong.com/posts/zpRhsdDkWygTDScxb/flexchunk-enabling-100m-100m-out-of-core-spmv-1-8-min-1-7-gb)
+    - View source code on [GitHub](https://github.com/DanielSwift1992/FlexChunk)
+    
+    ---
+    ### Benchmark Results
+    
+    Actual performance measurements from our tests:
+    
+    | Matrix Size     | Non-zero Elements | Total Time    | Peak RAM Usage |
+    |-----------------|-------------------|---------------|----------------|
+    | 1.0M × 1.0M     | 1.2M              | 1.07 s        | 17.00 MB       |
+    | 10.0M × 10.0M   | 12.0M             | 10.21 s       | 170.00 MB      |
+    | 50.0M × 50.0M   | 62.5M             | 55.27 s       | 850.00 MB      |
+    | 100.0M × 100.0M | 120.0M            | 1 min 47.1 s  | 1.70 GB        |
+    
+    The algorithm scales nearly linearly and can theoretically handle even larger matrices (up to 300M×300M), with proportionally increased processing time and memory usage.
+    """)
 
 # Launch the app
 if __name__ == "__main__":