Import NVIDIA/warp from GitHub.
#MIGRATION_3P_PY_WARP__DEFAULT - 23f1a48c1acbc1aa1f353423ec5f9182560c565d Upgrade to libmathdx 0.2.2 (GH-809) by Zach Corse <zcorse@nvidia.com> - a5b978abc2995c875498111d3cd4d307e621fe1f Fix GH-822 adj_tile_sum with shared tiles by Gilles Daviet <gdaviet@nvidia.com> - 4cd0fda5978c550665b54235b8f19cf87b844989 Disable parallel test execution for MJWarp by Eric Shi <ershi@nvidia.com> - 7408f297d03e7ed605f6d31fcfdaa288a9058ece Fix the OpenGL Renderer Box Colors (GH-810) by Christopher Crouzet <ccrouzet@nvidia.com> - 1f3f589ba6732a3272f9e1fe8cc2a292864df4cb Fix rendering of meshes with different scales (GH-828) by Eric Heiden <eheiden@nvidia.com> - a317a3d064f3518429e48858c5a99c7270934b5d Add a config file by Eric Shi <ershi@nvidia.com> - 61848a5f12b1b1b1275a7a57ae734318e1cbb7d6 Resolve GH-832 "Fix bsr_from_triplets prune_numerical_zer... by Gilles Daviet <gdaviet@nvidia.com> - f9fc3a1b9b3a69beee57766554ea10036ebe64b7 Extend the OpenGL example to show how to use ImGui (GH-833) by Tobias Widmer <twidmer@nvidia.com> - bd2ca117a003736d95ec9938764a93eea10086fb Add flag to enable support for mac os - it avoids OpenGL ... by Tobias Widmer <twidmer@nvidia.com> - 3c6388890e7a18c0164504b68b53ccb60859cbc0 Fix hashing of static expressions ([GH-830], partially re... by Eric Heiden <eheiden@nvidia.com> - 74cac74e57f4b9d90f77cbe0b2b01f1f5ac84c67 fix box-box collision by Zhenggehan <gehan.zheng@vanderbilt.edu> - 5d93b757c1febd64ff3e273e6c8225202391b5cc Remove a redundant assignment for query_b in box collision by Zhenggehan <gehan.zheng@vanderbilt.edu> - 486422dc20d5665d9622115eb34fc5a7b9ef831d Add missing changelog entries for GH-822 and GH-809 by Zach Corse <zcorse@nvidia.com> - 52be82e7bded3a496fbba04fb90e31b0c5512950 Add a missing changelog entry for GH-839 by Eric Shi <ershi@nvidia.com> - b266b6a1497dd910cb23cbc4f0c03035005da8a1 GitLab jobs now use the R570 runners by Eric Shi <ershi@nvidia.com> - 7368bf7d630c58ca9999f5284076cd5c0e9f2dc3 Fix potential issues in bitonic tile_sort ([GH-836]) by Kenny Vilella <kvilella@nvidia.com> - 31621ff99b90927e608f56c5165d8b6332ec5b2f Fix for wp.tile_min() when tile thread block includes emp... by Zach Corse <zcorse@nvidia.com> - 569b5a2b3672e0501958e6a3fbc6901df8800c49 Add more ImGui helper functions GH-844 by Tobias Widmer <twidmer@nvidia.com> - a5a746a14254ae8ef89197c6d1988e3d33013683 Arbitrary size blocks and tile-based computations in warp... by Gilles Daviet <gdaviet@nvidia.com> - 6ef9436700acb9f80cf7f3a6a34618a5e2da0a3d Fix `adj_div(scalar, vec/mat/quat)` by Christopher Crouzet <christopher@crouzet.pm> - 0044654bb97a9de8c75629e52fa04a04fa632c21 Fix miscellaneous issues with adjoints by Christopher Crouzet <christopher@crouzet.pm> - 42e374713557e3a1a85f23286ff1b02334eb15ea Fix the Possible Dead Loop in wp_balance_coloring (GH-816) by Anka Chen <ankac@nvidia.com> - b39954ece68e67c124e1ff856fd189cc366debb8 Graph-capturable linear solvers (GH-843) by Gilles Daviet <gdaviet@nvidia.com> - b2a9133d0125c2b973c8f0dbafaa31830e2af0ef Expose `array.ptr` in kernels (GH-819) by Gilles Daviet <gdaviet@nvidia.com> - 3fa6c9b19c9f4774bd3f61f102a0867136563804 Fix wrong version passed to cuDeviceGetUuid by Eric Shi <ershi@nvidia.com> - 195fd4e0ccb83ed7cfd74b716a6d33bfc9a35217 Add new publication: GeoWarp framework for geomechanics by Eric Shi <ershi@nvidia.com> - 514e8b5a8f6a80a5a6aa230d9d0f23e3f2e9aa9a Add edit page link on right sidebar by Eric Shi <ershi@nvidia.com> - 178e8430911b77b44bf1b0cd543f0304423c2ef2 Fix positioning of changelog entry of MR !1438 by Gilles Daviet <gdaviet@nvidia.com> - 907cbd4b572a937b8b1c9a5fbd4bc21976304596 Use latest clang for builds in CI/CD by Eric Shi <ershi@nvidia.com> - 649d206616735ec923d7a6bc1631bb03f4abcc6a Handle CUDA 12.8 + Thor combination by Eric Shi <ershi@nvidia.com> - 8775c618ec9b1a8efd91f18b48ef6acd15caad92 Allow arr = wp.array(dtype=MyStruct...) in kernels GH-853 by Tobias Widmer <twidmer@nvidia.com> - d3d0bdc67854c74045499253fa556369337925b8 Fix default PTX architecture selection to include only su... by Eric Shi <ershi@nvidia.com> - 0e2b1f9cd28db7fe366647b39a09e2f43ebc236b Fixed the Bug of Missing Cloth-body Contact in VBD Integr... by Anka Chen <ankac@nvidia.com> - 70479a95722e50ab6a32c477a238b801b0db5d4b warp.fem: leverage tiles to sum over quadrature points (G... by Gilles Daviet <gdaviet@nvidia.com> - 2446e4c06a1d6d50ab825f6991d402f7959745c2 Add Warp stuctured array NumPy interop examples to docs (... by Zach Corse <zcorse@nvidia.com> - bac82151b4cc46659f5278df9749cfd1d994ca61 Refactor some tests to reduce recompilation cost by Eric Shi <ershi@nvidia.com> - c0cc239b5fe457c8aedca636e7d2e3f32f060a62 Fix debug mode compilation (GH-856) by Alain Denzler <adenzler@nvidia.com> - 906495b742f583b5c5a524a23ffc7112b542286a Fix typo in wp.fem nodal field DG gradient by Gilles Daviet <gdaviet@nvidia.com> - 476765315edf7268274f8862af540a7ddb68ce4f Reenable parallel execution of mjwarp GitLab job by Eric Shi <ershi@nvidia.com> - 8f2bcfc503e82ad492dabc6cdbd9103495f0ff17 Disable Blackwell job by Eric Shi <ershi@nvidia.com> - 531e822b1939be95edb11c56db310c0ac4701e17 Support initializing fixed-size arrays in kernels by Christopher Crouzet <christopher@crouzet.pm> - 52f7f53f247302c666e73490e33228d0172e08b6 Support fixed-size arrays with structs as `dtype` by Christopher Crouzet <christopher@crouzet.pm> - af038bddb408acd76fe009fb0b548ec0da68b74b Fix keyword arguments not being reported by Christopher Crouzet <christopher@crouzet.pm> - 1bd4fac0637cd1eed3abc9af14699b5c9e05b35c JAX FFI graph capture using Warp (GH-848) by Lukasz Wawrzyniak <lwawrzyniak@nvidia.com> - ff63d5073ae453d04aa8f78dcf9ed90bf7d70e5b Split `assign_copy` tests by Christopher Crouzet <christopher@crouzet.pm> - 62daadda34f794bcb8c8ece4fe40c95224005b26 Add `wp.matrix_from_rows/cols()` to Python's API by Christopher Crouzet <christopher@crouzet.pm> - a19818854f716629028186a24815629cb5513145 Add support for negative indexing by Christopher Crouzet <christopher@crouzet.pm> - f440c5b15766943623870b8e887efe78bb9f9d62 Implement a slice object by Christopher Crouzet <christopher@crouzet.pm> - e26f23207acf09449deeeea7db2168d2123804cf Add support for slicing vec/quat/transform by Christopher Crouzet <christopher@crouzet.pm> - 28907cdfcae397d0e32b6f34d302d178677308e0 Add `wp.matrix_get/set_col()` to Python's API by Christopher Crouzet <christopher@crouzet.pm> - 303c0ac8b680f788700e3974aba161042fef449b Add support for slicing matrices by Christopher Crouzet <christopher@crouzet.pm> - bbc79c3dbcc9b37b00b9d61094c920e93515a2b1 Implement adjoints for slice indexing by Christopher Crouzet <christopher@crouzet.pm> - c1e872f74f886592490771fe37e410f7b3905701 Update the hash with the config for `assign_copy` by Christopher Crouzet <christopher@crouzet.pm> - b978cf94a87f66e90e418dda3c05d7c864b52849 Symmetrize operators by Christopher Crouzet <christopher@crouzet.pm> - fd71d280d7a9d390bd7b322519ade06be62ac875 Fix a documented return type by Christopher Crouzet <christopher@crouzet.pm> - 57fd611636f804cdaec3c8b69a0271f7717b770d Add a doc section for indexing/slicing by Christopher Crouzet <christopher@crouzet.pm> - 6aa03089a96ff18792fbc2a3cbf50e205cea7783 Reduce the code size by Christopher Crouzet <christopher@crouzet.pm> - 6c6c6e738972b25e69d171f6801b717ca921664b Fix wrong macros being used by Christopher Crouzet <christopher@crouzet.pm> - 43be725152909f6434f9bf6aa0cdcf5e77bcd8c8 Add post-loop checks by Christopher Crouzet <christopher@crouzet.pm> - 16a6a575ce6bee368b056a70c9d5bec8c6f1a409 Work around GH-637 in wp.fem integrands by Gilles Daviet <gdaviet@nvidia.com> - 637bafffb90d495b222bf1c79d7173f60aa1871c Resolve function return types for wp.map() and wp.tile_ma... by Zach Corse <zcorse@nvidia.com> - 18849b08e0809ccb357d22a6601139b9c3277cd2 Publish macos dev wheel to artifactory by Eric Shi <ershi@nvidia.com> - de7481a377d0380e13208a7261996daba69fd9ba Tentative fix for bsr_mm asv regression from c5fdbc20 by Gilles Daviet <gdaviet@nvidia.com> - eba4be3cf12ac3accdd0d48a265d0f8766ed8c99 Reenable Blackwell GitLab job by Eric Shi <ershi@nvidia.com> - ee012fe62c8cb47b498ac430652c234368d41723 GH-870 Detect and warn about long file paths in Windows by Gilles Daviet <gdaviet@nvidia.com> - fe5785431d4423a1faef64f2c810fa5ac18265d1 Reduce error messages when CUDA driver is not found by Eric Shi <ershi@nvidia.com> - e5904c3d78e9e77174ab84527575a06c3c9d9265 Clarify Python vs. Warp scope execution of @wp.func funct... by Eric Shi <ershi@nvidia.com> - b4ca2187d15f8b2c4220dc89b953670ab1b99913 Compile time improvements by Eric Shi <ershi@nvidia.com> - 3a07ca72d155ff8666fc68a7c51a8690d8fef19e Add a documentation section about indexed arrays by Christopher Crouzet <christopher@crouzet.pm> - 248fd07b67395323098231e64169e3d1c8216d89 Namespace the remaining warp-clang symbols by Christopher Crouzet <christopher@crouzet.pm> - c6af3501853d30bd059f3ac1c9c9a9057f72247b Fix adding superfluous inactive nodes to tet function spa... by Gilles Daviet <gdaviet@nvidia.com> - 62ac48861731a1800ad46548dc254d94362bdd18 Use usd-core==25.5.1 for Windows jobs to avoid failures by Eric Shi <ershi@nvidia.com> - 5d1c2ba11cb6a11c4e89b4f47980f3063e3e63fc Add support for `IntEnum`, `IntFlag` inside Warp kernels ... by Eric Heiden <eheiden@nvidia.com> - a4aba20a00e3d267d178a49b2ee145b82966f166 Update Warp version to 1.8.1 by Eric Shi <ershi@nvidia.com> PiperOrigin-RevId: 792191632 Change-Id: I227d605677c925ca5e7fa64948ef884ef25fa5ad
This commit is contained in:
Baruch Tabanpour
committed by
Copybara-Service
Copybara-Service
parent
eb1da7c43d
commit
bb4aeb44f1
+84
-14
@@ -16,6 +16,7 @@
|
||||
import ctypes
|
||||
import threading
|
||||
import traceback
|
||||
from enum import IntEnum
|
||||
from typing import Callable, Optional
|
||||
|
||||
import jax
|
||||
@@ -28,6 +29,12 @@ from warp.types import array_t, launch_bounds_t, strides_from_shape, type_to_war
|
||||
from .xla_ffi import *
|
||||
|
||||
|
||||
class GraphMode(IntEnum):
|
||||
NONE = 0 # don't capture a graph
|
||||
JAX = 1 # let JAX capture a graph
|
||||
WARP = 2 # let Warp capture a graph
|
||||
|
||||
|
||||
class FfiArg:
|
||||
def __init__(self, name, type, in_out=False):
|
||||
self.name = name
|
||||
@@ -331,15 +338,16 @@ class FfiKernel:
|
||||
class FfiCallDesc:
|
||||
def __init__(self, static_inputs):
|
||||
self.static_inputs = static_inputs
|
||||
self.captures = {}
|
||||
|
||||
|
||||
class FfiCallable:
|
||||
def __init__(self, func, num_outputs, graph_compatible, vmap_method, output_dims, in_out_argnames):
|
||||
def __init__(self, func, num_outputs, graph_mode, vmap_method, output_dims, in_out_argnames):
|
||||
self.func = func
|
||||
self.name = generate_unique_name(func)
|
||||
self.num_outputs = num_outputs
|
||||
self.vmap_method = vmap_method
|
||||
self.graph_compatible = graph_compatible
|
||||
self.graph_mode = graph_mode
|
||||
self.output_dims = output_dims
|
||||
self.first_array_arg = None
|
||||
self.call_id = 0
|
||||
@@ -396,6 +404,12 @@ class FfiCallable:
|
||||
self.input_args = self.args[: self.num_inputs] # includes in-out args
|
||||
self.output_args = self.args[self.num_inputs :] # pure output args
|
||||
|
||||
# Buffer indices for array arguments in callback.
|
||||
# In-out buffers are the same pointers in the XLA call frame,
|
||||
# so we only include them for inputs and skip them for outputs.
|
||||
self.array_input_indices = [i for i, arg in enumerate(self.input_args) if arg.is_array]
|
||||
self.array_output_indices = list(range(self.num_in_out, self.num_outputs))
|
||||
|
||||
# Build input output aliases.
|
||||
out_id = 0
|
||||
input_output_aliases = {}
|
||||
@@ -503,11 +517,10 @@ class FfiCallable:
|
||||
|
||||
def ffi_callback(self, call_frame):
|
||||
try:
|
||||
# TODO Try-catch around the body and return XLA_FFI_Error on error.
|
||||
extension = call_frame.contents.extension_start
|
||||
# On the first call, XLA runtime will query the API version and traits
|
||||
# metadata using the |extension| field. Let us respond to that query
|
||||
# if the metadata extension is present.
|
||||
extension = call_frame.contents.extension_start
|
||||
if extension:
|
||||
# Try to set the version metadata.
|
||||
if extension.contents.type == XLA_FFI_Extension_Type.Metadata:
|
||||
@@ -515,15 +528,19 @@ class FfiCallable:
|
||||
metadata_ext.contents.metadata.contents.api_version.major_version = 0
|
||||
metadata_ext.contents.metadata.contents.api_version.minor_version = 1
|
||||
# Turn on CUDA graphs for this handler.
|
||||
if self.graph_compatible:
|
||||
if self.graph_mode is GraphMode.JAX:
|
||||
metadata_ext.contents.metadata.contents.traits = (
|
||||
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
|
||||
)
|
||||
return None
|
||||
|
||||
# retrieve call info
|
||||
attrs = decode_attrs(call_frame.contents.attrs)
|
||||
call_id = int(attrs["call_id"])
|
||||
# NOTE: this assumes that there's only one attribute - call_id (int64).
|
||||
# A more general but slower approach is this:
|
||||
# attrs = decode_attrs(call_frame.contents.attrs)
|
||||
# call_id = int(attrs["call_id"])
|
||||
attr = ctypes.cast(call_frame.contents.attrs.attrs[0], ctypes.POINTER(XLA_FFI_Scalar)).contents
|
||||
call_id = ctypes.cast(attr.value, ctypes.POINTER(ctypes.c_int64)).contents.value
|
||||
call_desc = self.call_descriptors[call_id]
|
||||
|
||||
num_inputs = call_frame.contents.args.size
|
||||
@@ -535,8 +552,35 @@ class FfiCallable:
|
||||
assert num_inputs == self.num_inputs
|
||||
assert num_outputs == self.num_outputs
|
||||
|
||||
device = wp.device_from_jax(get_jax_device())
|
||||
cuda_stream = get_stream_from_callframe(call_frame.contents)
|
||||
|
||||
if self.graph_mode == GraphMode.WARP:
|
||||
# check if we already captured an identical call
|
||||
ip = [inputs[i].contents.data for i in self.array_input_indices]
|
||||
op = [outputs[i].contents.data for i in self.array_output_indices]
|
||||
buffer_hash = hash((*ip, *op))
|
||||
capture = call_desc.captures.get(buffer_hash)
|
||||
|
||||
# launch existing graph
|
||||
if capture is not None:
|
||||
# NOTE: We use the native graph API to avoid overhead with obtaining Stream and Device objects in Python.
|
||||
# This code should match wp.capture_launch().
|
||||
graph = capture.graph
|
||||
if graph.graph_exec is None:
|
||||
g = ctypes.c_void_p()
|
||||
if not wp.context.runtime.core.wp_cuda_graph_create_exec(
|
||||
graph.device.context, cuda_stream, graph.graph, ctypes.byref(g)
|
||||
):
|
||||
raise RuntimeError(f"Graph creation error: {wp.context.runtime.get_error_string()}")
|
||||
graph.graph_exec = g
|
||||
|
||||
if not wp.context.runtime.core.wp_cuda_graph_launch(graph.graph_exec, cuda_stream):
|
||||
raise RuntimeError(f"Graph launch error: {wp.context.runtime.get_error_string()}")
|
||||
|
||||
# early out
|
||||
return
|
||||
|
||||
device = wp.device_from_jax(get_jax_device())
|
||||
stream = wp.Stream(device, cuda_stream=cuda_stream)
|
||||
|
||||
# reconstruct the argument list
|
||||
@@ -564,11 +608,20 @@ 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) as capture:
|
||||
# capturing with JAX
|
||||
with wp.ScopedCapture(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
|
||||
elif self.graph_mode == GraphMode.WARP:
|
||||
# capturing with WARP
|
||||
with wp.ScopedCapture() as capture:
|
||||
self.func(*arg_list)
|
||||
wp.capture_launch(capture.graph)
|
||||
# keep a reference to the capture object and reuse it with same buffers
|
||||
call_desc.captures[buffer_hash] = capture
|
||||
else:
|
||||
# not capturing
|
||||
self.func(*arg_list)
|
||||
|
||||
except Exception as e:
|
||||
@@ -633,7 +686,8 @@ def jax_kernel(
|
||||
def jax_callable(
|
||||
func: Callable,
|
||||
num_outputs: int = 1,
|
||||
graph_compatible: bool = True,
|
||||
graph_compatible: Optional[bool] = None, # deprecated
|
||||
graph_mode: GraphMode = GraphMode.JAX,
|
||||
vmap_method: Optional[str] = "broadcast_all",
|
||||
output_dims=None,
|
||||
in_out_argnames=None,
|
||||
@@ -647,8 +701,15 @@ def jax_callable(
|
||||
Args:
|
||||
func: The Python function to call.
|
||||
num_outputs: Optional. Specify the number of output arguments if greater than 1.
|
||||
This must include the number of ``in_out_arguments``.
|
||||
This must include the number of ``in_out_arguments``.
|
||||
graph_compatible: Optional. Whether the function can be called during CUDA graph capture.
|
||||
This argument is deprecated, use ``graph_mode`` instead.
|
||||
graph_mode: Optional. CUDA graph capture mode.
|
||||
``GraphMode.JAX`` (default): Let JAX capture the graph, which may be used as a subgraph in an enclosing capture.
|
||||
``GraphMode.WARP``: Let Warp capture the graph. Use this mode when the callable cannot be used as a subraph,
|
||||
such as when the callable uses conditional graph nodes.
|
||||
``GraphMode.NONE``: Disable graph capture. Use when the callable performs operations that are not legal in a graph,
|
||||
such as host synchronization.
|
||||
vmap_method: Optional. String specifying how the callback transforms under ``vmap()``.
|
||||
This argument can also be specified for individual calls.
|
||||
output_dims: Optional. Specify the default dimensions of output arrays.
|
||||
@@ -663,17 +724,27 @@ def jax_callable(
|
||||
- There must be at least one output or input-output argument.
|
||||
- Only the CUDA backend is supported.
|
||||
"""
|
||||
|
||||
if graph_compatible is not None:
|
||||
wp.utils.warn(
|
||||
"The `graph_compatible` argument is deprecated, use `graph_mode` instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
if graph_compatible is False:
|
||||
graph_mode = GraphMode.NONE
|
||||
|
||||
key = (
|
||||
func,
|
||||
num_outputs,
|
||||
graph_compatible,
|
||||
graph_mode,
|
||||
vmap_method,
|
||||
tuple(sorted(output_dims.items())) if output_dims else output_dims,
|
||||
)
|
||||
|
||||
with _FFI_REGISTRY_LOCK:
|
||||
if key not in _FFI_CALLABLE_REGISTRY:
|
||||
new_callable = FfiCallable(func, num_outputs, graph_compatible, vmap_method, output_dims, in_out_argnames)
|
||||
new_callable = FfiCallable(func, num_outputs, graph_mode, vmap_method, output_dims, in_out_argnames)
|
||||
_FFI_CALLABLE_REGISTRY[key] = new_callable
|
||||
|
||||
return _FFI_CALLABLE_REGISTRY[key]
|
||||
@@ -704,7 +775,6 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr
|
||||
|
||||
def ffi_callback(call_frame):
|
||||
try:
|
||||
# TODO Try-catch around the body and return XLA_FFI_Error on error.
|
||||
extension = call_frame.contents.extension_start
|
||||
# On the first call, XLA runtime will query the API version and traits
|
||||
# metadata using the |extension| field. Let us respond to that query
|
||||
|
||||
Reference in New Issue
Block a user