Spaces:
Sleeping
Sleeping
File size: 4,900 Bytes
0a735c8 |
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 |
"""
Test for hyperrealistic multi-chip GPU system with full SM and tensor core realism,
using WebSocket-based storage for zero CPU usage.
"""
import time
import numpy as np
from gpu_arch import Chip, OpticalInterconnect
def test_multi_chip_gpu():
print("\n=== Multi-Chip GPU System with WebSocket Storage Test ===")
num_chips = 2 # Use 2 for realism, scale up as needed
num_sms = 4 # Use 4 for realism, scale up as needed
# Initialize WebSocket storage for all chips
from websocket_storage import WebSocketGPUStorage
storage = WebSocketGPUStorage()
if not storage.wait_for_connection():
raise RuntimeError("Could not connect to GPU storage server")
chips = [Chip(
chip_id=i,
num_sms=num_sms,
vram_size_gb=None # Use unlimited WebSocket storage
) for i in range(num_chips)]
print(f"Created {num_chips} chips with unlimited WebSocket storage, each with {num_sms} SMs.")
# Connect chips in a ring topology with optical interconnect
optical_link = OpticalInterconnect(bandwidth_tbps=800, latency_ns=1)
for i in range(num_chips):
chips[i].connect_chip(chips[(i+1)%num_chips], optical_link)
# Initialize shared WebSocket storage for cross-chip communication
for chip in chips:
chip_state = {
"chip_id": chip.chip_id,
"num_sms": num_sms,
"connected_chips": [(c.chip_id, "optical") for c in chip.connected_chips]
}
storage.store_state(f"chips/{chip.chip_id}", "config", chip_state)
# Run tensor core operations with WebSocket-backed storage
print("\n=== Testing WebSocket-backed Multi-Chip Operations ===")
# Create test matrices
matrix_a = [[1.0, 2.0], [3.0, 4.0]]
matrix_b = [[5.0, 6.0], [7.0, 8.0]]
for chip in chips:
print(f"\n--- Chip {chip.chip_id} ---")
# Store matrices in WebSocket storage for this chip
storage.store_tensor(f"chip_{chip.chip_id}/matrix_a", np.array(matrix_a))
storage.store_tensor(f"chip_{chip.chip_id}/matrix_b", np.array(matrix_b))
# Process using each SM
for sm_id in range(num_sms):
sm = chip.get_sm(sm_id)
# Load matrices from WebSocket storage
matrix_a_data = storage.load_tensor(f"chip_{chip.chip_id}/matrix_a")
matrix_b_data = storage.load_tensor(f"chip_{chip.chip_id}/matrix_b")
# Perform tensor core operation
result = sm.tensor_core_matmul(matrix_a_data.tolist(), matrix_b_data.tolist())
# Store result back in WebSocket storage
storage.store_tensor(f"chip_{chip.chip_id}/sm_{sm_id}/result", np.array(result))
print(f"SM {sm_id} tensor core matmul result: {result}")
# Test cross-chip communication
if len(chip.connected_chips) > 0:
next_chip, link = chip.connected_chips[0]
# Get result from this chip
result_data = storage.load_tensor(f"chip_{chip.chip_id}/sm_0/result")
# Transfer to next chip through optical link
transfer_id = f"transfer_chip_{chip.chip_id}_to_{next_chip.chip_id}"
storage.store_tensor(transfer_id, result_data)
print(f"Transferred result from Chip {chip.chip_id} to Chip {next_chip.chip_id} via {link.__class__.__name__}")
for i in range(len(sm.register_file)):
for j in range(len(sm.register_file[0])):
sm.register_file[i][j] = float(i + j)
for addr in range(sm.shared_mem.size):
sm.shared_mem.write(addr, float(addr % 10))
for addr in range(sm.global_mem.size_bytes if sm.global_mem else 0):
sm.global_mem.write(addr, float(addr % 100))
# Test tensor core matmul from registers
reg_result = sm.tensor_core_matmul_from_memory('register', 0, 'register', 0, (2,2), (2,2))
print(f"SM {sm.sm_id} tensor core matmul from registers: {reg_result}")
# Test tensor core matmul from shared memory
shared_result = sm.tensor_core_matmul_from_memory('shared', 0, 'shared', 0, (2,2), (2,2))
print(f"SM {sm.sm_id} tensor core matmul from shared memory: {shared_result}")
# Test tensor core matmul from global memory
global_result = sm.tensor_core_matmul_from_memory('global', 0, 'global', 0, (2,2), (2,2))
print(f"SM {sm.sm_id} tensor core matmul from global memory: {global_result}")
print("\n=== Multi-Chip GPU System Test Complete ===")
if __name__ == "__main__":
start = time.time()
test_multi_chip_gpu()
print(f"Test runtime: {time.time()-start:.3f} seconds")
|