warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/dlpack.py

@@ -108,9 +108,6 @@ def dtype_to_dlpack(wp_dtype) -> DLDataType:
         return (DLDataTypeCode.kDLFloat, 32, 1)
     elif wp_dtype == warp.float64:
         return (DLDataTypeCode.kDLFloat, 64, 1)
-    elif wp_dtype in warp.types.vector_types:
-        # treat vector/matrix arrays as regular nd-arrays with one dtype lane
-        return (DLDataTypeCode.kDLFloat, 32, 1)
     else:
         raise RuntimeError(f"No conversion from Warp type {wp_dtype} to DLPack type")
 
@@ -196,7 +193,7 @@ def to_dlpack(wp_array: warp.array):
     # DLPack does not support structured arrays
     if isinstance(wp_array.dtype, warp.codegen.Struct):
         raise RuntimeError("Cannot convert structured Warp arrays to DLPack.")
-    
+
     holder = _Holder(wp_array)
 
     # allocate DLManagedTensor
@@ -204,12 +201,12 @@ def to_dlpack(wp_array: warp.array):
     dl_managed_tensor = DLManagedTensor.from_address(ctypes.pythonapi.PyMem_RawMalloc(size))
 
     # handle vector types
-    if wp_array.dtype in warp.types.vector_types:
+    if hasattr(wp_array.dtype, "_wp_scalar_type_"):
         # vector type, flatten the dimensions into one tuple
-        target_dtype = warp.float32
+        target_dtype = wp_array.dtype._wp_scalar_type_
         target_ndim = wp_array.ndim + len(wp_array.dtype._shape_)
         target_shape = (*wp_array.shape, *wp_array.dtype._shape_)
-        dtype_strides = warp.types.strides_from_shape(wp_array.dtype._shape_, warp.float32)
+        dtype_strides = warp.types.strides_from_shape(wp_array.dtype._shape_, wp_array.dtype._wp_scalar_type_)
         target_strides = (*wp_array.strides, *dtype_strides)
     else:
         # scalar type
@@ -255,7 +252,7 @@ def dtype_is_compatible(dl_dtype, wp_dtype):
         if dl_dtype.bits == 16:
             return wp_dtype == warp.float16
         elif dl_dtype.bits == 32:
-            return wp_dtype == warp.float32 or wp_dtype in warp.types.vector_types
+            return wp_dtype == warp.float32
         elif dl_dtype.bits == 64:
             return wp_dtype == warp.float64
     elif dl_dtype.type_code.value == DLDataTypeCode.kDLInt or dl_dtype.type_code.value == DLDataTypeCode.kDLUInt:
@@ -320,30 +317,33 @@ def from_dlpack(pycapsule, dtype=None) -> warp.array:
         # automatically detect dtype
         dtype = dtype_from_dlpack(dlt.dtype)
 
-    elif dtype_is_compatible(dlt.dtype, dtype):
-        # handle vector types
-        if dtype in warp.types.vector_types:
-            dtype_shape = dtype._shape_
-            dtype_dims = len(dtype._shape_)
-            if dtype_dims > len(shape) or dtype_shape != shape[-dtype_dims:]:
-                raise RuntimeError(
-                    f"Could not convert DLPack tensor with shape {shape} to Warp array with dtype={dtype}, ensure that source inner shape is {dtype_shape}"
-                )
-
-            if strides is not None:
-                # ensure the inner strides are contiguous
-                stride = 4
-                for i in range(dtype_dims):
-                    if strides[-i - 1] != stride:
-                        raise RuntimeError(
-                            f"Could not convert DLPack tensor with shape {shape} to Warp array with dtype={dtype}, because the source inner strides are not contiguous"
-                        )
-                    stride *= dtype_shape[-i - 1]
-                strides = tuple(strides[:-dtype_dims])
-
-            shape = tuple(shape[:-dtype_dims])
+    elif hasattr(dtype, "_wp_scalar_type_"):
+        # handle vector/matrix types
 
