warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/{config.py → _src/config.py}

@@ -1,169 +1,178 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Optional
-
-version: str = "1.9.0"
-"""Warp version string"""
-
-verify_fp: bool = False
-"""Enable floating-point verification for inputs and outputs.
-
-When enabled, checks if all values are finite before and after operations.
-
-Note: Enabling this flag impacts performance.
-"""
-
-verify_cuda: bool = False
-"""Enable CUDA error checking after kernel launches.
-
-This setting cannot be used during graph capture
-
-Note: Enabling this flag impacts performance
-"""
-
-print_launches: bool = False
-"""Enable detailed kernel launch logging.
-
-Prints information about each kernel launch including:
-
-- Launch dimensions
-- Input/output parameters
-- Target device
-
-Note: Enabling this flag impacts performance.
-"""
-
-mode: str = "release"
-"""Compilation mode for Warp kernels.
-
-Args:
-    mode: Either ``"release"`` or ``"debug"``.
-
-Note: Debug mode may impact performance.
-
-This setting can be overridden at the module level by setting the ``"mode"`` module option.
-"""
-
-verbose: bool = False
-"""Enable detailed logging during code generation and compilation."""
-
-verbose_warnings: bool = False
-"""Enable extended warning messages with source location information."""
-
-quiet: bool = False
-"""Disable Warp module initialization messages.
-
-Error messages and warnings remain unaffected.
-"""
-
-verify_autograd_array_access: bool = False
-"""Enable warnings for array overwrites that may affect gradient computation."""
-
-enable_vector_component_overwrites: bool = False
-"""Allow multiple writes to vector/matrix/quaternion components.
-
-Note: Enabling this may significantly increase kernel compilation time.
-"""
-
-cache_kernels: bool = True
-"""Enable kernel caching between application launches."""
-
-kernel_cache_dir: Optional[str] = None
-"""Directory path for storing compiled kernel cache.
-
-If ``None``, the path is determined in the following order:
-
-1. ``WARP_CACHE_PATH`` environment variable.
-2. System's user cache directory (via ``appdirs.user_cache_directory``).
-
-Note: Subdirectories prefixed with ``wp_`` will be created in this location.
-"""
-
-cuda_output: Optional[str] = None
-"""Preferred CUDA output format for kernel compilation.
-
-Args:
-    cuda_output: One of {``None``, ``"ptx"``, ``"cubin"``}. If ``None``, format is auto-determined.
-"""
-
-ptx_target_arch: Optional[int] = None
-"""Target architecture version for PTX generation, e.g., ``ptx_target_arch = 75``.
-
-If ``None``, the architecture is determined by devices present in the system.
-"""
-
-lineinfo: bool = False
-"""Enable the compilation of modules with line information.
-
-Modules compiled for GPU execution will be compiled with the
-``--generate-line-info`` compiler option, which generates line-number
-information for device code. Line-number information is always included when
-compiling a module in ``"debug"`` mode regardless of this setting.
-
-This setting can be overridden at the module level by setting the ``"lineinfo"`` module option.
-"""
-
-line_directives: bool = True
-"""Enable Python source line mapping in generated code.
-
-If ``True``, ``#line`` directives are inserted in generated code for modules
-compiled with line information  to map back to the original Python source file.
-"""
-
-compile_time_trace: bool = False
-"""Enable the generation of Trace Event Format files for runtime module compilation.
-
-These are JSON files that can be opened by tools like ``edge://tracing/`` and
-``chrome://tracing/``.
-
-This setting is currently only effective when compiling modules for the GPU with NVRTC (CUDA 12.8+).
-
-This setting can be overridden at the module level by setting the ``"compile_time_trace"`` module option.
-"""
-
-enable_backward: bool = True
-"""Enable compilation of kernel backward passes.
-
-This setting can be overridden at the module level by setting the ``"enable_backward"`` module option.
-"""
-
-llvm_cuda: bool = False
-"""Use Clang/LLVM compiler instead of NVRTC for CUDA compilation."""
-
-enable_graph_capture_module_load_by_default: bool = True
-"""Enable automatic module loading before graph capture.
-
-Only affects systems with CUDA driver versions below 12.3.
-"""
-
-enable_mempools_at_init: bool = True
-"""Enable CUDA memory pools during device initialization when supported."""
-
-max_unroll: int = 16
-"""Maximum unroll factor for loops.
-
-Note that ``max_unroll`` does not consider the total number of iterations in
-nested loops. This can result in a large amount of automatically generated code
-if each nested loop is below the ``max_unroll`` threshold.
-
-This setting can be overridden at the module level by setting the ``"max_unroll"`` module option.
-"""
-
-_git_commit_hash: Optional[str] = None
-"""Git commit hash associated with the Warp installation.
-
-Set automatically by CI, do not modify.
-"""
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+version: str = "1.10.0rc2"
+"""Warp version string"""
+
+verify_fp: bool = False
+"""Enable floating-point verification for inputs and outputs.
+
+When enabled, checks if all values are finite before and after operations.
+
+Note: Enabling this flag impacts performance.
+"""
+
+verify_cuda: bool = False
+"""Enable CUDA error checking after kernel launches.
+
+This setting cannot be used during graph capture
+
+Note: Enabling this flag impacts performance
+"""
+
+print_launches: bool = False
+"""Enable detailed kernel launch logging.
+
+Prints information about each kernel launch including:
+
+- Launch dimensions
+- Input/output parameters
+- Target device
+
+Note: Enabling this flag impacts performance.
+"""
+
+mode: str = "release"
+"""Compilation mode for Warp kernels.
+
+Args:
+    mode: Either ``"release"`` or ``"debug"``.
+
+Note: Debug mode may impact performance.
+
+This setting can be overridden at the module level by setting the ``"mode"`` module option.
+"""
+
+verbose: bool = False
+"""Enable detailed logging during code generation and compilation."""
+
+verbose_warnings: bool = False
+"""Enable extended warning messages with source location information."""
+
+quiet: bool = False
+"""Disable Warp module initialization messages.
+
+Error messages and warnings remain unaffected.
+"""
+
+verify_autograd_array_access: bool = False
+"""Enable warnings for array overwrites that may affect gradient computation."""
+
+enable_vector_component_overwrites: bool = False
+"""Allow multiple writes to vector/matrix/quaternion components.
+
+Note: Enabling this may significantly increase kernel compilation time.
+"""
+
+cache_kernels: bool = True
+"""Enable kernel caching between application launches."""
+
+kernel_cache_dir: Optional[str] = None
+"""Directory path for storing compiled kernel cache.
+
+If ``None``, the path is determined in the following order:
+
+1. ``WARP_CACHE_PATH`` environment variable.
+2. System's user cache directory (via ``appdirs.user_cache_directory``).
+
+Note: Subdirectories prefixed with ``wp_`` will be created in this location.
+"""
+
+cuda_output: Optional[str] = None
+"""Preferred CUDA output format for kernel compilation.
+
+Args:
+    cuda_output: One of {``None``, ``"ptx"``, ``"cubin"``}. If ``None``, format is auto-determined.
+"""
+
+ptx_target_arch: Optional[int] = None
+"""Target architecture version for PTX generation, e.g., ``ptx_target_arch = 75``.
+
+If ``None``, the architecture is determined by devices present in the system.
+"""
+
+lineinfo: bool = False
+"""Enable the compilation of modules with line information.
+
+Modules compiled for GPU execution will be compiled with the
+``--generate-line-info`` compiler option, which generates line-number
+information for device code. Line-number information is always included when
+compiling a module in ``"debug"`` mode regardless of this setting.
+
+This setting can be overridden at the module level by setting the ``"lineinfo"`` module option.
+"""
+
+line_directives: bool = True
+"""Enable Python source line mapping in generated code.
+
+If ``True``, ``#line`` directives are inserted in generated code for modules
+compiled with line information to map back to the original Python source file.
+"""
+
+compile_time_trace: bool = False
+"""Enable the generation of Trace Event Format files for runtime module compilation.
+
+These are JSON files that can be opened by tools like ``edge://tracing/`` and
+``chrome://tracing/``.
+
+This setting is currently only effective when compiling modules for the GPU with NVRTC (CUDA 12.8+).
+
+This setting can be overridden at the module level by setting the ``"compile_time_trace"`` module option.
+"""
+
+enable_backward: bool = True
+"""Enable compilation of kernel backward passes.
+
+This setting can be overridden at the module level by setting the ``"enable_backward"`` module option.
+"""
+
+llvm_cuda: bool = False
+"""Use Clang/LLVM compiler instead of NVRTC for CUDA compilation."""
+
+enable_graph_capture_module_load_by_default: bool = True
+"""Enable automatic module loading before graph capture.
+
+Only affects systems with CUDA driver versions below 12.3.
+"""
+
+enable_mempools_at_init: bool = True
+"""Enable CUDA memory pools during device initialization when supported."""
+
+max_unroll: int = 16
+"""Maximum unroll factor for loops.
+
+Note that ``max_unroll`` does not consider the total number of iterations in
+nested loops. This can result in a large amount of automatically generated code
+if each nested loop is below the ``max_unroll`` threshold.
+
+This setting can be overridden at the module level by setting the ``"max_unroll"`` module option.
+"""
+
+enable_tiles_in_stack_memory: bool = True
+"""Use stack memory instead of static memory for tile allocations on the CPU.
+
+Static memory in kernels is not well supported on some architectures (notably AArch64). We work
+around it by reserving stack memory on kernel entry and pointing a reserved callee-saved
+register to it (AArch64) or a single static pointer (x86-64). To disable this and revert to relying
+on the JIT-compiler's support for static memory, set this flag to ``False``.
+"""
+
+_git_commit_hash: Optional[str] = None
+"""Git commit hash associated with the Warp installation.
+
+Set automatically by CI, do not modify.
+"""

