wafer-lsp 0.1.13__py3-none-any.whl

wafer_lsp/handlers/inlay_hint.py

@@ -0,0 +1,245 @@
+
+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)
+
+    if not language_info:
+        return []
+
+    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
+
+        layout_line = lines[layout.line] if layout.line < len(lines) else ""
+
+        if "=" in layout_line:
+            equals_pos = layout_line.find("=")
+            hint_text = ": Layout"
+            if layout.shape:
+                hint_text = f": Layout[Shape{layout.shape}]"
+
+            hint_position = Position(
+                line=layout.line,
+                character=equals_pos + 1
+            )
+
+            hints.append(InlayHint(
+                position=hint_position,
+                label=hint_text,
+                kind=InlayHintKind.Type,
+                padding_left=True,
+                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
+
+        kernel_line = lines[kernel.line] if kernel.line < len(lines) else ""
+
+        if "def " in kernel_line and "(" in kernel_line:
+            paren_pos = kernel_line.find("(")
+            hint_text = " -> Kernel"
+
+            hint_position = Position(
+                line=kernel.line,
+                character=paren_pos
+            )
+
+            hints.append(InlayHint(
+                position=hint_position,
+                label=hint_text,
+                kind=InlayHintKind.Type,
+                padding_left=True,
+                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

@@ -0,0 +1,224 @@
+
+import re
+
+from lsprotocol.types import SemanticTokens, SemanticTokensLegend
+
+from ..languages.registry import get_language_registry
+
+TOKEN_TYPES = [
+    "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 = [
+    "definition",
+    "declaration",
+]
+
+SEMANTIC_TOKENS_LEGEND = SemanticTokensLegend(
+    token_types=TOKEN_TYPES,
+    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()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return SemanticTokens(data=[])
+
+    # 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")
+
+    # Add kernel tokens
+    for kernel in language_info.kernels:
+        if kernel.line >= len(lines):
+            continue
+
+        kernel_line = lines[kernel.line]
+        name_start = kernel_line.find(kernel.name)
+
+        if name_start >= 0:
+            token_entries.append((
+                kernel.line,
+                name_start,
+                len(kernel.name),
+                TOKEN_TYPES.index("kernel"),
+                0
+            ))
+
+    # Add layout tokens
+    for layout in language_info.layouts:
+        if layout.line >= len(lines):
+            continue
+
+        layout_line = lines[layout.line]
+        name_start = layout_line.find(layout.name)
+
+        if name_start >= 0:
+            token_entries.append((
+                layout.line,
+                name_start,
+                len(layout.name),
+                TOKEN_TYPES.index("layout"),
+                0
+            ))
+
+    # Add struct tokens
+    for struct in language_info.structs:
+        if struct.line >= len(lines):
+            continue
+
+        struct_line = lines[struct.line]
+        name_start = struct_line.find(struct.name)
+
+        if name_start >= 0:
+            token_entries.append((
+                struct.line,
+                name_start,
+                len(struct.name),
+                TOKEN_TYPES.index("struct"),
+                0
+            ))
+
+    # Add CuTeDSL decorator tokens
+    for i, line in enumerate(lines):
+        if "@cute.kernel" in line or "@cute.struct" in line:
+            decorator_start = line.find("@")
+            if decorator_start >= 0:
+                decorator_end = line.find(" ", decorator_start)
+                if decorator_end == -1:
+                    decorator_end = len(line)
+
+                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))
+
+    # 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/handlers/workspace_symbol.py

@@ -0,0 +1,87 @@
+
+from lsprotocol.types import Location, Position, Range, SymbolKind, WorkspaceSymbol
+
+from ..languages.registry import get_language_registry
+
+
+def _matches_query(name: str, query: str) -> bool:
+    if not query:
+        return True
+
+    name_lower = name.lower()
+    query_lower = query.lower()
+
+    query_idx = 0
+    for char in name_lower:
+        if query_idx < len(query_lower) and char == query_lower[query_idx]:
+            query_idx += 1
+
+    return query_idx == len(query_lower)
+
+
+def handle_workspace_symbol(query: str) -> list[WorkspaceSymbol]:
+    registry = get_language_registry()
+    symbols: list[WorkspaceSymbol] = []
+
+    return symbols
+
+
+def handle_workspace_symbol_with_documents(
+    query: str,
+    document_contents: dict[str, str]
+) -> list[WorkspaceSymbol]:
+    registry = get_language_registry()
+    symbols: list[WorkspaceSymbol] = []
+
+    for uri, content in document_contents.items():
+        language_info = registry.parse_file(uri, content)
+
+        if not language_info:
+            continue
+
+        for kernel in language_info.kernels:
+            if _matches_query(kernel.name, query):
+                symbols.append(WorkspaceSymbol(
+                    name=kernel.name,
+                    kind=SymbolKind.Function,
+                    location=Location(
+                        uri=uri,
+                        range=Range(
+                            start=Position(line=kernel.line, character=0),
+                            end=Position(line=kernel.line, character=0)
+                        )
+                    ),
+                    container_name=f"GPU Kernel ({registry.get_language_name(kernel.language)})"
+                ))
+
+        for layout in language_info.layouts:
+            if _matches_query(layout.name, query):
+                symbols.append(WorkspaceSymbol(
+                    name=layout.name,
+                    kind=SymbolKind.Variable,
+                    location=Location(
+                        uri=uri,
+                        range=Range(
+                            start=Position(line=layout.line, character=0),
+                            end=Position(line=layout.line, character=0)
+                        )
+                    ),
+                    container_name="Layout"
+                ))
+
+        for struct in language_info.structs:
+            if _matches_query(struct.name, query):
+                symbols.append(WorkspaceSymbol(
+                    name=struct.name,
+                    kind=SymbolKind.Struct,
+                    location=Location(
+                        uri=uri,
+                        range=Range(
+                            start=Position(line=struct.line, character=0),
+                            end=Position(line=struct.line, character=0)
+                        )
+                    ),
+                    container_name=f"Struct ({registry.get_language_name(struct.language)})"
+                ))
+
+    return symbols