tensoracle 0.3.2__tar.gz → 0.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {tensoracle-0.3.2 → tensoracle-0.4.0}/PKG-INFO +19 -2
- {tensoracle-0.3.2 → tensoracle-0.4.0}/README.md +18 -1
- {tensoracle-0.3.2 → tensoracle-0.4.0}/cpp/include/tensoracle/cubin_plugin.hpp +24 -3
- {tensoracle-0.3.2 → tensoracle-0.4.0}/cpp/src/cubin_plugin.cpp +66 -15
- {tensoracle-0.3.2 → tensoracle-0.4.0}/docs/index.md +18 -1
- {tensoracle-0.3.2 → tensoracle-0.4.0}/docs/optimization.md +1 -1
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/plugin.cpp +8 -6
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/_version.py +2 -2
- tensoracle-0.4.0/tests/cubin_io/check.cpp +158 -0
- tensoracle-0.4.0/tests/cubin_io/kernels.cu +11 -0
- tensoracle-0.4.0/tests/cubin_io/plugin.cpp +106 -0
- tensoracle-0.4.0/tests/plugin_io_runtime.py +35 -0
- tensoracle-0.4.0/tests/test_plugin_io.py +230 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/.gitignore +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/LICENSE +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/docs/api/bundle.md +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/docs/api/index.md +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/docs/api/producer.md +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/docs/tensorrt-sdk.md +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/CMakeLists.txt +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/README.md +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/bundle.json +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/kernels.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/pyproject.toml +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/mkdocs.yml +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/pyproject.toml +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/__init__.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/_build.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/_cache.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/_comparison.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/_execution.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/_memory.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/_plugins.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/bundle.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/convert.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/src/tensoracle/producer.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/build_memory_runtime.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/example_runtime.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_build.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_build_memory.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_bundle.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_cache.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_comparison.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_convert.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_cpp_core.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_distribution.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_engine_io.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_example_plugin.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_execution.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_plugins.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_producer.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_public_mirror_spec.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_shape_tensor_outputs.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/tests/test_weight_hash.py +0 -0
- {tensoracle-0.3.2 → tensoracle-0.4.0}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: tensoracle
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.4.0
|
|
4
4
|
Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
|
|
5
5
|
Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
|
|
6
6
|
Project-URL: repository, https://github.com/sync-dev-org/tensoracle
|
|
@@ -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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
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&
|
|
130
|
-
return
|
|
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
|
|
139
|
-
|
|
140
|
-
|
|
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
|
|
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[
|
|
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
|
-
|
|
160
|
-
|
|
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
|
-
|
|
173
|
-
|
|
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](
|
|
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
|
|
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(
|
|
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
|
|
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(
|
|
201
|
+
return input_count == 1 && shape_matches(inputs[0], output, 16) ? 0 : -1;
|
|
200
202
|
}
|
|
201
203
|
|
|
202
204
|
private:
|
|
@@ -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.
|
|
22
|
-
__version_tuple__ = version_tuple = (0,
|
|
21
|
+
__version__ = version = '0.4.0'
|
|
22
|
+
__version_tuple__ = version_tuple = (0, 4, 0)
|
|
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
|
+
}
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
#include "tensoracle/cubin_plugin.hpp"
|
|
2
|
+
#include "kernel_bytes.hpp"
|
|
3
|
+
|
|
4
|
+
#include <climits>
|
|
5
|
+
#include <new>
|
|
6
|
+
#include <vector>
|
|
7
|
+
|
|
8
|
+
using namespace nvinfer1;
|
|
9
|
+
using namespace tensoracle::cubin;
|
|
10
|
+
|
|
11
|
+
namespace {
|
|
12
|
+
char const* const kNames[] = {"TensoracleHalfUnary", "TensoracleHalfResidual", "TensoracleHalfSelect"};
|
|
13
|
+
constexpr char kNamespace[] = "tensoracle.tests.io";
|
|
14
|
+
|
|
15
|
+
PluginIO declaration(int mode)
|
|
16
|
+
{
|
|
17
|
+
TensorSpec half{DataType::kHALF, TensorFormat::kHWC8};
|
|
18
|
+
return {std::vector<TensorSpec>(mode == 0 ? 1 : 2, half), half, mode == 2 ? 1 : 0};
|
|
19
|
+
}
|
|
20
|
+
|
|
21
|
+
class HalfPlugin final : public PluginBase {
|
|
22
|
+
public:
|
|
23
|
+
explicit HalfPlugin(int mode) : PluginBase(kNames[mode], kNamespace, {}, declaration(mode)), mode_(mode) {}
|
|
24
|
+
IPluginV3* clone() noexcept override
|
|
25
|
+
{
|
|
26
|
+
try { return new HalfPlugin(mode_); }
|
|
27
|
+
catch (...) { return nullptr; }
|
|
28
|
+
}
|
|
29
|
+
int32_t enqueue(PluginTensorDesc const* input_desc, PluginTensorDesc const* output_desc,
|
|
30
|
+
void const* const* inputs, void* const* outputs, void*, cudaStream_t stream) noexcept override
|
|
31
|
+
{
|
|
32
|
+
int input_count = mode_ == 0 ? 1 : 2;
|
|
33
|
+
Dims dims[2] = {input_desc[0].dims, input_desc[input_count - 1].dims};
|
|
34
|
+
if (validate_shape(dims, input_count, output_desc[0].dims) != 0) return -1;
|
|
35
|
+
auto const& shape = output_desc[0].dims;
|
|
36
|
+
int64_t count = shape.d[0] * ((shape.d[1] + 7) / 8 * 8) * shape.d[2] * shape.d[3];
|
|
37
|
+
if (count <= 0 || count > INT_MAX) return -1;
|
|
38
|
+
int size = static_cast<int>(count);
|
|
39
|
+
void const* x = inputs[0];
|
|
40
|
+
void const* residual = inputs[input_count - 1];
|
|
41
|
+
void* output = outputs[0];
|
|
42
|
+
void* arguments[] = {&x, &residual, &output, &size, &mode_};
|
|
43
|
+
CUmodule module{};
|
|
44
|
+
CUfunction function{};
|
|
45
|
+
auto status = cuModuleLoadData(&module, kKernelBytes);
|
|
46
|
+
if (status == CUDA_SUCCESS) status = cuModuleGetFunction(&function, module, "half_io");
|
|
47
|
+
auto native_stream = reinterpret_cast<CUstream>(stream);
|
|
48
|
+
if (status == CUDA_SUCCESS)
|
|
49
|
+
status = cuLaunchKernel(function, (size + 127) / 128, 1, 1, 128, 1, 1, 0,
|
|
50
|
+
native_stream, arguments, nullptr);
|
|
51
|
+
if (status == CUDA_SUCCESS) status = cuStreamSynchronize(native_stream);
|
|
52
|
+
if (module) static_cast<void>(cuModuleUnload(module));
|
|
53
|
+
return status == CUDA_SUCCESS ? 0 : -1;
|
|
54
|
+
}
|
|
55
|
+
|
|
56
|
+
protected:
|
|
57
|
+
int32_t validate_shape(Dims const* inputs, int32_t input_count,
|
|
58
|
+
Dims const& output) const noexcept override
|
|
59
|
+
{
|
|
60
|
+
if (input_count != (mode_ == 0 ? 1 : 2) || output.nbDims != 4) return -1;
|
|
61
|
+
for (int i = mode_ == 2 ? 1 : 0; i < input_count; ++i) {
|
|
62
|
+
if (inputs[i].nbDims != 4) return -1;
|
|
63
|
+
for (int j = 0; j < 4; ++j)
|
|
64
|
+
if (inputs[i].d[j] != -1 && output.d[j] != -1 &&
|
|
65
|
+
inputs[i].d[j] != output.d[j]) return -1;
|
|
66
|
+
}
|
|
67
|
+
return 0;
|
|
68
|
+
}
|
|
69
|
+
|
|
70
|
+
private:
|
|
71
|
+
int mode_;
|
|
72
|
+
};
|
|
73
|
+
|
|
74
|
+
class Creator final : public IPluginCreatorV3One {
|
|
75
|
+
public:
|
|
76
|
+
explicit Creator(int mode) : mode_(mode) {}
|
|
77
|
+
char const* getPluginName() const noexcept override { return kNames[mode_]; }
|
|
78
|
+
char const* getPluginVersion() const noexcept override { return "1"; }
|
|
79
|
+
char const* getPluginNamespace() const noexcept override { return kNamespace; }
|
|
80
|
+
PluginFieldCollection const* getFieldNames() noexcept override { return &fields_; }
|
|
81
|
+
IPluginV3* createPlugin(char const*, PluginFieldCollection const* fields,
|
|
82
|
+
TensorRTPhase) noexcept override
|
|
83
|
+
{
|
|
84
|
+
if (fields == nullptr || fields->nbFields != 0) return nullptr;
|
|
85
|
+
try { return new HalfPlugin(mode_); }
|
|
86
|
+
catch (...) { return nullptr; }
|
|
87
|
+
}
|
|
88
|
+
private:
|
|
89
|
+
int mode_;
|
|
90
|
+
PluginFieldCollection fields_{};
|
|
91
|
+
};
|
|
92
|
+
Creator gUnary(0), gResidual(1), gSelect(2);
|
|
93
|
+
}
|
|
94
|
+
|
|
95
|
+
#if defined(_WIN32)
|
|
96
|
+
#define IO_EXPORT __declspec(dllexport)
|
|
97
|
+
#else
|
|
98
|
+
#define IO_EXPORT __attribute__((visibility("default")))
|
|
99
|
+
#endif
|
|
100
|
+
extern "C" IO_EXPORT void setLoggerFinder(ILoggerFinder*) noexcept {}
|
|
101
|
+
extern "C" IO_EXPORT IPluginCreatorInterface* const* getCreators(int32_t& count) noexcept
|
|
102
|
+
{
|
|
103
|
+
static IPluginCreatorInterface* creators[] = {&gUnary, &gResidual, &gSelect};
|
|
104
|
+
count = 3;
|
|
105
|
+
return creators;
|
|
106
|
+
}
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
"""Build one invalid plugin graph in a separate TensorRT process."""
|
|
2
|
+
|
|
3
|
+
import sys
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
import tensorrt as trt
|
|
7
|
+
|
|
8
|
+
from test_plugin_io import _plan
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def main():
|
|
12
|
+
plugin, fault = sys.argv[1:]
|
|
13
|
+
shapes = {"x": (1, 48, 3, 5), "residual": (1, 48, 3, 5)}
|
|
14
|
+
if fault == "missing-input":
|
|
15
|
+
shapes.pop("residual")
|
|
16
|
+
elif fault == "extra-input":
|
|
17
|
+
shapes["extra"] = (1, 48, 3, 5)
|
|
18
|
+
elif fault == "shape":
|
|
19
|
+
shapes["residual"] = (1, 48, 4, 5)
|
|
20
|
+
try:
|
|
21
|
+
_plan(
|
|
22
|
+
Path(plugin), "Residual", shapes, dtype=trt.float32 if fault == "type" else trt.float16
|
|
23
|
+
)
|
|
24
|
+
except RuntimeError as error:
|
|
25
|
+
assert str(error) in {
|
|
26
|
+
"TensorRT rejected the plugin IO declaration",
|
|
27
|
+
"TensorRT failed to build the serialized engine",
|
|
28
|
+
}, str(error)
|
|
29
|
+
print("declaration rejected")
|
|
30
|
+
else:
|
|
31
|
+
raise AssertionError("TensorRT built a network outside the plugin declaration")
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
if __name__ == "__main__":
|
|
35
|
+
main()
|
|
@@ -0,0 +1,230 @@
|
|
|
1
|
+
"""Native cubin plugins declare tensor types, formats, and their output shape source."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import subprocess
|
|
5
|
+
import sys
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
import numpy as np
|
|
9
|
+
import pytest
|
|
10
|
+
import tensorrt as trt
|
|
11
|
+
from cuda.bindings import runtime as cudart
|
|
12
|
+
|
|
13
|
+
from tensoracle import BuilderSettings, EngineBuild, InputShape, build_serialized_engine
|
|
14
|
+
|
|
15
|
+
ROOT = Path(__file__).resolve().parents[1]
|
|
16
|
+
FIXTURES = ROOT / "tests/cubin_io"
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@pytest.fixture(scope="module")
|
|
20
|
+
def native_tools():
|
|
21
|
+
sdk = os.environ.get("TENSORACLE_SDK_ROOT")
|
|
22
|
+
cuda = os.environ.get("TENSORACLE_CUDA_ROOT")
|
|
23
|
+
if not sdk or not cuda:
|
|
24
|
+
pytest.skip("TensorRT SDK and CUDA toolkit were not supplied")
|
|
25
|
+
return Path(sdk), Path(cuda)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _compile(source, output, native_tools, *options):
|
|
29
|
+
sdk, cuda = native_tools
|
|
30
|
+
result = subprocess.run(
|
|
31
|
+
[
|
|
32
|
+
"c++",
|
|
33
|
+
"-std=c++17",
|
|
34
|
+
"-Wno-deprecated-declarations",
|
|
35
|
+
f"-I{ROOT / 'cpp/include'}",
|
|
36
|
+
f"-I{sdk / 'include'}",
|
|
37
|
+
f"-I{cuda / 'include'}",
|
|
38
|
+
str(source),
|
|
39
|
+
str(ROOT / "cpp/src/cubin_plugin.cpp"),
|
|
40
|
+
f"-L{cuda / 'lib64/stubs'}",
|
|
41
|
+
"-lcuda",
|
|
42
|
+
*options,
|
|
43
|
+
"-o",
|
|
44
|
+
str(output),
|
|
45
|
+
],
|
|
46
|
+
check=False,
|
|
47
|
+
capture_output=True,
|
|
48
|
+
text=True,
|
|
49
|
+
timeout=45,
|
|
50
|
+
)
|
|
51
|
+
assert result.returncode == 0, result.stderr
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@pytest.fixture(scope="module")
|
|
55
|
+
def native_contract(native_tools, tmp_path_factory):
|
|
56
|
+
work = tmp_path_factory.mktemp("plugin-io-contract")
|
|
57
|
+
executable = work / "check"
|
|
58
|
+
_compile(FIXTURES / "check.cpp", executable, native_tools)
|
|
59
|
+
return executable
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
@pytest.mark.req("REQ-TENSORACLE-003")
|
|
63
|
+
@pytest.mark.parametrize("case", ["formats", "counts", "types", "shapes", "defaults", "invalid"])
|
|
64
|
+
def test_native_plugin_checks_declared_io(native_contract, case):
|
|
65
|
+
"""Plugins enforce each tensor's declared format, type, count, and shape validation."""
|
|
66
|
+
subprocess.run([str(native_contract), case], check=True, timeout=10)
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def _cuda(result):
|
|
70
|
+
status, *values = result
|
|
71
|
+
assert status == cudart.cudaError_t.cudaSuccess, status
|
|
72
|
+
return values[0] if values else None
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
@pytest.fixture(scope="module")
|
|
76
|
+
def half_plugin(native_tools, tmp_path_factory):
|
|
77
|
+
_, cuda = native_tools
|
|
78
|
+
work = tmp_path_factory.mktemp("plugin-half-io")
|
|
79
|
+
device = _cuda(cudart.cudaGetDevice())
|
|
80
|
+
properties = _cuda(cudart.cudaGetDeviceProperties(device))
|
|
81
|
+
sm = properties.major * 10 + properties.minor
|
|
82
|
+
cubin = work / "half_io.cubin"
|
|
83
|
+
subprocess.run(
|
|
84
|
+
[
|
|
85
|
+
str(cuda / "bin/nvcc"),
|
|
86
|
+
"--cubin",
|
|
87
|
+
f"-arch=sm_{sm}",
|
|
88
|
+
str(FIXTURES / "kernels.cu"),
|
|
89
|
+
"-o",
|
|
90
|
+
str(cubin),
|
|
91
|
+
],
|
|
92
|
+
check=True,
|
|
93
|
+
capture_output=True,
|
|
94
|
+
text=True,
|
|
95
|
+
timeout=45,
|
|
96
|
+
)
|
|
97
|
+
(work / "kernel_bytes.hpp").write_text(
|
|
98
|
+
"alignas(16) unsigned char const kKernelBytes[] = {"
|
|
99
|
+
+ ",".join(str(byte) for byte in cubin.read_bytes())
|
|
100
|
+
+ "};\n",
|
|
101
|
+
encoding="utf-8",
|
|
102
|
+
)
|
|
103
|
+
plugin = work / "libhalf_io.so"
|
|
104
|
+
_compile(FIXTURES / "plugin.cpp", plugin, native_tools, "-shared", "-fPIC", f"-I{work}")
|
|
105
|
+
return plugin
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def _plan(plugin, kind, shapes, *, dtype=trt.float16):
|
|
109
|
+
def graph(network):
|
|
110
|
+
creator = trt.get_plugin_registry().get_creator(
|
|
111
|
+
f"TensoracleHalf{kind}", "1", "tensoracle.tests.io"
|
|
112
|
+
)
|
|
113
|
+
assert creator is not None
|
|
114
|
+
native = creator.create_plugin("io", trt.PluginFieldCollection([]), trt.TensorRTPhase.BUILD)
|
|
115
|
+
assert native is not None
|
|
116
|
+
inputs = [network.add_input(name, dtype, shape) for name, shape in shapes.items()]
|
|
117
|
+
layer = network.add_plugin_v3(inputs, [], native)
|
|
118
|
+
if layer is None:
|
|
119
|
+
raise RuntimeError("TensorRT rejected the plugin IO declaration")
|
|
120
|
+
result = layer.get_output(0)
|
|
121
|
+
result.name = "result"
|
|
122
|
+
network.mark_output(result)
|
|
123
|
+
|
|
124
|
+
return build_serialized_engine(
|
|
125
|
+
graph,
|
|
126
|
+
EngineBuild(
|
|
127
|
+
"fp16",
|
|
128
|
+
{name: InputShape.fixed(shape) for name, shape in shapes.items()},
|
|
129
|
+
BuilderSettings(workspace_limit=16 << 20, optimization_level=0),
|
|
130
|
+
),
|
|
131
|
+
plugins=(plugin,),
|
|
132
|
+
)
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
def _execute(plan, inputs):
|
|
136
|
+
logger = trt.Logger(trt.Logger.ERROR)
|
|
137
|
+
runtime = trt.Runtime(logger)
|
|
138
|
+
try:
|
|
139
|
+
engine = runtime.deserialize_cuda_engine(plan)
|
|
140
|
+
assert engine is not None
|
|
141
|
+
context = engine.create_execution_context()
|
|
142
|
+
assert context is not None
|
|
143
|
+
output = np.empty(tuple(context.get_tensor_shape("result")), dtype=np.float16)
|
|
144
|
+
buffers = {}
|
|
145
|
+
stream = _cuda(cudart.cudaStreamCreate())
|
|
146
|
+
try:
|
|
147
|
+
for name, value in {**inputs, "result": output}.items():
|
|
148
|
+
# Network I/O is LINEAR; TensorRT reformats to the plugin's HWC8 declaration.
|
|
149
|
+
assert engine.get_tensor_format(name) == trt.TensorFormat.LINEAR
|
|
150
|
+
buffers[name] = _cuda(cudart.cudaMalloc(value.nbytes))
|
|
151
|
+
assert context.set_tensor_address(name, int(buffers[name]))
|
|
152
|
+
if name in inputs:
|
|
153
|
+
_cuda(
|
|
154
|
+
cudart.cudaMemcpy(
|
|
155
|
+
buffers[name],
|
|
156
|
+
value.ctypes.data,
|
|
157
|
+
value.nbytes,
|
|
158
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
|
|
159
|
+
)
|
|
160
|
+
)
|
|
161
|
+
assert context.execute_async_v3(int(stream))
|
|
162
|
+
_cuda(cudart.cudaStreamSynchronize(stream))
|
|
163
|
+
_cuda(
|
|
164
|
+
cudart.cudaMemcpy(
|
|
165
|
+
output.ctypes.data,
|
|
166
|
+
buffers["result"],
|
|
167
|
+
output.nbytes,
|
|
168
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
|
|
169
|
+
)
|
|
170
|
+
)
|
|
171
|
+
return output
|
|
172
|
+
finally:
|
|
173
|
+
_cuda(cudart.cudaStreamSynchronize(stream))
|
|
174
|
+
del context, engine
|
|
175
|
+
for pointer in buffers.values():
|
|
176
|
+
_cuda(cudart.cudaFree(pointer))
|
|
177
|
+
_cuda(cudart.cudaStreamDestroy(stream))
|
|
178
|
+
finally:
|
|
179
|
+
del runtime
|
|
180
|
+
|
|
181
|
+
|
|
182
|
+
@pytest.mark.gpu
|
|
183
|
+
@pytest.mark.req("REQ-TENSORACLE-003")
|
|
184
|
+
@pytest.mark.parametrize("kind", ["Unary", "Residual"])
|
|
185
|
+
def test_half_hwc8_plugin_matches_numpy(half_plugin, kind, tmp_path, monkeypatch):
|
|
186
|
+
"""One-input and two-input fp16 HWC8 plugins execute and match NumPy's output."""
|
|
187
|
+
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
|
|
188
|
+
shape = (1, 48, 3, 5)
|
|
189
|
+
x = (np.arange(np.prod(shape)).reshape(shape) % 31 / 8).astype(np.float16)
|
|
190
|
+
inputs = {"x": x}
|
|
191
|
+
expected = x + x
|
|
192
|
+
if kind == "Residual":
|
|
193
|
+
residual = (np.flip(x, axis=1) / 2).astype(np.float16).copy()
|
|
194
|
+
inputs["residual"] = residual
|
|
195
|
+
expected = x + residual
|
|
196
|
+
plan = _plan(half_plugin, kind, {name: value.shape for name, value in inputs.items()})
|
|
197
|
+
np.testing.assert_array_equal(_execute(plan, inputs), expected)
|
|
198
|
+
|
|
199
|
+
|
|
200
|
+
@pytest.mark.gpu
|
|
201
|
+
@pytest.mark.req("REQ-TENSORACLE-003")
|
|
202
|
+
def test_plugin_output_shape_follows_selected_input(half_plugin, tmp_path, monkeypatch):
|
|
203
|
+
"""A plugin can copy input one's shape when input zero has a different shape."""
|
|
204
|
+
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
|
|
205
|
+
inputs = {
|
|
206
|
+
"x": np.zeros((1, 48, 1, 1), dtype=np.float16),
|
|
207
|
+
"residual": np.arange(720, dtype=np.float16).reshape(1, 48, 3, 5),
|
|
208
|
+
}
|
|
209
|
+
plan = _plan(half_plugin, "Select", {name: value.shape for name, value in inputs.items()})
|
|
210
|
+
np.testing.assert_array_equal(_execute(plan, inputs), inputs["residual"])
|
|
211
|
+
|
|
212
|
+
|
|
213
|
+
@pytest.mark.gpu
|
|
214
|
+
@pytest.mark.req("REQ-TENSORACLE-003")
|
|
215
|
+
@pytest.mark.parametrize("fault", ["missing-input", "extra-input", "shape", "type"])
|
|
216
|
+
def test_plugin_rejects_network_outside_declaration(half_plugin, fault, tmp_path):
|
|
217
|
+
"""A network with wrong input count, differing residual shapes, or wrong types fails."""
|
|
218
|
+
environment = dict(os.environ, XDG_CACHE_HOME=str(tmp_path / "cache"))
|
|
219
|
+
# Keep TensorRT's error state and any failed native graph in its own process.
|
|
220
|
+
result = subprocess.run(
|
|
221
|
+
[sys.executable, str(ROOT / "tests/plugin_io_runtime.py"), str(half_plugin), fault],
|
|
222
|
+
cwd=tmp_path,
|
|
223
|
+
env=environment,
|
|
224
|
+
check=False,
|
|
225
|
+
capture_output=True,
|
|
226
|
+
text=True,
|
|
227
|
+
timeout=45,
|
|
228
|
+
)
|
|
229
|
+
assert result.returncode == 0, result.stdout + result.stderr
|
|
230
|
+
assert result.stdout.splitlines()[-1] == "declaration rejected"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py
RENAMED
|
File without changes
|
{tensoracle-0.3.2 → tensoracle-0.4.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|