python/utils.py

# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import os
import multiprocessing
import subprocess
import sys
import shutil
from dataclasses import dataclass, field
import setuptools.command.build_ext

@dataclass
class BuildConfig:
    cmake_only: bool = False
    build_setup: bool = True
    no_python: bool = False
    no_cutlass: bool = False
    no_test: bool = False
    no_benchmark: bool = False


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import torch

from ._C_DIRECT import DataType

import ctypes
from typing import Type, Union, Tuple
import functools

NumberTypeType = Union[Type[bool], Type[int], Type[float], Type[complex]]

_torch_dtype_to_nvfuser_dtype_map = {
    torch.cdouble: DataType.ComplexDouble,
    torch.cfloat: DataType.ComplexFloat,
    torch.double: DataType.Double,
    torch.float: DataType.Float,
    torch.half: DataType.Half,
    torch.bfloat16: DataType.BFloat16,
    torch.float8_e4m3fn: DataType.Float8_e4m3fn,


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
from typing import Any
from .version import _version_str

__all__ = ["NvfuserVersion", "Version"]

class _LazyImport:
    """Wraps around classes lazy imported from packaging.version
    Output of the function v in following snippets are identical:
       from packaging.version import Version
       def v():
           return Version('1.2.3')
    and
       Version = _LazyImport('Version')
       def v():
           return Version('1.2.3')
    The difference here is that in later example imports
    do not happen until v is called


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import sys
import traceback
import warnings
from typing import Iterable, Optional
import functools

if "nvfuser" in sys.modules:
    warnings.warn(
        "Be careful! You've imported nvfuser_direct when the nvfuser module is already imported.",
        UserWarning,
    )

import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Placement, Shard, Replicate


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch
from cupti import cupti
import cxxfilt
import pytest
from ._C_DIRECT import get_fusion_profile

# Base class for all timers used by pytest-benchmark.
class Timer:
    def __init__(self):
        self.current_time = 0.0

    def _increment_global_time(self, elapsed_time: float) -> None:
        self.current_time += elapsed_time

    def __call__(self):
        raise NotImplementedError("Subclass must implement this method")


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

def get_available_memory_gb():
    """Returns the available memory in GB."""
    try:
        import psutil

        return psutil.virtual_memory().available / 1024 / 1024 / 1024
    except:  # noqa: E722
        pass

    try:
        with open("/proc/meminfo", "r") as f:
            while True:
                line = f.readline()
                if line.startswith("MemAvailable:"):
                    mem = line.split()[1]
                    assert line.split()[2] == "kB"


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause


#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
nvFuser Dependency Report Generator

Reads dependency data from JSON (generated by CMake) and prints
a comprehensive, user-friendly report with colored output and
actionable installation instructions for missing dependencies.

IMPORTANT: CMake is the source of truth for all dependency requirements.
This script only formats output and provides help text.
"""

import json
import sys
from pathlib import Path
from typing import Dict

from prereqs import detect_platform, format_platform_info


# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import subprocess
import sys
from pathlib import Path

UNKNOWN = "Unknown"
nvfuser_root = Path(__file__).parent.parent

# note that this root currently is still part of pytorch.
def get_sha() -> str:
    try:
        return (
            subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=nvfuser_root)
            .decode("ascii")
            .strip()
        )
    except Exception:
        import os


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
Color utilities for terminal output.
"""

import os

class Colors:
    """ANSI color codes for terminal output"""

    _codes = {
        "RESET": "\033[m",
        "BOLD": "\033[1m",
        # Regular colors
        "GREEN": "\033[32m",
        "YELLOW": "\033[33m",
        "CYAN": "\033[36m",
        "WHITE": "\033[37m",


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
Custom exceptions for nvFuser prerequisite validation.

These exceptions provide structured error handling for build prerequisite
checks, enabling clear and actionable error messages.
"""

class PrerequisiteMissingError(Exception):
    """
    Raised when a prerequisite for building nvFuser is missing or has an incorrect version.

    This exception should include:
    - What prerequisite is missing or incorrect
    - Why it's required
    - Exact commands to install or fix it
    - Platform-specific guidance when applicable
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
Platform detection utilities for nvFuser build system.

Detects OS, architecture, and Linux distribution to provide platform-specific
error messages and installation guidance.
"""

