tensoracle 0.3.1__tar.gz → 0.4.1__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.
Files changed (74) hide show
  1. {tensoracle-0.3.1 → tensoracle-0.4.1}/PKG-INFO +21 -2
  2. {tensoracle-0.3.1 → tensoracle-0.4.1}/README.md +20 -1
  3. {tensoracle-0.3.1 → tensoracle-0.4.1}/cpp/include/tensoracle/cubin_plugin.hpp +24 -3
  4. {tensoracle-0.3.1 → tensoracle-0.4.1}/cpp/src/cubin_plugin.cpp +66 -15
  5. {tensoracle-0.3.1 → tensoracle-0.4.1}/docs/index.md +20 -1
  6. {tensoracle-0.3.1 → tensoracle-0.4.1}/docs/optimization.md +1 -1
  7. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/plugin.cpp +8 -6
  8. {tensoracle-0.3.1 → tensoracle-0.4.1}/pyproject.toml +11 -2
  9. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/__init__.py +9 -1
  10. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/_build.py +35 -1
  11. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/_cache.py +30 -8
  12. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/_execution.py +3 -0
  13. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/_version.py +2 -2
  14. tensoracle-0.4.1/tests/cubin_io/check.cpp +158 -0
  15. tensoracle-0.4.1/tests/cubin_io/kernels.cu +11 -0
  16. tensoracle-0.4.1/tests/cubin_io/plugin.cpp +106 -0
  17. tensoracle-0.4.1/tests/plugin_io_runtime.py +35 -0
  18. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_build.py +7 -6
  19. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_build_memory.py +3 -3
  20. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_bundle.py +6 -6
  21. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_cache.py +18 -17
  22. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_comparison.py +8 -8
  23. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_convert.py +9 -9
  24. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_cpp_core.py +2 -2
  25. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_distribution.py +25 -16
  26. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_engine_io.py +3 -3
  27. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_example_plugin.py +17 -15
  28. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_execution.py +20 -20
  29. tensoracle-0.4.1/tests/test_plugin_io.py +230 -0
  30. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_plugins.py +6 -6
  31. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_producer.py +7 -7
  32. tensoracle-0.4.1/tests/test_shape_tensor_outputs.py +282 -0
  33. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/test_weight_hash.py +10 -10
  34. {tensoracle-0.3.1 → tensoracle-0.4.1}/uv.lock +14 -0
  35. tensoracle-0.3.1/tests/test_public_mirror_spec.py +0 -324
  36. {tensoracle-0.3.1 → tensoracle-0.4.1}/.gitignore +0 -0
  37. {tensoracle-0.3.1 → tensoracle-0.4.1}/LICENSE +0 -0
  38. {tensoracle-0.3.1 → tensoracle-0.4.1}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
  39. {tensoracle-0.3.1 → tensoracle-0.4.1}/docs/api/bundle.md +0 -0
  40. {tensoracle-0.3.1 → tensoracle-0.4.1}/docs/api/index.md +0 -0
  41. {tensoracle-0.3.1 → tensoracle-0.4.1}/docs/api/producer.md +0 -0
  42. {tensoracle-0.3.1 → tensoracle-0.4.1}/docs/tensorrt-sdk.md +0 -0
  43. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/CMakeLists.txt +0 -0
  44. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/README.md +0 -0
  45. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
  46. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/bundle.json +0 -0
  47. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
  48. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
  49. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
  50. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
  51. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
  52. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
  53. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
  54. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
  55. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
  56. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
  57. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
  58. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
  59. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
  60. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
  61. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
  62. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
  63. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
  64. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/kernels.py +0 -0
  65. {tensoracle-0.3.1 → tensoracle-0.4.1}/examples/cubin_plugin/pyproject.toml +0 -0
  66. {tensoracle-0.3.1 → tensoracle-0.4.1}/mkdocs.yml +0 -0
  67. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/_comparison.py +0 -0
  68. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/_memory.py +0 -0
  69. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/_plugins.py +0 -0
  70. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/bundle.py +0 -0
  71. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/convert.py +0 -0
  72. {tensoracle-0.3.1 → tensoracle-0.4.1}/src/tensoracle/producer.py +0 -0
  73. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/build_memory_runtime.py +0 -0
  74. {tensoracle-0.3.1 → tensoracle-0.4.1}/tests/example_runtime.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: tensoracle
