env_server.py

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()