Spaces:
Runtime error
Runtime error
| """ | |
| Interactive GEMM Memory Access Pattern Visualizer using Dash | |
| Hugging Face Space version | |
| A web-based interactive visualization tool for exploring different | |
| GEMM loop orderings and their memory access patterns. | |
| """ | |
| import os | |
| import dash | |
| from dash import dcc, html, Input, Output, State, callback_context | |
| import plotly.graph_objects as go | |
| from plotly.subplots import make_subplots | |
| import numpy as np | |
| from gemm_simulator import GEMMSimulator | |
| from cache_simulator import CacheSimulator | |
| # Initialize the Dash app | |
| app = dash.Dash(__name__, suppress_callback_exceptions=True) | |
| app.title = "GEMM Visualizer" | |
| server = app.server # For Hugging Face Spaces deployment | |
| # Global state for simulation data | |
| simulation_data = { | |
| 'simulator': None, | |
| 'tracks': [], | |
| 'current_frame': 0, | |
| 'cache_stats': {}, | |
| 'heatmaps': {} | |
| } | |
| # Color scheme for matrices | |
| COLORS = { | |
| 'A': 'rgba(255, 0, 0, 0.7)', # Red | |
| 'B': 'rgba(0, 255, 0, 0.7)', # Green | |
| 'C': 'rgba(0, 0, 255, 0.7)', # Blue | |
| } | |
| # App layout | |
| app.layout = html.Div([ | |
| html.Div([ | |
| html.H1("🔬 GEMM Memory Access Pattern Visualizer", | |
| style={'textAlign': 'center', 'color': '#2c3e50', 'marginBottom': 10}), | |
| html.P("Explore different loop orderings and their impact on cache performance", | |
| style={'textAlign': 'center', 'color': '#7f8c8d', 'fontSize': 16}) | |
| ], style={'backgroundColor': '#ecf0f1', 'padding': '20px', 'borderRadius': '10px'}), | |
| html.Div([ | |
| # Left panel - Controls | |
| html.Div([ | |
| html.H3("⚙️ Configuration", style={'color': '#34495e'}), | |
| html.Label("Matrix Size (n × n):", style={'fontWeight': 'bold', 'marginTop': 15}), | |
| dcc.Slider( | |
| id='matrix-size-slider', | |
| min=4, max=32, step=2, value=16, | |
| marks={i: str(i) for i in [4, 8, 16, 24, 32]}, | |
| tooltip={"placement": "bottom", "always_visible": True} | |
| ), | |
| html.Label("Block Size:", style={'fontWeight': 'bold', 'marginTop': 25}), | |
| dcc.Slider( | |
| id='block-size-slider', | |
| min=2, max=16, step=2, value=4, | |
| marks={i: str(i) for i in [2, 4, 8, 16]}, | |
| tooltip={"placement": "bottom", "always_visible": True} | |
| ), | |
| html.Label("Loop Order:", style={'fontWeight': 'bold', 'marginTop': 25}), | |
| dcc.Dropdown( | |
| id='loop-order-dropdown', | |
| options=[ | |
| {'label': 'IJK (row-column-depth)', 'value': 'ijk'}, | |
| {'label': 'IKJ (row-depth-column)', 'value': 'ikj'}, | |
| {'label': 'JIK (column-row-depth)', 'value': 'jik'}, | |
| {'label': 'JKI (column-depth-row)', 'value': 'jki'}, | |
| {'label': 'KIJ (depth-row-column)', 'value': 'kij'}, | |
| {'label': 'KJI (depth-column-row)', 'value': 'kji'}, | |
| ], | |
| value='kji', | |
| clearable=False | |
| ), | |
| html.Div([ | |
| html.Label("Blocking:", style={'fontWeight': 'bold', 'marginTop': 25}), | |
| dcc.RadioItems( | |
| id='blocking-radio', | |
| options=[ | |
| {'label': ' Blocked', 'value': True}, | |
| {'label': ' Unblocked', 'value': False} | |
| ], | |
| value=True, | |
| inline=True, | |
| style={'marginTop': 10} | |
| ) | |
| ]), | |
| html.Hr(style={'marginTop': 30, 'marginBottom': 20}), | |
| html.H3("🎬 Animation Controls", style={'color': '#34495e'}), | |
| html.Div([ | |
| html.Button('▶ Play', id='play-button', n_clicks=0, | |
| style={'marginRight': 10, 'padding': '10px 20px'}), | |
| html.Button('⏸ Pause', id='pause-button', n_clicks=0, | |
| style={'marginRight': 10, 'padding': '10px 20px'}), | |
| html.Button('🔄 Reset', id='reset-button', n_clicks=0, | |
| style={'padding': '10px 20px'}), | |
| ], style={'marginTop': 15, 'marginBottom': 15}), | |
| html.Label("Animation Speed:", style={'fontWeight': 'bold', 'marginTop': 15}), | |
| dcc.Slider( | |
| id='speed-slider', | |
| min=1, max=100, step=1, value=10, | |
| marks={1: 'Slow', 50: 'Medium', 100: 'Fast'}, | |
| tooltip={"placement": "bottom", "always_visible": True} | |
| ), | |
| html.Label("Frame:", style={'fontWeight': 'bold', 'marginTop': 25}), | |
| dcc.Slider( | |
| id='frame-slider', | |
| min=0, max=100, step=1, value=0, | |
| marks={}, | |
| tooltip={"placement": "bottom", "always_visible": True} | |
| ), | |
| html.Div(id='frame-info', style={'marginTop': 10, 'fontSize': 14, 'color': '#7f8c8d'}), | |
| # Hidden components for state management | |
| dcc.Interval(id='animation-interval', interval=100, disabled=True), | |
| dcc.Store(id='animation-state', data={'playing': False, 'frame': 0}), | |
| ], style={ | |
| 'width': '25%', | |
| 'padding': '20px', | |
| 'backgroundColor': '#f8f9fa', | |
| 'borderRadius': '10px', | |
| 'marginRight': '2%' | |
| }), | |
| # Right panel - Visualizations | |
| html.Div([ | |
| # Main animation display | |
| dcc.Graph(id='main-animation', style={'height': '450px'}), | |
| # Statistics and cache performance | |
| html.Div([ | |
| html.Div([ | |
| dcc.Graph(id='cache-performance', style={'height': '250px'}) | |
| ], style={'width': '48%', 'display': 'inline-block'}), | |
| html.Div([ | |
| html.Div(id='statistics-panel', style={ | |
| 'padding': '20px', | |
| 'backgroundColor': '#fff', | |
| 'borderRadius': '10px', | |
| 'height': '250px', | |
| 'overflowY': 'auto' | |
| }) | |
| ], style={'width': '48%', 'display': 'inline-block', 'marginLeft': '4%'}), | |
| ]), | |
| # Heatmaps | |
| html.H3("📊 Access Frequency Heatmaps", style={'marginTop': 30, 'color': '#34495e'}), | |
| dcc.Graph(id='heatmaps', style={'height': '300px'}), | |
| ], style={'width': '73%'}), | |
| ], style={'display': 'flex', 'marginTop': 20}), | |
| ], style={'padding': '20px', 'backgroundColor': '#ffffff', 'fontFamily': 'Arial, sans-serif'}) | |
| # Callback to initialize/update simulation | |
| def update_simulation(n, block_size, loop_order, blocked, reset_clicks): | |
| """Initialize or update the simulation when parameters change.""" | |
| # Create simulator | |
| sim = GEMMSimulator(n=n, block_size=block_size if blocked else None) | |
| tracks = sim.simulate(loop_order, blocked=blocked) | |
| # Run cache simulation | |
| # Use a smaller cache to better demonstrate the impact of different loop orderings | |
| # For n=16: 3 matrices * 16*16*8 bytes = 6KB total | |
| # Using 4KB cache creates realistic cache pressure | |
| cache_size = max(4096, n * n * 8) # Scale cache size with matrix size, minimum 4KB | |
| cache = CacheSimulator(cache_size=cache_size, line_size=64, associativity=4) | |
| cache_stats = cache.simulate_accesses(tracks, matrix_size=n) | |
| # Store in global state | |
| simulation_data['simulator'] = sim | |
| simulation_data['tracks'] = tracks | |
| simulation_data['cache_stats'] = cache_stats | |
| simulation_data['heatmaps'] = sim.get_heatmap_data() | |
| # Prepare slider marks | |
| max_frame = len(tracks) - 1 | |
| step = max(1, max_frame // 10) | |
| marks = {i: str(i) for i in range(0, max_frame + 1, step)} | |
| # Statistics panel content | |
| stats_content = html.Div([ | |
| html.H4("📈 Statistics", style={'color': '#34495e', 'marginBottom': 15}), | |
| html.P([html.Strong("Matrix Size: "), f"{n} × {n}"]), | |
| html.P([html.Strong("Block Size: "), f"{block_size}" if blocked else "Unblocked"]), | |
| html.P([html.Strong("Loop Order: "), loop_order.upper()]), | |
| html.Hr(), | |
| html.P([html.Strong("Total Operations: "), f"{len(tracks):,}"]), | |
| html.P([html.Strong("Memory Accesses: "), f"{cache_stats['total_accesses']:,}"]), | |
| html.P([html.Strong("Cache Size: "), f"{cache_size // 1024}KB"], style={'fontSize': 12, 'color': '#7f8c8d'}), | |
| html.Hr(), | |
| html.H5("🎯 Cache Performance", style={'color': '#34495e', 'marginTop': 15}), | |
| html.P([html.Strong("Hit Rate: "), | |
| html.Span(f"{cache_stats['hit_rate']:.2f}%", | |
| style={'color': '#27ae60' if cache_stats['hit_rate'] > 80 else '#e74c3c', | |
| 'fontSize': 18, 'fontWeight': 'bold'})]), | |
| html.P([html.Strong("Hits: "), f"{cache_stats['hits']:,}"]), | |
| html.P([html.Strong("Misses: "), f"{cache_stats['misses']:,}"]), | |
| ]) | |
| return max_frame, marks, stats_content, {'playing': False, 'frame': 0} | |
| # Callback to handle play/pause | |
| def control_animation(play_clicks, pause_clicks, state): | |
| """Control play/pause of animation.""" | |
| ctx = callback_context | |
| if not ctx.triggered: | |
| return True, state | |
| button_id = ctx.triggered[0]['prop_id'].split('.')[0] | |
| if button_id == 'play-button': | |
| state['playing'] = True | |
| return False, state # Enable interval | |
| else: # pause-button | |
| state['playing'] = False | |
| return True, state # Disable interval | |
| # Callback to update frame | |
| def update_frame(n_intervals, slider_value, reset_clicks, state, max_frame, speed): | |
| """Update the current frame.""" | |
| ctx = callback_context | |
| if not ctx.triggered: | |
| return 0, state | |
| trigger_id = ctx.triggered[0]['prop_id'].split('.')[0] | |
| if trigger_id == 'reset-button': | |
| state['frame'] = 0 | |
| state['playing'] = False | |
| return 0, state | |
| if trigger_id == 'frame-slider': | |
| state['frame'] = slider_value | |
| return slider_value, state | |
| if trigger_id == 'animation-interval' and state['playing']: | |
| # Calculate frame increment based on speed | |
| increment = max(1, speed // 10) | |
| new_frame = state['frame'] + increment | |
| if new_frame >= max_frame: | |
| new_frame = max_frame | |
| state['playing'] = False | |
| state['frame'] = new_frame | |
| return new_frame, state | |
| return state['frame'], state | |
| # Callback to update visualizations | |
| def update_visualizations(frame, n): | |
| """Update all visualizations based on current frame.""" | |
| tracks = simulation_data.get('tracks', []) | |
| cache_stats = simulation_data.get('cache_stats', {}) | |
| heatmaps = simulation_data.get('heatmaps', {}) | |
| if not tracks: | |
| # Return empty figures | |
| empty_fig = go.Figure() | |
| return empty_fig, "", empty_fig, empty_fig | |
| frame = min(frame, len(tracks) - 1) | |
| current_access = tracks[frame] | |
| # Main animation figure | |
| main_fig = make_subplots( | |
| rows=1, cols=3, | |
| subplot_titles=('Matrix A', 'Matrix B', 'Matrix C'), | |
| horizontal_spacing=0.1 | |
| ) | |
| # Create grid for each matrix | |
| for idx, (matrix_name, pos) in enumerate([('A', current_access[0]), | |
| ('B', current_access[1]), | |
| ('C', current_access[2])]): | |
| # Background grid | |
| grid = np.zeros((n, n)) | |
| grid[pos[0], pos[1]] = 1 | |
| main_fig.add_trace( | |
| go.Heatmap( | |
| z=grid, | |
| colorscale=[[0, 'white'], [1, COLORS[matrix_name]]], | |
| showscale=False, | |
| hovertemplate=f'{matrix_name}[%{{y}}, %{{x}}]<extra></extra>' | |
| ), | |
| row=1, col=idx + 1 | |
| ) | |
| main_fig.update_xaxes(showticklabels=False, showgrid=True, gridcolor='lightgray') | |
| main_fig.update_yaxes(showticklabels=False, showgrid=True, gridcolor='lightgray', autorange='reversed') | |
| main_fig.update_layout( | |
| title=f"Memory Access Pattern (Frame {frame}/{len(tracks)-1})", | |
| height=450, | |
| showlegend=False | |
| ) | |
| # Frame info | |
| a_pos, b_pos, c_pos = current_access | |
| frame_info = html.Div([ | |
| html.P(f"Frame: {frame} / {len(tracks)-1}", style={'fontWeight': 'bold'}), | |
| html.P(f"A[{a_pos[0]},{a_pos[1]}] • B[{b_pos[0]},{b_pos[1]}] • C[{c_pos[0]},{c_pos[1]}]") | |
| ]) | |
| # Cache performance chart | |
| hit_rate_history = cache_stats.get('hit_rate_history', []) | |
| cache_fig = go.Figure() | |
| cache_fig.add_trace(go.Scatter( | |
| y=hit_rate_history, | |
| mode='lines', | |
| line=dict(color='#27ae60', width=2), | |
| fill='tozeroy', | |
| name='Hit Rate' | |
| )) | |
| cache_fig.update_layout( | |
| title="Cache Hit Rate Over Time", | |
| xaxis_title="Access (×100)", | |
| yaxis_title="Hit Rate (%)", | |
| yaxis=dict(range=[0, 100]), | |
| height=250, | |
| margin=dict(l=40, r=20, t=40, b=40) | |
| ) | |
| # Heatmaps | |
| heatmap_fig = make_subplots( | |
| rows=1, cols=3, | |
| subplot_titles=('Matrix A Access Frequency', | |
| 'Matrix B Access Frequency', | |
| 'Matrix C Access Frequency'), | |
| horizontal_spacing=0.1 | |
| ) | |
| for idx, matrix_name in enumerate(['A', 'B', 'C']): | |
| heatmap_data = heatmaps.get(matrix_name, np.zeros((n, n))) | |
| heatmap_fig.add_trace( | |
| go.Heatmap( | |
| z=heatmap_data, | |
| colorscale='YlOrRd', | |
| showscale=(idx == 2), | |
| hovertemplate=f'{matrix_name}[%{{y}}, %{{x}}]: %{{z}} accesses<extra></extra>' | |
| ), | |
| row=1, col=idx + 1 | |
| ) | |
| heatmap_fig.update_xaxes(showticklabels=False) | |
| heatmap_fig.update_yaxes(showticklabels=False, autorange='reversed') | |
| heatmap_fig.update_layout(height=300) | |
| return main_fig, frame_info, cache_fig, heatmap_fig | |
| if __name__ == '__main__': | |
| # Get port from environment variable (for Hugging Face Spaces) | |
| port = int(os.environ.get('PORT', 7860)) | |
| print("🚀 Starting GEMM Interactive Visualizer...") | |
| print(f"📱 Server running on port {port}") | |
| print("Press Ctrl+C to stop the server\n") | |
| # Run with production settings for Hugging Face Spaces | |
| app.run(debug=False, host='0.0.0.0', port=port) | |