wafer-lsp 0.1.13__py3-none-any.whl

wafer_lsp/handlers/hip_diagnostics.py

@@ -0,0 +1,303 @@
+"""
+HIP Diagnostics Provider.
+
+Generates warnings for common wavefront-related mistakes in HIP code.
+AMD GPUs use 64-thread wavefronts (CDNA) vs CUDA's 32-thread warps,
+and code written for CUDA often has hard-coded assumptions that break on AMD.
+
+Diagnostic Codes:
+- HIP001: Incorrect wavefront size assumption (using 32 instead of 64)
+- HIP002: Wavefront intrinsic misuse
+- HIP003: Ballot result handling errors (32-bit vs 64-bit)
+- HIP004: Lane/wavefront index calculation errors
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+from lsprotocol.types import (
+    Diagnostic,
+    DiagnosticSeverity,
+    DiagnosticTag,
+    Position,
+    Range,
+)
+
+from ..parsers.hip_parser import HIPParser, WavefrontPattern
+
+
+@dataclass(frozen=True)
+class HIPDiagnosticInfo:
+    """Information about a HIP diagnostic."""
+    code: str
+    message: str
+    severity: DiagnosticSeverity
+    suggestion: str
+    doc_link: str | None = None
+
+
+# Diagnostic code to info mapping
+DIAGNOSTIC_INFO: dict[str, HIPDiagnosticInfo] = {
+    "HIP001": HIPDiagnosticInfo(
+        code="HIP001",
+        message="Potential wavefront size mismatch: AMD GPUs use 64-thread wavefronts, not 32.",
+        severity=DiagnosticSeverity.Warning,
+        suggestion="Use `warpSize` variable or `__AMDGCN_WAVEFRONT_SIZE__` macro instead of hard-coding 32.",
+        doc_link="https://rocm.docs.amd.com/en/latest/",
+    ),
+    "HIP002": HIPDiagnosticInfo(
+        code="HIP002",
+        message="Wavefront intrinsic may behave differently on AMD GPUs.",
+        severity=DiagnosticSeverity.Warning,
+        suggestion="AMD wavefronts have 64 threads. Shuffle/ballot operations operate on the full wavefront. Adjust your code accordingly.",
+    ),
+    "HIP003": HIPDiagnosticInfo(
+        code="HIP003",
+        message="Ballot result comparison uses 32-bit mask, but AMD's __ballot returns 64-bit.",
+        severity=DiagnosticSeverity.Warning,
+        suggestion="Use `0xFFFFFFFFFFFFFFFFull` for 64-bit comparison. Use `__popcll()` to count bits.",
+    ),
+    "HIP004": HIPDiagnosticInfo(
+        code="HIP004",
+        message="Lane/wavefront index calculation assumes 32-thread warps.",
+        severity=DiagnosticSeverity.Warning,
+        suggestion="Use `warpSize` (64 on AMD CDNA) instead of hard-coding 32 for lane/warp calculations.",
+    ),
+}
+
+
+# Map pattern types to diagnostic codes
+PATTERN_TO_DIAGNOSTIC: dict[str, str] = {
+    "warp_size_32": "HIP001",
+    "shuffle_mask": "HIP002",
+    "ballot_32bit": "HIP003",
+    "lane_calc_32": "HIP004",
+}
+
+
+class HIPDiagnosticsProvider:
+    """Provides diagnostics for HIP code.
+    
+    Detects patterns that indicate potential issues when running on AMD GPUs,
+    particularly around wavefront size differences (64 vs 32).
+    """
+
+    def __init__(self):
+        self._parser = HIPParser()
+
+    def get_diagnostics(
+        self,
+        content: str,
+        uri: str,
+        enable_wavefront_diagnostics: bool = True,
+    ) -> list[Diagnostic]:
+        """Generate diagnostics for HIP code.
+        
+        Args:
+            content: The document content
+            uri: The document URI
+            enable_wavefront_diagnostics: Whether to check for wavefront issues
+            
+        Returns:
+            List of LSP Diagnostic objects
+        """
+        diagnostics: list[Diagnostic] = []
+        
+        if not enable_wavefront_diagnostics:
+            return diagnostics
+
+        # Parse the file to find wavefront patterns
+        parsed = self._parser.parse_file(content)
+        wavefront_patterns = parsed.get("wavefront_patterns", [])
+        
+        lines = content.split("\n")
+        
+        for pattern in wavefront_patterns:
+            diagnostic = self._create_diagnostic_from_pattern(pattern, lines)
+            if diagnostic:
+                diagnostics.append(diagnostic)
+        
+        # Additional checks that require more context
+        diagnostics.extend(self._check_warp_reduction_patterns(content, lines))
+        diagnostics.extend(self._check_shuffle_width_parameters(content, lines))
+        
+        return diagnostics
+
+    def _create_diagnostic_from_pattern(
+        self,
+        pattern: WavefrontPattern,
+        lines: list[str],
+    ) -> Diagnostic | None:
+        """Create a diagnostic from a detected wavefront pattern."""
+        diag_code = PATTERN_TO_DIAGNOSTIC.get(pattern.pattern_type)
+        if not diag_code:
+            return None
+
+        diag_info = DIAGNOSTIC_INFO.get(diag_code)
+        if not diag_info:
+            return None
+
+        # Find the problematic value in the line to highlight it
+        line = lines[pattern.line] if pattern.line < len(lines) else ""
+        start_char = line.find(pattern.problematic_value)
+        if start_char == -1:
+            start_char = 0
+        end_char = start_char + len(pattern.problematic_value)
+
+        range_ = Range(
+            start=Position(line=pattern.line, character=start_char),
+            end=Position(line=pattern.line, character=end_char),
+        )
+
+        message = self._format_diagnostic_message(
+            diag_info,
+            pattern.problematic_value,
+            pattern.code_snippet,
+        )
+
+        return Diagnostic(
+            range=range_,
+            severity=diag_info.severity,
+            code=diag_info.code,
+            source="wafer-hip",
+            message=message,
+            data={"suggestion": diag_info.suggestion},
+        )
+
+    def _format_diagnostic_message(
+        self,
+        diag_info: HIPDiagnosticInfo,
+        problematic_value: str,
+        code_snippet: str,
+    ) -> str:
+        """Format the diagnostic message with context."""
+        msg = diag_info.message
+        
+        # Add specific guidance based on the problematic value
+        if problematic_value == "32":
+            msg += f"\n\nFound: `{problematic_value}` (assumes CUDA warp size)"
+            msg += "\nExpected: `warpSize` or `64` for AMD CDNA GPUs"
+        elif problematic_value == "31":
+            msg += f"\n\nFound: `& {problematic_value}` (assumes 32-thread warp)"
+            msg += "\nExpected: `& (warpSize - 1)` or `& 63` for AMD CDNA GPUs"
+        elif problematic_value == "0xFFFFFFFF":
+            msg += f"\n\nFound: `{problematic_value}` (32-bit mask)"
+            msg += "\nExpected: `0xFFFFFFFFFFFFFFFFull` (64-bit mask for AMD)"
+        
+        msg += f"\n\n💡 {diag_info.suggestion}"
+        
+        return msg
+
+    def _check_warp_reduction_patterns(
+        self,
+        content: str,
+        lines: list[str],
+    ) -> list[Diagnostic]:
+        """Check for warp reduction patterns that start from 16 instead of 32."""
+        diagnostics: list[Diagnostic] = []
+        
+        # Look for reduction loops that start from 16 (half of CUDA warp)
+        # This is a common CUDA pattern: for (int i = 16; i > 0; i >>= 1)
+        import re
+        
+        pattern = re.compile(
+            r'for\s*\(\s*(?:int|unsigned)\s+\w+\s*=\s*16\s*;[^;]*;\s*\w+\s*>>=\s*1\s*\)'
+        )
+        
+        for i, line in enumerate(lines):
+            match = pattern.search(line)
+            if match:
+                diag_info = DIAGNOSTIC_INFO["HIP001"]
+                
+                range_ = Range(
+                    start=Position(line=i, character=match.start()),
+                    end=Position(line=i, character=match.end()),
+                )
+                
+                diagnostics.append(Diagnostic(
+                    range=range_,
+                    severity=DiagnosticSeverity.Warning,
+                    code="HIP001",
+                    source="wafer-hip",
+                    message=(
+                        "Warp reduction loop starts from 16 (half CUDA warp).\n\n"
+                        "For AMD's 64-thread wavefronts, start from 32:\n"
+                        "`for (int offset = 32; offset > 0; offset >>= 1)`\n\n"
+                        "Or better, use `warpSize / 2` for portability."
+                    ),
+                ))
+        
+        return diagnostics
+
+    def _check_shuffle_width_parameters(
+        self,
+        content: str,
+        lines: list[str],
+    ) -> list[Diagnostic]:
+        """Check for shuffle operations with explicit width=32."""
+        diagnostics: list[Diagnostic] = []
+        
+        import re
+        
+        # Pattern for __shfl variants with explicit width parameter of 32
+        # __shfl_down(val, offset, 32)
+        patterns = [
+            (re.compile(r'__shfl_down\s*\([^)]*,\s*(\d+)\s*,\s*32\s*\)'), "__shfl_down"),
+            (re.compile(r'__shfl_up\s*\([^)]*,\s*(\d+)\s*,\s*32\s*\)'), "__shfl_up"),
+            (re.compile(r'__shfl_xor\s*\([^)]*,\s*(\d+)\s*,\s*32\s*\)'), "__shfl_xor"),
+            (re.compile(r'__shfl\s*\([^)]*,\s*(\d+)\s*,\s*32\s*\)'), "__shfl"),
+        ]
+        
+        for i, line in enumerate(lines):
+            for pattern, intrinsic_name in patterns:
+                match = pattern.search(line)
+                if match:
+                    # Find position of "32" in the match
+                    match_text = match.group(0)
+                    pos_32 = match_text.rfind("32")
+                    start_char = match.start() + pos_32
+                    
+                    range_ = Range(
+                        start=Position(line=i, character=start_char),
+                        end=Position(line=i, character=start_char + 2),
+                    )
+                    
+                    diagnostics.append(Diagnostic(
+                        range=range_,
+                        severity=DiagnosticSeverity.Warning,
+                        code="HIP002",
+                        source="wafer-hip",
+                        message=(
+                            f"`{intrinsic_name}` with explicit width=32 assumes CUDA warp size.\n\n"
+                            f"On AMD CDNA GPUs, wavefront size is 64. Either:\n"
+                            f"• Remove the width parameter (defaults to `warpSize`)\n"
+                            f"• Use `warpSize` as the width parameter\n\n"
+                            f"Example: `{intrinsic_name}(val, offset)` or `{intrinsic_name}(val, offset, warpSize)`"
+                        ),
+                    ))
+        
+        return diagnostics
+
+
+def create_hip_diagnostics_provider() -> HIPDiagnosticsProvider:
+    """Create a HIP diagnostics provider instance."""
+    return HIPDiagnosticsProvider()
+
+
+def get_hip_diagnostics(
+    content: str,
+    uri: str,
+    enable_wavefront_diagnostics: bool = True,
+) -> list[Diagnostic]:
+    """Convenience function to get HIP diagnostics for a file.
+    
+    Args:
+        content: The document content
+        uri: The document URI
+        enable_wavefront_diagnostics: Whether to check for wavefront issues
+        
+    Returns:
+        List of LSP Diagnostic objects
+    """
+    provider = create_hip_diagnostics_provider()
+    return provider.get_diagnostics(content, uri, enable_wavefront_diagnostics)

wafer_lsp/handlers/hover.py

@@ -0,0 +1,251 @@
+
+from lsprotocol.types import Hover, MarkupContent, MarkupKind, Position
+
+from ..analyzers.compiler_integration import get_analysis_for_kernel
+from ..analyzers.docs_index import DocsIndex
+from ..languages.registry import get_language_registry
+from ..languages.types import KernelInfo, LayoutInfo
+from ..services.hip_hover_service import create_hip_hover_service
+from ..utils.lsp_helpers import get_decorator_at_position, get_word_at_position
+
+# Lazy-initialized HIP hover service
+_hip_hover_service = None
+
+
+def _get_hip_hover_service():
+    """Get or create the HIP hover service singleton."""
+    global _hip_hover_service
+    if _hip_hover_service is None:
+        _hip_hover_service = create_hip_hover_service()
+    return _hip_hover_service
+
+
+def find_kernel_at_position(
+    content: str, position: Position, uri: str
+) -> KernelInfo | None:
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return None
+
+    word = get_word_at_position(content, position)
+
+    for kernel in language_info.kernels:
+        if kernel.name == word:
+            if position.line >= kernel.line:
+                return kernel
+
+    return None
+
+
+def find_layout_at_position(
+    content: str, position: Position, uri: str
+) -> LayoutInfo | None:
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return None
+
+    word = get_word_at_position(content, position)
+
+    for layout in language_info.layouts:
+        if layout.name == word:
+            return layout
+
+    return None
+
+
+def format_kernel_hover(kernel: KernelInfo, analysis: dict | None) -> str:
+    registry = get_language_registry()
+    language_name = registry.get_language_name(kernel.language) or kernel.language
+
+    lines = [f"**GPU Kernel: {kernel.name}**", f"*Language: {language_name}*", ""]
+
+    if kernel.docstring:
+        lines.append(kernel.docstring)
+        lines.append("")
+
+    if kernel.parameters:
+        params_str = ", ".join(kernel.parameters)
+        lines.append(f"**Parameters:** `{params_str}`")
+        lines.append("")
+
+    if analysis:
+        lines.append("**Analysis:**")
+        if "layouts" in analysis:
+            lines.append(f"- Layouts: {analysis['layouts']}")
+        if "memory_paths" in analysis:
+            lines.append(f"- Memory paths: {analysis['memory_paths']}")
+        if "pipeline_stages" in analysis:
+            lines.append(f"- Pipeline stages: {analysis['pipeline_stages']}")
+
+    return "\n".join(lines)
+
+
+def format_kernel_hover_basic(kernel: KernelInfo) -> str:
+    return format_kernel_hover(kernel, None)
+
+
+def format_decorator_hover(decorator_name: str, function_name: str | None = None) -> str:
+    lines = []
+
+    if decorator_name == "cute.kernel" or decorator_name == "kernel":
+        lines.append("**@cute.kernel**")
+        lines.append("")
+        lines.append("CuTeDSL kernel decorator. Marks a function as a GPU kernel.")
+        lines.append("")
+        lines.append("**Usage:**")
+        lines.append("```python")
+        lines.append("@cute.kernel")
+        lines.append("def my_kernel(a: cute.Tensor, b: cute.Tensor):")
+        lines.append("    # Kernel implementation")
+        lines.append("    pass")
+        lines.append("```")
+        lines.append("")
+        lines.append("**Features:**")
+        lines.append("- Automatic GPU code generation")
+        lines.append("- Tensor layout optimization")
+        lines.append("- Memory access pattern analysis")
+
+        if function_name:
+            lines.append("")
+            lines.append(f"Applied to: `{function_name}()`")
+
+    elif decorator_name == "cute.struct" or decorator_name == "struct":
+        lines.append("**@cute.struct**")
+        lines.append("")
+        lines.append("CuTeDSL struct decorator. Marks a class as a GPU struct.")
+        lines.append("")
+        lines.append("**Usage:**")
+        lines.append("```python")
+        lines.append("@cute.struct")
+        lines.append("class MyStruct:")
+        lines.append("    field1: int")
+        lines.append("    field2: float")
+        lines.append("```")
+
+        if function_name:
+            lines.append("")
+            lines.append(f"Applied to: `{function_name}`")
+
+    else:
+        lines.append(f"**{decorator_name}**")
+        lines.append("")
+        lines.append("CuTeDSL decorator")
+
+    docs = DocsIndex()
+    doc_link = docs.get_doc_for_concept("kernel" if "kernel" in decorator_name else "struct")
+    if doc_link:
+        lines.append("")
+        lines.append(f"[Documentation]({doc_link})")
+
+    return "\n".join(lines)
+
+
+def handle_hover(uri: str, position: Position, content: str) -> Hover | None:
+    """Handle hover requests for all supported languages.
+    
+    For HIP/CUDA/C++ files, provides rich documentation for:
+    - HIP API functions (hipMalloc, hipMemcpy, etc.)
+    - Memory qualifiers (__device__, __shared__, etc.)
+    - Wavefront intrinsics (__shfl, __ballot, etc.)
+    - Thread indexing (threadIdx, blockIdx, etc.)
+    - Kernel and device functions
+    
+    For CuTeDSL (Python) files:
+    - Decorators (@cute.kernel, @cute.struct)
+    - Layouts and kernel functions
+    """
+    # Detect language to route to appropriate hover handler
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+    
+    # For HIP/CUDA/C++ files, try HIP hover service first
+    if language_info and language_info.language in ("hip", "cuda", "cpp"):
+        hip_service = _get_hip_hover_service()
+        hip_hover = hip_service.get_hover(content, position, uri)
+        if hip_hover:
+            return hip_hover
+    
+    # CuTeDSL decorator handling
+    decorator_info = get_decorator_at_position(content, position)
+    if decorator_info:
+        decorator_name, function_line = decorator_info
+
+        function_name = None
+        lines = content.split("\n")
+        if function_line < len(lines):
+            func_line = lines[function_line].strip()
+            if func_line.startswith("def "):
+                func_name_start = func_line.find("def ") + 4
+                func_name_end = func_line.find("(", func_name_start)
+                if func_name_end > func_name_start:
+                    function_name = func_line[func_name_start:func_name_end].strip()
+            elif func_line.startswith("class "):
+                class_name_start = func_line.find("class ") + 6
+                class_name_end = func_line.find(":", class_name_start)
+                if class_name_end > class_name_start:
+                    function_name = func_line[class_name_start:class_name_end].strip()
+
+        hover_content = format_decorator_hover(decorator_name, function_name)
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    # CuTeDSL cute module
+    word = get_word_at_position(content, position)
+    if word == "cute" or word.startswith("cute."):
+        hover_lines = [
+            "**cutlass.cute**",
+            "",
+            "CuTeDSL (CUDA Unified Tensor Expression) library for GPU programming.",
+            "",
+            "**Key Features:**",
+            "- `@cute.kernel` - Define GPU kernels",
+            "- `@cute.struct` - Define GPU structs",
+            "- `cute.make_layout()` - Create tensor layouts",
+            "- `cute.Tensor` - Tensor type annotations",
+            "",
+            "[Documentation](https://github.com/NVIDIA/cutlass)"
+        ]
+        hover_content = "\n".join(hover_lines)
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    # Kernel hover (all languages)
+    kernel = find_kernel_at_position(content, position, uri)
+    if kernel:
+        analysis = get_analysis_for_kernel(uri, kernel.name)
+
+        if analysis:
+            hover_content = format_kernel_hover(kernel, analysis)
+        else:
+            hover_content = format_kernel_hover_basic(kernel)
+
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    # Layout hover (CuTeDSL)
+    layout = find_layout_at_position(content, position, uri)
+    if layout:
+        docs = DocsIndex()
+        doc_link = docs.get_doc_for_concept("layout")
+
+        hover_lines = [
+            f"**Layout: {layout.name}**",
+            ""
+        ]
+
+        if layout.shape:
+            hover_lines.append(f"Shape: `{layout.shape}`")
+        if layout.stride:
+            hover_lines.append(f"Stride: `{layout.stride}`")
+
+        if doc_link:
+            hover_lines.append("")
+            hover_lines.append(f"[Documentation]({doc_link})")
+
+        hover_content = "\n".join(hover_lines)
+
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    # No hover info found
+    return None