wafer-lsp 0.1.13__py3-none-any.whl

wafer_lsp/languages/README.md

@@ -0,0 +1,195 @@
+# Adding a New Language to Wafer LSP
+
+This guide explains how to add support for a new GPU programming language to the LSP server.
+
+## Architecture Overview
+
+The LSP server uses a modular, component-based architecture:
+
+- **`LanguageDetector`**: Detects language from file extensions/URIs
+- **`ParserManager`**: Manages parser instances for each language
+- **`ParserResultConverter`**: Converts parser-specific types to common types
+- **`LanguageRegistry`**: Coordinates all components (facade pattern)
+
+Each component is small, focused, and independently testable.
+
+## Overview
+
+To add a new language, you need to:
+
+1. Create a parser that implements `BaseParser`
+2. Register the parser with the language registry
+3. (Optional) Add language-specific handlers
+
+## Step 1: Create a Parser
+
+Create a new parser file in `wafer_lsp/parsers/`:
+
+```python
+# wafer_lsp/parsers/my_language_parser.py
+
+from typing import List, Dict, Any, Optional
+from dataclasses import dataclass
+from .base_parser import BaseParser
+
+@dataclass
+class MyLanguageKernel:
+    name: str
+    line: int
+    parameters: List[str]
+    docstring: Optional[str] = None
+
+class MyLanguageParser(BaseParser):
+    """Parser for MyLanguage GPU code."""
+    
+    def parse_file(self, content: str) -> Dict[str, Any]:
+        """Parse file and extract kernels, layouts, etc."""
+        kernels: List[MyLanguageKernel] = []
+        
+        # Your parsing logic here
+        # Example: find kernel definitions
+        for match in re.finditer(r'kernel\s+(\w+)', content):
+            kernels.append(MyLanguageKernel(
+                name=match.group(1),
+                line=content[:match.start()].count('\n'),
+                parameters=[],
+            ))
+        
+        return {
+            "kernels": kernels,
+            "layouts": [],  # If your language has layouts
+            "structs": [],  # If your language has structs
+        }
+```
+
+## Step 2: Register the Parser
+
+Update `wafer_lsp/languages/registry.py` to register your parser:
+
+```python
+from ..parsers.my_language_parser import MyLanguageParser
+
+class LanguageRegistry:
+    def _register_defaults(self):
+        # ... existing registrations ...
+        
+        # Register your new language
+        self.register_language(
+            language_id="mylang",
+            display_name="My Language",
+            parser=MyLanguageParser(),
+            extensions=[".mylang", ".ml"],
+            file_patterns=["*.mylang", "*.ml"]
+        )
+```
+
+## Step 3: Update VS Code Extension (if needed)
+
+If your language needs special handling in the VS Code extension, update `lspClient.ts`:
+
+```typescript
+const clientOptions: LanguageClientOptions = {
+    documentSelector: [
+        // ... existing selectors ...
+        { scheme: 'file', language: 'mylang' },  // Your language
+    ],
+    // ...
+};
+```
+
+## Parser Interface
+
+Your parser must implement `BaseParser`:
+
+```python
+class BaseParser(ABC):
+    @abstractmethod
+    def parse_file(self, content: str) -> Dict[str, Any]:
+        """Parse file content and extract language-specific constructs.
+        
+        Returns:
+            Dictionary with keys:
+            - "kernels": List of kernel objects (must have .name, .line, .parameters)
+            - "layouts": List of layout objects (must have .name, .line)
+            - "structs": List of struct objects (must have .name, .line)
+        """
+        pass
+```
+
+## Example: Adding OpenCL Support
+
+Here's a complete example for adding OpenCL support:
+
+```python
+# wafer_lsp/parsers/opencl_parser.py
+
+import re
+from typing import List, Dict, Any
+from dataclasses import dataclass
+from .base_parser import BaseParser
+
+@dataclass
+class OpenCLKernel:
+    name: str
+    line: int
+    parameters: List[str]
+
+class OpenCLParser(BaseParser):
+    def parse_file(self, content: str) -> Dict[str, Any]:
+        kernels: List[OpenCLKernel] = []
+        
+        # Pattern: __kernel void kernel_name(...)
+        pattern = r'__kernel\s+(?:__global\s+)?(?:void|.*?)\s+(\w+)\s*\('
+        
+        for match in re.finditer(pattern, content):
+            line = content[:match.start()].count('\n')
+            kernel_name = match.group(1)
+            params = self._extract_parameters(content, match.end())
+            
+            kernels.append(OpenCLKernel(
+                name=kernel_name,
+                line=line,
+                parameters=params
+            ))
+        
+        return {"kernels": kernels, "layouts": [], "structs": []}
+    
+    def _extract_parameters(self, content: str, start: int) -> List[str]:
+        # Extract parameter list logic
+        return []
+```
+
+Then register it:
+
+```python
+# In registry.py _register_defaults()
+self.register_language(
+    language_id="opencl",
+    display_name="OpenCL",
+    parser=OpenCLParser(),
+    extensions=[".cl"],
+    file_patterns=["*.cl"]
+)
+```
+
+## Benefits of This Architecture
+
+1. **Modular**: Each language parser is independent
+2. **Extensible**: Easy to add new languages without modifying existing code
+3. **Type-safe**: Common types (`KernelInfo`, `LayoutInfo`) ensure consistency
+4. **Language-agnostic handlers**: Handlers work with any registered language
+5. **Automatic detection**: Language is detected from file extension automatically
+
+## Testing
+
+Add tests for your parser:
+
+```python
+# tests/test_my_language_parser.py
+
+def test_parse_kernel():
+    parser = MyLanguageParser()
+    result = parser.parse_file("kernel my_kernel() { }")
+    assert len(result["kernels"]) == 1
+    assert result["kernels"][0].name == "my_kernel"
+```

