diff --git a/.ci/docker/ci_commit_pins/rocm-composable-kernel.txt b/.ci/docker/ci_commit_pins/rocm-composable-kernel.txt deleted file mode 100644 index c45f46af95d..00000000000 --- a/.ci/docker/ci_commit_pins/rocm-composable-kernel.txt +++ /dev/null @@ -1 +0,0 @@ -7fe50dc3da2069d6645d9deb8c017a876472a977 diff --git a/setup.py b/setup.py index 4b4704a9ca5..c0523a1b5c6 100644 --- a/setup.py +++ b/setup.py @@ -324,7 +324,6 @@ from tools.setup_helpers.env import ( IS_WINDOWS, ) from tools.setup_helpers.generate_linker_script import gen_linker_script -from tools.setup_helpers.rocm_env import get_ck_dependency_string, IS_ROCM def str2bool(value: str | None) -> bool: @@ -507,6 +506,7 @@ else: sysconfig.get_config_var("LIBDIR") ) / sysconfig.get_config_var("INSTSONAME") + ################################################################################ # Version, create_version_file, and package_name ################################################################################ @@ -1494,12 +1494,6 @@ def configure_extension_build() -> tuple[ map(str.strip, pytorch_extra_install_requires.split("|")) ) - # Adding extra requirements for ROCm builds - if IS_ROCM and platform.system() == "Linux": - extra_install_requires.append( - f"rocm-composable-kernel {get_ck_dependency_string()}" - ) - # Cross-compile for M1 if IS_DARWIN: macos_target_arch = os.getenv("CMAKE_OSX_ARCHITECTURES", "") diff --git a/test/inductor/test_ck_backend.py b/test/inductor/test_ck_backend.py index 079be79fcc9..f73a47e45a5 100644 --- a/test/inductor/test_ck_backend.py +++ b/test/inductor/test_ck_backend.py @@ -1,4 +1,5 @@ # Owner(s): ["module: inductor"] +import functools import logging import os import unittest @@ -12,7 +13,6 @@ except ImportError: import torch from torch._inductor import config from torch._inductor.test_case import run_tests, TestCase -from torch._inductor.utils import try_import_ck_lib from torch.testing._internal.common_cuda import tf32_off from torch.testing._internal.common_utils import ( instantiate_parametrized_tests, @@ -32,8 +32,20 @@ if HAS_CUDA_AND_TRITON: log = logging.getLogger(__name__) -# patch env for tests if needed -_test_env = {} +@functools.lru_cache(None) +def _get_path_without_sccache() -> str: + """ + Get the PATH environment variable without sccache. + """ + path_envs = os.environ.get("PATH", "").split(":") + path_envs = [env for env in path_envs if "/opt/cache/bin" not in env] + return ":".join(path_envs) + + +_test_env = { + "PATH": _get_path_without_sccache(), + "DISABLE_SCCACHE": "1", +} @instantiate_parametrized_tests @@ -49,10 +61,13 @@ class TestCKBackend(TestCase): ) torch.random.manual_seed(1234) + try: + import ck4inductor # @manual - self.ck_dir, _, _, _ = try_import_ck_lib() - if not self.ck_dir: - raise unittest.SkipTest("Composable Kernel library is not installed") + self.ck_dir = os.path.dirname(ck4inductor.__file__) + os.environ["TORCHINDUCTOR_CK_DIR"] = self.ck_dir + except ImportError as e: + raise unittest.SkipTest("Composable Kernel library not installed") from e try: os.environ["INDUCTOR_TEST_DISABLE_FRESH_CACHE"] = "1" @@ -273,9 +288,6 @@ class TestCKBackend(TestCase): torch.testing.assert_close(Y_compiled, Y_eager) - @unittest.skip( - "FIXME(tenpercent): kernel compilation errors on gfx942 as of 09/01/25" - ) @unittest.skipIf(not torch.version.hip, "ROCM only") @unittest.mock.patch.dict(os.environ, _test_env) @parametrize("max_autotune_gemm_backends", ("CK", "ATen,Triton,CK")) diff --git a/tools/setup_helpers/rocm_env.py b/tools/setup_helpers/rocm_env.py deleted file mode 100644 index 4bbc9667ca9..00000000000 --- a/tools/setup_helpers/rocm_env.py +++ /dev/null @@ -1,61 +0,0 @@ -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+@#egg=rocm-composable-kernel" - where: - - is the URL for ROCm Composable Kernel - - 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}" diff --git a/torch/_inductor/codegen/rocm/compile_command.py b/torch/_inductor/codegen/rocm/compile_command.py index aa935b14af2..b9cae55102b 100644 --- a/torch/_inductor/codegen/rocm/compile_command.py +++ b/torch/_inductor/codegen/rocm/compile_command.py @@ -4,7 +4,7 @@ import os from typing import Optional from torch._inductor import config -from torch._inductor.utils import is_linux, try_import_ck_lib +from torch._inductor.utils import is_linux log = logging.getLogger(__name__) @@ -18,23 +18,18 @@ def _rocm_include_paths(dst_file_ext: str) -> list[str]: if config.rocm.rocm_home else cpp_extension._join_rocm_home("include") ) + if not config.rocm.ck_dir: + log.warning("Unspecified Composable Kernel include dir") if config.is_fbcode(): from libfb.py import parutil ck_path = parutil.get_dir_path("composable-kernel-headers") else: - if not config.rocm.ck_dir: - ck_dir, _, _, _ = try_import_ck_lib() - if not ck_dir: - log.warning("Unspecified Composable Kernel directory") - config.rocm.ck_dir = ck_dir ck_path = config.rocm.ck_dir or cpp_extension._join_rocm_home( "composable_kernel" ) - log.debug("Using ck path %s", ck_path) - ck_include = os.path.join(ck_path, "include") ck_library_include = os.path.join(ck_path, "library", "include") diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py index 46e9bc9407a..69e6111817b 100644 --- a/torch/_inductor/utils.py +++ b/torch/_inductor/utils.py @@ -1979,7 +1979,16 @@ def use_ck_template(layout: Layout) -> bool: log.warning("Please pip install Composable Kernel package") return False - config.rocm.ck_dir = ck_package_dirname + if config.is_fbcode(): + config.rocm.ck_dir = ck_package_dirname + + if not config.rocm.ck_dir: + log.warning("Please set TORCHINDUCTOR_CK_DIR env variable") + return False + + if ck_package_dirname != config.rocm.ck_dir: + log.warning("Invalid path to CK library") + return False return True