Merge pull request #1491 from mfeliz-cruise/michael.feliz/1d_topk

[feat] Support 1D topk

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;

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

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());
 

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::Graph>();
+  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));
+}