Merge pull request #2006 from pytorch/swin_unet

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

Author: narendasan

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"
@@ -298,20 +299,22 @@ auto aten_registrations TORCHTRT_UNUSED =
                } else {
                  auto dim = args.at(n->input(1)).unwrapToInt();
                  if (tensor_var.isITensor()) {
-                   if (ctx->input_is_dynamic) {
+                   auto tensor = tensor_var.ITensor();
+                   auto dims = util::toVec(tensor->getDimensions());
+                   auto nbDims = tensor->getDimensions().nbDims;
+                   if (dim < 0) {
+                     dim += nbDims;
+                   }
+                   // Check if selected dimension size is -1 else return static size
+                   if (ctx->input_is_dynamic && dims[dim] == -1) {
                      if (ctx->settings.allow_shape_tensors) {
                        return dynamic_size_layer(ctx, n, args);
                      } else {
                        LOG_WARNING(
                            "There may be undefined behavior using dynamic shape and aten::size without setting allow_shape_tensors");
                      }
                    }
-                   auto tensor = tensor_var.ITensor();
-                   auto dims = util::toVec(tensor->getDimensions());
-                   auto nbDims = tensor->getDimensions().nbDims;
-                   if (dim < 0) {
-                     dim += nbDims;
-                   }
+
                    return dims[dim];
                  } else if (tensor_var.IValue()->isTensor()) {
                    auto tensor = tensor_var.unwrapToTensor();
@@ -677,6 +680,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(