Coverage for src/pymetallic/helpers.py: 73%
805 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-10 16:53 -0400
1#!/usr/bin/env python3
2"""
3PyMetallic Complete Demo and Documentation
4Comprehensive demonstration of the PyMetallic library capabilities
5"""
6import os
7import sys
8import time
9from typing import Any
11import numpy as np
13# Optional image display support
14_PIL_AVAILABLE = False
15try:
16 from PIL import Image
18 _PIL_AVAILABLE = True
19except ImportError:
20 pass
21try:
22 import pymetallic
23 from pymetallic import Device, Buffer
25except ImportError:
26 print("PyMetallic not available. Please build and install first:")
27 print(" make build && make install-dev")
28 sys.exit(1)
31class StopWatch:
32 def __init__(self):
33 self.elapsed_time = 0
34 self.clicks = 0
35 self._start = 0.0
37 def __enter__(self):
38 self._start = time.time()
40 def __exit__(self, exc_type, exc_val, exc_tb):
41 self.elapsed_time += time.time() - self._start
42 self.clicks += 1
44 def __str__(self):
45 if self.clicks == 0:
46 return "No clicks"
47 return f"⏱ Elapsed {self.elapsed_time:.3f} sec for {self.clicks} steps, average = {self.elapsed_time/self.clicks:.3f}"
50class AnimateIt:
51 _usable = False
53 def __init__(
54 self,
55 name="animation",
56 mode: str | None = None,
57 duration: int = 100,
58 seconds: float | None = None,
59 loop: bool = False,
60 out_path: str | None = None,
61 colorful: bool = False,
62 normalize: bool = True,
63 **kwargs,
64 ):
65 self._usable = _PIL_AVAILABLE
66 self.frames: list[Any] = []
67 self.times: list[float] = []
68 self.mode = mode
69 self.name = name
70 self.seconds = seconds
71 self.duration = duration
72 self.loops = 0 if loop else 1
73 self.path = os.path.join(out_path, f"{self.name}.gif") if out_path else None
74 self.save_args = kwargs
75 self.colorful = colorful
76 self.normalize = normalize
78 def colorize(self, frame: np.ndarray):
79 """
80 If frame is a 2D floating-point array, map it to RGB via a perceptual spectral colormap.
81 Otherwise, return frame unchanged.
82 """
83 arr = frame
84 if arr.ndim == 2 and np.issubdtype(arr.dtype, np.floating):
85 a = arr.astype(np.float32, copy=False)
87 # Normalize to [0, 1]
88 amin = float(np.nanmin(a))
89 amax = float(np.nanmax(a))
90 if self.normalize:
91 if amax > amin:
92 a = (a - amin) / (amax - amin)
93 else:
94 a = np.zeros_like(a, dtype=np.float32)
95 a = np.clip(a, 0.0, 1.0)
97 # Gentle gamma to emphasize mid-tones
98 a = a ** np.float32(0.9)
100 # Spectral mapping using HSV:
101 # Hue from blue (0.66) -> red (0.0), with saturation/value rising with intensity
102 h = (1.0 - a) * 0.66 # [0, 0.66]
103 s = np.clip(0.6 + 0.4 * a, 0.0, 1.0)
104 v = np.clip(0.75 + 0.25 * a, 0.0, 1.0)
106 # Vectorized HSV -> RGB
107 c = v * s
108 m = v - c
109 hp = (h * 6.0) % 6.0
110 x = c * (1.0 - np.abs((hp % 2.0) - 1.0))
112 # Initialize channels
113 r1 = np.zeros_like(a, dtype=np.float32)
114 g1 = np.zeros_like(a, dtype=np.float32)
115 b1 = np.zeros_like(a, dtype=np.float32)
117 # Masks for each sextant
118 m0 = (0.0 <= hp) & (hp < 1.0)
119 m1 = (1.0 <= hp) & (hp < 2.0)
120 m2 = (2.0 <= hp) & (hp < 3.0)
121 m3 = (3.0 <= hp) & (hp < 4.0)
122 m4 = (4.0 <= hp) & (hp < 5.0)
123 m5 = (5.0 <= hp) & (hp < 6.0)
125 r1[m0], g1[m0], b1[m0] = c[m0], x[m0], 0.0
126 r1[m1], g1[m1], b1[m1] = x[m1], c[m1], 0.0
127 r1[m2], g1[m2], b1[m2] = 0.0, c[m2], x[m2]
128 r1[m3], g1[m3], b1[m3] = 0.0, x[m3], c[m3]
129 r1[m4], g1[m4], b1[m4] = x[m4], 0.0, c[m4]
130 r1[m5], g1[m5], b1[m5] = c[m5], 0.0, x[m5]
132 rgb = np.stack([r1 + m, g1 + m, b1 + m], axis=-1)
133 return rgb
135 return frame
137 def add_frame(self, array: np.ndarray):
138 if self._usable:
139 self.times.append(time.time())
140 frame = np.clip(array, 0.0, 1.0)
141 if self.colorful:
142 frame = self.colorize(frame)
143 frame_image = Image.fromarray((frame * 255.0).astype(np.uint8))
144 self.frames.append(frame_image)
146 def save(self) -> str | None:
147 if self._usable and self.path:
148 duration = (
149 int(1000 * self.seconds / len(self.frames))
150 if self.seconds
151 else self.duration
152 )
153 self.frames[0].save(
154 self.path,
155 save_all=True,
156 append_images=self.frames[1:],
157 duration=duration,
158 loop=self.loops,
159 optimize=False,
160 **self.save_args,
161 )
162 print(f"🎞️ Saved {self.name} GIF to {self.path}")
163 return self.path
164 return None
166 def show(self):
167 path = self.save()
168 if path:
169 # Safari actually does pretty well at this
170 os.system(f"open {path} -a safari")
173def display_array(array, title=None, normalize=True):
174 """Display a 2D or 3D numpy array as an image if PIL is available.
176 - 2D arrays are shown as grayscale (L).
177 - 3D arrays with 3 or 4 channels are shown as RGB/RGBA.
178 - Values are normalized to 0..255 if normalize is True or dtype is not uint8.
179 Returns True if displayed, False otherwise.
180 """
181 if not _PIL_AVAILABLE:
182 return False
183 try:
184 arr = np.asarray(array)
185 if arr.ndim == 2:
186 a = arr
187 if normalize or a.dtype != np.uint8:
188 a = a.astype(np.float32)
189 amin = float(np.min(a))
190 amax = float(np.max(a))
191 if amax > amin:
192 a = (a - amin) / (amax - amin)
193 a = (a * 255.0).clip(0, 255).astype(np.uint8)
194 img = Image.fromarray(a, mode="L")
195 elif arr.ndim == 3 and arr.shape[2] in (3, 4):
196 a = arr
197 if a.dtype != np.uint8:
198 scale = 255.0 if normalize else 1.0
199 a = np.clip(a.astype(np.float32) * scale, 0, 255).astype(np.uint8)
200 mode = "RGB" if a.shape[2] == 3 else "RGBA"
201 img = Image.fromarray(a, mode=mode)
202 else:
203 return False
204 if title:
205 # TODO: render the title to the image
206 pass
207 img.show()
208 return True
209 except Exception:
210 return False
213class MetallicDemo:
214 """Complete demonstration of PyMetallic capabilities"""
216 def __init__(self, quiet: bool = False, out_path: str | None = None):
217 if not quiet:
218 print("🚀 PyMetallic Comprehensive Demo")
219 print("=" * 50)
220 self.device = None
221 self.quiet = quiet
222 self.out_path = out_path
223 global _PIL_AVAILABLE
224 if self.quiet:
225 AnimateIt._usable = _PIL_AVAILABLE = False
226 self.metal_ops = MetalMatrixOperations()
227 self.initialize_metal()
229 def print(self, *args, **kwargs):
230 if not self.quiet:
231 print(*args, **kwargs)
233 def initialize_metal(self):
234 """Initialize Metal and display system information"""
235 self.print("\n📱 Metal System Information")
236 self.print("-" * 30)
238 try:
239 # Get all devices
240 devices = Device.get_all_devices()
241 self.print(f"Available Metal devices: {len(devices)}")
243 for i, device in enumerate(devices):
244 self.print(f" Device {i}: {device.name}")
245 self.print(
246 f" Supports barycentric coordinates: {device.supports_shader_barycentric_coordinates()}"
247 )
249 # Use default device
250 self.device = Device.get_default_device()
251 self.print(f"\n✅ Using device: {self.device.name}")
253 except Exception as e:
254 self.print(f"❌ Failed to initialize Metal: {e}")
255 sys.exit(1)
257 def demo_basic_compute(self):
258 """Demonstrate basic compute operations"""
259 self.print("\n🔢 Basic Compute Operations")
260 self.print("-" * 30)
262 # Simple vector addition
263 self.print("Vector Addition:")
264 size = 10000
265 a = np.random.random(size).astype(np.float32)
266 b = np.random.random(size).astype(np.float32)
268 start_time = time.time()
270 # Create Metal resources
271 device = self.device
272 queue = device.make_command_queue()
273 buffer_a = device.make_buffer_from_numpy(a)
274 buffer_b = device.make_buffer_from_numpy(b)
275 buffer_result = device.make_buffer(size * 4)
277 # Compile shader
278 shader_source = """
279 #include <metal_stdlib>
280 using namespace metal;
282 kernel void vector_add(device float* a [[buffer(0)]],
283 device float* b [[buffer(1)]],
284 device float* result [[buffer(2)]],
285 uint index [[thread_position_in_grid]]) {
286 result[index] = a[index] + b[index];
287 }
288 """
290 library = pymetallic.Library(self.device, shader_source)
291 function = library.make_function("vector_add")
292 pipeline_state = pymetallic.ComputePipelineState(self.device, function)
294 # Execute
295 command_buffer = queue.make_command_buffer()
296 encoder = command_buffer.make_compute_command_encoder()
298 encoder.set_compute_pipeline_state(pipeline_state)
299 encoder.set_buffer(buffer_a, 0, 0)
300 encoder.set_buffer(buffer_b, 0, 1)
301 encoder.set_buffer(buffer_result, 0, 2)
303 encoder.dispatch_threads((size, 1, 1), (64, 1, 1))
304 encoder.end_encoding()
306 command_buffer.commit()
307 command_buffer.wait_until_completed()
309 metal_time = (time.time() - start_time) * 1000
310 result = buffer_result.to_numpy(np.float32, (size,))
312 # Verify
313 expected = a + b
314 is_correct = np.allclose(result, expected, rtol=1e-5)
316 self.print(f" Size: {size:,} elements")
317 self.print(f" Time: {metal_time:.2f}ms")
318 self.print(f" Correct: {'✅' if is_correct else '❌'}")
320 # Compare with NumPy
321 start_time = time.time()
322 numpy_result = a + b
323 numpy_time = (time.time() - start_time) * 1000
325 speedup = numpy_time / metal_time if metal_time > 0 else 0
326 self.print(f" vs NumPy: {numpy_time:.2f}ms (speedup: {speedup:.1f}x)")
328 def demo_matrix_operations(self):
329 """Demonstrate matrix operations"""
330 self.print("\n🧮 Matrix Operations")
331 self.print("-" * 30)
333 metal_ops = MetalMatrixOperations(self.device)
335 # Matrix multiplication
336 self.print("Matrix Multiplication:")
337 sizes = [(64, 32, 48), (128, 64, 96), (256, 128, 192)]
339 for m, k, n in sizes:
340 A = np.random.random((m, k)).astype(np.float32)
341 B = np.random.random((k, n)).astype(np.float32)
343 # Metal computation
344 start_time = time.time()
345 C_metal = metal_ops.matrix_multiply(A, B)
346 metal_time = (time.time() - start_time) * 1000
348 # NumPy comparison
349 start_time = time.time()
350 C_numpy = np.dot(A, B)
351 numpy_time = (time.time() - start_time) * 1000
353 is_correct = np.allclose(C_metal, C_numpy, rtol=1e-4)
354 speedup = numpy_time / metal_time if metal_time > 0 else 0
356 self.print(
357 f" {m}×{k} × {k}×{n}: Metal {metal_time:.1f}ms, "
358 f"NumPy {numpy_time:.1f}ms, speedup {speedup:.1f}x {'✅' if is_correct else '❌'}"
359 )
361 # Vector operations
362 self.print("\nVector Operations:")
363 vec_size = 500000
364 x = np.random.random(vec_size).astype(np.float32)
365 y = np.random.random(vec_size).astype(np.float32)
367 operations = [("add", lambda a, b: a + b), ("multiply", lambda a, b: a * b)]
369 for op_name, numpy_op in operations:
370 start_time = time.time()
371 metal_result = metal_ops.vector_operations(x, y, op_name)
372 metal_time = (time.time() - start_time) * 1000
374 start_time = time.time()
375 numpy_result = numpy_op(x, y)
376 numpy_time = (time.time() - start_time) * 1000
378 is_correct = np.allclose(metal_result, numpy_result, rtol=1e-5)
379 speedup = numpy_time / metal_time if metal_time > 0 else 0
381 self.print(
382 f" {op_name.capitalize()}: Metal {metal_time:.1f}ms, "
383 f"NumPy {numpy_time:.1f}ms, speedup {speedup:.1f}x {'✅' if is_correct else '❌'}"
384 )
386 def demo_advanced_features(self):
387 """Demonstrate advanced Metal features"""
388 self.print("\n🔬 Advanced Features")
389 self.print("-" * 30)
391 # Multi-dimensional dispatch
392 self.print("2D Grid Computation:")
394 shader_source = """
395 #include <metal_stdlib>
396 using namespace metal;
398 kernel void mandelbrot(device float* output [[buffer(0)]],
399 constant uint& width [[buffer(1)]],
400 constant uint& height [[buffer(2)]],
401 constant uint& max_iterations [[buffer(3)]],
402 uint2 gid [[thread_position_in_grid]]) {
404 if (gid.x >= width || gid.y >= height) return;
406 float x = (float(gid.x) / float(width)) * 3.5 - 2.5;
407 float y = (float(gid.y) / float(height)) * 2.0 - 1.0;
409 float zx = 0.0, zy = 0.0;
410 uint iter = 0;
412 while (iter < max_iterations && (zx*zx + zy*zy) < 4.0) {
413 float tmp = zx*zx - zy*zy + x;
414 zy = 2.0*zx*zy + y;
415 zx = tmp;
416 iter++;
417 }
419 output[gid.y * width + gid.x] = float(iter) / float(max_iterations);
420 }
421 """
423 width, height = 512, 512
424 max_iterations = 100
425 device = self.device
426 queue = device.make_command_queue()
427 buffer_output = device.make_buffer(width * height * 4)
428 buffer_width = device.make_buffer_from_numpy(np.array([width], dtype=np.uint32))
429 buffer_height = device.make_buffer_from_numpy(
430 np.array([height], dtype=np.uint32)
431 )
432 buffer_max_iter = device.make_buffer_from_numpy(
433 np.array([max_iterations], dtype=np.uint32)
434 )
436 library = pymetallic.Library(self.device, shader_source)
437 function = library.make_function("mandelbrot")
438 pipeline_state = pymetallic.ComputePipelineState(self.device, function)
440 start_time = time.time()
442 command_buffer = queue.make_command_buffer()
443 encoder = command_buffer.make_compute_command_encoder()
445 encoder.set_compute_pipeline_state(pipeline_state)
446 encoder.set_buffer(buffer_output, 0, 0)
447 encoder.set_buffer(buffer_width, 0, 1)
448 encoder.set_buffer(buffer_height, 0, 2)
449 encoder.set_buffer(buffer_max_iter, 0, 3)
451 encoder.dispatch_threads((width, height, 1), (16, 16, 1))
452 encoder.end_encoding()
454 command_buffer.commit()
455 command_buffer.wait_until_completed()
457 computation_time = (time.time() - start_time) * 1000
458 result = buffer_output.to_numpy(np.float32, (height, width))
460 self.print(f" Mandelbrot set: {width}×{height} in {computation_time:.1f}ms")
461 self.print(f" Generated {width*height:,} pixels")
462 self.print(
463 f" Performance: {(width*height*max_iterations)/(computation_time*1000)/1e9:.2f} GOP/s"
464 )
465 # Display image if PIL is available
466 display_array(result, title=f"Mandelbrot {width}x{height}")
468 def demo_memory_patterns(self):
469 """Demonstrate different memory access patterns"""
470 self.print("\n💾 Memory Access Patterns")
471 self.print("-" * 30)
472 device = self.device
473 # Test different buffer storage modes
474 test_data = np.random.random(10000).astype(np.float32)
476 storage_modes = [
477 (pymetallic.Buffer.STORAGE_SHARED, "Shared"),
478 (pymetallic.Buffer.STORAGE_MANAGED, "Managed"),
479 # These tests do not work for Private memory
480 # On macOS, MTLStorageModePrivate buffers are not CPU-accessible.
481 # Calling to_numpy on such buffers will not work and may return invalid data or crash.
482 # For CPU readback from private buffers, you need a blit/compute copy into a shared/managed buffer.
483 # (pymetallic.Buffer.STORAGE_PRIVATE, "Private"),
484 ]
486 for mode, name in storage_modes:
487 try:
488 start_time = time.time()
489 buffer = device.make_buffer_from_numpy(test_data, mode)
490 retrieved = buffer.to_numpy(np.float32, test_data.shape)
491 access_time = (time.time() - start_time) * 1000
493 is_correct = np.array_equal(test_data, retrieved)
494 self.print(
495 f" {name} memory: {access_time:.2f}ms {'✅' if is_correct else '❌'}"
496 )
497 except Exception as e:
498 self.print(f" {name} memory: Not supported ({e})")
500 def demo_image_processing(self):
501 """
502 HERO: Real-time Image Processing Pipeline
504 Simulates a complete image processing pipeline with multiple effects:
505 - Gaussian blur
506 - Edge detection
507 - Color correction
509 Tests compute shaders, 2D dispatch, texture-like operations, and pipeline chaining.
510 """
511 device = self.device
512 command_queue = device.make_command_queue()
513 self.print("\n🎯 HERO: Image Processing Pipeline")
514 animation = AnimateIt("image_processing", duration=2, out_path=self.out_path)
515 # Simulate a high-resolution image
516 width, height = 1024, 768
517 channels = 4 # RGBA
518 image_size = width * height * channels
520 # Generate test image data (RGBA format)
521 np.random.seed(123)
522 original_image = np.random.random((height, width, channels)).astype(np.float32)
523 colors = [[1, 1, 1, 1], [1, 0, 0, 1], [0, 1, 0, 1], [0, 0, 1, 1]]
524 nstripes = len(colors)
525 wstripe = width // nstripes
526 for i, color in enumerate(colors):
527 l = i * wstripe + 50
528 r = l + wstripe - 100
529 original_image[300:500, l:r, :] = color
530 animation.add_frame(original_image)
531 # Multi-stage image processing shader
532 processing_shader = f"""
533 #include <metal_stdlib>
534 using namespace metal;
536 kernel void image_processing_pipeline(device const float* input [[buffer(0)]],
537 device float* output [[buffer(1)]],
538 device float* temp_buffer [[buffer(2)]],
539 constant uint& width [[buffer(3)]],
540 constant uint& height [[buffer(4)]],
541 uint2 gid [[thread_position_in_grid]]) {{
543 if (gid.x >= width || gid.y >= height) return;
545 const uint idx = gid.y * width + gid.x;
546 const uint channels = 4;
547 const uint pixel_idx = idx * channels;
549 // Stage 1: Gaussian blur (simplified 3x3 kernel)
550 float4 blurred = float4(0.0);
551 float gaussian_kernel[9] = {{0.0625, 0.125, 0.0625, 0.125, 0.25, 0.125, 0.0625, 0.125, 0.0625}};
553 for (int dy = -1; dy <= 1; dy++) {{
554 for (int dx = -1; dx <= 1; dx++) {{
555 int nx = int(gid.x) + dx;
556 int ny = int(gid.y) + dy;
558 if (nx >= 0 && nx < int(width) && ny >= 0 && ny < int(height)) {{
559 uint neighbor_idx = (ny * int(width) + nx) * channels;
560 float weight = gaussian_kernel[(dy + 1) * 3 + (dx + 1)];
562 blurred.r += input[neighbor_idx + 0] * weight;
563 blurred.g += input[neighbor_idx + 1] * weight;
564 blurred.b += input[neighbor_idx + 2] * weight;
565 blurred.a += input[neighbor_idx + 3] * weight;
566 }}
567 }}
568 }}
570 // Stage 2: Edge detection (Sobel operator)
571 float4 edge_x = float4(0.0);
572 float4 edge_y = float4(0.0);
574 float sobel_x[9] = {{-1, 0, 1, -2, 0, 2, -1, 0, 1}};
575 float sobel_y[9] = {{-1, -2, -1, 0, 0, 0, 1, 2, 1}};
577 for (int dy = -1; dy <= 1; dy++) {{
578 for (int dx = -1; dx <= 1; dx++) {{
579 int nx = int(gid.x) + dx;
580 int ny = int(gid.y) + dy;
582 if (nx >= 0 && nx < int(width) && ny >= 0 && ny < int(height)) {{
583 uint neighbor_idx = (ny * int(width) + nx) * channels;
584 int kernel_idx = (dy + 1) * 3 + (dx + 1);
586 float wx = sobel_x[kernel_idx];
587 float wy = sobel_y[kernel_idx];
589 edge_x += float4(input[neighbor_idx + 0], input[neighbor_idx + 1],
590 input[neighbor_idx + 2], input[neighbor_idx + 3]) * wx;
591 edge_y += float4(input[neighbor_idx + 0], input[neighbor_idx + 1],
592 input[neighbor_idx + 2], input[neighbor_idx + 3]) * wy;
593 }}
594 }}
595 }}
597 float4 edge_magnitude = sqrt(edge_x * edge_x + edge_y * edge_y);
599 // Stage 3: Color correction and final composition
600 float4 final_color = blurred * 0.7 + edge_magnitude * 0.3;
602 // Apply gamma correction
603 final_color = pow(final_color, float4(0.8));
605 // Clamp values
606 final_color = clamp(final_color, 0.0, 1.0);
608 // Write result
609 output[pixel_idx + 0] = final_color.r;
610 output[pixel_idx + 1] = final_color.g;
611 output[pixel_idx + 2] = final_color.b;
612 output[pixel_idx + 3] = final_color.a;
613 }}
614 """
616 # Create buffers
617 input_buffer = device.make_buffer_from_numpy(original_image.flatten())
618 output_buffer = device.make_buffer(image_size * 4)
619 temp_buffer = device.make_buffer(image_size * 4)
620 width_buffer = device.make_buffer_from_numpy(np.array([width], dtype=np.uint32))
621 height_buffer = device.make_buffer_from_numpy(
622 np.array([height], dtype=np.uint32)
623 )
625 # Compile and execute
626 start_time = time.time()
628 library = device.make_library(processing_shader)
629 function = library.make_function("image_processing_pipeline")
630 pipeline_state = device.make_compute_pipeline_state(function)
632 command_buffer = command_queue.make_command_buffer()
633 encoder = command_buffer.make_compute_command_encoder()
635 encoder.set_compute_pipeline_state(pipeline_state)
636 encoder.set_buffer(input_buffer, 0, 0)
637 encoder.set_buffer(output_buffer, 0, 1)
638 encoder.set_buffer(temp_buffer, 0, 2)
639 encoder.set_buffer(width_buffer, 0, 3)
640 encoder.set_buffer(height_buffer, 0, 4)
642 # Use 2D dispatch for image processing
643 encoder.dispatch_threads((width, height, 1), (16, 16, 1))
644 encoder.end_encoding()
646 command_buffer.commit()
647 command_buffer.wait_until_completed()
649 processing_time = (time.time() - start_time) * 1000
651 # Retrieve and validate results
652 processed_image = output_buffer.to_numpy(np.float32, original_image.shape)
653 animation.add_frame(processed_image)
654 # Validation checks
655 assert (
656 processed_image.shape == original_image.shape
657 ), "Output shape should match input"
658 assert np.all(processed_image >= 0.0) and np.all(
659 processed_image <= 1.0
660 ), "Processed image values should be in [0, 1] range"
661 assert not np.array_equal(
662 processed_image, original_image
663 ), "Processed image should be different from original"
665 # Performance metrics
666 pixels_processed = width * height
667 megapixels_per_second = (pixels_processed / 1e6) / (processing_time / 1000)
669 self.print(f"✅ Processed {width}×{height} image in {processing_time:.1f}ms")
670 self.print(f"📊 Performance: {megapixels_per_second:.1f} Megapixels/second")
671 self.print(
672 f"🎨 Pipeline stages: Gaussian blur → Edge detection → Color correction"
673 )
674 animation.show()
675 # Performance assertion for hero test
676 assert (
677 processing_time < 500
678 ), f"Image processing should complete in <500ms, took {processing_time:.1f}ms"
679 assert (
680 megapixels_per_second > 1.0
681 ), f"Should process >1 MP/s, achieved {megapixels_per_second:.1f} MP/s"
683 def demo_cellular_automata(
684 self,
685 width: int = 512,
686 height: int = 512,
687 steps: int = 200,
688 seed_probability: float = 0.05,
689 threadgroup: tuple = (16, 16, 1),
690 ) -> np.ndarray:
691 """
692 HERO: Cellular Automata (Conway's Game of Life)
693 ------------------------------------------------
694 Runs a GPU-accelerated Conway's Game of Life simulation on a 2D grid.
696 Returns the final state as a (height, width) uint8 NumPy array.
697 """
698 """
699 """
700 sw = StopWatch()
701 device = self.device
702 queue = device.make_command_queue()
704 animation = AnimateIt(
705 "cellular", duration=5, pallette=2, out_path=self.out_path
706 )
707 self.print("\n🎯 HERO: Cellular Automata")
709 # Prepare initial random state (0 or 1), packed as uint8
710 rng = np.random.default_rng(1234)
711 init = (rng.random((height, width)) < seed_probability).astype(np.uint8)
712 animation.add_frame(init)
713 buf_a = device.make_buffer_from_numpy(init, storage_mode=Buffer.STORAGE_SHARED)
714 buf_b = device.make_buffer(init.size) # bytes; uint8 per cell
715 # Params buffer: [width, height] as uint32
716 params = np.array([np.uint32(width), np.uint32(height)], dtype=np.uint32)
717 buf_params = device.make_buffer_from_numpy(
718 params, storage_mode=Buffer.STORAGE_SHARED
719 )
721 # Metal kernel for one Life step (with wrap-around)
722 source = f"""
723 #include <metal_stdlib>
724 using namespace metal;
726 struct Params {{
727 uint width;
728 uint height;
729 }};
731 inline uint wrap_int(int v, int m) {{
732 int r = v % m;
733 return (uint)(r < 0 ? r + m : r);
734 }}
736 kernel void life_step(const device uchar* in_state [[buffer(0)]],
737 device uchar* out_state [[buffer(1)]],
738 const device Params* p [[buffer(2)]],
739 uint2 gid [[thread_position_in_grid]]) {{
741 uint W = p->width;
742 uint H = p->height;
743 if (gid.x >= W || gid.y >= H) return;
745 int x = (int)gid.x;
746 int y = (int)gid.y;
748 int count = 0;
749 // 8-neighborhood
750 for (int dy = -1; dy <= 1; ++dy) {{
751 for (int dx = -1; dx <= 1; ++dx) {{
752 if (dx == 0 && dy == 0) continue;
753 uint nx = wrap_int(x + dx, (int)W);
754 uint ny = wrap_int(y + dy, (int)H);
755 uint nidx = ny * W + nx;
756 count += in_state[nidx] > 0 ? 1 : 0;
757 }}
758 }}
760 uint idx = (uint)y * W + (uint)x;
761 bool alive = in_state[idx] > 0;
762 bool next_alive = (alive && (count == 2 || count == 3)) || (!alive && count == 3);
763 out_state[idx] = next_alive ? (uchar)1 : (uchar)0;
764 }}
765 """
767 lib = device.make_library(source)
768 fn = lib.make_function("life_step")
769 pso = device.make_compute_pipeline_state(fn)
771 curr, nxt = buf_a, buf_b
772 for _ in range(int(steps)):
773 with sw:
774 cb = queue.make_command_buffer()
775 enc = cb.make_compute_command_encoder()
776 enc.set_compute_pipeline_state(pso)
777 enc.set_buffer(curr, 0, 0)
778 enc.set_buffer(nxt, 0, 1)
779 enc.set_buffer(buf_params, 0, 2)
780 enc.dispatch_threads((width, height, 1), threadgroup)
781 enc.end_encoding()
782 cb.commit()
783 cb.wait_until_completed()
784 animation.add_frame(nxt.to_numpy(init.dtype, init.shape))
785 curr, nxt = nxt, curr
787 # Copy back to host
788 out = curr.to_numpy(np.uint8, (height * width,))
789 animation.show()
790 self.print(sw)
791 return out.reshape((height, width))
793 def demo_fluid_dynamics(
794 self,
795 width: int = 512,
796 height: int = 512,
797 steps: int = 100,
798 dt: float = 0.1,
799 visc: float = 0.0001,
800 threadgroup: tuple = (16, 16, 1),
801 ) -> np.ndarray:
802 """
803 HERO: 2D Stable Fluids / CFD Demo
804 ---------------------------------
805 Simulates a simple 2D incompressible fluid using semi-Lagrangian advection and a Jacobi pressure solve.
806 Returns a (height, width) float32 dye field that you can visualize.
807 """
809 animation = AnimateIt(
810 "fluid_dynamics",
811 duration=10,
812 out_path=self.out_path,
813 colorful=True,
814 normalize=False,
815 )
816 self.print("\n🎯 HERO: Computational Fluid Dynamics")
817 sw = StopWatch()
818 device = self.device
819 queue = device.make_command_queue()
821 W, H = int(width), int(height)
822 N = W * H
824 # Allocate fields
825 # velocity ping-pong (float2), pressure ping-pong (float), divergence (float), dye ping-pong (float)
826 pr0 = device.make_buffer(N * 4)
827 pr1 = device.make_buffer(N * 4)
828 div = device.make_buffer(N * 4)
829 dye0 = device.make_buffer(N * 4)
830 dye1 = device.make_buffer(N * 4)
832 # Initialize dye and velocity with complex patterns for a richer start
833 y, x = np.mgrid[0:H, 0:W]
834 cx, cy = W // 2, H // 2
836 # Composite dye: two Gaussians + a circular ring
837 def gauss(cx0, cy0, sigma):
838 return np.exp(-(((x - cx0) ** 2 + (y - cy0) ** 2) / (2.0 * sigma * sigma)))
840 sigma1 = 0.10 * min(W, H)
841 sigma2 = 0.07 * min(W, H)
842 blob1 = gauss(0.30 * W, 0.40 * H, sigma1)
843 blob2 = gauss(0.70 * W, 0.60 * H, sigma2) * 0.8
845 r = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
846 r0 = 0.28 * min(W, H)
847 ring_sigma = 0.04 * min(W, H)
848 ring = np.exp(-((r - r0) ** 2) / (2.0 * ring_sigma * ring_sigma)) * 0.6
850 dye_comp = np.clip(blob1 + blob2 + ring, 0.0, 1.0).astype(np.float32)
851 animation.add_frame(dye_comp)
853 # Complex initial velocity: decaying swirl + shear bands
854 dx = x.astype(np.float32) - np.float32(cx)
855 dy_ = y.astype(np.float32) - np.float32(cy)
856 r2 = dx * dx + dy_ * dy_
857 sigma_v = np.float32(0.25 * min(W, H))
858 swirl_mag = np.exp(-(r2) / (2.0 * sigma_v * sigma_v)).astype(np.float32)
859 eps = np.float32(1e-5)
860 inv_norm = 1.0 / np.sqrt(r2 + eps)
861 ox = (-dy_) * inv_norm
862 oy = (dx) * inv_norm
863 swirl_strength = np.float32(2.0)
864 vx_swirl = swirl_strength * swirl_mag * ox
865 vy_swirl = swirl_strength * swirl_mag * oy
867 vx_shear = np.zeros_like(vx_swirl, dtype=np.float32)
868 vy_shear = np.zeros_like(vy_swirl, dtype=np.float32)
869 band1 = (y >= int(0.25 * H)) & (y < int(0.35 * H))
870 band2 = (y >= int(0.65 * H)) & (y < int(0.75 * H))
871 vx_shear[band1] = 0.5
872 vx_shear[band2] = -0.5
874 vx = (vx_swirl + vx_shear).astype(np.float32)
875 vy = (vy_swirl + vy_shear).astype(np.float32)
877 vel_init = np.stack([vx, vy], axis=-1).astype(np.float32)
878 vel0 = device.make_buffer_from_numpy(vel_init)
879 vel1 = device.make_buffer_from_numpy(vel_init)
881 # Fill initial dye
882 dye_init = np.ascontiguousarray(dye_comp)
883 dye_init = dye_init.reshape(-1)
884 tmp_dye = device.make_buffer_from_numpy(dye_init)
885 # If we don't have a direct blit, just keep tmp_dye as dye0 initial
886 dye0 = tmp_dye
888 # Params buffer (match Metal struct layout exactly: uint,uint,float,float,float,float,float,float,float)
889 params_dtype = np.dtype(
890 [
891 ("width", np.uint32),
892 ("height", np.uint32),
893 ("dt", np.float32),
894 ("dx", np.float32),
895 ("visc", np.float32),
896 ("fx", np.float32),
897 ("fy", np.float32),
898 ("fr", np.float32),
899 ("fs", np.float32),
900 ],
901 align=False,
902 )
903 params_struct = np.zeros(1, dtype=params_dtype)
904 params_struct["width"] = np.uint32(W)
905 params_struct["height"] = np.uint32(H)
906 params_struct["dt"] = np.float32(dt)
907 params_struct["dx"] = np.float32(1.0) # dx = 1.0
908 params_struct["visc"] = np.float32(visc)
909 params_struct["fx"] = np.float32(cx)
910 params_struct["fy"] = np.float32(cy)
911 params_struct["fr"] = np.float32(0.15 * min(W, H)) # force radius
912 params_struct["fs"] = np.float32(200.0) # force strength
913 # Reinterpret as 32-bit words for raw byte copy
914 params_words = np.frombuffer(params_struct.tobytes(), dtype=np.uint32)
915 params_buf = device.make_buffer_from_numpy(params_words)
917 # Metal kernels
918 source = f"""
919 #include <metal_stdlib>
920 using namespace metal;
922 struct Params {{
923 uint width;
924 uint height;
925 float dt;
926 float dx;
927 float visc;
928 float fx;
929 float fy;
930 float fr;
931 float fs;
932 }};
934 inline uint idx(uint x, uint y, uint W) {{
935 return y * W + x;
936 }}
938 inline float clamp01(float v) {{ return clamp(v, 0.0f, 1.0f); }}
940 inline float2 sample_vel(const device float2* vel, float x, float y, uint W, uint H) {{
941 // bilinear sample at (x,y) in grid space
942 x = clamp(x, 0.0f, (float)(W-1));
943 y = clamp(y, 0.0f, (float)(H-1));
944 uint x0 = (uint)floor(x);
945 uint y0 = (uint)floor(y);
946 uint x1 = min(x0 + 1, W - 1);
947 uint y1 = min(y0 + 1, H - 1);
948 float tx = x - (float)x0;
949 float ty = y - (float)y0;
950 float2 v00 = vel[idx(x0,y0,W)];
951 float2 v10 = vel[idx(x1,y0,W)];
952 float2 v01 = vel[idx(x0,y1,W)];
953 float2 v11 = vel[idx(x1,y1,W)];
954 float2 vx0 = mix(v00, v10, tx);
955 float2 vx1 = mix(v01, v11, tx);
956 return mix(vx0, vx1, ty);
957 }}
959 inline float sample_s(const device float* s, float x, float y, uint W, uint H) {{
960 x = clamp(x, 0.0f, (float)(W-1));
961 y = clamp(y, 0.0f, (float)(H-1));
962 uint x0 = (uint)floor(x);
963 uint y0 = (uint)floor(y);
964 uint x1 = min(x0 + 1, W - 1);
965 uint y1 = min(y0 + 1, H - 1);
966 float tx = x - (float)x0;
967 float ty = y - (float)y0;
968 float s00 = s[idx(x0,y0,W)];
969 float s10 = s[idx(x1,y0,W)];
970 float s01 = s[idx(x0,y1,W)];
971 float s11 = s[idx(x1,y1,W)];
972 float sx0 = mix(s00, s10, tx);
973 float sx1 = mix(s01, s11, tx);
974 return mix(sx0, sx1, ty);
975 }}
977 kernel void add_force(const device Params* P [[buffer(3)]],
978 device float2* vel_out [[buffer(0)]],
979 uint2 gid [[thread_position_in_grid]]) {{
980 uint W = P->width, H = P->height;
981 if (gid.x >= W || gid.y >= H) return;
982 float2 pos = float2(P->fx, P->fy);
983 float r = P->fr;
984 float2 center = float2((float)gid.x, (float)gid.y);
985 float2 d = center - pos;
986 float dist2 = dot(d,d);
987 float influence = exp(-dist2 / (r*r));
988 float2 orth = float2(-d.y, d.x);
989 vel_out[idx(gid.x, gid.y, W)] += normalize(orth + 1e-5) * (P->fs * influence);
990 }}
992 kernel void advect_vel(const device Params* P [[buffer(3)]],
993 const device float2* vel_in [[buffer(0)]],
994 device float2* vel_out [[buffer(1)]],
995 uint2 gid [[thread_position_in_grid]]) {{
996 uint W = P->width, H = P->height;
997 if (gid.x >= W || gid.y >= H) return;
998 float2 v = vel_in[idx(gid.x, gid.y, W)];
999 float x = (float)gid.x - P->dt * v.x;
1000 float y = (float)gid.y - P->dt * v.y;
1001 vel_out[idx(gid.x, gid.y, W)] = sample_vel(vel_in, x, y, W, H);
1002 }}
1004 kernel void divergence(const device Params* P [[buffer(3)]],
1005 const device float2* vel [[buffer(0)]],
1006 device float* div_out [[buffer(1)]],
1007 uint2 gid [[thread_position_in_grid]]) {{
1008 uint W = P->width, H = P->height;
1009 if (gid.x >= W || gid.y >= H) return;
1010 uint x = gid.x, y = gid.y;
1011 uint xm = max(int(x)-1, 0);
1012 uint xp = min(x+1, W-1);
1013 uint ym = max(int(y)-1, 0);
1014 uint yp = min(y+1, H-1);
1015 float2 vxm = vel[idx(xm,y,W)];
1016 float2 vxp = vel[idx(xp,y,W)];
1017 float2 vym = vel[idx(x,ym,W)];
1018 float2 vyp = vel[idx(x,yp,W)];
1019 float div = 0.5f * ((vxp.x - vxm.x) + (vyp.y - vym.y));
1020 div_out[idx(x,y,W)] = div;
1021 }}
1023 kernel void jacobi_pressure(const device Params* P [[buffer(3)]],
1024 const device float* p_in [[buffer(0)]],
1025 const device float* b [[buffer(1)]],
1026 device float* p_out [[buffer(2)]],
1027 uint2 gid [[thread_position_in_grid]]) {{
1028 uint W = P->width, H = P->height;
1029 if (gid.x >= W || gid.y >= H) return;
1030 uint x = gid.x, y = gid.y;
1031 uint xm = max(int(x)-1, 0);
1032 uint xp = min(x+1, W-1);
1033 uint ym = max(int(y)-1, 0);
1034 uint yp = min(y+1, H-1);
1035 float pL = p_in[idx(xm,y,W)];
1036 float pR = p_in[idx(xp,y,W)];
1037 float pB = p_in[idx(x,ym,W)];
1038 float pT = p_in[idx(x,yp,W)];
1039 float rhs = b[idx(x,y,W)];
1040 // alpha = -dx*dx, rBeta = 0.25
1041 float p_new = (pL + pR + pB + pT - rhs) * 0.25f;
1042 p_out[idx(x,y,W)] = p_new;
1043 }}
1045 kernel void subtract_gradient(const device Params* P [[buffer(3)]],
1046 const device float2* vel_in [[buffer(0)]],
1047 const device float* p [[buffer(1)]],
1048 device float2* vel_out [[buffer(2)]],
1049 uint2 gid [[thread_position_in_grid]]) {{
1050 uint W = P->width, H = P->height;
1051 if (gid.x >= W || gid.y >= H) return;
1052 uint x = gid.x, y = gid.y;
1053 uint xm = max(int(x)-1, 0);
1054 uint xp = min(x+1, W-1);
1055 uint ym = max(int(y)-1, 0);
1056 uint yp = min(y+1, H-1);
1057 float pL = p[idx(xm,y,W)];
1058 float pR = p[idx(xp,y,W)];
1059 float pB = p[idx(x,ym,W)];
1060 float pT = p[idx(x,yp,W)];
1061 float2 v = vel_in[idx(x,y,W)];
1062 v -= 0.5f * float2(pR - pL, pT - pB);
1063 vel_out[idx(x,y,W)] = v;
1064 }}
1066 kernel void advect_scalar(const device Params* P [[buffer(3)]],
1067 const device float* s_in [[buffer(0)]],
1068 const device float2* vel [[buffer(1)]],
1069 device float* s_out [[buffer(2)]],
1070 uint2 gid [[thread_position_in_grid]]) {{
1071 uint W = P->width, H = P->height;
1072 if (gid.x >= W || gid.y >= H) return;
1073 float2 v = vel[idx(gid.x, gid.y, W)];
1074 float x = (float)gid.x - P->dt * v.x;
1075 float y = (float)gid.y - P->dt * v.y;
1076 s_out[idx(gid.x, gid.y, W)] = sample_s(s_in, x, y, W, H);
1077 }}
1078 """
1080 lib = device.make_library(source)
1081 fn_add_force = lib.make_function("add_force")
1082 fn_advect_v = lib.make_function("advect_vel")
1083 fn_div = lib.make_function("divergence")
1084 fn_jacobi = lib.make_function("jacobi_pressure")
1085 fn_subgrad = lib.make_function("subtract_gradient")
1086 fn_advect_s = lib.make_function("advect_scalar")
1088 p_add = device.make_compute_pipeline_state(fn_add_force)
1089 p_advv = device.make_compute_pipeline_state(fn_advect_v)
1090 p_div = device.make_compute_pipeline_state(fn_div)
1091 p_jac = device.make_compute_pipeline_state(fn_jacobi)
1092 p_sub = device.make_compute_pipeline_state(fn_subgrad)
1093 p_advs = device.make_compute_pipeline_state(fn_advect_s)
1095 # Simple simulation loop
1096 for _ in range(int(steps)):
1097 with sw:
1098 # Add swirling force into vel0 -> vel1
1099 cb = queue.make_command_buffer()
1100 enc = cb.make_compute_command_encoder()
1101 enc.set_compute_pipeline_state(p_add)
1102 enc.set_buffer(vel0, 0, 0) # out
1103 enc.set_buffer(params_buf, 0, 3)
1104 enc.dispatch_threads((W, H, 1), threadgroup)
1105 enc.end_encoding()
1106 cb.commit()
1107 cb.wait_until_completed()
1109 # Advect velocity: vel1 = advect_vel(vel0)
1110 cb = queue.make_command_buffer()
1111 enc = cb.make_compute_command_encoder()
1112 enc.set_compute_pipeline_state(p_advv)
1113 enc.set_buffer(vel0, 0, 0) # in
1114 enc.set_buffer(vel1, 0, 1) # out
1115 enc.set_buffer(params_buf, 0, 3)
1116 enc.dispatch_threads((W, H, 1), threadgroup)
1117 enc.end_encoding()
1118 cb.commit()
1119 cb.wait_until_completed()
1121 # Compute divergence of vel1 into div
1122 cb = queue.make_command_buffer()
1123 enc = cb.make_compute_command_encoder()
1124 enc.set_compute_pipeline_state(p_div)
1125 enc.set_buffer(vel1, 0, 0)
1126 enc.set_buffer(div, 0, 1)
1127 enc.set_buffer(params_buf, 0, 3)
1128 enc.dispatch_threads((W, H, 1), threadgroup)
1129 enc.end_encoding()
1130 cb.commit()
1131 cb.wait_until_completed()
1133 # Clear pressure buffers to zero (first few iterations rely on initial zeros)
1134 zero_np = np.zeros(N, dtype=np.float32)
1135 pr0 = device.make_buffer_from_numpy(zero_np)
1136 pr1 = device.make_buffer_from_numpy(zero_np)
1138 # Jacobi iterations to solve for pressure
1139 J_ITERS = 20
1140 pin, pout = pr0, pr1
1141 for __ in range(J_ITERS):
1142 cb = queue.make_command_buffer()
1143 enc = cb.make_compute_command_encoder()
1144 enc.set_compute_pipeline_state(p_jac)
1145 enc.set_buffer(pin, 0, 0) # p_in
1146 enc.set_buffer(div, 0, 1) # b
1147 enc.set_buffer(pout, 0, 2) # p_out
1148 enc.set_buffer(params_buf, 0, 3)
1149 enc.dispatch_threads((W, H, 1), threadgroup)
1150 enc.end_encoding()
1151 cb.commit()
1152 cb.wait_until_completed()
1153 pin, pout = pout, pin
1154 pressure = pin
1156 # Subtract gradient: vel0 = project(vel1, pressure)
1157 cb = queue.make_command_buffer()
1158 enc = cb.make_compute_command_encoder()
1159 enc.set_compute_pipeline_state(p_sub)
1160 enc.set_buffer(vel1, 0, 0)
1161 enc.set_buffer(pressure, 0, 1)
1162 enc.set_buffer(vel0, 0, 2)
1163 enc.set_buffer(params_buf, 0, 3)
1164 enc.dispatch_threads((W, H, 1), threadgroup)
1165 enc.end_encoding()
1166 cb.commit()
1167 cb.wait_until_completed()
1169 # Advect dye by velocity: dye1 = advect_scalar(dye0, vel0)
1170 cb = queue.make_command_buffer()
1171 enc = cb.make_compute_command_encoder()
1172 enc.set_compute_pipeline_state(p_advs)
1173 enc.set_buffer(dye0, 0, 0) # s_in
1174 enc.set_buffer(vel0, 0, 1) # vel
1175 enc.set_buffer(dye1, 0, 2) # s_out
1176 enc.set_buffer(params_buf, 0, 3)
1177 enc.dispatch_threads((W, H, 1), threadgroup)
1178 enc.end_encoding()
1179 cb.commit()
1180 cb.wait_until_completed()
1182 # Ping-pong dye
1183 dye0, dye1 = dye1, dye0
1184 frame = dye0.to_numpy(np.float32, (H, W))
1185 animation.add_frame(frame)
1187 # Read back dye
1188 dye_out = dye0.to_numpy(np.float32, (H, W))
1189 animation.add_frame(dye_out)
1190 animation.show()
1191 self.print(sw)
1192 return dye_out
1194 def benchmark_matrix_multiply(self, sizes: list):
1195 """Benchmark matrix multiplication performance"""
1196 print("Matrix Multiplication Benchmark")
1197 print("=" * 50)
1198 print(f"{'Size':<10} {'Metal (ms)':<12} {'NumPy (ms)':<12} {'Speedup':<10}")
1199 print("-" * 50)
1201 for size in sizes:
1202 # Generate random matrices
1203 A = np.random.random((size, size)).astype(np.float32)
1204 B = np.random.random((size, size)).astype(np.float32)
1206 # Benchmark Metal
1207 start_time = time.time()
1208 metal_result = self.metal_ops.matrix_multiply(A, B)
1209 metal_time = (time.time() - start_time) * 1000
1211 # Benchmark NumPy
1212 start_time = time.time()
1213 numpy_result = np.dot(A, B)
1214 numpy_time = (time.time() - start_time) * 1000
1216 # Verify correctness
1217 if np.allclose(metal_result, numpy_result, rtol=1e-4):
1218 speedup = numpy_time / metal_time
1219 print(
1220 f"{size:<10} {metal_time:<12.2f} {numpy_time:<12.2f} {speedup:<10.2f}x"
1221 )
1222 else:
1223 print(f"{size:<10} ERROR: Results don't match!")
1225 def benchmark_vector_operations(self, size: int = 1000000):
1226 """Benchmark vector operations"""
1227 self.print(f"\nVector Operations Benchmark (size: {size:,})")
1228 self.print("=" * 50)
1230 a = np.random.random(size).astype(np.float32)
1231 b = np.random.random(size).astype(np.float32)
1233 operations = [
1234 ("Addition", "add", lambda x, y: x + y),
1235 ("Multiplication", "multiply", lambda x, y: x * y),
1236 ]
1238 for op_name, metal_op, numpy_op in operations:
1239 # Metal
1240 start_time = time.time()
1241 metal_result = self.metal_ops.vector_operations(a, b, metal_op)
1242 metal_time = (time.time() - start_time) * 1000
1244 # NumPy
1245 start_time = time.time()
1246 numpy_result = numpy_op(a, b)
1247 numpy_time = (time.time() - start_time) * 1000
1249 if np.allclose(metal_result, numpy_result, rtol=1e-5):
1250 speedup = numpy_time / metal_time
1251 self.print(
1252 f"{op_name}: Metal {metal_time:.2f}ms, NumPy {numpy_time:.2f}ms, "
1253 f"Speedup: {speedup:.2f}x"
1254 )
1255 else:
1256 self.print(f"{op_name}: ERROR - Results don't match!")
1258 def demo_device_info(self):
1259 """Show information about available Metal devices"""
1260 self.print("Available Metal Devices")
1261 self.print("=" * 30)
1263 devices = pymetallic.Device.get_all_devices()
1264 for i, device in enumerate(devices):
1265 self.print(f"Device {i}: {device.name}")
1266 self.print(
1267 f" Supports barycentric coordinates: {device.supports_shader_barycentric_coordinates()}"
1268 )
1270 self.print(
1271 f"\nUsing default device: {pymetallic.Device.get_default_device().name}"
1272 )
1273 self.print()
1275 def demo_performance_benchmark(self):
1276 """Run comprehensive performance benchmarks"""
1277 self.print("\n⚡ Performance Benchmarks")
1278 self.print("-" * 30)
1279 # Matrix multiplication benchmark
1280 self.print("Matrix Multiplication Performance:")
1281 self.benchmark_matrix_multiply([128, 256, 512, 1024])
1283 # Vector operations benchmark
1284 self.print("\nVector Operations Performance:")
1285 self.benchmark_vector_operations(2000000)
1287 def demo_error_handling(self):
1288 """Demonstrate error handling"""
1289 self.print("\n🛡️ Error Handling")
1290 self.print("-" * 30)
1291 device = self.device
1292 # Test invalid shader compilation
1293 self.print("Invalid shader compilation:")
1294 try:
1295 invalid_shader = "This is not valid Metal code!"
1296 library = device.make_library(invalid_shader)
1297 self.print(" ❌ Should have failed!")
1298 except pymetallic.MetalError as e:
1299 self.print(" ✅ Correctly caught compilation error")
1301 # Test non-existent function
1302 self.print("Non-existent function access:")
1303 try:
1304 valid_shader = """
1305 #include <metal_stdlib>
1306 using namespace metal;
1307 kernel void test_function(device float* data [[buffer(0)]]) {}
1308 """
1309 library = device.make_library(valid_shader)
1310 function = library.make_function("nonexistent_function")
1311 self.print(" ❌ Should have failed!")
1312 except pymetallic.MetalError as e:
1313 self.print(" ✅ Correctly caught function not found error")
1315 def get_demos(self):
1316 demos = {
1317 attr.removeprefix("demo_"): getattr(self, attr)
1318 for attr in dir(self)
1319 if attr.startswith("demo_") and callable(getattr(self, attr))
1320 }
1321 return demos
1323 def run_complete_demo(self):
1324 """Run the complete demonstration"""
1325 try:
1326 for name, demo in self.get_demos().items():
1327 self.print(f"🚀 Running {name} demo")
1328 demo()
1329 self.print("\n🎉 PyMetallic Demo Complete!")
1330 self.print("=" * 50)
1331 self.print("✅ All features demonstrated successfully")
1332 self.print("📖 See the API summary above for usage patterns")
1333 self.print("🚀 Ready for high-performance GPU computing on macOS!")
1335 except KeyboardInterrupt:
1336 self.print("\n\n⏸️ Demo interrupted by user")
1337 except Exception as e:
1338 self.print(f"\n❌ Demo failed: {e}")
1339 import traceback
1341 traceback.print_exc()
1344class MetalMatrixOperations:
1345 """High-level matrix operations using Metal compute shaders"""
1347 def __init__(self, device=None, quiet=False):
1348 self.device = device or Device.get_default_device()
1349 self.queue = self.device.make_command_queue()
1350 self.quiet = quiet
1351 self._compile_shaders()
1353 def _compile_shaders(self):
1354 """Compile all Metal shaders used by this class"""
1355 device = self.device
1356 shader_source = """
1357 #include <metal_stdlib>
1358 using namespace metal;
1360 // Matrix multiplication kernel
1361 kernel void matrix_multiply(device const float* A [[buffer(0)]],
1362 device const float* B [[buffer(1)]],
1363 device float* C [[buffer(2)]],
1364 constant uint& M [[buffer(3)]],
1365 constant uint& N [[buffer(4)]],
1366 constant uint& K [[buffer(5)]],
1367 uint2 gid [[thread_position_in_grid]]) {
1368 uint row = gid.y;
1369 uint col = gid.x;
1371 if (row >= M || col >= N) return;
1373 float sum = 0.0;
1374 for (uint k = 0; k < K; k++) {
1375 sum += A[row * K + k] * B[k * N + col];
1376 }
1377 C[row * N + col] = sum;
1378 }
1380 // Vector operations
1381 kernel void vector_add(device const float* a [[buffer(0)]],
1382 device const float* b [[buffer(1)]],
1383 device float* result [[buffer(2)]],
1384 uint index [[thread_position_in_grid]]) {
1385 result[index] = a[index] + b[index];
1386 }
1388 kernel void vector_multiply(device const float* a [[buffer(0)]],
1389 device const float* b [[buffer(1)]],
1390 device float* result [[buffer(2)]],
1391 uint index [[thread_position_in_grid]]) {
1392 result[index] = a[index] * b[index];
1393 }
1395 kernel void vector_scale(device const float* input [[buffer(0)]],
1396 device float* output [[buffer(1)]],
1397 constant float& scale [[buffer(2)]],
1398 uint index [[thread_position_in_grid]]) {
1399 output[index] = input[index] * scale;
1400 }
1402 // Reduction operations
1403 kernel void reduce_sum(device const float* input [[buffer(0)]],
1404 device float* output [[buffer(1)]],
1405 constant uint& n [[buffer(2)]],
1406 uint index [[thread_position_in_grid]],
1407 uint threads_per_group [[threads_per_threadgroup]]) {
1409 threadgroup float shared_data[256];
1410 uint tid = index % threads_per_group;
1411 uint gid = index;
1413 // Load data into shared memory
1414 shared_data[tid] = (gid < n) ? input[gid] : 0.0;
1415 threadgroup_barrier(mem_flags::mem_threadgroup);
1417 // Reduction in shared memory
1418 for (uint s = threads_per_group / 2; s > 0; s >>= 1) {
1419 if (tid < s) {
1420 shared_data[tid] += shared_data[tid + s];
1421 }
1422 threadgroup_barrier(mem_flags::mem_threadgroup);
1423 }
1425 // Write result for this block
1426 if (tid == 0) {
1427 output[index / threads_per_group] = shared_data[0];
1428 }
1429 }
1431 // Image processing kernels
1432 kernel void gaussian_blur_3x3(texture2d<float, access::read> inputTexture [[texture(0)]],
1433 texture2d<float, access::write> outputTexture [[texture(1)]],
1434 uint2 gid [[thread_position_in_grid]]) {
1436 if (gid.x >= inputTexture.get_width() || gid.y >= inputTexture.get_height()) {
1437 return;
1438 }
1440 // 3x3 Gaussian kernel
1441 const float gaussian_kernel[9] = {
1442 1.0/16.0, 2.0/16.0, 1.0/16.0,
1443 2.0/16.0, 4.0/16.0, 2.0/16.0,
1444 1.0/16.0, 2.0/16.0, 1.0/16.0
1445 };
1447 float4 color = float4(0.0);
1448 for (int dy = -1; dy <= 1; dy++) {
1449 for (int dx = -1; dx <= 1; dx++) {
1450 uint2 coord = uint2(max(0, min(int(inputTexture.get_width() - 1), int(gid.x) + dx)),
1451 max(0, min(int(inputTexture.get_height() - 1), int(gid.y) + dy)));
1452 color += inputTexture.read(coord) * gaussian_kernel[(dy + 1) * 3 + (dx + 1)];
1453 }
1454 }
1456 outputTexture.write(color, gid);
1457 }
1459 // Buffer-based 5x5 Gaussian blur (separable weights, clamp-to-edge)
1460 kernel void gaussian_blur_5x5_buffer(device const float* input [[buffer(0)]],
1461 device float* output [[buffer(1)]],
1462 constant uint& width [[buffer(2)]],
1463 constant uint& height [[buffer(3)]],
1464 uint2 gid [[thread_position_in_grid]]) {
1465 if (gid.x >= width || gid.y >= height) { return; }
1467 const float k[5] = {1.0, 4.0, 6.0, 4.0, 1.0};
1468 float sum = 0.0;
1469 float norm = 0.0;
1471 for (int dy = -2; dy <= 2; ++dy) {
1472 int y = clamp(int(gid.y) + dy, 0, int(height) - 1);
1473 float wy = k[dy + 2];
1474 for (int dx = -2; dx <= 2; ++dx) {
1475 int x = clamp(int(gid.x) + dx, 0, int(width) - 1);
1476 float w = wy * k[dx + 2];
1477 sum += input[y * width + x] * w;
1478 norm += w;
1479 }
1480 }
1482 output[gid.y * width + gid.x] = sum / norm;
1483 }
1484 """
1486 self.library = device.make_library(shader_source)
1488 # Create pipeline states
1489 self.matrix_multiply_pipeline = device.make_compute_pipeline_state(
1490 self.library.make_function("matrix_multiply")
1491 )
1492 self.vector_add_pipeline = device.make_compute_pipeline_state(
1493 self.library.make_function("vector_add")
1494 )
1495 self.vector_multiply_pipeline = device.make_compute_pipeline_state(
1496 self.library.make_function("vector_multiply")
1497 )
1498 self.vector_scale_pipeline = device.make_compute_pipeline_state(
1499 self.library.make_function("vector_scale")
1500 )
1501 self.reduce_sum_pipeline = device.make_compute_pipeline_state(
1502 self.library.make_function("reduce_sum")
1503 )
1504 self.gaussian_blur_5x5_pipeline = device.make_compute_pipeline_state(
1505 self.library.make_function("gaussian_blur_5x5_buffer")
1506 )
1508 def matrix_multiply(self, A: np.ndarray, B: np.ndarray) -> np.ndarray:
1509 """Perform matrix multiplication using Metal compute shaders"""
1510 device = self.device
1511 if A.shape[1] != B.shape[0]:
1512 raise ValueError("Matrix dimensions don't match for multiplication")
1514 M, K = A.shape
1515 K2, N = B.shape
1517 # Create buffers
1518 buffer_A = device.make_buffer_from_numpy(A.astype(np.float32))
1519 buffer_B = device.make_buffer_from_numpy(B.astype(np.float32))
1520 buffer_C = device.make_buffer(M * N * 4) # float32 = 4 bytes
1522 # Create parameter buffers for matrix dimensions
1523 dims = np.array([M, N, K], dtype=np.uint32)
1524 buffer_M = device.make_buffer_from_numpy(np.array([M], dtype=np.uint32))
1525 buffer_N = device.make_buffer_from_numpy(np.array([N], dtype=np.uint32))
1526 buffer_K = device.make_buffer_from_numpy(np.array([K], dtype=np.uint32))
1528 # Execute kernel
1529 command_buffer = self.queue.make_command_buffer()
1530 encoder = command_buffer.make_compute_command_encoder()
1532 encoder.set_compute_pipeline_state(self.matrix_multiply_pipeline)
1533 encoder.set_buffer(buffer_A, 0, 0)
1534 encoder.set_buffer(buffer_B, 0, 1)
1535 encoder.set_buffer(buffer_C, 0, 2)
1536 encoder.set_buffer(buffer_M, 0, 3)
1537 encoder.set_buffer(buffer_N, 0, 4)
1538 encoder.set_buffer(buffer_K, 0, 5)
1540 # Dispatch with 2D grid
1541 threads_per_grid = (N, M, 1)
1542 threads_per_threadgroup = (16, 16, 1)
1543 encoder.dispatch_threads(threads_per_grid, threads_per_threadgroup)
1544 encoder.end_encoding()
1546 command_buffer.commit()
1547 command_buffer.wait_until_completed()
1549 # Get result
1550 result = buffer_C.to_numpy(np.float32, (M, N))
1551 return result
1553 def vector_operations(
1554 self, a: np.ndarray, b: np.ndarray, operation: str
1555 ) -> np.ndarray:
1556 """Perform element-wise vector operations"""
1557 device = self.device
1558 if a.shape != b.shape:
1559 raise ValueError("Vector shapes must match")
1561 # Flatten arrays for processing
1562 a_flat = a.flatten().astype(np.float32)
1563 b_flat = b.flatten().astype(np.float32)
1565 buffer_a = device.make_buffer_from_numpy(a_flat)
1566 buffer_b = device.make_buffer_from_numpy(b_flat)
1567 buffer_result = device.make_buffer(len(a_flat) * 4)
1569 # Select pipeline based on operation
1570 if operation == "add":
1571 pipeline = self.vector_add_pipeline
1572 elif operation == "multiply":
1573 pipeline = self.vector_multiply_pipeline
1574 else:
1575 raise ValueError(f"Unsupported operation: {operation}")
1577 # Execute
1578 command_buffer = self.queue.make_command_buffer()
1579 encoder = command_buffer.make_compute_command_encoder()
1581 encoder.set_compute_pipeline_state(pipeline)
1582 encoder.set_buffer(buffer_a, 0, 0)
1583 encoder.set_buffer(buffer_b, 0, 1)
1584 encoder.set_buffer(buffer_result, 0, 2)
1586 encoder.dispatch_threads((len(a_flat), 1, 1), (64, 1, 1))
1587 encoder.end_encoding()
1589 command_buffer.commit()
1590 command_buffer.wait_until_completed()
1592 result = buffer_result.to_numpy(np.float32, a.shape)
1593 return result
1595 def vector_scale(self, input_vector: np.ndarray, scale: float) -> np.ndarray:
1596 """Scale a vector by a constant"""
1597 device = self.device
1598 input_flat = input_vector.flatten().astype(np.float32)
1600 buffer_input = device.make_buffer_from_numpy(input_flat)
1601 buffer_output = device.make_buffer(len(input_flat) * 4)
1602 buffer_scale = device.make_buffer_from_numpy(
1603 np.array([scale], dtype=np.float32)
1604 )
1606 command_buffer = self.queue.make_command_buffer()
1607 encoder = command_buffer.make_compute_command_encoder()
1609 encoder.set_compute_pipeline_state(self.vector_scale_pipeline)
1610 encoder.set_buffer(buffer_input, 0, 0)
1611 encoder.set_buffer(buffer_output, 0, 1)
1612 encoder.set_buffer(buffer_scale, 0, 2)
1614 encoder.dispatch_threads((len(input_flat), 1, 1), (64, 1, 1))
1615 encoder.end_encoding()
1617 command_buffer.commit()
1618 command_buffer.wait_until_completed()
1620 result = buffer_output.to_numpy(np.float32, input_vector.shape)
1621 return result
1623 def gaussian_blur_5x5(self, image: np.ndarray) -> np.ndarray:
1624 """Apply a 5x5 Gaussian blur to a 2D float32 image using a compute kernel."""
1625 device = self.device
1626 if image.ndim != 2:
1627 raise ValueError("Input image must be a 2D array (grayscale).")
1628 img = image.astype(np.float32, copy=False)
1629 h, w = img.shape
1631 buffer_in = device.make_buffer_from_numpy(img)
1632 buffer_out = device.make_buffer(w * h * 4)
1633 buf_w = device.make_buffer_from_numpy(np.array([w], dtype=np.uint32))
1634 buf_h = device.make_buffer_from_numpy(np.array([h], dtype=np.uint32))
1636 command_buffer = self.queue.make_command_buffer()
1637 encoder = command_buffer.make_compute_command_encoder()
1639 encoder.set_compute_pipeline_state(self.gaussian_blur_5x5_pipeline)
1640 encoder.set_buffer(buffer_in, 0, 0)
1641 encoder.set_buffer(buffer_out, 0, 1)
1642 encoder.set_buffer(buf_w, 0, 2)
1643 encoder.set_buffer(buf_h, 0, 3)
1645 encoder.dispatch_threads((w, h, 1), (16, 16, 1))
1646 encoder.end_encoding()
1648 command_buffer.commit()
1649 command_buffer.wait_until_completed()
1651 return buffer_out.to_numpy(np.float32, (h, w))
1654def run_comprehensive_example():
1655 """Run a comprehensive example showing various features"""
1656 print("PyMetallic Comprehensive Example")
1657 print("=" * 40)
1658 md = MetallicDemo()
1659 try:
1660 md.demo_device_info()
1661 metal_ops = md.metal_ops
1663 print("1. Matrix Multiplication Example")
1664 print("-" * 30)
1665 A = np.random.random((128, 64)).astype(np.float32)
1666 B = np.random.random((64, 96)).astype(np.float32)
1668 start_time = time.time()
1669 C_metal = metal_ops.matrix_multiply(A, B)
1670 metal_time = time.time() - start_time
1672 C_numpy = np.dot(A, B)
1674 print(f"Matrix shapes: A{A.shape} × B{B.shape} = C{C_metal.shape}")
1675 print(f"Metal computation time: {metal_time*1000:.2f}ms")
1676 print(f"Results match NumPy: {np.allclose(C_metal, C_numpy, rtol=1e-4)}")
1678 print("\n2. Vector Operations Example")
1679 print("-" * 30)
1680 vec_size = 100000
1681 x = np.random.random(vec_size).astype(np.float32)
1682 y = np.random.random(vec_size).astype(np.float32)
1684 # Vector addition
1685 z_add = metal_ops.vector_operations(x, y, "add")
1686 print(f"Vector addition (size {vec_size:,}): {np.allclose(z_add, x + y)}")
1688 # Vector multiplication
1689 z_mul = metal_ops.vector_operations(x, y, "multiply")
1690 print(f"Vector multiplication: {np.allclose(z_mul, x * y)}")
1692 # Vector scaling
1693 scale_factor = 2.5
1694 z_scale = metal_ops.vector_scale(x, scale_factor)
1695 print(
1696 f"Vector scaling by {scale_factor}: {np.allclose(z_scale, x * scale_factor)}"
1697 )
1699 print("\n3. Performance Benchmark")
1700 print("-" * 30)
1701 md.benchmark_matrix_multiply([64, 128, 256, 512, 1024, 2048])
1702 md.benchmark_vector_operations(4000000)
1704 print("\n4. 5x5 Gaussian Blur Example")
1705 print("-" * 30)
1706 # Create a sample grayscale image
1707 height, width = 128, 128
1708 image = np.random.random((height, width)).astype(np.float32)
1710 # Metal blur
1711 t = time.time()
1712 blurred_metal = metal_ops.gaussian_blur_5x5(image)
1713 mtl_elapsed = time.time() - t
1715 # CPU reference (separable 5x5 with clamp-to-edge)
1716 k1 = np.array([1, 4, 6, 4, 1], dtype=np.float32)
1717 k1 /= k1.sum()
1718 t = time.time()
1719 tmp = np.empty_like(image)
1720 for y in range(height):
1721 for x in range(width):
1722 s = 0.0
1723 for dx in range(-2, 3):
1724 xx = min(max(0, x + dx), width - 1)
1725 s += image[y, xx] * k1[dx + 2]
1726 tmp[y, x] = s
1728 blurred_cpu = np.empty_like(image)
1729 for y in range(height):
1730 for x in range(width):
1731 s = 0.0
1732 for dy in range(-2, 3):
1733 yy = min(max(0, y + dy), height - 1)
1734 s += tmp[yy, x] * k1[dy + 2]
1735 blurred_cpu[y, x] = s
1736 np_elapsed = time.time() - t
1737 speedup = np_elapsed / mtl_elapsed
1739 print(
1740 f"Gaussian blur correctness: {np.allclose(blurred_metal, blurred_cpu, rtol=1e-5)}"
1741 )
1742 print(
1743 f"Performance speedup: {speedup:.2f}x numpy={np_elapsed*1000:.2f}ms metal={mtl_elapsed*1000:.2f}ms"
1744 )
1746 print("\nAll examples completed successfully!")
1748 except Exception as e:
1749 print(f"Error running example: {e}")
1750 import traceback
1752 traceback.print_exc()