feat: Implement dynamic shape support for floordiv, NumToTensor, layernorm

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/converters/impl/layer_norm.cpp

@@ -20,8 +20,8 @@ auto layer_norm_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns()
 
       /* Layer_Norm normalizes over last N dimensions.
          normalizaed_shape could be (C,H,W), (H,W), or (W). */
-      auto normalized_shape = args[1].unwrapToIntList();
-      auto normalized_shape_vec = util::toVec(util::toDims(normalized_shape));
+      // This could be an IntList or ITensorList. We only need the size of this list.
+      auto normalized_shape = args[1].IValue()->toList();
 
       // Unwrap eps.
       auto eps = args[4].unwrapToDouble();
@@ -30,7 +30,7 @@ auto layer_norm_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns()
 
       // Set up  axis_ask for E[x].
       uint32_t axis_mask = 0;
-      for (size_t i = 0; i < normalized_shape_vec.size(); i++) {
+      for (size_t i = 0; i < normalized_shape.size(); i++) {
         axis_mask |= 1 << (shape.size() - i - 1);
       }
       LOG_DEBUG("Axis Mask for E[x]" << std::bitset<32>(axis_mask));

core/conversion/evaluators/aten.cpp

@@ -9,6 +9,7 @@
 #include "torch/csrc/jit/ir/ir.h"
 #include "torch/torch.h"
 
+#include "core/conversion/converters/converter_util.h"
 #include "core/conversion/evaluators/eval_macros.h"
 #include "core/conversion/evaluators/eval_util.h"
 #include "core/conversion/evaluators/evaluators.h"
@@ -677,6 +678,25 @@ auto aten_registrations TORCHTRT_UNUSED =
         .evaluator(
             {c10::Symbol::fromQualString("aten::floordiv"),
              [](ConversionCtx* ctx, const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+               // Dynamic version of aten::floordiv
+               if (args.at(n->input(0)).isITensor()) {
+                 if (args.at(n->input(1)).IValue()->isInt()) {
+                   auto int_tensor = scalar_to_tensor(args.at(n->input(1)).IValue()->toInt());
+                   auto int_itensor = converters::tensor_to_const(ctx, int_tensor, util::node_info(n) + "_constant");
+                   auto elementwise_layer = converters::add_elementwise(
+                       ctx,
+                       nvinfer1::ElementWiseOperation::kFLOOR_DIV,
+                       args.at(n->input(0)).ITensor(),
+                       int_itensor,
+                       util::node_info(n));
+                   auto output_tensor = elementwise_layer->getOutput(0);
+                   auto tensor_holder = TensorContainer();
+                   tensor_holder.hold_tensor(output_tensor);
+                   auto output_ivalue = c10::IValue(std::move(c10::make_intrusive<TensorContainer>(tensor_holder)));
+                   return output_ivalue;
+                 }
+               }
+               // Static version
                if (args.at(n->input(0)).IValue()->isInt()) {
                  auto a = args.at(n->input(0)).unwrapToInt();
                  auto b = args.at(n->input(1)).unwrapToInt();

core/conversion/evaluators/prim.cpp

@@ -32,6 +32,10 @@ auto prim_registrations =
         .evaluator(
             {torch::jit::prim::NumToTensor,
              [](ConversionCtx* ctx, const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+               // Dynamic version receives an ITensor here so pass that as output directly.
+               if (args.at(n->input(0)).isITensor()) {
+                 return args.at(n->input(0)).ITensor();
+               }
                return evaluators::scalar_to_tensor(args.at(n->input(0)).IValue()->toScalar());
              }})
         .evaluator(

tests/core/conversion/converters/BUILD

@@ -224,33 +224,33 @@ test_suite(
         ":test_div",
         ":test_einsum",
         ":test_expand",
+        ":test_index",
         ":test_instance_norm",
         ":test_interpolate",
-        ":test_index",
         ":test_layer_norm",
         ":test_linear",
         ":test_lstm_cell",
-        ":test_matrix_multiply",
         ":test_masked_fill",
+        ":test_matrix_multiply",
         ":test_max",
         ":test_normalize",
         ":test_pooling",
         ":test_reduce",
-        ":test_roll",
         ":test_replication_pad",
+        ":test_roll",
         ":test_scatter",
         ":test_select",
         ":test_shuffle",
+        ":test_slice",
         ":test_softmax",
+        ":test_split",
         ":test_squeeze",
         ":test_stack",
-        ":test_split",
-        ":test_slice",
         ":test_topk",
         ":test_unary",
-        ":test_unsqueeze",
         ":test_unbind",
         ":test_unpack",
+        ":test_unsqueeze",
         ":test_where",
     ],
 )