tensoracle 0.3.2__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.2 → tensoracle-0.4.1}/PKG-INFO +19 -2
  2. {tensoracle-0.3.2 → tensoracle-0.4.1}/README.md +18 -1
  3. {tensoracle-0.3.2 → tensoracle-0.4.1}/cpp/include/tensoracle/cubin_plugin.hpp +24 -3
  4. {tensoracle-0.3.2 → tensoracle-0.4.1}/cpp/src/cubin_plugin.cpp +66 -15
  5. {tensoracle-0.3.2 → tensoracle-0.4.1}/docs/index.md +18 -1
  6. {tensoracle-0.3.2 → tensoracle-0.4.1}/docs/optimization.md +1 -1
  7. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/plugin.cpp +8 -6
  8. {tensoracle-0.3.2 → tensoracle-0.4.1}/pyproject.toml +11 -2
  9. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/_version.py +2 -2
  10. tensoracle-0.4.1/tests/cubin_io/check.cpp +158 -0
  11. tensoracle-0.4.1/tests/cubin_io/kernels.cu +11 -0
  12. tensoracle-0.4.1/tests/cubin_io/plugin.cpp +106 -0
  13. tensoracle-0.4.1/tests/plugin_io_runtime.py +35 -0
  14. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_build.py +6 -6
  15. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_build_memory.py +2 -2
  16. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_bundle.py +6 -6
  17. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_cache.py +13 -13
  18. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_comparison.py +8 -8
  19. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_convert.py +9 -9
  20. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_cpp_core.py +2 -2
  21. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_distribution.py +25 -16
  22. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_engine_io.py +3 -3
  23. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_example_plugin.py +17 -15
  24. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_execution.py +20 -20
  25. tensoracle-0.4.1/tests/test_plugin_io.py +230 -0
  26. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_plugins.py +6 -6
  27. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_producer.py +7 -7
  28. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_shape_tensor_outputs.py +17 -17
  29. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/test_weight_hash.py +10 -10
  30. {tensoracle-0.3.2 → tensoracle-0.4.1}/uv.lock +14 -0
  31. tensoracle-0.3.2/tests/test_public_mirror_spec.py +0 -324
  32. {tensoracle-0.3.2 → tensoracle-0.4.1}/.gitignore +0 -0
  33. {tensoracle-0.3.2 → tensoracle-0.4.1}/LICENSE +0 -0
  34. {tensoracle-0.3.2 → tensoracle-0.4.1}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
  35. {tensoracle-0.3.2 → tensoracle-0.4.1}/docs/api/bundle.md +0 -0
  36. {tensoracle-0.3.2 → tensoracle-0.4.1}/docs/api/index.md +0 -0
  37. {tensoracle-0.3.2 → tensoracle-0.4.1}/docs/api/producer.md +0 -0
  38. {tensoracle-0.3.2 → tensoracle-0.4.1}/docs/tensorrt-sdk.md +0 -0
  39. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/CMakeLists.txt +0 -0
  40. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/README.md +0 -0
  41. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
  42. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/bundle.json +0 -0
  43. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
  44. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
  45. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
  46. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
  47. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
  48. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
  49. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
  50. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
  51. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
  52. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
  53. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
  54. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
  55. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
  56. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
  57. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
  58. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
  59. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
  60. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/kernels.py +0 -0
  61. {tensoracle-0.3.2 → tensoracle-0.4.1}/examples/cubin_plugin/pyproject.toml +0 -0
  62. {tensoracle-0.3.2 → tensoracle-0.4.1}/mkdocs.yml +0 -0
  63. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/__init__.py +0 -0
  64. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/_build.py +0 -0
  65. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/_cache.py +0 -0
  66. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/_comparison.py +0 -0
  67. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/_execution.py +0 -0
  68. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/_memory.py +0 -0
  69. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/_plugins.py +0 -0
  70. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/bundle.py +0 -0
  71. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/convert.py +0 -0
  72. {tensoracle-0.3.2 → tensoracle-0.4.1}/src/tensoracle/producer.py +0 -0
  73. {tensoracle-0.3.2 → tensoracle-0.4.1}/tests/build_memory_runtime.py +0 -0
  74. {tensoracle-0.3.2 → 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.2
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
 
@@ -269,6 +269,23 @@ pip install 'tensoracle[producer]'
269
269
 
270
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.
271
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
+
272
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.
273
290
 
274
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
 
@@ -239,6 +239,23 @@ pip install 'tensoracle[producer]'
239
239
 
240
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.
241
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
+
242
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.
243
260
 
244
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
 
@@ -175,6 +175,23 @@ pip install 'tensoracle[producer]'
175
175
 
176
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.
177
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
+
178
195
  ## API reference
179
196
 
180
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]
@@ -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 = '0.3.2'
22
- __version_tuple__ = version_tuple = (0, 3, 2)
21
+ __version__ = version = '0.4.1'
22
+ __version_tuple__ = version_tuple = (0, 4, 1)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -0,0 +1,158 @@
1
+ #include "tensoracle/cubin_plugin.hpp"
2
+
3
+ #include <cassert>
4
+ #include <cstring>
5
+ #include <stdexcept>
6
+ #include <utility>
7
+
8
+ using namespace nvinfer1;
9
+ using namespace tensoracle::cubin;
10
+
11
+ class Probe final : public PluginBase {
12
+ public:
13
+ explicit Probe(PluginIO io = {}) : PluginBase("probe", "tests", {}, std::move(io)) {}
14
+ IPluginV3* clone() noexcept override { return nullptr; }
15
+ int32_t enqueue(PluginTensorDesc const*, PluginTensorDesc const*, void const* const*,
16
+ void* const*, void*, cudaStream_t) noexcept override { return -1; }
17
+
18
+ protected:
19
+ int32_t validate_shape(Dims const* inputs, int32_t count,
20
+ Dims const& output) const noexcept override
21
+ {
22
+ for (int32_t i = 0; i < count; ++i) {
23
+ if (inputs[i].nbDims != output.nbDims) return -1;
24
+ for (int32_t j = 0; j < output.nbDims; ++j)
25
+ if (inputs[i].d[j] != output.d[j]) return -1;
26
+ }
27
+ return 0;
28
+ }
29
+ };
30
+
31
+ PluginIO mixed_io()
32
+ {
33
+ return {{{DataType::kHALF, TensorFormat::kHWC8},
34
+ {DataType::kFLOAT, TensorFormat::kLINEAR}},
35
+ {DataType::kINT32, TensorFormat::kLINEAR}, 1};
36
+ }
37
+
38
+ int main(int argc, char** argv)
39
+ {
40
+ assert(argc == 2);
41
+ Probe plugin(mixed_io());
42
+ DynamicPluginTensorDesc descriptors[3]{};
43
+ DataType types[] = {DataType::kHALF, DataType::kFLOAT, DataType::kINT32};
44
+ TensorFormat formats[] = {TensorFormat::kHWC8, TensorFormat::kLINEAR,
45
+ TensorFormat::kLINEAR};
46
+ for (int i = 0; i < 3; ++i) {
47
+ descriptors[i].desc.type = types[i];
48
+ descriptors[i].desc.format = formats[i];
49
+ descriptors[i].desc.dims = Dims4(1, 48, 3, 5);
50
+ }
51
+ if (std::strcmp(argv[1], "formats") == 0) {
52
+ for (int i = 0; i < 3; ++i) {
53
+ assert(plugin.supportsFormatCombination(i, descriptors, 2, 1));
54
+ descriptors[i].desc.type = DataType::kINT8;
55
+ assert(!plugin.supportsFormatCombination(i, descriptors, 2, 1));
56
+ descriptors[i].desc.type = types[i];
57
+ descriptors[i].desc.format = TensorFormat::kCHW32;
58
+ assert(!plugin.supportsFormatCombination(i, descriptors, 2, 1));
59
+ descriptors[i].desc.format = formats[i];
60
+ }
61
+ // TensorRT has not initialized descriptors beyond the requested position.
62
+ descriptors[1].desc.type = DataType::kINT8;
63
+ descriptors[2].desc.format = TensorFormat::kCHW32;
64
+ assert(plugin.supportsFormatCombination(0, descriptors, 2, 1));
65
+ assert(!plugin.supportsFormatCombination(-1, descriptors, 2, 1));
66
+ assert(!plugin.supportsFormatCombination(3, descriptors, 2, 1));
67
+ assert(!plugin.supportsFormatCombination(0, nullptr, 2, 1));
68
+ } else if (std::strcmp(argv[1], "counts") == 0) {
69
+ DataType output{};
70
+ PluginTensorDesc runtime[] = {descriptors[0].desc, descriptors[1].desc};
71
+ assert(plugin.getNbOutputs() == 1);
72
+ for (int count : {-1, 0, 1, 3}) {
73
+ assert(!plugin.supportsFormatCombination(0, descriptors, count, 1));
74
+ assert(plugin.getOutputDataTypes(&output, 1, types, count) != 0);
75
+ assert(plugin.configurePlugin(descriptors, count, descriptors + 2, 1) != 0);
76
+ assert(plugin.onShapeChange(runtime, count, &descriptors[2].desc, 1) != 0);
77
+ }
78
+ for (int count : {-1, 0, 2}) {
79
+ assert(!plugin.supportsFormatCombination(0, descriptors, 2, count));
80
+ assert(plugin.getOutputDataTypes(&output, count, types, 2) != 0);
81
+ assert(plugin.configurePlugin(descriptors, 2, descriptors + 2, count) != 0);
82
+ assert(plugin.onShapeChange(runtime, 2, &descriptors[2].desc, count) != 0);
83
+ }
84
+ assert(plugin.getOutputDataTypes(nullptr, 1, types, 2) != 0);
85
+ assert(plugin.getOutputDataTypes(&output, 1, nullptr, 2) != 0);
86
+ assert(plugin.configurePlugin(nullptr, 2, descriptors + 2, 1) != 0);
87
+ assert(plugin.configurePlugin(descriptors, 2, nullptr, 1) != 0);
88
+ assert(plugin.onShapeChange(nullptr, 2, &descriptors[2].desc, 1) != 0);
89
+ assert(plugin.onShapeChange(runtime, 2, nullptr, 1) != 0);
90
+ } else if (std::strcmp(argv[1], "types") == 0) {
91
+ DataType output{};
92
+ assert(plugin.getOutputDataTypes(&output, 1, types, 2) == 0);
93
+ assert(output == DataType::kINT32);
94
+ for (int i = 0; i < 2; ++i) {
95
+ auto previous = types[i];
96
+ types[i] = DataType::kINT8;
97
+ assert(plugin.getOutputDataTypes(&output, 1, types, 2) != 0);
98
+ types[i] = previous;
99
+ }
100
+ } else if (std::strcmp(argv[1], "shapes") == 0) {
101
+ assert(plugin.configurePlugin(descriptors, 2, descriptors + 2, 1) == 0);
102
+ PluginTensorDesc runtime[] = {descriptors[0].desc, descriptors[1].desc};
103
+ assert(plugin.onShapeChange(runtime, 2, &descriptors[2].desc, 1) == 0);
104
+ descriptors[1].desc.dims.d[2] = 7;
105
+ runtime[1] = descriptors[1].desc;
106
+ assert(plugin.configurePlugin(descriptors, 2, descriptors + 2, 1) != 0);
107
+ assert(plugin.onShapeChange(runtime, 2, &descriptors[2].desc, 1) != 0);
108
+ descriptors[1].desc.dims.d[2] = 3;
109
+ runtime[1] = descriptors[1].desc;
110
+ for (int i = 0; i < 3; ++i) {
111
+ descriptors[i].desc.type = DataType::kINT8;
112
+ assert(plugin.configurePlugin(descriptors, 2, descriptors + 2, 1) != 0);
113
+ descriptors[i].desc.type = types[i];
114
+ descriptors[i].desc.format = TensorFormat::kCHW32;
115
+ assert(plugin.configurePlugin(descriptors, 2, descriptors + 2, 1) != 0);
116
+ descriptors[i].desc.format = formats[i];
117
+ auto saved = i < 2 ? runtime[i] : descriptors[2].desc;
118
+ auto& changed = i < 2 ? runtime[i] : descriptors[2].desc;
119
+ changed.type = DataType::kINT8;
120
+ assert(plugin.onShapeChange(runtime, 2, &descriptors[2].desc, 1) != 0);
121
+ changed = saved;
122
+ changed.format = TensorFormat::kCHW32;
123
+ assert(plugin.onShapeChange(runtime, 2, &descriptors[2].desc, 1) != 0);
124
+ changed = saved;
125
+ }
126
+ } else if (std::strcmp(argv[1], "defaults") == 0) {
127
+ Probe default_plugin;
128
+ descriptors[0].desc.type = descriptors[1].desc.type = DataType::kFLOAT;
129
+ descriptors[0].desc.format = descriptors[1].desc.format = TensorFormat::kLINEAR;
130
+ assert(default_plugin.getNbOutputs() == 1);
131
+ assert(default_plugin.supportsFormatCombination(0, descriptors, 1, 1));
132
+ assert(default_plugin.supportsFormatCombination(1, descriptors, 1, 1));
133
+ assert(!default_plugin.supportsFormatCombination(0, descriptors, 2, 1));
134
+ DataType input = DataType::kFLOAT, output{};
135
+ assert(default_plugin.getOutputDataTypes(&output, 1, &input, 1) == 0);
136
+ assert(output == DataType::kFLOAT);
137
+ assert(default_plugin.configurePlugin(descriptors, 1, descriptors + 1, 1) == 0);
138
+ assert(default_plugin.onShapeChange(&descriptors[0].desc, 1,
139
+ &descriptors[1].desc, 1) == 0);
140
+ descriptors[0].desc.type = DataType::kHALF;
141
+ assert(!default_plugin.supportsFormatCombination(0, descriptors, 1, 1));
142
+ } else if (std::strcmp(argv[1], "invalid") == 0) {
143
+ for (int shape_input : {-1, 2}) {
144
+ auto io = mixed_io();
145
+ io.shape_input_index = shape_input;
146
+ bool rejected = false;
147
+ try { Probe invalid(io); }
148
+ catch (std::invalid_argument const&) { rejected = true; }
149
+ assert(rejected);
150
+ }
151
+ auto io = mixed_io();
152
+ io.inputs.clear();
153
+ bool rejected = false;
154
+ try { Probe invalid(io); }
155
+ catch (std::invalid_argument const&) { rejected = true; }
156
+ assert(rejected);
157
+ } else { return 1; }
158
+ }
@@ -0,0 +1,11 @@
1
+ #include <cuda_fp16.h>
2
+
3
+ // All tensors use the same packed HWC8 storage. Channels are padded by the caller.
4
+ extern "C" __global__ void half_io(__half const* x, __half const* residual,
5
+ __half* output, int count, int mode)
6
+ {
7
+ int index = blockIdx.x * blockDim.x + threadIdx.x;
8
+ if (index >= count) return;
9
+ output[index] = mode == 0 ? __hadd(x[index], x[index])
10
+ : mode == 1 ? __hadd(x[index], residual[index]) : residual[index];
11
+ }