3
- Version: 0.3.1
3
+ Version: 0.4.1
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
@@ -138,7 +138,7 @@ handle = open_engine(
138
138
  engine = handle.engine # Keep handle.runtime alive while using the engine.
139
139
  ```
140
140
 
141
- Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
141
+ Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Include in `graph_version` any other caller value that changes the engine contents but is not part of the cache identity, such as a plugin manifest or a build recipe. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
142
142
 
143
143
  `plugins` names native TensorRT plugin libraries. Both `open_engine` and `build_serialized_engine` accept it. Each function checks that every library exports `getCreators`, loads it through TensorRT's registry, and retains it for the process before building the graph or deserializing an engine. Repeating the same library does not register it twice. `open_engine` also hashes the loaded library bytes into the engine identity; a changed library selects a different saved engine in a new process. A missing, unreadable, or incompatible library, or a library changed after loading in the same process, raises `PluginLibraryError` before a cached plan is opened. Keep the library and its cubin bundle together in the caller's wheel.
144
144
 
@@ -253,6 +253,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
253
253
 
254
254
  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. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
255
255
 
256
+ 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.
257
+
256
258
  ## Build a cubin plugin
257
259
 
258
260
  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.
@@ -267,6 +269,23 @@ pip install 'tensoracle[producer]'
267
269
 
268
270
  The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
269
271
 
272
+ `tensoracle::cubin::PluginBase` accepts a `PluginIO` declaration as its fourth constructor argument. `inputs` lists one `TensorSpec` per input in network order; each spec declares a TensorRT `type` and `format`. `output` declares the single output's type and format, and `shape_input_index` selects the input whose shape the output copies. For example, declare two fp16 / HWC8 inputs and one fp16 / HWC8 output with the shape of input 0:
273
+
274
+ ```cpp
275
+ #include "tensoracle/cubin_plugin.hpp"
276
+
277
+ tensoracle::cubin::PluginIO residual_io()
278
+ {
279
+ using namespace tensoracle::cubin;
280
+ TensorSpec half{nvinfer1::DataType::kHALF, nvinfer1::TensorFormat::kHWC8};
281
+ return PluginIO{{half, half}, half, 0};
282
+ }
283
+ ```
284
+
285
+ Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in the derived plugin's constructor. For one fp16 / HWC8 input, use `PluginIO{{half}, half, 0}`. The base class checks tensor counts, types, and formats during format negotiation, configuration, and shape changes. Implement `validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count, nvinfer1::Dims const& output) const noexcept` to check the dimensions of every input in declaration order and the output; configuration may include wildcard dimensions (`-1`). For a residual operation, this callback can reject different shapes for its two inputs. The caller's kernel must handle the declared storage format. Omitting the fourth argument keeps one fp32 / LINEAR input and output with input 0's shape, so `PluginBase(name, plugin_namespace, fields)` remains valid. An empty input list or an invalid `shape_input_index` throws `std::invalid_argument` at construction.
286
+
287
+ For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
288
+
270
289
  `tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records Triton, ptxas, and kernel source digests without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
271
290
 
272
291
  ## Tested environments
@@ -108,7 +108,7 @@ handle = open_engine(
108
108
  engine = handle.engine # Keep handle.runtime alive while using the engine.
109
109
  ```
110
110
 
111
- Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
111
+ Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, precision, input shapes, builder settings, GPU product name, OS, CPU architecture, and TensorRT version. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Include in `graph_version` any other caller value that changes the engine contents but is not part of the cache identity, such as a plugin manifest or a build recipe. Pass `force_rebuild=True` to replace a cached engine even when it is intact.
112
112
 
113
113
  `plugins` names native TensorRT plugin libraries. Both `open_engine` and `build_serialized_engine` accept it. Each function checks that every library exports `getCreators`, loads it through TensorRT's registry, and retains it for the process before building the graph or deserializing an engine. Repeating the same library does not register it twice. `open_engine` also hashes the loaded library bytes into the engine identity; a changed library selects a different saved engine in a new process. A missing, unreadable, or incompatible library, or a library changed after loading in the same process, raises `PluginLibraryError` before a cached plan is opened. Keep the library and its cubin bundle together in the caller's wheel.
114
114
 
@@ -223,6 +223,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
223
223
 
224
224
  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. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
225
225
 
226
+ 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.
227
+
226
228
  ## Build a cubin plugin
227
229
 
228
230
  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.
@@ -237,6 +239,23 @@ pip install 'tensoracle[producer]'
237
239
 
238
240
  The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
239
241
 
242
+ `tensoracle::cubin::PluginBase` accepts a `PluginIO` declaration as its fourth constructor argument. `inputs` lists one `TensorSpec` per input in network order; each spec declares a TensorRT `type` and `format`. `output` declares the single output's type and format, and `shape_input_index` selects the input whose shape the output copies. For example, declare two fp16 / HWC8 inputs and one fp16 / HWC8 output with the shape of input 0:
243
+
244
+ ```cpp
245
+ #include "tensoracle/cubin_plugin.hpp"
246
+
247
+ tensoracle::cubin::PluginIO residual_io()
248
+ {
249
+ using namespace tensoracle::cubin;
250
+ TensorSpec half{nvinfer1::DataType::kHALF, nvinfer1::TensorFormat::kHWC8};
251
+ return PluginIO{{half, half}, half, 0};
252
+ }
253
+ ```
254
+
255
+ Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in the derived plugin's constructor. For one fp16 / HWC8 input, use `PluginIO{{half}, half, 0}`. The base class checks tensor counts, types, and formats during format negotiation, configuration, and shape changes. Implement `validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count, nvinfer1::Dims const& output) const noexcept` to check the dimensions of every input in declaration order and the output; configuration may include wildcard dimensions (`-1`). For a residual operation, this callback can reject different shapes for its two inputs. The caller's kernel must handle the declared storage format. Omitting the fourth argument keeps one fp32 / LINEAR input and output with input 0's shape, so `PluginBase(name, plugin_namespace, fields)` remains valid. An empty input list or an invalid `shape_input_index` throws `std::invalid_argument` at construction.
256
+
257
+ For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
258
+
240
259
  `tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records Triton, ptxas, and kernel source digests without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
241
260
 
242
261
  ## Tested environments
