warp-lang 1.7.2__py3-none-win_amd64.whl → 1.8.0__py3-none-win_amd64.whl

warp/fem/space/topology.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Type
+from typing import ClassVar, Optional, Tuple, Type
 
 import warp as wp
 from warp.fem import cache
@@ -39,6 +39,12 @@ class SpaceTopology:
     .. note:: This will change to be defined per-element in future versions
     """
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "element_node_count": lambda obj: obj._make_constant_element_node_count(),
+        "element_node_sign": lambda obj: obj._make_constant_element_node_sign(),
+        "side_neighbor_node_counts": lambda obj: obj._make_constant_side_neighbor_node_counts(),
+    }
+
     @wp.struct
     class TopologyArg:
         """Structure containing arguments to be passed to device functions"""
@@ -51,8 +57,7 @@ class SpaceTopology:
         self.MAX_NODES_PER_ELEMENT = wp.constant(max_nodes_per_element)
         self.ElementArg = geometry.CellArg
 
-        self._make_constant_element_node_count()
-        self._make_constant_element_node_sign()
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     @property
     def geometry(self) -> Geometry:
@@ -67,6 +72,9 @@ class SpaceTopology:
         """Value of the topology argument structure to be passed to device functions"""
         return SpaceTopology.TopologyArg()
 
+    def fill_topo_arg(self, arg, device):
+        pass
+
     @property
     def name(self):
         return f"{self.__class__.__name__}_{self.MAX_NODES_PER_ELEMENT}"
@@ -182,6 +190,11 @@ class SpaceTopology:
         ):
             return NODES_PER_ELEMENT
 
+        return constant_element_node_count
+
+    def _make_constant_side_neighbor_node_counts(self):
+        NODES_PER_ELEMENT = wp.constant(self.MAX_NODES_PER_ELEMENT)
+
         @cache.dynamic_func(suffix=self.name)
         def constant_side_neighbor_node_counts(
             side_arg: self.geometry.SideArg,
@@ -189,8 +202,7 @@ class SpaceTopology:
         ):
             return NODES_PER_ELEMENT, NODES_PER_ELEMENT
 
-        self.element_node_count = constant_element_node_count
-        self.side_neighbor_node_counts = constant_side_neighbor_node_counts
+        return constant_side_neighbor_node_counts
 
     def _make_constant_element_node_sign(self):
         @cache.dynamic_func(suffix=self.name)
@@ -202,12 +214,21 @@ class SpaceTopology:
         ):
             return 1.0
 
-        self.element_node_sign = constant_element_node_sign
+        return constant_element_node_sign
 
 
 class TraceSpaceTopology(SpaceTopology):
     """Auto-generated trace topology defining the node indices associated to the geometry sides"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "inner_cell_index": lambda obj: obj._make_inner_cell_index(),
+        "outer_cell_index": lambda obj: obj._make_outer_cell_index(),
+        "neighbor_cell_index": lambda obj: obj._make_neighbor_cell_index(),
+        "element_node_index": lambda obj: obj._make_element_node_index(),
+        "element_node_count": lambda obj: obj._make_element_node_count(),
+        "element_node_sign": lambda obj: obj._make_element_node_sign(),
+    }
+
     def __init__(self, topo: SpaceTopology):
         self._topo = topo
 
@@ -218,14 +239,10 @@ class TraceSpaceTopology(SpaceTopology):
 
         self.TopologyArg = topo.TopologyArg
         self.topo_arg_value = topo.topo_arg_value
+        self.fill_topo_arg = topo.fill_topo_arg
 
-        self.inner_cell_index = self._make_inner_cell_index()
-        self.outer_cell_index = self._make_outer_cell_index()
-        self.neighbor_cell_index = self._make_neighbor_cell_index()
-
-        self.element_node_index = self._make_element_node_index()
-        self.element_node_count = self._make_element_node_count()
         self.side_neighbor_node_counts = None
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     def node_count(self) -> int:
         return self._topo.node_count()
