From 482c9367cef5aa337c9f1a01c77fbd04456bacd5 Mon Sep 17 00:00:00 2001 From: Brenton Chu Date: Mon, 15 Aug 2022 16:32:52 -0700 Subject: [PATCH 1/7] added interleave_repeat int repeats converter --- core/conversion/converters/BUILD | 1 + core/conversion/converters/impl/repeat.cpp | 89 ++++++++++++++++++++++ 2 files changed, 90 insertions(+) create mode 100644 core/conversion/converters/impl/repeat.cpp diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index 95dde838dc..a5651e9f51 100755 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -75,6 +75,7 @@ cc_library( "impl/quantization.cpp", "impl/reduce.cpp", "impl/reflection_pad.cpp", + "impl/repeat.cpp", "impl/replication_pad.cpp", "impl/select.cpp", "impl/shuffle.cpp", diff --git a/core/conversion/converters/impl/repeat.cpp b/core/conversion/converters/impl/repeat.cpp new file mode 100644 index 0000000000..6335d2afb9 --- /dev/null +++ b/core/conversion/converters/impl/repeat.cpp @@ -0,0 +1,89 @@ +#include "core/conversion/converters/converters.h" +#include "core/util/prelude.h" + +#include + +namespace torch_tensorrt { +namespace core { +namespace conversion { +namespace converters { +namespace impl { + +auto repeatinterleave TORCHTRT_UNUSED = + RegisterNodeConversionPatterns() + .pattern({ + "aten::repeat_interleave(Tensor self, int repeats, int? dim=None, int? output_size=None) -> Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto repeats = args[1].unwrapToScalar().to(); + + auto input_shape = self->getDimensions(); + + int dim; + if (args[2].IValue()->isNone()) { + dim = 0; + + // Flatten self tensor + int size = 0; + for (int i = 0; i < input_shape.nbDims; i++) { + size += input_shape.d[i]; + } + auto flatten = ctx->net->addShuffle(*self); + TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); + flatten.setReshapeDimensions(util::toDims(std::vector({size}))); + self = flatten->getOutput(0); + input_shape = self->getDimensions(); + } + else { + dim = args[2].unwrapToScalar().to(); + } + + // Insert singleton dimension after desired repeat dimension + std::vector repeat_shape_vec; + for (int j = 0; j < input_shape.nbDims; j++) { + repeat_shape_vec.push_back(input_shape.d[j]); + if (j == dim) { + repeate_shape_vec.push_back(1); + } + } + auto expand = ctx->net->addShuffle(*self); + TORCHTRT_CHECK(expand, "Unable to create shuffle layer from node: " << *n); + auto repeat_shape = util::toDims(repeat_shape_vec); + expand_shape.setReshapeDimensions(repeat_shape); + + // Expand on newly created singleton dimension + repeat_shape.d[dim + 1] = repeats; + std::vector start_vec(expanded_shape.nbDims, 0); + auto start = util::toDims(start_vec); + + std::vector strides_vec(expanded_shape.nbDims, 1); + strides_vec[dim + 1] = 0; + auto strides = util::toDims(strides_vec); + + auto slice = ctx->net->addSlice(*expand->getOutput(0), start, repeat_shape, strides); + + // Collapse repeated dimension back into desired dimension + std::vector collapse_shape_vec; + for (int k = 0; k < expand_shape.nbDim; k++) { + if (k == dim) { + collapse_shape_vec.push_back(expand_shape.d[k] * expand_shape.d[++k]); + } else { + collapse_shape_vec.push_back(expand_shape.d[k]); + } + } + auto collapse = ctx->net->addShuffle(*slice->getOutput(0)); + TORCHTRT_CHECK(collapse, "Unable to create shuffle layer from node: " << *n); + collapse.setReshapeDimensions(util::toDims(collapse_shape_vec)); + + collapse->setName(util::node_info(n).c_str()); + auto out_tensor = ctx->AssociateValueAnTensor(n->outputs()[0], collapse->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions()); + + return true; + + } +} // namespace impl +} // namespace impl +} // namespace converters +} // namespace conversion +} // namespace core From f1d492436c148806ab2efd1ba5ce50246b39cb27 Mon Sep 17 00:00:00 2001 From: Brenton Chu Date: Wed, 17 Aug 2022 18:55:13 -0700 Subject: [PATCH 2/7] fixed compile time errors --- core/conversion/converters/impl/repeat.cpp | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/core/conversion/converters/impl/repeat.cpp b/core/conversion/converters/impl/repeat.cpp index 6335d2afb9..29e3ad89dc 100644 --- a/core/conversion/converters/impl/repeat.cpp +++ b/core/conversion/converters/impl/repeat.cpp @@ -30,7 +30,7 @@ auto repeatinterleave TORCHTRT_UNUSED = } auto flatten = ctx->net->addShuffle(*self); TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); - flatten.setReshapeDimensions(util::toDims(std::vector({size}))); + flatten->setReshapeDimensions(util::toDims(std::vector({size}))); self = flatten->getOutput(0); input_shape = self->getDimensions(); } @@ -43,20 +43,20 @@ auto repeatinterleave TORCHTRT_UNUSED = for (int j = 0; j < input_shape.nbDims; j++) { repeat_shape_vec.push_back(input_shape.d[j]); if (j == dim) { - repeate_shape_vec.push_back(1); + repeat_shape_vec.push_back(1); } } auto expand = ctx->net->addShuffle(*self); TORCHTRT_CHECK(expand, "Unable to create shuffle layer from node: " << *n); auto repeat_shape = util::toDims(repeat_shape_vec); - expand_shape.setReshapeDimensions(repeat_shape); + expand->setReshapeDimensions(repeat_shape); // Expand on newly created singleton dimension repeat_shape.d[dim + 1] = repeats; - std::vector start_vec(expanded_shape.nbDims, 0); + std::vector start_vec(repeat_shape.nbDims, 0); auto start = util::toDims(start_vec); - std::vector strides_vec(expanded_shape.nbDims, 1); + std::vector strides_vec(repeat_shape.nbDims, 1); strides_vec[dim + 1] = 0; auto strides = util::toDims(strides_vec); @@ -64,24 +64,24 @@ auto repeatinterleave TORCHTRT_UNUSED = // Collapse repeated dimension back into desired dimension std::vector collapse_shape_vec; - for (int k = 0; k < expand_shape.nbDim; k++) { + for (int k = 0; k < repeat_shape.nbDims; k++) { if (k == dim) { - collapse_shape_vec.push_back(expand_shape.d[k] * expand_shape.d[++k]); + collapse_shape_vec.push_back(repeat_shape.d[k] * repeat_shape.d[++k]); } else { - collapse_shape_vec.push_back(expand_shape.d[k]); + collapse_shape_vec.push_back(repeat_shape.d[k]); } } auto collapse = ctx->net->addShuffle(*slice->getOutput(0)); TORCHTRT_CHECK(collapse, "Unable to create shuffle layer from node: " << *n); - collapse.setReshapeDimensions(util::toDims(collapse_shape_vec)); + collapse->setReshapeDimensions(util::toDims(collapse_shape_vec)); collapse->setName(util::node_info(n).c_str()); - auto out_tensor = ctx->AssociateValueAnTensor(n->outputs()[0], collapse->getOutput(0)); + auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], collapse->getOutput(0)); LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions()); return true; - } + }}); } // namespace impl } // namespace impl } // namespace converters From 214eb75b1752176dc638c7f8c8d97f5da0cc6d70 Mon Sep 17 00:00:00 2001 From: Brenton Chu Date: Fri, 19 Aug 2022 10:49:53 -0700 Subject: [PATCH 3/7] added repeat_interleave tests, moved converter to expand file --- core/conversion/converters/impl/expand.cpp | 74 +++++++++++++++ core/conversion/converters/impl/repeat.cpp | 89 ------------------- .../conversion/converters/test_expand.cpp | 56 ++++++++++++ 3 files changed, 130 insertions(+), 89 deletions(-) delete mode 100644 core/conversion/converters/impl/repeat.cpp diff --git a/core/conversion/converters/impl/expand.cpp b/core/conversion/converters/impl/expand.cpp index e090261a2a..3b2c02cee2 100644 --- a/core/conversion/converters/impl/expand.cpp +++ b/core/conversion/converters/impl/expand.cpp @@ -285,6 +285,80 @@ auto expand_registrations TORCHTRT_UNUSED = return true; }}); +auto repeatinterleave TORCHTRT_UNUSED = + RegisterNodeConversionPatterns() + .pattern({ + "aten::repeat_interleave.self_int(Tensor self, int repeats, int? dim=None, *, int? output_size=None) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto repeats = args[1].unwrapToScalar().to(); + + auto input_shape = self->getDimensions(); + + int dim; + if (args[2].IValue()->isNone()) { + dim = 0; + + // Flatten self tensor + int size = 0; + for (int i = 0; i < input_shape.nbDims; i++) { + size += input_shape.d[i]; + } + auto flatten = ctx->net->addShuffle(*self); + TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); + flatten->setReshapeDimensions(util::toDims(std::vector({size}))); + self = flatten->getOutput(0); + input_shape = self->getDimensions(); + } + else { + dim = args[2].unwrapToScalar().to(); + } + + // Insert singleton dimension after desired repeat dimension + std::vector repeat_shape_vec; + for (int j = 0; j < input_shape.nbDims; j++) { + repeat_shape_vec.push_back(input_shape.d[j]); + if (j == dim) { + repeat_shape_vec.push_back(1); + } + } + auto expand = ctx->net->addShuffle(*self); + TORCHTRT_CHECK(expand, "Unable to create shuffle layer from node: " << *n); + auto repeat_shape = util::toDims(repeat_shape_vec); + expand->setReshapeDimensions(repeat_shape); + + // Expand on newly created singleton dimension + repeat_shape.d[dim + 1] = repeats; + std::vector start_vec(repeat_shape.nbDims, 0); + auto start = util::toDims(start_vec); + + std::vector strides_vec(repeat_shape.nbDims, 1); + strides_vec[dim + 1] = 0; + auto strides = util::toDims(strides_vec); + + auto slice = ctx->net->addSlice(*expand->getOutput(0), start, repeat_shape, strides); + + // Collapse repeated dimension back into desired dimension + std::vector collapse_shape_vec; + for (int k = 0; k < repeat_shape.nbDims; k++) { + if (k == dim) { + collapse_shape_vec.push_back(repeat_shape.d[k] * repeat_shape.d[++k]); + } else { + collapse_shape_vec.push_back(repeat_shape.d[k]); + } + } + auto collapse = ctx->net->addShuffle(*slice->getOutput(0)); + TORCHTRT_CHECK(collapse, "Unable to create shuffle layer from node: " << *n); + collapse->setReshapeDimensions(util::toDims(collapse_shape_vec)); + + collapse->setName(util::node_info(n).c_str()); + auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], collapse->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions()); + + return true; + + }}); + } // namespace } // namespace impl } // namespace converters diff --git a/core/conversion/converters/impl/repeat.cpp b/core/conversion/converters/impl/repeat.cpp deleted file mode 100644 index 29e3ad89dc..0000000000 --- a/core/conversion/converters/impl/repeat.cpp +++ /dev/null @@ -1,89 +0,0 @@ -#include "core/conversion/converters/converters.h" -#include "core/util/prelude.h" - -#include - -namespace torch_tensorrt { -namespace core { -namespace conversion { -namespace converters { -namespace impl { - -auto repeatinterleave TORCHTRT_UNUSED = - RegisterNodeConversionPatterns() - .pattern({ - "aten::repeat_interleave(Tensor self, int repeats, int? dim=None, int? output_size=None) -> Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto repeats = args[1].unwrapToScalar().to(); - - auto input_shape = self->getDimensions(); - - int dim; - if (args[2].IValue()->isNone()) { - dim = 0; - - // Flatten self tensor - int size = 0; - for (int i = 0; i < input_shape.nbDims; i++) { - size += input_shape.d[i]; - } - auto flatten = ctx->net->addShuffle(*self); - TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); - flatten->setReshapeDimensions(util::toDims(std::vector({size}))); - self = flatten->getOutput(0); - input_shape = self->getDimensions(); - } - else { - dim = args[2].unwrapToScalar().to(); - } - - // Insert singleton dimension after desired repeat dimension - std::vector repeat_shape_vec; - for (int j = 0; j < input_shape.nbDims; j++) { - repeat_shape_vec.push_back(input_shape.d[j]); - if (j == dim) { - repeat_shape_vec.push_back(1); - } - } - auto expand = ctx->net->addShuffle(*self); - TORCHTRT_CHECK(expand, "Unable to create shuffle layer from node: " << *n); - auto repeat_shape = util::toDims(repeat_shape_vec); - expand->setReshapeDimensions(repeat_shape); - - // Expand on newly created singleton dimension - repeat_shape.d[dim + 1] = repeats; - std::vector start_vec(repeat_shape.nbDims, 0); - auto start = util::toDims(start_vec); - - std::vector strides_vec(repeat_shape.nbDims, 1); - strides_vec[dim + 1] = 0; - auto strides = util::toDims(strides_vec); - - auto slice = ctx->net->addSlice(*expand->getOutput(0), start, repeat_shape, strides); - - // Collapse repeated dimension back into desired dimension - std::vector collapse_shape_vec; - for (int k = 0; k < repeat_shape.nbDims; k++) { - if (k == dim) { - collapse_shape_vec.push_back(repeat_shape.d[k] * repeat_shape.d[++k]); - } else { - collapse_shape_vec.push_back(repeat_shape.d[k]); - } - } - auto collapse = ctx->net->addShuffle(*slice->getOutput(0)); - TORCHTRT_CHECK(collapse, "Unable to create shuffle layer from node: " << *n); - collapse->setReshapeDimensions(util::toDims(collapse_shape_vec)); - - collapse->setName(util::node_info(n).c_str()); - auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], collapse->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions()); - - return true; - - }}); -} // namespace impl -} // namespace impl -} // namespace converters -} // namespace conversion -} // namespace core diff --git a/tests/core/conversion/converters/test_expand.cpp b/tests/core/conversion/converters/test_expand.cpp index bf62266f32..7a329a37cf 100644 --- a/tests/core/conversion/converters/test_expand.cpp +++ b/tests/core/conversion/converters/test_expand.cpp @@ -445,3 +445,59 @@ TEST(Converters, ATenRepeatExtraDimsConvertsCorrectlyWithDynamicInput) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); } + +TEST(Converters, ATenRepeatInterleaveScalarDimConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : int = prim::Constant[value=1]() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {3, 2, 4}, {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::RunGraphEngineDynamic(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)); +} + +TEST(Converters, ATenRepeatInterleaveScalarNoDimConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : None = prim::Constant() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {3, 2, 4}, {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::RunGraphEngineDynamic(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)); +} From 04029397762e7c7331966d5697fe18e539a64bc7 Mon Sep 17 00:00:00 2001 From: Brenton Chu Date: Fri, 19 Aug 2022 11:37:16 -0700 Subject: [PATCH 4/7] repeat_interleave passing tests for static input --- core/conversion/converters/BUILD | 1 - core/conversion/converters/impl/expand.cpp | 4 +- .../conversion/converters/test_expand.cpp | 64 +++++++++++++++++-- 3 files changed, 62 insertions(+), 7 deletions(-) diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index a5651e9f51..95dde838dc 100755 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -75,7 +75,6 @@ cc_library( "impl/quantization.cpp", "impl/reduce.cpp", "impl/reflection_pad.cpp", - "impl/repeat.cpp", "impl/replication_pad.cpp", "impl/select.cpp", "impl/shuffle.cpp", diff --git a/core/conversion/converters/impl/expand.cpp b/core/conversion/converters/impl/expand.cpp index 3b2c02cee2..05b9e07bbf 100644 --- a/core/conversion/converters/impl/expand.cpp +++ b/core/conversion/converters/impl/expand.cpp @@ -300,9 +300,9 @@ auto repeatinterleave TORCHTRT_UNUSED = dim = 0; // Flatten self tensor - int size = 0; + int size = 1; for (int i = 0; i < input_shape.nbDims; i++) { - size += input_shape.d[i]; + size *= input_shape.d[i]; } auto flatten = ctx->net->addShuffle(*self); TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); diff --git a/tests/core/conversion/converters/test_expand.cpp b/tests/core/conversion/converters/test_expand.cpp index 7a329a37cf..b5278549e6 100644 --- a/tests/core/conversion/converters/test_expand.cpp +++ b/tests/core/conversion/converters/test_expand.cpp @@ -459,7 +459,7 @@ TEST(Converters, ATenRepeatInterleaveScalarDimConvertsCorrectly) { torch::jit::parseIR(graph, g.get()); - auto in = at::randint(1, 10, {3, 2, 4}, {at::kCUDA}); + auto in = at::randint(1, 10, {1, 3}, {at::kCUDA}); auto jit_in = at::clone(in); auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); @@ -467,7 +467,7 @@ TEST(Converters, ATenRepeatInterleaveScalarDimConvertsCorrectly) { auto trt_in = at::clone(in); params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); auto trt = trt_results[0].reshape(jit_results[0].sizes()); @@ -487,7 +487,7 @@ TEST(Converters, ATenRepeatInterleaveScalarNoDimConvertsCorrectly) { torch::jit::parseIR(graph, g.get()); - auto in = at::randint(1, 10, {3, 2, 4}, {at::kCUDA}); + auto in = at::randint(1, 10, {1, 3}, {at::kCUDA}); auto jit_in = at::clone(in); auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); @@ -495,7 +495,63 @@ TEST(Converters, ATenRepeatInterleaveScalarNoDimConvertsCorrectly) { auto trt_in = at::clone(in); params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}); + 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)); +} + +TEST(Converters, ATenRepeatInterleave3dScalarDimConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : int = prim::Constant[value=1]() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {2, 3, 2}, {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)); +} + +TEST(Converters, ATenRepeatInterleave3dScalarNoDimConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : None = prim::Constant() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {2, 3, 2}, {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()); From 47e0a54ab607e3dfa01560cbb19b9df0570cd296 Mon Sep 17 00:00:00 2001 From: Brenton Chu Date: Wed, 24 Aug 2022 17:26:21 -0700 Subject: [PATCH 5/7] implementation and tests for dynamic input repeat_interleave --- core/conversion/converters/impl/expand.cpp | 54 ++++++--- .../conversion/converters/test_expand.cpp | 112 ++++++++++++++++++ 2 files changed, 151 insertions(+), 15 deletions(-) diff --git a/core/conversion/converters/impl/expand.cpp b/core/conversion/converters/impl/expand.cpp index 05b9e07bbf..2501fd976d 100644 --- a/core/conversion/converters/impl/expand.cpp +++ b/core/conversion/converters/impl/expand.cpp @@ -300,9 +300,15 @@ auto repeatinterleave TORCHTRT_UNUSED = dim = 0; // Flatten self tensor - int size = 1; - for (int i = 0; i < input_shape.nbDims; i++) { - size *= input_shape.d[i]; + int size; + if (ctx->input_is_dynamic) { + // Set size to -1 if input is dynamic + size = -1; + } else { + size = 1; + for (int i = 0; i < input_shape.nbDims; i++) { + size *= input_shape.d[i]; + } } auto flatten = ctx->net->addShuffle(*self); TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); @@ -324,27 +330,45 @@ auto repeatinterleave TORCHTRT_UNUSED = } auto expand = ctx->net->addShuffle(*self); TORCHTRT_CHECK(expand, "Unable to create shuffle layer from node: " << *n); - auto repeat_shape = util::toDims(repeat_shape_vec); - expand->setReshapeDimensions(repeat_shape); + auto repeat_shape_dims = util::toDims(repeat_shape_vec); + expand->setReshapeDimensions(repeat_shape_dims); // Expand on newly created singleton dimension - repeat_shape.d[dim + 1] = repeats; - std::vector start_vec(repeat_shape.nbDims, 0); - auto start = util::toDims(start_vec); + repeat_shape_dims.d[dim + 1] = repeats; + std::vector start_vec(repeat_shape_dims.nbDims, 0); + auto start_dims = util::toDims(start_vec); - std::vector strides_vec(repeat_shape.nbDims, 1); + std::vector strides_vec(repeat_shape_dims.nbDims, 1); strides_vec[dim + 1] = 0; - auto strides = util::toDims(strides_vec); - - auto slice = ctx->net->addSlice(*expand->getOutput(0), start, repeat_shape, strides); + auto strides_dims = util::toDims(strides_vec); + + auto slice = ctx->net->addSlice(*expand->getOutput(0), start_dims, repeat_shape_dims, strides_dims); + + if (ctx->input_is_dynamic) { + auto start_tensor = tensor_to_const(ctx, torch::tensor(start_vec, torch::kInt32)); + + auto expand_output_shape = ctx->net->addShape(*expand->getOutput(0))->getOutput(0); + std::vector repeat_const_vec(repeat_shape_dims.nbDims, 1); + repeat_const_vec[dim + 1] = repeats; + auto repeat_const = tensor_to_const(ctx, torch::tensor(repeat_const_vec, torch::kInt32)); + auto repeat_shape_tensor = ctx->net->addElementWise(*expand_output_shape, *repeat_const, nvinfer1::ElementWiseOperation::kPROD)->getOutput(0); + + auto strides_tensor = tensor_to_const(ctx, torch::tensor(strides_vec, torch::kInt32)); + slice->setInput(1, *start_tensor); + slice->setInput(2, *repeat_shape_tensor); + slice->setInput(3, *strides_tensor); + } // Collapse repeated dimension back into desired dimension std::vector collapse_shape_vec; - for (int k = 0; k < repeat_shape.nbDims; k++) { + for (int k = 0; k < repeat_shape_dims.nbDims; k++) { if (k == dim) { - collapse_shape_vec.push_back(repeat_shape.d[k] * repeat_shape.d[++k]); + int64_t collapse_dim = repeat_shape_dims.d[k] * repeat_shape_dims.d[++k]; + // Set dim size to -1 if repeat is being done on dynamic dim + collapse_dim = std::max(collapse_dim, (int64_t)-1); + collapse_shape_vec.push_back(collapse_dim); } else { - collapse_shape_vec.push_back(repeat_shape.d[k]); + collapse_shape_vec.push_back(repeat_shape_dims.d[k]); } } auto collapse = ctx->net->addShuffle(*slice->getOutput(0)); diff --git a/tests/core/conversion/converters/test_expand.cpp b/tests/core/conversion/converters/test_expand.cpp index b5278549e6..53630b661a 100644 --- a/tests/core/conversion/converters/test_expand.cpp +++ b/tests/core/conversion/converters/test_expand.cpp @@ -474,6 +474,34 @@ TEST(Converters, ATenRepeatInterleaveScalarDimConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); } +TEST(Converters, ATenRepeatInterleaveScalarDimConvertsCorrectlyWithDynamicInput) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : int = prim::Constant[value=1]() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1, 3}, {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::RunGraphEngineDynamic(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)); +} + TEST(Converters, ATenRepeatInterleaveScalarNoDimConvertsCorrectly) { const auto graph = R"IR( graph(%x.1 : Tensor): @@ -502,6 +530,34 @@ TEST(Converters, ATenRepeatInterleaveScalarNoDimConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); } +TEST(Converters, ATenRepeatInterleaveScalarNoDimConvertsCorrectlyWithDynamicInput) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : None = prim::Constant() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1, 3}, {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::RunGraphEngineDynamic(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)); +} + TEST(Converters, ATenRepeatInterleave3dScalarDimConvertsCorrectly) { const auto graph = R"IR( graph(%x.1 : Tensor): @@ -530,6 +586,34 @@ TEST(Converters, ATenRepeatInterleave3dScalarDimConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); } +TEST(Converters, ATenRepeatInterleave3dScalarDimConvertsCorrectlyWithDynamicInput) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : int = prim::Constant[value=1]() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {2, 3, 2}, {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::RunGraphEngineDynamic(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)); +} + TEST(Converters, ATenRepeatInterleave3dScalarNoDimConvertsCorrectly) { const auto graph = R"IR( graph(%x.1 : Tensor): @@ -557,3 +641,31 @@ TEST(Converters, ATenRepeatInterleave3dScalarNoDimConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); } + +TEST(Converters, ATenRepeatInterleave3dScalarNoDimConvertsCorrectlyWithDynamicInput) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : None = prim::Constant() + %4 : None = prim::Constant() + %5 : Tensor = aten::repeat_interleave(%x.1, %2, %3, %4) + return (%5))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {2, 3, 2}, {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::RunGraphEngineDynamic(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)); +} From cbdde93576ee0acb14e70786d0f1066d2ee82984 Mon Sep 17 00:00:00 2001 From: Brenton Chu Date: Thu, 25 Aug 2022 21:33:14 +0000 Subject: [PATCH 6/7] dynamic shape checks --- core/conversion/converters/impl/expand.cpp | 23 ++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/core/conversion/converters/impl/expand.cpp b/core/conversion/converters/impl/expand.cpp index 2501fd976d..79c0703a69 100644 --- a/core/conversion/converters/impl/expand.cpp +++ b/core/conversion/converters/impl/expand.cpp @@ -283,12 +283,9 @@ auto expand_registrations TORCHTRT_UNUSED = LOG_DEBUG("Repeat layer output tensor shape: " << out->getDimensions()); return true; - }}); - -auto repeatinterleave TORCHTRT_UNUSED = - RegisterNodeConversionPatterns() - .pattern({ - "aten::repeat_interleave.self_int(Tensor self, int repeats, int? dim=None, *, int? output_size=None) -> (Tensor)", + }}) + .pattern( + {"aten::repeat_interleave.self_int(Tensor self, int repeats, int? dim=None, *, int? output_size=None) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { auto self = args[0].ITensorOrFreeze(ctx); auto repeats = args[1].unwrapToScalar().to(); @@ -320,6 +317,20 @@ auto repeatinterleave TORCHTRT_UNUSED = dim = args[2].unwrapToScalar().to(); } + if (ctx->input_is_dynamic) { + int dynamic_dims = 0; + for (int idx = 0; idx < input_shape.nbDims; idx++) { + if (input_shape.d[idx] == -1) { + dynamic_dims++; + } + } + + if (dynamic_dims > 1) { + TORCHTRT_THROW_ERROR( + "Repeat_interleave is currently not supported when target shape contains more than one dynamic dimension"); + } + } + // Insert singleton dimension after desired repeat dimension std::vector repeat_shape_vec; for (int j = 0; j < input_shape.nbDims; j++) { From ee7b5a7715fb0004bb79abf6b458b6b46048c58d Mon Sep 17 00:00:00 2001 From: Brenton Chu Date: Thu, 25 Aug 2022 15:21:37 -0700 Subject: [PATCH 7/7] reformatting --- core/conversion/converters/impl/expand.cpp | 215 +++++++++++---------- 1 file changed, 108 insertions(+), 107 deletions(-) diff --git a/core/conversion/converters/impl/expand.cpp b/core/conversion/converters/impl/expand.cpp index 79c0703a69..e379614ad3 100644 --- a/core/conversion/converters/impl/expand.cpp +++ b/core/conversion/converters/impl/expand.cpp @@ -286,113 +286,114 @@ auto expand_registrations TORCHTRT_UNUSED = }}) .pattern( {"aten::repeat_interleave.self_int(Tensor self, int repeats, int? dim=None, *, int? output_size=None) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto repeats = args[1].unwrapToScalar().to(); - - auto input_shape = self->getDimensions(); - - int dim; - if (args[2].IValue()->isNone()) { - dim = 0; - - // Flatten self tensor - int size; - if (ctx->input_is_dynamic) { - // Set size to -1 if input is dynamic - size = -1; - } else { - size = 1; - for (int i = 0; i < input_shape.nbDims; i++) { - size *= input_shape.d[i]; - } - } - auto flatten = ctx->net->addShuffle(*self); - TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); - flatten->setReshapeDimensions(util::toDims(std::vector({size}))); - self = flatten->getOutput(0); - input_shape = self->getDimensions(); - } - else { - dim = args[2].unwrapToScalar().to(); - } - - if (ctx->input_is_dynamic) { - int dynamic_dims = 0; - for (int idx = 0; idx < input_shape.nbDims; idx++) { - if (input_shape.d[idx] == -1) { - dynamic_dims++; - } - } - - if (dynamic_dims > 1) { - TORCHTRT_THROW_ERROR( - "Repeat_interleave is currently not supported when target shape contains more than one dynamic dimension"); - } - } - - // Insert singleton dimension after desired repeat dimension - std::vector repeat_shape_vec; - for (int j = 0; j < input_shape.nbDims; j++) { - repeat_shape_vec.push_back(input_shape.d[j]); - if (j == dim) { - repeat_shape_vec.push_back(1); - } - } - auto expand = ctx->net->addShuffle(*self); - TORCHTRT_CHECK(expand, "Unable to create shuffle layer from node: " << *n); - auto repeat_shape_dims = util::toDims(repeat_shape_vec); - expand->setReshapeDimensions(repeat_shape_dims); - - // Expand on newly created singleton dimension - repeat_shape_dims.d[dim + 1] = repeats; - std::vector start_vec(repeat_shape_dims.nbDims, 0); - auto start_dims = util::toDims(start_vec); - - std::vector strides_vec(repeat_shape_dims.nbDims, 1); - strides_vec[dim + 1] = 0; - auto strides_dims = util::toDims(strides_vec); - - auto slice = ctx->net->addSlice(*expand->getOutput(0), start_dims, repeat_shape_dims, strides_dims); - - if (ctx->input_is_dynamic) { - auto start_tensor = tensor_to_const(ctx, torch::tensor(start_vec, torch::kInt32)); - - auto expand_output_shape = ctx->net->addShape(*expand->getOutput(0))->getOutput(0); - std::vector repeat_const_vec(repeat_shape_dims.nbDims, 1); - repeat_const_vec[dim + 1] = repeats; - auto repeat_const = tensor_to_const(ctx, torch::tensor(repeat_const_vec, torch::kInt32)); - auto repeat_shape_tensor = ctx->net->addElementWise(*expand_output_shape, *repeat_const, nvinfer1::ElementWiseOperation::kPROD)->getOutput(0); - - auto strides_tensor = tensor_to_const(ctx, torch::tensor(strides_vec, torch::kInt32)); - slice->setInput(1, *start_tensor); - slice->setInput(2, *repeat_shape_tensor); - slice->setInput(3, *strides_tensor); - } - - // Collapse repeated dimension back into desired dimension - std::vector collapse_shape_vec; - for (int k = 0; k < repeat_shape_dims.nbDims; k++) { - if (k == dim) { - int64_t collapse_dim = repeat_shape_dims.d[k] * repeat_shape_dims.d[++k]; - // Set dim size to -1 if repeat is being done on dynamic dim - collapse_dim = std::max(collapse_dim, (int64_t)-1); - collapse_shape_vec.push_back(collapse_dim); - } else { - collapse_shape_vec.push_back(repeat_shape_dims.d[k]); - } - } - auto collapse = ctx->net->addShuffle(*slice->getOutput(0)); - TORCHTRT_CHECK(collapse, "Unable to create shuffle layer from node: " << *n); - collapse->setReshapeDimensions(util::toDims(collapse_shape_vec)); - - collapse->setName(util::node_info(n).c_str()); - auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], collapse->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions()); - - return true; - - }}); + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto repeats = args[1].unwrapToScalar().to(); + + auto input_shape = self->getDimensions(); + + int dim; + if (args[2].IValue()->isNone()) { + dim = 0; + + // Flatten self tensor + int size; + if (ctx->input_is_dynamic) { + // Set size to -1 if input is dynamic + size = -1; + } else { + size = 1; + for (int i = 0; i < input_shape.nbDims; i++) { + size *= input_shape.d[i]; + } + } + auto flatten = ctx->net->addShuffle(*self); + TORCHTRT_CHECK(flatten, "Unable to create shuffle layer from node: " << *n); + flatten->setReshapeDimensions(util::toDims(std::vector({size}))); + self = flatten->getOutput(0); + input_shape = self->getDimensions(); + } else { + dim = args[2].unwrapToScalar().to(); + } + + if (ctx->input_is_dynamic) { + int dynamic_dims = 0; + for (int idx = 0; idx < input_shape.nbDims; idx++) { + if (input_shape.d[idx] == -1) { + dynamic_dims++; + } + } + + if (dynamic_dims > 1) { + TORCHTRT_THROW_ERROR( + "Repeat_interleave is currently not supported when target shape contains more than one dynamic dimension"); + } + } + + // Insert singleton dimension after desired repeat dimension + std::vector repeat_shape_vec; + for (int j = 0; j < input_shape.nbDims; j++) { + repeat_shape_vec.push_back(input_shape.d[j]); + if (j == dim) { + repeat_shape_vec.push_back(1); + } + } + auto expand = ctx->net->addShuffle(*self); + TORCHTRT_CHECK(expand, "Unable to create shuffle layer from node: " << *n); + auto repeat_shape_dims = util::toDims(repeat_shape_vec); + expand->setReshapeDimensions(repeat_shape_dims); + + // Expand on newly created singleton dimension + repeat_shape_dims.d[dim + 1] = repeats; + std::vector start_vec(repeat_shape_dims.nbDims, 0); + auto start_dims = util::toDims(start_vec); + + std::vector strides_vec(repeat_shape_dims.nbDims, 1); + strides_vec[dim + 1] = 0; + auto strides_dims = util::toDims(strides_vec); + + auto slice = ctx->net->addSlice(*expand->getOutput(0), start_dims, repeat_shape_dims, strides_dims); + + if (ctx->input_is_dynamic) { + auto start_tensor = tensor_to_const(ctx, torch::tensor(start_vec, torch::kInt32)); + + auto expand_output_shape = ctx->net->addShape(*expand->getOutput(0))->getOutput(0); + std::vector repeat_const_vec(repeat_shape_dims.nbDims, 1); + repeat_const_vec[dim + 1] = repeats; + auto repeat_const = tensor_to_const(ctx, torch::tensor(repeat_const_vec, torch::kInt32)); + auto repeat_shape_tensor = + ctx->net + ->addElementWise(*expand_output_shape, *repeat_const, nvinfer1::ElementWiseOperation::kPROD) + ->getOutput(0); + + auto strides_tensor = tensor_to_const(ctx, torch::tensor(strides_vec, torch::kInt32)); + slice->setInput(1, *start_tensor); + slice->setInput(2, *repeat_shape_tensor); + slice->setInput(3, *strides_tensor); + } + + // Collapse repeated dimension back into desired dimension + std::vector collapse_shape_vec; + for (int k = 0; k < repeat_shape_dims.nbDims; k++) { + if (k == dim) { + int64_t collapse_dim = repeat_shape_dims.d[k] * repeat_shape_dims.d[++k]; + // Set dim size to -1 if repeat is being done on dynamic dim + collapse_dim = std::max(collapse_dim, (int64_t)-1); + collapse_shape_vec.push_back(collapse_dim); + } else { + collapse_shape_vec.push_back(repeat_shape_dims.d[k]); + } + } + auto collapse = ctx->net->addShuffle(*slice->getOutput(0)); + TORCHTRT_CHECK(collapse, "Unable to create shuffle layer from node: " << *n); + collapse->setReshapeDimensions(util::toDims(collapse_shape_vec)); + + collapse->setName(util::node_info(n).c_str()); + auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], collapse->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions()); + + return true; + }}); } // namespace } // namespace impl