nvfuser-cu121-torch25 0.2.25.dev20250201__cp312-cp312-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/validator_utils.h

@@ -0,0 +1,92 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <array>
+#include <unordered_map>
+#include <utility>
+
+#include <ATen/ArrayRef.h>
+
+#include <expr_evaluator.h>
+#include <fusion.h>
+#include <ir/interface_nodes.h>
+#include <iter_visitor.h>
+#include <runtime/executor_params.h>
+#include <type.h>
+
+namespace nvfuser {
+
+struct ValidationConstants {
+  // Tolerances generated from randn + add + sum fusion
+  // compared against double precision
+  std::array<std::array<double, 2>, 20> sum_tolerances_float = {
+      {{4, 1.68222e-06},      {8, 2.23704e-06},      {16, 2.95788e-06},
+       {32, 4.4778e-06},      {64, 6.75395e-06},     {128, 8.57934e-06},
+       {256, 1.30594e-05},    {512, 2.19122e-05},    {1024, 3.3451e-05},
+       {2048, 5.78476e-05},   {4096, 0.000108292},   {8192, 0.00012207},
+       {16384, 0.000136882},  {32768, 0.000248561},  {65536, 0.000407594},
+       {131072, 0.000500901}, {262144, 0.000923019}, {524288, 0.00156909},
+       {1048576, 0.00223107}, {2097152, 0.00343043}}};
+
+  // Tolerances generated from randn + add + sum fusion
+  // compared against double precision
+  std::array<std::array<double, 2>, 20> sum_tolerances_half = {
+      {{4, 0.00390625},    {8, 0.0078125},    {16, 0.0078125},
+       {32, 0.0155334},    {64, 0.0156269},   {128, 0.0312042},
+       {256, 0.0312548},   {512, 0.0619979},  {1024, 0.0625103},
+       {2048, 0.124686},   {4096, 0.12501},   {8192, 0.24945},
+       {16384, 0.250049},  {32768, 0.498946}, {65536, 0.500071},
+       {131072, 0.985087}, {262144, 1.00006}, {524288, 1.99234},
+       {1048576, 2.00032}, {2097152, 3.99073}}};
+
+  double base_half_abs_tol = -1;
+  double base_half_rel_tol = -1;
+  double base_float_abs_tol = -1;
+  double base_float_rel_tol = -1;
+};
+
+// Returns abs and relative values to use for validation.
+std::pair<double, double> getTolerance(
+    DataType dtype,
+    int64_t reduction_size,
+    const ValidationConstants& tolerances);
+
+class ReductionSizeMapper : private IterVisitor {
+ public:
+  //! Runs through the fusion and determines how many reductions were performed
+  //! to compute each tensorview.
+  static std::unordered_map<TensorView*, int64_t> computeReductionSizes(
+      Fusion* fusion,
+      ExpressionEvaluator& expr_eval);
+
+ private:
+  ReductionSizeMapper(Fusion* fusion, ExpressionEvaluator& expr_eval);
+
+  int64_t getReductionSize(const TensorView* tv);
+
+  void dispatch(Expr* expr) override;
+
+  using IterVisitor::handle;
+
+  std::unordered_map<TensorView*, int64_t> reduction_map;
+  ExpressionEvaluator& expr_eval_;
+};
+
+ExpressionEvaluator bindInputsAndLaunchParams(
+    Fusion* fusion,
+    const at::ArrayRef<c10::IValue>& aten_inputs,
+    const LaunchParams& launch_constraints);
+
+std::vector<std::pair<double, double>> get_val_constants(
+    Fusion* fusion,
+    const at::ArrayRef<c10::IValue>& aten_inputs,
+    const LaunchParams& lparams = LaunchParams(),
+    const ValidationConstants& tolerances = ValidationConstants());
+
+} // namespace nvfuser

nvfuser/include/nvfuser/vectorization_info.h

