wafer-cli 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

wafer/corpus.py

@@ -15,7 +15,16 @@ import httpx
 
 CACHE_DIR = Path.home() / ".cache" / "wafer" / "corpora"
 
-CorpusName = Literal["cuda", "cutlass", "hip"]
+CorpusName = Literal["cuda", "cutlass", "hip", "amd"]
+
+
+@dataclass
+class RepoSource:
+    """A single GitHub repo source within a corpus."""
+
+    repo: str
+    paths: list[str]
+    branch: str = "main"
 
 
 @dataclass
@@ -24,10 +33,11 @@ class CorpusConfig:
 
     name: CorpusName
     description: str
-    source_type: Literal["nvidia_md", "github_repo"]
+    source_type: Literal["nvidia_md", "github_repo", "github_multi_repo"]
     urls: list[str] | None = None
     repo: str | None = None
     repo_paths: list[str] | None = None
+    repos: list[RepoSource] | None = None  # For multi-repo corpora
 
 
 CORPORA: dict[CorpusName, CorpusConfig] = {
@@ -69,6 +79,74 @@ CORPORA: dict[CorpusName, CorpusConfig] = {
         repo="ROCm/HIP",
         repo_paths=["docs"],
     ),
+    "amd": CorpusConfig(
+        name="amd",
+        description="AMD GPU kernel development (rocWMMA, CK, AITER, rocBLAS, HipKittens, vLLM)",
+        source_type="github_multi_repo",
+        repos=[
+            # rocWMMA - wave matrix multiply-accumulate (WMMA) intrinsics
+            RepoSource(
+                repo="ROCm/rocWMMA",
+                paths=["docs", "samples", "library/include"],
+                branch="develop",
+            ),
+            # Composable Kernel - tile-based GPU programming
+            RepoSource(
+                repo="ROCm/composable_kernel",
+                paths=["docs", "example", "tutorial", "include/ck_tile"],
+                branch="develop",
+            ),
+            # AITER - AMD inference tensor runtime
+            RepoSource(
+                repo="ROCm/aiter",
+                paths=["docs", "aiter/ops"],
+            ),
+            # MIOpen - deep learning primitives (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/MIOpen",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # rocBLAS - BLAS library (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/rocBLAS",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # hipBLASLt - lightweight BLAS (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/hipBLASLt",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # Tensile - GEMM code generator (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/Tensile",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # HipKittens - high-performance AMD kernels
+            RepoSource(
+                repo="HazyResearch/HipKittens",
+                paths=["docs", "kernels", "include"],
+            ),
+            # vLLM AMD kernels
+            RepoSource(
+                repo="vllm-project/vllm",
+                paths=["csrc/rocm"],
+            ),
+            # SGLang AMD kernels
+            RepoSource(
+                repo="sgl-project/sglang",
+                paths=["3rdparty/amd"],
+            ),
+            # HuggingFace ROCm kernels
+            RepoSource(
+                repo="huggingface/hf-rocm-kernels",
+                paths=["csrc", "hf_rocm_kernels", "docs"],
+            ),
+        ],
+    ),
 }
 
 
@@ -113,41 +191,87 @@ def _download_nvidia_md(config: CorpusConfig, dest: Path, verbose: bool = True)
     return downloaded
 
 
