wafer-lsp 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

wafer_lsp/handlers/__init__.py

@@ -0,0 +1,30 @@
+from .code_action import handle_code_action
+from .code_lens import handle_code_lens
+from .completion import handle_completion
+from .diagnostics import handle_diagnostics
+from .document_symbol import handle_document_symbol
+from .hip_diagnostics import (
+    HIPDiagnosticsProvider,
+    create_hip_diagnostics_provider,
+    get_hip_diagnostics,
+)
+from .hover import handle_hover
+from .inlay_hint import handle_inlay_hint
+from .semantic_tokens import handle_semantic_tokens, SEMANTIC_TOKENS_LEGEND
+from .workspace_symbol import handle_workspace_symbol
+
+__all__ = [
+    "handle_code_action",
+    "handle_code_lens",
+    "handle_completion",
+    "handle_diagnostics",
+    "handle_document_symbol",
+    "handle_hover",
+    "handle_inlay_hint",
+    "handle_semantic_tokens",
+    "handle_workspace_symbol",
+    "SEMANTIC_TOKENS_LEGEND",
+    "HIPDiagnosticsProvider",
+    "create_hip_diagnostics_provider",
+    "get_hip_diagnostics",
+]

wafer_lsp/handlers/diagnostics.py

@@ -2,9 +2,24 @@
 from lsprotocol.types import Diagnostic
 
 from ..languages.registry import get_language_registry
+from .hip_diagnostics import get_hip_diagnostics
 
 
-def handle_diagnostics(uri: str, content: str) -> list[Diagnostic]:
+def handle_diagnostics(
+    uri: str,
+    content: str,
+    enable_wavefront_diagnostics: bool = True,
+) -> list[Diagnostic]:
+    """Handle diagnostics for a document.
+    
+    Args:
+        uri: Document URI
+        content: Document content
+        enable_wavefront_diagnostics: Whether to enable HIP wavefront warnings
+        
+    Returns:
+        List of diagnostics
+    """
     diagnostics: list[Diagnostic] = []
 
     registry = get_language_registry()
@@ -13,4 +28,14 @@ def handle_diagnostics(uri: str, content: str) -> list[Diagnostic]:
     if not language_info:
         return diagnostics
 
+    # Add HIP-specific diagnostics for HIP, CUDA, and C++ files
+    # (C++ files might contain HIP code if they have HIP markers)
+    if language_info.language in ("hip", "cuda", "cpp"):
+        hip_diagnostics = get_hip_diagnostics(
+            content,
+            uri,
+            enable_wavefront_diagnostics=enable_wavefront_diagnostics,
+        )
+        diagnostics.extend(hip_diagnostics)
+
     return diagnostics

wafer_lsp/handlers/document_symbol.py

@@ -12,9 +12,10 @@ def handle_document_symbol(uri: str, content: str) -> list[DocumentSymbol]:
         return []
 
     symbols: list[DocumentSymbol] = []
+    lines = content.split("\n")
 
+    # Kernels
     for kernel in language_info.kernels:
-        lines = content.split("\n")
         kernel_line = lines[kernel.line] if kernel.line < len(lines) else ""
         name_start = kernel_line.find(kernel.name)
         name_end = name_start + len(kernel.name) if name_start >= 0 else 0
@@ -28,16 +29,24 @@ def handle_document_symbol(uri: str, content: str) -> list[DocumentSymbol]:
             end=Position(line=min(kernel.line + 10, len(lines) - 1), character=0)
         )
 
+        # Different detail based on language
+        if kernel.language == "hip":
+            detail = "🚀 HIP Kernel (AMD GPU)"
+        elif kernel.language in ("cuda", "cpp"):
+            detail = "🚀 CUDA Kernel"
+        else:
+            detail = f"GPU Kernel ({registry.get_language_name(kernel.language)})"
+
         symbols.append(DocumentSymbol(
             name=kernel.name,
             kind=SymbolKind.Function,
             range=full_range,
             selection_range=selection_range,
-            detail=f"GPU Kernel ({registry.get_language_name(kernel.language)})",
+            detail=detail,
         ))
 
+    # Layouts (CuTeDSL)
     for layout in language_info.layouts:
-        lines = content.split("\n")
         layout_line = lines[layout.line] if layout.line < len(lines) else ""
         name_start = layout_line.find(layout.name)
         name_end = name_start + len(layout.name) if name_start >= 0 else 0
@@ -61,8 +70,8 @@ def handle_document_symbol(uri: str, content: str) -> list[DocumentSymbol]:
             detail=detail,
         ))
 
+    # Structs
     for struct in language_info.structs:
-        lines = content.split("\n")
         struct_line = lines[struct.line] if struct.line < len(lines) else ""
         name_start = struct_line.find(struct.name)
         name_end = name_start + len(struct.name) if name_start >= 0 else 0
