wafer-lsp 0.1.0__py3-none-any.whl

wafer_lsp/parsers/cuda_parser.py

@@ -0,0 +1,95 @@
+import re
+from dataclasses import dataclass
+from typing import Any
+
+from .base_parser import BaseParser
+
+
+@dataclass
+class CUDAKernel:
+    name: str
+    line: int
+    parameters: list[str]
+
+
+class CUDAParser(BaseParser):
+
+    def parse_file(self, content: str) -> dict[str, Any]:
+        kernels: list[CUDAKernel] = []
+
+        pattern = r'__global__\s+(?:__device__\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(CUDAKernel(
+                name=kernel_name,
+                line=line,
+                parameters=params
+            ))
+
+        return {"kernels": kernels}
+
+    def _extract_parameters(self, content: str, start: int) -> list[str]:
+        if start >= len(content):
+            return []
+
+        depth = 0
+        param_start = start
+        param_end = start
+
+        for i in range(start, len(content)):
+            char = content[i]
+            if char == '(':
+                depth += 1
+            elif char == ')':
+                depth -= 1
+                if depth == 0:
+                    param_end = i
+                    break
+
+        if param_end == start:
+            return []
+
+        param_str = content[param_start:param_end + 1]
+
+        params: list[str] = []
+        current_param = ""
+        template_depth = 0
+        paren_depth = 0
+
+        for char in param_str[1:-1]:
+            if char == '<':
+                template_depth += 1
+                current_param += char
+            elif char == '>':
+                template_depth -= 1
+                current_param += char
+            elif char == '(':
+                paren_depth += 1
+                current_param += char
+            elif char == ')':
+                paren_depth -= 1
+                current_param += char
+            elif char == ',' and template_depth == 0 and paren_depth == 0:
+                param_name = current_param.strip()
+                if param_name:
+                    parts = param_name.split()
+                    if parts:
+                        name = parts[-1].strip('*&')
+                        params.append(name)
+                current_param = ""
+            else:
+                current_param += char
+
+        if current_param.strip():
+            param_name = current_param.strip()
+            parts = param_name.split()
+            if parts:
+                name = parts[-1].strip('*&')
+                params.append(name)
+
+        return params

wafer_lsp/parsers/cutedsl_parser.py

@@ -0,0 +1,114 @@
+import ast
+from dataclasses import dataclass
+from typing import Any
+
+from .base_parser import BaseParser
+
+
+@dataclass
+class CuTeDSLKernel:
+    name: str
+    line: int
+    parameters: list[str]
+    docstring: str | None = None
+
+
+@dataclass
+class CuTeDSLLayout:
+    name: str
+    line: int
+    shape: str | None = None
+    stride: str | None = None
+
+
+@dataclass
+class CuTeDSLStruct:
+    name: str
+    line: int
+    docstring: str | None = None
+
+
+class CuTeDSLParser(BaseParser):
+
+    def parse_file(self, content: str) -> dict[str, Any]:
+        try:
+            tree = ast.parse(content)
+        except SyntaxError:
+            return {"kernels": [], "layouts": [], "structs": []}
+
+        kernels: list[CuTeDSLKernel] = []
+        layouts: list[CuTeDSLLayout] = []
+        structs: list[CuTeDSLStruct] = []
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef):
+                if self._has_decorator(node, "cute.kernel"):
+                    kernels.append(self._extract_kernel(node))
+            elif isinstance(node, ast.ClassDef):
+                if self._has_decorator(node, "cute.struct"):
+                    structs.append(self._extract_struct(node))
+            elif isinstance(node, ast.Assign):
+                layout = self._extract_layout(node, content)
+                if layout:
+                    layouts.append(layout)
+
+        return {"kernels": kernels, "layouts": layouts, "structs": structs}
+
+    def _has_decorator(self, node: ast.FunctionDef | ast.ClassDef, decorator: str) -> bool:
+        for dec in node.decorator_list:
+            if isinstance(dec, ast.Attribute):
+                if isinstance(dec.value, ast.Name) and dec.value.id == "cute":
+                    if dec.attr == decorator.split(".")[-1]:
+                        return True
+            elif isinstance(dec, ast.Name):
+                if dec.id == decorator.split(".")[-1]:
+                    return True
+        return False
+
+    def _extract_kernel(self, node: ast.FunctionDef) -> CuTeDSLKernel:
+        parameters = [arg.arg for arg in node.args.args]
+        docstring = ast.get_docstring(node)
+
+        return CuTeDSLKernel(
+            name=node.name,
+            line=node.lineno - 1,
+            parameters=parameters,
+            docstring=docstring,
+        )
+
+    def _extract_struct(self, node: ast.ClassDef) -> CuTeDSLStruct:
+        docstring = ast.get_docstring(node)
+
+        return CuTeDSLStruct(
+            name=node.name,
+            line=node.lineno - 1,
+            docstring=docstring,
+        )
+
+    def _extract_layout(self, node: ast.Assign, content: str) -> CuTeDSLLayout | None:
+        for target in node.targets:
+            if not isinstance(target, ast.Name):
+                continue
+
+            if isinstance(node.value, ast.Call):
+                call = node.value
+                if isinstance(call.func, ast.Attribute):
+                    if isinstance(call.func.value, ast.Name):
+                        if call.func.value.id == "cute" and call.func.attr == "make_layout":
+                            shape_str = None
+                            stride_str = None
+
+                            if call.args:
+                                try:
+                                    shape_str = ast.unparse(call.args[0])
+                                except AttributeError:
+                                    shape_str = str(call.args[0])
+
+                            return CuTeDSLLayout(
+                                name=target.id,
+                                line=node.lineno - 1,
+                                shape=shape_str,
+                                stride=stride_str,
+                            )
+
+        return None