@@ -354,21 +371,29 @@ class RegularDiscontinuousSpaceTopology(RegularDiscontinuousSpaceTopologyMixin,
 
 
 class DeformedGeometrySpaceTopology(SpaceTopology):
+    _dynamic_attribute_constructors: ClassVar = {
+        "element_node_index": lambda obj: obj._make_element_node_index(),
+        "element_node_count": lambda obj: obj._make_element_node_count(),
+        "element_node_sign": lambda obj: obj._make_element_node_sign(),
+        "side_neighbor_node_counts": lambda obj: obj._make_side_neighbor_node_counts(),
+    }
+
     def __init__(self, geometry: DeformedGeometry, base_topology: SpaceTopology):
         self.base = base_topology
         super().__init__(geometry, base_topology.MAX_NODES_PER_ELEMENT)
 
         self.node_count = self.base.node_count
         self.topo_arg_value = self.base.topo_arg_value
+        self.fill_topo_arg = self.base.fill_topo_arg
         self.TopologyArg = self.base.TopologyArg
 
-        self._make_passthrough_functions()
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     @property
     def name(self):
         return f"{self.base.name}_{self.geometry.field.name}"
 
-    def _make_passthrough_functions(self):
+    def _make_element_node_index(self):
         @cache.dynamic_func(suffix=self.name)
         def element_node_index(
             elt_arg: self.geometry.CellArg,
@@ -376,16 +401,22 @@ class DeformedGeometrySpaceTopology(SpaceTopology):
             element_index: ElementIndex,
             node_index_in_elt: int,
         ):
-            return self.base.element_node_index(elt_arg.elt_arg, topo_arg, element_index, node_index_in_elt)
+            return self.base.element_node_index(elt_arg.base_arg, topo_arg, element_index, node_index_in_elt)
+
+        return element_node_index
 
+    def _make_element_node_count(self):
         @cache.dynamic_func(suffix=self.name)
         def element_node_count(
             elt_arg: self.geometry.CellArg,
             topo_arg: self.TopologyArg,
             element_count: ElementIndex,
         ):
-            return self.base.element_node_count(elt_arg.elt_arg, topo_arg, element_count)
+            return self.base.element_node_count(elt_arg.base_arg, topo_arg, element_count)
+
+        return element_node_count
 
+    def _make_side_neighbor_node_counts(self):
         @cache.dynamic_func(suffix=self.name)
         def side_neighbor_node_counts(
             side_arg: self.geometry.SideArg,
@@ -394,6 +425,9 @@ class DeformedGeometrySpaceTopology(SpaceTopology):
             inner_count, outer_count = self.base.side_neighbor_node_counts(side_arg.base_arg, element_index)
             return inner_count, outer_count
 
+        return side_neighbor_node_counts
+
+    def _make_element_node_sign(self):
         @cache.dynamic_func(suffix=self.name)
         def element_node_sign(
             elt_arg: self.geometry.CellArg,
@@ -401,12 +435,9 @@ class DeformedGeometrySpaceTopology(SpaceTopology):
             element_index: ElementIndex,
             node_index_in_elt: int,
         ):
-            return self.base.element_node_sign(elt_arg.elt_arg, topo_arg, element_index, node_index_in_elt)
+            return self.base.element_node_sign(elt_arg.base_arg, topo_arg, element_index, node_index_in_elt)
 
-        self.element_node_index = element_node_index
-        self.element_node_count = element_node_count
-        self.element_node_sign = element_node_sign
-        self.side_neighbor_node_counts = side_neighbor_node_counts
+        return element_node_sign
 
 
 def forward_base_topology(topology_class: Type[SpaceTopology], geometry: Geometry, *args, **kwargs) -> SpaceTopology:

warp/fem/space/trimesh_function_space.py

@@ -50,12 +50,15 @@ class TrimeshSpaceTopology(SpaceTopology):
     @cache.cached_arg_value
     def topo_arg_value(self, device):
         arg = TrimeshTopologyArg()
+        self.fill_topo_arg(arg, device)
+        return arg
+
+    def fill_topo_arg(self, arg: TrimeshTopologyArg, device):
         arg.tri_edge_indices = self._tri_edge_indices.to(device)
         arg.edge_vertex_indices = self._mesh.edge_vertex_indices.to(device)
 
         arg.vertex_count = self._mesh.vertex_count()
         arg.edge_count = self._mesh.side_count()
-        return arg
 
     def _compute_tri_edge_indices(self):
         self._tri_edge_indices = wp.empty(

warp/fem/utils.py

@@ -19,6 +19,7 @@ import numpy as np
 
 import warp as wp
 import warp.fem.cache as cache
+import warp.types
 from warp.fem.linalg import (  # noqa: F401 (for backward compatibility, not part of public API but used in examples)
     array_axpy,
     inverse_qr,
@@ -28,6 +29,57 @@ from warp.fem.types import NULL_NODE_INDEX
 from warp.utils import array_scan, radix_sort_pairs, runlength_encode
 
 
+def type_zero_element(dtype):
+    suffix = warp.types.get_type_code(dtype)
+
+    if dtype in warp.types.scalar_types:
+
+        @cache.dynamic_func(suffix=suffix)
+        def zero_element():
+            return dtype(0.0)
+
+        return zero_element
+
+    @cache.dynamic_func(suffix=suffix)
+    def zero_element():
+        return dtype()
+
+    return zero_element
+
+
+def type_basis_element(dtype):
+    suffix = warp.types.get_type_code(dtype)
+
+    if dtype in warp.types.scalar_types:
+
+        @cache.dynamic_func(suffix=suffix)
+        def basis_element(coord: int):
+            return dtype(1.0)
+
+        return basis_element
+
+    if warp.types.type_is_matrix(dtype):
+        cols = dtype._shape_[1]
+
+        @cache.dynamic_func(suffix=suffix)
+        def basis_element(coord: int):
+            v = dtype()
+            i = coord // cols
+            j = coord - i * cols
+            v[i, j] = v.dtype(1.0)
+            return v
+
+        return basis_element
+
+    @cache.dynamic_func(suffix=suffix)
+    def basis_element(coord: int):
+        v = dtype()
+        v[coord] = v.dtype(1.0)
+        return v
+
+    return basis_element
+
+
 def compress_node_indices(
     node_count: int,
     node_indices: wp.array(dtype=int),
@@ -126,14 +178,7 @@ def host_read_at_index(array: wp.array, index: int = -1, temporary_store: cache.
 
     if index < 0:
         index += array.shape[0]
-
-    if array.device.is_cuda:
-        temp = cache.borrow_temporary(temporary_store, shape=1, dtype=int, pinned=True, device="cpu")
-        wp.copy(dest=temp.array, src=array, src_offset=index, count=1)
-        wp.synchronize_stream(wp.get_stream(array.device))
-        return temp.array.numpy()[0]
-
-    return array.numpy()[index]
+    return array[index : index + 1].numpy()[0]
 
 
 def masked_indices(

warp/jax.py

@@ -182,6 +182,5 @@ def from_jax(jax_array, dtype=None) -> warp.array:
     Returns:
         warp.array: The converted Warp array.
     """
-    import jax.dlpack
 
-    return warp.from_dlpack(jax.dlpack.to_dlpack(jax_array), dtype=dtype)
+    return warp.from_dlpack(jax_array, dtype=dtype)

warp/jax_experimental/ffi.py

@@ -306,7 +306,6 @@ class FfiCallable:
         self.graph_compatible = graph_compatible
         self.output_dims = output_dims
         self.first_array_arg = None
-        self.has_static_args = False
         self.call_id = 0
         self.call_descriptors = {}
 
@@ -335,8 +334,6 @@ class FfiCallable:
                 if arg.is_array:
                     if arg_idx < self.num_inputs and self.first_array_arg is None:
                         self.first_array_arg = arg_idx
-                else:
-                    self.has_static_args = True
                 self.args.append(arg)
             arg_idx += 1
 
@@ -425,14 +422,11 @@ class FfiCallable:
         module = wp.get_module(self.func.__module__)
         module.load(device)
 
-        if self.has_static_args:
-            # save call data to be retrieved by callback
-            call_id = self.call_id
-            self.call_descriptors[call_id] = FfiCallDesc(static_inputs)
-            self.call_id += 1
-            return call(*args, call_id=call_id)
-        else:
-            return call(*args)
+        # save call data to be retrieved by callback
+        call_id = self.call_id
+        self.call_descriptors[call_id] = FfiCallDesc(static_inputs)
+        self.call_id += 1
+        return call(*args, call_id=call_id)
 
     def ffi_callback(self, call_frame):
         try:
@@ -454,11 +448,10 @@ class FfiCallable:
                         )
                     return None
 
-            if self.has_static_args:
-                # retrieve call info
-                attrs = decode_attrs(call_frame.contents.attrs)
-                call_id = int(attrs["call_id"])
-                call_desc = self.call_descriptors[call_id]
+            # retrieve call info
+            attrs = decode_attrs(call_frame.contents.attrs)
+            call_id = int(attrs["call_id"])
+            call_desc = self.call_descriptors[call_id]
 
             num_inputs = call_frame.contents.args.size
             inputs = ctypes.cast(call_frame.contents.args.args, ctypes.POINTER(ctypes.POINTER(XLA_FFI_Buffer)))
@@ -500,8 +493,10 @@ class FfiCallable:
             # call the Python function with reconstructed arguments
             with wp.ScopedStream(stream, sync_enter=False):
                 if stream.is_capturing:
-                    with wp.ScopedCapture(stream=stream, external=True):
+                    with wp.ScopedCapture(stream=stream, external=True) as capture:
                         self.func(*arg_list)
+                    # keep a reference to the capture object to prevent required modules getting unloaded
+                    call_desc.capture = capture
                 else:
                     self.func(*arg_list)
 

warp/jax_experimental/xla_ffi.py

@@ -130,14 +130,14 @@ class XLA_FFI_DataType(enum.IntEnum):
 #   int64_t* dims;  // length == rank
 # };
 class XLA_FFI_Buffer(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("dtype", ctypes.c_int),  # XLA_FFI_DataType
         ("data", ctypes.c_void_p),
         ("rank", ctypes.c_int64),
         ("dims", ctypes.POINTER(ctypes.c_int64)),
-    ]
+    )
 
 
 # typedef enum {
@@ -162,13 +162,13 @@ class XLA_FFI_RetType(enum.IntEnum):
 #   void** args;             // length == size
 # };
 class XLA_FFI_Args(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("size", ctypes.c_int64),
         ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_ArgType*
         ("args", ctypes.POINTER(ctypes.c_void_p)),
-    ]
+    )
 
 
 # struct XLA_FFI_Rets {
@@ -179,13 +179,13 @@ class XLA_FFI_Args(ctypes.Structure):
 #   void** rets;             // length == size
 # };
 class XLA_FFI_Rets(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("size", ctypes.c_int64),
         ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_RetType*
         ("rets", ctypes.POINTER(ctypes.c_void_p)),
-    ]
+    )
 
 
 # typedef struct XLA_FFI_ByteSpan {
@@ -193,7 +193,10 @@ class XLA_FFI_Rets(ctypes.Structure):
 #   size_t len;
 # } XLA_FFI_ByteSpan;
 class XLA_FFI_ByteSpan(ctypes.Structure):
-    _fields_ = [("ptr", ctypes.POINTER(ctypes.c_char)), ("len", ctypes.c_size_t)]
+    _fields_ = (
+        ("ptr", ctypes.POINTER(ctypes.c_char)),
+        ("len", ctypes.c_size_t),
+    )
 
 
 # typedef struct XLA_FFI_Scalar {
@@ -201,7 +204,10 @@ class XLA_FFI_ByteSpan(ctypes.Structure):
 #   void* value;
 # } XLA_FFI_Scalar;
 class XLA_FFI_Scalar(ctypes.Structure):
-    _fields_ = [("dtype", ctypes.c_int), ("value", ctypes.c_void_p)]
+    _fields_ = (
+        ("dtype", ctypes.c_int),
+        ("value", ctypes.c_void_p),
+    )
 
 
 # typedef struct XLA_FFI_Array {
@@ -210,7 +216,11 @@ class XLA_FFI_Scalar(ctypes.Structure):
 #   void* data;
 # } XLA_FFI_Array;
 class XLA_FFI_Array(ctypes.Structure):
-    _fields_ = [("dtype", ctypes.c_int), ("size", ctypes.c_size_t), ("data", ctypes.c_void_p)]
+    _fields_ = (
+        ("dtype", ctypes.c_int),
+        ("size", ctypes.c_size_t),
+        ("data", ctypes.c_void_p),
+    )
 
 
 # typedef enum {
@@ -235,14 +245,14 @@ class XLA_FFI_AttrType(enum.IntEnum):
 #   void** attrs;              // length == size
 # };
 class XLA_FFI_Attrs(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("size", ctypes.c_int64),
         ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_AttrType*
         ("names", ctypes.POINTER(ctypes.POINTER(XLA_FFI_ByteSpan))),
         ("attrs", ctypes.POINTER(ctypes.c_void_p)),
-    ]
+    )
 
 
 # struct XLA_FFI_Api_Version {
@@ -252,12 +262,12 @@ class XLA_FFI_Attrs(ctypes.Structure):
 #   int minor_version;  // out
 # };
 class XLA_FFI_Api_Version(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("major_version", ctypes.c_int),
         ("minor_version", ctypes.c_int),
-    ]
+    )
 
 
 # enum XLA_FFI_Handler_TraitsBits {
@@ -276,11 +286,11 @@ class XLA_FFI_Handler_TraitsBits(enum.IntEnum):
 #   XLA_FFI_Handler_Traits traits;
 # };
 class XLA_FFI_Metadata(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("api_version", XLA_FFI_Api_Version),  # XLA_FFI_Extension_Type
         ("traits", ctypes.c_uint32),  # XLA_FFI_Handler_Traits
-    ]
+    )
 
 
 # struct XLA_FFI_Metadata_Extension {
@@ -288,7 +298,10 @@ class XLA_FFI_Metadata(ctypes.Structure):
 #   XLA_FFI_Metadata* metadata;
 # };
 class XLA_FFI_Metadata_Extension(ctypes.Structure):
-    _fields_ = [("extension_base", XLA_FFI_Extension_Base), ("metadata", ctypes.POINTER(XLA_FFI_Metadata))]
+    _fields_ = (
+        ("extension_base", XLA_FFI_Extension_Base),
+        ("metadata", ctypes.POINTER(XLA_FFI_Metadata)),
+    )
 
 
 # typedef enum {
@@ -337,12 +350,12 @@ class XLA_FFI_Error_Code(enum.IntEnum):
 #   XLA_FFI_Error_Code errc;
 # };
 class XLA_FFI_Error_Create_Args(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("message", ctypes.c_char_p),
         ("errc", ctypes.c_int),
-    ]  # XLA_FFI_Error_Code
+    )  # XLA_FFI_Error_Code
 
 
 XLA_FFI_Error_Create = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_Error_Create_Args))
@@ -355,12 +368,12 @@ XLA_FFI_Error_Create = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_
 #   void* stream;  // out
 # };
 class XLA_FFI_Stream_Get_Args(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("ctx", ctypes.c_void_p),  # XLA_FFI_ExecutionContext*
         ("stream", ctypes.c_void_p),
-    ]  # // out
+    )  # // out
 
 
 XLA_FFI_Stream_Get = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_Stream_Get_Args))
@@ -391,7 +404,7 @@ XLA_FFI_Stream_Get = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_St
 #   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Future_SetError);
 # };
 class XLA_FFI_Api(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("api_version", XLA_FFI_Api_Version),
@@ -412,7 +425,7 @@ class XLA_FFI_Api(ctypes.Structure):
         ("XLA_FFI_Future_Create", ctypes.c_void_p),  # XLA_FFI_Future_Create
         ("XLA_FFI_Future_SetAvailable", ctypes.c_void_p),  # XLA_FFI_Future_SetAvailable
         ("XLA_FFI_Future_SetError", ctypes.c_void_p),  # XLA_FFI_Future_SetError
-    ]
+    )
 
 
 # struct XLA_FFI_CallFrame {
@@ -431,7 +444,7 @@ class XLA_FFI_Api(ctypes.Structure):
 #   XLA_FFI_Future* future;  // out
 # };
 class XLA_FFI_CallFrame(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("api", ctypes.POINTER(XLA_FFI_Api)),
@@ -441,7 +454,7 @@ class XLA_FFI_CallFrame(ctypes.Structure):
         ("rets", XLA_FFI_Rets),
         ("attrs", XLA_FFI_Attrs),
         ("future", ctypes.c_void_p),  # XLA_FFI_Future* // out
-    ]
+    )
 
 
 _xla_data_type_to_constructor = {