-    else:
+        if not dtype_is_compatible(dlt.dtype, dtype._wp_scalar_type_):
+            raise RuntimeError(f"Incompatible data types: {dlt.dtype} and {dtype}")
+
+        dtype_shape = dtype._shape_
+        dtype_dims = len(dtype._shape_)
+        if dtype_dims > len(shape) or dtype_shape != shape[-dtype_dims:]:
+            raise RuntimeError(
+                f"Could not convert DLPack tensor with shape {shape} to Warp array with dtype={dtype}, ensure that source inner shape is {dtype_shape}"
+            )
+
+        if strides is not None:
+            # ensure the inner strides are contiguous
+            stride = itemsize
+            for i in range(dtype_dims):
+                if strides[-i - 1] != stride:
+                    raise RuntimeError(
+                        f"Could not convert DLPack tensor with shape {shape} to Warp array with dtype={dtype}, because the source inner strides are not contiguous"
+                    )
+                stride *= dtype_shape[-i - 1]
+            strides = tuple(strides[:-dtype_dims]) or (itemsize,)
+
+        shape = tuple(shape[:-dtype_dims]) or (1,)
+
+    elif not dtype_is_compatible(dlt.dtype, dtype):
         # incompatible dtype requested
         raise RuntimeError(f"Incompatible data types: {dlt.dtype} and {dtype}")
 

warp/fabric.py