wafer_lsp/languages/__init__.py

@@ -0,0 +1,17 @@
+from .converter import ParserResultConverter
+from .detector import LanguageDetector
+from .parser_manager import ParserManager
+from .registry import LanguageRegistry, get_language_registry
+from .types import KernelInfo, LanguageInfo, LayoutInfo, StructInfo
+
+__all__ = [
+    "KernelInfo",
+    "LanguageDetector",
+    "LanguageInfo",
+    "LanguageRegistry",
+    "LayoutInfo",
+    "ParserManager",
+    "ParserResultConverter",
+    "StructInfo",
+    "get_language_registry",
+]

wafer_lsp/languages/converter.py

@@ -0,0 +1,88 @@
+from typing import Any
+
+from .types import KernelInfo, LanguageInfo, LayoutInfo, StructInfo
+
+
+class ParserResultConverter:
+
+    def convert(
+        self,
+        parsed_data: dict[str, Any],
+        language_id: str
+    ) -> LanguageInfo:
+        kernels = self._convert_kernels(
+            parsed_data.get("kernels", []),
+            language_id
+        )
+        layouts = self._convert_layouts(
+            parsed_data.get("layouts", []),
+            language_id
+        )
+        structs = self._convert_structs(
+            parsed_data.get("structs", []),
+            language_id
+        )
+
+        return LanguageInfo(
+            kernels=kernels,
+            layouts=layouts,
+            structs=structs,
+            language=language_id,
+            raw_data=parsed_data
+        )
+
+    def _convert_kernels(
+        self,
+        kernels: list[Any],
+        language_id: str
+    ) -> list[KernelInfo]:
+        result: list[KernelInfo] = []
+
+        for kernel in kernels:
+            if hasattr(kernel, "name") and hasattr(kernel, "line"):
+                result.append(KernelInfo(
+                    name=kernel.name,
+                    line=kernel.line,
+                    parameters=getattr(kernel, "parameters", []),
+                    docstring=getattr(kernel, "docstring", None),
+                    language=language_id
+                ))
+
+        return result
+
+    def _convert_layouts(
+        self,
+        layouts: list[Any],
+        language_id: str
+    ) -> list[LayoutInfo]:
+        result: list[LayoutInfo] = []
+
+        for layout in layouts:
+            if hasattr(layout, "name") and hasattr(layout, "line"):
+                result.append(LayoutInfo(
+                    name=layout.name,
+                    line=layout.line,
+                    shape=getattr(layout, "shape", None),
+                    stride=getattr(layout, "stride", None),
+                    language=language_id
+                ))
+
+        return result
+
+    def _convert_structs(
+        self,
+        structs: list[Any],
+        language_id: str
+    ) -> list[StructInfo]:
+        result: list[StructInfo] = []
+
+        for struct in structs:
+            if hasattr(struct, "name") and hasattr(struct, "line"):
+                result.append(StructInfo(
+                    name=struct.name,
+                    line=struct.line,
+                    docstring=getattr(struct, "docstring", None),
+                    language=language_id
+                ))
+
+        return result

