warp-lang 1.1.0__py3-none-manylinux2014_aarch64.whl → 1.2.1__py3-none-manylinux2014_aarch64.whl

warp/build.py

@@ -50,10 +50,6 @@ def build_cpu(obj_path, cpp_path, mode="release", verify_fp=False, fast_math=Fal
             raise Exception(f"CPU kernel build failed with error code {err}")
 
 
-kernel_bin_dir = None
-kernel_gen_dir = None
-
-
 def init_kernel_cache(path=None):
     """Initialize kernel cache directory.
 
@@ -70,46 +66,23 @@ def init_kernel_cache(path=None):
     else:
         cache_root_dir = appdirs.user_cache_dir(appname="warp", appauthor="NVIDIA", version=warp.config.version)
 
-    cache_bin_dir = os.path.join(cache_root_dir, "bin")
-    cache_gen_dir = os.path.join(cache_root_dir, "gen")
-
-    if not os.path.isdir(cache_root_dir):
-        # print("Creating cache directory '%s'" % cache_root_dir)
-        os.makedirs(cache_root_dir, exist_ok=True)
-
-    if not os.path.isdir(cache_gen_dir):
-        # print("Creating codegen directory '%s'" % cache_gen_dir)
-        os.makedirs(cache_gen_dir, exist_ok=True)
-
-    if not os.path.isdir(cache_bin_dir):
-        # print("Creating binary directory '%s'" % cache_bin_dir)
-        os.makedirs(cache_bin_dir, exist_ok=True)
-
     warp.config.kernel_cache_dir = cache_root_dir
 
-    global kernel_bin_dir, kernel_gen_dir
-    kernel_bin_dir = cache_bin_dir
-    kernel_gen_dir = cache_gen_dir
+    os.makedirs(warp.config.kernel_cache_dir, exist_ok=True)
 
 
 def clear_kernel_cache():
     """Clear the kernel cache."""
 
-    is_intialized = kernel_bin_dir is not None and kernel_gen_dir is not None
-    assert is_intialized, "The kernel cache directory is not configured; wp.init() has not been called yet or failed."
-
-    import glob
+    warp.context.init()
 
-    paths = []
+    import shutil
 
-    if os.path.isdir(kernel_bin_dir):
-        pattern = os.path.join(kernel_bin_dir, "wp_*")
-        paths += glob.glob(pattern)
-
-    if os.path.isdir(kernel_gen_dir):
-        pattern = os.path.join(kernel_gen_dir, "wp_*")
-        paths += glob.glob(pattern)
+    is_intialized = warp.context.runtime is not None
+    assert is_intialized, "The kernel cache directory is not configured; wp.init() has not been called yet or failed."
 
-    for p in paths:
-        if os.path.isfile(p):
-            os.remove(p)
+    for item in os.listdir(warp.config.kernel_cache_dir):
+        item_path = os.path.join(warp.config.kernel_cache_dir, item)
+        if os.path.isdir(item_path) and item.startswith("wp_"):
+            # Remove the directory and its contents
+            shutil.rmtree(item_path, ignore_errors=True)

warp/build_dll.py

@@ -330,10 +330,10 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
             cu_out = cu_path + ".o"
 
             if mode == "debug":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" -g -G -O0 --compiler-options -fPIC,-fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{cutlass_enabled} {cutlass_includes} -o "{cu_out}" -c "{cu_path}"'
+                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -g -G -O0 --compiler-options -fPIC,-fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{cutlass_enabled} {cutlass_includes} -o "{cu_out}" -c "{cu_path}"'
 
             elif mode == "release":
-                cuda_cmd = f'"{cuda_home}/bin/nvcc" -O3 --compiler-options -fPIC,-fvisibility=hidden {" ".join(nvcc_opts)} -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{cutlass_enabled} {cutlass_includes} -o "{cu_out}" -c "{cu_path}"'
+                cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -O3 --compiler-options -fPIC,-fvisibility=hidden {" ".join(nvcc_opts)} -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{cutlass_enabled} {cutlass_includes} -o "{cu_out}" -c "{cu_path}"'
 
             with ScopedTimer("build_cuda", active=args.verbose):
                 run_cmd(cuda_cmd)

warp/builtins.py

@@ -284,6 +284,93 @@ add_builtin(
     In other words, it discards the integer part of x and is equivalent to ``x - trunc(x)``.""",
 )
 
