""" ROCm Forge — Code Refactorer Agent Takes analysis results and transforms CUDA code into ROCm/HIP-compatible code. Performs intelligent, context-aware code transformations. """ import re from typing import List, Tuple from knowledge.cuda_mappings import ( CUDA_TO_HIP_API, CUDA_TO_ROCM_LIBS, ENV_VAR_MAPPINGS, HEADER_MAPPINGS, DOCKER_IMAGE_MAPPINGS, CLI_TOOL_MAPPINGS, HARDWARE_AWARE_MAPPINGS, ) class RefactorerAgent: """ Code Refactorer Agent — Transforms CUDA code into ROCm/HIP-compatible code. Uses deterministic, rule-based transformations for reliable migration. """ def __init__(self): self.name = "Code Refactorer Agent" self.changes: List[dict] = [] self.trace_log: List[str] = [] @staticmethod def _confidence(note: str) -> str: """Assign confidence level to a transformation.""" safe = ["Env var", "Path:", "CLI tool", "URL:", "Message:", "Header:", "Base image:", "Docker:", "CLI:"] review = ["Inline CUDA", "API:", "Library:", "Attention:", "Package:", "Variable:"] # If the note matches a safe pattern, it's safe for s in safe: if s in note: return "✅ Safe" for r in review: if r in note: return "⚠️ Review" return "🔧 Manual" def refactor(self, code: str, analysis_result) -> Tuple[str, List[dict], List[str]]: """ Refactor the provided code from CUDA to ROCm. Returns: (refactored_code, list_of_changes, trace_log) """ self.changes = [] self.trace_log = [] refactored = code code_type = analysis_result.code_type self.trace_log.append(f"🔄 Starting refactoring for {code_type} code...") # Apply transformations in order if code_type == "python": refactored = self._refactor_python(refactored, analysis_result) elif code_type == "cpp": refactored = self._refactor_cpp(refactored, analysis_result) elif code_type == "dockerfile": refactored = self._refactor_dockerfile(refactored, analysis_result) elif code_type == "requirements": refactored = self._refactor_requirements(refactored, analysis_result) # Add ROCm header comment refactored = self._add_migration_header(refactored, code_type) self.trace_log.append(f"✅ Refactoring complete — {len(self.changes)} transformations applied") return refactored, self.changes, self.trace_log def _refactor_python(self, code: str, analysis) -> str: """Refactor Python/PyTorch code.""" lines = code.split("\n") result_lines = [] for i, line in enumerate(lines): original = line modified = line change_notes = [] # 1. Replace environment variables for cuda_var, rocm_var in ENV_VAR_MAPPINGS.items(): if cuda_var in modified: modified = modified.replace(cuda_var, rocm_var) change_notes.append(f"Env var: {cuda_var} → {rocm_var}") # 2. Replace CUDA CLI tools in comments/strings for nv_tool, rocm_tool in CLI_TOOL_MAPPINGS.items(): if nv_tool in modified: modified = modified.replace(nv_tool, rocm_tool) change_notes.append(f"CLI tool: {nv_tool} → {rocm_tool}") # 3. Replace CUDA_HOME paths if "/usr/local/cuda" in modified: modified = modified.replace("/usr/local/cuda-12.1", "/opt/rocm") modified = modified.replace("/usr/local/cuda-12.0", "/opt/rocm") modified = modified.replace("/usr/local/cuda-11.8", "/opt/rocm") modified = modified.replace("/usr/local/cuda", "/opt/rocm") change_notes.append("Path: /usr/local/cuda → /opt/rocm") # 4. Replace CUDA version references if "torch.version.cuda" in modified: modified = modified.replace("torch.version.cuda", "torch.version.hip") change_notes.append("API: torch.version.cuda → torch.version.hip") # 5. Replace cuDNN with MIOpen backend references if "torch.backends.cudnn" in modified: # cudnn.benchmark and cudnn.deterministic work on ROCm via MIOpen # but we add a helpful comment if not any("MIOpen" in l for l in result_lines[-3:] if result_lines): change_notes.append("Info: torch.backends.cudnn works on ROCm via MIOpen backend") # 6. Handle flash_attention_2 → sdpa if 'attn_implementation="flash_attention_2"' in modified or "attn_implementation='flash_attention_2'" in modified: modified = modified.replace('"flash_attention_2"', '"sdpa"') modified = modified.replace("'flash_attention_2'", "'sdpa'") change_notes.append("Attention: flash_attention_2 → sdpa (PyTorch native, ROCm compatible)") # 7. Handle enforce_eager for CUDA graphs if "enforce_eager=False" in modified and "CUDA graphs" in line.lower() or "cuda graph" in line.lower(): modified = modified.replace("enforce_eager=False", "enforce_eager=True") change_notes.append("Config: enforce_eager=True (CUDA graphs have limited ROCm support)") # 8. Replace cuda pip install URLs if "download.pytorch.org/whl/cu" in modified: modified = re.sub( r'download\.pytorch\.org/whl/cu\d+', 'download.pytorch.org/whl/rocm6.2', modified ) change_notes.append("URL: PyTorch CUDA wheel → ROCm 6.2 wheel") # 9. Handle bitsandbytes if "bitsandbytes" in modified and "rocm" not in modified.lower(): if "pip install bitsandbytes" in modified: modified = modified.replace("pip install bitsandbytes", "pip install bitsandbytes-rocm") change_notes.append("Package: bitsandbytes → bitsandbytes-rocm") # 10. Replace CUDA-specific error messages if "CUDA is not available" in modified: modified = modified.replace("CUDA is not available", "ROCm/HIP is not available") change_notes.append("Message: Updated error text for ROCm") if "No CUDA GPU" in modified: modified = modified.replace("No CUDA GPU", "No AMD GPU") change_notes.append("Message: Updated error text for AMD") if "CUDA not available" in modified: modified = modified.replace("CUDA not available", "ROCm/HIP not available") change_notes.append("Message: Updated error text for ROCm") if "Install CUDA" in modified: modified = modified.replace("Install CUDA drivers", "Install ROCm drivers") modified = modified.replace("Install CUDA toolkit", "Install ROCm toolkit") change_notes.append("Message: Updated install instructions for ROCm") if "CUDA GPU required" in modified: modified = modified.replace("CUDA GPU required", "AMD GPU with ROCm required") change_notes.append("Message: Updated GPU requirement text") # 11. Transform CUDA C code inside Python strings (inline kernels) cuda_c_inline = { "cuda_runtime.h": "hip/hip_runtime.h", "cuda.h": "hip/hip_runtime.h", "cudaMalloc": "hipMalloc", "cudaFree": "hipFree", "cudaMemcpy": "hipMemcpy", "cudaMemset": "hipMemset", "cudaDeviceSynchronize": "hipDeviceSynchronize", "cudaGetLastError": "hipGetLastError", "cudaGetErrorString": "hipGetErrorString", "cudaError_t": "hipError_t", "cudaSuccess": "hipSuccess", "cudaMemcpyHostToDevice": "hipMemcpyHostToDevice", "cudaMemcpyDeviceToHost": "hipMemcpyDeviceToHost", } for cuda_tok, hip_tok in cuda_c_inline.items(): if cuda_tok in modified: modified = modified.replace(cuda_tok, hip_tok) change_notes.append(f"Inline CUDA: {cuda_tok} → {hip_tok}") # 12. Rename cuda_sources variable to hip_sources if "cuda_sources" in modified and "load_inline" not in modified: modified = modified.replace("cuda_sources", "hip_sources") change_notes.append("Variable: cuda_sources → hip_sources") # Record changes if modified != original: for note in change_notes: self.changes.append({ "line": i + 1, "original": original.strip(), "modified": modified.strip(), "note": note, "confidence": self._confidence(note), }) result_lines.append(modified) self.trace_log.append(f"🐍 Python refactoring: {len(self.changes)} changes applied") return "\n".join(result_lines) def _refactor_cpp(self, code: str, analysis) -> str: """Refactor C/C++ CUDA code to HIP.""" lines = code.split("\n") result_lines = [] for i, line in enumerate(lines): original = line modified = line change_notes = [] # 1. Replace CUDA headers for cuda_header, hip_header in HEADER_MAPPINGS.items(): if cuda_header in modified: modified = modified.replace(cuda_header, hip_header) change_notes.append(f"Header: {cuda_header} → {hip_header}") # 2. Replace CUDA API calls for cuda_api, hip_api in CUDA_TO_HIP_API.items(): if cuda_api in modified: modified = modified.replace(cuda_api, hip_api) change_notes.append(f"API: {cuda_api} → {hip_api}") # 3. Replace library calls for cuda_lib, rocm_lib in CUDA_TO_ROCM_LIBS.items(): if cuda_lib in modified and cuda_lib not in ["nvcc", "NVCC"]: modified = modified.replace(cuda_lib, rocm_lib) change_notes.append(f"Library: {cuda_lib} → {rocm_lib}") # Record changes if modified != original: for note in change_notes: self.changes.append({ "line": i + 1, "original": original.strip(), "modified": modified.strip(), "note": note, "confidence": self._confidence(note), }) result_lines.append(modified) # === Hardware-Aware Refactoring Pass === # Second pass: detect and rewrite Tensor Core / Warp-level patterns final_lines = [] all_code = "\n".join(result_lines) for i, line in enumerate(result_lines): modified = line # WMMA/Tensor Core → rocWMMA/MFMA for hw_key, hw_info in HARDWARE_AWARE_MAPPINGS.items(): if hw_key == "32": # Context-dependent: only replace '32' when near warp operations contexts = hw_info.get("context", []) window = "\n".join(result_lines[max(0, i-2):min(len(result_lines), i+3)]) if any(ctx in window for ctx in contexts): # Replace hardcoded 32 with wavefront-safe value if re.search(r'\b32\b', modified): original = modified modified = re.sub(r'\b32\b', hw_info["replacement"], modified, count=1) if modified != original: self.changes.append({ "line": i + 1, "original": original.strip(), "modified": modified.strip(), "note": hw_info["note"], "confidence": "🔧 Manual", }) else: if hw_key in modified: original = modified modified = modified.replace(hw_key, hw_info["replacement"]) self.changes.append({ "line": i + 1, "original": original.strip(), "modified": modified.strip(), "note": hw_info["note"], "confidence": "🔧 Manual", }) final_lines.append(modified) self.trace_log.append(f"⚙️ C++/CUDA refactoring: {len(self.changes)} changes applied") self.trace_log.append(f"🔬 Hardware-aware pass: Tensor Core → MFMA, Warp → Wavefront checks applied") return "\n".join(final_lines) def _refactor_dockerfile(self, code: str, analysis) -> str: """Refactor Dockerfile from NVIDIA to ROCm.""" lines = code.split("\n") result_lines = [] skip_nvidia_env = False for i, line in enumerate(lines): original = line modified = line change_notes = [] # 1. Replace base images for nv_image, rocm_image in DOCKER_IMAGE_MAPPINGS.items(): if nv_image in modified: modified = modified.replace(nv_image, rocm_image) change_notes.append(f"Base image: NVIDIA → ROCm") # 2. Replace NVIDIA env vars if "NVIDIA_VISIBLE_DEVICES" in modified: modified = "# NVIDIA_VISIBLE_DEVICES not needed — use --device flags in docker run" change_notes.append("Docker: Removed NVIDIA_VISIBLE_DEVICES") if "NVIDIA_DRIVER_CAPABILITIES" in modified: modified = "# NVIDIA_DRIVER_CAPABILITIES not needed on ROCm" change_notes.append("Docker: Removed NVIDIA_DRIVER_CAPABILITIES") # 3. Replace CUDA paths if "CUDA_HOME" in modified: modified = modified.replace("CUDA_HOME", "ROCM_HOME") modified = modified.replace("/usr/local/cuda", "/opt/rocm") change_notes.append("Path: CUDA_HOME → ROCM_HOME, /usr/local/cuda → /opt/rocm") # 4. Replace nvidia-smi if "nvidia-smi" in modified: modified = modified.replace("nvidia-smi", "rocm-smi") change_notes.append("CLI: nvidia-smi → rocm-smi") # 5. Replace PyTorch install commands if "download.pytorch.org/whl/cu" in modified: modified = re.sub( r'download\.pytorch\.org/whl/cu\d+', 'download.pytorch.org/whl/rocm6.2', modified ) change_notes.append("URL: CUDA PyTorch wheel → ROCm wheel") # 6. Replace CUDA verification if 'assert torch.cuda.is_available()' in modified: modified = modified.replace( 'assert torch.cuda.is_available(), \'CUDA not available!\'', 'assert torch.cuda.is_available(), \'ROCm/HIP not available!\'' ) change_notes.append("Message: CUDA → ROCm/HIP") # Record changes if modified != original: for note in change_notes: self.changes.append({ "line": i + 1, "original": original.strip(), "modified": modified.strip(), "note": note, "confidence": self._confidence(note), }) result_lines.append(modified) # Add ROCm-specific Docker instructions at the end result_lines.append("") result_lines.append("# === ROCm Docker Run Command ===") result_lines.append("# docker run --device=/dev/kfd --device=/dev/dri \\") result_lines.append("# --group-add video --group-add render \\") result_lines.append("# --ipc=host --shm-size=16g \\") result_lines.append("# -e HIP_VISIBLE_DEVICES=0 \\") result_lines.append("# your-image-name") self.trace_log.append(f"🐳 Dockerfile refactoring: {len(self.changes)} changes applied") return "\n".join(result_lines) def _refactor_requirements(self, code: str, analysis) -> str: """Refactor requirements.txt for ROCm compatibility.""" lines = code.split("\n") result_lines = [ "# ============================================================", "# ROCm Forge — Migrated Requirements for AMD GPU", "# IMPORTANT: Install PyTorch first:", "# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2", "# Then: pip install -r requirements.txt", "# ============================================================", "", ] from knowledge.cuda_mappings import PIP_PACKAGE_MAPPINGS for i, line in enumerate(lines): stripped = line.strip() if not stripped or stripped.startswith("#"): continue pkg_name = re.split(r'[>=<~!\[]', stripped)[0].strip() if pkg_name in PIP_PACKAGE_MAPPINGS: info = PIP_PACKAGE_MAPPINGS[pkg_name] action = info["action"] if action == "remove": result_lines.append(f"# REMOVED: {stripped} — {info['note']}") self.changes.append({ "line": i + 1, "original": stripped, "modified": f"# REMOVED: {stripped}", "note": info["note"], "confidence": "✅ Safe", }) elif action == "replace": install_cmd = info.get("install", "") result_lines.append(f"# REPLACED: {stripped}") result_lines.append(f"# Install with: {install_cmd}") self.changes.append({ "line": i + 1, "original": stripped, "modified": f"# Install: {install_cmd}", "note": f"Replace with {info.get('replacement', 'ROCm build')}", "confidence": "⚠️ Review", }) elif action == "warning": result_lines.append(f"{stripped} # ⚠️ {info.get('note', 'Check ROCm compatibility')}") self.changes.append({ "line": i + 1, "original": stripped, "modified": f"{stripped} (warning added)", "note": info.get("note", ""), "confidence": "⚠️ Review", }) else: result_lines.append(f"{stripped} # ℹ️ {info.get('note', '')}") else: result_lines.append(stripped) self.trace_log.append(f"📦 Requirements refactoring: {len(self.changes)} changes applied") return "\n".join(result_lines) def _add_migration_header(self, code: str, code_type: str) -> str: """Add ROCm Forge migration header to the code.""" if code_type == "python": header = ( "# ============================================================\n" "# 🔥 Migrated by ROCm Forge — CUDA → AMD ROCm/HIP\n" "# This code has been automatically migrated for AMD GPU compatibility.\n" "# Target: AMD Instinct MI300X / ROCm 6.2\n" "# ============================================================\n\n" ) elif code_type == "cpp": header = ( "// ============================================================\n" "// 🔥 Migrated by ROCm Forge — CUDA → AMD HIP\n" "// This code has been automatically migrated for AMD GPU.\n" "// Compile with: hipcc instead of nvcc\n" "// Target: AMD Instinct MI300X / ROCm 6.2\n" "// ============================================================\n\n" ) elif code_type == "dockerfile": header = ( "# ============================================================\n" "# 🔥 Migrated by ROCm Forge — NVIDIA → AMD ROCm\n" "# This Dockerfile has been migrated for AMD GPU deployment.\n" "# ============================================================\n\n" ) else: header = ( "# ============================================================\n" "# 🔥 Migrated by ROCm Forge — CUDA → AMD ROCm\n" "# ============================================================\n\n" ) return header + code