wafer_lsp/languages/detector.py

@@ -0,0 +1,107 @@
+from pathlib import Path
+
+
+class LanguageDetector:
+    """Detects language based on file extension and content markers.
+    
+    Supports both extension-based detection (fast) and content-based detection
+    (for files that share extensions, e.g., .cpp files that could be HIP or CUDA).
+    """
+
+    def __init__(self):
+        self._extensions: dict[str, str] = {}
+        self._content_markers: dict[str, list[str]] = {}  # language_id -> markers
+        # Compound extensions like .hip.cpp need special handling
+        self._compound_extensions: dict[str, str] = {}
+
+    def register_extension(self, extension: str, language_id: str):
+        normalized_ext = extension if extension.startswith(".") else f".{extension}"
+        
+        # Check if this is a compound extension (e.g., .hip.cpp)
+        if normalized_ext.count(".") > 1:
+            self._compound_extensions[normalized_ext] = language_id
+        else:
+            self._extensions[normalized_ext] = language_id
+
+    def register_content_markers(self, language_id: str, markers: list[str]):
+        """Register content markers for content-based language detection."""
+        self._content_markers[language_id] = markers
+
+    def detect_from_uri(self, uri: str, content: str | None = None) -> str | None:
+        """Detect language from URI and optionally content.
+        
+        Args:
+            uri: File URI or path
+            content: Optional file content for content-based detection
+            
+        Returns:
+            Language ID or None
+        """
+        if uri.startswith("file://"):
+            file_path = uri[7:]
+        else:
+            file_path = uri
+
+        return self.detect_from_path(file_path, content)
+
+    def detect_from_path(self, file_path: str, content: str | None = None) -> str | None:
+        """Detect language from file path and optionally content.
+        
+        Order of detection:
+        1. Compound extensions (e.g., .hip.cpp) - most specific
+        2. Content markers (for shared extensions like .cpp)
+        3. Simple extension
+        """
+        path = Path(file_path)
+        
+        # 1. Check compound extensions first
+        # Get the last two suffixes for compound extension detection
+        suffixes = path.suffixes
+        if len(suffixes) >= 2:
+            compound_ext = "".join(suffixes[-2:]).lower()
+            if compound_ext in self._compound_extensions:
+                return self._compound_extensions[compound_ext]
+        
+        # 2. If content is provided, check content markers
+        if content:
+            content_lang = self._detect_from_content(content)
+            if content_lang:
+                return content_lang
+        
+        # 3. Fall back to simple extension
+        ext = path.suffix.lower()
+        return self._extensions.get(ext)
+
+    def _detect_from_content(self, content: str) -> str | None:
+        """Detect language based on content markers.
+        
+        Returns the language with the most matching markers.
+        """
+        best_match: str | None = None
+        best_count = 0
+        
+        for language_id, markers in self._content_markers.items():
+            match_count = sum(1 for marker in markers if marker in content)
+            if match_count > best_count:
+                best_count = match_count
+                best_match = language_id
+        
+        # Require at least one marker match
+        return best_match if best_count > 0 else None
+
+    def detect_from_extension(self, extension: str) -> str | None:
+        normalized_ext = extension if extension.startswith(".") else f".{extension}"
+        normalized_ext = normalized_ext.lower()  # Case insensitive
+        return self._extensions.get(normalized_ext)
+
+    def get_supported_extensions(self) -> list[str]:
+        all_extensions = list(self._extensions.keys())
+        all_extensions.extend(self._compound_extensions.keys())
+        return all_extensions
+
+    def is_supported(self, uri: str, content: str | None = None) -> bool:
+        return self.detect_from_uri(uri, content) is not None
+
+    def get_content_markers(self, language_id: str) -> list[str]:
+        """Get content markers for a language."""
+        return self._content_markers.get(language_id, [])

