wafer-lsp 0.1.7__tar.gz

wafer_lsp-0.1.7/.gitignore

@@ -0,0 +1,118 @@
+# 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/
+research/evals/retrieval_eval_project/results/
+research/evals/trace_analyze_eval/results/
+research/evals/**/uv.lock
+
+# Large trace files
+**/pmc_perf.csv
+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
+
+# ROCprofiler artifacts
+.rocprofv3/
+
+# 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
+research/evals/ask_domain_eval/results/
+research/evals/optimize_performance_eval/results/
+research/evals/**/eval_output*.log
+
+# Wafer extension trace files
+.wafer/traces/
+
+# Test coverage
+coverage/
+**/coverage/
+.coverage
+*.coverage
+coverage.xml
+*.lcov
+htmlcov/
+experiments/steve/amd-model-practice/workloads/**/*.csv
+experiments/steve/amd-model-practice/liquid_profile.json
+
+# Local files
+logs/
+todo.md
+
+# Claude Code context/handoffs
+.claude/context/
+.claude/handoffs/
+KernelBench/
+HIP-Benchmarks-Results/

wafer_lsp-0.1.7/PKG-INFO

@@ -0,0 +1,60 @@
+Metadata-Version: 2.4
+Name: wafer-lsp
+Version: 0.1.7
+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
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
+Requires-Dist: pytest>=8.0.0; extra == 'dev'
+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.7/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.7/pyproject.toml

@@ -0,0 +1,73 @@
+[project]
+name = "wafer-lsp"
+version = "0.1.7"
+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 core analysis utilities
+]
+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)",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
+]
+
+[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.7/pyproject.toml.bak

@@ -0,0 +1,73 @@
+[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 core analysis utilities
+]
+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)",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
+]
+
+[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.7/src/wafer_lsp/__init__.py

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

wafer_lsp-0.1.7/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.7/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.7/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.7/src/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-0.1.7/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.7/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.7/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.7/src/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