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

wafer_lsp/handlers/inlay_hint.py

@@ -1,9 +1,27 @@
 
+import re
+
 from lsprotocol.types import InlayHint, InlayHintKind, Position, Range
 
 from ..languages.registry import get_language_registry
 
 
+# HIP kernel launch pattern: kernel<<<grid, block, shared, stream>>>
+HIP_LAUNCH_PATTERN = re.compile(
+    r'(\w+)\s*<<<\s*'
+    r'([^,>]+)\s*,\s*'  # grid
+    r'([^,>]+)'          # block
+    r'(?:\s*,\s*([^,>]+))?'  # shared (optional)
+    r'(?:\s*,\s*([^>]+))?'   # stream (optional)
+    r'\s*>>>'
+)
+
+# Shared memory declaration pattern
+SHARED_MEM_PATTERN = re.compile(
+    r'__shared__\s+([\w\s:<>]+?)\s+(\w+)\s*\[([^\]]+)\]'
+)
+
+
 def handle_inlay_hint(uri: str, content: str, range: Range) -> list[InlayHint]:
     registry = get_language_registry()
     language_info = registry.parse_file(uri, content)
@@ -14,6 +32,7 @@ def handle_inlay_hint(uri: str, content: str, range: Range) -> list[InlayHint]:
     hints: list[InlayHint] = []
     lines = content.split("\n")
 
+    # Layout hints (CuTeDSL)
     for layout in language_info.layouts:
         if layout.line < range.start.line or layout.line > range.end.line:
             continue
@@ -39,6 +58,7 @@ def handle_inlay_hint(uri: str, content: str, range: Range) -> list[InlayHint]:
                 padding_right=False
             ))
 
+    # Kernel hints (CuTeDSL)
     for kernel in language_info.kernels:
         if kernel.line < range.start.line or kernel.line > range.end.line:
             continue
@@ -62,4 +82,164 @@ def handle_inlay_hint(uri: str, content: str, range: Range) -> list[InlayHint]:
                 padding_right=True
             ))
 
+    # HIP/CUDA-specific hints
+    if language_info.language in ("hip", "cuda", "cpp"):
+        hints.extend(_get_hip_inlay_hints(lines, range))
+
     return hints
