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

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