warp/_src/constants.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+from warp._src.types import constant
+
+__all__ = [  # noqa: RUF022
+    "E",
+    "e",
+    "INF",
+    "inf",
+    "LOG2E",
+    "log2e",
+    "LOG10E",
+    "log10e",
+    "LN2",
+    "ln2",
+    "LN10",
+    "ln10",
+    "NAN",
+    "nan",
+    "PHI",
+    "phi",
+    "PI",
+    "pi",
+    "HALF_PI",
+    "half_pi",
+    "TAU",
+    "tau",
+]
+
+E = e = constant(2.71828182845904523536)  # e
+LOG2E = log2e = constant(1.44269504088896340736)  # log2(e)
+LOG10E = log10e = constant(0.43429448190325182765)  # log10(e)
+LN2 = ln2 = constant(0.69314718055994530942)  # ln(2)
+LN10 = ln10 = constant(2.30258509299404568402)  # ln(10)
+PHI = phi = constant(1.61803398874989484820)  # golden constant
+PI = pi = constant(3.14159265358979323846)  # pi
+HALF_PI = half_pi = constant(1.57079632679489661923)  # half pi
+TAU = tau = constant(6.28318530717958647692)  # 2 * pi
+
+INF = inf = constant(math.inf)
+
+NAN = nan = constant(float("nan"))