import platform
from typing import Dict, Optional

def detect_platform() -> Dict[str, Optional[str]]:
    """
    Detect the current platform and return structured information.

    Returns:
        dict: Platform information with keys:
            - 'os': Operating system (Linux, Darwin, Windows, etc.)


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
Utility functions for nvFuser dependency reporting.

IMPORTANT: This module provides ONLY utility functions for formatting and URL generation.
Version requirements and dependency validation are handled by CMake.

CMake defines requirements in: cmake/DependencyRequirements.cmake
CMake exports status to: build/nvfuser_dependencies.json
Python reads JSON and uses these utilities to format help text.
"""

import platform
import re
from dataclasses import dataclass
from typing import Optional, Tuple

# =============================================================================


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
nvFuser prerequisite utilities package.

IMPORTANT DESIGN PRINCIPLE:
===========================
CMake is the source of truth for ALL dependency requirements and validation.

- CMake defines version requirements in: cmake/DependencyRequirements.cmake
- CMake finds dependencies and validates versions
- CMake exports all data to JSON: build/nvfuser_dependencies.json
- Python reads JSON and formats output with helpful instructions

This package provides ONLY:
- Platform detection utilities
- Version parsing/formatting utilities
- URL generators for downloads
- Utilities for formatting help text


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Git submodules dependency requirement."""

from typing import Dict
from .base import BooleanRequirement

class GitSubmodulesRequirement(BooleanRequirement):
    """
    Git submodules initialization check.

    CMake variables used:
    - GitSubmodules_FOUND: Whether submodules are initialized
    - NVFUSER_REQUIREMENT_GitSubmodules_STATUS: Validation status
    - NVFUSER_REQUIREMENT_GitSubmodules_OPTIONAL: Whether submodules are optional

    No version checking - simple pass/fail.
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""
Base classes for requirement types.

This module provides the foundational classes for all dependency requirements.
Each requirement knows how to format its status and determine if it represents a failure.
"""

from abc import ABC, abstractmethod
from typing import Optional, Dict
from dataclasses import dataclass

from ..colors import colorize

@dataclass
class RequirementStatus:
    """Validation status constants."""


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""pybind11 dependency requirement."""

from typing import Dict
from .base import VersionRequirement

class Pybind11Requirement(VersionRequirement):
    """
    pybind11 requirement for Python bindings.

    CMake variables used:
    - pybind11_FOUND: Whether pybind11 was found
    - pybind11_VERSION: Detected version (e.g., "3.0.1")
    - pybind11_DIR: Path to pybind11 CMake config
    - NVFUSER_REQUIREMENT_pybind11_STATUS: Validation status
    - NVFUSER_REQUIREMENT_pybind11_VERSION_MIN: Minimum required version
    - NVFUSER_REQUIREMENT_pybind11_OPTIONAL: Whether pybind11 is optional
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""CUDA Toolkit dependency requirement."""

from typing import Dict
from .base import VersionRequirement

class CUDAToolkitRequirement(VersionRequirement):
    """
    NVIDIA CUDA Toolkit requirement.

    CMake variables used:
    - CUDAToolkit_FOUND: Whether CUDA was found
    - CUDAToolkit_VERSION: Detected version (e.g., "13.1.80")
    - CUDAToolkit_ROOT: Path to CUDA installation
    - NVFUSER_REQUIREMENT_CUDAToolkit_STATUS: Validation status
    - NVFUSER_REQUIREMENT_CUDAToolkit_VERSION_MIN: Minimum required version
    - NVFUSER_REQUIREMENT_CUDAToolkit_OPTIONAL: Whether CUDA is optional
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""LLVM dependency requirement."""

from typing import Dict
from .base import VersionRequirement

class LLVMRequirement(VersionRequirement):
    """
    LLVM requirement for Host IR JIT compilation.

    CMake variables used:
    - LLVM_FOUND: Whether LLVM was found
    - LLVM_VERSION: Detected version (e.g., "18.1.3")
    - LLVM_DIR: Path to LLVM CMake config
    - NVFUSER_REQUIREMENT_LLVM_STATUS: Validation status
    - NVFUSER_REQUIREMENT_LLVM_VERSION_MIN: Minimum required version
    - NVFUSER_REQUIREMENT_LLVM_OPTIONAL: Whether LLVM is optional
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Compiler dependency requirement (GNU/Clang)."""