@@ -0,0 +1,326 @@
+import ctypes
+import math
+from typing import Any
+
+import warp
+from warp.types import *
+
+
+class fabricbucket_t(ctypes.Structure):
+    _fields_ = [
+        ("index_start", ctypes.c_size_t),
+        ("index_end", ctypes.c_size_t),
+        ("ptr", ctypes.c_void_p),
+        ("lengths", ctypes.c_void_p),
+    ]
+
+    def __init__(self, index_start=0, index_end=0, ptr=None, lengths=None):
+        self.index_start = index_start
+        self.index_end = index_end
+        self.ptr = ctypes.c_void_p(ptr)
+        self.lengths = ctypes.c_void_p(lengths)
+
+
+class fabricarray_t(ctypes.Structure):
+    _fields_ = [
+        ("buckets", ctypes.c_void_p),  # array of fabricbucket_t on the correct device
+        ("nbuckets", ctypes.c_size_t),
+        ("size", ctypes.c_size_t),
+    ]
+
+    def __init__(self, buckets=None, nbuckets=0, size=0):
+        self.buckets = ctypes.c_void_p(buckets)
+        self.nbuckets = nbuckets
+        self.size = size
+
+
+class indexedfabricarray_t(ctypes.Structure):
+    _fields_ = [
+        ("fa", fabricarray_t),
+        ("indices", ctypes.c_void_p),
+        ("size", ctypes.c_size_t),
+    ]
+
+    def __init__(self, fa=None, indices=None):
+        if fa is None:
+            self.fa = fabricarray_t()
+        else:
+            self.fa = fa.__ctype__()
+
+        if indices is None:
+            self.indices = ctypes.c_void_p(None)
+            self.size = 0
+        else:
+            self.indices = ctypes.c_void_p(indices.ptr)
+            self.size = indices.size
+
+
+def fabric_to_warp_dtype(type_info, attrib_name):
+    if not type_info[0]:
+        raise RuntimeError(f"Attribute '{attrib_name}' cannot be used in Warp")
+
+    base_type_dict = {
+        "b": warp.bool,  # boolean
+        "i1": warp.int8,
+        "i2": warp.int16,
+        "i4": warp.int32,
+        "i8": warp.int64,
+        "u1": warp.uint8,
+        "u2": warp.uint16,
+        "u4": warp.uint32,
+        "u8": warp.uint64,
+        "f2": warp.float16,
+        "f4": warp.float32,
+        "f8": warp.float64,
+    }
+
+    base_dtype = base_type_dict.get(type_info[1])
+    if base_dtype is None:
+        raise RuntimeError(f"Attribute '{attrib_name}' base data type '{type_info[1]}' is not supported in Warp")
+
+    elem_count = type_info[2]
+    role = type_info[4]
+
+    if role in ("text", "path"):
+        raise RuntimeError(f"Attribute '{attrib_name}' role '{role}' is not supported in Warp")
+
+    if elem_count > 1:
+        # vector or matrix type
+        if role == "quat" and elem_count == 4:
+            return quaternion(base_dtype)
+        elif role in ("matrix", "transform", "frame"):
+            # only square matrices are currently supported
+            mat_size = int(math.sqrt(elem_count))
+            assert mat_size * mat_size == elem_count
+            return matrix((mat_size, mat_size), base_dtype)
+        else:
+            return vector(elem_count, base_dtype)
+    else:
+        # scalar type
+        return base_dtype
+
+
+class fabricarray(noncontiguous_array_base[T]):
+    # member attributes available during code-gen (e.g.: d = arr.shape[0])
+    # (initialized when needed)
+    _vars = None
+
+    def __init__(self, data=None, attrib=None, dtype=Any, ndim=None):
+        super().__init__(ARRAY_TYPE_FABRIC)
+
+        if data is not None:
+            from .context import runtime
+
+            # ensure the attribute name was also specified
+            if not isinstance(attrib, str):
+                raise ValueError(f"Invalid attribute name: {attrib}")
+
+            # get the fabric interface dictionary
+            if isinstance(data, dict):
+                iface = data
+            elif hasattr(data, "__fabric_arrays_interface__"):
+                iface = data.__fabric_arrays_interface__
+            else:
+                raise ValueError(
+                    "Invalid data argument for fabricarray: expected dict or object with __fabric_arrays_interface__"
+                )
+
+            version = iface.get("version")
+            if version != 1:
+                raise ValueError(f"Unsupported Fabric interface version: {version}")
+
+            device = iface.get("device")
+            if not isinstance(device, str):
+                raise ValueError(f"Invalid Fabric interface device: {device}")
+
+            self.device = runtime.get_device(device)
+
+            attribs = iface.get("attribs")
+            if not isinstance(attribs, dict):
+                raise ValueError("Failed to get Fabric interface attributes")
+
+            # look up attribute info by name
+            attrib_info = attribs.get(attrib)
+            if not isinstance(attrib_info, dict):
+                raise ValueError(f"Failed to get attribute '{attrib}'")
+
+            type_info = attrib_info["type"]
+            assert len(type_info) == 5
+
+            self.dtype = fabric_to_warp_dtype(type_info, attrib)
+
+            self.access = attrib_info["access"]
+
+            pointers = attrib_info["pointers"]
+            counts = attrib_info["counts"]
+
+            if not (hasattr(pointers, "__len__") and hasattr(counts, "__len__") and len(pointers) == len(counts)):
+                raise RuntimeError("Attribute pointers and counts must be lists of the same size")
+
+            # check whether it's an array
+            array_depth = type_info[3]
+            if array_depth == 0:
+                self.ndim = 1
+                array_lengths = None
+            elif array_depth == 1:
+                self.ndim = 2
+                array_lengths = attrib_info["array_lengths"]
+                if not hasattr(array_lengths, "__len__") or len(array_lengths) != len(pointers):
+                    raise RuntimeError(
+                        "Attribute `array_lengths` must be a list of the same size as `pointers` and `counts`"
+                    )
+            else:
+                raise ValueError(f"Invalid attribute array depth: {array_depth}")
+
+            num_buckets = len(pointers)
+            size = 0
+
+            buckets = (fabricbucket_t * num_buckets)()
+            for i in range(num_buckets):
+                buckets[i].index_start = size
+                buckets[i].index_end = size + counts[i]
+                buckets[i].ptr = pointers[i]
+                if array_lengths:
+                    buckets[i].lengths = array_lengths[i]
+                size += counts[i]
+
+            if self.device.is_cuda:
+                # copy bucket info to device
+                with warp.ScopedStream(self.device.null_stream):
+                    buckets_size = ctypes.sizeof(buckets)
+                    buckets_ptr = self.device.allocator.alloc(buckets_size)
+                    runtime.core.memcpy_h2d(self.device.context, buckets_ptr, ctypes.addressof(buckets), buckets_size)
+            else:
+                buckets_ptr = ctypes.addressof(buckets)
+
+            self.buckets = buckets
+            self.size = size
+            self.shape = (size,)
+
+            self.ctype = fabricarray_t(buckets_ptr, num_buckets, size)
+
+        else:
+            # empty array or type annotation
+            self.dtype = dtype
+            self.ndim = ndim or 1
+            self.device = None
+            self.access = None
+            self.buckets = None
+            self.size = 0
+            self.shape = (0,)
+            self.ctype = fabricarray_t()
+
+    def __del__(self):
+        # release the GPU copy of bucket info
+        if self.buckets is not None and self.device.is_cuda:
+            buckets_size = ctypes.sizeof(self.buckets)
+            with self.device.context_guard:
+                self.device.allocator.free(self.ctype.buckets, buckets_size)
+
+    def __ctype__(self):
+        return self.ctype
+
+    def __len__(self):
+        return self.size
+
+    def __str__(self):
+        if self.device is None:
+            # type annotation
+            return f"fabricarray{self.dtype}"
+        else:
+            return str(self.numpy())
+
+    def __getitem__(self, key):
+        if isinstance(key, array):
+            return indexedfabricarray(fa=self, indices=key)
+        else:
+            raise ValueError(f"Fabric arrays only support indexing using index arrays, got key of type {type(key)}")
+
+    @property
+    def vars(self):
+        # member attributes available during code-gen (e.g.: d = arr.shape[0])
+        # Note: we use a shared dict for all fabricarray instances
+        if fabricarray._vars is None:
+            fabricarray._vars = {"size": warp.codegen.Var("size", uint64)}
+        return fabricarray._vars
+
+    def fill_(self, value):
+        # TODO?
+        # filling Fabric arrays of arrays is not supported, because they are jagged arrays of arbitrary lengths
+        if self.ndim > 1:
+            raise RuntimeError("Filling Fabric arrays of arrays is not supported")
+
+        super().fill_(value)
+
+
+# special case for fabric array of arrays
+# equivalent to calling fabricarray(..., ndim=2)
+def fabricarrayarray(**kwargs):
+    kwargs["ndim"] = 2
+    return fabricarray(**kwargs)
+
+
+class indexedfabricarray(noncontiguous_array_base[T]):
+    # member attributes available during code-gen (e.g.: d = arr.shape[0])
+    # (initialized when needed)
+    _vars = None
+
+    def __init__(self, fa=None, indices=None, dtype=None, ndim=None):
+        super().__init__(ARRAY_TYPE_FABRIC_INDEXED)
+
+        if fa is not None:
+            check_index_array(indices, fa.device)
+            self.fa = fa
+            self.indices = indices
+            self.dtype = fa.dtype
+            self.ndim = fa.ndim
+            self.device = fa.device
+            self.size = indices.size
+            self.shape = (indices.size,)
+            self.ctype = indexedfabricarray_t(fa, indices)
+        else:
+            # allow empty indexedarrays in type annotations
+            self.fa = None
+            self.indices = None
+            self.dtype = dtype
+            self.ndim = ndim or 1
+            self.device = None
+            self.size = 0
+            self.shape = (0,)
+            self.ctype = indexedfabricarray_t()
+
+    def __ctype__(self):
+        return self.ctype
+
+    def __len__(self):
+        return self.size
+
+    def __str__(self):
+        if self.device is None:
+            # type annotation
+            return f"indexedfabricarray{self.dtype}"
+        else:
+            return str(self.numpy())
+
+    @property
+    def vars(self):
+        # member attributes available during code-gen (e.g.: d = arr.shape[0])
+        # Note: we use a shared dict for all indexedfabricarray instances
+        if indexedfabricarray._vars is None:
+            indexedfabricarray._vars = {"size": warp.codegen.Var("size", uint64)}
+        return indexedfabricarray._vars
+
+    def fill_(self, value):
+        # TODO?
+        # filling Fabric arrays of arrays is not supported, because they are jagged arrays of arbitrary lengths
+        if self.ndim > 1:
+            raise RuntimeError("Filling indexed Fabric arrays of arrays is not supported")
+
+        super().fill_(value)
+
+
+# special case for indexed fabric array of arrays
+# equivalent to calling fabricarray(..., ndim=2)
+def indexedfabricarrayarray(**kwargs):
+    kwargs["ndim"] = 2
+    return indexedfabricarray(**kwargs)