wafer_lsp/server.py

@@ -0,0 +1,58 @@
+from lsprotocol.types import (
+    INITIALIZE,
+    TEXT_DOCUMENT_HOVER,
+)
+from pygls.lsp.server import LanguageServer
+
+from .services import (
+    create_analysis_service,
+    create_docs_service,
+    create_document_service,
+    create_hover_service,
+    create_language_registry_service,
+    create_position_service,
+)
+
+language_registry_service = create_language_registry_service()
+analysis_service = create_analysis_service()
+docs_service = create_docs_service()
+position_service = create_position_service()
+hover_service = create_hover_service(
+    language_registry_service,
+    analysis_service,
+    docs_service,
+    position_service
+)
+
+server = LanguageServer("wafer-lsp", "1.0.0")
+document_service = create_document_service(server)
+
+
+@server.feature(INITIALIZE)
+def initialize(params):
+    return {
+        "capabilities": {
+            "hoverProvider": True,
+        }
+    }
+
+
+@server.feature(TEXT_DOCUMENT_HOVER)
+def hover(params):
+    uri = params.text_document.uri
+    position = params.position
+    content = document_service.get_document_content(uri)
+    if not content:
+        test_message = "🎉🎉🎉 **HEYOOO!!! LSP IS DEFINITELY WORKING!!!** 🎉🎉🎉\n\n**THIS IS THE WAFER LSP SERVER!**\n\n**Document content not available, but LSP is running!**"
+        from lsprotocol.types import Hover, MarkupContent, MarkupKind
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=test_message))
+    result = hover_service.handle_hover(uri, position, content)
+    if not result:
+        test_message = "🎉🎉🎉 **HEYOOO!!! LSP IS DEFINITELY WORKING!!!** 🎉🎉🎉\n\n**THIS IS THE WAFER LSP SERVER!**\n\n**Hover service returned None, but LSP is running!**"
+        from lsprotocol.types import Hover, MarkupContent, MarkupKind
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=test_message))
+    return result
+
+
+if __name__ == "__main__":
+    server.start_io()

wafer_lsp/services/__init__.py

@@ -0,0 +1,21 @@
+from .analysis_service import AnalysisService, create_analysis_service
+from .docs_service import DocsService, create_docs_service
+from .document_service import DocumentService, create_document_service
+from .hover_service import HoverService, create_hover_service
+from .language_registry_service import LanguageRegistryService, create_language_registry_service
+from .position_service import PositionService, create_position_service
+
+__all__ = [
+    "AnalysisService",
+    "DocsService",
+    "DocumentService",
+    "HoverService",
+    "LanguageRegistryService",
+    "PositionService",
+    "create_analysis_service",
+    "create_docs_service",
+    "create_document_service",
+    "create_hover_service",
+    "create_language_registry_service",
+    "create_position_service",
+]

wafer_lsp/services/analysis_service.py

@@ -0,0 +1,22 @@
+from typing import Any
+
+
+class AnalysisService:
+
+    def __init__(self):
+        self._cache: dict[str, dict[str, Any]] = {}
+
+    def get_analysis_for_kernel(self, uri: str, kernel_name: str) -> dict[str, Any] | None:
+        cache_key = f"{uri}:{kernel_name}"
+        return self._cache.get(cache_key)
+
+    def set_analysis_for_kernel(self, uri: str, kernel_name: str, analysis: dict[str, Any]):
+        cache_key = f"{uri}:{kernel_name}"
+        self._cache[cache_key] = analysis
+
+    def clear_cache(self):
+        self._cache.clear()
+
+
+def create_analysis_service() -> AnalysisService:
+    return AnalysisService()

wafer_lsp/services/docs_service.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+
+
+class DocsService:
+
+    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):
+        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
+
+
+def create_docs_service(docs_path: str | None = None) -> DocsService:
+    return DocsService(docs_path)

wafer_lsp/services/document_service.py

@@ -0,0 +1,20 @@
+from pygls.lsp.server import LanguageServer
+
+
+class DocumentService:
+
+    def __init__(self, server: LanguageServer):
+        self._server = server
+
+    def get_document(self, uri: str):
+        return self._server.workspace.text_documents.get(uri)
+
+    def get_document_content(self, uri: str) -> str:
+        doc = self.get_document(uri)
+        if doc is None:
+            return ''
+        return getattr(doc, 'text', getattr(doc, 'source', ''))
+
+
+def create_document_service(server: LanguageServer) -> DocumentService:
+    return DocumentService(server)

