mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
019fed39aa
This is a part of our effort for integrating Composable Kernel library for Inductor backend. Currently we have a submodule, but would prefer to have commit pin control over the library as with Triton. We intentionally avoid putting all installation logic in CI scripts to allow locally built versions to have this functionality. The idea is to have CK as a pytorch dependency in pytorch 2.9 release to allow people to use it with inductor and AOT inductor and then gradually step away from submodule usage. Right now CK usage in SDPA/Gemm is tied to submodule files. This PR is a remake of due to branch error: https://github.com/pytorch/pytorch/pull/156192 Pull Request resolved: https://github.com/pytorch/pytorch/pull/158747 Approved by: https://github.com/jeffdaily Co-authored-by: Jithun Nair <37884920+jithunnair-amd@users.noreply.github.com> Co-authored-by: Jack Taylor <108682042+jataylo@users.noreply.github.com> Co-authored-by: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
import os
|
|
from pathlib import Path
|
|
|
|
|
|
def check_if_rocm() -> bool:
|
|
# If user defines USE_ROCM during PyTorch build, respect their intention
|
|
use_rocm_env = os.environ.get("USE_ROCM")
|
|
if use_rocm_env:
|
|
return bool(use_rocm_env)
|
|
# otherwise infer existence of ROCm installation as indication of ROCm build
|
|
rocm_path_env = os.environ.get("ROCM_PATH", "/opt/rocm")
|
|
if rocm_path_env and os.path.exists(rocm_path_env):
|
|
return True
|
|
return False
|
|
|
|
|
|
IS_ROCM = check_if_rocm()
|
|
|
|
SCRIPT_DIR = Path(__file__).parent
|
|
REPO_DIR = SCRIPT_DIR.parent.parent
|
|
|
|
|
|
# CK pin is read in a similar way that triton commit is
|
|
def read_ck_pin() -> str:
|
|
"""
|
|
Reads the CK (Composable Kernel) commit hash.
|
|
The hash is pinned to a known stable version of CK.
|
|
|
|
Returns:
|
|
str: The commit hash read from 'rocm-composable-kernel.txt'.
|
|
"""
|
|
ck_file = "rocm-composable-kernel.txt"
|
|
with open(REPO_DIR / ".ci" / "docker" / "ci_commit_pins" / ck_file) as f:
|
|
return f.read().strip()
|
|
|
|
|
|
# Prepares a dependency string for install_requires in setuptools
|
|
# in specific PEP 508 URL format
|
|
def get_ck_dependency_string() -> str:
|
|
"""
|
|
Generates a PEP 508-compliant dependency string for the ROCm Composable Kernel
|
|
to be used in setuptools' install_requires.
|
|
|
|
The returned string is EITHER in the format:
|
|
" @ git+<repo_url>@<commit_hash>#egg=rocm-composable-kernel"
|
|
where:
|
|
- <repo_url> is the URL for ROCm Composable Kernel
|
|
- <commit_hash> is read from the commit pin file
|
|
- "#egg=rocm-composable-kernel" specifies the package name for setuptools
|
|
OR an empty string, making use of the existing rocm-composable-kernel installation.
|
|
|
|
Returns:
|
|
str: The formatted dependency string for use in install_requires.
|
|
"""
|
|
egg_name = "#egg=rocm-composable-kernel"
|
|
commit_pin = f"@{read_ck_pin()}"
|
|
if os.getenv("TORCHINDUCTOR_CK_DIR"):
|
|
# we take non-empty env as an indicator that the package has already been installed and doesn't need to be re-installed
|
|
# this comes with a caveat that the pinned version is known to work while the preinstalled version might not
|
|
return ""
|
|
return f"@ git+https://github.com/ROCm/composable_kernel.git{commit_pin}{egg_name}"
|