warp/fem/__init__.py

@@ -0,0 +1,27 @@
+from .geometry import Geometry, Grid2D, Trimesh2D, Quadmesh2D, Grid3D, Tetmesh, Hexmesh
+from .geometry import GeometryPartition, LinearGeometryPartition, ExplicitGeometryPartition
+
+from .space import FunctionSpace, make_polynomial_space, ElementBasis
+from .space import BasisSpace, PointBasisSpace, make_polynomial_basis_space, make_collocated_function_space
+from .space import DofMapper, SkewSymmetricTensorMapper, SymmetricTensorMapper
+from .space import SpaceTopology, SpacePartition, SpaceRestriction, make_space_partition, make_space_restriction
+
+from .domain import GeometryDomain, Cells, Sides, BoundarySides, FrontierSides
+from .quadrature import Quadrature, RegularQuadrature, NodalQuadrature, ExplicitQuadrature, PicQuadrature
+from .polynomial import Polynomial
+
+from .field import FieldLike, DiscreteField, make_test, make_trial, make_restriction
+
+from .integrate import integrate, interpolate
+
+from .operator import integrand
+from .operator import position, normal, lookup, measure, measure_ratio, deformation_gradient
+from .operator import inner, grad, div, outer, grad_outer, div_outer
+from .operator import degree, at_node
+from .operator import D, curl, jump, average, grad_jump, grad_average
+
+from .types import Sample, Field, Domain, Coords, ElementIndex
+
+from .dirichlet import project_linear_system, normalize_dirichlet_projector
+
+from .cache import TemporaryStore, set_default_temporary_store, borrow_temporary, borrow_temporary_like