diff --git a/cpp/BUILD b/cpp/BUILD index ebaa458144..30619cda92 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -172,6 +172,25 @@ cc_library( }), ) +# Registration only, for the same reason as the allocator above: the kernel +# registry aborts when the same kernel is registered twice, so this cannot ride +# along with the backend. Executables that run a TensorRT delegated program +# depend on this directly. +# +# Not part of :executorch_backend_source_files. The CMake build of the reference +# runner links executorch::kernels, which already registers these. +cc_library( + name = "tensorrt_executorch_device_copy_kernels", + srcs = [ + "src/torch_tensorrt/executorch/RegisterDeviceCopyKernels.cpp", + ], + alwayslink = True, + deps = [ + "@executorch//:executorch_device_copy_kernels", + "@executorch//:executorch_headers", + ], +) + cc_library( name = "tensorrt_executorch_backend", srcs = [ diff --git a/cpp/src/torch_tensorrt/executorch/RegisterDeviceCopyKernels.cpp b/cpp/src/torch_tensorrt/executorch/RegisterDeviceCopyKernels.cpp new file mode 100644 index 0000000000..5f8284e63f --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/RegisterDeviceCopyKernels.cpp @@ -0,0 +1,32 @@ +#include + +namespace torch { +namespace executor { +namespace native { + +using executorch::aten::Tensor; +using executorch::runtime::KernelRuntimeContext; + +// ExecuTorch publishes no header for its portable kernel sources, and the +// Bazel target that owns this file compiles op__device_copy.cpp straight from +// the pinned tree, so the two entry points are declared here. +Tensor& _h2d_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out); +Tensor& _d2h_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out); + +} // namespace native +} // namespace executor +} // namespace torch + +// A program delegated to TensorRT still runs two ops outside the delegate: the +// host to device copy of its inputs and the device to host copy of its outputs, +// which the device placement pass inserts at export time. ExecuTorch registers +// their kernels from a generated kernel library, which this Bazel build does +// not produce, so a runner linking the core runtime alone fails every +// load_method with +// +// kernel 'et_copy::_h2d_copy.out' not found. +// +// Register those two rather than the whole portable kernel set, which a fully +// delegated program never calls. +EXECUTORCH_LIBRARY(et_copy, "_h2d_copy.out", torch::executor::native::_h2d_copy_out); +EXECUTORCH_LIBRARY(et_copy, "_d2h_copy.out", torch::executor::native::_d2h_copy_out); diff --git a/examples/executorch_reference_runner/BUILD b/examples/executorch_reference_runner/BUILD index cccba328ec..5d85db9f5b 100644 --- a/examples/executorch_reference_runner/BUILD +++ b/examples/executorch_reference_runner/BUILD @@ -26,6 +26,7 @@ cc_binary( deps = [ "//cpp:tensorrt_executorch_backend", "//cpp:tensorrt_executorch_cuda_device_allocator", + "//cpp:tensorrt_executorch_device_copy_kernels", "@executorch//:executorch_core", "@executorch//:executorch_file_data_loader", "@executorch//:extension_cuda", @@ -38,6 +39,7 @@ cc_binary( deps = [ "//cpp:tensorrt_executorch_backend", "//cpp:tensorrt_executorch_cuda_device_allocator", + "//cpp:tensorrt_executorch_device_copy_kernels", "@cuda//:cudart", "@executorch//:executorch_core", "@executorch//:executorch_file_data_loader", diff --git a/tests/py/dynamo/conversion/test_cumsum_aten.py b/tests/py/dynamo/conversion/test_cumsum_aten.py index ab71b4339f..5c2147e281 100644 --- a/tests/py/dynamo/conversion/test_cumsum_aten.py +++ b/tests/py/dynamo/conversion/test_cumsum_aten.py @@ -33,9 +33,7 @@ def forward(self, x): ) return - self.run_test( - Cumsum(), inputs, immutable_weights=False, use_dynamo_tracer=True - ) + self.run_test(Cumsum(), inputs, immutable_weights=False, use_dynamo_tracer=True) @parameterized.expand( [ @@ -108,10 +106,7 @@ def forward(self, x): == opt_shape[positive_dim] == max_shape[positive_dim] ) - if ( - has_static_trip_count - and not is_tensorrt_rtx_version_supported("1.7") - ): + if has_static_trip_count and not is_tensorrt_rtx_version_supported("1.7"): with self.assertRaises(UnsupportedOperatorException): self.run_test_with_dynamic_shape( Cumsum(), diff --git a/third_party/executorch/BUILD b/third_party/executorch/BUILD index 17c6e5f337..98d66cd02b 100644 --- a/third_party/executorch/BUILD +++ b/third_party/executorch/BUILD @@ -147,6 +147,21 @@ cc_library( ], ) +# The kernels behind the device copies that an exported program runs around the +# delegate. In the pinned ExecuTorch release these reach a binary only through +# the generated kernel library, which this overlay does not build, so compile +# the two implementations on their own. Registering them is left to +# //cpp:tensorrt_executorch_device_copy_kernels, because a binary that already +# gets them from ExecuTorch's own kernel library must not register them twice. +cc_library( + name = "executorch_device_copy_kernels", + srcs = ["executorch/kernels/portable/cpu/op__device_copy.cpp"], + deps = [ + ":executorch_core", + ":executorch_headers", + ], +) + cc_library( name = "executorch_headers", hdrs = glob(