wafer_lsp/services/hover_service.py

@@ -0,0 +1,237 @@
+
+from lsprotocol.types import Hover, MarkupContent, MarkupKind, Position
+
+from ..languages.types import KernelInfo, LayoutInfo
+from .analysis_service import AnalysisService
+from .docs_service import DocsService
+from .language_registry_service import LanguageRegistryService
+from .position_service import PositionService
+
+
+class HoverService:
+
+    def __init__(
+        self,
+        language_registry: LanguageRegistryService,
+        analysis_service: AnalysisService,
+        docs_service: DocsService,
+        position_service: PositionService
+    ):
+        self._language_registry = language_registry
+        self._analysis_service = analysis_service
+        self._docs_service = docs_service
+        self._position_service = position_service
+
+    def handle_hover(self, uri: str, position: Position, content: str) -> Hover | None:
+        test_message = "🎉🎉🎉 **HEYOOO!!! LSP IS DEFINITELY WORKING!!!** 🎉🎉🎉\n\n**THIS IS THE WAFER LSP SERVER!**\n\n"
+
+        decorator_info = self._position_service.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 + self._format_decorator_hover(decorator_name, function_name)
+            return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+        word = self._position_service.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 = self._find_kernel_at_position(content, position, uri)
+        if kernel:
+            analysis = self._analysis_service.get_analysis_for_kernel(uri, kernel.name)
+
+            if analysis:
+                hover_content = test_message + self._format_kernel_hover(kernel, analysis)
+            else:
+                hover_content = test_message + self._format_kernel_hover_basic(kernel)
+
+            return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+        layout = self._find_layout_at_position(content, position, uri)
+        if layout:
+            doc_link = self._docs_service.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!** 🚀\n\nMove your cursor over any symbol, decorator, or even empty space to see LSP information.\n\n**Try hovering over:**\n- `@cute.kernel` decorators\n- `cute` module name\n- Kernel function names\n- Layout variables"
+        return Hover(contents=MarkupContent(kind=MarkupKind.Markdown, value=hover_content))
+
+    def _find_kernel_at_position(
+        self, content: str, position: Position, uri: str
+    ) -> KernelInfo | None:
+        language_info = self._language_registry.parse_file(uri, content)
+
+        if not language_info:
+            return None
+
+        word = self._position_service.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(
+        self, content: str, position: Position, uri: str
+    ) -> LayoutInfo | None:
+        language_info = self._language_registry.parse_file(uri, content)
+
+        if not language_info:
+            return None
+
+        word = self._position_service.get_word_at_position(content, position)
+
+        for layout in language_info.layouts:
+            if layout.name == word:
+                return layout
+
+        return None
+
+    def _format_kernel_hover(self, kernel: KernelInfo, analysis: dict | None) -> str:
+        language_name = self._language_registry.get_language_name(kernel.language) or kernel.language
+
+        if kernel.language == "cuda" or kernel.language == "cpp":
+            lines = [f"**CUDA Kernel: {kernel.name}**", f"*Language: {language_name}*", ""]
+        else:
+            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 kernel.language == "cuda" or kernel.language == "cpp":
+            lines.append("**CUDA Features:**")
+            lines.append("- `__global__` function executed on GPU")
+            lines.append("- Can be launched with `<<<grid, block>>>` syntax")
+            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(self, kernel: KernelInfo) -> str:
+        return self._format_kernel_hover(kernel, None)
+
+    def _format_decorator_hover(self, 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")
+
+        doc_link = self._docs_service.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 create_hover_service(
+    language_registry: LanguageRegistryService,
+    analysis_service: AnalysisService,
+    docs_service: DocsService,
+    position_service: PositionService
+) -> HoverService:
+    return HoverService(language_registry, analysis_service, docs_service, position_service)

wafer_lsp/services/language_registry_service.py

@@ -0,0 +1,26 @@
+
+from ..languages.registry import LanguageRegistry, get_language_registry
+from ..languages.types import LanguageInfo
+
+
+class LanguageRegistryService:
+
+    def __init__(self, registry: LanguageRegistry):
+        self._registry = registry
+
+    def detect_language(self, uri: str) -> str | None:
+        return self._registry.detect_language(uri)
+
+    def parse_file(self, uri: str, content: str) -> LanguageInfo | None:
+        return self._registry.parse_file(uri, content)
+
+    def get_language_name(self, language_id: str) -> str | None:
+        return self._registry.get_language_name(language_id)
+
+    def get_supported_extensions(self) -> list[str]:
+        return self._registry.get_supported_extensions()
+
+
+def create_language_registry_service() -> LanguageRegistryService:
+    registry = get_language_registry()
+    return LanguageRegistryService(registry)