Spaces:
Runtime error
Runtime error
File size: 6,144 Bytes
684cc60 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
import asyncio
import time
import numpy as np
import os
from virtual_gpu.vram import VRAM
from virtual_gpu.virtual_ram import VirtualRAM
from virtual_gpu.vgpu import VirtualGPU
from virtual_gpu.driver import GPUDriver
from virtual_gpu.render import Renderer
from virtual_gpu.ai import AIAccelerator
from virtual_gpu.shader import ShaderManager
from virtual_gpu.display import DisplayManager, DisplayMode
async def test_virtual_ram_transfer():
print("Starting Virtual RAM to VRAM Transfer Test...")
print("=" * 50)
# 1. Initialize VirtualRAM(128) and VirtualVRAM(500)
print("Initializing VirtualRAM and VRAM...")
ram = VirtualRAM(capacity_gb=128)
vram = VRAM(memory_size_gb=500)
print("✓ VirtualRAM (128GB) and VRAM (500GB) initialized.")
# Initialize other vGPU components for a complete environment
renderer = Renderer(vram)
ai_accelerator = AIAccelerator(vram)
vgpu = VirtualGPU(num_sms=800, total_cores=50000)
vgpu.set_modules(vram, renderer, ai_accelerator, None)
driver = GPUDriver(vgpu)
vgpu.driver = driver
shader_manager = ShaderManager()
display_manager = DisplayManager(vram)
file_display = display_manager.add_display("file", DisplayMode.FILE, output_dir="./test_output_ram_transfer")
print("✓ Other vGPU components initialized.")
print("\n" + "=" * 50)
# 2. Allocate a 1MB tensor in RAM.
tensor_size_mb = 1 # Use 1MB for actual data allocation
tensor_size_bytes = tensor_size_mb * 1024 * 1024
print(f"Allocating a {tensor_size_mb}MB tensor in VirtualRAM...")
ram.allocate_block("my_ai_tensor", tensor_size_bytes, store_data=True) # Request real data
ram.print_info()
print("✓ 1MB tensor allocated with real data in VirtualRAM.")
print("\n" + "=" * 50)
# 3. Transfer that tensor into VRAM as
print(f"Transferring \'my_ai_tensor\' from VirtualRAM to VRAM as \'ai_tensor_buffer\'")
transfer_start_time = time.time()
vram_tensor_id = ram.transfer_to_vram("my_ai_tensor", vram, "ai_tensor_buffer")
transfer_end_time = time.time()
time_taken_to_transfer = transfer_end_time - transfer_start_time
if vram_tensor_id:
print(f"✓ Tensor transferred to VRAM. VRAM ID: {vram_tensor_id}")
print(f"Time taken to transfer: {time_taken_to_transfer:.4f} seconds")
else:
print("✗ Failed to transfer tensor to VRAM.")
return
print("\n" + "=" * 50)
# 4. Apply a real AI operation (matrix multiplication) on the VRAM buffer.
print("Applying a real AI operation (matrix multiplication) on the VRAM buffer...")
# Retrieve the transferred data from VRAM (now it should contain actual data)
vram_block_data = vram.get_texture(vram_tensor_id) # This should now return actual data
if vram_block_data is None:
print("✗ Could not retrieve actual tensor data from VRAM for AI operation.")
return
# Create a dummy matrix B with real data for multiplication
# Ensure dimensions are compatible for multiplication
# Let\'s assume the transferred data is a flattened 1MB (1024*1024 bytes) of uint8
# We\'ll reshape it to (1024, 1024) for a dummy matrix A
try:
matrix_a_reshaped = vram_block_data.reshape((1024, 1024)).astype(np.float32)
except ValueError:
print("✗ Could not reshape transferred data for matrix A. Data size might not be suitable.")
return
# Create a dummy matrix B (e.g., 1024x512) for multiplication
matrix_b_data = np.random.rand(1024, 512).astype(np.float32)
# Load these real matrices into AI Accelerator (which will then load them into VRAM)
ai_accelerator.set_vram(vram) # Ensure AI accelerator has VRAM reference
matrix_a_id = ai_accelerator.load_matrix(matrix_a_reshaped, "real_ai_matrix_a")
matrix_b_id = ai_accelerator.load_matrix(matrix_b_data, "real_ai_matrix_b")
# Perform real matrix multiplication
result_id = ai_accelerator.matrix_multiply(matrix_a_id, matrix_b_id, "real_ai_op_result")
if result_id:
result_matrix = ai_accelerator.get_matrix(result_id)
if result_matrix is not None:
print(f"✓ Real AI operation (matrix multiplication) completed. Result shape: {result_matrix.shape}")
# Optional: Verify a small part of the result
expected_result = np.dot(matrix_a_reshaped, matrix_b_data)
if np.allclose(result_matrix[:5,:5], expected_result[:5,:5]): # Compare a small section
print("✓ Result verification (partial) successful.")
else:
print("✗ Result verification (partial) failed.")
else:
print("✗ Could not retrieve result matrix from VRAM.")
else:
print("✗ Real AI operation failed.")
print("\n" + "=" * 50)
# 5. Output logs showing:
# RAM usage
# VRAM usage
# Time taken to transfer
# Operation success
print("FINAL LOGS AND STATISTICS")
print("=" * 50)
print("\nVirtual RAM Usage:")
ram.print_info()
print("\nVRAM Usage:")
vram_stats = vram.get_stats()
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/Buffers: {vram_stats["textures_count"]}')
print(f' Allocated bytes: {vram_stats["allocated_bytes"]:,} bytes')
print(f"\nTime taken to transfer tensor from RAM to VRAM: {time_taken_to_transfer:.4f} seconds")
success_status = "Success" if result_id else "Failed"
print(f"AI Operation Success: {success_status}")
print("\n" + "=" * 50)
print("Virtual RAM Transfer Test Complete!")
print("Check ./test_output_ram_transfer/ for any saved frame images (if applicable).")
print("=" * 50)
# Clean up (optional)
# ram.release_block("my_ai_tensor")
# vram.delete_texture(vram_tensor_id)
if __name__ == "__main__":
# Ensure the output directory exists
os.makedirs("./test_output_ram_transfer", exist_ok=True)
asyncio.run(test_virtual_ram_transfer())
|