diff --git a/core/conversion/converters/converter_util.cpp b/core/conversion/converters/converter_util.cpp index d66c4514e9..36356d9db6 100644 --- a/core/conversion/converters/converter_util.cpp +++ b/core/conversion/converters/converter_util.cpp @@ -13,7 +13,8 @@ nvinfer1::ITensor* addPadding( nvinfer1::ITensor* tensor, int nDim, bool trailing, - bool use_zeros) { + bool use_zeros, + const std::string& name) { const auto dims = tensor->getDimensions(); if (dims.nbDims < nDim) { @@ -27,7 +28,11 @@ nvinfer1::ITensor* addPadding( TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer"); shuffle_layer->setReshapeDimensions(newDims); shuffle_layer->setZeroIsPlaceholder(use_zeros); - shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + ']').c_str()); + if (name.size()) { + shuffle_layer->setName(name.c_str()); + } else { + shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + ']').c_str()); + } return shuffle_layer->getOutput(0); } else { return tensor; @@ -40,7 +45,8 @@ nvinfer1::ITensor* addUnpadding( nvinfer1::ITensor* tensor, int nDim, bool trailing, - bool use_zeros) { + bool use_zeros, + const std::string& name) { const auto dims = tensor->getDimensions(); if (dims.nbDims > nDim) { auto newDims = dims; @@ -52,7 +58,11 @@ nvinfer1::ITensor* addUnpadding( TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer"); shuffle_layer->setReshapeDimensions(newDims); shuffle_layer->setZeroIsPlaceholder(use_zeros); - shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + "]").c_str()); + if (name.size()) { + shuffle_layer->setName(name.c_str()); + } else { + shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + ']').c_str()); + } return shuffle_layer->getOutput(0); } else { return tensor; diff --git a/core/conversion/converters/converter_util.h b/core/conversion/converters/converter_util.h index 6c4e9d53f0..245b93f498 100644 --- a/core/conversion/converters/converter_util.h +++ b/core/conversion/converters/converter_util.h @@ -22,7 +22,8 @@ nvinfer1::ITensor* addPadding( nvinfer1::ITensor* tensor, int nDim, bool trailing = true, - bool use_zeros = true); + bool use_zeros = true, + const std::string& name = ""); // If nDim < tensor size, adds shuffle layer to un-pad tensor (at the end if trailing) and returns (nDim-dimensional) // shuffle layer's output Otherwise, does nothing and passes tensor through. use _zeros controls whether we should be @@ -33,7 +34,8 @@ nvinfer1::ITensor* addUnpadding( nvinfer1::ITensor* tensor, int nDim, bool trailing = true, - bool use_zeros = true); + bool use_zeros = true, + const std::string& name = ""); // TODO: Change add_elementwise schema to output nvinfer1::ITensor* instead of nvinfer1::ILayer*, // for consistency with other utils. Need to change schema and usage in all calling contexts diff --git a/core/conversion/converters/impl/topk.cpp b/core/conversion/converters/impl/topk.cpp index 0cc26b92c1..0632f0f74f 100644 --- a/core/conversion/converters/impl/topk.cpp +++ b/core/conversion/converters/impl/topk.cpp @@ -39,15 +39,29 @@ auto topk_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patte LOG_DEBUG("Output topk reduce dim: " << dim); + // The topk layer requires at least 2 input dimensions + auto nbDims = self->getDimensions().nbDims; + if (nbDims == 1) { + self = addPadding(ctx, n, self, 2, true, true); + } + auto TopKOperation = largest ? (nvinfer1::TopKOperation::kMAX) : (nvinfer1::TopKOperation::kMIN); auto new_layer = ctx->net->addTopK(*self, TopKOperation, k, shiftDim); TORCHTRT_CHECK(new_layer, "Unable to create topk layer from node: " << *n); - auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0)); - auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1)); + auto values = new_layer->getOutput(0); + auto indices = new_layer->getOutput(1); + + // If we expanded the input, squeeze the outputs + if (nbDims == 1) { + values = addUnpadding(ctx, n, values, 1, true, true, util::node_info(n) + "_squeeze_values"); + indices = addUnpadding(ctx, n, indices, 1, true, true, util::node_info(n) + "_squeeze_indices"); + } + auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], values); + auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], indices); LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions()); LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions()); diff --git a/tests/core/conversion/converters/test_topk.cpp b/tests/core/conversion/converters/test_topk.cpp index 2e7e490ab7..9117fe1bb2 100644 --- a/tests/core/conversion/converters/test_topk.cpp +++ b/tests/core/conversion/converters/test_topk.cpp @@ -30,3 +30,29 @@ TEST(Converters, ATenTopKConvertsCorrectly) { ASSERT_TRUE( torch_tensorrt::tests::util::almostEqual(jit_results[1], trt_results[1].reshape_as(jit_results[1]), 8e-5)); } + +TEST(Converters, ATen1DTopKConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %1 : int = prim::Constant[value=20]() + %2 : int = prim::Constant[value=-1]() + %3 : bool = prim::Constant[value=1]() + %4 : bool = prim::Constant[value=1]() + %5 : Tensor, %6 : Tensor = aten::topk(%0, %1, %2, %3, %4) + return (%5, %6))IR"; + torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level( + torch_tensorrt::core::util::logging::LogLevel::kGRAPH); + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in = at::rand({100}, {at::kCUDA}); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 8e-5)); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[1], trt_results[1], 8e-5)); +}