tensoracle 1.0.1__tar.gz → 1.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {tensoracle-1.0.1 → tensoracle-1.1.0}/PKG-INFO +67 -2
- {tensoracle-1.0.1 → tensoracle-1.1.0}/README.md +66 -1
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/index.md +66 -1
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_version.py +2 -2
- tensoracle-1.1.0/tests/test_cuda_graphs.py +191 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/.gitignore +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/LICENSE +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/build_backend.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/CMakeLists.txt +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/src/cubin_plugin.cpp +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/tests/core.cpp +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/api/bundle.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/api/index.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/api/producer.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/cpp-python-repositories.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/optimization.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/precision.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/tensorrt-sdk.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/CMakeLists.txt +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/README.md +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.json +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/kernels.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/plugin.cpp +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/pyproject.toml +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/mkdocs.yml +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/pyproject.toml +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/__init__.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_build.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_cache.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_comparison.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_execution.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_memory.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_plugins.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/bundle.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/convert.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/producer.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/build_memory_runtime.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/cubin_io/check.cpp +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/cubin_io/kernels.cu +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/cubin_io/plugin.cpp +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/example_runtime.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/native/CMakeLists.txt +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/native/export.h +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/native_build.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/plugin_io_runtime.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/precision_runtime.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_build.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_build_memory.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_bundle.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_cache.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_cmake_plugin.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_comparison.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_convert.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_cpp_core.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_distribution.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_engine_io.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_example_plugin.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_execution.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_gpu_cache.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_plugin_io.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_plugins.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_precision.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_producer.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_shape_tensor_outputs.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_weight_hash.py +0 -0
- {tensoracle-1.0.1 → tensoracle-1.1.0}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: tensoracle
|
|
3
|
-
Version: 1.0
|
|
3
|
+
Version: 1.1.0
|
|
4
4
|
Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
|
|
5
5
|
Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
|
|
6
6
|
Project-URL: repository, https://github.com/sync-dev-org/tensoracle
|
|
@@ -255,10 +255,75 @@ with handle.create_execution() as execution:
|
|
|
255
255
|
|
|
256
256
|
The caller selects the CUDA device and context, orders inputs written on other streams before enqueue, and keeps buffers, streams, plugin preparation resources, and the execution unit alive until completion. Use one execution unit serially; do not call it concurrently from multiple threads. Confirm completion before changing shapes or streams, reusing buffers, or closing. `close()` and context exit release the TensorRT context before the handle, without waiting, and repeated close calls do nothing. Operations on a closed unit raise `RuntimeError` before touching TensorRT. The unit does not retain caller arrays or expose its TensorRT context.
|
|
257
257
|
|
|
258
|
-
Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported.
|
|
258
|
+
Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. Profile switching and caller-owned TensorRT workspace memory are outside this API.
|
|
259
259
|
|
|
260
260
|
Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
|
|
261
261
|
|
|
262
|
+
### Capture and replay with CUDA Graphs
|
|
263
|
+
|
|
264
|
+
Tensoracle supports recording one `EngineExecution.enqueue` call with CUDA stream capture and replaying the resulting graph. The application owns graph creation, recapture, and destruction; tensoracle provides no capture or replay operations.
|
|
265
|
+
|
|
266
|
+
Use a dedicated execution unit (`handle.create_execution()`) for each graph. After setting input shapes, enqueue once and wait for completion before starting capture, so TensorRT finishes deferred initialization. Capture and replay on a nondefault stream; the example uses a nonblocking stream. Keep synchronous CUDA calls and work on the legacy default stream out of capture.
|
|
267
|
+
|
|
268
|
+
The graph fixes the input shapes, tensor addresses, and execution unit. Keep the execution unit, its handle, buffers, and plugin preparation resources alive until the graph is destroyed and outstanding work has finished. To use different shapes or addresses, capture a new graph for that combination or select another graph with its own execution unit. Input contents may change at the same addresses between completed replays.
|
|
269
|
+
|
|
270
|
+
For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
|
|
271
|
+
|
|
272
|
+
```python
|
|
273
|
+
from contextlib import ExitStack
|
|
274
|
+
|
|
275
|
+
import numpy as np
|
|
276
|
+
import tensorrt as trt
|
|
277
|
+
from cuda.bindings import runtime as cudart
|
|
278
|
+
|
|
279
|
+
with handle.create_execution() as execution:
|
|
280
|
+
execution.set_input_shapes({"features": (1, 4, 2)})
|
|
281
|
+
# Omit set_input_shapes for fixed inputs and single-point profiles.
|
|
282
|
+
features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
|
|
283
|
+
output = np.empty(
|
|
284
|
+
execution.get_tensor_shape("result"),
|
|
285
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
286
|
+
)
|
|
287
|
+
with ExitStack() as resources:
|
|
288
|
+
stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
289
|
+
resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
|
|
290
|
+
source = cuda(cudart.cudaMalloc(features.nbytes))
|
|
291
|
+
resources.callback(lambda: cuda(cudart.cudaFree(source)))
|
|
292
|
+
target = cuda(cudart.cudaMalloc(output.nbytes))
|
|
293
|
+
resources.callback(lambda: cuda(cudart.cudaFree(target)))
|
|
294
|
+
addresses = {"features": int(source), "result": int(target)}
|
|
295
|
+
try:
|
|
296
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
297
|
+
source, features.ctypes.data, features.nbytes,
|
|
298
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
|
|
299
|
+
))
|
|
300
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
301
|
+
cuda(cudart.cudaStreamSynchronize(stream)) # Warm up before capture.
|
|
302
|
+
|
|
303
|
+
cuda(cudart.cudaStreamBeginCapture(
|
|
304
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
|
|
305
|
+
))
|
|
306
|
+
try:
|
|
307
|
+
execution.enqueue(addresses, stream=int(stream)) # Exactly one call.
|
|
308
|
+
finally:
|
|
309
|
+
graph = cuda(cudart.cudaStreamEndCapture(stream))
|
|
310
|
+
resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
|
|
311
|
+
instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
312
|
+
resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
313
|
+
|
|
314
|
+
for _ in range(3):
|
|
315
|
+
cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
316
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
317
|
+
output.ctypes.data, target, output.nbytes,
|
|
318
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
|
|
319
|
+
))
|
|
320
|
+
finally:
|
|
321
|
+
# Wait before ExitStack destroys the graph and releases its buffers.
|
|
322
|
+
cuda(cudart.cudaStreamSynchronize(stream))
|
|
323
|
+
```
|
|
324
|
+
|
|
325
|
+
Engines with loops or conditionals, or layers or plugins that synchronize with the host during inference, cannot be captured. CUDA reports capture failures through `cudaErrorStreamCapture*` errors. Always end an attempted capture and release any graph resources before falling back to ordinary `enqueue`. See [TensorRT's CUDA Graph limitations](https://docs.nvidia.com/deeplearning/tensorrt/latest/performance/optimization.html#inference-with-cuda-graphs).
|
|
326
|
+
|
|
262
327
|
## Build a cubin plugin
|
|
263
328
|
|
|
264
329
|
Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
|
|
@@ -224,10 +224,75 @@ with handle.create_execution() as execution:
|
|
|
224
224
|
|
|
225
225
|
The caller selects the CUDA device and context, orders inputs written on other streams before enqueue, and keeps buffers, streams, plugin preparation resources, and the execution unit alive until completion. Use one execution unit serially; do not call it concurrently from multiple threads. Confirm completion before changing shapes or streams, reusing buffers, or closing. `close()` and context exit release the TensorRT context before the handle, without waiting, and repeated close calls do nothing. Operations on a closed unit raise `RuntimeError` before touching TensorRT. The unit does not retain caller arrays or expose its TensorRT context.
|
|
226
226
|
|
|
227
|
-
Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported.
|
|
227
|
+
Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. Profile switching and caller-owned TensorRT workspace memory are outside this API.
|
|
228
228
|
|
|
229
229
|
Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
|
|
230
230
|
|
|
231
|
+
### Capture and replay with CUDA Graphs
|
|
232
|
+
|
|
233
|
+
Tensoracle supports recording one `EngineExecution.enqueue` call with CUDA stream capture and replaying the resulting graph. The application owns graph creation, recapture, and destruction; tensoracle provides no capture or replay operations.
|
|
234
|
+
|
|
235
|
+
Use a dedicated execution unit (`handle.create_execution()`) for each graph. After setting input shapes, enqueue once and wait for completion before starting capture, so TensorRT finishes deferred initialization. Capture and replay on a nondefault stream; the example uses a nonblocking stream. Keep synchronous CUDA calls and work on the legacy default stream out of capture.
|
|
236
|
+
|
|
237
|
+
The graph fixes the input shapes, tensor addresses, and execution unit. Keep the execution unit, its handle, buffers, and plugin preparation resources alive until the graph is destroyed and outstanding work has finished. To use different shapes or addresses, capture a new graph for that combination or select another graph with its own execution unit. Input contents may change at the same addresses between completed replays.
|
|
238
|
+
|
|
239
|
+
For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
|
|
240
|
+
|
|
241
|
+
```python
|
|
242
|
+
from contextlib import ExitStack
|
|
243
|
+
|
|
244
|
+
import numpy as np
|
|
245
|
+
import tensorrt as trt
|
|
246
|
+
from cuda.bindings import runtime as cudart
|
|
247
|
+
|
|
248
|
+
with handle.create_execution() as execution:
|
|
249
|
+
execution.set_input_shapes({"features": (1, 4, 2)})
|
|
250
|
+
# Omit set_input_shapes for fixed inputs and single-point profiles.
|
|
251
|
+
features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
|
|
252
|
+
output = np.empty(
|
|
253
|
+
execution.get_tensor_shape("result"),
|
|
254
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
255
|
+
)
|
|
256
|
+
with ExitStack() as resources:
|
|
257
|
+
stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
258
|
+
resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
|
|
259
|
+
source = cuda(cudart.cudaMalloc(features.nbytes))
|
|
260
|
+
resources.callback(lambda: cuda(cudart.cudaFree(source)))
|
|
261
|
+
target = cuda(cudart.cudaMalloc(output.nbytes))
|
|
262
|
+
resources.callback(lambda: cuda(cudart.cudaFree(target)))
|
|
263
|
+
addresses = {"features": int(source), "result": int(target)}
|
|
264
|
+
try:
|
|
265
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
266
|
+
source, features.ctypes.data, features.nbytes,
|
|
267
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
|
|
268
|
+
))
|
|
269
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
270
|
+
cuda(cudart.cudaStreamSynchronize(stream)) # Warm up before capture.
|
|
271
|
+
|
|
272
|
+
cuda(cudart.cudaStreamBeginCapture(
|
|
273
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
|
|
274
|
+
))
|
|
275
|
+
try:
|
|
276
|
+
execution.enqueue(addresses, stream=int(stream)) # Exactly one call.
|
|
277
|
+
finally:
|
|
278
|
+
graph = cuda(cudart.cudaStreamEndCapture(stream))
|
|
279
|
+
resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
|
|
280
|
+
instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
281
|
+
resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
282
|
+
|
|
283
|
+
for _ in range(3):
|
|
284
|
+
cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
285
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
286
|
+
output.ctypes.data, target, output.nbytes,
|
|
287
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
|
|
288
|
+
))
|
|
289
|
+
finally:
|
|
290
|
+
# Wait before ExitStack destroys the graph and releases its buffers.
|
|
291
|
+
cuda(cudart.cudaStreamSynchronize(stream))
|
|
292
|
+
```
|
|
293
|
+
|
|
294
|
+
Engines with loops or conditionals, or layers or plugins that synchronize with the host during inference, cannot be captured. CUDA reports capture failures through `cudaErrorStreamCapture*` errors. Always end an attempted capture and release any graph resources before falling back to ordinary `enqueue`. See [TensorRT's CUDA Graph limitations](https://docs.nvidia.com/deeplearning/tensorrt/latest/performance/optimization.html#inference-with-cuda-graphs).
|
|
295
|
+
|
|
231
296
|
## Build a cubin plugin
|
|
232
297
|
|
|
233
298
|
Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
|
|
@@ -155,12 +155,77 @@ with handle.create_execution() as execution:
|
|
|
155
155
|
|
|
156
156
|
The caller selects the CUDA device and context, orders inputs written on other streams before enqueue, and keeps buffers, streams, plugin preparation resources, and the execution unit alive until completion. Use one execution unit serially; do not call it concurrently from multiple threads. Confirm completion before changing shapes or streams, reusing buffers, or closing. `close()` and context exit release the TensorRT context before the handle, without waiting, and repeated close calls do nothing. Operations on a closed unit raise `RuntimeError` before touching TensorRT. The unit does not retain caller arrays or expose its TensorRT context.
|
|
157
157
|
|
|
158
|
-
Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported.
|
|
158
|
+
Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. Profile switching and caller-owned TensorRT workspace memory are outside this API.
|
|
159
159
|
|
|
160
160
|
Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
|
|
161
161
|
|
|
162
162
|
`TensorSpec` accepts optional `location=trt.TensorLocation.DEVICE` and `format=trt.TensorFormat.LINEAR`. `open_engine` checks these only when declared, together with tensor names, modes, dtypes, build-time shapes, and profiles. An `EngineIOError` names the tensor, field, observed value, expected value, and saved file path. The execution unit does not accept a second declaration.
|
|
163
163
|
|
|
164
|
+
#### Capture and replay with CUDA Graphs
|
|
165
|
+
|
|
166
|
+
Tensoracle supports recording one `EngineExecution.enqueue` call with CUDA stream capture and replaying the resulting graph. The application owns graph creation, recapture, and destruction; tensoracle provides no capture or replay operations.
|
|
167
|
+
|
|
168
|
+
Use a dedicated execution unit (`handle.create_execution()`) for each graph. After setting input shapes, enqueue once and wait for completion before starting capture, so TensorRT finishes deferred initialization. Capture and replay on a nondefault stream; the example uses a nonblocking stream. Keep synchronous CUDA calls and work on the legacy default stream out of capture.
|
|
169
|
+
|
|
170
|
+
The graph fixes the input shapes, tensor addresses, and execution unit. Keep the execution unit, its handle, buffers, and plugin preparation resources alive until the graph is destroyed and outstanding work has finished. To use different shapes or addresses, capture a new graph for that combination or select another graph with its own execution unit. Input contents may change at the same addresses between completed replays.
|
|
171
|
+
|
|
172
|
+
For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
|
|
173
|
+
|
|
174
|
+
```python
|
|
175
|
+
from contextlib import ExitStack
|
|
176
|
+
|
|
177
|
+
import numpy as np
|
|
178
|
+
import tensorrt as trt
|
|
179
|
+
from cuda.bindings import runtime as cudart
|
|
180
|
+
|
|
181
|
+
with handle.create_execution() as execution:
|
|
182
|
+
execution.set_input_shapes({"features": (1, 4, 2)})
|
|
183
|
+
# Omit set_input_shapes for fixed inputs and single-point profiles.
|
|
184
|
+
features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
|
|
185
|
+
output = np.empty(
|
|
186
|
+
execution.get_tensor_shape("result"),
|
|
187
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
188
|
+
)
|
|
189
|
+
with ExitStack() as resources:
|
|
190
|
+
stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
191
|
+
resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
|
|
192
|
+
source = cuda(cudart.cudaMalloc(features.nbytes))
|
|
193
|
+
resources.callback(lambda: cuda(cudart.cudaFree(source)))
|
|
194
|
+
target = cuda(cudart.cudaMalloc(output.nbytes))
|
|
195
|
+
resources.callback(lambda: cuda(cudart.cudaFree(target)))
|
|
196
|
+
addresses = {"features": int(source), "result": int(target)}
|
|
197
|
+
try:
|
|
198
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
199
|
+
source, features.ctypes.data, features.nbytes,
|
|
200
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
|
|
201
|
+
))
|
|
202
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
203
|
+
cuda(cudart.cudaStreamSynchronize(stream)) # Warm up before capture.
|
|
204
|
+
|
|
205
|
+
cuda(cudart.cudaStreamBeginCapture(
|
|
206
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
|
|
207
|
+
))
|
|
208
|
+
try:
|
|
209
|
+
execution.enqueue(addresses, stream=int(stream)) # Exactly one call.
|
|
210
|
+
finally:
|
|
211
|
+
graph = cuda(cudart.cudaStreamEndCapture(stream))
|
|
212
|
+
resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
|
|
213
|
+
instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
214
|
+
resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
215
|
+
|
|
216
|
+
for _ in range(3):
|
|
217
|
+
cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
218
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
219
|
+
output.ctypes.data, target, output.nbytes,
|
|
220
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
|
|
221
|
+
))
|
|
222
|
+
finally:
|
|
223
|
+
# Wait before ExitStack destroys the graph and releases its buffers.
|
|
224
|
+
cuda(cudart.cudaStreamSynchronize(stream))
|
|
225
|
+
```
|
|
226
|
+
|
|
227
|
+
Engines with loops or conditionals, or layers or plugins that synchronize with the host during inference, cannot be captured. CUDA reports capture failures through `cudaErrorStreamCapture*` errors. Always end an attempted capture and release any graph resources before falling back to ordinary `enqueue`. See [TensorRT's CUDA Graph limitations](https://docs.nvidia.com/deeplearning/tensorrt/latest/performance/optimization.html#inference-with-cuda-graphs).
|
|
228
|
+
|
|
164
229
|
### Build a standard graph
|
|
165
230
|
|
|
166
231
|
Define `graph(network)` with TensorRT's Python network API and load model weights from safetensors. Describe each graph input with `InputShape.fixed` or `InputShape.ranged`, and pass an `EngineBuild` declaration to `build_serialized_engine`. `EngineBuild` declares input shapes and builder settings; the graph function determines precision through the tensor types it places in the network. The supported precisions are fp32 (with TF32 disabled), fp16, bf16, int8, and fp8. The [bf16, int8, and fp8 guide](precision.md) shows tensor casts, explicit Q/DQ layers, and scales supplied through safetensors; fp8 targets sm89 and sm120. `open_engine` uses the same declaration plus a name, graph version, weight files, and `EngineIO` to save and reuse a checked engine. Use a different `name` or `graph_version` for engines with different tensor types, such as fp32 and fp16. The [graph example](https://github.com/sync-dev-org/tensoracle#define-and-build-a-standard-graph) shows the exact calls.
|
|
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
|
|
|
18
18
|
commit_id: str | None
|
|
19
19
|
__commit_id__: str | None
|
|
20
20
|
|
|
21
|
-
__version__ = version = '1.0
|
|
22
|
-
__version_tuple__ = version_tuple = (1,
|
|
21
|
+
__version__ = version = '1.1.0'
|
|
22
|
+
__version_tuple__ = version_tuple = (1, 1, 0)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -0,0 +1,191 @@
|
|
|
1
|
+
"""Capture one enqueue and replay it with caller-owned CUDA graphs."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from contextlib import ExitStack, contextmanager
|
|
5
|
+
|
|
6
|
+
import numpy as np
|
|
7
|
+
import pytest
|
|
8
|
+
import tensorrt as trt
|
|
9
|
+
from cuda.bindings import runtime as cudart
|
|
10
|
+
|
|
11
|
+
import tensoracle
|
|
12
|
+
from test_build import _check_cuda
|
|
13
|
+
from test_execution import _open_sum
|
|
14
|
+
|
|
15
|
+
pytest_plugins = ("test_example_plugin",)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@contextmanager
|
|
19
|
+
def _captured(handle, batches, *, shapes):
|
|
20
|
+
"""Keep one execution, its addresses, and its captured graph alive together."""
|
|
21
|
+
with handle.create_execution() as execution, ExitStack() as resources:
|
|
22
|
+
execution.set_input_shapes(shapes)
|
|
23
|
+
output = np.empty(
|
|
24
|
+
execution.get_tensor_shape("result"),
|
|
25
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
26
|
+
)
|
|
27
|
+
stream = _check_cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
28
|
+
resources.callback(lambda: _check_cuda(cudart.cudaStreamDestroy(stream)))
|
|
29
|
+
addresses = {}
|
|
30
|
+
for name, value in {**batches[0], "result": output}.items():
|
|
31
|
+
pointer = _check_cuda(cudart.cudaMalloc(value.nbytes))
|
|
32
|
+
resources.callback(lambda pointer=pointer: _check_cuda(cudart.cudaFree(pointer)))
|
|
33
|
+
addresses[name] = int(pointer)
|
|
34
|
+
|
|
35
|
+
def copy_inputs(batch):
|
|
36
|
+
for name, value in batch.items():
|
|
37
|
+
_check_cuda(
|
|
38
|
+
cudart.cudaMemcpyAsync(
|
|
39
|
+
addresses[name],
|
|
40
|
+
value.ctypes.data,
|
|
41
|
+
value.nbytes,
|
|
42
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
|
|
43
|
+
stream,
|
|
44
|
+
)
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
def read_output():
|
|
48
|
+
_check_cuda(
|
|
49
|
+
cudart.cudaMemcpyAsync(
|
|
50
|
+
output.ctypes.data,
|
|
51
|
+
addresses["result"],
|
|
52
|
+
output.nbytes,
|
|
53
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
|
|
54
|
+
stream,
|
|
55
|
+
)
|
|
56
|
+
)
|
|
57
|
+
_check_cuda(cudart.cudaStreamSynchronize(stream))
|
|
58
|
+
return output.copy()
|
|
59
|
+
|
|
60
|
+
try:
|
|
61
|
+
# The first ordinary enqueue flushes deferred shape and plugin initialization.
|
|
62
|
+
ordinary = []
|
|
63
|
+
for batch in batches:
|
|
64
|
+
copy_inputs(batch)
|
|
65
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
66
|
+
ordinary.append(read_output())
|
|
67
|
+
|
|
68
|
+
_check_cuda(
|
|
69
|
+
cudart.cudaStreamBeginCapture(
|
|
70
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal
|
|
71
|
+
)
|
|
72
|
+
)
|
|
73
|
+
try:
|
|
74
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
75
|
+
finally:
|
|
76
|
+
graph = _check_cuda(cudart.cudaStreamEndCapture(stream))
|
|
77
|
+
resources.callback(lambda: _check_cuda(cudart.cudaGraphDestroy(graph)))
|
|
78
|
+
instance = _check_cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
79
|
+
resources.callback(lambda: _check_cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
80
|
+
|
|
81
|
+
def replay(index):
|
|
82
|
+
copy_inputs(batches[index])
|
|
83
|
+
# An empty graph or a launch writing to the wrong address must fail.
|
|
84
|
+
_check_cuda(
|
|
85
|
+
cudart.cudaMemsetAsync(addresses["result"], 0xFF, output.nbytes, stream)
|
|
86
|
+
)
|
|
87
|
+
_check_cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
88
|
+
actual = read_output()
|
|
89
|
+
np.testing.assert_array_equal(actual, ordinary[index])
|
|
90
|
+
return actual
|
|
91
|
+
|
|
92
|
+
yield replay
|
|
93
|
+
finally:
|
|
94
|
+
# Finish work before graphs, buffers, streams, and the execution are released.
|
|
95
|
+
_check_cuda(cudart.cudaStreamSynchronize(stream))
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def _sum_batches(rows):
|
|
99
|
+
return [
|
|
100
|
+
{
|
|
101
|
+
"features": np.arange(rows * 2, dtype=np.float32).reshape(rows, 2) / 4 + offset,
|
|
102
|
+
"extra": np.full((rows, 2), -0.25 - offset / 2, dtype=np.float32),
|
|
103
|
+
}
|
|
104
|
+
for offset in (0, 2)
|
|
105
|
+
]
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
@pytest.mark.gpu
|
|
109
|
+
@pytest.mark.req("SP-TENSORACLE-028")
|
|
110
|
+
def test_standard_engine_graph_replay_matches_ordinary_enqueue(tmp_path, monkeypatch):
|
|
111
|
+
"""One captured enqueue on a nondefault stream reproduces ordinary sums for new input data."""
|
|
112
|
+
handle = _open_sum(tmp_path, monkeypatch, fixed=True)
|
|
113
|
+
batches = _sum_batches(3)
|
|
114
|
+
with _captured(handle, batches, shapes={}) as replay:
|
|
115
|
+
for index in (0, 1, 0):
|
|
116
|
+
np.testing.assert_array_equal(
|
|
117
|
+
replay(index), batches[index]["features"] + batches[index]["extra"]
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
@pytest.mark.gpu
|
|
122
|
+
@pytest.mark.req("SP-TENSORACLE-028")
|
|
123
|
+
def test_plugin_engine_graph_replay_matches_ordinary_enqueue(installed_example, tmp_path):
|
|
124
|
+
"""A captured cubin plugin runs repeatedly and matches ordinary enqueue with new inputs."""
|
|
125
|
+
library = installed_example["library"]
|
|
126
|
+
manifest = json.loads((library.parent / "bundle/bundle.json").read_text())
|
|
127
|
+
digest = np.frombuffer(manifest["bundle_sha256"].encode(), dtype=np.uint8)
|
|
128
|
+
bias = np.arange(16, dtype=np.float32) / 8
|
|
129
|
+
|
|
130
|
+
def graph(network):
|
|
131
|
+
creator = trt.get_plugin_registry().get_creator(
|
|
132
|
+
"TensoracleExampleBias", "1", "tensoracle.example"
|
|
133
|
+
)
|
|
134
|
+
assert creator is not None
|
|
135
|
+
fields = trt.PluginFieldCollection(
|
|
136
|
+
[
|
|
137
|
+
trt.PluginField("bundle_sha256", digest, trt.PluginFieldType.CHAR),
|
|
138
|
+
trt.PluginField("bias", bias, trt.PluginFieldType.FLOAT32),
|
|
139
|
+
]
|
|
140
|
+
)
|
|
141
|
+
plugin = creator.create_plugin("bias", fields, trt.TensorRTPhase.BUILD)
|
|
142
|
+
assert plugin is not None
|
|
143
|
+
features = network.add_input("features", trt.float32, (3, 16))
|
|
144
|
+
result = network.add_plugin_v3([features], [], plugin).get_output(0)
|
|
145
|
+
result.name = "result"
|
|
146
|
+
network.mark_output(result)
|
|
147
|
+
|
|
148
|
+
plan = tensoracle.build_serialized_engine(
|
|
149
|
+
graph,
|
|
150
|
+
tensoracle.EngineBuild(
|
|
151
|
+
{"features": tensoracle.InputShape.fixed((3, 16))},
|
|
152
|
+
tensoracle.BuilderSettings(optimization_level=0),
|
|
153
|
+
),
|
|
154
|
+
plugins=(library,),
|
|
155
|
+
)
|
|
156
|
+
logger = trt.Logger(trt.Logger.ERROR)
|
|
157
|
+
runtime = trt.Runtime(logger)
|
|
158
|
+
engine = runtime.deserialize_cuda_engine(plan)
|
|
159
|
+
assert engine is not None
|
|
160
|
+
handle = tensoracle.EngineHandle(engine, runtime, tmp_path / "bias.engine")
|
|
161
|
+
batches = [
|
|
162
|
+
{"features": np.arange(48, dtype=np.float32).reshape(3, 16) / 4 + offset}
|
|
163
|
+
for offset in (0, 2)
|
|
164
|
+
]
|
|
165
|
+
with _captured(handle, batches, shapes={}) as replay:
|
|
166
|
+
for index in (0, 1, 0):
|
|
167
|
+
np.testing.assert_array_equal(replay(index), batches[index]["features"] + bias)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
@pytest.mark.gpu
|
|
171
|
+
@pytest.mark.req("SP-TENSORACLE-028")
|
|
172
|
+
def test_ranged_engine_keeps_a_graph_and_execution_per_shape(tmp_path, monkeypatch):
|
|
173
|
+
"""Two shapes keep dedicated executions and graphs that replay correctly in alternation."""
|
|
174
|
+
handle = _open_sum(tmp_path, monkeypatch)
|
|
175
|
+
batches = [_sum_batches(rows) for rows in (2, 5)]
|
|
176
|
+
with ExitStack() as resources:
|
|
177
|
+
replays = [
|
|
178
|
+
resources.enter_context(
|
|
179
|
+
_captured(
|
|
180
|
+
handle,
|
|
181
|
+
values,
|
|
182
|
+
shapes={name: value.shape for name, value in values[0].items()},
|
|
183
|
+
)
|
|
184
|
+
)
|
|
185
|
+
for values in batches
|
|
186
|
+
]
|
|
187
|
+
for index in (0, 1, 0):
|
|
188
|
+
for replay, values in zip(replays, batches, strict=True):
|
|
189
|
+
np.testing.assert_array_equal(
|
|
190
|
+
replay(index), values[index]["features"] + values[index]["extra"]
|
|
191
|
+
)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py
RENAMED
|
File without changes
|
{tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|