wafer_lsp/languages/parser_manager.py

@@ -0,0 +1,33 @@
+
+from ..parsers.base_parser import BaseParser
+
+
+class ParserManager:
+
+    def __init__(self):
+        self._parsers: dict[str, BaseParser] = {}
+        self._language_names: dict[str, str] = {}
+
+    def register_parser(
+        self,
+        language_id: str,
+        display_name: str,
+        parser: BaseParser
+    ):
+        assert language_id not in self._parsers, \
+            f"Language {language_id} already registered"
+
+        self._parsers[language_id] = parser
+        self._language_names[language_id] = display_name
+
+    def get_parser(self, language_id: str) -> BaseParser | None:
+        return self._parsers.get(language_id)
+
+    def get_language_name(self, language_id: str) -> str | None:
+        return self._language_names.get(language_id)
+
+    def has_parser(self, language_id: str) -> bool:
+        return language_id in self._parsers
+
+    def list_languages(self) -> list[str]:
+        return list(self._parsers.keys())

wafer_lsp/languages/registry.py

@@ -0,0 +1,120 @@
+
+from ..parsers.cuda_parser import CUDAParser
+from ..parsers.cutedsl_parser import CuTeDSLParser
+from ..parsers.hip_parser import HIPParser
+from .converter import ParserResultConverter
+from .detector import LanguageDetector
+from .parser_manager import ParserManager
+from .types import LanguageInfo
+
+
+class LanguageRegistry:
+
+    def __init__(self):
+        self._detector = LanguageDetector()
+        self._parser_manager = ParserManager()
+        self._converter = ParserResultConverter()
+
+        self._register_defaults()
+
+    def _register_defaults(self):
+        self.register_language(
+            language_id="cutedsl",
+            display_name="CuTeDSL",
+            parser=CuTeDSLParser(),
+            extensions=[".py"],
+            file_patterns=["*.py"]
+        )
+
+        self.register_language(
+            language_id="cuda",
+            display_name="CUDA",
+            parser=CUDAParser(),
+            extensions=[".cu", ".cuh"],
+            file_patterns=["*.cu", "*.cuh"]
+        )
+
+        # HIP (AMD GPU) - Register before cpp so .hip.cpp files get detected as HIP
+        self.register_language(
+            language_id="hip",
+            display_name="HIP (AMD GPU)",
+            parser=HIPParser(),
+            extensions=[".hip", ".hip.cpp", ".hip.hpp", ".hipcc"],
+            file_patterns=["*.hip", "*.hip.cpp", "*.hip.hpp", "*.hipcc"],
+            content_markers=[
+                "#include <hip/hip_runtime.h>",
+                "#include \"hip/hip_runtime.h\"",
+                "hipMalloc",
+                "hipLaunchKernelGGL",
+                "__HIP_PLATFORM_AMD__",
+            ]
+        )
+
+        self.register_language(
+            language_id="cpp",
+            display_name="C++",
+            parser=CUDAParser(),
+            extensions=[".cpp", ".hpp", ".cc", ".cxx"],
+            file_patterns=["*.cpp", "*.hpp", "*.cc", "*.cxx"]
+        )
+
+    def register_language(
+        self,
+        language_id: str,
+        display_name: str,
+        parser,
+        extensions: list[str],
+        file_patterns: list[str] | None = None,
+        content_markers: list[str] | None = None
+    ):
+        self._parser_manager.register_parser(language_id, display_name, parser)
+
+        for ext in extensions:
+            self._detector.register_extension(ext, language_id)
+        
+        if content_markers:
+            self._detector.register_content_markers(language_id, content_markers)
+
+    def detect_language(self, uri: str) -> str | None:
+        return self._detector.detect_from_uri(uri)
+
+    def get_parser(self, language_id: str):
+        return self._parser_manager.get_parser(language_id)
+
+    def parse_file(self, uri: str, content: str) -> LanguageInfo | None:
+        language_id = self.detect_language(uri)
+        if not language_id:
+            return None
+
+        parser = self.get_parser(language_id)
+        if not parser:
+            return None
+
+        try:
+            parsed_data = parser.parse_file(content)
+        except Exception:
+            return LanguageInfo(
+                kernels=[],
+                layouts=[],
+                structs=[],
+                language=language_id,
+                raw_data={}
+            )
+
+        return self._converter.convert(parsed_data, language_id)
+
+    def get_supported_extensions(self) -> list[str]:
+        return self._detector.get_supported_extensions()
+
+    def get_language_name(self, language_id: str) -> str | None:
+        return self._parser_manager.get_language_name(language_id)
+
+
+_registry: LanguageRegistry | None = None
+
+
+def get_language_registry() -> LanguageRegistry:
+    global _registry
+    if _registry is None:
+        _registry = LanguageRegistry()
+    return _registry