@@ -84,4 +93,84 @@ def handle_document_symbol(uri: str, content: str) -> list[DocumentSymbol]:
             detail=f"Struct ({registry.get_language_name(struct.language)})",
         ))
 
+    # HIP-specific: Device functions and shared memory
+    if language_info.language in ("hip", "cuda", "cpp"):
+        symbols.extend(_get_hip_symbols(language_info.raw_data, lines))
+
+    return symbols
+
+
+def _get_hip_symbols(raw_data: dict, lines: list[str]) -> list[DocumentSymbol]:
+    """Extract HIP-specific symbols: device functions, shared memory allocations."""
+    symbols: list[DocumentSymbol] = []
+    
+    # Device functions (from HIP parser)
+    device_functions = raw_data.get("device_functions", [])
+    for func in device_functions:
+        if not hasattr(func, "line") or not hasattr(func, "name"):
+            continue
+            
+        func_line = lines[func.line] if func.line < len(lines) else ""
+        name_start = func_line.find(func.name)
+        name_end = name_start + len(func.name) if name_start >= 0 else 0
+        
+        end_line = getattr(func, "end_line", func.line + 10)
+        
+        selection_range = Range(
+            start=Position(line=func.line, character=max(0, name_start)),
+            end=Position(line=func.line, character=name_end)
+        )
+        full_range = Range(
+            start=Position(line=func.line, character=0),
+            end=Position(line=min(end_line, len(lines) - 1), character=0)
+        )
+        
+        return_type = getattr(func, "return_type", "void")
+        detail = f"⚡ Device Function -> {return_type}"
+        
+        symbols.append(DocumentSymbol(
+            name=func.name,
+            kind=SymbolKind.Method,
+            range=full_range,
+            selection_range=selection_range,
+            detail=detail,
+        ))
+    
+    # Shared memory allocations
+    shared_memory = raw_data.get("shared_memory", [])
+    for shared in shared_memory:
+        if not hasattr(shared, "line") or not hasattr(shared, "name"):
+            continue
+            
+        shared_line = lines[shared.line] if shared.line < len(lines) else ""
+        name_start = shared_line.find(shared.name)
+        name_end = name_start + len(shared.name) if name_start >= 0 else 0
+        
+        selection_range = Range(
+            start=Position(line=shared.line, character=max(0, name_start)),
+            end=Position(line=shared.line, character=name_end)
+        )
+        full_range = Range(
+            start=Position(line=shared.line, character=0),
+            end=Position(line=shared.line, character=len(shared_line))
+        )
+        
+        type_str = getattr(shared, "type_str", "")
+        size_bytes = getattr(shared, "size_bytes", None)
+        if size_bytes:
+            if size_bytes >= 1024:
+                detail = f"📦 __shared__ {type_str} ({size_bytes / 1024:.1f} KB)"
+            else:
+                detail = f"📦 __shared__ {type_str} ({size_bytes} bytes)"
+        else:
+            detail = f"📦 __shared__ {type_str}"
+        
+        symbols.append(DocumentSymbol(
+            name=shared.name,
+            kind=SymbolKind.Variable,
+            range=full_range,
+            selection_range=selection_range,
+            detail=detail,
+        ))
+    
     return symbols

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

@@ -5,8 +5,20 @@ 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
@@ -133,8 +145,31 @@ def format_decorator_hover(decorator_name: str, function_name: str | None = None
 
 
 def handle_hover(uri: str, position: Position, content: str) -> Hover | None:
-    test_message = "**HEYOOOOOO** 🎉\n\n"
-
+    """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
@@ -154,13 +189,13 @@ def handle_hover(uri: str, position: Position, content: str) -> Hover | None:
                 if class_name_end > class_name_start:
                     function_name = func_line[class_name_start:class_name_end].strip()
 
-        hover_content = test_message + format_decorator_hover(decorator_name, function_name)
+        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 = [
-            test_message,
             "**cutlass.cute**",
             "",
             "CuTeDSL (CUDA Unified Tensor Expression) library for GPU programming.",
@@ -176,24 +211,25 @@ def handle_hover(uri: str, position: Position, content: str) -> Hover | None:
         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 = test_message + format_kernel_hover(kernel, analysis)
+            hover_content = format_kernel_hover(kernel, analysis)
         else:
-            hover_content = test_message + format_kernel_hover_basic(kernel)
+            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 = [
-            test_message,
             f"**Layout: {layout.name}**",
             ""
         ]
@@ -211,5 +247,5 @@ def handle_hover(uri: str, position: Position, content: str) -> Hover | None:
 
         return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
 
-    hover_content = test_message + "Hover is working! Move your cursor over symbols to see more info."
-    return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+    # No hover info found
+    return None