virtual_gpu/examples/test_basic_rendering.py

"""
Basic Rendering Test Example

This example demonstrates basic rendering capabilities of the vGPU,
including creating framebuffers, drawing primitives, and displaying output.
"""

import asyncio
import sys
import os

# Add parent directory to path to import vGPU modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from vgpu import VirtualGPU, TaskType
from vram import VRAM
from driver import GPUDriver, CommandType
from render import Renderer
from ai import AIAccelerator
from display import DisplayManager, DisplayMode
from shader import ShaderManager
import numpy as np


async def basic_rendering_test():
    """Test basic rendering functionality of the vGPU."""
    print("Starting Basic Rendering Test...")
    print("=" * 50)
    
    # Initialize components
    print("Initializing vGPU components...")
    
    # Create VRAM (1GB for testing)
    vram = VRAM(memory_size_gb=1)
    print(f"✓ VRAM initialized: {vram.get_stats()['total_memory_gb']} GB")
    
    # Create renderer
    renderer = Renderer(vram)
    print("✓ Renderer initialized")
    
    # Create AI accelerator
    ai_accelerator = AIAccelerator(vram)
    print("✓ AI Accelerator initialized")
    
    # Create vGPU with 800 SMs and 50,000 cores
    vgpu = VirtualGPU(num_sms=800, total_cores=50000)
    vgpu.set_modules(vram, renderer, ai_accelerator, None)
    print(f"✓ vGPU initialized: {vgpu.num_sms} SMs, {vgpu.total_cores} cores")
    
    # Create driver
    driver = GPUDriver(vgpu)
    vgpu.driver = driver
    print("✓ GPU Driver initialized")
    
    # Create display manager
    display_manager = DisplayManager(vram)
    print("✓ Display Manager initialized")
    
    # Create shader manager
    shader_manager = ShaderManager()
    print("✓ Shader Manager initialized")
    
    print("\nComponent initialization complete!")
    print("=" * 50)
    
    # Create framebuffer
    print("\nCreating framebuffer...")
    fb_id = driver.create_framebuffer(800, 600, 3, "main_framebuffer")
    await driver.process_commands()
    
    framebuffer = vram.get_framebuffer("main_framebuffer")
    if framebuffer:
        print(f"✓ Framebuffer created: {framebuffer.width}x{framebuffer.height}")
        display_manager.set_active_framebuffer("main_framebuffer")
    else:
        print("✗ Failed to create framebuffer")
        return
        
    # Add file display for saving frames
    file_display = display_manager.add_display("file", DisplayMode.FILE, 
                                              output_dir="./test_output")
    print("✓ File display added")
    
    # Test 1: Clear screen
    print("\nTest 1: Clear screen")
    driver.clear_screen((64, 128, 255))  # Blue background
    await driver.process_commands()
    await vgpu.tick()
    await display_manager.update_displays()
    print("✓ Screen cleared to blue")
    
    # Test 2: Draw rectangles
    print("\nTest 2: Draw rectangles")
    driver.draw_rectangle(100, 100, 200, 150, (255, 0, 0))  # Red rectangle
    driver.draw_rectangle(300, 200, 150, 100, (0, 255, 0))  # Green rectangle
    driver.draw_rectangle(500, 50, 100, 200, (255, 255, 0))  # Yellow rectangle
    await driver.process_commands()
    await vgpu.tick()
    await display_manager.update_displays()
    print("✓ Rectangles drawn")
    
    # Test 3: Draw with shader
    print("\nTest 3: Apply shader effects")
    
    # Set grayscale shader
    grayscale_shader = shader_manager.get_shader("grayscale")
    renderer.set_shader(grayscale_shader)
    
    # Draw some shapes with shader
    driver.draw_rectangle(50, 400, 100, 100, (255, 0, 255))  # Magenta (will be grayscale)
    driver.draw_rectangle(200, 450, 80, 80, (0, 255, 255))   # Cyan (will be grayscale)
    await driver.process_commands()
    await vgpu.tick()
    await display_manager.update_displays()
    print("✓ Shapes drawn with grayscale shader")
    
    # Test 4: Remove shader and draw more
    print("\nTest 4: Remove shader and draw more shapes")
    renderer.set_shader(None)  # Remove shader
    
    # Draw circles using line approximation
    center_x, center_y = 400, 400
    radius = 50
    for angle in range(0, 360, 10):
        x1 = center_x + int(radius * np.cos(np.radians(angle)))
        y1 = center_y + int(radius * np.sin(np.radians(angle)))
        x2 = center_x + int(radius * np.cos(np.radians(angle + 10)))
        y2 = center_y + int(radius * np.sin(np.radians(angle + 10)))
        
        # Draw line segment
        renderer.draw_line("main_framebuffer", x1, y1, x2, y2, (255, 255, 255))
        
    await display_manager.update_displays()
    print("✓ Circle drawn using line segments")
    
    # Test 5: Performance test
    print("\nTest 5: Performance test - drawing many pixels")
    start_time = asyncio.get_event_loop().time()
    
    # Draw a pattern of pixels
    for y in range(500, 550):
        for x in range(600, 700):
            color_r = (x - 600) * 255 // 100
            color_g = (y - 500) * 255 // 50
            color_b = 128
            driver.draw_pixel(x, y, (color_r, color_g, color_b))
            
    await driver.process_commands()
    
    # Process multiple ticks to handle all the pixel drawing tasks
    for _ in range(10):
        await vgpu.tick()
        
    end_time = asyncio.get_event_loop().time()
    await display_manager.update_displays()
    
    pixels_drawn = 50 * 100  # 5000 pixels
    time_taken = end_time - start_time
    print(f"✓ Drew {pixels_drawn:,} pixels in {time_taken:.3f}s "
          f"({pixels_drawn/time_taken:.0f} pixels/sec)")
    
    # Test 6: AI operations
    print("\nTest 6: AI matrix operations")
    
    # Create test matrices
    matrix_a = np.random.rand(100, 50).astype(np.float32)
    matrix_b = np.random.rand(50, 75).astype(np.float32)
    
    # Load matrices into vGPU
    a_id = ai_accelerator.load_matrix(matrix_a, "test_matrix_a")
    b_id = ai_accelerator.load_matrix(matrix_b, "test_matrix_b")
    
    # Perform matrix multiplication
    result_id = ai_accelerator.matrix_multiply(a_id, b_id, "result_matrix")
    
    if result_id:
        result = ai_accelerator.get_matrix(result_id)
        expected = np.dot(matrix_a, matrix_b)
        
        if np.allclose(result, expected, rtol=1e-5):
            print(f"✓ Matrix multiplication successful: {matrix_a.shape} x {matrix_b.shape} = {result.shape}")
        else:
            print("✗ Matrix multiplication result incorrect")
    else:
        print("✗ Matrix multiplication failed")
    
    # Final statistics
    print("\n" + "=" * 50)
    print("FINAL STATISTICS")
    print("=" * 50)
    
    # vGPU stats
    vgpu_stats = vgpu.get_stats()
    print(f"vGPU Statistics:")
    print(f"  Clock cycles: {vgpu_stats['clock_cycle']:,}")
    print(f"  Tasks processed: {vgpu_stats['total_tasks_processed']:,}")
    print(f"  Busy SMs: {vgpu_stats['busy_sms']}/{vgpu_stats['total_sms']}")
    
    # Renderer stats
    render_stats = renderer.get_stats()
    print(f"\nRenderer Statistics:")
    print(f"  Pixels drawn: {render_stats['pixels_drawn']:,}")
    print(f"  Draw calls: {render_stats['draw_calls']:,}")
    print(f"  Render time: {render_stats['total_render_time']:.3f}s")
    print(f"  Pixels/second: {render_stats['pixels_per_second']:,.0f}")
    
    # VRAM stats
    vram_stats = vram.get_stats()
    print(f"\nVRAM Statistics:")
    print(f"  Total memory: {vram_stats['total_memory_gb']:.1f} GB")
    print(f"  Utilization: {vram_stats['utilization_percent']:.2f}%")
    print(f"  Framebuffers: {vram_stats['framebuffers_count']}")
    print(f"  Textures: {vram_stats['textures_count']}")
    
    # AI stats
    ai_stats = ai_accelerator.get_stats()
    print(f"\nAI Accelerator Statistics:")
    print(f"  Operations: {ai_stats['operations_performed']}")
    print(f"  FLOPs performed: {ai_stats['flops_performed']:,}")
    print(f"  FLOPs/second: {ai_stats['flops_per_second']:,.0f}")
    print(f"  Matrices in memory: {ai_stats['matrices_in_memory']}")
    
    # Display stats
    display_stats = display_manager.get_stats()
    print(f"\nDisplay Statistics:")
    print(f"  Frames displayed: {display_stats['frames_displayed']}")
    print(f"  Display time: {display_stats['total_display_time']:.3f}s")
    print(f"  Active displays: {display_stats['active_displays']}")
    
    # Driver stats
    driver_stats = driver.get_stats()
    print(f"\nDriver Statistics:")
    print(f"  Commands processed: {driver_stats['commands_processed']}")
    print(f"  Commands failed: {driver_stats['commands_failed']}")
    
    print("\n" + "=" * 50)
    print("Basic Rendering Test Complete!")
    print("Check ./test_output/ for saved frame images.")
    print("=" * 50)


if __name__ == "__main__":
    # Run the basic rendering test
    asyncio.run(basic_rendering_test())