+def _extract_matching_files(
+    tar: tarfile.TarFile,
+    repo_paths: list[str],
+    dest: Path,
+    verbose: bool,
+) -> int:
+    """Extract files matching repo_paths from tarball."""
+    downloaded = 0
+    for member in tar.getmembers():
+        if not member.isfile():
+            continue
+        rel_path = "/".join(member.name.split("/")[1:])
+        if not any(rel_path.startswith(rp) for rp in repo_paths):
+            continue
+        target = dest / rel_path
+        target.parent.mkdir(parents=True, exist_ok=True)
+        src = tar.extractfile(member)
+        if src:
+            target.write_bytes(src.read())
+            downloaded += 1
+            if verbose:
+                print(f"  ✓ {rel_path}")
+    return downloaded
+
+
+def _download_single_github_repo(
+    client: httpx.Client,
+    repo: str,
+    repo_paths: list[str],
+    dest: Path,
+    branch: str = "main",
+    verbose: bool = True,
+) -> int:
+    """Download specific paths from a single GitHub repo."""
+    tarball_url = f"https://api.github.com/repos/{repo}/tarball/{branch}"
+    if verbose:
+        print(f"  Fetching {repo}...")
+    resp = client.get(tarball_url)
+    resp.raise_for_status()
+    with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
+        tmp.write(resp.content)
+        tmp_path = Path(tmp.name)
+    try:
+        with tarfile.open(tmp_path, "r:gz") as tar:
+            return _extract_matching_files(tar, repo_paths, dest, verbose)
+    finally:
+        tmp_path.unlink()
+
+
 def _download_github_repo(config: CorpusConfig, dest: Path, verbose: bool = True) -> int:
     """Download specific paths from GitHub repo."""
     assert config.repo is not None
     assert config.repo_paths is not None
-    downloaded = 0
     with httpx.Client(timeout=60.0, follow_redirects=True) as client:
-        tarball_url = f"https://api.github.com/repos/{config.repo}/tarball/main"
-        if verbose:
-            print(f"  Fetching {config.repo}...")
-        resp = client.get(tarball_url)
-        resp.raise_for_status()
-        with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
-            tmp.write(resp.content)
-            tmp_path = Path(tmp.name)
-        try:
-            with tarfile.open(tmp_path, "r:gz") as tar:
-                members = tar.getmembers()
-                root_prefix = members[0].name.split("/")[0] if members else ""
-                for member in members:
-                    if not member.isfile():
-                        continue
-                    rel_path = "/".join(member.name.split("/")[1:])
-                    for repo_path in config.repo_paths:
-                        if rel_path.startswith(repo_path):
-                            target = dest / rel_path
-                            target.parent.mkdir(parents=True, exist_ok=True)
-                            with tar.extractfile(member) as src:
-                                if src:
-                                    target.write_bytes(src.read())
-                                    downloaded += 1
-                                    if verbose:
-                                        print(f"  ✓ {rel_path}")
-                            break
-        finally:
-            tmp_path.unlink()
+        return _download_single_github_repo(
+            client, config.repo, config.repo_paths, dest, verbose=verbose
+        )
+
+
+def _download_github_multi_repo(config: CorpusConfig, dest: Path, verbose: bool = True) -> int:
+    """Download specific paths from multiple GitHub repos."""
+    assert config.repos is not None
+    downloaded = 0
+    with httpx.Client(timeout=120.0, follow_redirects=True) as client:
+        for repo_source in config.repos:
+            repo_name = repo_source.repo.split("/")[-1]
+            repo_dest = dest / repo_name
+            repo_dest.mkdir(parents=True, exist_ok=True)
+            try:
+                count = _download_single_github_repo(
+                    client,
+                    repo_source.repo,
+                    repo_source.paths,
+                    repo_dest,
+                    branch=repo_source.branch,
+                    verbose=verbose,
+                )
+                downloaded += count
+            except httpx.HTTPError as e:
+                if verbose:
+                    print(f"  ✗ {repo_source.repo}: {e}")
     return downloaded
 
 
@@ -185,6 +309,8 @@ def download_corpus(name: CorpusName, force: bool = False, verbose: bool = True)
         count = _download_nvidia_md(config, dest, verbose)
     elif config.source_type == "github_repo":
         count = _download_github_repo(config, dest, verbose)
+    elif config.source_type == "github_multi_repo":
+        count = _download_github_multi_repo(config, dest, verbose)
     else:
         raise ValueError(f"Unknown source type: {config.source_type}")
     if verbose: