""" Memory Benchmark Module 内存性能测试：带宽测试、延迟测试、缓存性能 Optimized with ctypes for raw C-level performance """ import time import ctypes import multiprocessing import mmap import os import numpy as np # Keep for latency/cache tests from concurrent.futures import ProcessPoolExecutor from typing import Dict, Any # Load C standard library try: libc = ctypes.CDLL("libc.so.6") libc.memset.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t] libc.memcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t] libc.memchr.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t] except Exception: libc = None # --- C Extension Handling --- C_LIB_PATH = os.path.join(os.path.dirname(__file__), "_memory_bench_c.so") C_SRC_PATH = os.path.join(os.path.dirname(__file__), "memory_bench_c.c") def _compile_c_helper(): """Compiles the C helper library if it doesn't exist or is outdated.""" if not os.path.exists(C_SRC_PATH): return None needs_compile = False if not os.path.exists(C_LIB_PATH): needs_compile = True else: # Check timestamps if os.path.getmtime(C_SRC_PATH) > os.path.getmtime(C_LIB_PATH): needs_compile = True if needs_compile: # User requested max optimization cmd = f"gcc -O3 -shared -fPIC -o {C_LIB_PATH} {C_SRC_PATH}" if os.system(cmd) != 0: print("Failed to compile C helper.") return None try: lib = ctypes.CDLL(C_LIB_PATH) lib.measure_latency_random.argtypes = [ctypes.c_size_t, ctypes.c_size_t] lib.measure_latency_random.restype = ctypes.c_double lib.measure_latency_sequential.argtypes = [ctypes.c_size_t, ctypes.c_size_t] lib.measure_latency_sequential.restype = ctypes.c_double lib.measure_alloc_rate.argtypes = [ctypes.c_size_t, ctypes.c_size_t] lib.measure_alloc_rate.restype = ctypes.c_double return lib except Exception as e: print(f"Failed to load C helper: {e}") return None c_lib = _compile_c_helper() def _raw_memory_worker(args): """ Worker process for memory bandwidth test using raw C calls. Equivalent to sysbench memory test. """ block_size_mb, duration, mode = args block_size = block_size_mb * 1024 * 1024 # Use mmap for aligned, raw memory allocation (no Python object overhead) # Anonymous mapping src_map = mmap.mmap(-1, block_size) dst_map = None # For copy mode, we need a destination if mode == 'copy': dst_map = mmap.mmap(-1, block_size) # Get raw pointers src_addr = ctypes.addressof(ctypes.c_char.from_buffer(src_map)) dst_addr = ctypes.addressof(ctypes.c_char.from_buffer(dst_map)) if dst_map else 0 # Prepare C function calls memset = libc.memset memcpy = libc.memcpy memchr = libc.memchr start_time = time.time() iterations = 0 while time.time() - start_time < duration: if mode == 'read': # Scan memory (read access) # Find a byte that (likely) isn't there to force full scan memchr(src_addr, 1, block_size) elif mode == 'write': # Write memory memset(src_addr, 0, block_size) elif mode == 'copy': # Copy memory memcpy(dst_addr, src_addr, block_size) iterations += 1 elapsed = time.time() - start_time # Cleanup src_map.close() if dst_map: dst_map.close() return iterations, elapsed def benchmark_memory_bandwidth(block_size_mb: int = 4) -> Dict[str, Any]: """ 内存带宽测试 (Raw C Performance) Uses multiprocessing + ctypes to bypass Python overhead. """ if not libc: return {"error": "libc not found, cannot run optimized benchmark"} num_cores = multiprocessing.cpu_count() duration = 3.0 # sysbench defaults to 1KB-1MB blocks. User mentioned 1MB. # We use a slightly larger buffer per thread to amortize loop overhead if needed, # but 1MB-4MB is usually good for L3/RAM cache thrashing. # Let's stick to 4MB per thread to ensure we hit RAM. modes = ['read', 'write', 'copy'] results = {} with ProcessPoolExecutor(max_workers=num_cores) as executor: for mode in modes: # Submit tasks futures = [executor.submit(_raw_memory_worker, (block_size_mb, duration, mode)) for _ in range(num_cores)] total_iterations = 0 max_elapsed = 0 for f in futures: iters, elapsed = f.result() total_iterations += iters max_elapsed = max(max_elapsed, elapsed) # Calculate Bandwidth # Data transferred per iteration = block_size bytes_per_iter = block_size_mb * 1024 * 1024 total_bytes = total_iterations * bytes_per_iter # Note: For 'copy', sysbench counts read+write? # Usually bandwidth is defined as bytes processed. # If we copy 1GB, we read 1GB and write 1GB. # sysbench memory test reports "transferred". # For copy, let's report the amount of data moved (Payload). # Or if user wants bus bandwidth, it's 2x. # Benchmarks usually report the size of the buffer processed. # However, previous impl multiplied by 2. Let's stick to total bytes moved over bus. if mode == 'copy': total_bytes *= 2 # Avoid division by zero if max_elapsed > 0: bandwidth_gb_s = total_bytes / max_elapsed / (1024**3) else: bandwidth_gb_s = 0 results[f"{mode}_bandwidth_gb_s"] = round(bandwidth_gb_s, 3) return { "test": "memory_bandwidth", "description": f"Memory bandwidth test (Multi-core C-level, {num_cores} threads)", "block_size_mb": block_size_mb, "read_bandwidth_gb_s": results['read_bandwidth_gb_s'], "write_bandwidth_gb_s": results['write_bandwidth_gb_s'], "copy_bandwidth_gb_s": results['copy_bandwidth_gb_s'], "score": round((results['read_bandwidth_gb_s'] + results['write_bandwidth_gb_s'] + results['copy_bandwidth_gb_s']) * 10, 2), } def benchmark_memory_latency(iterations: int = 10000000) -> Dict[str, Any]: """ 内存延迟测试（随机访问） Uses C helper for precise pointer chasing. """ if not c_lib: return {"error": "C helper not available"} # Test random access latency on a large block (64MB) to hit RAM array_size_bytes = 64 * 1024 * 1024 elapsed = c_lib.measure_latency_random(array_size_bytes, iterations) if elapsed <= 0: return {"error": "Benchmark failed"} latency_ns = (elapsed / iterations) * 1e9 return { "test": "memory_latency_random", "description": "Random access latency (64MB working set, Pointer Chasing)", "iterations": iterations, "total_time_seconds": round(elapsed, 4), "average_latency_ns": round(latency_ns, 2), "score": round(100 / latency_ns * 1000, 2), # Adjusted score scale } def benchmark_sequential_latency(iterations: int = 10000000) -> Dict[str, Any]: """ 内存延迟测试（顺序访问） Uses C helper. """ if not c_lib: return {"error": "C helper not available"} # Same 64MB block array_size_bytes = 64 * 1024 * 1024 elapsed = c_lib.measure_latency_sequential(array_size_bytes, iterations) if elapsed <= 0: return {"error": "Benchmark failed"} latency_ns = (elapsed / iterations) * 1e9 return { "test": "memory_latency_sequential", "description": "Sequential access latency (64MB working set, Strided Read)", "iterations": iterations, "total_time_seconds": round(elapsed, 4), "average_latency_ns": round(latency_ns, 2), "score": round(100 / latency_ns * 1000, 2), } def benchmark_alloc_rate(iterations: int = 1000000) -> Dict[str, Any]: """ 内存分配/释放速率测试 """ if not c_lib: return {"error": "C helper not available"} # Test small allocations (e.g. 1KB) which are common alloc_size = 1024 elapsed = c_lib.measure_alloc_rate(alloc_size, iterations) if elapsed <= 0: return {"error": "Benchmark failed"} ops_per_sec = iterations / elapsed return { "test": "memory_alloc_rate", "description": f"Malloc/Free rate (Size: {alloc_size} bytes)", "iterations": iterations, "ops_per_sec": round(ops_per_sec, 2), "score": round(ops_per_sec / 10000, 2), } def benchmark_cache_latency() -> Dict[str, Any]: """ 缓存层级延迟测试 (L1/L2/L3) Uses C helper pointer chasing with smaller working sets. """ if not c_lib: return {"error": "C helper not available"} results = {} # Approximate sizes. # Must be small enough to fit in cache, but large enough to measure. # Typical: L1=32KB(use 16KB), L2=256KB(use 128KB), L3=8MB+(use 4MB) levels = [ ("L1", 16 * 1024), ("L2", 128 * 1024), ("L3", 4 * 1024 * 1024) ] iterations = 10000000 # 10M iterations for name, size in levels: elapsed = c_lib.measure_latency_random(size, iterations) latency_ns = (elapsed / iterations) * 1e9 for name, size in levels: elapsed = c_lib.measure_latency_random(size, iterations) if elapsed <= 0: # Fallback or error latency_ns = 0.0 else: latency_ns = (elapsed / iterations) * 1e9 results[name] = { "size_bytes": size, "latency_ns": round(latency_ns, 2) } l1_lat = results["L1"]["latency_ns"] score = 0 if l1_lat > 0: score = round(100 / l1_lat * 500, 2) return { "test": "cache_latency", "description": "Cache hierarchy latency (Pointer Chasing)", "levels": results, "score": score } def run_all_memory_benchmarks() -> Dict[str, Any]: """运行所有内存基准测试""" results = { "bandwidth": benchmark_memory_bandwidth(), "latency_random": benchmark_memory_latency(), "latency_sequential": benchmark_sequential_latency(), "cache_latency": benchmark_cache_latency(), "alloc_rate": benchmark_alloc_rate(), } # 计算总分 total_score = sum(r.get("score", 0) for r in results.values()) results["total_score"] = round(total_score, 2) return results