@@ -0,0 +1,31 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <ir/all_nodes.h>
+
+namespace nvfuser {
+
+struct VectorizedSetInfo {
+  //! Producer of a vectorized set
+  TensorView* producer_tv = nullptr;
+  //! Consumer of a vectorized set
+  TensorView* consumer_tv = nullptr;
+  //! Number of elements to vectorize
+  int64_t word_size = -1;
+  //! Vectorized domain
+  IterDomain* vectorized_loop_id = nullptr;
+  //! Right-most allocation dependent domain of the loop domain for consumer
+  IterDomain* vectorized_consumer_alloc_id = nullptr;
+  //! Right-most allocation dependent domain of the loop domain for producer
+  IterDomain* vectorized_producer_alloc_id = nullptr;
+  //! All of the dependent allocation domains that are contiguously merged
+  std::unordered_set<IterDomain*> contig_alloc_ids;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/visibility.h

@@ -0,0 +1,21 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+/// Defines the NVF_API macro, which should be added on methods or classes
+/// that are used outside of nvFuser. See doc/dev/visibility.md for details.
+
+#if defined _WIN32 || defined __CYGWIN__
+#ifdef BUILDING_DLL
+#define NVF_API __declspec(dllexport)
+#else
+#define NVF_API __declspec(dllimport)
+#endif
+#else
+#define NVF_API __attribute__((visibility("default")))
+#endif

nvfuser/nvfuser_version.py

@@ -0,0 +1,69 @@
+# 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
+    """
+
+    def __init__(self, cls_name: str) -> None:
+        self._cls_name = cls_name
+
+    def get_cls(self):
+        try:
+            import packaging.version  # type: ignore[import]
+        except ImportError:
+            # If packaging isn't installed, try and use the vendored copy
+            # in pkg_resources
+            from pkg_resources import packaging  # type: ignore[attr-defined, no-redef]
+        return getattr(packaging.version, self._cls_name)
+
+    def __call__(self, *args, **kwargs):
+        return self.get_cls()(*args, **kwargs)
+
+    def __instancecheck__(self, obj):
+        return isinstance(obj, self.get_cls())
+
+
+Version = _LazyImport("Version")
+
+
+class NvfuserVersion(str):
+    @classmethod
+    def _convert_to_version(cls, ver: Any) -> Version:
+        if isinstance(ver, str):
+            return Version(ver.split("+")[0])
+        elif isinstance(ver, Version.get_cls()):
+            return ver
+        else:
+            raise ValueError("can't convert {} to Version".format(ver))
+
+    def _cmp_version(self, other: Any, method: str) -> Version:
+        return getattr(NvfuserVersion._convert_to_version(self), method)(
+            NvfuserVersion._convert_to_version(other)
+        )
+
+
+for cmp_method in ["__gt__", "__lt__", "__eq__", "__ge__", "__le__"]:
+    setattr(
+        NvfuserVersion,
+        cmp_method,
+        lambda x, y, method=cmp_method: x._cmp_version(y, method),
+    )
+
+__version__ = NvfuserVersion(_version_str)

nvfuser/pytorch_utils.py

@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+import torch
+
+from ._C import DataType
+
+import ctypes
+import functools
+import gc
+from typing import Type, Union, Tuple
+
+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,
+    torch.float8_e5m2: DataType.Float8_e5m2,
+    torch.long: DataType.Int,
+    torch.int: DataType.Int32,
+    torch.bool: DataType.Bool,
+    # Python scalars
+    complex: DataType.ComplexDouble,
+    float: DataType.Double,
+    int: DataType.Int,
+    bool: DataType.Bool,
+}
+
+
+def python_scalar_to_nvfuser_dtype(a: Union[int, float, complex, bool]):
+    return _torch_dtype_to_nvfuser_dtype_map[type(a)]
+
+
+def torch_dtype_to_nvfuser_dtype(dtype: Union[torch.dtype, NumberTypeType]):
+    """
+    Translates from torch.dtype to nvFuser's DataType enum
+    """
+    return _torch_dtype_to_nvfuser_dtype_map[dtype]
+
+
+def get_device_properties() -> Tuple[int, float]:
+    """
+    Computes device properties using ctypes and cuda.
+    Note: Consider using CUDA-Python when CUDA support >= 12.0.
+    """
+    libnames = ("libcuda.so", "libcuda.dylib", "nvcuda.dll", "cuda.dll")
+    for libname in libnames:
+        try:
+            cuda = ctypes.CDLL(libname)
+        except OSError:
+            continue
+        else:
+            break
+    else:
+        raise OSError("could not load any of: " + " ".join(libnames))
+
+    # Device attribute enums (taken from cuda.h)
+    # https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TYPES.html#group__CUDA__TYPES_1ge12b8a782bebe21b1ac0091bf9f4e2a3
+
+    CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1
+    CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8
+    CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK = 12
+    CU_DEVICE_ATTRIBUTE_CLOCK_RATE = 13
+    CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE = 36
+    CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH = 37
+    CU_DEVICE_ATTRIBUTE_L2_CACHE_SIZE = 38
+    CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR = 39
+
+    device_properties = {}
+    device = torch.cuda.current_device()
+    cuda_properties = torch.cuda.get_device_properties(device)
+
+    device_properties["gpu_name"] = cuda_properties.name
+    device_properties["gpu_compute_capability_major"] = cuda_properties.major
+    device_properties["gpu_compute_capability_minor"] = cuda_properties.minor
+    device_properties["gpu_gmem_bytes"] = cuda_properties.total_memory
+    device_properties["gpu_sm_count"] = cuda_properties.multi_processor_count
+
+    max_threads_per_block = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(max_threads_per_block),
+        CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK,
+        device,
+    )
+    device_properties["gpu_max_threads_per_block"] = max_threads_per_block.value
+
+    smem_per_block = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(smem_per_block),
+        CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK,
+        device,
+    )
+    device_properties["gpu_smem_bytes_per_block"] = smem_per_block.value
+
+    max_reg_per_block = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(max_reg_per_block),
+        CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK,
+        device,
+    )
+    device_properties["gpu_regs_per_block"] = max_reg_per_block.value
+
+    max_clock_khz = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(max_clock_khz),
+        CU_DEVICE_ATTRIBUTE_CLOCK_RATE,
+        device,
+    )
+    device_properties["gpu_clock_rate_khz"] = max_clock_khz.value
+
+    l2_cache_size = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(l2_cache_size), CU_DEVICE_ATTRIBUTE_L2_CACHE_SIZE, device
+    )
+    device_properties["gpu_l2_bytes"] = l2_cache_size.value
+
+    memory_clock_rate = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(memory_clock_rate), CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE, device
+    )
+    device_properties["gpu_mem_clock_khz"] = memory_clock_rate.value
+
+    memory_bus_width = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(memory_bus_width),
+        CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH,
+        device,
+    )
+    device_properties["gpu_mem_bus_width"] = memory_bus_width.value
+
+    max_threads_per_sm = ctypes.c_int()
+    cuda.cuDeviceGetAttribute(
+        ctypes.byref(max_threads_per_sm),
+        CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR,
+        device,
+    )
+    device_properties["gpu_max_threads_per_sm"] = max_threads_per_sm.value
+
+    # Compute peak bandwidth in GBps
+    peak_bandwidth = (2 * memory_bus_width.value * memory_clock_rate.value) / (1e6 * 8)
+    device_properties["gpu_peak_bandwidth_gbps"] = peak_bandwidth
+
+    return device_properties
+
+
+DEVICE_PROPERTIES = None
+if torch.cuda.is_available():
+    # Loading libraries will raise errors on non-CUDA machines.
+    DEVICE_PROPERTIES = get_device_properties()
+
+
+def retry_on_oom_or_skip_test(func):
+    """Decorator: upon torch.OutOfMemoryError clear the cache and retry test"""
+
+    @functools.wraps(func)
+    def retried_func(*args, **kwargs):
+        try:
+            output = func(*args, **kwargs)
+        except torch.OutOfMemoryError:
+            pass
+        else:
+            return output
+
+        # We have hit an OOM error, so clear the cache and retry
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        try:
+            output = func(*args, **kwargs)
+        except torch.OutOfMemoryError as e:
+            # If we hit an OOM this time, then skip the test
+            import pytest
+
+            pytest.skip(f"Test failed due to OutOfMemoryError: {e}")
+            return
+
+        return output
+
+    return retried_func

nvfuser/share/cmake/nvfuser/NvfuserConfig-release.cmake

@@ -0,0 +1,20 @@
+#----------------------------------------------------------------
+# Generated CMake target import file for configuration "Release".
+#----------------------------------------------------------------
+
+# Commands may need to know the format version.
+set(CMAKE_IMPORT_FILE_VERSION 1)
+
+# Import target "nvfuser_codegen" for configuration "Release"
+set_property(TARGET nvfuser_codegen APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE)
+set_target_properties(nvfuser_codegen PROPERTIES
+  IMPORTED_LINK_DEPENDENT_LIBRARIES_RELEASE "torch"
+  IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib/libnvfuser_codegen.so"
+  IMPORTED_SONAME_RELEASE "libnvfuser_codegen.so"
+  )
+
+list(APPEND _cmake_import_check_targets nvfuser_codegen )
+list(APPEND _cmake_import_check_files_for_nvfuser_codegen "${_IMPORT_PREFIX}/lib/libnvfuser_codegen.so" )
+
+# Commands beyond this point should not need to know the version.
+set(CMAKE_IMPORT_FILE_VERSION)

nvfuser/share/cmake/nvfuser/NvfuserConfig.cmake

@@ -0,0 +1,106 @@
+# Generated by CMake
+
+if("${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION}" LESS 2.8)
+   message(FATAL_ERROR "CMake >= 2.8.3 required")
+endif()
+if(CMAKE_VERSION VERSION_LESS "2.8.3")
+   message(FATAL_ERROR "CMake >= 2.8.3 required")
+endif()
+cmake_policy(PUSH)
+cmake_policy(VERSION 2.8.3...3.29)
+#----------------------------------------------------------------
+# Generated CMake target import file.
+#----------------------------------------------------------------
+
+# Commands may need to know the format version.
+set(CMAKE_IMPORT_FILE_VERSION 1)
+
+# Protect against multiple inclusion, which would fail when already imported targets are added once more.
+set(_cmake_targets_defined "")
+set(_cmake_targets_not_defined "")
+set(_cmake_expected_targets "")
+foreach(_cmake_expected_target IN ITEMS nvfuser_codegen)
+  list(APPEND _cmake_expected_targets "${_cmake_expected_target}")
+  if(TARGET "${_cmake_expected_target}")
+    list(APPEND _cmake_targets_defined "${_cmake_expected_target}")
+  else()
+    list(APPEND _cmake_targets_not_defined "${_cmake_expected_target}")
+  endif()
+endforeach()
+unset(_cmake_expected_target)
+if(_cmake_targets_defined STREQUAL _cmake_expected_targets)
+  unset(_cmake_targets_defined)
+  unset(_cmake_targets_not_defined)
+  unset(_cmake_expected_targets)
+  unset(CMAKE_IMPORT_FILE_VERSION)
+  cmake_policy(POP)
+  return()
+endif()
+if(NOT _cmake_targets_defined STREQUAL "")
+  string(REPLACE ";" ", " _cmake_targets_defined_text "${_cmake_targets_defined}")
+  string(REPLACE ";" ", " _cmake_targets_not_defined_text "${_cmake_targets_not_defined}")
+  message(FATAL_ERROR "Some (but not all) targets in this export set were already defined.\nTargets Defined: ${_cmake_targets_defined_text}\nTargets not yet defined: ${_cmake_targets_not_defined_text}\n")
+endif()
+unset(_cmake_targets_defined)
+unset(_cmake_targets_not_defined)
+unset(_cmake_expected_targets)
+
+
+# Compute the installation prefix relative to this file.
+get_filename_component(_IMPORT_PREFIX "${CMAKE_CURRENT_LIST_FILE}" PATH)
+get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH)
+get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH)
+get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH)
+if(_IMPORT_PREFIX STREQUAL "/")
+  set(_IMPORT_PREFIX "")
+endif()
+
+# Create imported target nvfuser_codegen
+add_library(nvfuser_codegen SHARED IMPORTED)
+
+set_target_properties(nvfuser_codegen PROPERTIES
+  INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include/nvfuser"
+)
+
+# Load information for each installed configuration.
+file(GLOB _cmake_config_files "${CMAKE_CURRENT_LIST_DIR}/NvfuserConfig-*.cmake")
+foreach(_cmake_config_file IN LISTS _cmake_config_files)
+  include("${_cmake_config_file}")
+endforeach()
+unset(_cmake_config_file)
+unset(_cmake_config_files)
+
+# Cleanup temporary variables.
+set(_IMPORT_PREFIX)
+
+# Loop over all imported files and verify that they actually exist
+foreach(_cmake_target IN LISTS _cmake_import_check_targets)
+  if(CMAKE_VERSION VERSION_LESS "3.28"
+      OR NOT DEFINED _cmake_import_check_xcframework_for_${_cmake_target}
+      OR NOT IS_DIRECTORY "${_cmake_import_check_xcframework_for_${_cmake_target}}")
+    foreach(_cmake_file IN LISTS "_cmake_import_check_files_for_${_cmake_target}")
+      if(NOT EXISTS "${_cmake_file}")
+        message(FATAL_ERROR "The imported target \"${_cmake_target}\" references the file
+   \"${_cmake_file}\"
+but this file does not exist.  Possible reasons include:
+* The file was deleted, renamed, or moved to another location.
+* An install or uninstall procedure did not complete successfully.
+* The installation package was faulty and contained
+   \"${CMAKE_CURRENT_LIST_FILE}\"
+but not all the files it references.
+")
+      endif()
+    endforeach()
+  endif()
+  unset(_cmake_file)
+  unset("_cmake_import_check_files_for_${_cmake_target}")
+endforeach()
+unset(_cmake_target)
+unset(_cmake_import_check_targets)
+
+# This file does not depend on other imported targets which have
+# been exported from the same project but in a separate export set.
+
+# Commands beyond this point should not need to know the version.
+set(CMAKE_IMPORT_FILE_VERSION)
+cmake_policy(POP)

nvfuser/utils.py

@@ -0,0 +1,18 @@
+# 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",
+    "share",
+    "cmake",
+    "nvfuser",
+)

nvfuser/version.py

@@ -0,0 +1 @@
+_version_str = '0.2.25+git93b68e0'