wafer-lsp 0.1.0__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/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,16 @@
+
+from lsprotocol.types import Diagnostic
+
+from ..languages.registry import get_language_registry
+
+
+def handle_diagnostics(uri: str, content: str) -> list[Diagnostic]:
+    diagnostics: list[Diagnostic] = []
+
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return diagnostics
+
+    return diagnostics

wafer_lsp/handlers/document_symbol.py

@@ -0,0 +1,87 @@
+
+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] = []
+
+    for kernel in language_info.kernels:
+        lines = content.split("\n")
+        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)
+        )
+
+        symbols.append(DocumentSymbol(
+            name=kernel.name,
+            kind=SymbolKind.Function,
+            range=full_range,
+            selection_range=selection_range,
+            detail=f"GPU Kernel ({registry.get_language_name(kernel.language)})",
+        ))
+
+    for layout in language_info.layouts:
+        lines = content.split("\n")
+        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,
+        ))
+
+    for struct in language_info.structs:
+        lines = content.split("\n")
+        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)})",
+        ))
+
+    return symbols

wafer_lsp/handlers/hover.py

@@ -0,0 +1,215 @@
+
+from lsprotocol.types import Hover, MarkupContent, MarkupKind, Position
+
+from ..analyzers.compiler_integration import get_analysis_for_kernel
+from ..analyzers.docs_index import DocsIndex
+from ..languages.registry import get_language_registry
+from ..languages.types import KernelInfo, LayoutInfo
+from ..utils.lsp_helpers import get_decorator_at_position, get_word_at_position
+
+
+def find_kernel_at_position(
+    content: str, position: Position, uri: str
+) -> KernelInfo | None:
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return None
+
+    word = 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(
+    content: str, position: Position, uri: str
+) -> LayoutInfo | None:
+    registry = get_language_registry()
+    language_info = registry.parse_file(uri, content)
+
+    if not language_info:
+        return None
+
+    word = get_word_at_position(content, position)
+
+    for layout in language_info.layouts:
+        if layout.name == word:
+            return layout
+
+    return None
+
+
+def format_kernel_hover(kernel: KernelInfo, analysis: dict | None) -> str:
+    registry = get_language_registry()
+    language_name = registry.get_language_name(kernel.language) or kernel.language
+
+    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 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(kernel: KernelInfo) -> str:
+    return format_kernel_hover(kernel, None)
+
+
+def format_decorator_hover(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")
+
+    docs = DocsIndex()
+    doc_link = docs.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 handle_hover(uri: str, position: Position, content: str) -> Hover | None:
+    test_message = "**HEYOOOOOO** 🎉\n\n"
+
+    decorator_info = 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 + format_decorator_hover(decorator_name, function_name)
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    word = 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 = find_kernel_at_position(content, position, uri)
+    if kernel:
+        analysis = get_analysis_for_kernel(uri, kernel.name)
+
+        if analysis:
+            hover_content = test_message + format_kernel_hover(kernel, analysis)
+        else:
+            hover_content = test_message + format_kernel_hover_basic(kernel)
+
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    layout = find_layout_at_position(content, position, uri)
+    if layout:
+        docs = DocsIndex()
+        doc_link = docs.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! Move your cursor over symbols to see more info."
+    return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))

wafer_lsp/handlers/inlay_hint.py

@@ -0,0 +1,65 @@
+
+from lsprotocol.types import InlayHint, InlayHintKind, Position, Range
+
+from ..languages.registry import get_language_registry
+
+
+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")
+
+    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
+            ))
+
+    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
+            ))
+
+    return hints

wafer_lsp/handlers/semantic_tokens.py

@@ -0,0 +1,124 @@
+
+from lsprotocol.types import SemanticTokens, SemanticTokensLegend
+
+from ..languages.registry import get_language_registry
+
+TOKEN_TYPES = [
+    "kernel",
+    "layout",
+    "struct",
+    "decorator",
+]
+
+TOKEN_MODIFIERS = [
+    "definition",
+    "declaration",
+]
+
+SEMANTIC_TOKENS_LEGEND = SemanticTokensLegend(
+    token_types=TOKEN_TYPES,
+    token_modifiers=TOKEN_MODIFIERS
+)
+
+
+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=[])
+
+    tokens: list[int] = []
+    lines = content.split("\n")
+    prev_line = 0
+    prev_char = 0
+
+    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:
+            delta_line = kernel.line - prev_line
+            delta_char = name_start - (prev_char if delta_line == 0 else 0)
+
+            tokens.extend([
+                delta_line,
+                delta_char,
+                len(kernel.name),
+                TOKEN_TYPES.index("kernel"),
+                0
+            ])
+
+            prev_line = kernel.line
+            prev_char = name_start + len(kernel.name)
+
+    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:
+            delta_line = layout.line - prev_line
+            delta_char = name_start - (prev_char if delta_line == 0 else 0)
+
+            tokens.extend([
+                delta_line,
+                delta_char,
+                len(layout.name),
+                TOKEN_TYPES.index("layout"),
+                0
+            ])
+
+            prev_line = layout.line
+            prev_char = name_start + len(layout.name)
+
+    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:
+            delta_line = struct.line - prev_line
+            delta_char = name_start - (prev_char if delta_line == 0 else 0)
+
+            tokens.extend([
+                delta_line,
+                delta_char,
+                len(struct.name),
+                TOKEN_TYPES.index("struct"),
+                0
+            ])
+
+            prev_line = struct.line
+            prev_char = name_start + len(struct.name)
+
+    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)
+
+                delta_line = i - prev_line
+                delta_char = decorator_start - (prev_char if delta_line == 0 else 0)
+
+                tokens.extend([
+                    delta_line,
+                    delta_char,
+                    decorator_end - decorator_start,
+                    TOKEN_TYPES.index("decorator"),
+                    0
+                ])
+
+                prev_line = i
+                prev_char = decorator_end
+
+    return SemanticTokens(data=tokens)