wafer-lsp 0.1.13__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