+
+
+def _get_hip_inlay_hints(lines: list[str], range: Range) -> list[InlayHint]:
+    """Generate HIP-specific inlay hints.
+    
+    - Kernel launch dimension annotations
+    - Shared memory size annotations
+    """
+    hints: list[InlayHint] = []
+
+    for i in range(range.start.line, min(range.end.line + 1, len(lines))):
+        line = lines[i]
+        
+        # Kernel launch hints
+        for match in HIP_LAUNCH_PATTERN.finditer(line):
+            kernel_name = match.group(1)
+            grid_dim = match.group(2).strip()
+            block_dim = match.group(3).strip()
+            shared_mem = match.group(4)
+            stream = match.group(5)
+            
+            # Add hint after >>> showing launch configuration
+            hint_parts = []
+            
+            # Try to parse and annotate dimensions
+            grid_info = _parse_dim(grid_dim)
+            block_info = _parse_dim(block_dim)
+            
+            if grid_info:
+                hint_parts.append(f"{grid_info} blocks")
+            if block_info:
+                hint_parts.append(f"{block_info} threads/block")
+                # Calculate wavefronts (AMD uses 64-thread wavefronts)
+                try:
+                    total_threads = _eval_dim(block_dim)
+                    if total_threads:
+                        wavefronts = (total_threads + 63) // 64
+                        hint_parts.append(f"{wavefronts} wavefront{'s' if wavefronts != 1 else ''}")
+                except (ValueError, SyntaxError):
+                    pass
+            
+            if hint_parts:
+                hint_text = " // " + ", ".join(hint_parts)
+                
+                # Position after >>>
+                hint_pos = match.end()
+                
+                hints.append(InlayHint(
+                    position=Position(line=i, character=hint_pos),
+                    label=hint_text,
+                    kind=InlayHintKind.Parameter,
+                    padding_left=True,
+                    padding_right=False
+                ))
+        
+        # Shared memory size hints
+        for match in SHARED_MEM_PATTERN.finditer(line):
+            type_str = match.group(1).strip()
+            var_name = match.group(2)
+            array_size = match.group(3).strip()
+            
+            size_bytes = _estimate_size(type_str, array_size)
+            if size_bytes:
+                if size_bytes >= 1024:
+                    size_str = f" // {size_bytes / 1024:.1f} KB LDS"
+                else:
+                    size_str = f" // {size_bytes} bytes LDS"
+                
+                # Position at end of declaration
+                hint_pos = match.end()
+                
+                hints.append(InlayHint(
+                    position=Position(line=i, character=hint_pos),
+                    label=size_str,
+                    kind=InlayHintKind.Type,
+                    padding_left=True,
+                    padding_right=False
+                ))
+
+    return hints
+
+
+def _parse_dim(dim_str: str) -> str | None:
+    """Parse a dimension string and return a human-readable description."""
+    dim_str = dim_str.strip()
+    
+    # Simple number
+    if dim_str.isdigit():
+        return dim_str
+    
+    # dim3(x, y, z)
+    if dim_str.startswith("dim3("):
+        return dim_str
+    
+    # Variable or expression
+    if re.match(r'^[\w_]+$', dim_str):
+        return dim_str
+    
+    return None
+
+
+def _eval_dim(dim_str: str) -> int | None:
+    """Try to evaluate a dimension to an integer."""
+    dim_str = dim_str.strip()
+    
+    # Simple number
+    if dim_str.isdigit():
+        return int(dim_str)
+    
+    # dim3(x) or dim3(x, y) or dim3(x, y, z) - try to multiply
+    if dim_str.startswith("dim3(") and dim_str.endswith(")"):
+        inner = dim_str[5:-1]
+        parts = [p.strip() for p in inner.split(",")]
+        try:
+            total = 1
+            for p in parts:
+                if p.isdigit():
+                    total *= int(p)
+                else:
+                    return None  # Can't evaluate variable
+            return total
+        except (ValueError, SyntaxError):
+            return None
+    
+    return None
+
+
+def _estimate_size(type_str: str, array_size: str) -> int | None:
+    """Estimate size in bytes for a shared memory allocation."""
+    type_sizes = {
+        'char': 1, 'int8_t': 1, 'uint8_t': 1,
+        'short': 2, 'int16_t': 2, 'uint16_t': 2, 'half': 2, '__half': 2,
+        'int': 4, 'int32_t': 4, 'uint32_t': 4, 'float': 4, 'unsigned': 4,
+        'long': 8, 'int64_t': 8, 'uint64_t': 8, 'double': 8,
+        'float4': 16, 'float2': 8, 'int4': 16, 'int2': 8,
+        'double2': 16, 'double4': 32,
+    }
+    
+    # Find base type
+    base_type = type_str.strip()
+    type_size = None
+    for known_type, size in type_sizes.items():
+        if known_type in base_type:
+            type_size = size
+            break
+    
+    if type_size is None:
+        type_size = 4  # Default to 4 bytes
+    
+    # Try to evaluate array size
+    try:
+        # Handle simple expressions
+        arr_size = eval(array_size.replace('*', ' * '))
+        return type_size * arr_size
+    except (ValueError, SyntaxError, NameError):
+        return None

wafer_lsp/handlers/semantic_tokens.py

@@ -1,13 +1,20 @@
 
