wafer-lsp 0.1.13__py3-none-any.whl

wafer_lsp/services/hip_hover_service.py

@@ -0,0 +1,412 @@
+"""
+HIP Hover Service.
+
+Provides rich hover documentation for HIP code including:
+- HIP API functions (hipMalloc, hipMemcpy, etc.)
+- Memory qualifiers (__device__, __shared__, __constant__)
+- Wavefront intrinsics (__shfl, __ballot, etc.)
+- Thread indexing (threadIdx, blockIdx, etc.)
+- Kernel function information
+"""
+
+from lsprotocol.types import Hover, MarkupContent, MarkupKind, Position
+
+from ..parsers.hip_parser import (
+    HIPKernel,
+    HIPDeviceFunction,
+    SharedMemoryAllocation,
+    KernelLaunchSite,
+    HIPParser,
+)
+from .hip_docs import (
+    HIPDocsService,
+    HIPAPIDoc,
+    MemoryQualifierDoc,
+    IntrinsicDoc,
+    ThreadIndexDoc,
+    create_hip_docs_service,
+)
+
+
+class HIPHoverService:
+    """Provides hover documentation for HIP code."""
+
+    def __init__(self, docs_service: HIPDocsService | None = None):
+        self._docs = docs_service or create_hip_docs_service()
+        self._parser = HIPParser()
+
+    def get_hover(self, content: str, position: Position, uri: str) -> Hover | None:
+        """Get hover information for a position in HIP code.
+        
+        Args:
+            content: The document content
+            position: The cursor position
+            uri: The document URI
+            
+        Returns:
+            Hover information or None
+        """
+        word = self._get_word_at_position(content, position)
+        if not word:
+            return None
+
+        # Try different types of hover in order of specificity
+        
+        # 1. Check for HIP API functions
+        api_doc = self._docs.get_api_doc(word)
+        if api_doc:
+            return self._format_api_hover(api_doc)
+
+        # 2. Check for memory qualifiers (including in context)
+        qualifier_doc = self._docs.get_memory_qualifier_doc(word)
+        if qualifier_doc:
+            return self._format_qualifier_hover(qualifier_doc)
+
+        # 3. Check for intrinsics
+        intrinsic_doc = self._docs.get_intrinsic_doc(word)
+        if intrinsic_doc:
+            return self._format_intrinsic_hover(intrinsic_doc)
+
+        # 4. Check for thread indexing variables
+        thread_doc = self._docs.get_thread_index_doc(word)
+        if thread_doc:
+            return self._format_thread_index_hover(thread_doc)
+
+        # 5. Check for kernels/device functions in the file
+        parsed = self._parser.parse_file(content)
+        
+        for kernel in parsed.get("kernels", []):
+            if kernel.name == word:
+                return self._format_kernel_hover(kernel)
+
+        for device_func in parsed.get("device_functions", []):
+            if device_func.name == word:
+                return self._format_device_function_hover(device_func)
+
+        # 6. Check for shared memory variables
+        for shared_var in parsed.get("shared_memory", []):
+            if shared_var.name == word:
+                return self._format_shared_memory_hover(shared_var)
+
+        # 7. Check for kernel launch (when hovering on kernel name at launch site)
+        for launch in parsed.get("launch_sites", []):
+            if launch.kernel_name == word and position.line == launch.line:
+                return self._format_launch_site_hover(launch)
+
+        return None
+
+    def _get_word_at_position(self, content: str, position: Position) -> str:
+        """Extract the word at the given position."""
+        lines = content.split("\n")
+        if position.line >= len(lines):
+            return ""
+
+        line = lines[position.line]
+        if position.character >= len(line):
+            return ""
+
+        # Find word boundaries (include underscores for __global__ etc.)
+        start = position.character
+        while start > 0 and (line[start - 1].isalnum() or line[start - 1] == "_"):
+            start -= 1
+
+        end = position.character
+        while end < len(line) and (line[end].isalnum() or line[end] == "_"):
+            end += 1
+
+        return line[start:end]
+
+    def _format_api_hover(self, doc: HIPAPIDoc) -> Hover:
+        """Format hover content for a HIP API function."""
+        lines = [
+            f"### `{doc.name}`",
+            "",
+            f"```cpp",
+            doc.signature,
+            "```",
+            "",
+            doc.description,
+            "",
+        ]
+
+        if doc.parameters:
+            lines.append("**Parameters:**")
+            for param_name, param_desc in doc.parameters:
+                lines.append(f"- `{param_name}`: {param_desc}")
+            lines.append("")
+
+        lines.append(f"**Returns:** {doc.return_value}")
+        lines.append("")
+
+        if doc.amd_notes:
+            lines.append("**AMD Notes:**")
+            lines.append(f"> {doc.amd_notes}")
+            lines.append("")
+
+        if doc.example:
+            lines.append("**Example:**")
+            lines.append("```cpp")
+            lines.extend(doc.example.replace("\\n", "\n").split("\n"))
+            lines.append("```")
+            lines.append("")
+
+        if doc.related:
+            lines.append(f"**Related:** {', '.join(f'`{r}`' for r in doc.related)}")
+
+        if doc.doc_url:
+            lines.append("")
+            lines.append(f"[📖 Documentation]({doc.doc_url})")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _format_qualifier_hover(self, doc: MemoryQualifierDoc) -> Hover:
+        """Format hover content for a memory qualifier."""
+        lines = [
+            f"### `{doc.name}`",
+            "",
+            doc.description,
+            "",
+            "**AMD Architecture:**",
+            doc.amd_details,
+            "",
+            "**Performance Tips:**",
+            doc.performance_tips,
+        ]
+
+        if doc.example:
+            lines.append("")
+            lines.append("**Example:**")
+            lines.append("```cpp")
+            lines.extend(doc.example.split("\n"))
+            lines.append("```")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _format_intrinsic_hover(self, doc: IntrinsicDoc) -> Hover:
+        """Format hover content for a wavefront intrinsic."""
+        lines = [
+            f"### `{doc.name}`",
+            "",
+            f"```cpp",
+            doc.signature,
+            "```",
+            "",
+            doc.description,
+            "",
+            "**⚠️ AMD Wavefront Behavior:**",
+            f"> {doc.amd_behavior}",
+            "",
+        ]
+
+        if doc.parameters:
+            lines.append("**Parameters:**")
+            for param_name, param_desc in doc.parameters:
+                lines.append(f"- `{param_name}`: {param_desc}")
+            lines.append("")
+
+        lines.append(f"**Returns:** {doc.return_value}")
+        lines.append("")
+
+        if doc.example:
+            lines.append("**Example:**")
+            lines.append("```cpp")
+            lines.extend(doc.example.replace("\\n", "\n").split("\n"))
+            lines.append("```")
+            lines.append("")
+
+        if doc.cuda_equivalent:
+            lines.append(f"**CUDA Equivalent:** `{doc.cuda_equivalent}`")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _format_thread_index_hover(self, doc: ThreadIndexDoc) -> Hover:
+        """Format hover content for thread indexing variables."""
+        lines = [
+            f"### `{doc.name}`",
+            "",
+            doc.description,
+            "",
+            "**AMD Context:**",
+            f"> {doc.amd_context}",
+            "",
+        ]
+
+        if doc.common_patterns:
+            lines.append("**Common Patterns:**")
+            lines.append("```cpp")
+            for pattern in doc.common_patterns:
+                lines.append(pattern)
+            lines.append("```")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _format_kernel_hover(self, kernel: HIPKernel) -> Hover:
+        """Format hover content for a kernel function."""
+        lines = [
+            f"### 🚀 HIP Kernel: `{kernel.name}`",
+            "",
+        ]
+
+        if kernel.docstring:
+            lines.append(kernel.docstring)
+            lines.append("")
+
+        # Build signature
+        params_str = ", ".join(kernel.parameters) if kernel.parameters else ""
+        lines.append("```cpp")
+        if kernel.attributes:
+            for attr in kernel.attributes:
+                lines.append(attr)
+        lines.append(f"__global__ void {kernel.name}({params_str})")
+        lines.append("```")
+        lines.append("")
+
+        if kernel.parameter_info:
+            lines.append("**Parameters:**")
+            for param in kernel.parameter_info:
+                type_info = f" (`{param.type_str}`)" if param.type_str else ""
+                lines.append(f"- `{param.name}`{type_info}")
+            lines.append("")
+
+        lines.append(f"**Location:** Lines {kernel.line + 1} - {kernel.end_line + 1}")
+        lines.append("")
+        lines.append("**AMD GPU Execution:**")
+        lines.append("- Executed on GPU Compute Units")
+        lines.append("- Threads grouped into 64-thread wavefronts (CDNA)")
+        lines.append("- Use `<<<grid, block>>>` or `hipLaunchKernelGGL` to launch")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _format_device_function_hover(self, func: HIPDeviceFunction) -> Hover:
+        """Format hover content for a device function."""
+        lines = [
+            f"### ⚡ Device Function: `{func.name}`",
+            "",
+        ]
+
+        # Build signature
+        params_str = ", ".join(func.parameters) if func.parameters else ""
+        inline_str = "__forceinline__ " if func.is_inline else ""
+        lines.append("```cpp")
+        lines.append(f"__device__ {inline_str}{func.return_type} {func.name}({params_str})")
+        lines.append("```")
+        lines.append("")
+
+        if func.parameter_info:
+            lines.append("**Parameters:**")
+            for param in func.parameter_info:
+                type_info = f" (`{param.type_str}`)" if param.type_str else ""
+                lines.append(f"- `{param.name}`{type_info}")
+            lines.append("")
+
+        lines.append(f"**Returns:** `{func.return_type}`")
+        lines.append("")
+        lines.append(f"**Location:** Lines {func.line + 1} - {func.end_line + 1}")
+        lines.append("")
+        lines.append("**Note:** Device functions can only be called from kernel or other device functions.")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _format_shared_memory_hover(self, shared: SharedMemoryAllocation) -> Hover:
+        """Format hover content for a shared memory allocation."""
+        lines = [
+            f"### 📦 Shared Memory (LDS): `{shared.name}`",
+            "",
+            f"**Type:** `{shared.type_str}`",
+        ]
+
+        if shared.array_size:
+            lines.append(f"**Array Size:** `[{shared.array_size}]`")
+
+        if shared.size_bytes:
+            if shared.size_bytes >= 1024:
+                size_str = f"{shared.size_bytes / 1024:.1f} KB"
+            else:
+                size_str = f"{shared.size_bytes} bytes"
+            lines.append(f"**Size:** {size_str}")
+
+        if shared.is_dynamic:
+            lines.append("**Allocation:** Dynamic (extern)")
+        else:
+            lines.append("**Allocation:** Static")
+
+        lines.append("")
+        lines.append("**AMD LDS Details:**")
+        lines.append("- On-chip memory with ~100x lower latency than HBM")
+        lines.append("- 64 KB per Compute Unit")
+        lines.append("- Shared by all threads in the block")
+        lines.append("- 32 banks of 4 bytes each")
+        lines.append("")
+        lines.append("💡 **Tip:** Avoid bank conflicts by using padding: `[SIZE + 1]`")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _format_launch_site_hover(self, launch: KernelLaunchSite) -> Hover:
+        """Format hover content for a kernel launch site."""
+        lines = [
+            f"### 🎯 Kernel Launch: `{launch.kernel_name}`",
+            "",
+        ]
+
+        if launch.is_hip_launch_kernel_ggl:
+            lines.append("**Launch Method:** `hipLaunchKernelGGL`")
+        else:
+            lines.append("**Launch Method:** `<<<>>>` syntax")
+
+        lines.append("")
+
+        if launch.grid_dim:
+            lines.append(f"**Grid Dimensions:** `{launch.grid_dim}`")
+        if launch.block_dim:
+            lines.append(f"**Block Dimensions:** `{launch.block_dim}`")
+            # Try to parse block dimensions for wavefront info
+            self._add_wavefront_info(lines, launch.block_dim)
+
+        if launch.shared_mem_bytes:
+            lines.append(f"**Dynamic Shared Memory:** `{launch.shared_mem_bytes}`")
+        if launch.stream:
+            lines.append(f"**Stream:** `{launch.stream}`")
+
+        return Hover(contents=MarkupContent(
+            kind=MarkupKind.Markdown,
+            value="\n".join(lines)
+        ))
+
+    def _add_wavefront_info(self, lines: list[str], block_dim: str) -> None:
+        """Add wavefront information based on block dimensions.
+        
+        Only adds info when block_dim is a simple numeric literal we can parse.
+        For complex expressions (variables, dim3, etc.), we don't display wavefront count
+        because we can't determine the value at parse time.
+        """
+        # Only handle simple numeric block size - we don't guess at complex expressions
+        if block_dim.isdigit():
+            block_size = int(block_dim)
+            wavefronts = (block_size + 63) // 64
+            lines.append(f"**Wavefronts per Block:** {wavefronts} (64 threads each on CDNA)")
+
+
+def create_hip_hover_service(docs_service: HIPDocsService | None = None) -> HIPHoverService:
+    """Create a HIP hover service instance."""
+    return HIPHoverService(docs_service)

