wafer-lsp 0.1.13__py3-none-any.whl

wafer_lsp/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

wafer_lsp/__main__.py

@@ -0,0 +1,9 @@
+from wafer_lsp.server import server
+
+
+def main():
+    server.start_io()
+
+
+if __name__ == "__main__":
+    main()

wafer_lsp/analyzers/compiler_integration.py

@@ -0,0 +1,16 @@
+from typing import Any
+
+_analysis_cache: dict[str, dict[str, Any]] = {}
+
+
+def get_analysis_for_kernel(uri: str, kernel_name: str) -> dict[str, Any] | None:
+    cache_key = f"{uri}:{kernel_name}"
+
+    if cache_key in _analysis_cache:
+        return _analysis_cache[cache_key]
+
+    return None
+
+
+def clear_cache():
+    _analysis_cache.clear()

wafer_lsp/analyzers/docs_index.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+
+class DocsIndex:
+
+    def __init__(self, docs_path: str | None = None):
+        if docs_path:
+            self.docs_path = Path(docs_path)
+        else:
+            self.docs_path = Path(__file__).parent.parent.parent.parent.parent / \
+                "curriculum" / "cutlass-docs" / "cutedsl-docs"
+
+        self.index = self._build_index()
+
+    def _build_index(self) -> dict[str, list[str]]:
+        return {
+            "layout": [
+                "intro-to-cutedsl.md",
+                "partitioning-strategies-inner-outer-threadvalue.md"
+            ],
+            "TMA": ["blackwell-tutorial-fp16-gemm-0.md"],
+            "TMEM": ["colfax-blackwell-umma-tensor-memory-part1.md"],
+            "kernel": ["blackwell-tutorial-fp16-gemm-0.md"],
+            "struct": ["blackwell-tutorial-fp16-gemm-0.md"],
+            "pipeline": ["blackwell-tutorial-fp16-gemm-0.md"],
+            "MMA": ["mma-atoms-fundamentals-sm70-example.md"],
+        }
+
+    def get_doc_for_concept(self, concept: str) -> str | None:
+        concept_lower = concept.lower()
+        if self.index.get(concept_lower):
+            doc_file = self.index[concept_lower][0]
+            doc_path = self.docs_path / doc_file
+            if doc_path.exists():
+                return str(doc_path)
+        return None

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/code_action.py

@@ -0,0 +1,48 @@
+
+from lsprotocol.types import CodeAction, CodeActionKind, Command, Range
+
+from ..languages.registry import get_language_registry
+from ..languages.types import KernelInfo
+
+
+def find_kernel_at_range(content: str, range: Range, uri: str) -> KernelInfo | None:
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return None
+
+    for kernel in language_info.kernels:
+        if kernel.line <= range.start.line <= kernel.line + 50:
+            return kernel
+
+    return None
+
+
+def handle_code_action(uri: str, range: Range, content: str) -> list[CodeAction]:
+    kernel = find_kernel_at_range(content, range, uri)
+    if not kernel:
+        return []
+
+    actions: list[CodeAction] = [
+        CodeAction(
+            title=f"Analyze Kernel: {kernel.name}",
+            kind=CodeActionKind.Source,
+            command=Command(
+                title=f"Analyze Kernel: {kernel.name}",
+                command="wafer.analyzeKernel",
+                arguments=[uri, kernel.name]
+            )
+        ),
+        CodeAction(
+            title=f"Profile Kernel: {kernel.name}",
+            kind=CodeActionKind.Source,
+            command=Command(
+                title=f"Profile Kernel: {kernel.name}",
+                command="wafer.profileKernel",
+                arguments=[uri, kernel.name]
+            )
+        ),
+    ]
+
+    return actions

wafer_lsp/handlers/code_lens.py

@@ -0,0 +1,48 @@
+
+from lsprotocol.types import CodeLens, Command, Position, Range
+
+from ..languages.registry import get_language_registry
+
+
+def handle_code_lens(uri: str, content: str) -> list[CodeLens]:
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return []
+
+    lenses: list[CodeLens] = []
+
+    for kernel in language_info.kernels:
+        lens_range = Range(
+            start=Position(line=kernel.line, character=0),
+            end=Position(line=kernel.line, character=0)
+        )
+
+        analyze_command = Command(
+            title=f"Analyze {kernel.name}",
+            command="wafer.analyzeKernel",
+            arguments=[uri, kernel.name]
+        )
+
+        profile_command = Command(
+            title=f"Profile {kernel.name}",
+            command="wafer.profileKernel",
+            arguments=[uri, kernel.name]
+        )
+
+        lenses.append(CodeLens(
+            range=lens_range,
+            command=analyze_command
+        ))
+
+        profile_range = Range(
+            start=Position(line=kernel.line, character=20),
+            end=Position(line=kernel.line, character=20)
+        )
+        lenses.append(CodeLens(
+            range=profile_range,
+            command=profile_command
+        ))
+
+    return lenses

wafer_lsp/handlers/completion.py

@@ -0,0 +1,6 @@
+
+from lsprotocol.types import CompletionItem
+
+
+def handle_completion(uri: str, position, content: str) -> list[CompletionItem]:
+    return []

wafer_lsp/handlers/diagnostics.py

@@ -0,0 +1,41 @@
+
+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,
+    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()
+    language_info = registry.parse_file(uri, content)
+
+    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

@@ -0,0 +1,176 @@
+
+from lsprotocol.types import DocumentSymbol, Position, Range, SymbolKind
+
+from ..languages.registry import get_language_registry
+
+
+def handle_document_symbol(uri: str, content: str) -> list[DocumentSymbol]:
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return []
+
+    symbols: list[DocumentSymbol] = []
+    lines = content.split("\n")
+
+    # Kernels
+    for kernel in language_info.kernels:
+        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
+
+        selection_range = Range(
+            start=Position(line=kernel.line, character=max(0, name_start)),
+            end=Position(line=kernel.line, character=name_end)
+        )
+        full_range = Range(
+            start=Position(line=kernel.line, character=0),
+            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=detail,
+        ))
+
+    # Layouts (CuTeDSL)
+    for layout in language_info.layouts:
+        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
+
+        detail = f"Layout: {layout.shape}" if layout.shape else "Layout"
+
+        selection_range = Range(
+            start=Position(line=layout.line, character=max(0, name_start)),
+            end=Position(line=layout.line, character=name_end)
+        )
+        full_range = Range(
+            start=Position(line=layout.line, character=0),
+            end=Position(line=layout.line, character=len(layout_line))
+        )
+
+        symbols.append(DocumentSymbol(
+            name=layout.name,
+            kind=SymbolKind.Variable,
+            range=full_range,
+            selection_range=selection_range,
+            detail=detail,
+        ))
+
+    # Structs
+    for struct in language_info.structs:
+        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
+
+        selection_range = Range(
+            start=Position(line=struct.line, character=max(0, name_start)),
+            end=Position(line=struct.line, character=name_end)
+        )
+        full_range = Range(
+            start=Position(line=struct.line, character=0),
+            end=Position(line=min(struct.line + 10, len(lines) - 1), character=0)
+        )
+
+        symbols.append(DocumentSymbol(
+            name=struct.name,
+            kind=SymbolKind.Struct,
+            range=full_range,
+            selection_range=selection_range,
+            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