+import re
+
 from lsprotocol.types import SemanticTokens, SemanticTokensLegend
 
 from ..languages.registry import get_language_registry
 
 TOKEN_TYPES = [
-    "kernel",
-    "layout",
-    "struct",
-    "decorator",
+    "kernel",           # 0: GPU kernel functions (__global__)
+    "layout",           # 1: Layout variables
+    "struct",           # 2: Structs
+    "decorator",        # 3: Python decorators (@cute.kernel)
+    "keyword_gpu",      # 4: GPU keywords (__global__, __device__, __shared__)
+    "keyword_memory",   # 5: Memory qualifiers (__shared__, __constant__)
+    "function_hip_api", # 6: HIP API calls (hipMalloc, etc.)
+    "function_intrinsic", # 7: Wavefront intrinsics (__shfl, __ballot)
+    "device_function",  # 8: __device__ functions
 ]
 
 TOKEN_MODIFIERS = [
@@ -20,6 +27,31 @@ SEMANTIC_TOKENS_LEGEND = SemanticTokensLegend(
     token_modifiers=TOKEN_MODIFIERS
 )
 
+# HIP-specific patterns for semantic highlighting
+HIP_KEYWORD_PATTERN = re.compile(r'\b(__global__|__device__|__host__|__forceinline__)\b')
+HIP_MEMORY_KEYWORD_PATTERN = re.compile(r'\b(__shared__|__constant__|__restrict__)\b')
+HIP_LAUNCH_BOUNDS_PATTERN = re.compile(r'__launch_bounds__\s*\([^)]+\)')
+HIP_API_PATTERN = re.compile(
+    r'\b(hipMalloc|hipMallocManaged|hipMallocAsync|hipFree|hipFreeAsync|'
+    r'hipMemcpy|hipMemcpyAsync|hipMemset|hipMemsetAsync|'
+    r'hipHostMalloc|hipHostFree|hipHostRegister|hipHostUnregister|'
+    r'hipDeviceSynchronize|hipStreamSynchronize|'
+    r'hipStreamCreate|hipStreamDestroy|hipStreamCreateWithFlags|'
+    r'hipEventCreate|hipEventDestroy|hipEventRecord|hipEventSynchronize|hipEventElapsedTime|'
+    r'hipSetDevice|hipGetDevice|hipGetDeviceCount|hipGetDeviceProperties|'
+    r'hipLaunchKernelGGL|hipLaunchCooperativeKernel|'
+    r'hipGetLastError|hipPeekAtLastError|hipGetErrorString|hipGetErrorName)\b'
+)
+HIP_INTRINSIC_PATTERN = re.compile(
+    r'\b(__shfl|__shfl_down|__shfl_up|__shfl_xor|__shfl_sync|'
+    r'__ballot|__any|__all|__activemask|'
+    r'__syncthreads|__syncwarp|__threadfence|__threadfence_block|__threadfence_system|'
+    r'atomicAdd|atomicSub|atomicMax|atomicMin|atomicExch|atomicCAS|atomicAnd|atomicOr|atomicXor|'
+    r'__popc|__popcll|__clz|__clzll|__ffs|__ffsll|'
+    r'__float2half|__half2float|__float2int_rn|__int2float_rn|'
+    r'__ldg|__ldcg|__ldca|__ldcs)\b'
+)
+
 
 def handle_semantic_tokens(uri: str, content: str) -> SemanticTokens:
     registry = get_language_registry()
@@ -28,11 +60,12 @@ def handle_semantic_tokens(uri: str, content: str) -> SemanticTokens:
     if not language_info:
         return SemanticTokens(data=[])
 
-    tokens: list[int] = []
+    # Collect all tokens with their positions
+    # We'll sort them later to ensure proper delta calculation
+    token_entries: list[tuple[int, int, int, int, int]] = []  # (line, char, length, type, modifier)
     lines = content.split("\n")
-    prev_line = 0
-    prev_char = 0
 
+    # Add kernel tokens
     for kernel in language_info.kernels:
         if kernel.line >= len(lines):
             continue
@@ -41,20 +74,15 @@ def handle_semantic_tokens(uri: str, content: str) -> SemanticTokens:
         name_start = kernel_line.find(kernel.name)
 
         if name_start >= 0:
-            delta_line = kernel.line - prev_line
-            delta_char = name_start - (prev_char if delta_line == 0 else 0)
-
-            tokens.extend([
-                delta_line,
-                delta_char,
+            token_entries.append((
+                kernel.line,
+                name_start,
                 len(kernel.name),
                 TOKEN_TYPES.index("kernel"),
                 0
-            ])
-
-            prev_line = kernel.line
-            prev_char = name_start + len(kernel.name)
+            ))
 
+    # Add layout tokens
     for layout in language_info.layouts:
         if layout.line >= len(lines):
             continue
@@ -63,20 +91,15 @@ def handle_semantic_tokens(uri: str, content: str) -> SemanticTokens:
         name_start = layout_line.find(layout.name)
 
         if name_start >= 0:
-            delta_line = layout.line - prev_line
-            delta_char = name_start - (prev_char if delta_line == 0 else 0)
-
-            tokens.extend([
-                delta_line,
-                delta_char,
+            token_entries.append((
+                layout.line,
+                name_start,
                 len(layout.name),
                 TOKEN_TYPES.index("layout"),
                 0
-            ])
-
-            prev_line = layout.line
-            prev_char = name_start + len(layout.name)
+            ))
 
+    # Add struct tokens
     for struct in language_info.structs:
         if struct.line >= len(lines):
             continue
@@ -85,20 +108,15 @@ def handle_semantic_tokens(uri: str, content: str) -> SemanticTokens:
         name_start = struct_line.find(struct.name)
 
         if name_start >= 0:
-            delta_line = struct.line - prev_line
-            delta_char = name_start - (prev_char if delta_line == 0 else 0)
-
-            tokens.extend([
-                delta_line,
-                delta_char,
+            token_entries.append((
+                struct.line,
+                name_start,
                 len(struct.name),
                 TOKEN_TYPES.index("struct"),
                 0
-            ])
-
-            prev_line = struct.line
-            prev_char = name_start + len(struct.name)
+            ))
 
+    # Add CuTeDSL decorator tokens
     for i, line in enumerate(lines):
         if "@cute.kernel" in line or "@cute.struct" in line:
             decorator_start = line.find("@")
@@ -107,18 +125,100 @@ def handle_semantic_tokens(uri: str, content: str) -> SemanticTokens:
                 if decorator_end == -1:
                     decorator_end = len(line)
 
-                delta_line = i - prev_line
-                delta_char = decorator_start - (prev_char if delta_line == 0 else 0)
-
-                tokens.extend([
-                    delta_line,
-                    delta_char,
+                token_entries.append((
+                    i,
+                    decorator_start,
                     decorator_end - decorator_start,
                     TOKEN_TYPES.index("decorator"),
                     0
-                ])
+                ))
+
+    # Add HIP-specific tokens if this is a HIP or CUDA file
+    if language_info.language in ("hip", "cuda", "cpp"):
+        token_entries.extend(_get_hip_tokens(lines))
 
-                prev_line = i
-                prev_char = decorator_end
+    # Sort tokens by position (line, then character)
+    token_entries.sort(key=lambda x: (x[0], x[1]))
+
+    # Convert to delta-encoded format
+    tokens: list[int] = []
+    prev_line = 0
+    prev_char = 0
+
+    for line, char, length, token_type, modifier in token_entries:
+        delta_line = line - prev_line
+        delta_char = char - (prev_char if delta_line == 0 else 0)
+
+        tokens.extend([
+            delta_line,
+            delta_char,
+            length,
+            token_type,
+            modifier
+        ])
+
+        prev_line = line
+        prev_char = char
 
     return SemanticTokens(data=tokens)
+
+
+def _get_hip_tokens(lines: list[str]) -> list[tuple[int, int, int, int, int]]:
+    """Extract HIP-specific semantic tokens from code.
+    
+    Returns list of (line, char, length, token_type, modifier) tuples.
+    """
+    token_entries: list[tuple[int, int, int, int, int]] = []
+    
+    for i, line in enumerate(lines):
+        # GPU keywords (__global__, __device__, etc.)
+        for match in HIP_KEYWORD_PATTERN.finditer(line):
+            token_entries.append((
+                i,
+                match.start(),
+                len(match.group()),
+                TOKEN_TYPES.index("keyword_gpu"),
+                0
+            ))
+        
+        # Memory keywords (__shared__, __constant__)
+        for match in HIP_MEMORY_KEYWORD_PATTERN.finditer(line):
+            token_entries.append((
+                i,
+                match.start(),
+                len(match.group()),
+                TOKEN_TYPES.index("keyword_memory"),
+                0
+            ))
+        
+        # __launch_bounds__
+        for match in HIP_LAUNCH_BOUNDS_PATTERN.finditer(line):
+            token_entries.append((
+                i,
+                match.start(),
+                len(match.group()),
+                TOKEN_TYPES.index("keyword_gpu"),
+                0
+            ))
+        
+        # HIP API functions
+        for match in HIP_API_PATTERN.finditer(line):
+            token_entries.append((
+                i,
+                match.start(),
+                len(match.group()),
+                TOKEN_TYPES.index("function_hip_api"),
+                0
+            ))
+        
+        # Wavefront intrinsics
+        for match in HIP_INTRINSIC_PATTERN.finditer(line):
+            token_entries.append((
+                i,
+                match.start(),
+                len(match.group()),
+                TOKEN_TYPES.index("function_intrinsic"),
+                0
+            ))
+    
+    return token_entries

wafer_lsp/languages/detector.py

@@ -2,33 +2,106 @@ 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".{ext}"
-        self._extensions[normalized_ext] = language_id
+        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) -> str | None:
+    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)
+        return self.detect_from_path(file_path, content)
 
