wafer-lsp 0.1.0__tar.gz

wafer_lsp-0.1.0/.gitignore

@@ -0,0 +1,85 @@
+# Python
+*.pyc
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+*.egg-info/
+
+# Virtual environments
+.venv/
+venv/
+env/
+ENV/
+
+# Distribution / packaging
+dist/
+build/
+*.egg
+
+# IDE
+.vscode/
+.cursor/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Temporary files
+*.tmp
+tests.txt
+output.txt
+
+# Evaluation results
+experiments/emilio/cutedsl-quack-eval/eval_results/*.txt
+experiments/emilio/cutedsl-quack-eval/eval_results/*.json
+experiments/emilio/cutedsl-quack-eval/eval_results/*.py
+research/evals/kernelbench/results/
+wafer_artifacts/
+**/wafer_artifacts/
+/.deps/
+
+# Local config files
+services/docs-tool/config.yaml
+
+# Environment files with secrets
+.env
+.env.local
+.env.*.local
+node_modules/
+.vscode-test/
+
+# npm artifacts (using yarn)
+package-lock.json
+npm-debug.log*
+.npm/
+.npmrc
+
+# Yarn artifacts
+.yarn/install-state.gz
+
+# NSYS report files (can be very large)
+*.nsys-rep
+
+# Turbo build cache
+.turbo/
+**/.turbo/
+
+# Supabase CLI temp files
+supabase/.branches/
+supabase/.temp/
+**/supabase/.branches/
+**/supabase/.temp/
+
+# Eval results (generated, large)
+research/evals/kernelbench/kernelbench/results/
+research/evals/retrieval_eval_project/retrieval_eval/results/
+research/evals/retrieval_eval_project/*.log
+
+# Wafer extension trace files
+.wafer/traces/

wafer_lsp-0.1.0/PKG-INFO

@@ -0,0 +1,57 @@
+Metadata-Version: 2.4
+Name: wafer-lsp
+Version: 0.1.0
+Summary: Language Server Protocol server for GPU programming languages
+Author-email: Wafer <support@wafer.ai>
+License: MIT
+Keywords: c++,cuda,cutedsl,gpu,language-server,lsp
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: Text Editors :: Integrated Development Environments (IDE)
+Requires-Python: >=3.12
+Requires-Dist: lsprotocol>=2024.0.0
+Requires-Dist: pygls>=1.0.0
+Requires-Dist: wafer-core
+Description-Content-Type: text/markdown
+
+# Wafer LSP
+
+Language Server Protocol server for CuTeDSL (Python GPU programming).
+
+**Beta Feature**: Currently only available when Beta Mode is enabled in VS Code settings.
+
+## Features
+
+- **Hover Information**: Shows kernel and layout information with compiler analysis when hovering over CuTeDSL code
+
+## Installation
+
+```bash
+pip install wafer-lsp
+```
+
+## Usage
+
+### VS Code Extension
+
+The LSP server is integrated into the `wevin-extension` VS Code extension. It starts automatically when Beta Mode is enabled.
+
+### Standalone
+
+For Neovim or other editors:
+
+```bash
+python -m wafer_lsp
+```
+
+## Supported Languages
+
+- **CuTeDSL**: Python files with `@cute.kernel` decorators
+
+## Architecture
+
+The LSP server uses a modular language registry system. Currently supports CuTeDSL only. See `languages/README.md` for details on adding more languages.

wafer_lsp-0.1.0/README.md

@@ -0,0 +1,37 @@
+# Wafer LSP
+
+Language Server Protocol server for CuTeDSL (Python GPU programming).
+
+**Beta Feature**: Currently only available when Beta Mode is enabled in VS Code settings.
+
+## Features
+
+- **Hover Information**: Shows kernel and layout information with compiler analysis when hovering over CuTeDSL code
+
+## Installation
+
+```bash
+pip install wafer-lsp
+```
+
+## Usage
+
+### VS Code Extension
+
+The LSP server is integrated into the `wevin-extension` VS Code extension. It starts automatically when Beta Mode is enabled.
+
+### Standalone
+
+For Neovim or other editors:
+
+```bash
+python -m wafer_lsp
+```
+
+## Supported Languages
+
+- **CuTeDSL**: Python files with `@cute.kernel` decorators
+
+## Architecture
+
+The LSP server uses a modular language registry system. Currently supports CuTeDSL only. See `languages/README.md` for details on adding more languages.

wafer_lsp-0.1.0/pyproject.toml

@@ -0,0 +1,67 @@
+[project]
+name = "wafer-lsp"
+version = "0.1.0"
+description = "Language Server Protocol server for GPU programming languages"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "pygls>=1.0.0",
+    "lsprotocol>=2024.0.0",
+    "wafer-core",  # For compiler analyzer
+]
+authors = [
+    { name = "Wafer", email = "support@wafer.ai" }
+]
+license = { text = "MIT" }
+keywords = ["lsp", "language-server", "gpu", "cuda", "cutedsl", "c++"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: Text Editors :: Integrated Development Environments (IDE)",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project.scripts]
+wafer-lsp = "wafer_lsp.__main__:main"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/wafer_lsp"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py312"
+preview = true
+
+[tool.ruff.lint]
+select = [
+    "E",       # pycodestyle errors
+    "W",       # pycodestyle warnings
+    "F",       # pyflakes
+    "I",       # isort
+    "B",       # flake8-bugbear
+    "C4",      # flake8-comprehensions
+    "UP",      # pyupgrade
+    "ARG",     # flake8-unused-arguments
+    "SIM",     # flake8-simplify
+    "PLR",     # pylint refactor
+    "PLW",     # pylint warnings
+    "PT",      # flake8-pytest-style
+    "RUF",     # Ruff-specific rules
+]
+ignore = [
+    "E501",    # line too long (handled by formatter)
+    "B008",    # do not perform function calls in argument defaults
+    "PLR0913", # too many arguments
+    "PLR2004", # magic value used in comparison
+]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"

wafer_lsp-0.1.0/src/wafer_lsp/__init__.py

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

wafer_lsp-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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