import gradio as gr import numpy as np import scipy.sparse as sparse import time import os import shutil import math import sys from pathlib import Path # Assuming flex_chunk.py and matrix_multiply.py are in the same directory from flex_chunk import FlexChunk, save_chunk, load_chunk from matrix_multiply import prepare_chunks, load_chunks, matrix_vector_multiply # --- Matrix Generation (copied from test_vs_scipy.py) --- def generate_sparse_matrix(size, density, challenging=False): """ Generate a sparse test matrix with optional challenging patterns. Args: size: Matrix size (n x n) density: Target density challenging: Whether to include challenging patterns and extreme values Returns: A scipy.sparse.csr_matrix """ # Calculate number of non-zeros nnz = int(size * size * density) if nnz == 0: # Ensure at least one non-zero element if density is very low nnz = 1 if not challenging: # Simple random matrix rows = np.random.randint(0, size, nnz) cols = np.random.randint(0, size, nnz) data = np.random.rand(nnz) # Ensure the matrix actually has the specified size if nnz is small if nnz < size: # Add diagonal elements to ensure size diag_indices = np.arange(min(nnz, size)) rows = np.concatenate([rows, diag_indices]) cols = np.concatenate([cols, diag_indices]) data = np.concatenate([data, np.ones(len(diag_indices))]) # Use 1 for diagonal matrix = sparse.csr_matrix((data, (rows, cols)), shape=(size, size)) matrix.sum_duplicates() # Consolidate duplicate entries return matrix # --- Challenging matrix with specific patterns --- # Base random matrix (80% of non-zeros) base_nnz = int(nnz * 0.8) rows = np.random.randint(0, size, base_nnz) cols = np.random.randint(0, size, base_nnz) data = np.random.rand(base_nnz) # Add diagonal elements (10% of non-zeros) diag_nnz = int(nnz * 0.1) diag_indices = np.random.choice(size, diag_nnz, replace=False) # Add extreme values (10% of non-zeros) extreme_nnz = max(0, nnz - base_nnz - diag_nnz) # Ensure non-negative extreme_rows = np.random.randint(0, size, extreme_nnz) extreme_cols = np.random.randint(0, size, extreme_nnz) # Mix of very large and very small values extreme_data = np.concatenate([ np.random.uniform(1e6, 1e9, extreme_nnz // 2), np.random.uniform(1e-9, 1e-6, extreme_nnz - extreme_nnz // 2) ]) if extreme_nnz > 0 else np.array([]) if extreme_nnz > 0: np.random.shuffle(extreme_data) # Combine all components all_rows = np.concatenate([rows, diag_indices, extreme_rows]) all_cols = np.concatenate([cols, diag_indices, extreme_cols]) all_data = np.concatenate([data, np.random.rand(diag_nnz), extreme_data]) matrix = sparse.csr_matrix((all_data, (all_rows, all_cols)), shape=(size, size)) matrix.sum_duplicates() # Consolidate duplicate entries return matrix # --- Benchmark Function (Placeholder) --- def run_benchmark(size, density, num_chunks, challenging, 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...") time.sleep(1) # Placeholder # 1. Setup storage storage_dir = Path("./flex_chunk_temp_space") if storage_dir.exists(): shutil.rmtree(storage_dir) storage_dir.mkdir(exist_ok=True) progress(0.1, desc="Generating Matrix...") # 2. Generate matrix and vector matrix = generate_sparse_matrix(size, density, challenging) vector = np.random.rand(size) actual_nnz = matrix.nnz actual_density = actual_nnz / (size * size) if size > 0 else 0 matrix_info = f"Matrix: {size}x{size}, Target Density: {density:.6f}, Actual Density: {actual_density:.6f}, NNZ: {actual_nnz}" print(matrix_info) # For debugging in Hugging Face console # --- FlexChunk Run --- progress(0.2, desc="Preparing FlexChunks...") prepare_start = time.time() prepare_chunks(matrix, num_chunks, str(storage_dir), verbose=False) prepare_time = time.time() - prepare_start progress(0.4, desc="Loading FlexChunks...") load_start = time.time() chunks = load_chunks(str(storage_dir), verbose=False) load_time = time.time() - load_start progress(0.6, desc="Running FlexChunk SpMV...") flex_compute_start = time.time() flex_result = matrix_vector_multiply(chunks, vector, verbose=False) 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})" # --- Cleanup --- shutil.rmtree(storage_dir) progress(1.0, desc="Benchmark Complete") # --- Format Output --- results_summary = f""" {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** **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** **Comparison:** {comparison_result} """ return results_summary # --- Gradio Interface --- with gr.Blocks() as demo: gr.Markdown(""" # FlexChunk: Out-of-Core Sparse Matrix-Vector Multiplication (SpMV) Demo 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. 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. """) 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.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] ) 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).") # Launch the app if __name__ == "__main__": demo.launch()