wafer_lsp/languages/types.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class KernelInfo:
+    name: str
+    line: int
+    parameters: list[str]
+    docstring: str | None = None
+    language: str = ""
+
+
+@dataclass
+class LayoutInfo:
+    name: str
+    line: int
+    shape: str | None = None
+    stride: str | None = None
+    language: str = ""
+
+
+@dataclass
+class StructInfo:
+    name: str
+    line: int
+    docstring: str | None = None
+    language: str = ""
+
+
+@dataclass
+class LanguageInfo:
+    kernels: list[KernelInfo]
+    layouts: list[LayoutInfo]
+    structs: list[StructInfo]
+    language: str
+    raw_data: dict[str, Any]

wafer_lsp/parsers/__init__.py

@@ -0,0 +1,36 @@
+from .base_parser import BaseParser
+from .cuda_parser import CUDAKernel, CUDAParser
+from .cutedsl_parser import (
+    CuTeDSLKernel,
+    CuTeDSLLayout,
+    CuTeDSLParser,
+    CuTeDSLStruct,
+)
+from .hip_parser import (
+    HIPKernel,
+    HIPDeviceFunction,
+    HIPParameter,
+    HIPParser,
+    KernelLaunchSite,
+    SharedMemoryAllocation,
+    WavefrontPattern,
+    is_hip_file,
+)
+
+__all__ = [
+    "BaseParser",
+    "CUDAKernel",
+    "CUDAParser",
+    "CuTeDSLKernel",
+    "CuTeDSLLayout",
+    "CuTeDSLParser",
+    "CuTeDSLStruct",
+    "HIPKernel",
+    "HIPDeviceFunction",
+    "HIPParameter",
+    "HIPParser",
+    "KernelLaunchSite",
+    "SharedMemoryAllocation",
+    "WavefrontPattern",
+    "is_hip_file",
+]

wafer_lsp/parsers/base_parser.py

@@ -0,0 +1,9 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseParser(ABC):
+
+    @abstractmethod
+    def parse_file(self, content: str) -> dict[str, Any]:
+        pass