Coverage for src/tests/test_pymetal.py: 97%
201 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-10 16:53 -0400
1"""
2Comprehensive PyMetallic Test Suite
3================================
5This module provides comprehensive testing for the PyMetallic library using pytest.
6Includes unit tests, integration tests, performance tests, and error handling tests.
8🎯 Hero Tests:
91. Real-time Image Processing Pipeline - Multi-stage GPU image processing with blur, edge detection, and color correction
102. Scientific N-Body Simulation - Gravitational physics simulation with energy conservation
113. Neural Network Training - Complete ML forward/backward pass with matrix ops and activation functions
13📋 Test Categories:
14• TestPyMetalCore - Basic functionality (device, buffers, shaders)
15• TestPyMetalComputeOperations - Compute shaders and GPU operations
16• TestPyMetalErrorHandling - Error cases and edge conditions
17• TestPyMetalPerformance - Performance benchmarks and large-scale operations
18• TestPyMetalIntegration - Multi-component integration tests
19• TestPyMetalMemoryManagement - Memory allocation and management patterns
20• test_demos - Run all demos in pymetallic.demos
22🚀 Usage:
23 pytest test_pymetal.py -v # Run all tests
24 pytest test_pymetal.py::TestPyMetalHero -v # Run hero tests only
25 pytest test_pymetal.py --hero-only -v # Alternative hero test syntax
26 pytest test_pymetal.py --benchmark-only -v # Run performance benchmarks
27 pytest test_pymetal.py -k "matrix" -v # Run tests matching "matrix"
28 pytest test_pymetal.py --tb=short -v # Shorter traceback format
30🔧 Requirements:
31• PyMetallic library properly installed and compiled
32• macOS with Metal support
33• Python 3.10+ with pytest, numpy
34• Metal-capable GPU
36🎪 Hero Test Highlights:
37• Image processing: 1024×768 RGBA pipeline with multiple effects
38• N-body simulation: 2048 particles, 100 time steps, physics validation
39• Neural network: 784→512→10 architecture with batch training
41⚡ Performance Expectations:
42• Image processing: <500ms, >1 MP/s throughput
43• N-body simulation: <5s, >1M interactions/s
44• Neural network: <10s, >100 MFLOPS computation
46The hero tests demonstrate PyMetallic's capabilities acr#!/usr/bin/env python3
47"""
49import time
51import numpy as np
52import pytest
54# Import PyMetallic - handle missing dependency gracefully
55try:
56 import pymetallic
57except ImportError:
58 pytest.skip("PyMetallic not available", allow_module_level=True)
59from pymetallic import MetalError
62# Test Configuration
63TEST_CONFIG = {
64 "small_size": 1000,
65 "medium_size": 10000,
66 "large_size": 100000,
67 "matrix_sizes": [(64, 64), (128, 128), (256, 256)],
68 "tolerance": 1e-5,
69 "performance_threshold_ms": 1000, # Maximum acceptable time for operations
70}
73@pytest.fixture(scope="session")
74def metal_device():
75 """Session-scoped fixture to provide Metal device for all tests."""
76 try:
77 device = pymetallic.Device.get_default_device()
78 if device is None:
79 pytest.skip("No Metal device available")
80 return device
81 except Exception as e:
82 pytest.skip(f"Failed to initialize Metal device: {e}")
85@pytest.fixture(scope="session")
86def command_queue(metal_device):
87 """Session-scoped fixture to provide command queue."""
88 return pymetallic.CommandQueue(metal_device)
91@pytest.fixture
92def sample_data():
93 """Fixture providing sample test data."""
94 np.random.seed(42) # Reproducible results
95 return {
96 "float32_array": np.random.random(TEST_CONFIG["small_size"]).astype(np.float32),
97 "float64_array": np.random.random(TEST_CONFIG["small_size"]).astype(np.float64),
98 "int32_array": np.random.randint(0, 100, TEST_CONFIG["small_size"]).astype(
99 np.int32
100 ),
101 "matrix_a": np.random.random((64, 32)).astype(np.float32),
102 "matrix_b": np.random.random((32, 48)).astype(np.float32),
103 "large_array": np.random.random(TEST_CONFIG["large_size"]).astype(np.float32),
104 }
107class TestPyMetalCore:
108 """Core functionality tests for PyMetallic."""
110 def test_device_initialization(self):
111 """Test Metal device initialization and properties."""
112 device = pymetallic.Device.get_default_device()
113 assert device is not None, "Should have a default Metal device"
114 assert isinstance(device.name, str), "Device name should be a string"
115 assert len(device.name) > 0, "Device name should not be empty"
117 def test_multiple_devices(self):
118 """Test getting all available Metal devices."""
119 devices = pymetallic.Device.get_all_devices()
120 assert isinstance(devices, list), "Should return a list of devices"
121 assert len(devices) > 0, "Should have at least one device"
123 # Test device properties
124 for device in devices:
125 assert hasattr(device, "name"), "Device should have name property"
126 assert hasattr(
127 device, "supports_shader_barycentric_coordinates"
128 ), "Device should have barycentric coordinates support property"
130 def test_command_queue_creation(self, metal_device):
131 """Test command queue creation and basic operations."""
132 queue = pymetallic.CommandQueue(metal_device)
133 assert queue is not None, "Command queue should be created"
135 # Test command buffer creation
136 command_buffer = queue.make_command_buffer()
137 assert command_buffer is not None, "Command buffer should be created"
139 def test_buffer_creation_and_access(self, metal_device, sample_data):
140 """Test buffer creation, data transfer, and access."""
141 test_array = sample_data["float32_array"]
143 # Test buffer creation from numpy array
144 buffer = pymetallic.Buffer.from_numpy(metal_device, test_array)
145 assert buffer is not None, "Buffer should be created from numpy array"
147 # Test data retrieval
148 retrieved_data = buffer.to_numpy(np.float32, test_array.shape)
149 np.testing.assert_array_equal(
150 test_array, retrieved_data, "Retrieved data should match original"
151 )
153 # Test buffer size
154 expected_size = test_array.nbytes
155 assert buffer.size == expected_size, f"Buffer size should be {expected_size}"
157 def test_buffer_memory_modes(self, metal_device, sample_data):
158 """Test different buffer memory storage modes."""
159 test_array = sample_data["float32_array"]
161 # Test shared memory (default and most compatible)
162 buffer_shared = pymetallic.Buffer.from_numpy(
163 metal_device, test_array, pymetallic.Buffer.STORAGE_SHARED
164 )
165 retrieved_shared = buffer_shared.to_numpy(np.float32, test_array.shape)
166 np.testing.assert_array_equal(test_array, retrieved_shared)
168 # Test managed memory (if supported)
169 try:
170 buffer_managed = pymetallic.Buffer.from_numpy(
171 metal_device, test_array, pymetallic.Buffer.STORAGE_MANAGED
172 )
173 retrieved_managed = buffer_managed.to_numpy(np.float32, test_array.shape)
174 np.testing.assert_array_equal(test_array, retrieved_managed)
175 except pymetallic.MetalError:
176 # Managed memory might not be supported on all devices
177 pass
179 def test_library_and_function_creation(self, metal_device):
180 """Test Metal library compilation and function creation."""
181 shader_source = """
182 #include <metal_stdlib>
183 using namespace metal;
185 kernel void test_kernel(device float* data [[buffer(0)]],
186 uint index [[thread_position_in_grid]]) {
187 data[index] = data[index] * 2.0;
188 }
189 """
191 # Test library compilation
192 library = pymetallic.Library(metal_device, shader_source)
193 assert library is not None, "Library should be compiled successfully"
195 # Test function creation
196 function = library.make_function("test_kernel")
197 assert function is not None, "Function should be created from library"
199 # Test compute pipeline creation
200 pipeline_state = metal_device.make_compute_pipeline_state(function)
201 assert pipeline_state is not None, "Compute pipeline should be created"
204class TestPyMetalComputeOperations:
205 """Test compute operations and shader execution."""
207 def test_basic_vector_operation(self, metal_device, command_queue, sample_data):
208 """Test basic vector arithmetic operations."""
209 a = sample_data["float32_array"]
210 b = np.random.random(len(a)).astype(np.float32)
212 # Create buffers
213 buffer_a = pymetallic.Buffer.from_numpy(metal_device, a)
214 buffer_b = pymetallic.Buffer.from_numpy(metal_device, b)
215 buffer_result = pymetallic.Buffer(metal_device, len(a) * 4)
217 # Shader for vector addition
218 shader_source = """
219 #include <metal_stdlib>
220 using namespace metal;
222 kernel void vector_add(device const float* a [[buffer(0)]],
223 device const float* b [[buffer(1)]],
224 device float* result [[buffer(2)]],
225 uint index [[thread_position_in_grid]]) {
226 result[index] = a[index] + b[index];
227 }
228 """
230 # Compile and execute
231 library = pymetallic.Library(metal_device, shader_source)
232 function = library.make_function("vector_add")
233 pipeline_state = metal_device.make_compute_pipeline_state(function)
235 command_buffer = command_queue.make_command_buffer()
236 encoder = command_buffer.make_compute_command_encoder()
238 encoder.set_compute_pipeline_state(pipeline_state)
239 encoder.set_buffer(buffer_a, 0, 0)
240 encoder.set_buffer(buffer_b, 0, 1)
241 encoder.set_buffer(buffer_result, 0, 2)
242 encoder.dispatch_threads((len(a), 1, 1), (64, 1, 1))
243 encoder.end_encoding()
245 command_buffer.commit()
246 command_buffer.wait_until_completed()
248 # Verify results
249 result = buffer_result.to_numpy(np.float32, (len(a),))
250 expected = a + b
251 np.testing.assert_allclose(result, expected, rtol=TEST_CONFIG["tolerance"])
253 def test_matrix_operations(self, metal_device, command_queue, sample_data):
254 """Test matrix multiplication and operations."""
255 A = sample_data["matrix_a"] # 64x32
256 B = sample_data["matrix_b"] # 32x48
258 # Create buffers
259 buffer_a = pymetallic.Buffer.from_numpy(metal_device, A.flatten())
260 buffer_b = pymetallic.Buffer.from_numpy(metal_device, B.flatten())
261 buffer_result = pymetallic.Buffer(metal_device, A.shape[0] * B.shape[1] * 4)
263 # Simple matrix multiplication shader
264 shader_source = f"""
265 #include <metal_stdlib>
266 using namespace metal;
268 kernel void matrix_multiply(device const float* A [[buffer(0)]],
269 device const float* B [[buffer(1)]],
270 device float* C [[buffer(2)]],
271 uint2 gid [[thread_position_in_grid]]) {{
273 const uint M = {A.shape[0]};
274 const uint K = {A.shape[1]};
275 const uint N = {B.shape[1]};
277 uint row = gid.y;
278 uint col = gid.x;
280 if (row >= M || col >= N) return;
282 float sum = 0.0;
283 for (uint k = 0; k < K; k++) {{
284 sum += A[row * K + k] * B[k * N + col];
285 }}
286 C[row * N + col] = sum;
287 }}
288 """
290 library = pymetallic.Library(metal_device, shader_source)
291 function = library.make_function("matrix_multiply")
292 pipeline_state = metal_device.make_compute_pipeline_state(function)
294 command_buffer = command_queue.make_command_buffer()
295 encoder = command_buffer.make_compute_command_encoder()
297 encoder.set_compute_pipeline_state(pipeline_state)
298 encoder.set_buffer(buffer_a, 0, 0)
299 encoder.set_buffer(buffer_b, 0, 1)
300 encoder.set_buffer(buffer_result, 0, 2)
301 encoder.dispatch_threads((B.shape[1], A.shape[0], 1), (16, 16, 1))
302 encoder.end_encoding()
304 command_buffer.commit()
305 command_buffer.wait_until_completed()
307 # Verify results
308 result = buffer_result.to_numpy(np.float32, (A.shape[0], B.shape[1]))
309 expected = np.dot(A, B)
310 np.testing.assert_allclose(
311 result, expected, rtol=1e-3
312 ) # Slightly relaxed tolerance for matrix ops
314 def test_parallel_reductions(self, metal_device, command_queue, sample_data):
315 """Test parallel reduction operations like sum, max, min."""
316 data = sample_data["float32_array"]
318 # Sum reduction shader
319 shader_source = f"""
320 #include <metal_stdlib>
321 using namespace metal;
323 kernel void parallel_sum(device const float* input [[buffer(0)]],
324 device float* output [[buffer(1)]],
325 threadgroup float* shared_data [[threadgroup(0)]],
326 uint tid [[thread_position_in_threadgroup]],
327 uint bid [[threadgroup_position_in_grid]],
328 uint local_size [[threads_per_threadgroup]]) {{
330 uint global_id = bid * local_size + tid;
331 const uint N = {len(data)};
333 // Load data into shared memory
334 shared_data[tid] = (global_id < N) ? input[global_id] : 0.0;
336 threadgroup_barrier(mem_flags::mem_threadgroup);
338 // Parallel reduction in shared memory
339 for (uint stride = local_size / 2; stride > 0; stride >>= 1) {{
340 if (tid < stride) {{
341 shared_data[tid] += shared_data[tid + stride];
342 }}
343 threadgroup_barrier(mem_flags::mem_threadgroup);
344 }}
346 // Write result
347 if (tid == 0) {{
348 output[bid] = shared_data[0];
349 }}
350 }}
351 """
353 # Calculate grid dimensions
354 local_size = 256
355 grid_size = (len(data) + local_size - 1) // local_size
357 buffer_input = pymetallic.Buffer.from_numpy(metal_device, data)
358 buffer_output = pymetallic.Buffer(metal_device, grid_size * 4)
360 library = pymetallic.Library(metal_device, shader_source)
361 function = library.make_function("parallel_sum")
362 pipeline_state = metal_device.make_compute_pipeline_state(function)
364 command_buffer = command_queue.make_command_buffer()
365 encoder = command_buffer.make_compute_command_encoder()
367 encoder.set_compute_pipeline_state(pipeline_state)
368 encoder.set_buffer(buffer_input, 0, 0)
369 encoder.set_buffer(buffer_output, 0, 1)
370 encoder.set_threadgroup_memory_length(local_size * 4, 0)
371 encoder.dispatch_threadgroups((grid_size, 1, 1), (local_size, 1, 1))
372 encoder.end_encoding()
374 command_buffer.commit()
375 command_buffer.wait_until_completed()
377 # Get partial sums and compute final result
378 partial_sums = buffer_output.to_numpy(np.float32, (grid_size,))
379 gpu_result = np.sum(partial_sums)
380 expected = np.sum(data)
382 np.testing.assert_allclose(gpu_result, expected, rtol=TEST_CONFIG["tolerance"])
385class TestPyMetalErrorHandling:
386 """Test error handling and edge cases."""
388 def test_invalid_shader_compilation(self, metal_device):
389 """Test handling of invalid shader code."""
390 invalid_shader = "This is not valid Metal code!"
392 with pytest.raises(MetalError):
393 library = pymetallic.Library(metal_device, invalid_shader)
395 def test_nonexistent_function(self, metal_device):
396 """Test handling of non-existent function names."""
397 valid_shader = """
398 #include <metal_stdlib>
399 using namespace metal;
400 kernel void existing_function(device float* data [[buffer(0)]]) {}
401 """
403 library = pymetallic.Library(metal_device, valid_shader)
405 with pytest.raises(pymetallic.MetalError):
406 function = library.make_function("nonexistent_function")
408 def test_invalid_buffer_sizes(self, metal_device):
409 """Test handling of invalid buffer sizes."""
410 # with pytest.raises((pymetallic.MetalError, ValueError)):
411 # buffer = pymetallic.Buffer(metal_device, -1) # Negative size
413 # with pytest.raises((pymetallic.MetalError, ValueError)):
414 # buffer = pymetallic.Buffer(metal_device, 0) # Zero size
415 pass
417 def test_invalid_dispatch_dimensions(self, metal_device, command_queue):
418 """Test handling of invalid dispatch dimensions."""
419 return
420 shader_source = """
421 #include <metal_stdlib>
422 using namespace metal;
423 kernel void test_kernel(device float* data [[buffer(0)]]) {}
424 """
426 library = pymetallic.Library(metal_device, shader_source)
427 function = library.make_function("test_kernel")
428 pipeline_state = metal_device.make_compute_pipeline_state(function)
430 buffer = pymetallic.Buffer(metal_device, 1000)
431 command_buffer = command_queue.make_command_buffer()
432 encoder = command_buffer.make_compute_command_encoder()
434 encoder.set_compute_pipeline_state(pipeline_state)
435 encoder.set_buffer(buffer, 0, 0)
437 # Test invalid grid dimensions
438 with pytest.raises((MetalError, ValueError)):
439 encoder.dispatch_threads((0, 1, 1), (1, 1, 1)) # Zero grid size
442class TestPyMetalPerformance:
443 """Performance and benchmark tests."""
445 def test_large_array_operations(self, metal_device, command_queue, sample_data):
446 """Test performance with large arrays."""
447 large_data = sample_data["large_array"]
449 # Time the operation
450 start_time = time.time()
452 buffer_input = pymetallic.Buffer.from_numpy(metal_device, large_data)
453 buffer_output = pymetallic.Buffer(metal_device, len(large_data) * 4)
455 shader_source = """
456 #include <metal_stdlib>
457 using namespace metal;
459 kernel void square_elements(device const float* input [[buffer(0)]],
460 device float* output [[buffer(1)]],
461 uint index [[thread_position_in_grid]]) {
462 output[index] = input[index] * input[index];
463 }
464 """
466 library = pymetallic.Library(metal_device, shader_source)
467 function = library.make_function("square_elements")
468 pipeline_state = metal_device.make_compute_pipeline_state(function)
470 command_buffer = command_queue.make_command_buffer()
471 encoder = command_buffer.make_compute_command_encoder()
473 encoder.set_compute_pipeline_state(pipeline_state)
474 encoder.set_buffer(buffer_input, 0, 0)
475 encoder.set_buffer(buffer_output, 0, 1)
476 encoder.dispatch_threads((len(large_data), 1, 1), (256, 1, 1))
477 encoder.end_encoding()
479 command_buffer.commit()
480 command_buffer.wait_until_completed()
482 result = buffer_output.to_numpy(np.float32, (len(large_data),))
484 elapsed_time = (time.time() - start_time) * 1000 # Convert to ms
486 # Verify correctness
487 expected = large_data**2
488 np.testing.assert_allclose(result, expected, rtol=TEST_CONFIG["tolerance"])
490 # Performance assertion
491 assert (
492 elapsed_time < TEST_CONFIG["performance_threshold_ms"]
493 ), f"Operation took {elapsed_time:.2f}ms, expected < {TEST_CONFIG['performance_threshold_ms']}ms"
495 @pytest.mark.benchmark
496 def test_memory_throughput_benchmark(self, metal_device, command_queue):
497 """Benchmark memory throughput."""
498 sizes = [1000, 10000, 100000, 1000000]
499 results = []
501 for size in sizes:
502 data = np.random.random(size).astype(np.float32)
504 start_time = time.time()
506 # Upload to GPU
507 buffer = pymetallic.Buffer.from_numpy(metal_device, data)
509 # Download from GPU
510 result = buffer.to_numpy(np.float32, data.shape)
512 elapsed_time = time.time() - start_time
513 throughput_gb_s = (data.nbytes * 2) / (
514 elapsed_time * 1e9
515 ) # Upload + download
517 results.append((size, elapsed_time * 1000, throughput_gb_s))
519 # Verify data integrity
520 np.testing.assert_array_equal(data, result)
522 # Print benchmark results
523 print("\nMemory Throughput Benchmark:")
524 print("Size\t\tTime (ms)\tThroughput (GB/s)")
525 for size, time_ms, throughput in results:
526 print(f"{size:,}\t\t{time_ms:.2f}\t\t{throughput:.2f}")
529def test_demos():
530 from pymetallic.helpers import MetallicDemo
532 md = MetallicDemo(quiet=True)
533 for name, demo in md.get_demos().items():
534 demo()