File size: 9,866 Bytes
c780f59 aebc8f0 c780f59 aebc8f0 c780f59 aebc8f0 c780f59 d5c6f39 aebc8f0 d5c6f39 c780f59 d5c6f39 c780f59 7e03473 c780f59 5fe4705 c780f59 d5c6f39 | 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 | from typing import List, Optional, Dict, Any
from models import Action, StepResult, ResetRequest, StepRequest, EnvState, Observation, Reward
from fastapi import FastAPI, HTTPException
import random
TASKS: Dict[str, Dict[str, Any]] ={
"vector_add_easy": {
"name": "Vector Addition Kernel Optimization",
"difficulty": "easy",
"grader": "deterministic_rule_based",
"max_steps": 5,
"target_speedup": 1.8,
"baseline_code": """extern "C" __global__ void vector_add(const float* a, const float* b, float* c, int n)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) c[idx] = a[idx] + b[idx];
}""",
"checks": {
"coalesced_memory": "Use memory-coalesced indexing",
"vectorized_loads": "Use vectorized loads/stores (float2/float4)",
"bounds_safe": "Keep safe boundary checks",
},
},
"matmul_medium": {
"name": "Matrix Multiplication Kernel Optimization",
"difficulty": "medium",
"grader": "deterministic_rule_based",
"max_steps": 6,
"target_speedup": 3.0,
"baseline_code": """extern "C" __global__ void matmul(const float* A, const float* B, float* C, int N)
{
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < N && col < N) {
float sum = 0.0f;
for (int k = 0; k < N; k++) sum += A[row * N + k] * B[k * N + col];
C[row * N + col] = sum;
}
}""",
"checks": {
"shared_tiling": "Use shared-memory tiling",
"synchronization": "Synchronize tiles with __syncthreads",
"register_accumulation": "Accumulate partial sums in registers",
},
},
"reduction_hard": {
"name": "Reduction Kernel Optimization",
"difficulty": "hard",
"grader": "deterministic_rule_based",
"max_steps":7,
"target_speedup": 3.5,
"baseline_code": """extern "C" __global__ void reduce_sum(const float* input, float* output, int n)
{
extern __shared__ float sdata[];
int tid = threadIdx.x;
int i = blockIdx.x * blockDim.x + threadIdx.x;
sdata[tid] = (i < n) ? input[i] : 0.0f;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) sdata[tid] += sdata[tid + s];
__syncthreads();
}
if (tid == 0) output[blockIdx.x] = sdata[0];
}""",
"checks": {
"warp_primitive": "Use warp-level primitive (e.g., __shfl_down_sync)",
"bank_conflict_reduction": "Reduce shared-memory bank conflicts",
"unrolled_reduction": "Use partial unrolling for final reduction",
},
}
}
def check_passed(check_id:str, code_lower:str) ->bool:
if check_id =="coalesced_memory":
return "idx" in code_lower and ("blockidx.x" in code_lower or "threadidx.x" in code_lower)
if check_id == "vectorized_loads":
return "float4" in code_lower or "float2" in code_lower
if check_id == "bounds_safe":
return "if" in code_lower and "< n" in code_lower
if check_id == "shared_tiling":
return "__shared__" in code_lower
if check_id == "synchronization":
return "__syncthreads" in code_lower
if check_id == "register_accumulation":
return "sum" in code_lower or "acc" in code_lower
if check_id == "warp_primitive":
return "__shfl_down_sync" in code_lower or "__shfl_sync" in code_lower
if check_id =="bank_conflict_reduction":
return "pad" in code_lower or "bank" in code_lower or "+ 1" in code_lower
if check_id == "unrolled_reduction":
return "#pragma unroll" in code_lower or "unroll" in code_lower
return False
def to_observation(task_id:str, state:EnvState)->Observation:
task = TASKS[task_id]
pending = [desc for cid, desc in task["checks"].items() if cid not in set(state.completed_checks)]
return Observation(task_id=task_id, task_name=task["name"], difficulty=task["difficulty"], baseline_code=task["baseline_code"], current_best_code=state.best_code or task["baseline_code"], current_best_speedup=state.best_speedup, step_count=state.step_count, max_steps=state.max_steps, pending_checks=pending, completed_checks=[task["checks"][cid] for cid in state.completed_checks if cid in task["checks"]], done=(len(pending) == 0 or state.step_count >= state.max_steps))
def grade_episode(task_id:str, completed_checks:List[str], best_speedup:float, step_count:int, max_steps:int)->float:
task=TASKS[task_id]
completion =len(completed_checks) / max(len(task["checks"]),1)
speedup_score = min(best_speedup /task["target_speedup"],1.0)
efficiency = max(0.0, 1.0 - ((step_count - 1) / max(max_steps, 1)))
return round(max(0.0, min(1.0, 0.5 * completion + 0.35 * speedup_score + 0.15 * efficiency)), 4)
class KernelOptimization_env:
def __init__(self):
self.state =EnvState(initialized=False)
self.current_task_id: Optional[str]=None
def reset(self, task_id:Optional[str]=None)->Dict[str, Any]:
if task_id and task_id not in TASKS:
raise HTTPException(status_code=400, detail=f"unknown task_id: {task_id}")
self.current_task_id =task_id or random.choice(list(TASKS.keys()))
task= TASKS[self.current_task_id]
self.state =EnvState(initialized=True, task_id=self.current_task_id, step_count=0, max_steps=task["max_steps"], total_reward=0.0, best_code=task["baseline_code"], best_speedup=1.0, completed_checks=[], action_history=[])
return {
"observation": to_observation(self.current_task_id, self.state).model_dump(),
"info": {
"task_id": self.current_task_id,
"task_name": task["name"],
"difficulty": task["difficulty"],
"grader": task["grader"],
"max_steps": task["max_steps"],
"target_speedup": task["target_speedup"],
"checks": task["checks"],
},
}
def step(self, action:Action) ->StepResult:
if not self.state.initialized or not self.current_task_id:
raise HTTPException(status_code=400, detail="Environment not initialized. Call /reset first.")
self.state.step_count += 1
code = action.optimized_code or ""
code_lower = code.lower()
compile_ok = "__global__" in code_lower and "{" in code and "}" in code
completed = set(self.state.completed_checks)
newly_completed = {cid for cid in TASKS[self.current_task_id]["checks"] if cid not in completed and check_passed(cid, code_lower)}
completed.update(newly_completed)
self.state.completed_checks = sorted(completed)
completion_ratio = len(completed) / max(len(TASKS[self.current_task_id]["checks"]), 1)
max_reasonable_speedup = 1.0 + completion_ratio * 3.0
if action.expected_speedup is None:
est_speedup = round(max_reasonable_speedup, 3)
else:
est_speedup = round(max(1.0, min(action.expected_speedup, max_reasonable_speedup)), 3)
if est_speedup > self.state.best_speedup:
self.state.best_speedup = est_speedup
self.state.best_code = code
progress = 0.22 * len(newly_completed)
quality = 0.18 * min(self.state.best_speedup / TASKS[self.current_task_id]["target_speedup"], 1.0)
penalty = 0.0
if not compile_ok:
penalty -= 0.25
if not newly_completed:
penalty -= 0.08
reward_value = max(0.0, min(1.0, progress + quality + penalty))
self.state.total_reward += reward_value
self.state.action_history.append(
{
"step": self.state.step_count,
"newly_completed": sorted(newly_completed),
"compile_ok": compile_ok,
"estimated_speedup": est_speedup,
"reward": reward_value,
}
)
obs =to_observation(self.current_task_id, self.state)
info: Dict[str, Any] = { "compile_ok": compile_ok, "estimated_speedup": est_speedup}
if obs.done:
info["final_score"] = grade_episode(
self.current_task_id, self.state.completed_checks, self.state.best_speedup, self.state.step_count, self.state.max_steps
)
return StepResult(
observation=obs,
reward=Reward(
value=round(reward_value, 4),
components={"progress": round(progress, 4), "quality": round(quality, 4), "penalty": round(penalty, 4)},
),
done=obs.done,
info=info,
)
def state_dict(self)->Dict[str, Any]:
data = self.state.model_dump()
if self.current_task_id:
data["task_name"] = TASKS[self.current_task_id]["name"]
data["difficulty"] = TASKS[self.current_task_id]["difficulty"]
data["grader_score"] = grade_episode(
self.current_task_id, self.state.completed_checks, self.state.best_speedup, self.state.step_count, self.state.max_steps
)
return data
env=KernelOptimization_env()
app=FastAPI(title="Kernel Optimization", version="1.0.0")
@app.get("/health")
def health_check():
return {"status": "healthy", "service": "kernel-optimization-openenv"}
@app.post("/reset")
def reset(request: ResetRequest | None = None):
return env.reset(task_id=request.task_id if request else None)
@app.post("/step")
def step(request: StepRequest):
return env.step(request.action).model_dump()
@app.get("/state")
def state():
return env.state_dict()
|