526 lines
16 KiB
Rust
526 lines
16 KiB
Rust
/*!
|
|
* Integration tests for RustyTorch++ runtime stack
|
|
*
|
|
* These tests exercise the complete integration between allocator, device,
|
|
* streams, scheduler, and kernel launch systems using real GPU resources
|
|
* when available.
|
|
*/
|
|
|
|
use rtx_runtime::*;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
/// Environment-based backend selection for testing
|
|
fn get_test_backend() -> Backend {
|
|
match std::env::var("RTX_BACKEND").as_deref() {
|
|
Ok("cuda") => Backend::Cuda,
|
|
Ok("rocm") => Backend::ROCm,
|
|
Ok("metal") => Backend::Metal,
|
|
_ => Backend::CPU,
|
|
}
|
|
}
|
|
|
|
/// Get device count from environment or default to 1
|
|
fn get_device_count() -> usize {
|
|
std::env::var("RTX_DEVICE_COUNT")
|
|
.ok()
|
|
.and_then(|s| s.parse().ok())
|
|
.unwrap_or(1)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_full_stack_initialization() -> Result<()> {
|
|
// Test complete stack initialization with real GPU if available
|
|
let backend = get_test_backend();
|
|
let device_count = get_device_count();
|
|
|
|
// Create allocator with production settings
|
|
let allocator_config = AllocatorConfig {
|
|
arena_size_bytes: 256 * 1024 * 1024, // 256MB
|
|
max_fragmentation_percent: 15,
|
|
enable_statistics: true,
|
|
};
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
// Initialize devices
|
|
let mut devices = Vec::new();
|
|
for device_id in 0..device_count {
|
|
let device = Device::new(backend, device_id, allocator.clone())?;
|
|
assert_eq!(device.backend(), backend);
|
|
assert_eq!(device.device_id(), device_id);
|
|
devices.push(device);
|
|
}
|
|
|
|
// Create scheduler with all devices
|
|
let scheduler = Scheduler::new(devices)?;
|
|
assert_eq!(scheduler.device_count(), device_count);
|
|
|
|
println!("Successfully initialized full stack with {} {} devices",
|
|
device_count,
|
|
match backend {
|
|
Backend::CUDA => "CUDA",
|
|
Backend::ROCm => "ROCm",
|
|
Backend::Metal => "Metal",
|
|
Backend::CPU => "CPU",
|
|
}
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_memory_lifecycle_integration() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let allocator_config = AllocatorConfig {
|
|
arena_size_bytes: 128 * 1024 * 1024, // 128MB
|
|
max_fragmentation_percent: 20,
|
|
enable_statistics: true,
|
|
};
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
let device = Device::new(backend, 0, allocator.clone())?;
|
|
let scheduler = Scheduler::new(vec![device])?;
|
|
|
|
// Test various allocation sizes across multiple size classes
|
|
let test_sizes = vec![
|
|
1024, // 1KB
|
|
64 * 1024, // 64KB
|
|
1024 * 1024, // 1MB
|
|
16 * 1024 * 1024, // 16MB
|
|
];
|
|
|
|
let mut handles = Vec::new();
|
|
|
|
// Allocate memory blocks
|
|
for &size in &test_sizes {
|
|
let handle = allocator.allocate(size)?;
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Check allocator statistics
|
|
let stats = allocator.statistics();
|
|
assert!(stats.total_allocated_bytes > 0);
|
|
assert_eq!(stats.active_allocations, test_sizes.len());
|
|
|
|
// Free half the blocks
|
|
for handle in handles.drain(0..test_sizes.len() / 2) {
|
|
allocator.deallocate(handle)?;
|
|
}
|
|
|
|
// Verify statistics updated correctly
|
|
let stats_after = allocator.statistics();
|
|
assert_eq!(stats_after.active_allocations, test_sizes.len() / 2);
|
|
assert!(stats_after.total_freed_bytes > 0);
|
|
|
|
// Free remaining blocks
|
|
for handle in handles {
|
|
allocator.deallocate(handle)?;
|
|
}
|
|
|
|
let final_stats = allocator.statistics();
|
|
assert_eq!(final_stats.active_allocations, 0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_stream_scheduling_integration() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let allocator_config = AllocatorConfig::default();
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
let device = Device::new(backend, 0, allocator.clone())?;
|
|
let mut scheduler = Scheduler::new(vec![device])?;
|
|
|
|
// Create a complex dependency graph with multiple operations
|
|
let op1 = ScheduledOperation::MemoryTransfer {
|
|
src: 0x1000,
|
|
dst: 0x2000,
|
|
size: 1024 * 1024,
|
|
device_id: 0
|
|
};
|
|
let op2 = ScheduledOperation::KernelLaunch {
|
|
kernel_id: 1,
|
|
device_id: 0
|
|
};
|
|
let op3 = ScheduledOperation::MemoryTransfer {
|
|
src: 0x2000,
|
|
dst: 0x3000,
|
|
size: 1024 * 1024,
|
|
device_id: 0
|
|
};
|
|
let op4 = ScheduledOperation::KernelLaunch {
|
|
kernel_id: 2,
|
|
device_id: 0
|
|
};
|
|
|
|
// Schedule operations with dependencies: op1 -> op2 -> op3 -> op4
|
|
let task1 = scheduler.schedule(op1, Vec::new()).await?;
|
|
let task2 = scheduler.schedule(op2, vec![task1]).await?;
|
|
let task3 = scheduler.schedule(op3, vec![task2]).await?;
|
|
let task4 = scheduler.schedule(op4, vec![task3]).await?;
|
|
|
|
// Wait for all operations to complete
|
|
scheduler.wait_for_completion(task4).await?;
|
|
|
|
// Verify scheduler statistics
|
|
let stats = scheduler.statistics();
|
|
assert_eq!(stats.total_scheduled, 4);
|
|
assert_eq!(stats.completed_tasks, 4);
|
|
assert!(stats.average_scheduling_time_ns < 1_000_000); // < 1ms
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_kernel_launch_integration() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let allocator_config = AllocatorConfig::default();
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
let device = Device::new(backend, 0, allocator.clone())?;
|
|
let mut scheduler = Scheduler::new(vec![device])?;
|
|
|
|
// Create a kernel launch system
|
|
let mut kernel_system = KernelLaunchSystem::new();
|
|
|
|
// Mock PTX source for a simple vector add kernel
|
|
let ptx_source = r#"
|
|
.version 7.0
|
|
.target sm_120
|
|
.address_size 64
|
|
|
|
.visible .entry vector_add(
|
|
.param .u64 vector_add_param_0,
|
|
.param .u64 vector_add_param_1,
|
|
.param .u64 vector_add_param_2,
|
|
.param .u32 vector_add_param_3
|
|
)
|
|
{
|
|
.reg .u32 %tid;
|
|
.reg .u64 %rd<4>;
|
|
|
|
mov.u32 %tid, %tid.x;
|
|
|
|
ld.param.u64 %rd0, [vector_add_param_0];
|
|
ld.param.u64 %rd1, [vector_add_param_1];
|
|
ld.param.u64 %rd2, [vector_add_param_2];
|
|
|
|
ret;
|
|
}
|
|
"#;
|
|
|
|
// Load and compile the kernel
|
|
let kernel_id = kernel_system.load_kernel("vector_add", ptx_source)?;
|
|
|
|
// Allocate input/output arrays
|
|
let n = 1024;
|
|
let size_bytes = n * std::mem::size_of::<f32>();
|
|
|
|
let input_a = allocator.allocate(size_bytes)?;
|
|
let input_b = allocator.allocate(size_bytes)?;
|
|
let output = allocator.allocate(size_bytes)?;
|
|
|
|
// Configure kernel launch parameters
|
|
let params = KernelParams {
|
|
grid_dim: (16, 1, 1), // 16 blocks
|
|
block_dim: (64, 1, 1), // 64 threads per block
|
|
shared_memory_bytes: 0,
|
|
parameters: vec![
|
|
input_a.device_ptr() as u64,
|
|
input_b.device_ptr() as u64,
|
|
output.device_ptr() as u64,
|
|
n as u32,
|
|
],
|
|
};
|
|
|
|
// Launch kernel through scheduler
|
|
let launch_op = ScheduledOperation::KernelLaunch {
|
|
kernel_id,
|
|
device_id: 0
|
|
};
|
|
|
|
let start_time = std::time::Instant::now();
|
|
let task = scheduler.schedule(launch_op, Vec::new()).await?;
|
|
scheduler.wait_for_completion(task).await?;
|
|
let elapsed = start_time.elapsed();
|
|
|
|
// Verify kernel execution completed
|
|
println!("Kernel execution completed in {:?}", elapsed);
|
|
assert!(elapsed < Duration::from_millis(100)); // Should complete quickly
|
|
|
|
// Check kernel system statistics
|
|
let stats = kernel_system.statistics();
|
|
assert_eq!(stats.kernels_loaded, 1);
|
|
assert_eq!(stats.total_launches, 1);
|
|
assert!(stats.average_launch_time_ns > 0);
|
|
|
|
// Clean up allocated memory
|
|
allocator.deallocate(input_a)?;
|
|
allocator.deallocate(input_b)?;
|
|
allocator.deallocate(output)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_device_coordination() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let device_count = get_device_count().max(1); // Use at least 1 device
|
|
|
|
if device_count == 1 {
|
|
println!("Skipping multi-device test - only 1 device available");
|
|
return Ok(());
|
|
}
|
|
|
|
let allocator_config = AllocatorConfig::default();
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
// Create multiple devices
|
|
let mut devices = Vec::new();
|
|
for device_id in 0..device_count {
|
|
let device = Device::new(backend, device_id, allocator.clone())?;
|
|
devices.push(device);
|
|
}
|
|
|
|
let mut scheduler = Scheduler::new(devices)?;
|
|
|
|
// Schedule operations across different devices
|
|
let mut tasks = Vec::new();
|
|
for device_id in 0..device_count {
|
|
let op = ScheduledOperation::MemoryTransfer {
|
|
src: 0x1000 + (device_id as u64 * 0x1000),
|
|
dst: 0x2000 + (device_id as u64 * 0x1000),
|
|
size: 1024 * 1024,
|
|
device_id,
|
|
};
|
|
|
|
let task = scheduler.schedule(op, Vec::new()).await?;
|
|
tasks.push(task);
|
|
}
|
|
|
|
// Wait for all cross-device operations to complete
|
|
for task in tasks {
|
|
scheduler.wait_for_completion(task).await?;
|
|
}
|
|
|
|
let stats = scheduler.statistics();
|
|
assert_eq!(stats.completed_tasks, device_count);
|
|
|
|
println!("Successfully coordinated operations across {} devices", device_count);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_error_recovery_integration() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let allocator_config = AllocatorConfig {
|
|
arena_size_bytes: 64 * 1024 * 1024, // Smaller arena to trigger OOM
|
|
max_fragmentation_percent: 15,
|
|
enable_statistics: true,
|
|
};
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
let device = Device::new(backend, 0, allocator.clone())?;
|
|
let mut scheduler = Scheduler::new(vec![device])?;
|
|
|
|
// Try to allocate more memory than available to test error handling
|
|
let large_size = 128 * 1024 * 1024; // 128MB (larger than 64MB arena)
|
|
|
|
let result = allocator.allocate(large_size);
|
|
match result {
|
|
Ok(handle) => {
|
|
// If allocation succeeded, clean up
|
|
allocator.deallocate(handle)?;
|
|
println!("Large allocation succeeded (arena may be larger than expected)");
|
|
}
|
|
Err(Error::OutOfMemory) => {
|
|
println!("Expected OOM error occurred - error recovery working correctly");
|
|
}
|
|
Err(e) => {
|
|
panic!("Unexpected error type: {:?}", e);
|
|
}
|
|
}
|
|
|
|
// Verify that smaller allocations still work after OOM
|
|
let small_handle = allocator.allocate(1024)?;
|
|
allocator.deallocate(small_handle)?;
|
|
|
|
// Test scheduler with invalid device ID
|
|
let invalid_op = ScheduledOperation::KernelLaunch {
|
|
kernel_id: 999,
|
|
device_id: 999 // Invalid device
|
|
};
|
|
|
|
let schedule_result = scheduler.schedule(invalid_op, Vec::new()).await;
|
|
assert!(schedule_result.is_err(), "Expected error for invalid device ID");
|
|
|
|
println!("Error recovery integration tests passed");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_performance_baseline() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let allocator_config = AllocatorConfig::default();
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
let device = Device::new(backend, 0, allocator.clone())?;
|
|
let mut scheduler = Scheduler::new(vec![device])?;
|
|
|
|
// Performance test: rapid allocation/deallocation
|
|
let iterations = 1000;
|
|
let allocation_size = 64 * 1024; // 64KB
|
|
|
|
let start_time = std::time::Instant::now();
|
|
|
|
for _ in 0..iterations {
|
|
let handle = allocator.allocate(allocation_size)?;
|
|
allocator.deallocate(handle)?;
|
|
}
|
|
|
|
let elapsed = start_time.elapsed();
|
|
let avg_per_cycle = elapsed / iterations;
|
|
|
|
println!("Allocation/deallocation performance: {:?} per cycle", avg_per_cycle);
|
|
|
|
// Should be sub-millisecond per allocation cycle
|
|
assert!(avg_per_cycle < Duration::from_micros(100),
|
|
"Allocation cycle too slow: {:?}", avg_per_cycle);
|
|
|
|
// Performance test: scheduling overhead
|
|
let scheduling_iterations = 100;
|
|
let start_time = std::time::Instant::now();
|
|
|
|
for i in 0..scheduling_iterations {
|
|
let op = ScheduledOperation::MemoryTransfer {
|
|
src: 0x1000 + (i as u64 * 1024),
|
|
dst: 0x2000 + (i as u64 * 1024),
|
|
size: 1024,
|
|
device_id: 0,
|
|
};
|
|
|
|
let task = scheduler.schedule(op, Vec::new()).await?;
|
|
scheduler.wait_for_completion(task).await?;
|
|
}
|
|
|
|
let scheduling_elapsed = start_time.elapsed();
|
|
let avg_schedule_time = scheduling_elapsed / scheduling_iterations;
|
|
|
|
println!("Average scheduling time: {:?}", avg_schedule_time);
|
|
|
|
// Target: sub-microsecond scheduling overhead
|
|
assert!(avg_schedule_time < Duration::from_micros(10),
|
|
"Scheduling overhead too high: {:?}", avg_schedule_time);
|
|
|
|
let final_stats = scheduler.statistics();
|
|
assert!(final_stats.average_scheduling_time_ns < 10_000); // < 10μs
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Helper function to run a comprehensive system stress test
|
|
#[tokio::test]
|
|
async fn test_system_stress() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let allocator_config = AllocatorConfig {
|
|
arena_size_bytes: 512 * 1024 * 1024, // 512MB for stress test
|
|
max_fragmentation_percent: 20,
|
|
enable_statistics: true,
|
|
};
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
let device = Device::new(backend, 0, allocator.clone())?;
|
|
let mut scheduler = Scheduler::new(vec![device])?;
|
|
|
|
println!("Starting system stress test...");
|
|
|
|
// Create multiple concurrent workloads
|
|
let num_concurrent_tasks = 50;
|
|
let mut tasks = Vec::new();
|
|
|
|
for i in 0..num_concurrent_tasks {
|
|
let op = if i % 2 == 0 {
|
|
ScheduledOperation::MemoryTransfer {
|
|
src: 0x1000 + (i as u64 * 4096),
|
|
dst: 0x10000 + (i as u64 * 4096),
|
|
size: 4096 * (1 + i % 10), // Variable sizes
|
|
device_id: 0,
|
|
}
|
|
} else {
|
|
ScheduledOperation::KernelLaunch {
|
|
kernel_id: (i % 5) as u32, // Cycle through 5 kernel IDs
|
|
device_id: 0,
|
|
}
|
|
};
|
|
|
|
// Create dependency chains
|
|
let deps = if i > 0 && i % 10 == 0 {
|
|
vec![tasks[i - 1]] // Every 10th task depends on previous
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
let task = scheduler.schedule(op, deps).await?;
|
|
tasks.push(task);
|
|
}
|
|
|
|
// Wait for all tasks to complete
|
|
let start_wait = std::time::Instant::now();
|
|
for task in tasks {
|
|
scheduler.wait_for_completion(task).await?;
|
|
}
|
|
let wait_duration = start_wait.elapsed();
|
|
|
|
println!("Stress test completed in {:?}", wait_duration);
|
|
|
|
// Verify system statistics
|
|
let scheduler_stats = scheduler.statistics();
|
|
assert_eq!(scheduler_stats.completed_tasks, num_concurrent_tasks);
|
|
|
|
let allocator_stats = allocator.statistics();
|
|
println!("Final allocator stats: {:?}", allocator_stats);
|
|
|
|
// Should complete stress test within reasonable time
|
|
assert!(wait_duration < Duration::from_secs(30),
|
|
"Stress test took too long: {:?}", wait_duration);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test that verifies numerical determinism across runs
|
|
#[tokio::test]
|
|
async fn test_deterministic_execution() -> Result<()> {
|
|
let backend = get_test_backend();
|
|
let allocator_config = AllocatorConfig::default();
|
|
let allocator = Arc::new(Allocator::new(allocator_config)?);
|
|
|
|
let device = Device::new(backend, 0, allocator.clone())?;
|
|
let scheduler = Scheduler::new(vec![device])?;
|
|
|
|
// For now, just verify that identical operations produce identical results
|
|
// (This would be expanded with actual kernel execution and result comparison)
|
|
|
|
let op1 = ScheduledOperation::MemoryTransfer {
|
|
src: 0x1000,
|
|
dst: 0x2000,
|
|
size: 1024,
|
|
device_id: 0,
|
|
};
|
|
let op2 = ScheduledOperation::MemoryTransfer {
|
|
src: 0x1000,
|
|
dst: 0x2000,
|
|
size: 1024,
|
|
device_id: 0,
|
|
};
|
|
|
|
// Both operations should behave identically
|
|
assert_eq!(op1, op2);
|
|
|
|
println!("Deterministic execution test passed (basic)");
|
|
|
|
Ok(())
|
|
}
|