File size: 16,343 Bytes
7a0c684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 |
from typing import Dict, List, Tuple, Any, Optional
import numpy as np
import threading
import time
import json
from dataclasses import dataclass
from enum import Enum
import duckdb
from huggingface_hub import HfApi, HfFileSystem
from .memory import RegisterFile
from config import get_hf_token_cached
# Initialize token from .env
@dataclass
class WarpBarrier:
"""Represents a synchronization barrier for warps"""
barrier_id: str
num_warps: int
arrived: int = 0
completed: bool = False
lock: threading.Lock = threading.Lock()
condition: threading.Condition = threading.Condition()
class ShuffleMode(Enum):
"""Different modes for warp shuffle operations"""
UP = "up" # Shuffle up relative to caller
DOWN = "down" # Shuffle down relative to caller
XOR = "xor" # Butterfly shuffle pattern
IDX = "idx" # Direct index-based shuffle
BCAST = "bcast" # Broadcast from source lane
class VotingMode(Enum):
"""Different modes for warp voting operations"""
ALL = "all" # True if predicate is true for all active threads
ANY = "any" # True if predicate is true for any active thread
BALLOT = "ballot" # Returns bitmask of true predicates
COUNT = "count" # Returns count of true predicates
class WarpState(Enum):
"""Possible states for a warp"""
READY = "ready" # Ready to execute
RUNNING = "running" # Currently executing
BLOCKED = "blocked" # Waiting for synchronization
YIELDED = "yielded" # Voluntarily yielded execution
COMPLETED = "completed" # Finished execution
class Warp:
"""Represents a group of threads that execute together with advanced synchronization"""
DB_URL = "hf://datasets/Fred808/helium/storage.json"
def __init__(self, warp_id: int, num_threads: int = 32, db_url: Optional[str] = None):
self.warp_id = warp_id
self.num_threads = min(num_threads, 32) # Max 32 threads per warp
self.active_mask = (1 << self.num_threads) - 1 # All threads active initially
self.predicate_mask = (1 << self.num_threads) - 1 # For predicated execution
self.registers = [RegisterFile() for _ in range(self.num_threads)]
self.state = WarpState.READY
# Initialize database connection
self.db_url = db_url or self.DB_URL
self.max_retries = 3
self._connect_with_retries()
self._setup_database()
# Register warp in database
self._register_warp()
def _connect_with_retries(self):
"""Establish database connection with retry logic"""
for attempt in range(self.max_retries):
try:
self.conn = self._init_db_connection()
return
except Exception as e:
if attempt == self.max_retries - 1:
raise RuntimeError(f"Failed to initialize database after {self.max_retries} attempts: {str(e)}")
time.sleep(1)
def _init_db_connection(self) -> duckdb.DuckDBPyConnection:
"""Initialize database connection with HuggingFace configuration"""
# Convert HF URL to S3 path
_, _, owner, dataset, db_file = self.db_url.split('/', 4)
db_path = f"s3://datasets-cached/{owner}/{dataset}/{db_file}"
# Connect to remote database
conn = duckdb.connect(db_path)
conn.execute("INSTALL httpfs;")
conn.execute("LOAD httpfs;")
conn.execute("SET s3_endpoint='s3.us-east-1.amazonaws.com';")
conn.execute("SET s3_use_ssl=true;")
conn.execute("SET s3_url_style='path';")
conn.execute(f"SET s3_access_key_id='{self.HF_TOKEN}';")
conn.execute(f"SET s3_secret_access_key='{self.HF_TOKEN}';")
return conn
def _setup_database(self):
"""Initialize database tables"""
# Warp state table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS warps (
warp_id VARCHAR PRIMARY KEY,
num_threads INTEGER,
active_mask BIGINT,
predicate_mask BIGINT,
state VARCHAR,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP,
state_json JSON
)
""")
# Barrier table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS warp_barriers (
barrier_id VARCHAR PRIMARY KEY,
num_warps INTEGER,
arrived_count INTEGER DEFAULT 0,
completed BOOLEAN DEFAULT false,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
completed_at TIMESTAMP,
state_json JSON
)
""")
# Register table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS warp_registers (
warp_id VARCHAR,
thread_id INTEGER,
register_id INTEGER,
value BLOB,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (warp_id, thread_id, register_id)
)
""")
def _register_warp(self):
"""Register warp in database"""
self.conn.execute("""
INSERT INTO warps (
warp_id, num_threads, active_mask, predicate_mask,
state, state_json
) VALUES (?, ?, ?, ?, ?, ?)
""", [
str(self.warp_id),
self.num_threads,
self.active_mask,
self.predicate_mask,
self.state.value,
{"status": "initialized"}
])
self.pc = 0 # Program counter
# Synchronization
self.barriers: Dict[str, WarpBarrier] = {}
self.lock = threading.Lock()
# Performance tracking
self.cycles_executed = 0
self.last_active_time = time.time()
def get_active_threads(self) -> List[int]:
"""Get indices of currently active threads"""
return [i for i in range(self.num_threads)
if self.active_mask & (1 << i)]
def get_predicated_threads(self) -> List[int]:
"""Get indices of threads that pass predication"""
return [i for i in range(self.num_threads)
if (self.active_mask & (1 << i)) and (self.predicate_mask & (1 << i))]
def set_active_mask(self, mask: int):
"""Set which threads are active"""
with self.lock:
self.active_mask = mask & ((1 << self.num_threads) - 1)
def set_predicate_mask(self, mask: int):
"""Set predication mask for conditional execution"""
with self.lock:
self.predicate_mask = mask & ((1 << self.num_threads) - 1)
def sync(self, barrier_id: str = None):
"""Synchronize all threads in the warp at a barrier"""
if not barrier_id:
barrier_id = f"warp_{self.warp_id}_barrier_{time.time_ns()}"
with self.lock:
if barrier_id not in self.barriers:
self.barriers[barrier_id] = WarpBarrier(
barrier_id=barrier_id,
num_warps=1
)
barrier = self.barriers[barrier_id]
with barrier.lock:
barrier.arrived += 1
if barrier.arrived == barrier.num_warps:
barrier.completed = True
with barrier.condition:
barrier.condition.notify_all()
else:
while not barrier.completed:
with barrier.condition:
barrier.condition.wait()
def vote(self, predicate: List[bool], mode: VotingMode = VotingMode.ALL) -> Any:
"""Perform voting operation across threads"""
active_threads = self.get_predicated_threads()
if not active_threads:
return False if mode != VotingMode.BALLOT else 0
if mode == VotingMode.ALL:
return all(predicate[i] for i in active_threads)
elif mode == VotingMode.ANY:
return any(predicate[i] for i in active_threads)
elif mode == VotingMode.BALLOT:
return sum(1 << i for i in active_threads if predicate[i])
elif mode == VotingMode.COUNT:
return sum(1 for i in active_threads if predicate[i])
def shuffle(self, var: List[Any], mode: ShuffleMode, offset: int) -> List[Any]:
"""Exchange variables between threads using different shuffle patterns"""
active_threads = self.get_predicated_threads()
result = list(var) # Create copy to store results
if mode == ShuffleMode.UP:
# Shift values up by offset
for i in active_threads:
src_idx = (i - offset) % self.num_threads
if src_idx in active_threads:
result[i] = var[src_idx]
elif mode == ShuffleMode.DOWN:
# Shift values down by offset
for i in active_threads:
src_idx = (i + offset) % self.num_threads
if src_idx in active_threads:
result[i] = var[src_idx]
elif mode == ShuffleMode.XOR:
# Butterfly pattern exchange
for i in active_threads:
src_idx = i ^ offset # XOR with offset
if src_idx < self.num_threads and src_idx in active_threads:
result[i] = var[src_idx]
elif mode == ShuffleMode.IDX:
# Direct index-based shuffle
for i in active_threads:
if offset < self.num_threads and offset in active_threads:
result[i] = var[offset]
elif mode == ShuffleMode.BCAST:
# Broadcast from source lane
if offset < self.num_threads and offset in active_threads:
src_val = var[offset]
for i in active_threads:
result[i] = src_val
return result
def execute(self, func: callable, *args, **kwargs):
"""Execute a function across all active threads"""
active_threads = self.get_active_threads()
results = []
for thread_idx in active_threads:
# Set up thread context
ctx = {
'thread_idx': thread_idx,
'warp_id': self.warp_id,
'registers': self.registers[thread_idx]
}
# Execute thread
result = func(ctx, *args, **kwargs)
results.append(result)
return results
class WarpScheduler:
"""Advanced warp scheduler with priority and dependency handling"""
def __init__(self, max_warps: int = 32, max_active_warps: int = 16):
self.max_warps = max_warps
self.max_active_warps = max_active_warps
self.warps: List[Warp] = []
self.active_warps: Dict[int, bool] = {}
self.warp_priorities: Dict[int, int] = {}
self.warp_dependencies: Dict[int, List[int]] = {}
self.lock = threading.Lock()
def create_warp(self, num_threads: int = 32, priority: int = 0) -> Warp:
"""Create a new warp with specified priority"""
with self.lock:
if len(self.warps) >= self.max_warps:
raise RuntimeError("Maximum number of warps reached")
warp_id = len(self.warps)
warp = Warp(warp_id, num_threads)
self.warps.append(warp)
self.active_warps[warp_id] = True
self.warp_priorities[warp_id] = priority
self.warp_dependencies[warp_id] = []
return warp
def set_warp_priority(self, warp_id: int, priority: int):
"""Set execution priority for a warp"""
with self.lock:
if 0 <= warp_id < len(self.warps):
self.warp_priorities[warp_id] = priority
def add_warp_dependency(self, warp_id: int, depends_on: int):
"""Add execution dependency between warps"""
with self.lock:
if 0 <= warp_id < len(self.warps) and 0 <= depends_on < len(self.warps):
self.warp_dependencies[warp_id].append(depends_on)
def remove_warp_dependency(self, warp_id: int, depends_on: int):
"""Remove execution dependency between warps"""
with self.lock:
if 0 <= warp_id < len(self.warps):
try:
self.warp_dependencies[warp_id].remove(depends_on)
except ValueError:
pass
def suspend_warp(self, warp_id: int):
"""Suspend a warp from execution"""
with self.lock:
if 0 <= warp_id < len(self.warps):
self.active_warps[warp_id] = False
self.warps[warp_id].state = WarpState.BLOCKED
def resume_warp(self, warp_id: int):
"""Resume a suspended warp"""
with self.lock:
if 0 <= warp_id < len(self.warps):
self.active_warps[warp_id] = True
self.warps[warp_id].state = WarpState.READY
def synchronize_warps(self, warp_ids: List[int], barrier_id: str = None):
"""Synchronize a group of warps"""
if not barrier_id:
barrier_id = f"barrier_{time.time_ns()}"
# Create barrier
barrier = WarpBarrier(barrier_id=barrier_id, num_warps=len(warp_ids))
# Register barrier with each warp
for warp_id in warp_ids:
if 0 <= warp_id < len(self.warps):
warp = self.warps[warp_id]
warp.barriers[barrier_id] = barrier
# Wait for all warps to reach barrier
for warp_id in warp_ids:
if 0 <= warp_id < len(self.warps):
self.warps[warp_id].sync(barrier_id)
def schedule(self) -> List[Warp]:
"""Schedule warps for execution based on priority and dependencies"""
with self.lock:
ready_warps = []
# Get warps that are ready to execute
for warp_id, warp in enumerate(self.warps):
if not self.active_warps.get(warp_id, False):
continue
# Check dependencies
dependencies_met = all(
self.warps[dep_id].state == WarpState.COMPLETED
for dep_id in self.warp_dependencies.get(warp_id, [])
)
if dependencies_met and warp.state == WarpState.READY:
ready_warps.append((warp_id, self.warp_priorities.get(warp_id, 0)))
# Sort by priority (higher numbers = higher priority)
ready_warps.sort(key=lambda x: x[1], reverse=True)
# Return warps up to max_active_warps
return [self.warps[warp_id]
for warp_id, _ in ready_warps[:self.max_active_warps]]
def execute_warps(self, func: callable, *args, **kwargs):
"""Execute function across all active warps with scheduling"""
results = []
scheduled_warps = self.schedule()
for warp in scheduled_warps:
warp.state = WarpState.RUNNING
result = warp.execute(func, *args, **kwargs)
results.extend(result)
warp.last_active_time = time.time()
warp.cycles_executed += 1
warp.state = WarpState.READY
return results
|