@@ -41,13 +41,31 @@ private:
41
41
  nvinfer1::PluginFieldCollection collection_{};
42
42
  };
43
43
 
44
- // Implements TensorRT V3 lifecycle and the format/serialization boilerplate.
44
+ struct TensorSpec {
45
+ nvinfer1::DataType type{nvinfer1::DataType::kFLOAT};
46
+ nvinfer1::TensorFormat format{nvinfer1::TensorFormat::kLINEAR};
47
+ };
48
+
49
+ // Declares the input count, each tensor's type and format, and the input whose
50
+ // shape is copied to the single output. For configurations this cannot express
51
+ // (multiple outputs or computed output shapes), developers can override virtual
52
+ // V3 methods or ask tensoracle's maintainers to extend the declaration.
53
+ struct PluginIO {
54
+ std::vector<TensorSpec> inputs{TensorSpec{}};
55
+ TensorSpec output{};
56
+ std::int32_t shape_input_index{0};
57
+ };
58
+
59
+ // Implements TensorRT V3 lifecycle and format/serialization boilerplate using
60
+ // a PluginIO declaration. For other configurations, developers can override
61
+ // virtual V3 methods or ask tensoracle's maintainers to extend the declaration.
45
62
  class PluginBase : public nvinfer1::IPluginV3,
46
63
  public nvinfer1::IPluginV3OneCore,
47
64
  public nvinfer1::IPluginV3OneBuild,
48
65
  public nvinfer1::IPluginV3OneRuntime {
49
66
  public:
50
- PluginBase(char const* name, char const* plugin_namespace, Fields fields);
67
+ // Throws std::invalid_argument for no inputs or an out-of-range shape input.
68
+ PluginBase(char const* name, char const* plugin_namespace, Fields fields, PluginIO io = {});
51
69
  ~PluginBase() noexcept override = default;
52
70
 
53
71
  nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override;
@@ -79,13 +97,16 @@ public:
79
97
 
80
98
  protected:
81
99
  Fields const& fields() const noexcept { return fields_; }
82
- virtual std::int32_t validate_shape(nvinfer1::Dims const& input,
100
+ // Inputs are in declaration order; the array is valid only during this call.
101
+ // configurePlugin may provide wildcard (-1) dimensions.
102
+ virtual std::int32_t validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count,
83
103
  nvinfer1::Dims const& output) const noexcept = 0;
84
104
 
85
105
  private:
86
106
  char const* name_;
87
107
  char const* namespace_;
88
108
  Fields fields_;
109
+ PluginIO io_;
89
110
  };
90
111
 
91
112
  enum class ArgumentKind : std::uint8_t { kPointer, kInt32, kFloat32 };
@@ -4,11 +4,48 @@
4
4
  #include <array>
5
5
  #include <cstring>
6
6
  #include <limits>
7
+ #include <stdexcept>
7
8
  #include <utility>
8
9
 