-    def detect_from_path(self, file_path: str) -> str | None:
+    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".{ext}"
+        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]:
-        return list(self._extensions.keys())
+        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 is_supported(self, uri: str) -> bool:
-        return self.detect_from_uri(uri) 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/registry.py

@@ -1,6 +1,7 @@
 
 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
@@ -33,6 +34,22 @@ class LanguageRegistry:
             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++",
@@ -47,12 +64,16 @@ class LanguageRegistry:
         display_name: str,
         parser,
         extensions: list[str],
-        file_patterns: list[str] | None = None
+        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)

wafer_lsp/parsers/__init__.py

@@ -6,6 +6,16 @@ from .cutedsl_parser import (
     CuTeDSLParser,
     CuTeDSLStruct,
 )
+from .hip_parser import (
+    HIPKernel,
+    HIPDeviceFunction,
+    HIPParameter,
+    HIPParser,
+    KernelLaunchSite,
+    SharedMemoryAllocation,
+    WavefrontPattern,
+    is_hip_file,
+)
 
 __all__ = [
     "BaseParser",
@@ -15,4 +25,12 @@ __all__ = [
     "CuTeDSLLayout",
     "CuTeDSLParser",
     "CuTeDSLStruct",
+    "HIPKernel",
+    "HIPDeviceFunction",
+    "HIPParameter",
+    "HIPParser",
+    "KernelLaunchSite",
+    "SharedMemoryAllocation",
+    "WavefrontPattern",
+    "is_hip_file",
 ]