wafer_lsp/services/hover_service.py

@@ -0,0 +1,237 @@
+
+from lsprotocol.types import Hover, MarkupContent, MarkupKind, Position
+
+from ..languages.types import KernelInfo, LayoutInfo
+from .analysis_service import AnalysisService
+from .docs_service import DocsService
+from .language_registry_service import LanguageRegistryService
+from .position_service import PositionService
+
+
+class HoverService:
+
+    def __init__(
+        self,
+        language_registry: LanguageRegistryService,
+        analysis_service: AnalysisService,
+        docs_service: DocsService,
+        position_service: PositionService
+    ):
+        self._language_registry = language_registry
+        self._analysis_service = analysis_service
+        self._docs_service = docs_service
+        self._position_service = position_service
+
+    def handle_hover(self, uri: str, position: Position, content: str) -> Hover | None:
+        test_message = "🎉🎉🎉 **HEYOOO!!! LSP IS DEFINITELY WORKING!!!** 🎉🎉🎉\n\n**THIS IS THE WAFER LSP SERVER!**\n\n"
+
+        decorator_info = self._position_service.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 = test_message + self._format_decorator_hover(decorator_name, function_name)
+            return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+        word = self._position_service.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.",
+                "",
+                "**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 = self._find_kernel_at_position(content, position, uri)
+        if kernel:
+            analysis = self._analysis_service.get_analysis_for_kernel(uri, kernel.name)
+
+            if analysis:
+                hover_content = test_message + self._format_kernel_hover(kernel, analysis)
+            else:
+                hover_content = test_message + self._format_kernel_hover_basic(kernel)
+
+            return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+        layout = self._find_layout_at_position(content, position, uri)
+        if layout:
+            doc_link = self._docs_service.get_doc_for_concept("layout")
+
+            hover_lines = [
+                test_message,
+                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))
+
+        hover_content = test_message + "**HOVER IS WORKING!** 🚀\n\nMove your cursor over any symbol, decorator, or even empty space to see LSP information.\n\n**Try hovering over:**\n- `@cute.kernel` decorators\n- `cute` module name\n- Kernel function names\n- Layout variables"
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    def _find_kernel_at_position(
+        self, content: str, position: Position, uri: str
+    ) -> KernelInfo | None:
+        language_info = self._language_registry.parse_file(uri, content)
+
+        if not language_info:
+            return None
+
+        word = self._position_service.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(
+        self, content: str, position: Position, uri: str
+    ) -> LayoutInfo | None:
+        language_info = self._language_registry.parse_file(uri, content)
+
+        if not language_info:
+            return None
+
+        word = self._position_service.get_word_at_position(content, position)
+
+        for layout in language_info.layouts:
+            if layout.name == word:
+                return layout
+
+        return None
+
+    def _format_kernel_hover(self, kernel: KernelInfo, analysis: dict | None) -> str:
+        language_name = self._language_registry.get_language_name(kernel.language) or kernel.language
+
+        if kernel.language == "cuda" or kernel.language == "cpp":
+            lines = [f"**CUDA Kernel: {kernel.name}**", f"*Language: {language_name}*", ""]
+        else:
+            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 kernel.language == "cuda" or kernel.language == "cpp":
+            lines.append("**CUDA Features:**")
+            lines.append("- `__global__` function executed on GPU")
+            lines.append("- Can be launched with `<<<grid, block>>>` syntax")
+            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(self, kernel: KernelInfo) -> str:
+        return self._format_kernel_hover(kernel, None)
+
+    def _format_decorator_hover(self, 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")
+
+        doc_link = self._docs_service.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 create_hover_service(
+    language_registry: LanguageRegistryService,
+    analysis_service: AnalysisService,
+    docs_service: DocsService,
+    position_service: PositionService
+) -> HoverService:
+    return HoverService(language_registry, analysis_service, docs_service, position_service)

wafer_lsp/services/language_registry_service.py

@@ -0,0 +1,26 @@
+
+from ..languages.registry import LanguageRegistry, get_language_registry
+from ..languages.types import LanguageInfo
+
+
+class LanguageRegistryService:
+
+    def __init__(self, registry: LanguageRegistry):
+        self._registry = registry
+
+    def detect_language(self, uri: str) -> str | None:
+        return self._registry.detect_language(uri)
+
+    def parse_file(self, uri: str, content: str) -> LanguageInfo | None:
+        return self._registry.parse_file(uri, content)
+
+    def get_language_name(self, language_id: str) -> str | None:
+        return self._registry.get_language_name(language_id)
+
+    def get_supported_extensions(self) -> list[str]:
+        return self._registry.get_supported_extensions()
+
+
+def create_language_registry_service() -> LanguageRegistryService:
+    registry = get_language_registry()
+    return LanguageRegistryService(registry)