9
10
  namespace tensoracle::cubin {
10
11
  namespace {
11
12
 
13
+ nvinfer1::PluginTensorDesc const& tensor_desc(nvinfer1::PluginTensorDesc const& value) noexcept
14
+ {
15
+ return value;
16
+ }
17
+
18
+ nvinfer1::PluginTensorDesc const& tensor_desc(nvinfer1::DynamicPluginTensorDesc const& value) noexcept
19
+ {
20
+ return value.desc;
21
+ }
22
+
23
+ bool matches_tensor(nvinfer1::PluginTensorDesc const& descriptor, TensorSpec const& spec) noexcept
24
+ {
25
+ return descriptor.type == spec.type && descriptor.format == spec.format;
26
+ }
27
+
28
+ template <typename Descriptor, typename Validate>
29
+ std::int32_t validate_tensors(Descriptor const* inputs, std::int32_t input_count,
30
+ Descriptor const* outputs, std::int32_t output_count, PluginIO const& io,
31
+ Validate validate) noexcept
32
+ {
33
+ if (inputs == nullptr || outputs == nullptr || input_count < 0 || output_count != 1 ||
34
+ static_cast<std::size_t>(input_count) != io.inputs.size()) return -1;
35
+ for (std::int32_t index = 0; index < input_count; ++index)
36
+ if (!matches_tensor(tensor_desc(inputs[index]), io.inputs[index])) return -1;
37
+ if (!matches_tensor(tensor_desc(outputs[0]), io.output)) return -1;
38
+ try {
39
+ std::vector<nvinfer1::Dims> dims;
40
+ dims.reserve(io.inputs.size());
41
+ for (std::int32_t index = 0; index < input_count; ++index)
42
+ dims.push_back(tensor_desc(inputs[index]).dims);
43
+ return validate(dims.data(), input_count, tensor_desc(outputs[0]).dims);
44
+ } catch (...) {
45
+ return -1;
46
+ }
47
+ }
48
+
12
49
  std::size_t field_size(nvinfer1::PluginFieldType type) noexcept
13
50
  {
14
51
  using nvinfer1::PluginFieldType;
@@ -102,8 +139,14 @@ void Fields::rebuild()
102
139
 
103
140
  nvinfer1::PluginFieldCollection const* Fields::collection() noexcept { return &collection_; }
104
141
 
105
- PluginBase::PluginBase(char const* name, char const* plugin_namespace, Fields fields)
106
- : name_(name), namespace_(plugin_namespace), fields_(std::move(fields)) {}
142
+ PluginBase::PluginBase(char const* name, char const* plugin_namespace, Fields fields, PluginIO io)
143
+ : name_(name), namespace_(plugin_namespace), fields_(std::move(fields)), io_(std::move(io))
144
+ {
145
+ if (io_.inputs.empty() || io_.inputs.size() > static_cast<std::size_t>(
146
+ std::numeric_limits<std::int32_t>::max()) || io_.shape_input_index < 0 ||
147
+ static_cast<std::size_t>(io_.shape_input_index) >= io_.inputs.size())
148
+ throw std::invalid_argument("PluginIO requires inputs and a valid shape_input_index");
149
+ }
107
150
 
108
151
  nvinfer1::IPluginCapability* PluginBase::getCapabilityInterface(
109
152
  nvinfer1::PluginCapabilityType type) noexcept
@@ -124,20 +167,23 @@ bool PluginBase::supportsFormatCombination(std::int32_t position,
124
167
  nvinfer1::DynamicPluginTensorDesc const* in_out,
125
168
  std::int32_t input_count, std::int32_t output_count) noexcept
126
169
  {
127
- if (in_out == nullptr || input_count != 1 || output_count != 1 || position < 0 || position > 1)
170
+ if (in_out == nullptr || input_count < 0 ||
171
+ static_cast<std::size_t>(input_count) != io_.inputs.size() ||
172
+ output_count != 1 || position < 0 || position > input_count)
128
173
  return false;
129
- auto const& descriptor = in_out[position].desc;
130
- return descriptor.type == nvinfer1::DataType::kFLOAT &&
131
- descriptor.format == nvinfer1::TensorFormat::kLINEAR;
174
+ auto const& spec = position < input_count ? io_.inputs[position] : io_.output;
175
+ return matches_tensor(in_out[position].desc, spec);
132
176
  }
133
177
 
134
178
  std::int32_t PluginBase::getOutputDataTypes(nvinfer1::DataType* output_types,
135
179
  std::int32_t output_count, nvinfer1::DataType const* input_types,
136
180
  std::int32_t input_count) const noexcept
137
181
  {
138
- if (output_types == nullptr || input_types == nullptr || output_count != 1 || input_count != 1 ||
139
- input_types[0] != nvinfer1::DataType::kFLOAT) return -1;
140
- output_types[0] = nvinfer1::DataType::kFLOAT;
182
+ if (output_types == nullptr || input_types == nullptr || output_count != 1 || input_count < 0 ||
183
+ static_cast<std::size_t>(input_count) != io_.inputs.size()) return -1;
184
+ for (std::int32_t index = 0; index < input_count; ++index)
185
+ if (input_types[index] != io_.inputs[index].type) return -1;
186
+ output_types[0] = io_.output.type;
141
187
  return 0;
142
188
  }
143
189
 
@@ -146,9 +192,10 @@ std::int32_t PluginBase::getOutputShapes(nvinfer1::DimsExprs const* inputs,
146
192
  nvinfer1::DimsExprs* outputs, std::int32_t output_count,
147
193
  nvinfer1::IExprBuilder&) noexcept
148
194
  {
149
- if (inputs == nullptr || outputs == nullptr || input_count != 1 ||
195
+ if (inputs == nullptr || outputs == nullptr || input_count < 0 ||
196
+ static_cast<std::size_t>(input_count) != io_.inputs.size() ||
150
197
  shape_input_count != 0 || output_count != 1) return -1;
151
- outputs[0] = inputs[0];
198
+ outputs[0] = inputs[io_.shape_input_index];
152
199
  return 0;
153
200
  }
154
201
 
@@ -156,8 +203,10 @@ std::int32_t PluginBase::configurePlugin(nvinfer1::DynamicPluginTensorDesc const
156
203
  std::int32_t input_count, nvinfer1::DynamicPluginTensorDesc const* outputs,
157
204
  std::int32_t output_count) noexcept
158
205
  {
159
- if (inputs == nullptr || outputs == nullptr || input_count != 1 || output_count != 1) return -1;
160
- return validate_shape(inputs[0].desc.dims, outputs[0].desc.dims);
206
+ return validate_tensors(inputs, input_count, outputs, output_count, io_,
207
+ [this](nvinfer1::Dims const* dims, std::int32_t count, nvinfer1::Dims const& output) {
208
+ return validate_shape(dims, count, output);
209
+ });
161
210
  }
162
211
 
163
212
  std::size_t PluginBase::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const*, std::int32_t,
@@ -169,8 +218,10 @@ std::int32_t PluginBase::onShapeChange(nvinfer1::PluginTensorDesc const* inputs,
169
218
  std::int32_t input_count, nvinfer1::PluginTensorDesc const* outputs,
170
219
  std::int32_t output_count) noexcept
171
220
  {
172
- if (inputs == nullptr || outputs == nullptr || input_count != 1 || output_count != 1) return -1;
173
- return validate_shape(inputs[0].dims, outputs[0].dims);
221
+ return validate_tensors(inputs, input_count, outputs, output_count, io_,
222
+ [this](nvinfer1::Dims const* dims, std::int32_t count, nvinfer1::Dims const& output) {
223
+ return validate_shape(dims, count, output);
224
+ });
174
225
  }
175
226
 
176
227
  nvinfer1::IPluginV3* PluginBase::attachToContext(nvinfer1::IPluginResourceContext*) noexcept
@@ -70,7 +70,7 @@ handle = open_engine(
70
70
  )
71
71
  ```
72
72
 
73
- `my_model` represents the application's model package. On a valid cache hit, `open_engine` does not call `graph`; on a miss, it builds and saves a new engine. `open_engine` itself does not execute the engine or manage GPU buffers. Create an execution unit from its returned handle to enqueue inference. Saved engines live under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle` when `XDG_CACHE_HOME` is unset. See the [complete cache example](https://github.com/sync-dev-org/tensoracle#save-and-reopen-an-engine) for the declaration and expected input and output types.
73
+ `my_model` represents the application's model package. On a valid cache hit, `open_engine` does not call `graph`; on a miss, it builds and saves a new engine. Raise `graph_version` when the graph code changes or when any other caller value that changes the engine contents but is not part of the cache identity changes, such as a plugin manifest or a build recipe. `open_engine` itself does not execute the engine or manage GPU buffers. Create an execution unit from its returned handle to enqueue inference. Saved engines live under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle` when `XDG_CACHE_HOME` is unset. See the [complete cache example](https://github.com/sync-dev-org/tensoracle#save-and-reopen-an-engine) for the declaration and expected input and output types.
74
74
 
75
75
  ### Execute an engine
76
76
 
@@ -155,6 +155,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
155
155
 
156
156
  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. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
157
157
 
158
+ 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.
159
+
158
160
  `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.
159
161
 
160
162
  ### Build a standard graph
@@ -173,6 +175,23 @@ pip install 'tensoracle[producer]'
173
175
 
174
176
  `tensoracle.producer.compile_bundle` uses the extra's Triton 3.8.0, `ptxas` 12.9.86, and `cuobjdump` 13.3 to produce cubins for sm86, sm89, and sm120. It finds the NVIDIA wheel executables by default and accepts explicit paths as overrides. Cubin production does not need a GPU or import torch; run it without importing torch first. Build the native plugin with a TensorRT 11 SDK, CUDA toolkit, CMake, and a C++17 compiler. Inspect the resulting bundle or wheel with `tensoracle.bundle`. The [cubin plugin example](https://github.com/sync-dev-org/tensoracle/tree/main/examples/cubin_plugin) covers rebuilding from a repository clone, packaging, and runtime use. Normal import, engine building, and engine reuse do not import the producer or Triton, and do not need the producer extra or native build tools.
175
177
 
178
+ Pass a `tensoracle::cubin::PluginIO` declaration as the fourth argument to `PluginBase(name, plugin_namespace, fields, io)`. The `inputs` vector determines the input count and declares each input's TensorRT `type` and `format` through `TensorSpec`; `output` declares the single output's type and format. `shape_input_index` selects the input whose shape the output copies. This example declares two fp16 / HWC8 inputs and an fp16 / HWC8 output with input 0's shape:
179
+
180
+ ```cpp
181
+ #include "tensoracle/cubin_plugin.hpp"
182
+
183
+ tensoracle::cubin::PluginIO residual_io()
184
+ {
185
+ using namespace tensoracle::cubin;
186
+ TensorSpec half{nvinfer1::DataType::kHALF, nvinfer1::TensorFormat::kHWC8};
187
+ return PluginIO{{half, half}, half, 0};
188
+ }
189
+ ```
190
+
191
+ Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in the derived plugin's constructor; use `PluginIO{{half}, half, 0}` for one fp16 / HWC8 input. The common implementation negotiates the declared types and formats, sets the output type and shape, and checks counts, types, and formats during configuration and shape changes. Implement `validate_shape(nvinfer1::Dims const* inputs, std::int32_t input_count, nvinfer1::Dims const& output) const noexcept` to check all input dimensions in declaration order and the output dimensions, including wildcard dimensions (`-1`) during configuration. A residual plugin can reject mismatched input shapes in this callback. The kernel must handle the declared storage format. Omitting `io` keeps one fp32 / LINEAR input and output with input 0's shape; the original three-argument constructor call remains valid. An empty input list or an out-of-range `shape_input_index` throws `std::invalid_argument` at construction.
192
+
193
+ For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
194
+
176
195
  ## API reference
177
196
 
178
197
  - [Engine building, cache reuse, execution, and output comparison](api/index.md)
@@ -18,7 +18,7 @@ Use the same representative inputs, shapes, and corresponding outputs throughout
18
18
 
19
19
  First check whether a combination of standard TensorRT layers can preserve the operation's meaning, shapes, control flow, and output quality. For example, a deformable convolution may be expressible through grid sampling, mask multiplication, and convolution. If a plugin bundled with TensorRT matches the required behavior, compare that option too. Use a custom cubin plugin only when the suitable alternatives still leave a bottleneck.
20
20
 
21
- - **tensoracle tools:** `tensoracle.producer.compile_bundle` generates cubins and a manifest for supported GPU generations. The shared C++ plugin helpers support a caller-built native plugin, and `tensoracle.bundle.inspect_bundle` or `inspect_wheel` checks the cubins and their match with the plugin. Pass the plugin library through `plugins=` to `build_serialized_engine` or `open_engine`. Follow [Build a cubin plugin](https://github.com/sync-dev-org/tensoracle#build-a-cubin-plugin) and the [cubin plugin example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md) for the generation, build, and inspection operations.
21
+ - **tensoracle tools:** `tensoracle.producer.compile_bundle` generates cubins and a manifest for supported GPU generations. The shared C++ plugin helpers support a caller-built native plugin with declared tensor types, formats, and an input whose shape the output copies; `tensoracle.bundle.inspect_bundle` or `inspect_wheel` checks the cubins and their match with the plugin. Pass the plugin library through `plugins=` to `build_serialized_engine` or `open_engine`. Follow [Build a cubin plugin](index.md#build-a-cubin-plugin) for the `PluginIO` declaration, including fp16 / HWC8 and two-input examples, and the [cubin plugin example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md) for the generation, build, and inspection operations.
22
22
  - **You provide:** Identify the bottleneck, write the kernel and model-specific plugin behavior, choose its supported shapes and types, build and package its native library, and run the resulting engines. An fp32 kernel must not use TF32.
23
23
  - **Gate:** For a custom cubin plugin, the bundle and plugin inspections must pass. Reapply stage 1 to the replacement's fp32 engine against the same fixed PyTorch reference. If using fp16, reapply stage 2 against an accepted fp32 engine with the same limits and model-specific measurements. Advance only when the replacement meets those quality gates; otherwise retain the accepted implementation.
24
24
 
@@ -117,7 +117,7 @@ public:
117
117
  {
118
118
  if (!valid_ || input_descriptors == nullptr || output_descriptors == nullptr ||
119
119
  inputs == nullptr || outputs == nullptr || inputs[0] == nullptr || outputs[0] == nullptr ||
120
- validate_shape(input_descriptors[0].dims, output_descriptors[0].dims) != 0) return -1;
120
+ validate_shape(&input_descriptors[0].dims, 1, output_descriptors[0].dims) != 0) return -1;
121
121
  int rows = input_descriptors[0].dims.d[0];
122
122
  if (rows < 1 || rows > INT_MAX / 16) return -1;
123
123
  int const index = sm_index();
@@ -141,9 +141,10 @@ public:
141
141
  }
142
142
 
143
143
  protected:
144
- std::int32_t validate_shape(Dims const& input, Dims const& output) const noexcept override
144
+ std::int32_t validate_shape(Dims const* inputs, std::int32_t input_count,
145
+ Dims const& output) const noexcept override
145
146
  {
146
- return shape_matches(input, output, 32) ? 0 : -1;
147
+ return input_count == 1 && shape_matches(inputs[0], output, 32) ? 0 : -1;
147
148
  }
148
149
 
149
150
  private:
@@ -178,7 +179,7 @@ public:
178
179
  {
179
180
  if (!valid_ || input_descriptors == nullptr || output_descriptors == nullptr ||
180
181
  inputs == nullptr || outputs == nullptr || inputs[0] == nullptr || outputs[0] == nullptr ||
181
- validate_shape(input_descriptors[0].dims, output_descriptors[0].dims) != 0) return -1;
182
+ validate_shape(&input_descriptors[0].dims, 1, output_descriptors[0].dims) != 0) return -1;
182
183
  int const rows = input_descriptors[0].dims.d[0];
183
184
  if (rows < 1 || rows > INT_MAX / 16) return -1;
184
185
  int const index = sm_index();
@@ -194,9 +195,10 @@ public:
194
195
  }
195
196
 
196
197
  protected:
197
- std::int32_t validate_shape(Dims const& input, Dims const& output) const noexcept override
198
+ std::int32_t validate_shape(Dims const* inputs, std::int32_t input_count,
199
+ Dims const& output) const noexcept override
198
200
  {
199
- return shape_matches(input, output, 16) ? 0 : -1;
201
+ return input_count == 1 && shape_matches(inputs[0], output, 16) ? 0 : -1;
200
202
  }
201
203
 
202
204
  private:
@@ -52,7 +52,15 @@ local_scheme = "no-local-version"
52
52
  version-file = "src/tensoracle/_version.py"
53
53
 
54
54
  [tool.hatch.build.targets.sdist]
55
- exclude = ["/.nf/", "/.github/", "/AGENTS.md", "/CLAUDE.md", "/tools/"]
55
+ exclude = [
56
+ "/.nf/",
57
+ "/.github/",
58
+ "/AGENTS.md",
59
+ "/CLAUDE.md",
60
+ "/tools/",
61
+ "/noren.yaml",
62
+ "/noren-words.txt",
63
+ ]
56
64
 
57
65
  [tool.hatch.build.targets.wheel]
58
66
  packages = ["src/tensoracle"]
@@ -66,6 +74,7 @@ dev = [
66
74
  "cupy-cuda13x>=14.2.0",
67
75
  "mkdocs-material>=9.7.7",
68
76
  "mkdocstrings[python]>=1.0.6",
77
+ "noren==0.2.0",
69
78
  "pytest>=8.0",
70
79
  "ruff>=0.8",
71
80
  ]
@@ -75,7 +84,7 @@ testpaths = ["tests"]
75
84
  addopts = ["--strict-markers"]
76
85
  markers = [
77
86
  "gpu: tests that build or load a TensorRT engine on a CUDA GPU",
78
- "req(id): parent requirement ID in .nf/docs/requirements.md",
87
+ "req(id): ID of the adopted spec item in .nf/docs/spec.md that the test verifies",
79
88
  ]
80
89
 
81
90
  [tool.ruff]
@@ -15,6 +15,7 @@ __all__ = [
15
15
  "InputShape",
16
16
  "OutputComparison",
17
17
  "PluginLibraryError",
18
+ "ShapeTensorOutputWarning",
18
19
  "TensorSpec",
19
20
  "build_serialized_engine",
20
21
  "compare_outputs",
@@ -23,7 +24,14 @@ __all__ = [
23
24
 
24
25
  _MODULES = {
25
26
  **dict.fromkeys(
26
- ("BuilderSettings", "EngineBuild", "InputShape", "build_serialized_engine"), "._build"
27
+ (
28
+ "BuilderSettings",
29
+ "EngineBuild",
30
+ "InputShape",
31
+ "ShapeTensorOutputWarning",
32
+ "build_serialized_engine",
33
+ ),
34
+ "._build",
27
35
  ),
28
36
  **dict.fromkeys(
29
37
  (
@@ -1,5 +1,6 @@
1
1
  """Build serialized TensorRT engines from caller-defined network graphs."""
2
2
 
3
+ import warnings
3
4
  from collections.abc import Callable, Mapping, Sequence
4
5
  from dataclasses import dataclass, field
5
6
  from pathlib import Path
@@ -11,6 +12,21 @@ from ._memory import release_host_memory
11
12
  from ._plugins import register_libraries
12
13
 
13
14
 
15
+ class ShapeTensorOutputWarning(UserWarning):
16
+ """Shape tensor outputs cannot be handled by tensoracle execution units."""
17
+
18
+
19
+ def _warn_shape_tensor_outputs(names: tuple[str, ...], *, stacklevel: int) -> None:
20
+ if names:
21
+ warnings.warn(
22
+ f"TensorRT outputs {', '.join(repr(name) for name in names)} are shape tensors; "
23
+ "tensoracle.create_execution() cannot handle this engine and calling it will "
24
+ "crash the process.",
25
+ ShapeTensorOutputWarning,
26
+ stacklevel=stacklevel + 1,
27
+ )
28
+
29
+
14
30
  @dataclass(frozen=True)
15
31
  class InputShape:
16
32
  """The minimum, preferred, and maximum dimensions of one input."""
@@ -76,6 +92,17 @@ def build_serialized_engine(
76
92
  logger: trt.ILogger | None = None,
77
93
  ) -> bytes:
78
94
  """Register ``plugins`` before calling ``graph`` and serializing its engine."""
95
+ plan, _ = _build_serialized_engine(graph, declaration, plugins=plugins, logger=logger)
96
+ return plan
97
+
98
+
99
+ def _build_serialized_engine(
100
+ graph: Callable[[trt.INetworkDefinition], None],
101
+ declaration: EngineBuild,
102
+ *,
103
+ plugins: Sequence[Path] = (),
104
+ logger: trt.ILogger | None = None,
105
+ ) -> tuple[bytes, tuple[str, ...]]:
79
106
  if declaration.precision not in ("fp32", "fp16"):
80
107
  raise ValueError(f"unsupported precision: {declaration.precision}")
81
108
  if declaration.settings.tf32:
@@ -93,6 +120,13 @@ def build_serialized_engine(
93
120
  builder = trt.Builder(logger or trt.Logger(trt.Logger.ERROR))
94
121
  network = builder.create_network(0)
95
122
  graph(network)
123
+ shape_tensor_outputs = []
124
+ for index in range(network.num_outputs):
125
+ tensor = network.get_output(index)
126
+ if tensor.is_shape_tensor:
127
+ shape_tensor_outputs.append(tensor.name)
128
+ output_names = tuple(shape_tensor_outputs)
129
+ _warn_shape_tensor_outputs(output_names, stacklevel=3)
96
130
  config = builder.create_builder_config()
97
131
  config.clear_flag(trt.BuilderFlag.TF32)
98
132
  config.set_memory_pool_limit(
@@ -129,7 +163,7 @@ def build_serialized_engine(
129
163
  serialized = builder.build_serialized_network(network, config)
130
164
  if serialized is None:
131
165
  raise RuntimeError("TensorRT failed to build the serialized engine")
132
- return bytes(serialized)
166
+ return bytes(serialized), output_names
133
167
  finally:
134
168
  # Tensor references also keep the network alive, including in exception frames.
135
169
  inputs.clear()
@@ -9,12 +9,13 @@ import platform
9
9
  import tempfile
10
10
  from collections.abc import Callable, Mapping, Sequence
11
11
  from dataclasses import asdict, dataclass
12
+ from dataclasses import field as dataclass_field
12
13
  from pathlib import Path
13
14
  from typing import TYPE_CHECKING
14
15
 
15
16
  import tensorrt as trt
16
17
 
17
- from ._build import EngineBuild, build_serialized_engine
18
+ from ._build import EngineBuild, _build_serialized_engine
18
19
  from ._plugins import register_libraries
19
20
 
20
21
  if TYPE_CHECKING:
@@ -49,6 +50,7 @@ class EngineHandle:
49
50
  engine: trt.ICudaEngine
50
51
  runtime: trt.Runtime
51
52
  path: Path
53
+ _shape_tensor_outputs: tuple[str, ...] = dataclass_field(default=(), repr=False)
52
54
 
53
55
  def create_execution(self) -> "EngineExecution":
54
56
  """Create an independent execution context with explicit close semantics."""
@@ -103,16 +105,29 @@ def open_engine(
103
105
  try:
104
106
  fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
105
107
  except FileExistsError:
106
- plan = _read_plan(path, conditions)
107
- return _open_plan(plan, path, declaration, io, actual_logger)
108
+ plan, shape_tensor_outputs = _read_plan(path, conditions)
109
+ return _open_plan(
110
+ plan,
111
+ path,
112
+ declaration,
113
+ io,
114
+ actual_logger,
115
+ shape_tensor_outputs=shape_tensor_outputs,
116
+ )
108
117
  else:
109
118
  os.close(fd)
110
119
 
111
120
  temporary: Path | None = None
112
121
  try:
113
- plan = build_serialized_engine(graph, declaration, plugins=plugins, logger=actual_logger)
114
- handle = _open_plan(plan, path, declaration, io, actual_logger)
122
+ plan, shape_tensor_outputs = _build_serialized_engine(
123
+ graph, declaration, plugins=plugins, logger=actual_logger
124
+ )
125
+ handle = _open_plan(
126
+ plan, path, declaration, io, actual_logger, shape_tensor_outputs=shape_tensor_outputs
127
+ )
115
128
  header = {"conditions": conditions, "plan_sha256": hashlib.sha256(plan).hexdigest()}
129
+ if shape_tensor_outputs:
130
+ header["shape_tensor_outputs"] = list(shape_tensor_outputs)
116
131
  header_bytes = _json_bytes(header)
117
132
  fd, filename = tempfile.mkstemp(prefix=f".{digest}.", dir=path.parent)
118
133
  temporary = Path(filename)
@@ -279,7 +294,7 @@ def _conditions(
279
294
  }
280
295
 
281
296
 
282
- def _read_plan(path: Path, conditions: dict) -> memoryview:
297
+ def _read_plan(path: Path, conditions: dict) -> tuple[memoryview, tuple[str, ...]]:
283
298
  data = path.read_bytes()
284
299
  offset = len(_MAGIC)
285
300
  if not data.startswith(_MAGIC) or len(data) < offset + 8:
@@ -294,12 +309,17 @@ def _read_plan(path: Path, conditions: dict) -> memoryview:
294
309
  raise CorruptEngineError(f"invalid TensorRT engine header: {path}") from error
295
310
  if not isinstance(header, dict) or not isinstance(header.get("conditions"), dict):
296
311
  raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
312
+ shape_tensor_outputs = header.get("shape_tensor_outputs", [])
313
+ if not isinstance(shape_tensor_outputs, list) or any(
314
+ not isinstance(name, str) for name in shape_tensor_outputs
315
+ ):
316
+ raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
297
317
  if header["conditions"] != conditions:
298
318
  raise EngineConditionError(f"TensorRT engine conditions differ: {path}")
299
319
  plan = memoryview(data)[header_end:]
300
320
  if header.get("plan_sha256") != hashlib.sha256(plan).hexdigest():
301
321
  raise CorruptEngineError(f"TensorRT engine checksum differs: {path}")
302
- return plan
322
+ return plan, tuple(shape_tensor_outputs)
303
323
 
304
324
 
305
325
  def _open_plan(
@@ -308,6 +328,8 @@ def _open_plan(
308
328
  declaration: EngineBuild,
309
329
  io: EngineIO,
310
330
  logger: trt.ILogger,
331
+ *,
332
+ shape_tensor_outputs: tuple[str, ...] = (),
311
333
  ) -> EngineHandle:
312
334
  runtime = trt.Runtime(logger)
313
335
  engine = runtime.deserialize_cuda_engine(plan)
@@ -349,7 +371,7 @@ def _open_plan(
349
371
  actual_range = (io.inputs[name].shape,) * 3
350
372
  if actual_range != expected_range:
351
373
  _io_mismatch(path, name, "profile range", actual_range, expected_range)
352
- return EngineHandle(engine, runtime, path)
374
+ return EngineHandle(engine, runtime, path, _shape_tensor_outputs=shape_tensor_outputs)
353
375
 
354
376
 
355
377
  def _io_mismatch(path: Path, name: str, field: str, observed: object, expected: object) -> None:
@@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Self
7
7
 
8
8
  import tensorrt as trt
9
9
 
10
+ from ._build import _warn_shape_tensor_outputs
11
+
10
12
  if TYPE_CHECKING:
11
13
  from ._cache import EngineHandle
12
14
 
@@ -51,6 +53,7 @@ class EngineExecution:
51
53
  self._context = engine.create_execution_context()
52
54
  if self._context is None:
53
55
  raise RuntimeError("TensorRT failed to create an execution context")
56
+ _warn_shape_tensor_outputs(handle._shape_tensor_outputs, stacklevel=2)
54
57
  # Probe output dimensions before accepting the engine. A ranged engine
55
58
  # still requires an explicit caller setting after this capability check.
56
59
  self._set_shapes(initial_shapes, creating=True)