+add_builtin(
+    "isfinite",
+    input_types={"x": Scalar},
+    value_type=builtins.bool,
+    group="Scalar Math",
+    doc="""Return ``True`` if x is a finite number, otherwise return ``False``.""",
+)
+add_builtin(
+    "isfinite",
+    input_types={"x": vector(length=Any, dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if all elements of the vector ``x`` are finite, otherwise return ``False``.",
+)
+add_builtin(
+    "isfinite",
+    input_types={"x": quaternion(dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if all elements of the quaternion ``x`` are finite, otherwise return ``False``.",
+)
+add_builtin(
+    "isfinite",
+    input_types={"m": matrix(shape=(Any, Any), dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if all elements of the matrix ``m`` are finite, otherwise return ``False``.",
+)
+
+add_builtin(
+    "isnan",
+    input_types={"x": Scalar},
+    value_type=builtins.bool,
+    doc="Return ``True`` if ``x`` is NaN, otherwise return ``False``.",
+    group="Scalar Math",
+)
+add_builtin(
+    "isnan",
+    input_types={"x": vector(length=Any, dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if any element of the vector ``x`` is NaN, otherwise return ``False``.",
+)
+add_builtin(
+    "isnan",
+    input_types={"x": quaternion(dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if any element of the quaternion ``x`` is NaN, otherwise return ``False``.",
+)
+add_builtin(
+    "isnan",
+    input_types={"m": matrix(shape=(Any, Any), dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if any element of the matrix ``m`` is NaN, otherwise return ``False``.",
+)
+
+add_builtin(
+    "isinf",
+    input_types={"x": Scalar},
+    value_type=builtins.bool,
+    group="Scalar Math",
+    doc="""Return ``True`` if x is positive or negative infinity, otherwise return ``False``.""",
+)
+add_builtin(
+    "isinf",
+    input_types={"x": vector(length=Any, dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if any element of the vector ``x`` is positive or negative infinity, otherwise return ``False``.",
+)
+add_builtin(
+    "isinf",
+    input_types={"x": quaternion(dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if any element of the quaternion ``x`` is positive or negative infinity, otherwise return ``False``.",
+)
+add_builtin(
+    "isinf",
+    input_types={"m": matrix(shape=(Any, Any), dtype=Scalar)},
+    value_type=builtins.bool,
+    group="Vector Math",
+    doc="Return ``True`` if any element of the matrix ``m`` is positive or negative infinity, otherwise return ``False``.",
+)
+
 
 def infer_scalar_type(arg_types):
     if arg_types is None:
@@ -2016,6 +2103,112 @@ add_builtin("iter_next", input_types={"query": mesh_query_aabb_t}, value_type=in
 # ---------------------------------
 # Volumes
 
+_volume_supported_value_types = {
+    int32,
+    int64,
+    uint32,
+    float32,
+    float64,
+    vec3f,
+    vec3d,
+    vec4f,
+    vec4d,
+}
+
+
+def volume_value_func(arg_types, kwds, templates):
+    try:
+        dtype = kwds["dtype"]
+    except KeyError as err:
+        raise RuntimeError(
+            "'dtype' keyword argument must be specified when calling generic volume lookup or sampling functions"
+        ) from err
+
+    if dtype not in _volume_supported_value_types:
+        raise RuntimeError(f"Unsupported volume type '{type_repr(dtype)}'")
+
+    templates.append(dtype)
+
+    return dtype
+
+
+add_builtin(
+    "volume_sample",
+    input_types={"id": uint64, "uvw": vec3, "sampling_mode": int, "dtype": Any},
+    value_func=volume_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume of type `dtype` given by ``id`` at the volume local-space point ``uvw``.
+
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`""",
+)
+
+
+def check_volume_value_grad_compatibility(dtype, grad_dtype):
+    if type_is_vector(dtype):
+        expected = matrix(shape=(type_length(dtype), 3), dtype=type_scalar_type(dtype))
+    else:
+        expected = vector(length=3, dtype=dtype)
+
+    if not types_equal(grad_dtype, expected):
+        raise RuntimeError(f"Incompatible gradient type, expected {type_repr(expected)}, got {type_repr(grad_dtype)}")
+
+
+def volume_sample_grad_value_func(arg_types, kwds, templates):
+    dtype = volume_value_func(arg_types, kwds, templates)
+
+    if len(arg_types) < 4:
+        raise RuntimeError("'volume_sample_grad' requires 4 positional arguments")
+
+    grad_type = arg_types[3]
+    check_volume_value_grad_compatibility(dtype, grad_type)
+    return dtype
+
+
+add_builtin(
+    "volume_sample_grad",
+    input_types={"id": uint64, "uvw": vec3, "sampling_mode": int, "grad": Any, "dtype": Any},
+    value_func=volume_sample_grad_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
+
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`""",
+)
+
+add_builtin(
+    "volume_lookup",
+    input_types={"id": uint64, "i": int, "j": int, "k": int, "dtype": Any},
+    value_type=int,
+    value_func=volume_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Returns the value of voxel with coordinates ``i``, ``j``, ``k`` for a volume of type type `dtype`.
+
+    If the voxel at this index does not exist, this function returns the background value.""",
+)
+
+
+def volume_store_value_func(arg_types, kwds, templates):
+    if len(arg_types) < 4:
+        raise RuntimeError("'volume_store' requires 5 positional arguments")
+
+    dtype = arg_types[4]
+    if dtype not in _volume_supported_value_types:
+        raise RuntimeError(f"Unsupported volume type '{type_repr(dtype)}'")
+
+    return None
+
+
+add_builtin(
+    "volume_store",
+    value_func=volume_store_value_func,
+    input_types={"id": uint64, "i": int, "j": int, "k": int, "value": Any},
+    export=False,
+    group="Volumes",
+    doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
+)
+
 add_builtin(
     "volume_sample_f",
     input_types={"id": uint64, "uvw": vec3, "sampling_mode": int},
@@ -2105,6 +2298,81 @@ add_builtin(
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
 )
 
+
+def volume_sample_index_value_func(arg_types, kwds, templates):
+    if len(arg_types) != 5:
+        raise RuntimeError("'volume_sample_index' requires 5 positional arguments")
+
+    dtype = arg_types[3].dtype
+
+    if not types_equal(dtype, arg_types[4]):
+        raise RuntimeError("The 'voxel_data' array and the 'background' value must have the same dtype")
+
+    return dtype
+
+
+add_builtin(
+    "volume_sample_index",
+    input_types={"id": uint64, "uvw": vec3, "sampling_mode": int, "voxel_data": array(dtype=Any), "background": Any},
+    value_func=volume_sample_index_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume given by ``id`` at the volume local-space point ``uvw``.
+
+    Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
+    This function is available for both index grids and classical volumes.
+    """,
+)
+
+
+def volume_sample_grad_index_value_func(arg_types, kwds, templates):
+    if len(arg_types) != 6:
+        raise RuntimeError("'volume_sample_grad_index' requires 6 positional arguments")
+
+    dtype = arg_types[3].dtype
+
+    if not types_equal(dtype, arg_types[4]):
+        raise RuntimeError("The 'voxel_data' array and the 'background' value must have the same dtype")
+
+    grad_type = arg_types[5]
+    check_volume_value_grad_compatibility(dtype, grad_type)
+    return dtype
+
+
+add_builtin(
+    "volume_sample_grad_index",
+    input_types={
+        "id": uint64,
+        "uvw": vec3,
+        "sampling_mode": int,
+        "voxel_data": array(dtype=Any),
+        "background": Any,
+        "grad": Any,
+    },
+    value_func=volume_sample_grad_index_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
+
+    Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
+    This function is available for both index grids and classical volumes.
+   """,
+)
+
+add_builtin(
+    "volume_lookup_index",
+    input_types={"id": uint64, "i": int, "j": int, "k": int},
+    value_type=int32,
+    group="Volumes",
+    doc="""Returns the index associated to the voxel with coordinates ``i``, ``j``, ``k``.
+
+    If the voxel at this index does not exist, this function returns -1.
+    This function is available for both index grids and classical volumes.
+    """,
+)
+
 add_builtin(
     "volume_index_to_world",
     input_types={"id": uint64, "uvw": vec3},
@@ -3067,7 +3335,7 @@ add_builtin(
 # ---------------------------------
 # Operators
 
-add_builtin("add", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar))
+add_builtin("add", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar), group="Operators")
 add_builtin(
     "add",
     input_types={"x": vector(length=Any, dtype=Scalar), "y": vector(length=Any, dtype=Scalar)},
@@ -3099,7 +3367,7 @@ add_builtin(
     group="Operators",
 )
 
-add_builtin("sub", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar))
+add_builtin("sub", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar), group="Operators")
 add_builtin(
     "sub",
     input_types={"x": vector(length=Any, dtype=Scalar), "y": vector(length=Any, dtype=Scalar)},
@@ -3214,7 +3482,7 @@ def mul_matmat_value_func(arg_types, kwds, _):
     return matrix(shape=(arg_types[0]._shape_[0], arg_types[1]._shape_[1]), dtype=arg_types[0]._wp_scalar_type_)
 
 
-add_builtin("mul", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar))
+add_builtin("mul", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar), group="Operators")
 add_builtin(
     "mul",
     input_types={"x": vector(length=Any, dtype=Scalar), "y": Scalar},
@@ -3311,7 +3579,7 @@ add_builtin(
     group="Operators",
 )
 
-add_builtin("mod", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar))
+add_builtin("mod", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar), group="Operators")
 
 add_builtin(
     "div",
@@ -3372,7 +3640,7 @@ add_builtin(
     group="Operators",
 )
 
-add_builtin("pos", input_types={"x": Scalar}, value_func=sametype_value_func(Scalar))
+add_builtin("pos", input_types={"x": Scalar}, value_func=sametype_value_func(Scalar), group="Operators")
 add_builtin(
     "pos",
     input_types={"x": vector(length=Any, dtype=Scalar)},
@@ -3394,7 +3662,7 @@ add_builtin(
     doc="",
     group="Operators",
 )
-add_builtin("neg", input_types={"x": Scalar}, value_func=sametype_value_func(Scalar))
+add_builtin("neg", input_types={"x": Scalar}, value_func=sametype_value_func(Scalar), group="Operators")
 add_builtin(
     "neg",
     input_types={"x": vector(length=Any, dtype=Scalar)},

warp/codegen.py

@@ -16,8 +16,7 @@ import re
 import sys
 import textwrap
 import types
-from types import ModuleType
-from typing import Any, Callable, Mapping
+from typing import Any, Callable, Dict, Mapping
 
 import warp.config
 from warp.types import *
@@ -554,7 +553,7 @@ class Adjoint:
 
     def __init__(
         adj,
-        func,
+        func: Callable[..., Any],
         overload_annotations=None,
         is_user_function=False,
         skip_forward_codegen=False,
@@ -1370,7 +1369,7 @@ class Adjoint:
             return obj
         if isinstance(obj, type):
             return obj
-        if isinstance(obj, ModuleType):
+        if isinstance(obj, types.ModuleType):
             return obj
 
         raise RuntimeError("Cannot reference a global variable from a kernel unless `wp.constant()` is being used")
@@ -2203,6 +2202,51 @@ class Adjoint:
         # return the Python code corresponding to the given AST node
         return ast.get_source_segment(adj.source, node)
 
+    def get_constant_references(adj) -> Dict[str, Any]:
+        """Traverses ``adj.tree`` and returns a dictionary containing constant variable names and values.
+
+        This function is meant to be used to populate a module's constants dictionary, which then feeds
+        into the computation of the module's ``content_hash``.
+        """
+
+        local_variables = set()  # Track local variables appearing on the LHS so we know when variables are shadowed
+        constants_dict = {}
+
+        for node in ast.walk(adj.tree):
+            if isinstance(node, ast.Name) and node.id not in local_variables:
+                # This node could be a reference to a wp.constant defined or imported into the current namespace
+
+                # try and resolve the name using the function's globals context (used to lookup constants + functions)
+                obj = adj.func.__globals__.get(node.id)
+
+                if obj is None:
+                    # Lookup constant in captured contents
+                    capturedvars = dict(
+                        zip(adj.func.__code__.co_freevars, [c.cell_contents for c in (adj.func.__closure__ or [])])
+                    )
+                    obj = capturedvars.get(str(node.id), None)
+
+                if warp.types.is_value(obj):
+                    constants_dict[node.id] = obj
+
+            elif isinstance(node, ast.Attribute):
+                obj, path = adj.resolve_static_expression(node, eval_types=False)
+
+                if warp.types.is_value(obj):
+                    constants_dict[".".join(path)] = obj
+            elif isinstance(node, ast.Assign):
+                # Add the LHS names to the local_variables so we know any subsequent uses are shadowed
+                lhs = node.targets[0]
+
+                if isinstance(lhs, ast.Tuple):
+                    for v in lhs.elts:
+                        if isinstance(v, ast.Name):
+                            local_variables.add(v.id)
+                elif isinstance(lhs, ast.Name):
+                    local_variables.add(lhs.id)
+
+        return constants_dict
+
 
 # ----------------
 # code generation
@@ -2454,6 +2498,9 @@ def constant_str(value):
     elif value == math.inf:
         return "INFINITY"
 
+    elif math.isnan(value):
+        return "NAN"
+
     else:
         # otherwise just convert constant to string
         return str(value)

warp/config.py

@@ -7,7 +7,7 @@
 
 from typing import Optional
 
-version: str = "1.1.0"
+version: str = "1.2.1"
 
 verify_fp: bool = False  # verify inputs and outputs are finite after each launch
 verify_cuda: bool = False  # if true will check CUDA errors after each kernel launch / memory operation
@@ -25,7 +25,7 @@ cuda_output: Optional[str] = (
     None  # preferred CUDA output format for kernels ("ptx" or "cubin"), determined automatically if unspecified
 )
 
-ptx_target_arch: int = 70  # target architecture for PTX generation, defaults to the lowest architecture that supports all of Warp's features
+ptx_target_arch: int = 75  # target architecture for PTX generation, defaults to the lowest architecture that supports all of Warp's features
 
 enable_backward: bool = True  # whether to compiler the backward passes of the kernels
 

warp/constants.py

@@ -22,6 +22,8 @@ __all__ = [
     "ln2",
     "LN10",
     "ln10",
+    "NAN",
+    "nan",
     "PHI",
     "phi",
     "PI",
@@ -43,3 +45,5 @@ 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"))