from typing import Dict
from .base import VersionRequirement

class CompilerRequirement(VersionRequirement):
    """
    C++ compiler requirement with name mapping.

    CMake variables used:
    - CMAKE_CXX_COMPILER_ID: Compiler name (GNU or Clang)
    - Compiler_FOUND: Whether compiler is available (always TRUE)
    - CMAKE_CXX_COMPILER_VERSION: Detected compiler version
    - CMAKE_CXX_COMPILER: Path to compiler executable
    - NVFUSER_REQUIREMENT_Compiler_STATUS: Validation status
    - NVFUSER_REQUIREMENT_Compiler_VERSION_MIN: Minimum required version (set based on compiler ID)
    - NVFUSER_REQUIREMENT_Compiler_OPTIONAL: Whether compiler is optional


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Ninja build system dependency requirement."""

from typing import Dict
from .base import BooleanRequirement

class NinjaRequirement(BooleanRequirement):
    """
    Ninja build system check.

    CMake variables used:
    - Ninja_FOUND: Whether Ninja is available
    - NVFUSER_REQUIREMENT_Ninja_STATUS: Validation status
    - NVFUSER_REQUIREMENT_Ninja_OPTIONAL: Whether Ninja is optional

    No version checking - just verifies Ninja is available.
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Requirement class registry and factory."""

from .base import Requirement, VersionRequirement, BooleanRequirement, RequirementStatus
from .python import PythonRequirement
from .torch import TorchRequirement
from .llvm import LLVMRequirement
from .cuda_toolkit import CUDAToolkitRequirement
from .pybind11 import Pybind11Requirement
from .compiler import CompilerRequirement
from .git_submodules import GitSubmodulesRequirement
from .ninja import NinjaRequirement
from .nvmmh import NVMMHRequirement

__all__ = [
    # Base classes
    "Requirement",
    "VersionRequirement",
    "BooleanRequirement",


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""PyTorch dependency requirement with CUDA constraint validation."""

from typing import Dict
from .base import VersionRequirement
from ..colors import colorize

class TorchRequirement(VersionRequirement):
    """
    PyTorch requirement with CUDA version constraint checking.

    CMake variables used:
    - Torch_FOUND: Whether Torch was found
    - Torch_VERSION: Detected PyTorch version
    - Torch_DIR: Path to Torch CMake config
    - NVFUSER_REQUIREMENT_Torch_STATUS: Validation status
    - NVFUSER_REQUIREMENT_Torch_VERSION_MIN: Minimum required version
    - NVFUSER_REQUIREMENT_Torch_OPTIONAL: Whether Torch is optional


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""nvidia-matmul-heuristics dependency requirement."""

from typing import Dict
from .base import BooleanRequirement

class NVMMHRequirement(BooleanRequirement):
    """
    nvidia-matmul-heuristics check.

    CMake variables used:
    - NVMMH_FOUND: Whether nvidia-matmul-heuristics is available
    - NVFUSER_REQUIREMENT_NVMMH_STATUS: Validation status
    - NVFUSER_REQUIREMENT_NVMMH_OPTIONAL: Whether NVMMH is optional

    No version checking - just verifies nvidia-matmul-heuristics headers are available.
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Python dependency requirement."""

from typing import Dict
from .base import VersionRequirement

class PythonRequirement(VersionRequirement):
    """
    Python interpreter requirement.

    CMake variables used:
    - Python_FOUND: Whether Python was found
    - Python_VERSION: Detected version (e.g., "3.12.3")
    - Python_EXECUTABLE: Path to python binary
    - NVFUSER_REQUIREMENT_Python_STATUS: Validation status
    - NVFUSER_REQUIREMENT_Python_VERSION_MIN: Minimum required version
    - NVFUSER_REQUIREMENT_Python_OPTIONAL: Whether Python is optional
    """


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import os

__all__ = [
    "cmake_prefix_path",
]

cmake_prefix_path = os.path.join(
    os.path.dirname(os.path.dirname(__file__)),
    "nvfuser_common",
    "share",
    "cmake",
    "nvfuser",
)