wafer-lsp 0.1.13__py3-none-any.whl

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,38 @@
+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 .hip_docs import (
+    HIPDocsService,
+    HIPAPIDoc,
+    MemoryQualifierDoc,
+    IntrinsicDoc,
+    ThreadIndexDoc,
+    create_hip_docs_service,
+)
+from .hip_hover_service import HIPHoverService, create_hip_hover_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",
+    "HIPDocsService",
+    "HIPAPIDoc",
+    "HIPHoverService",
+    "HoverService",
+    "IntrinsicDoc",
+    "LanguageRegistryService",
+    "MemoryQualifierDoc",
+    "PositionService",
+    "ThreadIndexDoc",
+    "create_analysis_service",
+    "create_docs_service",
+    "create_document_service",
+    "create_hip_docs_service",
+    "create_hip_hover_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)