Merge pull request #2115 from mfeliz-cruise/michael.feliz/dynamic_select_and_masked_fill

[feat] TS: Add support for dynamic select and masked_fill

Author: peri044

core/conversion/converters/impl/select.cpp

@@ -165,10 +165,7 @@ auto select_registrations TORCHTRT_UNUSED =
                  }
 
                  shuffle_layer->setReshapeDimensions(util::squeezeDims(
-                     out->getDimensions(),
-                     dim,
-                     ctx->input_is_dynamic,
-                     ctx->input_is_dynamic && (num_zero_dimensions > 0)));
+                     out->getDimensions(), dim, false, ctx->input_is_dynamic && (num_zero_dimensions > 0)));
                  shuffle_layer->setName(util::node_info(n).c_str());
                  out = shuffle_layer->getOutput(0);
                }
@@ -710,9 +707,8 @@ auto select_registrations TORCHTRT_UNUSED =
                auto val_t_dtype = util::TRTDataTypeToScalarType(self->getType());
 
                // Initialize contant tensor for fill with the inherited data type
-               auto val_t = tensor_to_const(
-                   ctx, torch::full(util::toVec(self->getDimensions()), val, {torch::dtype(val_t_dtype)}));
-
+               std::vector<int64_t> singleton_dims(self->getDimensions().nbDims, 1);
+               auto val_t = tensor_to_const(ctx, torch::full(singleton_dims, val, {torch::dtype(val_t_dtype)}));
                TORCHTRT_CHECK(
                    util::broadcastable(self->getDimensions(), mask->getDimensions(), /*multidirectional=*/false),
                    "Self and mask tensors are not broadcastable");

tests/core/conversion/converters/test_masked_fill.cpp

@@ -43,6 +43,40 @@ TEST(Converters, ATenMaskedFillZerosConvertsCorrectly) {
       torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }
 
+TEST(Converters, ATenMaskedFillZerosDynamicConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%x.1 : Tensor):
+      %44 : Device = prim::Constant[value="cuda"]()
+      %8 : bool = prim::Constant[value=0]()
+      %7 : None = prim::Constant()
+      %f32_dtype: int = prim::Constant[value=11]()
+      %1 : int = prim::Constant[value=0]() # bert.py:5:26
+      %2 : int = prim::Constant[value=1]() # bert.py:5:32
+      %33 : int = prim::Constant[value=2]() # bert.py:6:31
+      %3 : int[] = prim::ListConstruct(%1, %1, %2)
+      %9 : Tensor = aten::tensor(%3, %f32_dtype, %7, %8) # bert.py:5:11
+      %mask.1 : Tensor = aten::to(%9, %44, %7, %8, %8) # bert.py:5:11
+      %mask.2 : Tensor = trt::const(%mask.1)
+      %34 : Tensor = aten::masked_fill(%x.1, %mask.1, %33) # bert.py:6:11
+      return (%34, %mask.2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::zeros({1, 2, 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);
+  torch_tensorrt::core::lowering::passes::RemoveNOPs(g);
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0]));
+}
+
 TEST(Converters, ATenMaskedFillMixedTypesFloatIntConvertsCorrectly) {
   const auto graph = R"IR(
     graph(%x.1 : Tensor, %x.2 : Tensor):

tests/core/conversion/converters/test_select.cpp

@@ -32,6 +32,32 @@ TEST(Converters, ATenSelectIntConvertsCorrectly) {
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
 
+TEST(Converters, ATenSelectIntDynamicConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=0]()
+        %3 : Tensor = aten::select(%0, %2, %2)
+        return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::randint(1, 10, {5, 7, 9}, {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, ATenSelectIntDimIsOneConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor):