From d9ca78817cf891fe9d052f3f7b1f16709430d777 Mon Sep 17 00:00:00 2001 From: apbose Date: Mon, 3 Apr 2023 10:56:18 -0700 Subject: [PATCH 1/7] Adding converter for chunk in torchscript --- core/conversion/converters/impl/chunk.cpp | 69 +++++++++++++++++++ tests/core/conversion/converters/BUILD | 5 ++ .../core/conversion/converters/test_chunk.cpp | 34 +++++++++ 3 files changed, 108 insertions(+) create mode 100644 core/conversion/converters/impl/chunk.cpp create mode 100644 tests/core/conversion/converters/test_chunk.cpp diff --git a/core/conversion/converters/impl/chunk.cpp b/core/conversion/converters/impl/chunk.cpp new file mode 100644 index 0000000000..21bf7848ec --- /dev/null +++ b/core/conversion/converters/impl/chunk.cpp @@ -0,0 +1,69 @@ +#include "core/conversion/converters/converters.h" +#include "core/conversion/tensorcontainer/TensorContainer.h" +#include "core/util/prelude.h" + +namespace torch_tensorrt { +namespace core { +namespace conversion { +namespace converters { +namespace impl { +namespace { + +// clang-format off +auto cat_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() + .pattern({"aten::chunk(Tensor(a -> *) self, int chunks, int dim=0) -> Tensor(a)[]", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto in = args[0].ITensorOrFreeze(ctx); + auto chunks = args[1].unwrapToInt(); + auto dim = args[2].unwrapToInt(); + bool dynamic_shape = ctx->input_is_dynamic; + int size = in->getDimensions().nbDims; + int maxDim = static_cast(in->getDimensions().d[axis]); + + if(dim < 0) { + dim = tensors[0]->getDimensions().nbDims + dim; + } + if (dynamic_shape) { + TORCHTRT_ASSERT(in.d[dim] != -1, "Can't chunk on dynamic shape dimension!"); + } + if (chunks > in.d[dim]) { + LOG_WARNING("The chunks size" << chunks << "along dimension" << dim << "is greater than tensor with size" << in->getDimensions + << "it will default to dimension" << in.d[dim]) + } + int step = (input_val.shape[dim] + chunks - 1) / chunks + nvinfer1::Dims start_, size_, stride_; + start_.nbDims = nbdims; + size_.nbDims = nbdims; + stride_.nbDims = nbdims; + + int startIdx = 0; + int endIdx = maxDim; + + for (int i = 0; i < nbdims; i++) { + if (i == axis) { + start_.d[i] = startIdx; + size_.d[i] = (endIdx - startIdx - 1) / step + 1; + stride_.d[i] = step; + } else { + start_.d[i] = 0; + size_.d[i] = in.d[i]; // for static + stride_.d[i] = 1; + } + } + // update slice layer + if(!dynamic_shape): + auto slice_layer = ctx->net->addSlice(*in, start_, size_, stride_); + LOG_DEBUG("start_:" << start_); + LOG_DEBUG("size_:" << size_); + LOG_DEBUG("stride_:" << stride_); + auto slice_out = slice_layer->getOutput(0); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], slice_out); + LOG_DEBUG("Slice layer output shape: " << out->getDimensions()); + }}); +// clang-format on +} // namespace +} // namespace impl +} // namespace converters +} // namespace conversion +} // namespace core +} // namespace torch_tensorrt diff --git a/tests/core/conversion/converters/BUILD b/tests/core/conversion/converters/BUILD index a8c57b1b41..1973c112fd 100644 --- a/tests/core/conversion/converters/BUILD +++ b/tests/core/conversion/converters/BUILD @@ -35,6 +35,10 @@ converter_test( name = "test_cast", ) +converter_test( + name = "test_chunk", +) + converter_test( name = "test_clone", ) @@ -208,6 +212,7 @@ test_suite( ":test_batch_norm", ":test_bitwise", ":test_cast", + ":test_chunk", ":test_clamp", ":test_clone", ":test_comparators", diff --git a/tests/core/conversion/converters/test_chunk.cpp b/tests/core/conversion/converters/test_chunk.cpp new file mode 100644 index 0000000000..bd6b7cec1e --- /dev/null +++ b/tests/core/conversion/converters/test_chunk.cpp @@ -0,0 +1,34 @@ +#include +#include +#include "core/compiler.h" +#include "core/lowering/passes/passes.h" +#include "gtest/gtest.h" +#include "tests/util/util.h" +#include "torch/csrc/jit/ir/irparser.h" + +TEST(Converters, ATenChunkConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=3] + %4 : Tensor = aten::select(%0, %2, %3) + return (%3))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1,100}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} \ No newline at end of file From 2d610776ac3d847a788b88543932ae952af1b65f Mon Sep 17 00:00:00 2001 From: apbose Date: Mon, 3 Apr 2023 12:00:05 -0700 Subject: [PATCH 2/7] chunk test linting error --- core/conversion/converters/impl/chunk.cpp | 1 - tests/core/conversion/converters/test_chunk.cpp | 1 - 2 files changed, 2 deletions(-) diff --git a/core/conversion/converters/impl/chunk.cpp b/core/conversion/converters/impl/chunk.cpp index 21bf7848ec..753f69325c 100644 --- a/core/conversion/converters/impl/chunk.cpp +++ b/core/conversion/converters/impl/chunk.cpp @@ -19,7 +19,6 @@ auto cat_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() bool dynamic_shape = ctx->input_is_dynamic; int size = in->getDimensions().nbDims; int maxDim = static_cast(in->getDimensions().d[axis]); - if(dim < 0) { dim = tensors[0]->getDimensions().nbDims + dim; } diff --git a/tests/core/conversion/converters/test_chunk.cpp b/tests/core/conversion/converters/test_chunk.cpp index bd6b7cec1e..67e7388940 100644 --- a/tests/core/conversion/converters/test_chunk.cpp +++ b/tests/core/conversion/converters/test_chunk.cpp @@ -17,7 +17,6 @@ TEST(Converters, ATenChunkConvertsCorrectly) { auto g = std::make_shared(); torch::jit::parseIR(graph, g.get()); - auto in = at::randint(1, 10, {1,100}, {at::kCUDA}); auto jit_in = at::clone(in); From 170e56641614676bd5495575692f50f6d2e97525 Mon Sep 17 00:00:00 2001 From: apbose Date: Mon, 3 Apr 2023 12:02:05 -0700 Subject: [PATCH 3/7] chunk test linting error --- core/conversion/converters/impl/chunk.cpp | 2 +- tests/core/conversion/converters/test_chunk.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/core/conversion/converters/impl/chunk.cpp b/core/conversion/converters/impl/chunk.cpp index 753f69325c..862e626c81 100644 --- a/core/conversion/converters/impl/chunk.cpp +++ b/core/conversion/converters/impl/chunk.cpp @@ -26,7 +26,7 @@ auto cat_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() TORCHTRT_ASSERT(in.d[dim] != -1, "Can't chunk on dynamic shape dimension!"); } if (chunks > in.d[dim]) { - LOG_WARNING("The chunks size" << chunks << "along dimension" << dim << "is greater than tensor with size" << in->getDimensions + LOG_WARNING("The chunks size" << chunks << "along dimension" << dim << "is greater than tensor with size" << in->getDimensions << "it will default to dimension" << in.d[dim]) } int step = (input_val.shape[dim] + chunks - 1) / chunks diff --git a/tests/core/conversion/converters/test_chunk.cpp b/tests/core/conversion/converters/test_chunk.cpp index 67e7388940..f76684812e 100644 --- a/tests/core/conversion/converters/test_chunk.cpp +++ b/tests/core/conversion/converters/test_chunk.cpp @@ -17,7 +17,7 @@ TEST(Converters, ATenChunkConvertsCorrectly) { auto g = std::make_shared(); torch::jit::parseIR(graph, g.get()); - auto in = at::randint(1, 10, {1,100}, {at::kCUDA}); + auto in = at::randint(1, 10, {1, 100}, {at::kCUDA}); auto jit_in = at::clone(in); auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); From 25c2bb9aab4e36f55638d9c90eda4f79e7d1451c Mon Sep 17 00:00:00 2001 From: apbose Date: Mon, 3 Apr 2023 17:16:33 -0700 Subject: [PATCH 4/7] aten::chunk converter correction and test changes --- core/conversion/converters/impl/chunk.cpp | 64 ++++++++++++------- .../core/conversion/converters/test_chunk.cpp | 11 ++-- 2 files changed, 46 insertions(+), 29 deletions(-) diff --git a/core/conversion/converters/impl/chunk.cpp b/core/conversion/converters/impl/chunk.cpp index 862e626c81..ce4c996073 100644 --- a/core/conversion/converters/impl/chunk.cpp +++ b/core/conversion/converters/impl/chunk.cpp @@ -18,19 +18,25 @@ auto cat_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() auto dim = args[2].unwrapToInt(); bool dynamic_shape = ctx->input_is_dynamic; int size = in->getDimensions().nbDims; - int maxDim = static_cast(in->getDimensions().d[axis]); + int maxDim = static_cast(in->getDimensions().d[dim]); + + c10::ListTypePtr lt = n->output()->type()->expect(); + c10::TypePtr elementType = lt->getElementType(); + + int offset = 0; if(dim < 0) { - dim = tensors[0]->getDimensions().nbDims + dim; + dim = in->getDimensions().nbDims + dim; } if (dynamic_shape) { - TORCHTRT_ASSERT(in.d[dim] != -1, "Can't chunk on dynamic shape dimension!"); + TORCHTRT_ASSERT(in->getDimensions().d[dim] != -1, "Can't chunk on dynamic shape dimension!"); } - if (chunks > in.d[dim]) { - LOG_WARNING("The chunks size" << chunks << "along dimension" << dim << "is greater than tensor with size" << in->getDimensions - << "it will default to dimension" << in.d[dim]) + if (chunks > in->getDimensions().d[dim]) { + LOG_WARNING("The chunks size" << chunks << "along dimension" << dim << "is greater than tensor with size" << in->getDimensions().d[dim] + << "it will default to dimension" << in->getDimensions().d[dim]); } - int step = (input_val.shape[dim] + chunks - 1) / chunks + int step = (maxDim + chunks - 1) / chunks; nvinfer1::Dims start_, size_, stride_; + int nbdims = in->getDimensions().nbDims; start_.nbDims = nbdims; size_.nbDims = nbdims; stride_.nbDims = nbdims; @@ -39,25 +45,35 @@ auto cat_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() int endIdx = maxDim; for (int i = 0; i < nbdims; i++) { - if (i == axis) { - start_.d[i] = startIdx; - size_.d[i] = (endIdx - startIdx - 1) / step + 1; - stride_.d[i] = step; - } else { - start_.d[i] = 0; - size_.d[i] = in.d[i]; // for static - stride_.d[i] = 1; - } + start_.d[i] = 0; + size_.d[i] = 0; + stride_.d[i] = 1; } // update slice layer - if(!dynamic_shape): - auto slice_layer = ctx->net->addSlice(*in, start_, size_, stride_); - LOG_DEBUG("start_:" << start_); - LOG_DEBUG("size_:" << size_); - LOG_DEBUG("stride_:" << stride_); - auto slice_out = slice_layer->getOutput(0); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], slice_out); - LOG_DEBUG("Slice layer output shape: " << out->getDimensions()); + auto list = c10::impl::GenericList(elementType); + list.reserve(chunks); + if(!dynamic_shape) { + for (int chunk = 0; chunk < chunks; chunk++) { + for (int i = 0; i < nbdims; i++) { + if (i == dim) { + start_.d[i] = offset; + size_.d[i] = std::min(step, maxDim - offset); + } + } + LOG_DEBUG("start_:" << start_); + LOG_DEBUG("size_:" << size_); + LOG_DEBUG("stride_:" << stride_); + auto slice_layer = ctx->net->addSlice(*in, start_, size_, stride_); + auto tensor_holder = TensorContainer(); + tensor_holder.hold_tensor(slice_layer->getOutput(0)); + auto ival = c10::IValue(std::move(c10::make_intrusive(tensor_holder))); + list.emplace_back(ival); + offset = offset + step; + } + } + auto split_output_ivalue = std::move(torch::jit::IValue(list)); + ctx->AssociateValueAndIValue(n->outputs()[0], split_output_ivalue); + return true; }}); // clang-format on } // namespace diff --git a/tests/core/conversion/converters/test_chunk.cpp b/tests/core/conversion/converters/test_chunk.cpp index f76684812e..eccefbea99 100644 --- a/tests/core/conversion/converters/test_chunk.cpp +++ b/tests/core/conversion/converters/test_chunk.cpp @@ -9,15 +9,16 @@ TEST(Converters, ATenChunkConvertsCorrectly) { const auto graph = R"IR( graph(%0 : Tensor): - %2 : int = prim::Constant[value=1]() - %3 : int = prim::Constant[value=3] - %4 : Tensor = aten::select(%0, %2, %3) - return (%3))IR"; + %2 : int = prim::Constant[value=6]() + %3 : int = prim::Constant[value=0]() + %4 : Tensor[] = aten::chunk(%0, %2, %3) + %5 : Tensor, %6 : Tensor, %7 : Tensor, %8 : Tensor, %9 : Tensor, %10 : Tensor = prim::ListUnpack(%4) + return (%5, %6, %7, %8, %9, %10))IR"; auto g = std::make_shared(); torch::jit::parseIR(graph, g.get()); - auto in = at::randint(1, 10, {1, 100}, {at::kCUDA}); + auto in = at::randint(1, 10, {12}, {at::kCUDA}); auto jit_in = at::clone(in); auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); From fffe1b3783ddc743090963c48082a19f214d35d6 Mon Sep 17 00:00:00 2001 From: apbose Date: Mon, 3 Apr 2023 17:18:07 -0700 Subject: [PATCH 5/7] aten::chunk converter correction and test changes --- core/conversion/converters/impl/chunk.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/conversion/converters/impl/chunk.cpp b/core/conversion/converters/impl/chunk.cpp index ce4c996073..a7191133fb 100644 --- a/core/conversion/converters/impl/chunk.cpp +++ b/core/conversion/converters/impl/chunk.cpp @@ -66,7 +66,7 @@ auto cat_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() auto slice_layer = ctx->net->addSlice(*in, start_, size_, stride_); auto tensor_holder = TensorContainer(); tensor_holder.hold_tensor(slice_layer->getOutput(0)); - auto ival = c10::IValue(std::move(c10::make_intrusive(tensor_holder))); + auto ival = c10::IValue(std::move(c10::make_intrusive(tensor_holder))); list.emplace_back(ival); offset = offset + step; } From c3ef998212e6ba03039244832bef5708e3382f5c Mon Sep 17 00:00:00 2001 From: apbose Date: Mon, 3 Apr 2023 17:19:35 -0700 Subject: [PATCH 6/7] adding aten::chunk in build --- core/conversion/converters/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index 354f17a734..518c6a2ded 100755 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -57,6 +57,7 @@ cc_library( "impl/batch_norm.cpp", "impl/bitwise.cpp", "impl/cast.cpp", + "impl/chunk.cpp", "impl/concat.cpp", "impl/constant.cpp", "impl/constant_pad.cpp", From 19706d38446e2599178a5dcad0e2142ad07d5cea Mon Sep 17 00:00:00 2001 From: apbose Date: Fri, 2 Jun 2023 12:18:31 -0700 Subject: [PATCH 7/7] adding chunk implementation to cmakelist --- core/conversion/converters/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/core/conversion/converters/CMakeLists.txt b/core/conversion/converters/CMakeLists.txt index c90a81a6cc..392a82b744 100644 --- a/core/conversion/converters/CMakeLists.txt +++ b/core/conversion/converters/CMakeLists.txt @@ -7,6 +7,7 @@ target_sources(${lib_name} "${CMAKE_CURRENT_SOURCE_DIR}/impl/activation.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/impl/batch_norm.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/impl/cast.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/impl/chunk.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/impl/concat.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/impl/constant.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/impl/constant_pad.cpp"